Changing encryptKey and decryptKey methods to use key bytes instead of SecretKey

Author: WesleyRosenblum

src/main/java/com/amazonaws/encryptionsdk/internal/JceKeyCipher.java

@@ -81,16 +81,15 @@ abstract Cipher buildUnwrappingCipher(Key key, byte[] extraInfo, int offset,
      *                         during encryption and decryption to provide additional authenticated data (AAD).
      * @return The encrypted data key.
      */
-    public EncryptedDataKey encryptKey(final SecretKey key, final String keyName,
+    public EncryptedDataKey encryptKey(final byte[] key, final String keyName,
                                        final Map<String, String> encryptionContext) {
 
-        final byte[] keyBytes = key.getEncoded();
         final byte[] keyNameBytes = keyName.getBytes(KEY_NAME_ENCODING);
 
         try {
             final JceKeyCipher.WrappingData wData = buildWrappingCipher(wrappingKey, encryptionContext);
             final Cipher cipher = wData.cipher;
-            final byte[] encryptedKey = cipher.doFinal(keyBytes);
+            final byte[] encryptedKey = cipher.doFinal(key);
 
             final byte[] provInfo = new byte[keyNameBytes.length + wData.extraInfo.length];
             System.arraycopy(keyNameBytes, 0, provInfo, 0, keyNameBytes.length);
@@ -105,27 +104,20 @@ public EncryptedDataKey encryptKey(final SecretKey key, final String keyName,
     /**
      * Decrypts the given encrypted data key.
      *
-     * @param algorithm The algorithm that encrypted the data key.
      * @param edk The encrypted data key.
      * @param keyName A UTF-8 encoded representing a name for the key.
      * @param encryptionContext A key-value mapping of arbitrary, non-secret, UTF-8 encoded strings used
      *                          during encryption and decryption to provide additional authenticated data (AAD).
      * @return The decrypted key.
      * @throws GeneralSecurityException If a problem occurred decrypting the key.
      */
-    public SecretKey decryptKey(final CryptoAlgorithm algorithm, final EncryptedDataKey edk, final String keyName,
+    public byte[] decryptKey(final EncryptedDataKey edk, final String keyName,
                               final Map<String, String> encryptionContext) throws GeneralSecurityException {
         final byte[] keyNameBytes = keyName.getBytes(KEY_NAME_ENCODING);
 
         final Cipher cipher = buildUnwrappingCipher(unwrappingKey, edk.getProviderInformation(),
                 keyNameBytes.length, encryptionContext);
-        final byte[] rawKey = cipher.doFinal(edk.getEncryptedDataKey());
-        if (rawKey.length != algorithm.getDataKeyLength()) {
-            // Something's wrong here. Assume that the decryption is invalid.
-            return null;
-        }
-
-        return new SecretKeySpec(rawKey, algorithm.getDataKeyAlgo());
+        return cipher.doFinal(edk.getEncryptedDataKey());
     }
 
     static class WrappingData {

src/main/java/com/amazonaws/encryptionsdk/jce/JceMasterKey.java

@@ -110,8 +110,7 @@ public DataKey<JceMasterKey> generateDataKey(final CryptoAlgorithm algorithm,
             final Map<String, String> encryptionContext) {
         final byte[] rawKey = new byte[algorithm.getDataKeyLength()];
         rnd.nextBytes(rawKey);
-        EncryptedDataKey encryptedDataKey = jceKeyCipher_.encryptKey(new SecretKeySpec(rawKey, algorithm.getDataKeyAlgo()), keyId_,
-                encryptionContext);
+        EncryptedDataKey encryptedDataKey = jceKeyCipher_.encryptKey(rawKey, keyId_, encryptionContext);
         return new DataKey<>(new SecretKeySpec(rawKey, algorithm.getDataKeyAlgo()),
                 encryptedDataKey.getEncryptedDataKey(), encryptedDataKey.getProviderInformation(), this);
     }
@@ -129,7 +128,7 @@ public DataKey<JceMasterKey> encryptDataKey(final CryptoAlgorithm algorithm,
             throw new IllegalArgumentException("Incorrect key algorithm. Expected " + key.getAlgorithm()
                     + " but got " + algorithm.getKeyAlgo());
         }
-        EncryptedDataKey encryptedDataKey = jceKeyCipher_.encryptKey(key, keyId_, encryptionContext);
+        EncryptedDataKey encryptedDataKey = jceKeyCipher_.encryptKey(key.getEncoded(), keyId_, encryptionContext);
         return new DataKey<>(key, encryptedDataKey.getEncryptedDataKey(), encryptedDataKey.getProviderInformation(), this);
     }
 
@@ -144,10 +143,12 @@ public DataKey<JceMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
             try {
                 if (edk.getProviderId().equals(getProviderId())
                         && arrayPrefixEquals(edk.getProviderInformation(), keyIdBytes_, keyIdBytes_.length)) {
-                    final SecretKey decryptedKey = jceKeyCipher_.decryptKey(algorithm, edk, keyId_, encryptionContext);
+                    final byte[] decryptedKey = jceKeyCipher_.decryptKey(edk, keyId_, encryptionContext);
 
-                    if(decryptedKey != null) {
-                        return new DataKey<>(decryptedKey, edk.getEncryptedDataKey(), edk.getProviderInformation(), this);
+                    // Validate that the decrypted key length is as expected
+                    if (decryptedKey.length == algorithm.getDataKeyLength()) {
+                        return new DataKey<>(new SecretKeySpec(decryptedKey, algorithm.getDataKeyAlgo()),
+                                edk.getEncryptedDataKey(), edk.getProviderInformation(), this);
                     }
                 }
             } catch (final Exception ex) {