Adding validation to AesGcm key cipher and moving ArrayPrefixEquals to Utils

Author: WesleyRosenblum

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

@@ -22,6 +22,7 @@
 import java.io.DataOutputStream;
 import java.io.IOException;
 import java.security.GeneralSecurityException;
+import java.security.InvalidKeyException;
 import java.security.Key;
 import java.security.SecureRandom;
 import java.util.Map;
@@ -54,11 +55,20 @@ private static byte[] specToBytes(final GCMParameterSpec spec) {
         return baos.toByteArray();
     }
 
-    private static GCMParameterSpec bytesToSpec(final byte[] data, final int offset) {
+    private static GCMParameterSpec bytesToSpec(final byte[] data, final int offset) throws GeneralSecurityException {
         final ByteArrayInputStream bais = new ByteArrayInputStream(data, offset, data.length - offset);
         try (final DataInputStream dis = new DataInputStream(bais)) {
             final int tagLen = dis.readInt();
             final int nonceLen = dis.readInt();
+
+            if(tagLen != TAG_LENGTH) {
+                throw new InvalidKeyException(String.format("Authentication tag length must be %s", TAG_LENGTH));
+            }
+
+            if(nonceLen != NONCE_LENGTH) {
+                throw new InvalidKeyException(String.format("Initialization vector (IV) length must be %s", NONCE_LENGTH));
+            }
+
             final byte[] nonce = new byte[nonceLen];
             dis.readFully(nonce);
             return new GCMParameterSpec(tagLen, nonce);

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

@@ -311,4 +311,24 @@ public static byte[] bigIntegerToByteArray(final BigInteger bigInteger, final in
         System.arraycopy(rawBytes, 0, paddedResult, length - rawBytes.length, rawBytes.length);
         return paddedResult;
     }
+
+    /**
+     * Returns true if the prefix of the given length for the input arrays are equal.
+     *
+     * @param a      The first array.
+     * @param b      The second array.
+     * @param length The length of the prefix to compare.
+     * @return True if the prefixes are equal, false otherwise.
+     */
+    public static boolean arrayPrefixEquals(final byte[] a, final byte[] b, final int length) {
+        if (a == null || b == null || a.length < length || b.length < length) {
+            return false;
+        }
+        for (int x = 0; x < length; x++) {
+            if (a[x] != b[x]) {
+                return false;
+            }
+        }
+        return true;
+    }
 }

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

@@ -20,6 +20,7 @@
 import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
 import com.amazonaws.encryptionsdk.exception.UnsupportedProviderException;
 import com.amazonaws.encryptionsdk.internal.JceKeyCipher;
+import com.amazonaws.encryptionsdk.internal.Utils;
 
 import javax.crypto.SecretKey;
 import javax.crypto.spec.SecretKeySpec;
@@ -142,7 +143,7 @@ public DataKey<JceMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
         for (final EncryptedDataKey edk : encryptedDataKeys) {
             try {
                 if (edk.getProviderId().equals(getProviderId())
-                        && arrayPrefixEquals(edk.getProviderInformation(), keyIdBytes_, keyIdBytes_.length)) {
+                        && Utils.arrayPrefixEquals(edk.getProviderInformation(), keyIdBytes_, keyIdBytes_.length)) {
                     final byte[] decryptedKey = jceKeyCipher_.decryptKey(edk, keyId_, encryptionContext);
 
                     // Validate that the decrypted key length is as expected
@@ -157,16 +158,4 @@ && arrayPrefixEquals(edk.getProviderInformation(), keyIdBytes_, keyIdBytes_.leng
         }
         throw buildCannotDecryptDksException(exceptions);
     }
-
-    private static boolean arrayPrefixEquals(final byte[] a, final byte[] b, final int len) {
-        if (a == null || b == null || a.length < len || b.length < len) {
-            return false;
-        }
-        for (int x = 0; x < len; x++) {
-            if (a[x] != b[x]) {
-                return false;
-            }
-        }
-        return true;
-    }
 }

src/test/java/com/amazonaws/encryptionsdk/UtilsTest.java

@@ -2,6 +2,7 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -121,5 +122,17 @@ public void testBigIntegerToByteArray_InvalidLength() {
                 Utils.bigIntegerToByteArray(new BigInteger(bytes), 3));
     }
 
+    @Test
+    public void testArrayPrefixEquals() {
+        byte[] a = new byte[] {10, 11, 12, 13, 14, 15};
+        byte[] b = new byte[] {10, 11, 12, 13, 20, 21, 22};
+
+        assertFalse(Utils.arrayPrefixEquals(null, b, 4));
+        assertFalse(Utils.arrayPrefixEquals(a, null, 4));
+        assertFalse(Utils.arrayPrefixEquals(a, b, a.length + 1));
+        assertTrue(Utils.arrayPrefixEquals(a, b, 4));
+        assertFalse(Utils.arrayPrefixEquals(a, b, 5));
+    }
+
 }