added kms parameters to specify kms key for cross-account tasks

awslabs · May 17, 2024 · 9d59768 · 9d59768
1 parent 2d4d922
commit 9d59768
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 8 deletions.
diff --git a/api.go b/api.go
@@ -137,6 +137,19 @@ func WithTarFormat(format string) func(options *S3TarS3Options) {
 	}
 }
 
+func WithKMS(kmsKeyID, sseAlgo string) func(options *S3TarS3Options) {
+	return func(opts *S3TarS3Options) {
+		if kmsKeyID == "" {
+			return
+		}
+		if sseAlgo != "aws:kms" && sseAlgo != "AES256" && sseAlgo != "aws:kms:dsse" {
+			Fatalf(context.TODO(), "unknown sseAlgo")
+		}
+		opts.KMSKeyID = kmsKeyID
+		opts.SSEAlgo = types.ServerSideEncryption(sseAlgo)
+	}
+}
+
 func checkCreateArgs(opts *S3TarS3Options) error {
 	if opts.SrcBucket == "" && opts.SrcManifest == "" {
 		return fmt.Errorf("src bucket or src manifest required")

diff --git a/cmd/s3tar/main.go b/cmd/s3tar/main.go
@@ -67,6 +67,8 @@ func run(args []string) error {
 	var userPartMaxSize int64
 	var awsProfile string
 	var tagSetInput string
+	var kmsKeyID string
+	var sseAlgo string
 
 	var tagSet types.Tagging
 	var err error
@@ -238,6 +240,16 @@ func run(args []string) error {
 				Usage:       "pass a tag value following awscli syntax: --tagging='{\"TagSet\": [{ \"Key\": \"transition-to\", \"Value\": \"GDA\" }]}'",
 				Destination: &tagSetInput,
 			},
+			&cli.StringFlag{
+				Name:        "sse-kms-key-id",
+				Usage:       "",
+				Destination: &kmsKeyID,
+			},
+			&cli.StringFlag{
+				Name:        "sse-algo",
+				Usage:       "aws:kms or AES256",
+				Destination: &sseAlgo,
+			},
 		},
 		Action: func(cCtx *cli.Context) error {
 			logLevel := parseLogLevel(cCtx.Count("verbose"))
@@ -340,7 +352,8 @@ func run(args []string) error {
 						s3opts.DstPrefix = filepath.Dir(s3opts.DstKey)
 						err := archiveClient.CreateFromList(ctx, archive, s3opts,
 							s3tar.WithStorageClass(storageClass),
-							s3tar.WithTarFormat(tarFormat))
+							s3tar.WithTarFormat(tarFormat),
+							s3tar.WithKMS(kmsKeyID, sseAlgo))
 						if err != nil {
 							return err
 						}
@@ -349,7 +362,8 @@ func run(args []string) error {
 				} else {
 					return archiveClient.CreateFromList(ctx, objectList, s3opts,
 						s3tar.WithStorageClass(storageClass),
-						s3tar.WithTarFormat(tarFormat))
+						s3tar.WithTarFormat(tarFormat),
+						s3tar.WithKMS(kmsKeyID, sseAlgo))
 				}
 
 			} else if extract {

diff --git a/mem_concat.go b/mem_concat.go
@@ -51,12 +51,14 @@ func buildInMemoryConcat(ctx context.Context, client *s3.Client, objectList []*S
 
 		// create MPU
 		mpu, err := client.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
-			Bucket:            &opts.DstBucket,
-			Key:               &opts.DstKey,
-			StorageClass:      opts.storageClass,
-			ChecksumAlgorithm: types.ChecksumAlgorithmSha256,
-			Tagging:           &tags,
-			ACL:               types.ObjectCannedACLBucketOwnerFullControl,
+			Bucket:               &opts.DstBucket,
+			Key:                  &opts.DstKey,
+			StorageClass:         opts.storageClass,
+			ChecksumAlgorithm:    types.ChecksumAlgorithmSha256,
+			Tagging:              &tags,
+			ACL:                  types.ObjectCannedACLBucketOwnerFullControl,
+			SSEKMSKeyId:          &opts.KMSKeyID,
+			ServerSideEncryption: opts.SSEAlgo,
 		})
 		if err != nil {
 			Errorf(ctx, "unable to create multipart")

diff --git a/utils.go b/utils.go
@@ -57,6 +57,8 @@ type S3TarS3Options struct {
 	UrlDecode          bool
 	UserMaxPartSize    int64
 	ObjectTags         types.Tagging
+	KMSKeyID           string
+	SSEAlgo            types.ServerSideEncryption
 }
 
 func TagsToUrlEncodedString(tagging types.Tagging) string {