Avoid creating multiple clients if KmsMasterKeys are created before client initialization

Note that, while previously we would throw an exception when a MK was requested
for a region that the MKP cannot service immediately on MK creation, we now
throw on first use of the MK.

Author: Bryan Donlan

CHANGELOG.md

@@ -5,6 +5,10 @@
 ### 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
 

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,6 +25,7 @@
 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;
@@ -517,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;
     }