Support S3 server side encryption headers for Write and Copy

Adds support to s3blob for setting AWS S3 server side encryption headers when making requests that require such headers. The additional settings can be specified with the `ssetype` and `kmskeyid` URL params, similar to the other configuration settings
google · Nov 21, 2023 · 570ae46 · 570ae46
1 parent 6d5d289
commit 570ae46
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 5 deletions.
diff --git a/aws/aws.go b/aws/aws.go
@@ -103,6 +103,8 @@ func ConfigFromURLParams(q url.Values) (*aws.Config, error) {
 			}
 			cfg.S3ForcePathStyle = aws.Bool(b)
 		case "awssdk":
+		case "ssetype":
+		case "kmskeyid":
 			// ignore, should be handled before this
 		default:
 			return nil, fmt.Errorf("unknown query parameter %q", param)
@@ -198,6 +200,8 @@ func V2ConfigFromURLParams(ctx context.Context, q url.Values) (awsv2.Config, err
 		case "profile":
 			opts = append(opts, awsv2cfg.WithSharedConfigProfile(value))
 		case "awssdk":
+		case "ssetype":
+		case "kmskeyid":
 			// ignore, should be handled before this
 		default:
 			return awsv2.Config{}, fmt.Errorf("unknown query parameter %q", param)

diff --git a/blob/s3blob/s3blob.go b/blob/s3blob/s3blob.go
@@ -152,6 +152,10 @@ func (o *URLOpener) OpenBucketURL(ctx context.Context, u *url.URL) (*blob.Bucket
 			return nil, fmt.Errorf("open bucket %v: %v", u, err)
 		}
 		clientV2 := s3v2.NewFromConfig(cfg)
+
+		o.Options.EncryptionType = u.Query().Get("ssetype")
+		o.Options.KMSEncryptionId = u.Query().Get("kmskeyid")
+
 		return OpenBucketV2(ctx, clientV2, u.Host, &o.Options)
 	}
 	configProvider := &gcaws.ConfigOverrider{
@@ -162,6 +166,10 @@ func (o *URLOpener) OpenBucketURL(ctx context.Context, u *url.URL) (*blob.Bucket
 		return nil, fmt.Errorf("open bucket %v: %v", u, err)
 	}
 	configProvider.Configs = append(configProvider.Configs, overrideCfg)
+
+	o.Options.EncryptionType = u.Query().Get("ssetype")
+	o.Options.KMSEncryptionId = u.Query().Get("kmskeyid")
+
 	return OpenBucket(ctx, configProvider, u.Host, &o.Options)
 }
 
@@ -171,6 +179,10 @@ type Options struct {
 	// Some S3-compatible services (like CEPH) do not currently support
 	// ListObjectsV2.
 	UseLegacyList bool
+
+	EncryptionType string
+
+	KMSEncryptionId string
 }
 
 // openBucket returns an S3 Bucket.
@@ -193,11 +205,13 @@ func openBucket(ctx context.Context, useV2 bool, sess client.ConfigProvider, cli
 		client = s3.New(sess)
 	}
 	return &bucket{
-		useV2:         useV2,
-		name:          bucketName,
-		client:        client,
-		clientV2:      clientV2,
-		useLegacyList: opts.UseLegacyList,
+		useV2:          useV2,
+		name:           bucketName,
+		client:         client,
+		clientV2:       clientV2,
+		useLegacyList:  opts.UseLegacyList,
+		kmsKeyId:       opts.KMSEncryptionId,
+		encryptionType: opts.EncryptionType,
 	}, nil
 }
 
@@ -365,6 +379,9 @@ type bucket struct {
 	client        *s3.S3
 	clientV2      *s3v2.Client
 	useLegacyList bool
+
+	encryptionType string
+	kmsKeyId       string
 }
 
 func (b *bucket) Close() error {
@@ -973,6 +990,12 @@ func (b *bucket) NewTypedWriter(ctx context.Context, key string, contentType str
 		if len(opts.ContentMD5) > 0 {
 			reqV2.ContentMD5 = aws.String(base64.StdEncoding.EncodeToString(opts.ContentMD5))
 		}
+		if b.encryptionType != "" {
+			reqV2.ServerSideEncryption = typesv2.ServerSideEncryption(b.encryptionType)
+		}
+		if b.kmsKeyId != "" {
+			reqV2.SSEKMSKeyId = aws.String(b.kmsKeyId)
+		}
 		if opts.BeforeWrite != nil {
 			asFunc := func(i interface{}) bool {
 				// Note that since the Go CDK Blob
@@ -1046,6 +1069,12 @@ func (b *bucket) NewTypedWriter(ctx context.Context, key string, contentType str
 		if len(opts.ContentMD5) > 0 {
 			req.ContentMD5 = aws.String(base64.StdEncoding.EncodeToString(opts.ContentMD5))
 		}
+		if b.encryptionType != "" {
+			req.ServerSideEncryption = aws.String(b.encryptionType)
+		}
+		if b.kmsKeyId != "" {
+			req.SSEKMSKeyId = aws.String(b.kmsKeyId)
+		}
 		if opts.BeforeWrite != nil {
 			asFunc := func(i interface{}) bool {
 				pu, ok := i.(**s3manager.Uploader)
@@ -1083,6 +1112,12 @@ func (b *bucket) Copy(ctx context.Context, dstKey, srcKey string, opts *driver.C
 			CopySource: aws.String(b.name + "/" + srcKey),
 			Key:        aws.String(dstKey),
 		}
+		if b.encryptionType != "" {
+			input.ServerSideEncryption = typesv2.ServerSideEncryption(b.encryptionType)
+		}
+		if b.kmsKeyId != "" {
+			input.SSEKMSKeyId = aws.String(b.kmsKeyId)
+		}
 		if opts.BeforeCopy != nil {
 			asFunc := func(i interface{}) bool {
 				switch v := i.(type) {
@@ -1104,6 +1139,12 @@ func (b *bucket) Copy(ctx context.Context, dstKey, srcKey string, opts *driver.C
 			CopySource: aws.String(b.name + "/" + srcKey),
 			Key:        aws.String(dstKey),
 		}
+		if b.encryptionType != "" {
+			input.ServerSideEncryption = aws.String(b.encryptionType)
+		}
+		if b.kmsKeyId != "" {
+			input.SSEKMSKeyId = aws.String(b.kmsKeyId)
+		}
 		if opts.BeforeCopy != nil {
 			asFunc := func(i interface{}) bool {
 				switch v := i.(type) {

diff --git a/blob/s3blob/s3blob_test.go b/blob/s3blob/s3blob_test.go
@@ -466,6 +466,8 @@ func TestOpenBucketFromURL(t *testing.T) {
 		{"s3://mybucket?profile=main&region=us-west-1", false},
 		// OK, use V2.
 		{"s3://mybucket?awssdk=v2", false},
+		// OK, use KMS Server Side Encryption
+		{"s3://mybucket?ssetype=aws:kms&kmskeyid=arn:aws:us-east-1:12345:key/1-a-2-b", false},
 		// Invalid parameter together with a valid one.
 		{"s3://mybucket?profile=main&param=value", true},
 		// Invalid parameter.