fix(storage): preserve custom host in S3 presigned URLs

Author: You

drivers/s3/util.go

@@ -14,6 +14,7 @@ import (
 	awsconfig "github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	awss3 "github.com/aws/aws-sdk-go-v2/service/s3"
+	smithyendpoints "github.com/aws/smithy-go/endpoints"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -43,7 +44,10 @@ func (d *S3) initSession(ctx context.Context) error {
 	d.client = d.newClient(d.Endpoint)
 	d.linkClient = d.client
 	if d.CustomHost != "" && d.EnableCustomHostPresign {
-		d.linkClient = d.newClient(normalizeEndpoint(d.CustomHost))
+		d.linkClient, err = d.newCustomHostClient()
+		if err != nil {
+			return err
+		}
 	}
 	return nil
 }
@@ -57,6 +61,33 @@ func (d *S3) newClient(endpoint string) *awss3.Client {
 	})
 }
 
+type customHostEndpointResolver struct {
+	endpoint url.URL
+}
+
+func (r customHostEndpointResolver) ResolveEndpoint(
+	_ context.Context,
+	params awss3.EndpointParameters,
+) (smithyendpoints.Endpoint, error) {
+	endpoint := r.endpoint
+	if aws.ToBool(params.ForcePathStyle) && params.Bucket != nil {
+		endpoint = *endpoint.JoinPath(*params.Bucket)
+	}
+	return smithyendpoints.Endpoint{URI: endpoint}, nil
+}
+
+func (d *S3) newCustomHostClient() (*awss3.Client, error) {
+	endpoint, err := url.Parse(normalizeEndpoint(d.CustomHost))
+	if err != nil {
+		return nil, err
+	}
+	client := awss3.NewFromConfig(d.awsConfig, func(options *awss3.Options) {
+		options.EndpointResolverV2 = customHostEndpointResolver{endpoint: *endpoint}
+		options.UsePathStyle = d.ForcePathStyle
+	})
+	return client, nil
+}
+
 func normalizeEndpoint(endpoint string) string {
 	if strings.Contains(endpoint, "://") {
 		return endpoint

drivers/s3/util_test.go

@@ -1,8 +1,11 @@
 package s3
 
 import (
+	"context"
+	"net/url"
 	"testing"
 
+	"github.com/AlliotTech/openalist/internal/model"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -51,3 +54,57 @@ func TestCustomObjectURL(t *testing.T) {
 		})
 	}
 }
+
+func TestCustomHostPresign(t *testing.T) {
+	tests := []struct {
+		name           string
+		forcePathStyle bool
+		removeBucket   bool
+		wantHost       string
+		wantPath       string
+	}{
+		{
+			name:     "virtual host style keeps custom host",
+			wantHost: "cdn.example.com",
+			wantPath: "/base/dir/file.txt",
+		},
+		{
+			name:           "path style keeps bucket in path",
+			forcePathStyle: true,
+			wantHost:       "cdn.example.com",
+			wantPath:       "/base/bucket/dir/file.txt",
+		},
+		{
+			name:           "remove bucket keeps custom host",
+			forcePathStyle: true,
+			removeBucket:   true,
+			wantHost:       "cdn.example.com",
+			wantPath:       "/base/dir/file.txt",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			d := S3{Addition: Addition{
+				AccessKeyID:             "access-key",
+				SecretAccessKey:         "secret-key",
+				Region:                  "us-east-1",
+				Bucket:                  "bucket",
+				CustomHost:              "https://cdn.example.com/base",
+				EnableCustomHostPresign: true,
+				ForcePathStyle:          tt.forcePathStyle,
+				RemoveBucket:            tt.removeBucket,
+				SignURLExpire:           1,
+			}}
+			require.NoError(t, d.Init(context.Background()))
+
+			link, err := d.Link(context.Background(), &model.Object{Path: "/dir/file.txt"}, model.LinkArgs{})
+			require.NoError(t, err)
+			parsed, err := url.Parse(link.URL)
+			require.NoError(t, err)
+			assert.Equal(t, tt.wantHost, parsed.Host)
+			assert.Equal(t, tt.wantPath, parsed.Path)
+			assert.NotEmpty(t, parsed.Query().Get("X-Amz-Signature"))
+		})
+	}
+}

go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.19.14
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.22.13
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.99.0
+	github.com/aws/smithy-go v1.24.3
 	github.com/blevesearch/bleve/v2 v2.6.0
 	github.com/caarlos0/env/v9 v9.0.0
 	github.com/charmbracelet/bubbles v1.0.0
@@ -105,7 +106,6 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/sso v1.30.15 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.19 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.41.10 // indirect
-	github.com/aws/smithy-go v1.24.3 // indirect
 	github.com/blevesearch/zapx/v17 v17.1.2 // indirect
 	github.com/bytedance/gopkg v0.1.3 // indirect
 	github.com/charmbracelet/colorprofile v0.4.1 // indirect