Create AwsKmsCmkId type to represent AWS KMS Key Ids

Author: WesleyRosenblum

src/examples/java/com/amazonaws/crypto/examples/BasicEncryptionExample.java

@@ -24,6 +24,7 @@
 import com.amazonaws.encryptionsdk.EncryptRequest;
 import com.amazonaws.encryptionsdk.keyrings.Keyring;
 import com.amazonaws.encryptionsdk.keyrings.StandardKeyrings;
+import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
 
 /**
  * <p>
@@ -41,12 +42,10 @@ public class BasicEncryptionExample {
     private static final byte[] EXAMPLE_DATA = "Hello World".getBytes(StandardCharsets.UTF_8);
 
     public static void main(final String[] args) {
-        final String keyArn = args[0];
-
-        encryptAndDecrypt(keyArn);
+        encryptAndDecrypt(AwsKmsCmkId.fromString(args[0]));
     }
 
-    static void encryptAndDecrypt(final String keyArn) {
+    static void encryptAndDecrypt(final AwsKmsCmkId keyArn) {
         // 1. Instantiate the SDK
         final AwsCrypto crypto = new AwsCrypto();
 
@@ -81,7 +80,7 @@ static void encryptAndDecrypt(final String keyArn) {
 
         // 6. Before verifying the plaintext, inspect the Keyring Trace to verify that the CMK used
         //    to decrypt the encrypted data key was the CMK in the encryption keyring.
-        if(!decryptResult.getKeyringTrace().getEntries().get(0).getKeyName().equals(keyArn)) {
+        if(!decryptResult.getKeyringTrace().getEntries().get(0).getKeyName().equals(keyArn.toString())) {
             throw new IllegalStateException("Wrong key ID!");
         }
 

src/examples/java/com/amazonaws/crypto/examples/EscrowedEncryptExample.java

@@ -18,6 +18,7 @@
 import com.amazonaws.encryptionsdk.EncryptRequest;
 import com.amazonaws.encryptionsdk.keyrings.Keyring;
 import com.amazonaws.encryptionsdk.keyrings.StandardKeyrings;
+import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
 
 import java.nio.charset.StandardCharsets;
 import java.security.GeneralSecurityException;
@@ -55,12 +56,10 @@ public class EscrowedEncryptExample {
     private static final byte[] EXAMPLE_DATA = "Hello World".getBytes(StandardCharsets.UTF_8);
 
     public static void main(final String[] args) throws GeneralSecurityException {
-        final String kmsArn = args[0];
-
-        escrowEncryptAndDecrypt(kmsArn);
+        escrowEncryptAndDecrypt(AwsKmsCmkId.fromString(args[0]));
     }
 
-    static void escrowEncryptAndDecrypt(String kmsArn) throws GeneralSecurityException {
+    static void escrowEncryptAndDecrypt(AwsKmsCmkId kmsArn) throws GeneralSecurityException {
         // This sample generates a new random key for each operation.
         // In practice, you would distribute the public key and save the private key in secure storage.
         final KeyPair escrowKeyPair = generateEscrowKeyPair();
@@ -79,7 +78,7 @@ static void escrowEncryptAndDecrypt(String kmsArn) throws GeneralSecurityExcepti
         assert Arrays.equals(escrowedDecryptedData, EXAMPLE_DATA);
     }
 
-    private static byte[] standardEncrypt(final String kmsArn, final PublicKey publicEscrowKey) {
+    private static byte[] standardEncrypt(final AwsKmsCmkId kmsArn, final PublicKey publicEscrowKey) {
         // Encrypt with the KMS CMK and the escrowed public key
 
         // 1. Instantiate the SDK
@@ -110,7 +109,7 @@ private static byte[] standardEncrypt(final String kmsArn, final PublicKey publi
                 .getResult();
     }
 
-    private static byte[] standardDecrypt(final String kmsArn, final byte[] cipherText) {
+    private static byte[] standardDecrypt(final AwsKmsCmkId kmsArn, final byte[] cipherText) {
         // Decrypt with the KMS CMK
 
         // 1. Instantiate the SDK

src/main/java/com/amazonaws/encryptionsdk/internal/Constants.java

@@ -62,4 +62,9 @@ private Constants() {
     public static final long MAX_FRAME_NUMBER = (1L << 32) - 1;
 
     public static final String EC_PUBLIC_KEY_FIELD = "aws-crypto-public-key";
+
+    /**
+     * The provider ID used for the AwsKmsKeyring
+     */
+    public static final String AWS_KMS_PROVIDER_ID = "aws-kms";
 }

