blobstore: Add SSE in multi-part upload (salesforce#112)

Author: sandeepvinayak

blob/blob-ali/src/main/java/com/salesforce/multicloudj/blob/ali/AliBlobStore.java

@@ -361,7 +361,7 @@ protected ListBlobsPageResponse doListPage(ListBlobsPageRequest request) {
     protected MultipartUpload doInitiateMultipartUpload(final MultipartUploadRequest request){
         InitiateMultipartUploadRequest initiateMultipartUploadRequest = transformer.toInitiateMultipartUploadRequest(request);
         InitiateMultipartUploadResult initiateMultipartUploadResult = ossClient.initiateMultipartUpload(initiateMultipartUploadRequest);
-        return transformer.toMultipartUpload(initiateMultipartUploadResult, request.getMetadata());
+        return transformer.toMultipartUpload(initiateMultipartUploadResult, request.getMetadata(), request.getKmsKeyId());
     }
 
     /**

blob/blob-ali/src/main/java/com/salesforce/multicloudj/blob/ali/AliTransformer.java

@@ -202,15 +202,22 @@ public BlobMetadata toBlobMetadata(String key, ObjectMetadata metadata) {
     public InitiateMultipartUploadRequest toInitiateMultipartUploadRequest(MultipartUploadRequest request) {
         ObjectMetadata metadata = new ObjectMetadata();
         metadata.setUserMetadata(request.getMetadata());
+
+        if (request.getKmsKeyId() != null && !request.getKmsKeyId().isEmpty()) {
+            metadata.setServerSideEncryption(ObjectMetadata.KMS_SERVER_SIDE_ENCRYPTION);
+            metadata.setHeader(OSSHeaders.OSS_SERVER_SIDE_ENCRYPTION_KEY_ID, request.getKmsKeyId());
+        }
+
         return new InitiateMultipartUploadRequest(getBucket(), request.getKey(), metadata);
     }
 
-    public MultipartUpload toMultipartUpload(InitiateMultipartUploadResult initiateMultipartUploadResult, Map<String, String> metadata) {
+    public MultipartUpload toMultipartUpload(InitiateMultipartUploadResult initiateMultipartUploadResult, Map<String, String> metadata, String kmsKeyId) {
         return MultipartUpload.builder()
                 .bucket(initiateMultipartUploadResult.getBucketName())
                 .key(initiateMultipartUploadResult.getKey())
                 .id(initiateMultipartUploadResult.getUploadId())
                 .metadata(metadata)
+                .kmsKeyId(kmsKeyId)
                 .build();
     }
 

blob/blob-ali/src/test/java/com/salesforce/multicloudj/blob/ali/AliBlobStoreTest.java

@@ -542,6 +542,36 @@ void testDoInitiateMultipartUpload() {
         assertEquals(metadata, actualRequest.getObjectMetadata().getUserMetadata());
     }
 
+    @Test
+    void testDoInitiateMultipartUploadWithKms() {
+        InitiateMultipartUploadResult mockResponse = mock(InitiateMultipartUploadResult.class);
+        doReturn("bucket-1").when(mockResponse).getBucketName();
+        doReturn("object-1").when(mockResponse).getKey();
+        doReturn("mpu-id").when(mockResponse).getUploadId();
+        when(mockOssClient.initiateMultipartUpload((InitiateMultipartUploadRequest) any())).thenReturn(mockResponse);
+        Map<String, String> metadata = Map.of("key-1", "value-1");
+        String kmsKeyId = "test-kms-key-id";
+        MultipartUploadRequest request = new MultipartUploadRequest.Builder()
+                .withKey("object-1")
+                .withMetadata(metadata)
+                .withKmsKeyId(kmsKeyId)
+                .build();
+
+        MultipartUpload response = ali.initiateMultipartUpload(request);
+
+        ArgumentCaptor<InitiateMultipartUploadRequest> requestCaptor = ArgumentCaptor.forClass(InitiateMultipartUploadRequest.class);
+        verify(mockOssClient, times(1)).initiateMultipartUpload(requestCaptor.capture());
+        InitiateMultipartUploadRequest actualRequest = requestCaptor.getValue();
+        assertEquals("object-1", actualRequest.getKey());
+        assertEquals("bucket-1", actualRequest.getBucketName());
+        assertEquals(metadata, actualRequest.getObjectMetadata().getUserMetadata());
+        assertEquals(ObjectMetadata.KMS_SERVER_SIDE_ENCRYPTION, actualRequest.getObjectMetadata().getServerSideEncryption());
+        assertEquals(kmsKeyId, actualRequest.getObjectMetadata().getRawMetadata().get(OSSHeaders.OSS_SERVER_SIDE_ENCRYPTION_KEY_ID));
+
+        // Verify the response has KMS key
+        assertEquals(kmsKeyId, response.getKmsKeyId());
+    }
+
     @Test
     void testDoUploadMultipartPart() {
         UploadPartResult mockResponse = mock(UploadPartResult.class);

blob/blob-ali/src/test/java/com/salesforce/multicloudj/blob/ali/AliTransformerTest.java

@@ -350,14 +350,32 @@ void testToMultipartUpload() {
         doReturn("uploadId").when(initiateMultipartUploadResult).getUploadId();
         Map<String, String> metadata = Map.of("key1", "value1", "key2", "value2");
 
-        var actual = transformer.toMultipartUpload(initiateMultipartUploadResult, metadata);
+        var actual = transformer.toMultipartUpload(initiateMultipartUploadResult, metadata, null);
 
         assertEquals(BUCKET, actual.getBucket());
         assertEquals("key", actual.getKey());
         assertEquals("uploadId", actual.getId());
         assertEquals(metadata, actual.getMetadata());
     }
 
+    @Test
+    void testToMultipartUploadWithKms() {
+        InitiateMultipartUploadResult initiateMultipartUploadResult = mock(InitiateMultipartUploadResult.class);
+        doReturn(BUCKET).when(initiateMultipartUploadResult).getBucketName();
+        doReturn("key").when(initiateMultipartUploadResult).getKey();
+        doReturn("uploadId").when(initiateMultipartUploadResult).getUploadId();
+        Map<String, String> metadata = Map.of("key1", "value1", "key2", "value2");
+        String kmsKeyId = "test-kms-key-id";
+
+        var actual = transformer.toMultipartUpload(initiateMultipartUploadResult, metadata, kmsKeyId);
+
+        assertEquals(BUCKET, actual.getBucket());
+        assertEquals("key", actual.getKey());
+        assertEquals("uploadId", actual.getId());
+        assertEquals(metadata, actual.getMetadata());
+        assertEquals(kmsKeyId, actual.getKmsKeyId());
+    }
+
     @Test
     void testToUploadPartRequest() {
         Map<String, String> metadata = Map.of("key1", "value1", "key2", "value2");

blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/AwsBlobStore.java

@@ -369,6 +369,7 @@ protected MultipartUpload doInitiateMultipartUpload(final MultipartUploadRequest
                 .key(createMultipartUploadResponse.key())
                 .id(createMultipartUploadResponse.uploadId())
                 .metadata(request.getMetadata())
+                .kmsKeyId(request.getKmsKeyId())
                 .build();
     }
 

blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/AwsTransformer.java

@@ -150,7 +150,7 @@ public PutObjectRequest toRequest(UploadRequest request) {
                 .tagging(Tagging.builder().tagSet(tags).build());
 
         if (request.getKmsKeyId() != null && !request.getKmsKeyId().isEmpty()) {
-            builder.serverSideEncryption("aws:kms")
+            builder.serverSideEncryption(ServerSideEncryption.AWS_KMS)
                    .ssekmsKeyId(request.getKmsKeyId());
         }
 
@@ -292,11 +292,17 @@ byte[] eTagToMD5(String eTag) {
     }
 
     public CreateMultipartUploadRequest toCreateMultipartUploadRequest(MultipartUploadRequest request) {
-        return CreateMultipartUploadRequest.builder()
+        CreateMultipartUploadRequest.Builder builder = CreateMultipartUploadRequest.builder()
                 .bucket(getBucket())
                 .key(request.getKey())
-                .metadata(request.getMetadata())
-                .build();
+                .metadata(request.getMetadata());
+
+        if (request.getKmsKeyId() != null && !request.getKmsKeyId().isEmpty()) {
+            builder.serverSideEncryption(ServerSideEncryption.AWS_KMS)
+                   .ssekmsKeyId(request.getKmsKeyId());
+        }
+
+        return builder.build();
     }
 
     public UploadPartRequest toUploadPartRequest(MultipartUpload mpu, MultipartPart mpp) {

blob/blob-aws/src/test/java/com/salesforce/multicloudj/blob/aws/AwsBlobStoreTest.java

@@ -70,6 +70,7 @@
 import software.amazon.awssdk.services.s3.model.PutObjectTaggingResponse;
 import software.amazon.awssdk.services.s3.model.S3Exception;
 import software.amazon.awssdk.services.s3.model.S3Object;
+import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
 import software.amazon.awssdk.services.s3.model.Tag;
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
@@ -687,6 +688,40 @@ void testDoInitiateMultipartUpload() {
         assertEquals("mpu-id", response.getId());
     }
 
+    @Test
+    void testDoInitiateMultipartUploadWithKms() {
+        CreateMultipartUploadResponse mockResponse = mock(CreateMultipartUploadResponse.class);
+        doReturn("bucket-1").when(mockResponse).bucket();
+        doReturn("object-1").when(mockResponse).key();
+        doReturn("mpu-id").when(mockResponse).uploadId();
+        when(mockS3Client.createMultipartUpload((CreateMultipartUploadRequest) any())).thenReturn(mockResponse);
+        Map<String, String> metadata = Map.of("key-1", "value-1");
+        String kmsKeyId = "arn:aws:kms:us-east-1:123456789012:key/12345678-1234-1234-1234-123456789012";
+        MultipartUploadRequest request = new MultipartUploadRequest.Builder()
+                .withKey("object-1")
+                .withMetadata(metadata)
+                .withKmsKeyId(kmsKeyId)
+                .build();
+
+        MultipartUpload response = aws.initiateMultipartUpload(request);
+
+        // Verify the request is mapped to the SDK with KMS encryption
+        ArgumentCaptor<CreateMultipartUploadRequest> requestCaptor = ArgumentCaptor.forClass(CreateMultipartUploadRequest.class);
+        verify(mockS3Client, times(1)).createMultipartUpload(requestCaptor.capture());
+        CreateMultipartUploadRequest actualRequest = requestCaptor.getValue();
+        assertEquals("object-1", actualRequest.key());
+        assertEquals("bucket-1", actualRequest.bucket());
+        assertEquals(metadata, actualRequest.metadata());
+        assertEquals(ServerSideEncryption.AWS_KMS, actualRequest.serverSideEncryption());
+        assertEquals(kmsKeyId, actualRequest.ssekmsKeyId());
+
+        // Verify the response is mapped back properly with KMS key
+        assertEquals("object-1", response.getKey());
+        assertEquals("bucket-1", response.getBucket());
+        assertEquals("mpu-id", response.getId());
+        assertEquals(kmsKeyId, response.getKmsKeyId());
+    }
+
     @Test
     void testDoUploadMultipartPart() {
         UploadPartResponse mockResponse = mock(UploadPartResponse.class);

blob/blob-aws/src/test/resources/mappings/delete-4lfap3itgx.json

@@ -0,0 +1,20 @@
+{
+  "id" : "db22fc9f-3cea-41b0-a18f-6dc0b9985c0c",
+  "name" : "chameleon-jcloud_conformance-tests_multipart-withkms",
+  "request" : {
+    "url" : "/chameleon-jcloud/conformance-tests/multipart-withKms",
+    "method" : "DELETE"
+  },
+  "response" : {
+    "status" : 204,
+    "headers" : {
+      "Server" : "AmazonS3",
+      "x-amz-request-id" : "QEXVC59YXBQPMRTM",
+      "x-amz-id-2" : "+N1BJ/b+RhzccdF12ZsqyJoepI5nCc2kEh7xEBuVFq3VzqwJ2NjDrAWrwP0LXcOqz6FKCgtYyuA=",
+      "Date" : "Thu, 30 Oct 2025 03:16:06 GMT"
+    }
+  },
+  "uuid" : "db22fc9f-3cea-41b0-a18f-6dc0b9985c0c",
+  "persistent" : true,
+  "insertionIndex" : 541
+}

blob/blob-aws/src/test/resources/mappings/delete-rxwqfxe3lf.json

@@ -0,0 +1,20 @@
+{
+  "id" : "1771bde0-65c2-481e-9256-2f4318863074",
+  "name" : "chameleon-jcloud_conformance-tests_multipart-withkms",
+  "request" : {
+    "url" : "/chameleon-jcloud/conformance-tests/multipart-withKms?uploadId=HP3yS0r4PKOhIfexxxwxciB6pikz.Mk0xzdlaZdnwSCFX9NNc2EUsYmNL_RS77aRIcK3.taL6VBdjpgwVFpkb6CKVP1vWsn2wTNKAtm_P3c9G_7smpYdQerWRyT4seAKvmaaeVOi4HbR2Pidr_sEPkI3WwVrVc59Yqu7IuVfzO9FW_VUXOBvYnk5qsKuQSTUh6p5UmvZXrYTOGtJHq0C9w--",
+    "method" : "DELETE"
+  },
+  "response" : {
+    "status" : 204,
+    "headers" : {
+      "Server" : "AmazonS3",
+      "x-amz-request-id" : "QEXKDM9VY5PD64D6",
+      "x-amz-id-2" : "5rVY8XuJLUAYgfYr1eRG8hujkZlDHlnLtKiMUW6uslOvv955eC4sb3QkF7EDLBB3BJV2Q1OuF+M=",
+      "Date" : "Thu, 30 Oct 2025 03:16:06 GMT"
+    }
+  },
+  "uuid" : "1771bde0-65c2-481e-9256-2f4318863074",
+  "persistent" : true,
+  "insertionIndex" : 540
+}

blob/blob-aws/src/test/resources/mappings/get-ukvs6x01b5.json

@@ -0,0 +1,22 @@
+{
+  "id" : "19ea6fb7-e8e3-4194-a40a-83a5373fd128",
+  "name" : "chameleon-jcloud_conformance-tests_multipart-withkms",
+  "request" : {
+    "url" : "/chameleon-jcloud/conformance-tests/multipart-withKms?uploadId=HP3yS0r4PKOhIfexxxwxciB6pikz.Mk0xzdlaZdnwSCFX9NNc2EUsYmNL_RS77aRIcK3.taL6VBdjpgwVFpkb6CKVP1vWsn2wTNKAtm_P3c9G_7smpYdQerWRyT4seAKvmaaeVOi4HbR2Pidr_sEPkI3WwVrVc59Yqu7IuVfzO9FW_VUXOBvYnk5qsKuQSTUh6p5UmvZXrYTOGtJHq0C9w--",
+    "method" : "GET"
+  },
+  "response" : {
+    "status" : 200,
+    "body" : "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n<ListPartsResult xmlns=\"http://s3.amazonaws.com/doc/2006-03-01/\"><Bucket>chameleon-jcloud</Bucket><Key>conformance-tests/multipart-withKms</Key><UploadId>HP3yS0r4PKOhIfexxxwxciB6pikz.Mk0xzdlaZdnwSCFX9NNc2EUsYmNL_RS77aRIcK3.taL6VBdjpgwVFpkb6CKVP1vWsn2wTNKAtm_P3c9G_7smpYdQerWRyT4seAKvmaaeVOi4HbR2Pidr_sEPkI3WwVrVc59Yqu7IuVfzO9FW_VUXOBvYnk5qsKuQSTUh6p5UmvZXrYTOGtJHq0C9w--</UploadId><Initiator><ID>arn:aws:sts::654654370895:assumed-role/PCSKAdministratorAccessRole/PCSK-sandeeppal@d2999721-cb95-450c-8455-96b96fdcdc84</ID><DisplayName>PCSKAdministratorAccessRole/PCSK-sandeeppal@d2999721-cb95-450c-8455-96b96fdcdc84</DisplayName></Initiator><Owner><ID>b6eaab2af32ba61bce00b8c9aceaa1c649844388ec2c3827bbd3e2b06f798cf1</ID></Owner><StorageClass>STANDARD</StorageClass><PartNumberMarker>0</PartNumberMarker><NextPartNumberMarker>2</NextPartNumberMarker><MaxParts>1000</MaxParts><IsTruncated>false</IsTruncated><Part><PartNumber>1</PartNumber><LastModified>2025-10-30T03:16:01.000Z</LastModified><ETag>&quot;ddbbf3e1d34182a6a0233f9e3b607419&quot;</ETag><Size>5242880</Size></Part><Part><PartNumber>2</PartNumber><LastModified>2025-10-30T03:16:03.000Z</LastModified><ETag>&quot;8582535128d55f015eb60cc3d05b3355&quot;</ETag><Size>5242880</Size></Part></ListPartsResult>",
+    "headers" : {
+      "Server" : "AmazonS3",
+      "x-amz-request-id" : "D13SSBSNEPX0SRBQ",
+      "x-amz-id-2" : "mERBuJQtkoMlMy1yy6m4LemDsZlQzIpXwywQrfWgwDWaatzKG/g6HfyWsFlqOjZVD9X7bqgQi7o=",
+      "Date" : "Thu, 30 Oct 2025 03:16:04 GMT",
+      "Content-Type" : "application/xml"
+    }
+  },
+  "uuid" : "19ea6fb7-e8e3-4194-a40a-83a5373fd128",
+  "persistent" : true,
+  "insertionIndex" : 544
+}