Make estimateCiphertextSize work with cached keys

estimateCiphertextSize was not passing a plaintext size down to the CMM, so the
caching CMM assumed that a streaming encryption of unknown size was being
performed, and bypassed the cache entirely.

This change passes a plaintext size of zero instead to allow cached keys to be
used; since we don't actually encrypt any data it's safe to not consume any of
the byte limit.

Note that this may not behave quite right if the CMM does more clever things
with plaintext size. In general we should probably resolve this by moving
estimateCiphertextSize over to the Cipher*Stream objects instead, where it can
know the actual DataKeys in use instead of hoping they're reasonably
consistently sized.

Fixes: #29

Author: Bryan Donlan

src/main/java/com/amazonaws/encryptionsdk/AwsCrypto.java

@@ -171,6 +171,10 @@ public long estimateCiphertextSize(
         EncryptionMaterialsRequest request = EncryptionMaterialsRequest.newBuilder()
                                                                        .setContext(encryptionContext)
                                                                        .setRequestedAlgorithm(getEncryptionAlgorithm())
+        // We're not actually encrypting any data, so don't consume any bytes from the cache's limits. We do need to
+        // pass /something/ though, or the cache will be bypassed (as it'll assume this is a streaming encrypt of
+        // unknown size).
+                                                                       .setPlaintextSize(0)
                                                                        .build();
 
         final MessageCryptoHandler cryptoHandler = new EncryptionHandler(

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

@@ -36,12 +36,16 @@
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 
+import com.amazonaws.encryptionsdk.caching.CachingCryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.caching.LocalCryptoMaterialsCache;
 import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
 import com.amazonaws.encryptionsdk.exception.BadCiphertextException;
 import com.amazonaws.encryptionsdk.internal.StaticMasterKey;
@@ -61,7 +65,7 @@ public void init() {
         masterKeyProvider = spy(new StaticMasterKey("testmaterial"));
 
         encryptionClient_ = new AwsCrypto();
-        encryptionClient_.setEncryptionAlgorithm(CryptoAlgorithm.ALG_AES_192_GCM_IV12_TAG16_NO_KDF);
+        encryptionClient_.setEncryptionAlgorithm(CryptoAlgorithm.ALG_AES_128_GCM_IV12_TAG16_HKDF_SHA256);
     }
 
     private void doEncryptDecrypt(final CryptoAlgorithm cryptoAlg, final int byteSize, final int frameSize) {
@@ -374,6 +378,28 @@ public void estimateCiphertextSizeWithoutEncContext() {
         assertTrue(errMsg, estimatedCiphertextSize - cipherText.length <= 16);
     }
 
+    @Test
+    public void estimateCiphertextSize_usesCachedKeys() throws Exception {
+        // Make sure estimateCiphertextSize works with cached CMMs
+        CryptoMaterialsManager cmm = spy(new DefaultCryptoMaterialsManager(masterKeyProvider));
+
+        CachingCryptoMaterialsManager cache = CachingCryptoMaterialsManager.newBuilder()
+                .withBackingMaterialsManager(cmm)
+                .withMaxAge(Long.MAX_VALUE, TimeUnit.SECONDS)
+                .withCache(new LocalCryptoMaterialsCache(1))
+                .withMessageUseLimit(9999)
+                .withByteUseLimit(501)
+                .build();
+
+        // These estimates should be cached, and should not consume any bytes from the byte use limit.
+        encryptionClient_.estimateCiphertextSize(cache, 500, new HashMap<>());
+        encryptionClient_.estimateCiphertextSize(cache, 500, new HashMap<>());
+
+        encryptionClient_.encryptData(cache, new byte[500]);
+
+        verify(cmm, times(1)).getMaterialsForEncrypt(any());
+    }
+
     @Test
     public void encryptDecryptWithoutEncContext() {
         final int ptSize = 1000000; // 1MB