src/main/java/com/amazonaws/encryptionsdk/keyrings/AwsKmsKeyring.java

@@ -16,11 +16,10 @@
 import com.amazonaws.encryptionsdk.EncryptedDataKey;
 import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
 import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
-import com.amazonaws.encryptionsdk.exception.MalformedArnException;
+import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
 import com.amazonaws.encryptionsdk.kms.DataKeyEncryptionDao;
 import com.amazonaws.encryptionsdk.kms.DataKeyEncryptionDao.DecryptDataKeyResult;
 import com.amazonaws.encryptionsdk.kms.DataKeyEncryptionDao.GenerateDataKeyResult;
-import com.amazonaws.encryptionsdk.kms.KmsUtils;
 import com.amazonaws.encryptionsdk.model.DecryptionMaterials;
 import com.amazonaws.encryptionsdk.model.EncryptionMaterials;
 import com.amazonaws.encryptionsdk.model.KeyBlob;
@@ -31,8 +30,8 @@
 import java.util.Set;
 
 import static com.amazonaws.encryptionsdk.EncryptedDataKey.PROVIDER_ENCODING;
-import static com.amazonaws.encryptionsdk.kms.KmsUtils.KMS_PROVIDER_ID;
-import static com.amazonaws.encryptionsdk.kms.KmsUtils.isArnWellFormed;
+import static com.amazonaws.encryptionsdk.internal.Constants.AWS_KMS_PROVIDER_ID;
+import static com.amazonaws.encryptionsdk.kms.AwsKmsCmkId.isKeyIdWellFormed;
 import static java.util.Collections.emptyList;
 import static java.util.Collections.unmodifiableList;
 import static java.util.Objects.requireNonNull;
