V2 migration KMS (#2107)

awslabs · Jul 22, 2024 · 3f2a330 · 3f2a330
1 parent 40a3a3f
commit 3f2a330
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 21 deletions.
diff --git a/athena-federation-sdk/pom.xml b/athena-federation-sdk/pom.xml
@@ -201,11 +201,24 @@
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-s3</artifactId>
             <version>${aws-sdk.version}</version>
+            <exclusions>
+                <exclusion>
+                    <!--REMOVE ONCE MIGRATED TO v2-->
+                    <groupId>com.amazonaws</groupId>
+                    <artifactId>aws-java-sdk-kms</artifactId>
+                </exclusion>
+            </exclusions>
         </dependency>
         <dependency>
-            <groupId>com.amazonaws</groupId>
-            <artifactId>aws-java-sdk-kms</artifactId>
-            <version>${aws-sdk.version}</version>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>kms</artifactId>
+            <version>${aws-sdk-v2.version}</version>
+            <exclusions>
+                <exclusion>
+                    <groupId>software.amazon.awssdk</groupId>
+                    <artifactId>netty-nio-client</artifactId>
+                </exclusion>
+            </exclusions>
         </dependency>
         <dependency>
             <groupId>com.google.guava</groupId>

diff --git a/...ion-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java b/...ion-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/MetadataHandler.java
@@ -58,7 +58,6 @@
 import com.amazonaws.athena.connector.lambda.security.KmsKeyFactory;
 import com.amazonaws.athena.connector.lambda.security.LocalKeyFactory;
 import com.amazonaws.athena.connector.lambda.serde.VersionedObjectMapperFactory;
-import com.amazonaws.services.kms.AWSKMSClientBuilder;
 import com.amazonaws.services.lambda.runtime.Context;
 import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
 import com.amazonaws.services.s3.AmazonS3ClientBuilder;
@@ -69,6 +68,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.athena.AthenaClient;
+import software.amazon.awssdk.services.kms.KmsClient;
 import software.amazon.awssdk.services.secretsmanager.SecretsManagerClient;
 
 import java.io.IOException;
@@ -144,7 +144,7 @@ public MetadataHandler(String sourceType, java.util.Map<String, String> configOp
         }
         else {
             this.encryptionKeyFactory = (this.configOptions.get(KMS_KEY_ID_ENV) != null) ?
-                    new KmsKeyFactory(AWSKMSClientBuilder.standard().build(), this.configOptions.get(KMS_KEY_ID_ENV)) :
+                    new KmsKeyFactory(KmsClient.create(), this.configOptions.get(KMS_KEY_ID_ENV)) :
                     new LocalKeyFactory();
             logger.debug("ENABLE_SPILL_ENCRYPTION with encryption factory: " + encryptionKeyFactory.getClass().getSimpleName());
         }

diff --git a/...ation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/KmsKeyFactory.java b/...ation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/KmsKeyFactory.java
@@ -20,12 +20,12 @@
  * #L%
  */
 
-import com.amazonaws.services.kms.AWSKMS;
-import com.amazonaws.services.kms.model.DataKeySpec;
-import com.amazonaws.services.kms.model.GenerateDataKeyRequest;
-import com.amazonaws.services.kms.model.GenerateDataKeyResult;
-import com.amazonaws.services.kms.model.GenerateRandomRequest;
-import com.amazonaws.services.kms.model.GenerateRandomResult;
+import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.kms.model.DataKeySpec;
+import software.amazon.awssdk.services.kms.model.GenerateDataKeyRequest;
+import software.amazon.awssdk.services.kms.model.GenerateDataKeyResponse;
+import software.amazon.awssdk.services.kms.model.GenerateRandomRequest;
+import software.amazon.awssdk.services.kms.model.GenerateRandomResponse;
 
 /**
  * An EncryptionKeyFactory that is backed by AWS KMS.
@@ -35,10 +35,10 @@
 public class KmsKeyFactory
         implements EncryptionKeyFactory
 {
-    private final AWSKMS kmsClient;
+    private final KmsClient kmsClient;
     private final String masterKeyId;
 
-    public KmsKeyFactory(AWSKMS kmsClient, String masterKeyId)
+    public KmsKeyFactory(KmsClient kmsClient, String masterKeyId)
     {
         this.kmsClient = kmsClient;
         this.masterKeyId = masterKeyId;
@@ -49,16 +49,18 @@ public KmsKeyFactory(AWSKMS kmsClient, String masterKeyId)
      */
     public EncryptionKey create()
     {
-        GenerateDataKeyResult dataKeyResult =
+        GenerateDataKeyResponse dataKeyResponse =
                 kmsClient.generateDataKey(
-                        new GenerateDataKeyRequest()
-                                .withKeyId(masterKeyId)
-                                .withKeySpec(DataKeySpec.AES_128));
+                        GenerateDataKeyRequest.builder()
+                                .keyId(masterKeyId)
+                                .keySpec(DataKeySpec.AES_128)
+                                .build());
 
-        GenerateRandomRequest randomRequest = new GenerateRandomRequest()
-                .withNumberOfBytes(AesGcmBlockCrypto.NONCE_BYTES);
-        GenerateRandomResult randomResult = kmsClient.generateRandom(randomRequest);
+        GenerateRandomRequest randomRequest = GenerateRandomRequest.builder()
+                .numberOfBytes(AesGcmBlockCrypto.NONCE_BYTES)
+                .build();
+        GenerateRandomResponse randomResponse = kmsClient.generateRandom(randomRequest);
 
-        return new EncryptionKey(dataKeyResult.getPlaintext().array(), randomResult.getPlaintext().array());
+        return new EncryptionKey(dataKeyResponse.plaintext().asByteArray(), randomResponse.plaintext().asByteArray());
     }
 }