fix: Update DecryptionMaterials code to support legacy custom CMMs

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

@@ -0,0 +1,124 @@
+// 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.CryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.CryptoResult;
+import com.amazonaws.encryptionsdk.DefaultCryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.MasterKeyProvider;
+import com.amazonaws.encryptionsdk.kmssdkv2.KmsMasterKeyProvider;
+import com.amazonaws.encryptionsdk.model.DecryptionMaterials;
+import com.amazonaws.encryptionsdk.model.DecryptionMaterialsRequest;
+import com.amazonaws.encryptionsdk.model.EncryptionMaterials;
+import com.amazonaws.encryptionsdk.model.EncryptionMaterialsRequest;
+
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * <p>
+ * Creates a custom implementation of the CryptoMaterialsManager interface,
+ * then uses that implementation to encrypt and decrypt a file using an AWS KMS CMK.
+ *
+ * <p>
+ * Arguments:
+ * <ol>
+ * <li>Key ARN: For help finding the Amazon Resource Name (ARN) of your AWS KMS customer master
+ *    key (CMK), see 'Viewing Keys' at http://docs.aws.amazon.com/kms/latest/developerguide/viewing-keys.html
+ * </ol>
+ */
+public class CustomCMMExample {
+
+    private static final byte[] EXAMPLE_DATA = "Hello World".getBytes(StandardCharsets.UTF_8);
+
+    public static void main(final String[] args) {
+        final String keyArn = args[0];
+
+        CryptoMaterialsManager cmm = new SigningSuiteOnlyCMM(
+            KmsMasterKeyProvider.builder().buildStrict(keyArn)
+        );
+
+        encryptAndDecryptWithCMM(cmm);
+    }
+
+    static void encryptAndDecryptWithCMM(final CryptoMaterialsManager cmm) {
+        // 1. Instantiate the SDK
+        // This builds the AwsCrypto client with the RequireEncryptRequireDecrypt commitment policy,
+        // which enforces that this client only encrypts using committing algorithm suites and enforces
+        // that this client will only decrypt encrypted messages that were created with a committing algorithm suite.
+        // This is the default commitment policy if you build the client with `AwsCrypto.builder().build()`
+        // or `AwsCrypto.standard()`.
+        final AwsCrypto crypto = AwsCrypto.builder()
+                .withCommitmentPolicy(CommitmentPolicy.RequireEncryptRequireDecrypt)
+                .build();
+
+        // 2. Create an encryption context
+        // Most encrypted data should have an associated encryption context
+        // to protect integrity. This sample uses placeholder values.
+        // For more information see:
+        // blogs.aws.amazon.com/security/post/Tx2LZ6WBJJANTNW/How-to-Protect-the-Integrity-of-Your-Encrypted-Data-by-Using-AWS-Key-Management
+        final Map<String, String> encryptionContext = Collections.singletonMap("ExampleContextKey", "ExampleContextValue");
+
+        // 3. Encrypt the data with the provided CMM
+        final CryptoResult<byte[], ?> encryptResult = crypto.encryptData(cmm, EXAMPLE_DATA, encryptionContext);
+        final byte[] ciphertext = encryptResult.getResult();
+
+        // 4. Decrypt the data
+        final CryptoResult<byte[], ?> decryptResult = crypto.decryptData(cmm, ciphertext);
+
+        // 5. Verify that the encryption context in the result contains the
+        // encryption context supplied to the encryptData method. Because the
+        // SDK can add values to the encryption context, don't require that
+        // the entire context matches.
+        if (!encryptionContext.entrySet().stream()
+                .allMatch(e -> e.getValue().equals(decryptResult.getEncryptionContext().get(e.getKey())))) {
+            throw new IllegalStateException("Wrong Encryption Context!");
+        }
+
+        // 6. Verify that the decrypted plaintext matches the original plaintext
+        assert Arrays.equals(decryptResult.getResult(), EXAMPLE_DATA);
+    }
+
+    // Custom CMM implementation.
+    // This CMM only allows encryption/decryption using signing algorithms.
+    // It wraps an underlying CMM implementation and checks its materials
+    // to ensure that it is only using signed encryption algorithms.
+    public static class SigningSuiteOnlyCMM implements CryptoMaterialsManager {
+
+        // The underlying CMM.
+        private CryptoMaterialsManager underlyingCMM;
+
+        // If only a MasterKeyProvider is constructed, the underlying CMM is the default CMM.
+        public SigningSuiteOnlyCMM(MasterKeyProvider<?> mkp) {
+            this.underlyingCMM = new DefaultCryptoMaterialsManager(mkp);
+        }
+
+        // This CMM can wrap any other CMM implementation.
+        public SigningSuiteOnlyCMM(CryptoMaterialsManager underlyingCMM) {
+            this.underlyingCMM = underlyingCMM;
+        }
+
+        @Override
+        public EncryptionMaterials getMaterialsForEncrypt(EncryptionMaterialsRequest request) {
+            EncryptionMaterials materials = underlyingCMM.getMaterialsForEncrypt(request);
+            if (materials.getAlgorithm().getTrailingSignatureAlgo() == null) {
+                throw new IllegalArgumentException("Algorithm provided to SigningSuiteOnlyCMM is not a supported signing algorithm: " + materials.getAlgorithm());
+            }
+            return materials;
+        }
+
+        @Override
+        public DecryptionMaterials decryptMaterials(DecryptionMaterialsRequest request) {
+            if (request.getAlgorithm().getTrailingSignatureAlgo() == null) {
+                throw new IllegalArgumentException("Algorithm provided to SigningSuiteOnlyCMM is not a supported signing algorithm: " + request.getAlgorithm());
+            }
+            return underlyingCMM.decryptMaterials(request);
+        }
+    }
+
+}

