Restore KMS caching logic

Bryan Donlan · Bryan Donlan · commit 4b85273b0ae6 · 2018-08-03T13:34:48.000-07:00
We now verify that the requested region is reachable before caching the KMS
client.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,7 +2,9 @@
 
 ## 1.3.5
 
-(nothing yet)
+### Minor Changes
+
+* Restored the KMS client cache with a fix for the memory leak.
 
 ## 1.3.4
 
diff --git a/src/main/java/com/amazonaws/encryptionsdk/kms/KmsMasterKeyProvider.java b/src/main/java/com/amazonaws/encryptionsdk/kms/KmsMasterKeyProvider.java
@@ -26,7 +26,10 @@
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 
+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 +74,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 +266,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() {
diff --git a/src/test/java/com/amazonaws/encryptionsdk/AllTestsSuite.java b/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({
diff --git a/src/test/java/com/amazonaws/encryptionsdk/IntegrationTestSuite.java b/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({
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/KMSProviderBuilderIntegrationTests.java b/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();
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/KMSProviderBuilderMockTests.java b/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;
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/KMSTestFixtures.java b/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() {
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/LegacyKMSMasterKeyProviderTests.java b/src/test/java/com/amazonaws/encryptionsdk/kms/LegacyKMSMasterKeyProviderTests.java
@@ -11,7 +11,7 @@
  * specific language governing permissions and limitations under the License.
  */
 
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 import static com.amazonaws.encryptionsdk.CryptoAlgorithm.ALG_AES_128_GCM_IV12_TAG16_NO_KDF;
 import static com.amazonaws.encryptionsdk.internal.RandomBytesGenerator.generate;
@@ -34,11 +34,10 @@
 import com.amazonaws.encryptionsdk.MasterKeyProvider;
 import com.amazonaws.encryptionsdk.MasterKeyRequest;
 import com.amazonaws.encryptionsdk.jce.JceMasterKey;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKey;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
 import com.amazonaws.encryptionsdk.multi.MultipleProviderFactory;
 import com.amazonaws.regions.Region;
 import com.amazonaws.regions.Regions;
+import com.amazonaws.services.kms.AWSKMS;
 
 public class LegacyKMSMasterKeyProviderTests {
     private static final String WRAPPING_ALG = "AES/GCM/NoPadding";
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/MockKMSClient.java b/src/test/java/com/amazonaws/encryptionsdk/kms/MockKMSClient.java
@@ -11,7 +11,7 @@
  * specific language governing permissions and limitations under the License.
  */
 
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 import java.nio.ByteBuffer;
 import java.security.SecureRandom;
@@ -29,6 +29,7 @@
 import com.amazonaws.ResponseMetadata;
 import com.amazonaws.regions.Region;
 import com.amazonaws.regions.Regions;
+import com.amazonaws.services.kms.AWSKMSClient;
 import com.amazonaws.services.kms.model.CreateAliasRequest;
 import com.amazonaws.services.kms.model.CreateAliasResult;
 import com.amazonaws.services.kms.model.CreateGrantRequest;
diff --git a/src/test/java/com/amazonaws/encryptionsdk/kms/XCompatKmsDecryptTest.java b/src/test/java/com/amazonaws/encryptionsdk/kms/XCompatKmsDecryptTest.java
@@ -11,7 +11,7 @@
  * specific language governing permissions and limitations under the License.
  */
 
-package com.amazonaws.services.kms;
+package com.amazonaws.encryptionsdk.kms;
 
 import static org.junit.Assert.assertArrayEquals;
 
@@ -32,7 +32,6 @@
 
 import com.amazonaws.encryptionsdk.AwsCrypto;
 import com.amazonaws.encryptionsdk.CryptoResult;
-import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
 

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package com.amazonaws.services.kms;`
	`1`	`+package com.amazonaws.encryptionsdk.kms;`
`2`	`2`
`3`	`3`	`final class KMSTestFixtures {`
`4`	`4`	`private KMSTestFixtures() {`