feat(Examples): Wrap H-Keyring as a Master Key

src/examples/java/com/amazonaws/crypto/examples/v2/HKeyringMasterKey.java

@@ -0,0 +1,202 @@
+package com.amazonaws.crypto.examples.v2;
+
+import com.amazonaws.encryptionsdk.CryptoAlgorithm;
+import com.amazonaws.encryptionsdk.DataKey;
+import com.amazonaws.encryptionsdk.EncryptedDataKey;
+import com.amazonaws.encryptionsdk.MasterKey;
+import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
+import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
+import com.amazonaws.encryptionsdk.exception.UnsupportedProviderException;
+import software.amazon.cryptography.materialproviders.IKeyring;
+import software.amazon.cryptography.materialproviders.MaterialProviders;
+import software.amazon.cryptography.materialproviders.model.*;
+
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+import java.nio.ByteBuffer;
+import java.util.*;
+
+public class HKeyringMasterKey extends MasterKey<HKeyringMasterKey> {
+
+  final private IKeyring hKeyring;
+  final private CreateAwsKmsHierarchicalKeyringInput hKeyringInput;
+  final private MaterialProviders mpl;
+
+  public HKeyringMasterKey(
+          CreateAwsKmsHierarchicalKeyringInput input
+  ) {
+    if (input.branchKeyIdSupplier() != null) throw new UnsupportedProviderException("branchKeyIdSupplier must be null");
+    if (input.branchKeyId() == null) throw new UnsupportedProviderException("branchKeyId cannot be null");
+    hKeyringInput = input;
+    mpl =
+            MaterialProviders.builder()
+                    .MaterialProvidersConfig(MaterialProvidersConfig.builder().build())
+                    .build();
+    hKeyring = mpl.CreateAwsKmsHierarchicalKeyring(hKeyringInput);
+  }
+
+  @Override
+  public String getProviderId() {
+    return "aws-kms-hierarchy";
+  }
+
+  @Override
+  public String getKeyId() {
+    return this.hKeyringInput.branchKeyId();
+  }
+
+  /**
+   * Generates a new {@link DataKey} which is protected by this {@link MasterKey} for use with
+   * {@code algorithm} and associated with the provided {@code encryptionContext}.
+   *
+   * @param algorithm
+   * @param encryptionContext
+   */
+  @Override
+  public DataKey<HKeyringMasterKey> generateDataKey(CryptoAlgorithm algorithm, Map<String, String> encryptionContext) {
+    AlgorithmSuiteInfo algorithmSuiteInfo = ValidateAndConvertAlgo(algorithm);
+    EncryptionMaterials encryptionMaterials = EncryptionMaterials.builder()
+            .algorithmSuite(algorithmSuiteInfo)
+            .encryptionContext(encryptionContext)
+            .encryptedDataKeys(Collections.emptyList())
+            .requiredEncryptionContextKeys(Collections.emptyList())
+            .build();
+    OnEncryptInput eInput = OnEncryptInput.builder()
+            .materials(encryptionMaterials)
+            .build();
+    OnEncryptOutput onEncryptOutput = hKeyring.OnEncrypt(eInput);
+    software.amazon.cryptography.materialproviders.model.EncryptedDataKey encryptedDataKey = onEncryptOutput.materials().encryptedDataKeys().get(0);
+    return new DataKey<>(
+            new SecretKeySpec(onEncryptOutput.materials().plaintextDataKey().array(), algorithm.getDataKeyAlgo()),
+            encryptedDataKey.ciphertext().array(),
+            encryptedDataKey.keyProviderInfo().array(),
+            this);
+  }
+
+  /**
+   * Returns a new copy of the provided {@code dataKey} which is protected by this {@link MasterKey}
+   * for use with {@code algorithm} and associated with the provided {@code encryptionContext}.
+   *
+   * @param algorithm
+   * @param encryptionContext
+   * @param dataKey
+   */
+  @Override
+  public DataKey<HKeyringMasterKey> encryptDataKey(
+          CryptoAlgorithm algorithm,
+          Map<String, String> encryptionContext,
+          DataKey<?> dataKey
+  ) {
+    AlgorithmSuiteInfo algorithmSuiteInfo = ValidateAndConvertAlgo(algorithm);
+    final SecretKey key = dataKey.getKey();
+    if (!key.getFormat().equals("RAW")) {
+      throw new IllegalArgumentException(
+              "Can only re-encrypt data keys which are in RAW format, not "
+                      + dataKey.getKey().getFormat());
+    }
+    EncryptionMaterials encryptionMaterials = EncryptionMaterials.builder()
+            .algorithmSuite(algorithmSuiteInfo)
+            .encryptionContext(encryptionContext)
+            .encryptedDataKeys(Collections.emptyList())
+            .requiredEncryptionContextKeys(Collections.emptyList())
+            .plaintextDataKey(ByteBuffer.wrap(dataKey.getKey().getEncoded()))
+            .build();
+    OnEncryptInput eInput = OnEncryptInput.builder()
+            .materials(encryptionMaterials)
+            .build();
+    OnEncryptOutput onEncryptOutput = hKeyring.OnEncrypt(eInput);
+    software.amazon.cryptography.materialproviders.model.EncryptedDataKey encryptedDataKey = onEncryptOutput.materials().encryptedDataKeys().get(0);
+    return new DataKey<>(
+            key,
+            encryptedDataKey.ciphertext().array(),
+            encryptedDataKey.keyProviderInfo().array(),
+            this);
+  }
+
+  /**
+   * Iterates through {@code encryptedDataKeys} and returns the first one which can be successfully
+   * decrypted.
+   *
+   * @param algorithm
+   * @param encryptedDataKeys
+   * @param encryptionContext
+   * @return a DataKey if one can be decrypted, otherwise returns {@code null}
+   * @throws UnsupportedProviderException if the {@code encryptedDataKey} is associated with an
+   *                                      unsupported provider
+   * @throws CannotUnwrapDataKeyException if the {@code encryptedDataKey} cannot be decrypted
+   */
+  @Override
+  public DataKey<HKeyringMasterKey> decryptDataKey(
+          CryptoAlgorithm algorithm,
+          Collection<? extends EncryptedDataKey> encryptedDataKeys,
+          Map<String, String> encryptionContext
+  ) throws UnsupportedProviderException, AwsCryptoException
+  {
+    AlgorithmSuiteInfo algorithmSuiteInfo = ValidateAndConvertAlgo(algorithm);
+    if (encryptedDataKeys.size() != 1) {
+      // TODO: If needed, we could refactor this to properly support multiple EDKs, it would not be hard.
+      throw new UnsupportedProviderException("Alas, this Master Key Provider can work with one (1) Encrypted Data Key; got " + encryptedDataKeys.size());
+    }
+    List<EncryptedDataKey> nativeEDKS = EDKCollectionToNative(encryptedDataKeys);
+    List<software.amazon.cryptography.materialproviders.model.EncryptedDataKey> mplEDKS = EDKCollectionToMPL(encryptedDataKeys);
+    DecryptionMaterials decryptionMaterials = DecryptionMaterials.builder()
+            .algorithmSuite(algorithmSuiteInfo)
+            .encryptionContext(encryptionContext)
+            .requiredEncryptionContextKeys(Collections.emptyList())
+            .build();
+    OnDecryptInput onDecryptInput = OnDecryptInput.builder()
+            .encryptedDataKeys(mplEDKS)
+            .materials(decryptionMaterials)
+            .build();
+    OnDecryptOutput onDecryptOutput = this.hKeyring.OnDecrypt(onDecryptInput);
+    if (onDecryptOutput.materials().plaintextDataKey() != null) {
+      if (onDecryptOutput.materials().plaintextDataKey().array().length != algorithm.getDataKeyLength())
+        throw new AwsCryptoException("Decrypted Data Key is incorrect length!");
+      return new DataKey<>(
+              new SecretKeySpec(onDecryptOutput.materials().plaintextDataKey().array(), algorithm.getDataKeyAlgo()),
+              nativeEDKS.get(0).getEncryptedDataKey(),
+              nativeEDKS.get(0).getProviderInformation(),
+              this
+      );
+    }
+    return null;
+  }
+
+  private List<software.amazon.cryptography.materialproviders.model.EncryptedDataKey> EDKCollectionToMPL(
+          Collection<? extends EncryptedDataKey> encryptedDataKeys
+  ) {
+    List<software.amazon.cryptography.materialproviders.model.EncryptedDataKey> mplEDKS =
+            new ArrayList<>(encryptedDataKeys.size());
+
+    for (EncryptedDataKey keyBlob : encryptedDataKeys) {
+      mplEDKS.add(
+              software.amazon.cryptography.materialproviders.model.EncryptedDataKey.builder()
+                      .keyProviderId(keyBlob.getProviderId())
+                      .keyProviderInfo(
+                              ByteBuffer.wrap(
+                                      keyBlob.getProviderInformation(), 0, keyBlob.getProviderInformation().length))
+                      .ciphertext(
+                              ByteBuffer.wrap(
+                                      keyBlob.getEncryptedDataKey(), 0, keyBlob.getEncryptedDataKey().length))
+                      .build());
+    }
+    return mplEDKS;
+  }
+
+  private List<EncryptedDataKey> EDKCollectionToNative(
+          Collection<? extends EncryptedDataKey> encryptedDataKeys
+  ) {
+    List<EncryptedDataKey> nativeEDKS = new ArrayList<>(encryptedDataKeys.size());
+    nativeEDKS.addAll(encryptedDataKeys);
+    return nativeEDKS;
+  }
+
+  private AlgorithmSuiteInfo ValidateAndConvertAlgo(
+          CryptoAlgorithm algorithm
+  ) {
+    if (algorithm.getTrailingSignatureAlgo() != null) {
+      throw new UnsupportedProviderException("The HKeyringMasterKey provider does not support trailing signature algorithms!");
+    }
+    return mpl.GetAlgorithmSuiteInfo(ByteBuffer.allocate(2).putShort((short) algorithm.getValue()));
+  }
+}

