Merge pull request #67 from bdonlan/kms-cache-restore

Restore KMS caching logic

Author: SalusaSecondus

CHANGELOG.md

@@ -1,9 +1,21 @@
 # Changelog
 
-## 1.3.5
+## 1.3.6
+
+### Minor Changes
 
 (nothing yet)
 
+## 1.3.5
+
+### Minor Changes
+
+* Restored the KMS client cache with a fix for the memory leak.
+* When using a master key provider that can only service a subset of regions
+(e.g. using the deprecated constructors), and requesting a master key from a
+region not servicable by that MKP, the exception will now be thrown on first
+use of the MK, rather than at getMasterKey time.
+
 ## 1.3.4
 
 ### Minor Changes

README.md

@@ -45,7 +45,7 @@ You can get the latest release from Maven:
 <dependency>
   <groupId>com.amazonaws</groupId>
   <artifactId>aws-encryption-sdk-java</artifactId>
-  <version>1.3.4</version>
+  <version>1.3.5</version>
 </dependency>
 ```
 

pom.xml

@@ -4,7 +4,7 @@
 
     <groupId>com.amazonaws</groupId>
     <artifactId>aws-encryption-sdk-java</artifactId>
-    <version>1.3.5-SNAPSHOT</version>
+    <version>1.3.6-SNAPSHOT</version>
     <packaging>jar</packaging>
 
     <name>aws-encryption-sdk-java</name>

src/main/java/com/amazonaws/encryptionsdk/kms/KmsMasterKey.java

@@ -21,6 +21,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 import com.amazonaws.AmazonServiceException;
 import com.amazonaws.AmazonWebServiceRequest;
@@ -48,7 +49,7 @@
  * {@link AwsCrypto}.
  */
 public final class KmsMasterKey extends MasterKey<KmsMasterKey> implements KmsMethods {
-    private final AWSKMS kms_;
+    private final Supplier<AWSKMS> kms_;
     private final MasterKeyProvider<KmsMasterKey> sourceProvider_;
     private final String id_;
     private final List<String> grantTokens_ = new ArrayList<>();
@@ -77,12 +78,12 @@ public static KmsMasterKey getInstance(final AWSCredentialsProvider creds, final
         return new KmsMasterKeyProvider(creds, keyId).getMasterKey(keyId);
     }
 
-    static KmsMasterKey getInstance(final AWSKMS kms, final String id,
+    static KmsMasterKey getInstance(final Supplier<AWSKMS> kms, final String id,
             final MasterKeyProvider<KmsMasterKey> provider) {
         return new KmsMasterKey(kms, id, provider);
     }
 
-    private KmsMasterKey(final AWSKMS kms, final String id, final MasterKeyProvider<KmsMasterKey> provider) {
+    private KmsMasterKey(final Supplier<AWSKMS> kms, final String id, final MasterKeyProvider<KmsMasterKey> provider) {
         kms_ = kms;
         id_ = id;
         sourceProvider_ = provider;
@@ -101,7 +102,7 @@ public String getKeyId() {
     @Override
     public DataKey<KmsMasterKey> generateDataKey(final CryptoAlgorithm algorithm,
             final Map<String, String> encryptionContext) {
-        final GenerateDataKeyResult gdkResult = kms_.generateDataKey(updateUserAgent(
+        final GenerateDataKeyResult gdkResult = kms_.get().generateDataKey(updateUserAgent(
                 new GenerateDataKeyRequest()
                         .withKeyId(getKeyId())
                         .withNumberOfBytes(algorithm.getDataKeyLength())
@@ -145,7 +146,7 @@ public DataKey<KmsMasterKey> encryptDataKey(final CryptoAlgorithm algorithm,
             throw new IllegalArgumentException("Only RAW encoded keys are supported");
         }
         try {
-            final EncryptResult encryptResult = kms_.encrypt(updateUserAgent(
+            final EncryptResult encryptResult = kms_.get().encrypt(updateUserAgent(
                     new EncryptRequest()
                             .withKeyId(id_)
                             .withPlaintext(ByteBuffer.wrap(key.getEncoded()))
@@ -167,7 +168,7 @@ public DataKey<KmsMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
         final List<Exception> exceptions = new ArrayList<>();
         for (final EncryptedDataKey edk : encryptedDataKeys) {
             try {
-                final DecryptResult decryptResult = kms_.decrypt(updateUserAgent(
+                final DecryptResult decryptResult = kms_.get().decrypt(updateUserAgent(
                         new DecryptRequest()
                                 .withCiphertextBlob(ByteBuffer.wrap(edk.getEncryptedDataKey()))
                                 .withEncryptionContext(encryptionContext)

src/main/java/com/amazonaws/encryptionsdk/kms/KmsMasterKeyProvider.java

@@ -25,8 +25,12 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 
+import com.amazonaws.AmazonServiceException;
 import com.amazonaws.ClientConfiguration;
+import com.amazonaws.Request;
+import com.amazonaws.Response;
 import com.amazonaws.auth.AWSCredentials;
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.AWSStaticCredentialsProvider;
@@ -71,12 +75,16 @@ public interface RegionalClientSupplier {
         AWSKMS getClient(String regionName);
     }
 
-    public static final class Builder implements Cloneable {
+    public static class Builder implements Cloneable {
         private String defaultRegion_ = null;
         private RegionalClientSupplier regionalClientSupplier_ = null;
         private AWSKMSClientBuilder templateBuilder_ = null;
         private List<String> keyIds_ = new ArrayList<>();
 
+        Builder() {
+            // Default access: Don't allow outside classes to extend this class
+        }
+
         public Builder clone() {
             try {
                 Builder cloned = (Builder) super.clone();
@@ -259,11 +267,68 @@ private RegionalClientSupplier clientFactory() {
             AWSKMSClientBuilder builder = templateBuilder_ != null ? cloneClientBuilder(templateBuilder_)
                                                                    : AWSKMSClientBuilder.standard();
 
+            ConcurrentHashMap<String, AWSKMS> clientCache = new ConcurrentHashMap<>();
+            snoopClientCache(clientCache);
+
             return region -> {
-                // Clone yet again as we're going to change the region field.
-                return cloneClientBuilder(builder).withRegion(region).build();
+                AWSKMS kms = clientCache.get(region);
+
+                if (kms != null) return kms;
+
+                // We can't just use computeIfAbsent as we need to avoid leaking KMS clients if we're asked to decrypt
+                // an EDK with a bogus region in its ARN. So we'll install a request handler to identify the first
+                // successful call, and cache it when we see that.
+                SuccessfulRequestCacher cacher = new SuccessfulRequestCacher(clientCache, region);
+                ArrayList<RequestHandler2> handlers = new ArrayList<>();
+                if (builder.getRequestHandlers() != null) {
+                    handlers.addAll(builder.getRequestHandlers());
+                }
+                handlers.add(cacher);
+
+                kms = cloneClientBuilder(builder)
+                        .withRegion(region)
+                        .withRequestHandlers(handlers.toArray(new RequestHandler2[handlers.size()]))
+                        .build();
+                cacher.client_ = kms;
+
+                return kms;
             };
         }
+
+        protected void snoopClientCache(ConcurrentHashMap<String, AWSKMS> map) {
+            // no-op - this is a test hook
+        }
+    }
+
+    private static class SuccessfulRequestCacher extends RequestHandler2 {
+        private final ConcurrentHashMap<String, AWSKMS> cache_;
+        private final String region_;
+        private AWSKMS client_;
+
+        volatile boolean ranBefore_ = false;
+
+        private SuccessfulRequestCacher(
+                final ConcurrentHashMap<String, AWSKMS> cache,
+                final String region
+        ) {
+            this.region_ = region;
+            this.cache_ = cache;
+        }
+
+        @Override public void afterResponse(final Request<?> request, final Response<?> response) {
+            if (ranBefore_) return;
+            ranBefore_ = true;
+
+            cache_.putIfAbsent(region_, client_);
+        }
+
+        @Override public void afterError(final Request<?> request, final Response<?> response, final Exception e) {
+            if (ranBefore_) return;
+            if (e instanceof AmazonServiceException) {
+                ranBefore_ = true;
+                cache_.putIfAbsent(region_, client_);
+            }
+        }
     }
 
     public static Builder builder() {
@@ -453,12 +518,17 @@ public KmsMasterKey getMasterKey(final String provider, final String keyId) thro
             regionName = defaultRegion_;
         }
 
-        AWSKMS kms = regionalClientSupplier_.getClient(regionName);
-        if (kms == null) {
-            throw new AwsCryptoException("Can't use keys from region " + regionName);
-        }
+        String regionName_ = regionName;
+
+        Supplier<AWSKMS> kmsSupplier = () -> {
+            AWSKMS kms = regionalClientSupplier_.getClient(regionName_);
+            if (kms == null) {
+                throw new AwsCryptoException("Can't use keys from region " + regionName_);
+            }
+            return kms;
+        };
 
-        final KmsMasterKey result = KmsMasterKey.getInstance(kms, keyId, this);
+        final KmsMasterKey result = KmsMasterKey.getInstance(kmsSupplier, keyId, this);
         result.setGrantTokens(grantTokens_);
         return result;
     }

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

@@ -23,7 +23,7 @@
 import com.amazonaws.encryptionsdk.model.CipherFrameHeadersTest;
 import com.amazonaws.encryptionsdk.model.KeyBlobTest;
 import com.amazonaws.encryptionsdk.multi.MultipleMasterKeyTest;
-import com.amazonaws.services.kms.KMSProviderBuilderMockTests;
+import com.amazonaws.encryptionsdk.kms.KMSProviderBuilderMockTests;
 
 @RunWith(Suite.class)
 @Suite.SuiteClasses({

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

@@ -3,8 +3,8 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.Suite;
 
-import com.amazonaws.services.kms.KMSProviderBuilderIntegrationTests;
-import com.amazonaws.services.kms.XCompatKmsDecryptTest;
+import com.amazonaws.encryptionsdk.kms.KMSProviderBuilderIntegrationTests;
+import com.amazonaws.encryptionsdk.kms.XCompatKmsDecryptTest;
 
 @RunWith(Suite.class)
 @Suite.SuiteClasses({

src/test/java/com/amazonaws/services/kms/KMSProviderBuilderIntegrationTests.java renamed to src/test/java/com/amazonaws/encryptionsdk/kms/KMSProviderBuilderIntegrationTests.java

@@ -1,6 +1,8 @@
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
@@ -10,29 +12,91 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 
 import com.amazonaws.AbortedException;
-import com.amazonaws.AmazonWebServiceRequest;
 import com.amazonaws.ClientConfiguration;
 import com.amazonaws.Request;
 import com.amazonaws.auth.AWSCredentials;
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.client.builder.AwsClientBuilder;
 import com.amazonaws.encryptionsdk.AwsCrypto;
+import com.amazonaws.encryptionsdk.CryptoAlgorithm;
 import com.amazonaws.encryptionsdk.CryptoResult;
 import com.amazonaws.encryptionsdk.MasterKeyProvider;
 import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
 import com.amazonaws.encryptionsdk.internal.VersionInfo;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
+import com.amazonaws.encryptionsdk.model.KeyBlob;
 import com.amazonaws.handlers.RequestHandler2;
 import com.amazonaws.http.exception.HttpRequestTimeoutException;
+import com.amazonaws.services.kms.AWSKMS;
+import com.amazonaws.services.kms.AWSKMSClientBuilder;
 
 public class KMSProviderBuilderIntegrationTests {
+    @Test
+    public void whenBogusRegionsDecrypted_doesNotLeakClients() throws Exception {
+        AtomicReference<ConcurrentHashMap<String, AWSKMS>> kmsCache = new AtomicReference<>();
+
+        KmsMasterKeyProvider mkp = (new KmsMasterKeyProvider.Builder() {
+            @Override protected void snoopClientCache(
+                    final ConcurrentHashMap<String, AWSKMS> map
+            ) {
+                kmsCache.set(map);
+            }
+        }).build();
+
+        try {
+            mkp.decryptDataKey(
+                    CryptoAlgorithm.ALG_AES_128_GCM_IV12_TAG16_HKDF_SHA256,
+                    Collections.singleton(
+                            new KeyBlob("aws-kms",
+                                        "arn:aws:kms:us-bogus-1:123456789010:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f"
+                                                .getBytes(StandardCharsets.UTF_8),
+                                        new byte[40]
+                            )
+                    ),
+                    new HashMap<>()
+            );
+            fail("Expected CannotUnwrapDataKeyException");
+        } catch (CannotUnwrapDataKeyException e) {
+            // ok
+        }
+
+        assertTrue(kmsCache.get().isEmpty());
+    }
+
+    @Test
+    public void whenOperationSuccessful_clientIsCached() {
+        AtomicReference<ConcurrentHashMap<String, AWSKMS>> kmsCache = new AtomicReference<>();
+
+        KmsMasterKeyProvider mkp = (new KmsMasterKeyProvider.Builder() {
+            @Override protected void snoopClientCache(
+                    final ConcurrentHashMap<String, AWSKMS> map
+            ) {
+                kmsCache.set(map);
+            }
+        }).withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
+          .build();
+
+        new AwsCrypto().encryptData(mkp, new byte[1]);
+
+        AWSKMS kms = kmsCache.get().get("us-west-2");
+        assertNotNull(kms);
+
+        new AwsCrypto().encryptData(mkp, new byte[1]);
+
+        // Cache entry should stay the same
+        assertEquals(kms, kmsCache.get().get("us-west-2"));
+    }
 
     @Test
     public void whenConstructedWithoutArguments_canUseMultipleRegions() throws Exception {
@@ -75,7 +139,7 @@ public void whenHandlerConfigured_handlerIsInvoked() throws Exception {
                 KmsMasterKeyProvider.builder()
                                     .withClientBuilder(
                                             AWSKMSClientBuilder.standard()
-                                                .withRequestHandlers(handler)
+                                                               .withRequestHandlers(handler)
                                     )
                                     .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                     .build();

src/test/java/com/amazonaws/services/kms/KMSProviderBuilderMockTests.java renamed to src/test/java/com/amazonaws/encryptionsdk/kms/KMSProviderBuilderMockTests.java

@@ -1,8 +1,7 @@
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 import static com.amazonaws.encryptionsdk.multi.MultipleProviderFactory.buildMultiProvider;
 import static com.amazonaws.regions.Region.getRegion;
-import static com.amazonaws.regions.Regions.DEFAULT_REGION;
 import static com.amazonaws.regions.Regions.fromName;
 import static java.util.Collections.singletonList;
 import static org.junit.Assert.assertEquals;
@@ -30,8 +29,6 @@
 import com.amazonaws.encryptionsdk.AwsCrypto;
 import com.amazonaws.encryptionsdk.MasterKeyProvider;
 import com.amazonaws.encryptionsdk.internal.VersionInfo;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKey;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
 import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider.RegionalClientSupplier;
 import com.amazonaws.regions.Region;
 import com.amazonaws.regions.Regions;

src/test/java/com/amazonaws/services/kms/KMSTestFixtures.java renamed to src/test/java/com/amazonaws/encryptionsdk/kms/KMSTestFixtures.java

@@ -1,4 +1,4 @@
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 final class KMSTestFixtures {
     private KMSTestFixtures() {