Fix GCMShortBuffer test (#367)

Co-authored-by: Greg Rubin <ga_rubin@apple.com>
corretto · Mar 12, 2024 · 31fdb38 · 31fdb38
1 parent 9d1f26d
commit 31fdb38
Show file tree

Hide file tree

Showing 2 changed files with 91 additions and 18 deletions.
diff --git a/src/com/amazon/corretto/crypto/provider/AesGcmSpi.java b/src/com/amazon/corretto/crypto/provider/AesGcmSpi.java
@@ -564,17 +564,19 @@ protected void engineUpdateAAD(ByteBuffer byteBuffer) {
   private int engineEncryptFinal(
       byte[] input, final int inputOffset, int inputLen, final byte[] output, int outputOffset)
       throws ShortBufferException {
-    try {
-      if (opMode != NATIVE_MODE_ENCRYPT) {
-        throw new IllegalStateException("Cipher not initialized for encryption");
-      }
-      if (input == null) {
-        input = EMPTY_ARRAY;
-      }
+    // The following failures should not trigger reset
+    if (opMode != NATIVE_MODE_ENCRYPT) {
+      throw new IllegalStateException("Cipher not initialized for encryption");
+    }
+    if (input == null) {
+      input = EMPTY_ARRAY;
+    }
 
-      checkOutputBuffer(inputLen, output, outputOffset, true);
-      checkArrayLimits(input, inputOffset, inputLen);
+    checkOutputBuffer(inputLen, output, outputOffset, true);
+    checkArrayLimits(input, inputOffset, inputLen);
 
+    // Any future success or failure should trigger reset
+    try {
       final boolean clobbers =
           Utils.outputClobbersInput(input, inputOffset, inputLen, output, outputOffset);
 
@@ -705,17 +707,19 @@ private int engineDecryptFinal(
       byte[] output,
       final int outputOffset)
       throws AEADBadTagException, ShortBufferException {
-    try {
-      if (opMode != NATIVE_MODE_DECRYPT) {
-        throw new IllegalStateException("Cipher not initialized for decryption");
-      }
-      if (input == null) {
-        input = EMPTY_ARRAY;
-      }
+    // The following failures should not trigger reset
+    if (opMode != NATIVE_MODE_DECRYPT) {
+      throw new IllegalStateException("Cipher not initialized for decryption");
+    }
+    if (input == null) {
+      input = EMPTY_ARRAY;
+    }
 
-      checkOutputBuffer(inputLen, output, outputOffset, true);
-      checkArrayLimits(input, inputOffset, inputLen);
+    checkOutputBuffer(inputLen, output, outputOffset, true);
+    checkArrayLimits(input, inputOffset, inputLen);
 
+    // Any future failure (or success) should trigger reset
+    try {
       final byte[] workingInputArray;
       final int workingInputOffset;
       final int workingInputLength;

diff --git a/tst/com/amazon/corretto/crypto/provider/test/AesTest.java b/tst/com/amazon/corretto/crypto/provider/test/AesTest.java
@@ -3,6 +3,7 @@
 package com.amazon.corretto.crypto.provider.test;
 
 import static com.amazon.corretto.crypto.provider.test.TestUtil.NATIVE_PROVIDER;
+import static com.amazon.corretto.crypto.provider.test.TestUtil.assertArraysHexEquals;
 import static com.amazon.corretto.crypto.provider.test.TestUtil.assertThrows;
 import static com.amazon.corretto.crypto.provider.test.TestUtil.assumeMinimumVersion;
 import static com.amazon.corretto.crypto.provider.test.TestUtil.sneakyConstruct;
@@ -1314,6 +1315,74 @@ public void unintCipherReturnsParameters() throws GeneralSecurityException {
     assertFalse(Arrays.equals(new byte[12], iv));
   }
 
+  @Test
+  public void shortBufferDoesNotResetDecrypt() throws GeneralSecurityException {
+    final GCMParameterSpec spec = new GCMParameterSpec(128, randomIV());
+    amznC.init(Cipher.ENCRYPT_MODE, key, spec);
+    final byte[] plaintext = new byte[32];
+    final byte[] ciphertext = amznC.doFinal(plaintext);
+
+    amznC.init(Cipher.DECRYPT_MODE, key, spec);
+    amznC.update(ciphertext, 0, 16);
+
+    assertThrows(ShortBufferException.class, () -> amznC.doFinal(ciphertext, 8, 8, new byte[4]));
+
+    assertArraysHexEquals(plaintext, amznC.doFinal(ciphertext, 16, ciphertext.length - 16));
+  }
+
+  @Test
+  public void arrayIndexDoesNotResetDecrypt() throws GeneralSecurityException {
+    final GCMParameterSpec spec = new GCMParameterSpec(128, randomIV());
+    amznC.init(Cipher.ENCRYPT_MODE, key, spec);
+    final byte[] plaintext = new byte[32];
+    final byte[] ciphertext = amznC.doFinal(plaintext);
+
+    amznC.init(Cipher.DECRYPT_MODE, key, spec);
+    amznC.update(ciphertext, 0, 16);
+
+    assertThrows(
+        ArrayIndexOutOfBoundsException.class,
+        () -> amznC.doFinal(ciphertext, 8, 8, new byte[32], 36));
+
+    assertArraysHexEquals(plaintext, amznC.doFinal(ciphertext, 16, ciphertext.length - 16));
+  }
+
+  @Test
+  public void shortBufferDoesNotResetEncrypt() throws Exception {
+    final GCMParameterSpec spec = new GCMParameterSpec(128, randomIV());
+    amznC.init(Cipher.ENCRYPT_MODE, key, spec);
+    final byte[] plaintext = new byte[32];
+    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    baos.write(amznC.update(plaintext, 0, 16));
+
+    assertThrows(ShortBufferException.class, () -> amznC.doFinal(plaintext, 16, 16, new byte[4]));
+    baos.write(amznC.doFinal(plaintext, 16, 16));
+    final byte[] ciphertext = baos.toByteArray();
+
+    amznC.init(Cipher.DECRYPT_MODE, key, spec);
+
+    assertArraysHexEquals(plaintext, amznC.doFinal(ciphertext));
+  }
+
+  @Test
+  public void arrayIndexDoesNotResetEncrypt() throws Exception {
+    final GCMParameterSpec spec = new GCMParameterSpec(128, randomIV());
+    amznC.init(Cipher.ENCRYPT_MODE, key, spec);
+    final byte[] plaintext = new byte[32];
+    final ByteArrayOutputStream baos = new ByteArrayOutputStream();
+    baos.write(amznC.update(plaintext, 0, 16));
+
+    assertThrows(
+        ArrayIndexOutOfBoundsException.class,
+        () -> amznC.doFinal(plaintext, 16, 16, new byte[32], 36));
+    baos.write(amznC.doFinal(plaintext, 16, 16));
+    final byte[] ciphertext = baos.toByteArray();
+
+    amznC.init(Cipher.DECRYPT_MODE, key, spec);
+
+    assertArraysHexEquals(plaintext, amznC.doFinal(ciphertext));
+  }
+
   private byte[] randomIV() {
     return TestUtil.getRandomBytes(16);
   }