Merge pull request #71 from slyubomirsky/keyblob-parse-exception

Keyblob should parse field lengths as unsigned shorts

Author: SalusaSecondus

src/main/java/com/amazonaws/encryptionsdk/model/KeyBlob.java

@@ -20,6 +20,7 @@
 import com.amazonaws.encryptionsdk.EncryptedDataKey;
 import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
 import com.amazonaws.encryptionsdk.exception.ParseException;
+import com.amazonaws.encryptionsdk.internal.Constants;
 import com.amazonaws.encryptionsdk.internal.PrimitivesParser;
 
 /**
@@ -41,11 +42,11 @@
  * </ol>
  */
 public final class KeyBlob implements EncryptedDataKey {
-    private short keyProviderIdLen_ = -1;
+    private int keyProviderIdLen_ = -1;
     private byte[] keyProviderId_;
-    private short keyProviderInfoLen_ = -1;
+    private int keyProviderInfoLen_ = -1;
     private byte[] keyProviderInfo_;
-    private short encryptedKeyLen_ = -1;
+    private int encryptedKeyLen_ = -1;
     private byte[] encryptedKey_;
 
     private boolean isComplete_ = false;
@@ -77,6 +78,7 @@ public KeyBlob(final EncryptedDataKey edk) {
         setKeyProviderId(edk.getProviderId());
         setKeyProviderInfo(edk.getProviderInformation());
     }
+
     /**
      * Parse the key provider identifier length in the provided bytes. It looks
      * for 2 bytes representing a short primitive type in the provided bytes
@@ -98,7 +100,7 @@ public KeyBlob(final EncryptedDataKey edk) {
      *             length.
      */
     private int parseKeyProviderIdLen(final byte[] b, final int off) throws ParseException {
-        keyProviderIdLen_ = PrimitivesParser.parseShort(b, off);
+        keyProviderIdLen_ = PrimitivesParser.parseUnsignedShort(b, off);
         return Short.SIZE / Byte.SIZE;
     }
 
@@ -152,7 +154,7 @@ private int parseKeyProviderId(final byte[] b, final int off) throws ParseExcept
      *             length.
      */
     private int parseKeyProviderInfoLen(final byte[] b, final int off) throws ParseException {
-        keyProviderInfoLen_ = PrimitivesParser.parseShort(b, off);
+        keyProviderInfoLen_ = PrimitivesParser.parseUnsignedShort(b, off);
         return Short.SIZE / Byte.SIZE;
     }
 
@@ -205,7 +207,7 @@ private int parseKeyProviderInfo(final byte[] b, final int off) throws ParseExce
      *             if there are not sufficient bytes to parse the key length.
      */
     private int parseKeyLen(final byte[] b, final int off) throws ParseException {
-        encryptedKeyLen_ = PrimitivesParser.parseShort(b, off);
+        encryptedKeyLen_ = PrimitivesParser.parseUnsignedShort(b, off);
         return Short.SIZE / Byte.SIZE;
     }
 
@@ -302,13 +304,13 @@ public byte[] toByteArray() {
         final int outLen = 3 * (Short.SIZE / Byte.SIZE) + keyProviderIdLen_ + keyProviderInfoLen_ + encryptedKeyLen_;
         final ByteBuffer out = ByteBuffer.allocate(outLen);
 
-        out.putShort(keyProviderIdLen_);
+        out.putShort((short) keyProviderIdLen_);
         out.put(keyProviderId_, 0, keyProviderIdLen_);
 
-        out.putShort(keyProviderInfoLen_);
+        out.putShort((short) keyProviderInfoLen_);
         out.put(keyProviderInfo_, 0, keyProviderInfoLen_);
 
-        out.putShort(encryptedKeyLen_);
+        out.putShort((short) encryptedKeyLen_);
         out.put(encryptedKey_, 0, encryptedKeyLen_);
 
         return out.array();
@@ -332,7 +334,7 @@ public boolean isComplete() {
      * @return
      *         the length of the key provider identifier.
      */
-    public short getKeyProviderIdLen() {
+    public int getKeyProviderIdLen() {
         return keyProviderIdLen_;
     }
 
@@ -353,7 +355,7 @@ public String getProviderId() {
      * @return
      *         the length of the key provider info.
      */
-    public short getKeyProviderInfoLen() {
+    public int getKeyProviderInfoLen() {
         return keyProviderInfoLen_;
     }
 
@@ -374,7 +376,7 @@ public byte[] getProviderInformation() {
      * @return
      *         the length of the encrypted data key.
      */
-    public short getEncryptedDataKeyLen() {
+    public int getEncryptedDataKeyLen() {
         return encryptedKeyLen_;
     }
 
@@ -397,12 +399,12 @@ public byte[] getEncryptedDataKey() {
      */
     public void setKeyProviderId(final String keyProviderId) {
         final byte[] keyProviderIdBytes = keyProviderId.getBytes(StandardCharsets.UTF_8);
-        if (keyProviderIdBytes.length > Short.MAX_VALUE) {
+        if (keyProviderIdBytes.length > Constants.UNSIGNED_SHORT_MAX_VAL) {
             throw new AwsCryptoException(
-                    "Key provider identifier length exceeds the max value of a short primitive.");
+                    "Key provider identifier length exceeds the max value of an unsigned short primitive.");
         }
         keyProviderId_ = keyProviderIdBytes;
-        keyProviderIdLen_ = (short) keyProviderId_.length;
+        keyProviderIdLen_ = keyProviderId_.length;
     }
 
     /**
@@ -413,12 +415,12 @@ public void setKeyProviderId(final String keyProviderId) {
      *            identifier.
      */
     public void setKeyProviderInfo(final byte[] keyProviderInfo) {
-        if (keyProviderInfo.length > Short.MAX_VALUE) {
+        if (keyProviderInfo.length > Constants.UNSIGNED_SHORT_MAX_VAL) {
             throw new AwsCryptoException(
-                    "Key provider identifier information length exceeds the max value of a short primitive.");
+                    "Key provider identifier information length exceeds the max value of an unsigned short primitive.");
         }
         keyProviderInfo_ = keyProviderInfo.clone();
-        keyProviderInfoLen_ = (short) keyProviderInfo.length;
+        keyProviderInfoLen_ = keyProviderInfo.length;
     }
 
     /**
@@ -428,10 +430,10 @@ public void setKeyProviderInfo(final byte[] keyProviderInfo) {
      *            the bytes containing the encrypted data key.
      */
     public void setEncryptedDataKey(final byte[] encryptedDataKey) {
-        if (encryptedDataKey.length > Short.MAX_VALUE) {
-            throw new AwsCryptoException("Key length exceeds the max value of a short primitive.");
+        if (encryptedDataKey.length > Constants.UNSIGNED_SHORT_MAX_VAL) {
+            throw new AwsCryptoException("Key length exceeds the max value of an unsigned short primitive.");
         }
         encryptedKey_ = encryptedDataKey.clone();
-        encryptedKeyLen_ = (short) encryptedKey_.length;
+        encryptedKeyLen_ = encryptedKey_.length;
     }
-}
+}

src/test/java/com/amazonaws/encryptionsdk/model/KeyBlobTest.java

@@ -17,6 +17,7 @@
 import static org.junit.Assert.assertEquals;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -26,6 +27,7 @@
 import com.amazonaws.encryptionsdk.CryptoAlgorithm;
 import com.amazonaws.encryptionsdk.DataKey;
 import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
+import com.amazonaws.encryptionsdk.internal.Constants;
 import com.amazonaws.encryptionsdk.internal.RandomBytesGenerator;
 import com.amazonaws.encryptionsdk.internal.StaticMasterKey;
 
@@ -76,7 +78,7 @@ public void overlyLargeKeyProviderIdLen() {
 
         final DataKey<StaticMasterKey> mockDataKey = masterKeyProvider_.generateDataKey(ALGORITHM, encryptionContext);
 
-        final int providerId_Len = Short.MAX_VALUE + 1;
+        final int providerId_Len = Constants.UNSIGNED_SHORT_MAX_VAL + 1;
         final byte[] providerId_Bytes = RandomBytesGenerator.generate(providerId_Len);
         final String providerId_ = new String(providerId_Bytes, StandardCharsets.UTF_8);
 
@@ -93,15 +95,15 @@ public void overlyLargeKeyProviderInfoLen() {
 
         final DataKey<StaticMasterKey> mockDataKey = masterKeyProvider_.generateDataKey(ALGORITHM, encryptionContext);
 
-        final int providerInfo_Len = Short.MAX_VALUE + 1;
+        final int providerInfo_Len = Constants.UNSIGNED_SHORT_MAX_VAL + 1;
         final byte[] providerInfo_ = RandomBytesGenerator.generate(providerInfo_Len);
 
         new KeyBlob(providerId_, providerInfo_, mockDataKey.getEncryptedDataKey());
     }
 
     @Test(expected = AwsCryptoException.class)
     public void overlyLargeKey() {
-        final int keyLen = Short.MAX_VALUE + 1;
+        final int keyLen = Constants.UNSIGNED_SHORT_MAX_VAL + 1;
         final byte[] encryptedKeyBytes = RandomBytesGenerator.generate(keyLen);
 
         new KeyBlob(providerId_, providerInfo_.getBytes(StandardCharsets.UTF_8), encryptedKeyBytes);
@@ -170,4 +172,48 @@ public void checkKeyLen() {
 
         assertEquals(mockDataKey_.getEncryptedDataKey().length, reconstructedKeyBlob.getEncryptedDataKeyLen());
     }
+
+    private KeyBlob generateRandomKeyBlob(int idLen, int infoLen, int keyLen) {
+        final byte[] idBytes = new byte[idLen];
+        Arrays.fill(idBytes, (byte) 'A');
+
+        final byte[] infoBytes = RandomBytesGenerator.generate(infoLen);
+        final byte[] keyBytes = RandomBytesGenerator.generate(keyLen);
+
+        return new KeyBlob(new String(idBytes, StandardCharsets.UTF_8), infoBytes, keyBytes);
+    }
+
+    private void assertKeyBlobsEqual(KeyBlob b1, KeyBlob b2) {
+        assertArrayEquals(b1.getProviderId().getBytes(StandardCharsets.UTF_8),
+                          b2.getProviderId().getBytes(StandardCharsets.UTF_8));
+        assertArrayEquals(b1.getProviderInformation(), b2.getProviderInformation());
+        assertArrayEquals(b1.getEncryptedDataKey(), b2.getEncryptedDataKey());
+    }
+    
+    @Test
+    public void checkKeyProviderIdLenUnsigned() {
+        // provider id length is too large for a signed short but fits in unsigned
+        final KeyBlob blob = generateRandomKeyBlob(Constants.UNSIGNED_SHORT_MAX_VAL, Short.MAX_VALUE, Short.MAX_VALUE);
+        final byte[] arr = blob.toByteArray();
+
+        assertKeyBlobsEqual(deserialize(arr), blob);
+    }
+
+    @Test
+    public void checkKeyProviderInfoLenUnsigned() {
+        // provider info length is too large for a signed short but fits in unsigned
+        final KeyBlob blob = generateRandomKeyBlob(Short.MAX_VALUE, Constants.UNSIGNED_SHORT_MAX_VAL, Short.MAX_VALUE);
+        final byte[] arr = blob.toByteArray();
+
+        assertKeyBlobsEqual(deserialize(arr), blob);
+    }
+
+    @Test
+    public void checkKeyLenUnsigned() {
+        // key length is too large for a signed short but fits in unsigned
+        final KeyBlob blob = generateRandomKeyBlob(Short.MAX_VALUE, Short.MAX_VALUE, Constants.UNSIGNED_SHORT_MAX_VAL);
+        final byte[] arr = blob.toByteArray();
+
+        assertKeyBlobsEqual(deserialize(arr), blob);
+    }
 }