src/examples/java/com/amazonaws/crypto/examples/v2/HKeyringMasterKeyWithCachingCMM.java

@@ -0,0 +1,116 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package com.amazonaws.crypto.examples.v2;
+
+import com.amazonaws.encryptionsdk.AwsCrypto;
+import com.amazonaws.encryptionsdk.CommitmentPolicy;
+import com.amazonaws.encryptionsdk.CryptoAlgorithm;
+import com.amazonaws.encryptionsdk.CryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.CryptoResult;
+import com.amazonaws.encryptionsdk.caching.CachingCryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.caching.CryptoMaterialsCache;
+import com.amazonaws.encryptionsdk.caching.LocalCryptoMaterialsCache;
+import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
+import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.cryptography.keystore.KeyStore;
+import software.amazon.cryptography.keystore.model.CreateKeyInput;
+import software.amazon.cryptography.keystore.model.KMSConfiguration;
+import software.amazon.cryptography.keystore.model.KeyStoreConfig;
+import software.amazon.cryptography.materialproviders.IKeyring;
+import software.amazon.cryptography.materialproviders.MaterialProviders;
+import software.amazon.cryptography.materialproviders.model.*;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+public class HKeyringMasterKeyWithCachingCMM {
+  private static final byte[] EXAMPLE_DATA = "Hello World".getBytes(StandardCharsets.UTF_8);
+
+  public static void encryptAndDecryptWithKeyringViaCachingCMM(
+      String keyStoreTableName, String logicalKeyStoreName, String kmsKeyId) {
+    final AwsCrypto crypto = AwsCrypto.builder()
+            .withCommitmentPolicy(CommitmentPolicy.RequireEncryptAllowDecrypt)
+            .withEncryptionAlgorithm(CryptoAlgorithm.ALG_AES_256_GCM_HKDF_SHA512_COMMIT_KEY)
+            .withMaxEncryptedDataKeys(1)
+            .build();
+
+    // Configure your KeyStore resource.
+    //    This SHOULD be the same configuration that you used
+    //    to initially create and populate your KeyStore.
+    final KeyStore keystore =
+        KeyStore.builder()
+            .KeyStoreConfig(
+                KeyStoreConfig.builder()
+                    .ddbClient(DynamoDbClient.create())
+                    .ddbTableName(keyStoreTableName)
+                    .logicalKeyStoreName(logicalKeyStoreName)
+                    .kmsClient(KmsClient.create())
+                    .kmsConfiguration(KMSConfiguration.builder().kmsKeyArn(kmsKeyId).build())
+                    .build())
+            .build();
+
+    // Call CreateKey to create two new active branch keys
+    final String branchKeyIdA =
+        keystore.CreateKey(CreateKeyInput.builder().build()).branchKeyIdentifier();
+
+    // 4. Create the Hierarchical Keyring.
+    final MaterialProviders matProv =
+        MaterialProviders.builder()
+            .MaterialProvidersConfig(MaterialProvidersConfig.builder().build())
+            .build();
+    final CreateAwsKmsHierarchicalKeyringInput keyringInput =
+        CreateAwsKmsHierarchicalKeyringInput.builder()
+            .keyStore(keystore)
+            .branchKeyId(branchKeyIdA)
+            .ttlSeconds(600)
+            .cache(
+                CacheType.builder()
+                        .MultiThreaded(MultiThreadedCache.builder().entryCapacity(100).build())// OPTIONAL
+                    .build())
+            .build();
+
+    Map<String, String> encryptionContextA = new HashMap<>();
+    encryptionContextA.put("tenant", "TenantA");
+    HKeyringMasterKey masterKey = new HKeyringMasterKey(keyringInput);
+    final IKeyring hierarchicalKeyringA = matProv.CreateAwsKmsHierarchicalKeyring(keyringInput);
+
+    // Create a cache
+    CryptoMaterialsCache cache = new LocalCryptoMaterialsCache(100);
+
+    // Create a caching CMM
+    CryptoMaterialsManager cachingCmm =
+            CachingCryptoMaterialsManager.newBuilder().withMasterKeyProvider(masterKey)
+                    .withCache(cache)
+                    .withMaxAge(100, TimeUnit.SECONDS)
+                    .withMessageUseLimit(1000)
+                    .build();
+
+    // Encrypt the data for encryptionContextA & encryptionContextB
+    final CryptoResult<byte[], ?> encryptResultA =
+        crypto.encryptData(cachingCmm, EXAMPLE_DATA, encryptionContextA);
+
+    // OK. Can the Keyring decrypt it?
+    CryptoResult<byte[], ?> decryptResultA =
+            crypto.decryptData(hierarchicalKeyringA, encryptResultA.getResult());
+    assert Arrays.equals(decryptResultA.getResult(), EXAMPLE_DATA);
+
+    // OK. Can the Caching CMM Decrypt it?
+    decryptResultA = crypto.decryptData(cachingCmm, encryptResultA.getResult());
+    assert Arrays.equals(decryptResultA.getResult(), EXAMPLE_DATA);
+  }
+
+  public static void main(final String[] args) {
+    if (args.length <= 0) {
+      throw new IllegalArgumentException(
+          "To run this example, include the keyStoreTableName, logicalKeyStoreName, and kmsKeyId in args");
+    }
+    final String keyStoreTableName = args[0];
+    final String logicalKeyStoreName = args[1];
+    final String kmsKeyId = args[2];
+    encryptAndDecryptWithKeyringViaCachingCMM(keyStoreTableName, logicalKeyStoreName, kmsKeyId);
+  }
+}