@@ -44,28 +43,19 @@
 class AwsKmsKeyring implements Keyring {
 
     private final DataKeyEncryptionDao dataKeyEncryptionDao;
-    private final List<String> keyIds;
-    private final String generatorKeyId;
+    private final List<AwsKmsCmkId> keyIds;
+    private final AwsKmsCmkId generatorKeyId;
     private final boolean isDiscovery;
 
-    AwsKmsKeyring(DataKeyEncryptionDao dataKeyEncryptionDao, List<String> keyIds, String generatorKeyId) {
+    AwsKmsKeyring(DataKeyEncryptionDao dataKeyEncryptionDao, List<AwsKmsCmkId> keyIds, AwsKmsCmkId generatorKeyId) {
         requireNonNull(dataKeyEncryptionDao, "dataKeyEncryptionDao is required");
         this.dataKeyEncryptionDao = dataKeyEncryptionDao;
         this.keyIds = keyIds == null ? emptyList() : unmodifiableList(new ArrayList<>(keyIds));
         this.generatorKeyId = generatorKeyId;
         this.isDiscovery = this.generatorKeyId == null && this.keyIds.isEmpty();
 
-        if (!this.keyIds.stream().allMatch(KmsUtils::isArnWellFormed)) {
-            throw new MalformedArnException("keyIds must contain only CMK aliases and well formed ARNs");
-        }
-
-        if (generatorKeyId != null) {
-            if (!isArnWellFormed(generatorKeyId)) {
-                throw new MalformedArnException("generatorKeyId must be either a CMK alias or a well formed ARN");
-            }
-            if (this.keyIds.contains(generatorKeyId)) {
-                throw new IllegalArgumentException("KeyIds should not contain the generatorKeyId");
-            }
+        if (this.keyIds.contains(generatorKeyId)) {
+            throw new IllegalArgumentException("KeyIds should not contain the generatorKeyId");
         }
     }
 
@@ -86,7 +76,7 @@ public EncryptionMaterials onEncrypt(EncryptionMaterials encryptionMaterials) {
             throw new AwsCryptoException("Encryption materials must contain either a plaintext data key or a generator");
         }
 
-        final List<String> keyIdsToEncrypt = new ArrayList<>(keyIds);
+        final List<AwsKmsCmkId> keyIdsToEncrypt = new ArrayList<>(keyIds);
 
         // If the input encryption materials do not contain a plaintext data key and a generator is defined onEncrypt
         // MUST attempt to generate a new plaintext data key and encrypt that data key by calling KMS GenerateDataKey.
@@ -100,7 +90,7 @@ public EncryptionMaterials onEncrypt(EncryptionMaterials encryptionMaterials) {
 
         // Given a plaintext data key in the encryption materials, OnEncrypt MUST attempt
         // to encrypt the plaintext data key using each CMK specified in it's key IDs list.
-        for (String keyId : keyIdsToEncrypt) {
+        for (AwsKmsCmkId keyId : keyIdsToEncrypt) {
             resultMaterials = encryptDataKey(keyId, resultMaterials);
         }
 
@@ -113,17 +103,20 @@ private EncryptionMaterials generateDataKey(final EncryptionMaterials encryption
 
         return encryptionMaterials
                 .withCleartextDataKey(result.getPlaintextDataKey(),
-                        new KeyringTraceEntry(KMS_PROVIDER_ID, generatorKeyId, KeyringTraceFlag.GENERATED_DATA_KEY))
+                        new KeyringTraceEntry(AWS_KMS_PROVIDER_ID, generatorKeyId.toString(),
+                                KeyringTraceFlag.GENERATED_DATA_KEY))
                 .withEncryptedDataKey(new KeyBlob(result.getEncryptedDataKey()),
-                        new KeyringTraceEntry(KMS_PROVIDER_ID, generatorKeyId, KeyringTraceFlag.ENCRYPTED_DATA_KEY, KeyringTraceFlag.SIGNED_ENCRYPTION_CONTEXT));
+                        new KeyringTraceEntry(AWS_KMS_PROVIDER_ID, generatorKeyId.toString(),
+                                KeyringTraceFlag.ENCRYPTED_DATA_KEY, KeyringTraceFlag.SIGNED_ENCRYPTION_CONTEXT));
     }
 
-    private EncryptionMaterials encryptDataKey(final String keyId, final EncryptionMaterials encryptionMaterials) {
+    private EncryptionMaterials encryptDataKey(final AwsKmsCmkId keyId, final EncryptionMaterials encryptionMaterials) {
         final EncryptedDataKey encryptedDataKey = dataKeyEncryptionDao.encryptDataKey(keyId,
                 encryptionMaterials.getCleartextDataKey(), encryptionMaterials.getEncryptionContext());
 
         return encryptionMaterials.withEncryptedDataKey(new KeyBlob(encryptedDataKey),
-                new KeyringTraceEntry(KMS_PROVIDER_ID, keyId, KeyringTraceFlag.ENCRYPTED_DATA_KEY, KeyringTraceFlag.SIGNED_ENCRYPTION_CONTEXT));
+                new KeyringTraceEntry(AWS_KMS_PROVIDER_ID, keyId.toString(),
+                        KeyringTraceFlag.ENCRYPTED_DATA_KEY, KeyringTraceFlag.SIGNED_ENCRYPTION_CONTEXT));
     }
 
     @Override
@@ -135,7 +128,7 @@ public DecryptionMaterials onDecrypt(DecryptionMaterials decryptionMaterials, Li
             return decryptionMaterials;
         }
 
-        final Set<String> configuredKeyIds = new HashSet<>(keyIds);
+        final Set<AwsKmsCmkId> configuredKeyIds = new HashSet<>(keyIds);
 
         if (generatorKeyId != null) {
             configuredKeyIds.add(generatorKeyId);
@@ -148,7 +141,7 @@ public DecryptionMaterials onDecrypt(DecryptionMaterials decryptionMaterials, Li
                             decryptionMaterials.getAlgorithm(), decryptionMaterials.getEncryptionContext());
 
                     return decryptionMaterials.withCleartextDataKey(result.getPlaintextDataKey(),
-                            new KeyringTraceEntry(KMS_PROVIDER_ID, result.getKeyArn(),
+                            new KeyringTraceEntry(AWS_KMS_PROVIDER_ID, result.getKeyArn(),
                                     KeyringTraceFlag.DECRYPTED_DATA_KEY, KeyringTraceFlag.VERIFIED_ENCRYPTION_CONTEXT));
                 } catch (CannotUnwrapDataKeyException e) {
                     continue;
@@ -159,14 +152,14 @@ public DecryptionMaterials onDecrypt(DecryptionMaterials decryptionMaterials, Li
         return decryptionMaterials;
     }
 
-    private boolean okToDecrypt(EncryptedDataKey encryptedDataKey, Set<String> configuredKeyIds) {
+    private boolean okToDecrypt(EncryptedDataKey encryptedDataKey, Set<AwsKmsCmkId> configuredKeyIds) {
         // Only attempt to decrypt keys provided by KMS
-        if (!encryptedDataKey.getProviderId().equals(KMS_PROVIDER_ID)) {
+        if (!encryptedDataKey.getProviderId().equals(AWS_KMS_PROVIDER_ID)) {
             return false;
         }
 
-        // If the key ARN cannot be parsed, skip it
-        if(!isArnWellFormed(new String(encryptedDataKey.getProviderInformation(), PROVIDER_ENCODING)))
+        // If the key ID cannot be parsed, skip it
+        if(!isKeyIdWellFormed(new String(encryptedDataKey.getProviderInformation(), PROVIDER_ENCODING)))
         {
             return false;
         }
@@ -180,6 +173,7 @@ private boolean okToDecrypt(EncryptedDataKey encryptedDataKey, Set<String> confi
         // OnDecrypt MUST attempt to decrypt each input encrypted data key in the input
         // encrypted data key list where the key provider info has a value equal to one
         // of the ARNs in this keyring's key IDs or the generator
-        return configuredKeyIds.contains(new String(encryptedDataKey.getProviderInformation(), PROVIDER_ENCODING));
+        return configuredKeyIds.contains(
+                AwsKmsCmkId.fromString(new String(encryptedDataKey.getProviderInformation(), PROVIDER_ENCODING)));
     }
 }

src/main/java/com/amazonaws/encryptionsdk/keyrings/AwsKmsKeyringBuilder.java

@@ -14,15 +14,16 @@
 package com.amazonaws.encryptionsdk.keyrings;
 
 import com.amazonaws.encryptionsdk.kms.AwsKmsClientSupplier;
+import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
 import com.amazonaws.encryptionsdk.kms.DataKeyEncryptionDao;
 
 import java.util.List;
 
 public class AwsKmsKeyringBuilder {
     private AwsKmsClientSupplier awsKmsClientSupplier;
     private List<String> grantTokens;
-    private List<String> keyIds;
-    private String generatorKeyId;
+    private List<AwsKmsCmkId> keyIds;
+    private AwsKmsCmkId generatorKeyId;
 
     AwsKmsKeyringBuilder() {
         // Use StandardKeyrings.awsKms() to instantiate
@@ -53,25 +54,28 @@ public AwsKmsKeyringBuilder grantTokens(List<String> grantTokens) {
     }
 
     /**
-     * A list of strings identifying AWS KMS CMKs used for encrypting and decrypting data keys
-     * in ARN, CMK Alias, or ARN Alias format.
+     * A list of {@link AwsKmsCmkId}s in ARN, CMK Alias, or ARN Alias format identifying AWS KMS CMKs
+     * used for encrypting and decrypting data keys.
      *
      * @param keyIds The list of AWS KMS CMKs
      * @return The AwsKmsKeyringBuilder, for method chaining
      */
-    public AwsKmsKeyringBuilder keyIds(List<String> keyIds) {
+    public AwsKmsKeyringBuilder keyIds(List<AwsKmsCmkId> keyIds) {
         this.keyIds = keyIds;
         return this;
     }
 
     /**
-     * A string that identifies a AWS KMS CMK responsible for generating a data key,
-     * as well as encrypting and decrypting data keys in ARN, CMK Alias, or ARN Alias format.
+     * An {@link AwsKmsCmkId} in ARN, CMK Alias, or ARN Alias format that identifies a
+     * AWS KMS CMK responsible for generating a data key, as well as encrypting and
+     * decrypting data keys .
      *
-     * @param generatorKeyId The generator AWS KMS CMK
+     * @param generatorKeyId An {@link AwsKmsCmkId} in ARN, CMK Alias, or ARN Alias format that identifies a
+     *                       AWS KMS CMK responsible for generating a data key, as well as encrypting and
+     *                       decrypting data keys.
      * @return The AwsKmsKeyringBuilder, for method chaining
      */
-    public AwsKmsKeyringBuilder generatorKeyId(String generatorKeyId) {
+    public AwsKmsKeyringBuilder generatorKeyId(AwsKmsCmkId generatorKeyId) {
         this.generatorKeyId = generatorKeyId;
         return this;
     }

src/main/java/com/amazonaws/encryptionsdk/keyrings/StandardKeyrings.java

@@ -13,6 +13,8 @@
 
 package com.amazonaws.encryptionsdk.keyrings;
 
+import com.amazonaws.encryptionsdk.kms.AwsKmsCmkId;
+
 import java.util.Arrays;
 import java.util.List;
 
@@ -50,11 +52,12 @@ public static RawRsaKeyringBuilder rawRsa() {
      * encrypt, and decrypt data keys using the supplied AWS KMS defined Customer Master Key (CMK).
      * Use {@link #awsKms()} for more advanced configuration using a {@link AwsKmsKeyringBuilder}/
      *
-     * @param generatorKeyId    A string that identifies a AWS KMS CMK responsible for generating a data key,
-     *                          as well as encrypting and decrypting data keys in ARN, CMK Alias, or ARN Alias format.
+     * @param generatorKeyId    An {@link AwsKmsCmkId} in ARN, CMK Alias, ARN Alias or Key Id format that identifies a
+     *                          AWS KMS CMK responsible for generating a data key, as well as encrypting and
+     *                          decrypting data keys .
      * @return The {@code Keyring}
      */
-    public static Keyring awsKms(String generatorKeyId) {
+    public static Keyring awsKms(AwsKmsCmkId generatorKeyId) {
         return new AwsKmsKeyringBuilder()
                 .generatorKeyId(generatorKeyId)
                 .build();

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

@@ -14,6 +14,7 @@
 package com.amazonaws.encryptionsdk.kms;
 
 import com.amazonaws.ClientConfiguration;
+import com.amazonaws.arn.Arn;
 import com.amazonaws.auth.AWSCredentialsProvider;
 import com.amazonaws.encryptionsdk.exception.UnsupportedRegionException;
 import com.amazonaws.services.kms.AWSKMS;
@@ -59,6 +60,25 @@ static Builder builder() {
         return new Builder(AWSKMSClientBuilder.standard());
     }
 
+    /**
+     * Parses region from the given key id (if possible) and passes that region to the
+     * given clientSupplier to produce an {@code AWSKMS} client.
+     *
+     * @param keyId          The Amazon Resource Name, Key Alias, Alias ARN or KeyId
+     * @param clientSupplier The client supplier
+     * @return AWSKMS The client
+     */
+    static AWSKMS getClientByKeyId(AwsKmsCmkId keyId, AwsKmsClientSupplier clientSupplier) {
+        requireNonNull(keyId, "keyId is required");
+        requireNonNull(clientSupplier, "clientSupplier is required");
+
+        if(keyId.isArn()) {
+            return clientSupplier.getClient(Arn.fromString(keyId.toString()).getRegion());
+        }
+
+        return clientSupplier.getClient(null);
+    }
+
     /**
      * Builder to construct an AwsKmsClientSupplier given various
      * optional settings.