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,11 +4,13 @@
 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;
 import com.amazonaws.encryptionsdk.model.EncryptionMaterialsRequest;
 import com.amazonaws.encryptionsdk.model.KeyBlob;
+
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.List;
@@ -63,7 +65,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,31 @@
+// 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);
+  }
+}
+
+