src/test/java/com/amazonaws/crypto/examples/keyrings/AwsKmsHierarchicalKeyringExampleTest.java

@@ -4,6 +4,7 @@
 package com.amazonaws.crypto.examples.keyrings;
 
 import com.amazonaws.crypto.examples.keyrings.hierarchical.AwsKmsHierarchicalKeyringExample;
+import com.amazonaws.crypto.examples.v2.HKeyringMasterKeyWithCachingCMM;
 import com.amazonaws.encryptionsdk.kms.KMSTestFixtures;
 import org.junit.Test;
 
@@ -23,4 +24,12 @@ public void testEncryptAndDecryptWithKeyringThreadSafe() {
         KMSTestFixtures.TEST_LOGICAL_KEYSTORE_NAME,
         KMSTestFixtures.TEST_KEYSTORE_KMS_KEY_ID);
   }
+
+  @Test
+  public void testCrazy() {
+   HKeyringMasterKeyWithCachingCMM.encryptAndDecryptWithKeyringViaCachingCMM(
+           KMSTestFixtures.TEST_KEYSTORE_NAME,
+           KMSTestFixtures.TEST_LOGICAL_KEYSTORE_NAME,
+           KMSTestFixtures.TEST_KEYSTORE_KMS_KEY_ID);
+  }
 }