feat: Add missed KMS support in msq durable storage and exports

docs/development/extensions-core/s3.md

@@ -153,3 +153,7 @@ You can enable [server-side encryption](https://docs.aws.amazon.com/AmazonS3/lat
 * s3: [Server-side encryption with S3-managed encryption keys](https://docs.aws.amazon.com/AmazonS3/latest/dev/UsingServerSideEncryption)
 * kms: [Server-side encryption with AWS KMS–Managed Keys](https://docs.aws.amazon.com/AmazonS3/latest/dev/UsingKMSEncryption)
 * custom: [Server-side encryption with Customer-Provided Encryption Keys](https://docs.aws.amazon.com/AmazonS3/latest/dev/ServerSideEncryptionCustomerKeys)
+
+Druid applies the configured encryption to deep storage, task logs, MSQ durable storage, and S3 exports. This includes multipart uploads.
+
+When you use `kms`, the KMS key policy must grant `kms:GenerateDataKey` and `kms:Decrypt` to the roles that write to S3, because AWS requires both for KMS-encrypted multipart uploads. Grant `kms:Decrypt` to every role that reads from S3, including Brokers that read MSQ query results from durable storage and any systems outside Druid that read exported files.

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/CustomServerSideEncryption.java

@@ -22,9 +22,11 @@
 import com.fasterxml.jackson.annotation.JacksonInject;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
+import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 
 class CustomServerSideEncryption implements ServerSideEncryption
 {
@@ -73,4 +75,20 @@ public CopyObjectRequest.Builder decorate(CopyObjectRequest.Builder builder)
         .copySourceSSECustomerAlgorithm(ALGORITHM)
         .copySourceSSECustomerKey(base64EncodedKey);
   }
+
+  @Override
+  public CreateMultipartUploadRequest.Builder decorate(CreateMultipartUploadRequest.Builder builder)
+  {
+    return builder
+        .sseCustomerAlgorithm(ALGORITHM)
+        .sseCustomerKey(base64EncodedKey);
+  }
+
+  @Override
+  public UploadPartRequest.Builder decorate(UploadPartRequest.Builder builder)
+  {
+    return builder
+        .sseCustomerAlgorithm(ALGORITHM)
+        .sseCustomerKey(base64EncodedKey);
+  }
 }

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/KmsServerSideEncryption.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.annotation.JacksonInject;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
+import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.ServerSideEncryption;
 
@@ -57,4 +58,14 @@ public CopyObjectRequest.Builder decorate(CopyObjectRequest.Builder builder)
     }
     return builder;
   }
+
+  @Override
+  public CreateMultipartUploadRequest.Builder decorate(CreateMultipartUploadRequest.Builder builder)
+  {
+    builder.serverSideEncryption(ServerSideEncryption.AWS_KMS);
+    if (keyId != null) {
+      builder.ssekmsKeyId(keyId);
+    }
+    return builder;
+  }
 }

extensions-core/s3-extensions/src/main/java/org/apache/druid/storage/s3/S3ServerSideEncryption.java

@@ -20,6 +20,7 @@
 package org.apache.druid.storage.s3;
 
 import software.amazon.awssdk.services.s3.model.CopyObjectRequest;
+import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 
 class S3ServerSideEncryption implements ServerSideEncryption
@@ -38,4 +39,10 @@ public CopyObjectRequest.Builder decorate(CopyObjectRequest.Builder builder)
   {
     return builder.serverSideEncryption(SSE_AES256);
   }
+
+  @Override
+  public CreateMultipartUploadRequest.Builder decorate(CreateMultipartUploadRequest.Builder builder)
+  {
+    return builder.serverSideEncryption(SSE_AES256);
+  }
 }

extensions-core/s3-extensions/src/test/java/org/apache/druid/storage/s3/ServerSideEncryptionTest.java

