Add Builder.clone() and KMSMKP.withGrantTokens()

Author: Bryan Donlan

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

@@ -13,11 +13,12 @@
 
 package com.amazonaws.encryptionsdk.kms;
 
+import static java.util.Arrays.asList;
+import static java.util.Collections.emptyList;
 import static java.util.Collections.singletonList;
 
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -39,7 +40,6 @@
 import com.amazonaws.encryptionsdk.exception.NoSuchMasterKeyException;
 import com.amazonaws.encryptionsdk.exception.UnsupportedProviderException;
 import com.amazonaws.handlers.RequestHandler2;
-import com.amazonaws.regions.DefaultAwsRegionProviderChain;
 import com.amazonaws.regions.Region;
 import com.amazonaws.regions.Regions;
 import com.amazonaws.services.kms.AWSKMS;
@@ -70,12 +70,27 @@ public interface RegionalClientSupplier {
         AWSKMS getClient(String regionName);
     }
 
-    public static class Builder implements Cloneable {
+    public static final class Builder implements Cloneable {
         private String defaultRegion_ = null;
         private RegionalClientSupplier regionalClientSupplier_ = null;
         private AWSKMSClientBuilder templateBuilder_ = null;
         private List<String> keyIds_ = new ArrayList<>();
-        private List<String> grantTokens_ = new ArrayList<>();
+
+        public Builder clone() {
+            try {
+                Builder cloned = (Builder)super.clone();
+
+                if (templateBuilder_ != null) {
+                    cloned.templateBuilder_ = cloneClientBuilder(templateBuilder_);
+                }
+
+                cloned.keyIds_ = new ArrayList<>(keyIds_);
+
+                return cloned;
+            } catch (CloneNotSupportedException e) {
+                throw new Error("Impossible: CloneNotSupportedException", e);
+            }
+        }
 
         /**
          * Adds key ID(s) to the list of keys to use on encryption.
@@ -84,7 +99,7 @@ public static class Builder implements Cloneable {
          * @return
          */
         public Builder withKeysForEncryption(String... keyIds) {
-            keyIds_.addAll(Arrays.asList(keyIds));
+            keyIds_.addAll(asList(keyIds));
             return this;
         }
 
@@ -99,17 +114,6 @@ public Builder withKeysForEncryption(List<String> keyIds) {
             return this;
         }
 
-        /**
-         * Adds grant tokens to the provider under construction.
-         *
-         * @param tokens
-         * @return
-         */
-        public Builder withGrantTokens(List<String> tokens) {
-            this.grantTokens_.addAll(tokens);
-            return this;
-        }
-
         /**
          * Sets the default region. This region will be used when specifying key IDs for encryption or in
          * {@link KmsMasterKeyProvider#getMasterKey(String)} that are not full ARNs, but are instead bare key IDs or
@@ -230,9 +234,19 @@ private AWSKMSClientBuilder cloneClientBuilder(final AWSKMSClientBuilder builder
          * @return
          */
         public KmsMasterKeyProvider build() {
+            // If we don't have a default region, we need to check that all key IDs will be usable
+            if (defaultRegion_ == null) {
+                for (String keyId : keyIds_) {
+                    if (parseRegionfromKeyArn(keyId) == null) {
+                        throw new AwsCryptoException("Can't use non-ARN key identifiers or aliases when " +
+                                                             "no default region is set");
+                    }
+                }
+            }
+
             RegionalClientSupplier supplier = clientFactory();
 
-            return new KmsMasterKeyProvider(supplier, defaultRegion_, keyIds_, grantTokens_, false);
+            return new KmsMasterKeyProvider(supplier, defaultRegion_, keyIds_, emptyList(), false);
         }
 
         private RegionalClientSupplier clientFactory() {
@@ -283,17 +297,12 @@ private KmsMasterKeyProvider(
         this.defaultRegion_ = defaultRegion;
         this.keyIds_ = Collections.unmodifiableList(new ArrayList<>(keyIds));
 
-        if (grantTokens != null) {
-            this.grantTokens_ = Collections.unmodifiableList(new ArrayList<>(grantTokens));
-        } else {
-            // Legacy support for the mutating API
-            this.grantTokens_ = new ArrayList<>();
-        }
+        this.grantTokens_ = grantTokens;
     }
 
     // Helper ctor for legacy ctors
     private KmsMasterKeyProvider(RegionalClientSupplier supplier, String defaultRegion, List<String> keyIds) {
-        this(supplier, defaultRegion, keyIds, null, true);
+        this(supplier, defaultRegion, keyIds, new ArrayList<>(), true);
     }
 
     private static RegionalClientSupplier defaultProvider() {
@@ -309,7 +318,7 @@ private static RegionalClientSupplier defaultProvider() {
      */
     @Deprecated
     public KmsMasterKeyProvider() {
-        this(defaultProvider(), Regions.DEFAULT_REGION.getName(), Collections.emptyList());
+        this(defaultProvider(), Regions.DEFAULT_REGION.getName(), emptyList());
     }
 
 
@@ -439,7 +448,7 @@ public KmsMasterKey getMasterKey(final String provider, final String keyId) thro
             throw new UnsupportedProviderException();
         }
 
-        String regionName = identifyKeyRegion(keyId);
+        String regionName = parseRegionfromKeyArn(keyId);
         AWSKMS kms = regionalClientSupplier_.getClient(regionName);
         if (kms == null) {
             throw new AwsCryptoException("Can't use keys from region " + regionName);
@@ -456,7 +465,7 @@ public KmsMasterKey getMasterKey(final String provider, final String keyId) thro
     @Override
     public List<KmsMasterKey> getMasterKeysForEncryption(final MasterKeyRequest request) {
         if (keyIds_ == null) {
-            return Collections.emptyList();
+            return emptyList();
         }
         List<KmsMasterKey> result = new ArrayList<>(keyIds_.size());
         for (String id : keyIds_) {
@@ -485,7 +494,7 @@ public DataKey<KmsMasterKey> decryptDataKey(final CryptoAlgorithm algorithm,
     }
 
     /**
-     * @deprecated This method is inherently not thread safe. Use {@link Builder#setGrantTokens(List)} instead.
+     * @deprecated This method is inherently not thread safe. Use {@link KmsMasterKey#setGrantTokens(List)} instead.
      * {@link KmsMasterKeyProvider}s constructed using the builder will throw an exception on attempts to modify the
      * list of grant tokens.
      */
@@ -496,9 +505,7 @@ public void setGrantTokens(final List<String> grantTokens) {
             this.grantTokens_.clear();
             this.grantTokens_.addAll(grantTokens);
         } catch (UnsupportedOperationException e) {
-            throw new IllegalStateException(
-                    "Changing grant tokens are not supported when constructing using the new builder API. Set them " +
-                            "at construction time instead.");
+            throw grantTokenError();
         }
     }
 
@@ -508,22 +515,48 @@ public List<String> getGrantTokens() {
     }
 
     /**
-     * @deprecated This method is inherently not thread safe. Use {@link Builder#setGrantTokens(List)} instead.
-     * {@link KmsMasterKeyProvider}s constructed using the builder will throw an exception on attempts to modify the
-     * list of grant tokens.
+     * @deprecated This method is inherently not thread safe. Use {@link #withGrantTokens(List)} or
+     * {@link KmsMasterKey#setGrantTokens(List)} instead. {@link KmsMasterKeyProvider}s constructed using the builder
+     * will throw an exception on attempts to modify the list of grant tokens.
      */
     @Deprecated
     @Override
     public void addGrantToken(final String grantToken) {
         try {
             grantTokens_.add(grantToken);
         } catch (UnsupportedOperationException e) {
-            throw new IllegalStateException(
-                    "Changing grant tokens are not supported when constructing using the new builder API. Set them " +
-                            "at construction time instead.");
+            throw grantTokenError();
         }
     }
 
+    private RuntimeException grantTokenError() {
+        return new IllegalStateException("This master key provider is immutable. Use withGrantTokens instead.");
+    }
+
+    /**
+     * Returns a new {@link KmsMasterKeyProvider} that is configured identically to this one, except with the given list
+     * of grant tokens. The grant token list in the returned provider is immutable (but can be further overridden by
+     * invoking withGrantTokens again).
+     * @param grantTokens
+     * @return
+     */
+    public KmsMasterKeyProvider withGrantTokens(List<String> grantTokens) {
+        grantTokens = Collections.unmodifiableList(new ArrayList<>(grantTokens));
+
+        return new KmsMasterKeyProvider(regionalClientSupplier_, defaultRegion_, keyIds_, grantTokens, false);
+    }
+
+    /**
+     * Returns a new {@link KmsMasterKeyProvider} that is configured identically to this one, except with the given list
+     * of grant tokens. The grant token list in the returned provider is immutable (but can be further overridden by
+     * invoking withGrantTokens again).
+     * @param grantTokens
+     * @return
+     */
+    public KmsMasterKeyProvider withGrantTokens(String... grantTokens) {
+        return withGrantTokens(asList(grantTokens));
+    }
+
     private static Region getStartingRegion(final String keyArn) {
         final String region = parseRegionfromKeyArn(keyArn);
         if (region != null) {
@@ -537,20 +570,6 @@ private static Region getStartingRegion(final String keyArn) {
         return Region.getRegion(Regions.DEFAULT_REGION);
     }
 
-    private String identifyKeyRegion(final String keyArn) {
-        String region = parseRegionfromKeyArn(keyArn);
-
-        if (region != null) {
-            return region;
-        }
-
-        if (defaultRegion_ == null) {
-            throw new AwsCryptoException("Can't use non-ARN key identifiers or aliases when no default region is set");
-        }
-
-        return defaultRegion_;
-    }
-
     private static String parseRegionfromKeyArn(final String keyArn) {
         final String[] parts = keyArn.split(":", 5);
 

src/test/java/com/amazonaws/services/kms/KMSProviderBuilderIntegrationTests.java

@@ -1,16 +1,17 @@
 package com.amazonaws.services.kms;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 
-import java.lang.reflect.Field;
 import java.util.Arrays;
 
-import org.junit.Assume;
 import org.junit.Test;
 
 import com.amazonaws.AbortedException;
@@ -20,23 +21,20 @@
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.client.builder.AwsClientBuilder;
 import com.amazonaws.encryptionsdk.AwsCrypto;
+import com.amazonaws.encryptionsdk.CryptoResult;
+import com.amazonaws.encryptionsdk.MasterKeyProvider;
 import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
 import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
 import com.amazonaws.handlers.RequestHandler2;
 import com.amazonaws.http.exception.HttpRequestTimeoutException;
-import com.amazonaws.regions.DefaultAwsRegionProviderChain;
 
 public class KMSProviderBuilderIntegrationTests {
-    public String[] keyIds = new String[] {
-            "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f",
-            "arn:aws:kms:eu-central-1:658956600833:key/75414c93-5285-4b57-99c9-30c1cf0a22c2"
-    };
 
     @Test
     public void whenConstructedWithoutArguments_canUseMultipleRegions() throws Exception {
         KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder().build();
 
-        for (String key : keyIds) {
+        for (String key : KMSTestFixtures.TEST_KEY_IDS) {
             byte[] ciphertext =
                     new AwsCrypto().encryptData(
                             KmsMasterKeyProvider.builder()
@@ -53,7 +51,7 @@ public void whenConstructedWithoutArguments_canUseMultipleRegions() throws Excep
     public void whenLegacyConstructorsUsed_multiRegionDecryptIsNotSupported() throws Exception {
         KmsMasterKeyProvider mkp = new KmsMasterKeyProvider();
 
-        for (String key : keyIds) {
+        for (String key : KMSTestFixtures.TEST_KEY_IDS) {
             byte[] ciphertext =
                     new AwsCrypto().encryptData(
                             KmsMasterKeyProvider.builder()
@@ -75,7 +73,7 @@ public void whenHandlerConfigured_handlerIsInvoked() throws Exception {
                                             AWSKMSClientBuilder.standard()
                                                 .withRequestHandlers(handler)
                                     )
-                                    .withKeysForEncryption(keyIds[0])
+                                    .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                     .build();
 
         new AwsCrypto().encryptData(mkp, new byte[1]);
@@ -95,7 +93,7 @@ public void whenShortTimeoutSet_timesOut() throws Exception {
                                                                             .withRequestTimeout(1)
                                                                 )
                                                        )
-                                                       .withKeysForEncryption(Arrays.asList(keyIds))
+                                                       .withKeysForEncryption(Arrays.asList(KMSTestFixtures.TEST_KEY_IDS))
                                                        .build();
 
         try {
@@ -118,7 +116,7 @@ public void whenCustomCredentialsSet_theyAreUsed() throws Exception {
 
         KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
                                                        .withCredentials(customProvider)
-                                                       .withKeysForEncryption(keyIds[0])
+                                                       .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                        .build();
 
         new AwsCrypto().encryptData(mkp, new byte[1]);
@@ -129,14 +127,63 @@ public void whenCustomCredentialsSet_theyAreUsed() throws Exception {
 
         mkp = KmsMasterKeyProvider.builder()
                                                        .withCredentials(customCredentials)
-                                                       .withKeysForEncryption(keyIds[0])
+                                                       .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
                                                        .build();
 
         new AwsCrypto().encryptData(mkp, new byte[1]);
 
         verify(customCredentials, atLeastOnce()).getAWSSecretKey();
     }
 
+    @Test
+    public void whenBuilderCloned_credentialsAndConfigurationAreRetained() throws Exception {
+        AWSCredentialsProvider customProvider1 = spy(new DefaultAWSCredentialsProviderChain());
+        AWSCredentialsProvider customProvider2 = spy(new DefaultAWSCredentialsProviderChain());
+
+        KmsMasterKeyProvider.Builder builder = KmsMasterKeyProvider.builder()
+                .withCredentials(customProvider1)
+                .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0]);
+
+        KmsMasterKeyProvider.Builder builder2 = builder.clone();
+
+        // This will mutate the first builder to add the new key and change the creds, but leave the clone unchanged.
+        MasterKeyProvider<?> mkp2 = builder.withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[1]).withCredentials(customProvider2).build();
+        MasterKeyProvider<?> mkp1 = builder2.build();
+
+        CryptoResult<byte[], ?> result = new AwsCrypto().encryptData(mkp1, new byte[0]);
+
+        assertEquals(KMSTestFixtures.TEST_KEY_IDS[0], result.getMasterKeyIds().get(0));
+        assertEquals(1, result.getMasterKeyIds().size());
+        verify(customProvider1, atLeastOnce()).getCredentials();
+        verify(customProvider2, never()).getCredentials();
+
+        reset(customProvider1, customProvider2);
+
+        result = new AwsCrypto().encryptData(mkp2, new byte[0]);
+
+        assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[0]));
+        assertTrue(result.getMasterKeyIds().contains(KMSTestFixtures.TEST_KEY_IDS[1]));
+        assertEquals(2, result.getMasterKeyIds().size());
+        verify(customProvider1, never()).getCredentials();
+        verify(customProvider2, atLeastOnce()).getCredentials();
+    }
+
+    @Test
+    public void whenBuilderCloned_clientBuilderCustomizationIsRetained() throws Exception {
+        RequestHandler2 handler = spy(new RequestHandler2() {});
+
+        KmsMasterKeyProvider mkp = KmsMasterKeyProvider.builder()
+                .withClientBuilder(
+                        AWSKMSClientBuilder.standard().withRequestHandlers(handler)
+                )
+                .withKeysForEncryption(KMSTestFixtures.TEST_KEY_IDS[0])
+                .clone().build();
+
+        new AwsCrypto().encryptData(mkp, new byte[0]);
+
+        verify(handler, atLeastOnce()).beforeRequest(any());
+    }
+
     @Test(expected = IllegalArgumentException.class)
     public void whenBogusEndpointIsSet_constructionFails() throws Exception {
         KmsMasterKeyProvider.builder()
@@ -155,3 +202,4 @@ public void whenDefaultRegionSet_itIsUsedForBareKeyIds() throws Exception {
         // TODO: Need to set up a role to assume as bare key IDs are relative to the caller account
     }
 }
+