Updating EncryptionKeyFactory to add overirde AWS request configuration in KMS calls made generate encryption key (#3103)

Author: DilipReddyGaddam

athena-dynamodb/src/main/java/com/amazonaws/athena/connectors/dynamodb/DynamoDBMetadataHandler.java

@@ -559,6 +559,8 @@ public GetSplitsResponse doGetSplits(BlockAllocator allocator, GetSplitsRequest
             logger.info("QPT Split Requested");
             return setupQueryPassthroughSplit(request);
         }
+        FederatedIdentity federatedIdentity = request.getIdentity();
+        AwsRequestOverrideConfiguration overrideConfig = getRequestOverrideConfig(federatedIdentity.getConfigOptions());
 
         int partitionContd = decodeContinuationToken(request);
         Set<Split> splits = new HashSet<>();
@@ -584,7 +586,7 @@ public GetSplitsResponse doGetSplits(BlockAllocator allocator, GetSplitsRequest
                 Object hashKeyValue = DDBTypeUtils.convertArrowTypeIfNecessary(hashKeyName, hashKeyValueReader.readObject());
                 splitMetadata.put(hashKeyName, DDBTypeUtils.attributeToJson(DDBTypeUtils.toAttributeValue(hashKeyValue), hashKeyName));
 
-                splits.add(new Split(spillLocation, makeEncryptionKey(), splitMetadata));
+                splits.add(new Split(spillLocation, makeEncryptionKey(overrideConfig), splitMetadata));
 
                 if (splits.size() == MAX_SPLITS_PER_REQUEST && curPartition != partitions.getRowCount() - 1) {
                     // We've reached max page size and this is not the last partition
@@ -609,7 +611,7 @@ else if (SCAN_PARTITION_TYPE.equals(partitionType)) {
                 splitMetadata.put(SEGMENT_ID_PROPERTY, String.valueOf(curPartition));
                 splitMetadata.put(SEGMENT_COUNT_METADATA, String.valueOf(segmentCount));
 
-                splits.add(new Split(spillLocation, makeEncryptionKey(), splitMetadata));
+                splits.add(new Split(spillLocation, makeEncryptionKey(overrideConfig), splitMetadata));
 
                 if (splits.size() == MAX_SPLITS_PER_REQUEST && curPartition != segmentCount - 1) {
                     // We've reached max page size and this is not the last partition
@@ -740,13 +742,15 @@ else if (useQueryPlan && filterPredicates.containsKey(rangeKeyName)) {
      */
     private GetSplitsResponse setupQueryPassthroughSplit(GetSplitsRequest request)
     {
+        FederatedIdentity federatedIdentity = request.getIdentity();
+        AwsRequestOverrideConfiguration overrideConfig = getRequestOverrideConfig(federatedIdentity.getConfigOptions());
         //Every split must have a unique location if we wish to spill to avoid failures
         SpillLocation spillLocation = makeSpillLocation(request);
 
         //Since this is QPT query we return a fixed split.
         Map<String, String> qptArguments = request.getConstraints().getQueryPassthroughArguments();
         return new GetSplitsResponse(request.getCatalogName(),
-                Split.newBuilder(spillLocation, makeEncryptionKey())
+                Split.newBuilder(spillLocation, makeEncryptionKey(overrideConfig))
                         .applyProperties(qptArguments)
                         .build());
     }

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java

@@ -233,6 +233,11 @@ protected EncryptionKey makeEncryptionKey()
         return (encryptionKeyFactory != null) ? encryptionKeyFactory.create() : null;
     }
 
+    protected EncryptionKey makeEncryptionKey(AwsRequestOverrideConfiguration awsRequestOverrideConfiguration)
+    {
+        return (encryptionKeyFactory != null) ? encryptionKeyFactory.create(awsRequestOverrideConfiguration) : null;
+    }
+
     /**
      * Used to make a spill location for a split. Each split should have a unique spill location, so be sure
      * to call this method once per split!

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/EncryptionKeyFactory.java

@@ -20,6 +20,8 @@
  * #L%
  */
 
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
+
 /**
  * Defines a factory that can be used to create AES-GCM compatible encryption keys.
  */
@@ -29,4 +31,9 @@ public interface EncryptionKeyFactory
      * @return A key that satisfies the specification defined in BlockCrypto
      */
     EncryptionKey create();
+
+    default EncryptionKey create(AwsRequestOverrideConfiguration awsRequestOverrideConfiguration)
+    {
+        return create();
+    }
 }

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/KmsKeyFactory.java

@@ -21,6 +21,9 @@
  */
 
 import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.services.glue.model.ErrorDetails;
 import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.kms.KmsClient;
@@ -39,6 +42,7 @@
 public class KmsKeyFactory
         implements EncryptionKeyFactory
 {
+    private static final Logger logger = LoggerFactory.getLogger(KmsKeyFactory.class);
     private final KmsClient kmsClient;
     private final String masterKeyId;
 
@@ -53,21 +57,30 @@ public KmsKeyFactory(KmsClient kmsClient, String masterKeyId)
      */
     public EncryptionKey create()
     {
+        return create(null);
+    }
+
+    @Override
+    public EncryptionKey create(AwsRequestOverrideConfiguration awsRequestOverrideConfiguration)
+    {
+        GenerateDataKeyRequest.Builder dataKeyBuilder = GenerateDataKeyRequest.builder()
+                .keyId(masterKeyId)
+                .keySpec(DataKeySpec.AES_128);
+        if (awsRequestOverrideConfiguration != null) {
+            logger.info("Using AWS KMS Request Override Configuration:");
+            dataKeyBuilder.overrideConfiguration(awsRequestOverrideConfiguration);
+        }
+
         GenerateDataKeyResponse dataKeyResponse;
         try {
-            dataKeyResponse = kmsClient.generateDataKey(
-                            GenerateDataKeyRequest.builder()
-                                    .keyId(masterKeyId)
-                                    .keySpec(DataKeySpec.AES_128)
-                                    .build());
+            dataKeyResponse = kmsClient.generateDataKey(dataKeyBuilder.build());
         }
         catch (NotFoundException e) {
             throw new AthenaConnectorException(e.getMessage(), ErrorDetails.builder().errorCode(FederationSourceErrorCode.ENTITY_NOT_FOUND_EXCEPTION.toString()).build());
         }
 
         GenerateRandomRequest randomRequest = GenerateRandomRequest.builder()
-                .numberOfBytes(AesGcmBlockCrypto.NONCE_BYTES)
-                .build();
+                .numberOfBytes(AesGcmBlockCrypto.NONCE_BYTES).build();
         GenerateRandomResponse randomResponse = kmsClient.generateRandom(randomRequest);
 
         return new EncryptionKey(dataKeyResponse.plaintext().asByteArray(), randomResponse.plaintext().asByteArray());

athena-federation-sdk/src/test/java/com/amazonaws/athena/connector/lambda/security/KmsKeyFactoryTest.java

@@ -32,6 +32,7 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.services.kms.KmsClient;
 import software.amazon.awssdk.services.kms.model.GenerateDataKeyRequest;
@@ -95,4 +96,51 @@ public void testNotFoundException() {
             kmsKeyFactory.create();
         });
     }
+
+    @Test
+    public void testCreateWithOverrideConfiguration() {
+        byte[] testPlaintextKey = new byte[] { 1, 2, 3, 4, 5 };
+        byte[] testNonce = new byte[] { 9, 8, 7 };
+        AwsRequestOverrideConfiguration overrideConfig = AwsRequestOverrideConfiguration.builder().build();
+
+        GenerateDataKeyResponse dataKeyResponse = GenerateDataKeyResponse.builder()
+                .plaintext(SdkBytes.fromByteArray(testPlaintextKey))
+                .build();
+
+        GenerateRandomResponse randomResponse = GenerateRandomResponse.builder()
+                .plaintext(SdkBytes.fromByteArray(testNonce))
+                .build();
+
+        when(mockKmsClient.generateDataKey((GenerateDataKeyRequest) any())).thenReturn(dataKeyResponse);
+        when(mockKmsClient.generateRandom((GenerateRandomRequest) any())).thenReturn(randomResponse);
+
+        EncryptionKey result = kmsKeyFactory.create(overrideConfig);
+
+        assertArrayEquals(testPlaintextKey, result.getKey());
+        assertArrayEquals(testNonce, result.getNonce());
+        verify(mockKmsClient).generateDataKey((GenerateDataKeyRequest) any());
+        verify(mockKmsClient).generateRandom((GenerateRandomRequest) any());
+    }
+
+    @Test
+    public void testCreateWithNullOverrideConfiguration() {
+        byte[] testPlaintextKey = new byte[] { 1, 2, 3, 4, 5 };
+        byte[] testNonce = new byte[] { 9, 8, 7 };
+
+        GenerateDataKeyResponse dataKeyResponse = GenerateDataKeyResponse.builder()
+                .plaintext(SdkBytes.fromByteArray(testPlaintextKey))
+                .build();
+
+        GenerateRandomResponse randomResponse = GenerateRandomResponse.builder()
+                .plaintext(SdkBytes.fromByteArray(testNonce))
+                .build();
+
+        when(mockKmsClient.generateDataKey((GenerateDataKeyRequest) any())).thenReturn(dataKeyResponse);
+        when(mockKmsClient.generateRandom((GenerateRandomRequest) any())).thenReturn(randomResponse);
+
+        EncryptionKey result = kmsKeyFactory.create(null);
+
+        assertArrayEquals(testPlaintextKey, result.getKey());
+        assertArrayEquals(testNonce, result.getNonce());
+    }
 }