src/main/java/com/amazonaws/encryptionsdk/CMMHandler.java

@@ -4,6 +4,7 @@
 package com.amazonaws.encryptionsdk;
 
 import com.amazonaws.encryptionsdk.internal.Utils;
+import com.amazonaws.encryptionsdk.model.DecryptionMaterials;
 import com.amazonaws.encryptionsdk.model.DecryptionMaterialsHandler;
 import com.amazonaws.encryptionsdk.model.DecryptionMaterialsRequest;
 import com.amazonaws.encryptionsdk.model.EncryptionMaterialsHandler;
@@ -63,7 +64,37 @@ private GetEncryptionMaterialsInput getEncryptionMaterialsRequestInput(
   public DecryptionMaterialsHandler decryptMaterials(
       DecryptionMaterialsRequest request, CommitmentPolicy commitmentPolicy) {
     if (cmm != null && mplCMM == null) {
-      return new DecryptionMaterialsHandler(cmm.decryptMaterials(request));
+      // This is an implementation of the legacy native CryptoMaterialsManager interface from
+      // ESDK-Java.
+      DecryptionMaterials materials = cmm.decryptMaterials(request);
+      if (materials.getEncryptionContext().isEmpty() && !request.getEncryptionContext().isEmpty()) {
+        // If the request specified an encryption context,
+        // and we are using the legacy native CMM,
+        // add the encryptionContext to the materials.
+        //
+        // ESDK-Java 3.0 changed internals of decrypt behavior,
+        // This code makes earlier CMM implementations compatible with post-3.0 behavior.
+        //
+        // Version 3.0 assumes that CMMs' implementations of decryptMaterials
+        // will set an encryptionContext attribute on returned DecryptionMaterials.
+        // The DefaultCryptoMaterialsManager's behavior was changed in 3.0.
+        // It now sets the encryptionContext attribute with the value from the ciphertext's headers.
+        //
+        // But custom CMMs' behavior was not updated.
+        // However, there is no custom CMM before version 3.0 that could set an encryptionContext
+        // attribute.
+        // The encryptionContext attribute was only introduced to decryptMaterials objects
+        // in ESDK 3.0, so no CMM could have configured this attribute before 3.0.
+        // As a result, the ESDK assumes that any native CMM
+        // that does not add encryptionContext to its decryptMaterials
+        // SHOULD add encryptionContext to its decryptMaterials,
+        //
+        // If a custom CMM implementation conflicts with this assumption.
+        // that CMM implementation MUST move to the MPL.
+        materials =
+            materials.toBuilder().setEncryptionContext(request.getEncryptionContext()).build();
+      }
+      return new DecryptionMaterialsHandler(materials);
     } else {
       DecryptMaterialsInput input = getDecryptMaterialsInput(request, commitmentPolicy);
       DecryptMaterialsOutput output = mplCMM.DecryptMaterials(input);

src/main/java/com/amazonaws/encryptionsdk/model/DecryptionMaterials.java

@@ -2,6 +2,7 @@
 
 import com.amazonaws.encryptionsdk.DataKey;
 import java.security.PublicKey;
+import java.util.Collections;
 import java.util.Map;
 
 public final class DecryptionMaterials {
@@ -12,7 +13,11 @@ public final class DecryptionMaterials {
   private DecryptionMaterials(Builder b) {
     dataKey = b.getDataKey();
     trailingSignatureKey = b.getTrailingSignatureKey();
-    encryptionContext = b.getEncryptionContext();
+    if (b.getEncryptionContext() != null) {
+      encryptionContext = b.getEncryptionContext();
+    } else {
+      encryptionContext = Collections.emptyMap();
+    }
   }
 
   public DataKey<?> getDataKey() {

src/test/java/com/amazonaws/crypto/examples/v2/CustomCMMExampleTest.java

@@ -0,0 +1,28 @@
+// 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.CryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.kms.KMSTestFixtures;
+import com.amazonaws.encryptionsdk.kms.KmsMasterKeyProvider;
+import org.junit.Test;
+
+public class CustomCMMExampleTest {
+
+  @Test
+  public void testCustomCMMExample() {
+    CryptoMaterialsManager cmm =
+        new CustomCMMExample.SigningSuiteOnlyCMM(
+            KmsMasterKeyProvider.builder().buildStrict(KMSTestFixtures.US_WEST_2_KEY_ID));
+    CustomCMMExample.encryptAndDecryptWithCMM(cmm);
+  }
+
+  @Test
+  public void testV2Cmm() {
+    V2DefaultCryptoMaterialsManager cmm =
+        new V2DefaultCryptoMaterialsManager(
+            KmsMasterKeyProvider.builder().buildStrict(KMSTestFixtures.US_WEST_2_KEY_ID));
+    CustomCMMExample.encryptAndDecryptWithCMM(cmm);
+  }
+}

src/test/java/com/amazonaws/crypto/examples/v2/V2DefaultCryptoMaterialsManager.java

@@ -0,0 +1,171 @@
+package com.amazonaws.crypto.examples.v2;
+
+import static com.amazonaws.encryptionsdk.internal.Utils.assertNonNull;
+
+import com.amazonaws.encryptionsdk.CommitmentPolicy;
+import com.amazonaws.encryptionsdk.CryptoAlgorithm;
+import com.amazonaws.encryptionsdk.CryptoMaterialsManager;
+import com.amazonaws.encryptionsdk.DataKey;
+import com.amazonaws.encryptionsdk.MasterKey;
+import com.amazonaws.encryptionsdk.MasterKeyProvider;
+import com.amazonaws.encryptionsdk.MasterKeyRequest;
+import com.amazonaws.encryptionsdk.exception.AwsCryptoException;
+import com.amazonaws.encryptionsdk.exception.CannotUnwrapDataKeyException;
+import com.amazonaws.encryptionsdk.internal.Constants;
+import com.amazonaws.encryptionsdk.internal.TrailingSignatureAlgorithm;
+import com.amazonaws.encryptionsdk.model.DecryptionMaterials;
+import com.amazonaws.encryptionsdk.model.DecryptionMaterialsRequest;
+import com.amazonaws.encryptionsdk.model.EncryptionMaterials;
+import com.amazonaws.encryptionsdk.model.EncryptionMaterialsRequest;
+import com.amazonaws.encryptionsdk.model.KeyBlob;
+import java.security.GeneralSecurityException;
+import java.security.KeyPair;
+import java.security.PublicKey;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/*
+ This is a copy-paste of the DefaultCryptoMaterialsManager implementation
+ from the final commit of the V2 ESDK: 1870a082358d59e32c60d74116d6f43c0efa466b
+ ESDK V3 implicitly changed the contract between CMMs and the ESDK.
+ After V3, DecryptMaterials has an `encryptionContext` attribute,
+ and CMMs are expected to set this attribute.
+ The V3 commit modified this DefaultCMM's `decryptMaterials` implementation
+ to set encryptionContext on returned DecryptionMaterials objects.
+ However, there are custom implementations of the legacy native CMM
+ that do not set encryptionContext.
+ This CMM is used to explicitly assert that the V2 implementation of
+ the DefaultCMM is compatible with V3 logic,
+ which implicitly asserts that custom implementations of V2-compatible CMMs
+ are also compatible with V3 logic.
+*/
+public class V2DefaultCryptoMaterialsManager implements CryptoMaterialsManager {
+  private final MasterKeyProvider<?> mkp;
+
+  private final CryptoAlgorithm DEFAULT_CRYPTO_ALGORITHM =
+      CryptoAlgorithm.ALG_AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384;
+
+  /** @param mkp The master key provider to delegate to */
+  public V2DefaultCryptoMaterialsManager(MasterKeyProvider<?> mkp) {
+    assertNonNull(mkp, "mkp");
+    this.mkp = mkp;
+  }
+
+  @Override
+  public EncryptionMaterials getMaterialsForEncrypt(EncryptionMaterialsRequest request) {
+    Map<String, String> context = request.getContext();
+
+    CryptoAlgorithm algo = request.getRequestedAlgorithm();
+    CommitmentPolicy commitmentPolicy = request.getCommitmentPolicy();
+    // Set default according to commitment policy
+    if (algo == null && commitmentPolicy == CommitmentPolicy.ForbidEncryptAllowDecrypt) {
+      algo = CryptoAlgorithm.ALG_AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384;
+    } else if (algo == null) {
+      algo = CryptoAlgorithm.ALG_AES_256_GCM_HKDF_SHA512_COMMIT_KEY_ECDSA_P384;
+    }
+
+    KeyPair trailingKeys = null;
+    if (algo.getTrailingSignatureLength() > 0) {
+      try {
+        trailingKeys = generateTrailingSigKeyPair(algo);
+        if (context.containsKey(Constants.EC_PUBLIC_KEY_FIELD)) {
+          throw new IllegalArgumentException(
+              "EncryptionContext contains reserved field " + Constants.EC_PUBLIC_KEY_FIELD);
+        }
+        // make mutable
+        context = new HashMap<>(context);
+        context.put(Constants.EC_PUBLIC_KEY_FIELD, serializeTrailingKeyForEc(algo, trailingKeys));
+      } catch (final GeneralSecurityException ex) {
+        throw new AwsCryptoException(ex);
+      }
+    }
+
+    final MasterKeyRequest.Builder mkRequestBuilder = MasterKeyRequest.newBuilder();
+    mkRequestBuilder.setEncryptionContext(context);
+
+    mkRequestBuilder.setStreaming(request.getPlaintextSize() == -1);
+    if (request.getPlaintext() != null) {
+      mkRequestBuilder.setPlaintext(request.getPlaintext());
+    } else {
+      mkRequestBuilder.setSize(request.getPlaintextSize());
+    }
+
+    @SuppressWarnings("unchecked")
+    final List<MasterKey> mks =
+        (List<MasterKey>)
+            assertNonNull(mkp, "provider").getMasterKeysForEncryption(mkRequestBuilder.build());
+
+    if (mks.isEmpty()) {
+      throw new IllegalArgumentException("No master keys provided");
+    }
+
+    DataKey<?> dataKey = mks.get(0).generateDataKey(algo, context);
+
+    List<KeyBlob> keyBlobs = new ArrayList<>(mks.size());
+    keyBlobs.add(new KeyBlob(dataKey));
+
+    for (int i = 1; i < mks.size(); i++) {
+      //noinspection unchecked
+      keyBlobs.add(new KeyBlob(mks.get(i).encryptDataKey(algo, context, dataKey)));
+    }
+
+    //noinspection unchecked
+    return EncryptionMaterials.newBuilder()
+        .setAlgorithm(algo)
+        .setCleartextDataKey(dataKey.getKey())
+        .setEncryptedDataKeys(keyBlobs)
+        .setEncryptionContext(context)
+        .setTrailingSignatureKey(trailingKeys == null ? null : trailingKeys.getPrivate())
+        .setMasterKeys(mks)
+        .build();
+  }
+
+  @Override
+  public DecryptionMaterials decryptMaterials(DecryptionMaterialsRequest request) {
+    DataKey<?> dataKey =
+        mkp.decryptDataKey(
+            request.getAlgorithm(), request.getEncryptedDataKeys(), request.getEncryptionContext());
+
+    if (dataKey == null) {
+      throw new CannotUnwrapDataKeyException("Could not decrypt any data keys");
+    }
+
+    PublicKey pubKey = null;
+    if (request.getAlgorithm().getTrailingSignatureLength() > 0) {
+      try {
+        String serializedPubKey = request.getEncryptionContext().get(Constants.EC_PUBLIC_KEY_FIELD);
+
+        if (serializedPubKey == null) {
+          throw new AwsCryptoException("Missing trailing signature public key");
+        }
+
+        pubKey = deserializeTrailingKeyFromEc(request.getAlgorithm(), serializedPubKey);
+      } catch (final IllegalStateException ex) {
+        throw new AwsCryptoException(ex);
+      }
+    } else if (request.getEncryptionContext().containsKey(Constants.EC_PUBLIC_KEY_FIELD)) {
+      throw new AwsCryptoException("Trailing signature public key found for non-signed algorithm");
+    }
+
+    return DecryptionMaterials.newBuilder()
+        .setDataKey(dataKey)
+        .setTrailingSignatureKey(pubKey)
+        .build();
+  }
+
+  private PublicKey deserializeTrailingKeyFromEc(CryptoAlgorithm algo, String pubKey) {
+    return TrailingSignatureAlgorithm.forCryptoAlgorithm(algo).deserializePublicKey(pubKey);
+  }
+
+  private static String serializeTrailingKeyForEc(CryptoAlgorithm algo, KeyPair trailingKeys) {
+    return TrailingSignatureAlgorithm.forCryptoAlgorithm(algo)
+        .serializePublicKey(trailingKeys.getPublic());
+  }
+
+  private static KeyPair generateTrailingSigKeyPair(CryptoAlgorithm algo)
+      throws GeneralSecurityException {
+    return TrailingSignatureAlgorithm.forCryptoAlgorithm(algo).generateKey();
+  }
+}