@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.storage.s3;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.druid.jackson.DefaultObjectMapper;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.UploadPartRequest;
+
+import java.io.IOException;
+
+/**
+ * Multipart uploads (used by MSQ durable storage and export) must carry the same encryption settings as single PUTs.
+ */
+public class ServerSideEncryptionTest
+{
+  private static final ObjectMapper MAPPER = new DefaultObjectMapper();
+  private static final String KEY_ID = "arn:aws:kms:us-east-1:123456789012:key/test-key";
+  private static final String CUSTOM_KEY = "dGVzdC1rZXktdGVzdC1rZXktdGVzdC1rZXktMTIzNA==";
+
+  @Test
+  public void testKmsDecoratesCreateMultipartUploadWithKeyId() throws IOException
+  {
+    final KmsServerSideEncryption sse = new KmsServerSideEncryption(
+        MAPPER.readValue("{\"keyId\":\"" + KEY_ID + "\"}", S3SSEKmsConfig.class)
+    );
+
+    final CreateMultipartUploadRequest request = sse.decorate(multipartBuilder()).build();
+    Assertions.assertEquals(software.amazon.awssdk.services.s3.model.ServerSideEncryption.AWS_KMS, request.serverSideEncryption());
+    Assertions.assertEquals(KEY_ID, request.ssekmsKeyId());
+
+    // Must match what a single PUT gets.
+    final PutObjectRequest putRequest = sse.decorate(PutObjectRequest.builder().bucket("b").key("k")).build();
+    Assertions.assertEquals(putRequest.serverSideEncryption(), request.serverSideEncryption());
+    Assertions.assertEquals(putRequest.ssekmsKeyId(), request.ssekmsKeyId());
+  }
+
+  @Test
+  public void testKmsWithoutKeyIdDecoratesCreateMultipartUpload() throws IOException
+  {
+    final KmsServerSideEncryption sse = new KmsServerSideEncryption(MAPPER.readValue("{}", S3SSEKmsConfig.class));
+
+    final CreateMultipartUploadRequest request = sse.decorate(multipartBuilder()).build();
+    Assertions.assertEquals(software.amazon.awssdk.services.s3.model.ServerSideEncryption.AWS_KMS, request.serverSideEncryption());
+    Assertions.assertNull(request.ssekmsKeyId());
+  }
+
+  @Test
+  public void testS3DecoratesCreateMultipartUpload()
+  {
+    final CreateMultipartUploadRequest request = new S3ServerSideEncryption().decorate(multipartBuilder()).build();
+    Assertions.assertEquals(software.amazon.awssdk.services.s3.model.ServerSideEncryption.AES256, request.serverSideEncryption());
+  }
+
+  @Test
+  public void testCustomDecoratesCreateMultipartUploadAndUploadPart() throws IOException
+  {
+    final CustomServerSideEncryption sse = new CustomServerSideEncryption(
+        MAPPER.readValue("{\"base64EncodedKey\":\"" + CUSTOM_KEY + "\"}", S3SSECustomConfig.class)
+    );
+
+    final CreateMultipartUploadRequest createRequest = sse.decorate(multipartBuilder()).build();
+    Assertions.assertEquals("AES256", createRequest.sseCustomerAlgorithm());
+    Assertions.assertEquals(CUSTOM_KEY, createRequest.sseCustomerKey());
+
+    final UploadPartRequest partRequest = sse.decorate(
+        UploadPartRequest.builder().bucket("b").key("k").uploadId("u").partNumber(1)
+    ).build();
+    Assertions.assertEquals("AES256", partRequest.sseCustomerAlgorithm());
+    Assertions.assertEquals(CUSTOM_KEY, partRequest.sseCustomerKey());
+  }
+
+  @Test
+  public void testNoopLeavesCreateMultipartUploadUntouched()
+  {
+    final CreateMultipartUploadRequest request = new NoopServerSideEncryption().decorate(multipartBuilder()).build();
+    Assertions.assertNull(request.serverSideEncryption());
+    Assertions.assertNull(request.ssekmsKeyId());
+  }
+
+  private static CreateMultipartUploadRequest.Builder multipartBuilder()
+  {
+    return CreateMultipartUploadRequest.builder().bucket("b").key("k");
+  }
+}