improve encryption key length check (#2044)

awslabs · Jun 26, 2024 · 6b72835 · 6b72835
1 parent 2567609
commit 6b72835
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
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
@@ -141,12 +141,14 @@ public MetadataHandler(String sourceType, java.util.Map<String, String> configOp
         this.spillPrefix = this.configOptions.getOrDefault(SPILL_PREFIX_ENV, DEFAULT_SPILL_PREFIX);
 
         if (DISABLE_ENCRYPTION.equalsIgnoreCase(this.configOptions.getOrDefault(DISABLE_SPILL_ENCRYPTION, "false"))) {
+            logger.debug("DISABLE_SPILL_ENCRYPTION");
             this.encryptionKeyFactory = null;
         }
         else {
             this.encryptionKeyFactory = (this.configOptions.get(KMS_KEY_ID_ENV) != null) ?
                     new KmsKeyFactory(AWSKMSClientBuilder.standard().build(), this.configOptions.get(KMS_KEY_ID_ENV)) :
                     new LocalKeyFactory();
+            logger.debug("ENABLE_SPILL_ENCRYPTION with encryption factory: " + encryptionKeyFactory.getClass().getSimpleName());
         }
 
         this.secretsManager = new CachableSecretsManager(AWSSecretsManagerClientBuilder.defaultClient());

diff --git a/...n-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/AesGcmBlockCrypto.java b/...n-sdk/src/main/java/com/amazonaws/athena/connector/lambda/security/AesGcmBlockCrypto.java
@@ -25,6 +25,8 @@
 import com.amazonaws.athena.connector.lambda.data.RecordBatchSerDe;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import javax.crypto.BadPaddingException;
 import javax.crypto.Cipher;
@@ -49,6 +51,7 @@
 public class AesGcmBlockCrypto
         implements BlockCrypto
 {
+    private static final Logger logger = LoggerFactory.getLogger(AesGcmBlockCrypto.class);
     protected static final int GCM_TAG_LENGTH_BITS = 16 * 8;
     protected static final int NONCE_BYTES = 12;
     protected static final int KEY_BYTES = 16;
@@ -116,8 +119,9 @@ private Cipher makeCipher(int mode, EncryptionKey key)
             throw new RuntimeException("Expected " + NONCE_BYTES + " nonce bytes but found " + key.getNonce().length);
         }
 
-        if (key.getKey().length != KEY_BYTES) {
-            throw new RuntimeException("Expected " + KEY_BYTES + " key bytes but found " + key.getKey().length);
+        logger.debug("Cipher key bytes length: " + key.getKey().length);
+        if (key.getKey() == null || key.getKey().length == 0) {
+            throw new RuntimeException("Invalid key");
         }
 
         GCMParameterSpec spec = new GCMParameterSpec(GCM_TAG_LENGTH_BITS, key.getNonce());