feat: allow raw keyrings to decrypt with multiple wrapping keys

src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java

@@ -68,6 +68,7 @@
 import software.amazon.encryption.s3.materials.EncryptionMaterials;
 import software.amazon.encryption.s3.materials.Keyring;
 import software.amazon.encryption.s3.materials.KmsKeyring;
+import software.amazon.encryption.s3.materials.MaterialsDescription;
 import software.amazon.encryption.s3.materials.MultipartConfiguration;
 import software.amazon.encryption.s3.materials.PartialRsaKeyPair;
 import software.amazon.encryption.s3.materials.RawKeyring;
@@ -246,7 +247,7 @@ public ReEncryptInstructionFileResponse reEncryptInstructionFile(ReEncryptInstru
         //Extract cryptographic parameters from the current instruction file that MUST be preserved during re-encryption
         final AlgorithmSuite algorithmSuite = contentMetadata.algorithmSuite();
         final EncryptedDataKey originalEncryptedDataKey = contentMetadata.encryptedDataKey();
-        final Map<String, String> currentKeyringMaterialsDescription = contentMetadata.encryptedDataKeyMatDescOrContext();
+        final MaterialsDescription currentKeyringMaterialsDescription = contentMetadata.materialsDescription();
         final byte[] iv = contentMetadata.contentIv();
 
         //Decrypt the data key using the current keyring
@@ -257,6 +258,7 @@ public ReEncryptInstructionFileResponse reEncryptInstructionFile(ReEncryptInstru
           DecryptMaterialsRequest.builder()
             .algorithmSuite(algorithmSuite)
             .encryptedDataKeys(Collections.singletonList(originalEncryptedDataKey))
+            .materialsDescription(contentMetadata.materialsDescription())
             .s3Request(request)
             .build()
         );
@@ -277,7 +279,7 @@ public ReEncryptInstructionFileResponse reEncryptInstructionFile(ReEncryptInstru
         RawKeyring newKeyring = reEncryptInstructionFileRequest.newKeyring();
         EncryptionMaterials encryptedMaterials = newKeyring.onEncrypt(encryptionMaterials);
 
-        final Map<String, String> newMaterialsDescription = encryptedMaterials.materialsDescription().getMaterialsDescription();
+        final MaterialsDescription newMaterialsDescription = encryptedMaterials.materialsDescription();
         //Validate that the new keyring has different materials description than the old keyring
         if (newMaterialsDescription.equals(currentKeyringMaterialsDescription)) {
             throw new S3EncryptionClientException("New keyring must have new materials description!");

src/main/java/software/amazon/encryption/s3/internal/ContentMetadata.java

@@ -5,6 +5,7 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
 import software.amazon.encryption.s3.materials.EncryptedDataKey;
+import software.amazon.encryption.s3.materials.MaterialsDescription;
 
 import java.util.Collections;
 import java.util.Map;
@@ -17,11 +18,14 @@ public class ContentMetadata {
     private final String _encryptedDataKeyAlgorithm;
 
     /**
-     * This field stores either encryption context or material description.
-     * We use a single field to store both in order to maintain backwards
-     * compatibility with V2, which treated both as the same.
+     * This field stores the encryption context.
      */
-    private final Map<String, String> _encryptionContextOrMatDesc;
+    private final Map<String, String> _encryptionContext;
+
+    /**
+     * This field stores the materials description used for RSA and AES keyrings.
+     */
+    private final MaterialsDescription _materialsDescription;
 
     private final byte[] _contentIv;
     private final String _contentCipher;
@@ -33,7 +37,8 @@ private ContentMetadata(Builder builder) {
 
         _encryptedDataKey = builder._encryptedDataKey;
         _encryptedDataKeyAlgorithm = builder._encryptedDataKeyAlgorithm;
-        _encryptionContextOrMatDesc = builder._encryptionContextOrMatDesc;
+        _encryptionContext = builder._encryptionContext;
+        _materialsDescription = builder._materialsDescription;
 
         _contentIv = builder._contentIv;
         _contentCipher = builder._contentCipher;
@@ -64,8 +69,16 @@ public String encryptedDataKeyAlgorithm() {
      */
     @SuppressFBWarnings(value = "EI_EXPOSE_REP", justification = "False positive; underlying"
         + " implementation is immutable")
-    public Map<String, String> encryptedDataKeyMatDescOrContext() {
-        return _encryptionContextOrMatDesc;
+    public Map<String, String> encryptionContext() {
+        return _encryptionContext;
+    }
+
+    /**
+     * Returns the materials description used for RSA and AES keyrings.
+     * @return the materials description
+     */
+    public MaterialsDescription materialsDescription() {
+        return _materialsDescription;
     }
 
     public byte[] contentIv() {
@@ -92,7 +105,8 @@ public static class Builder {
 
         private EncryptedDataKey _encryptedDataKey;
         private String _encryptedDataKeyAlgorithm;
-        private Map<String, String> _encryptionContextOrMatDesc;
+        private Map<String, String> _encryptionContext;
+        private MaterialsDescription _materialsDescription = MaterialsDescription.builder().build();
 
         private byte[] _contentIv;
         private String _contentCipher;
@@ -118,8 +132,15 @@ public Builder encryptedDataKeyAlgorithm(String encryptedDataKeyAlgorithm) {
             return this;
         }
 
-        public Builder encryptionContextOrMatDesc(Map<String, String> encryptionContextOrMatDesc) {
-            _encryptionContextOrMatDesc = Collections.unmodifiableMap(encryptionContextOrMatDesc);
+        public Builder encryptionContext(Map<String, String> encryptionContext) {
+            _encryptionContext = Collections.unmodifiableMap(encryptionContext);
+            return this;
+        }
+
+        public Builder materialsDescription(MaterialsDescription materialsDescription) {
+            _materialsDescription = materialsDescription == null
+                    ? MaterialsDescription.builder().build()
+                    : materialsDescription;
             return this;
         }
 

src/main/java/software/amazon/encryption/s3/internal/ContentMetadataDecodingStrategy.java

@@ -13,6 +13,7 @@
 import software.amazon.encryption.s3.S3EncryptionClientException;
 import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
 import software.amazon.encryption.s3.materials.EncryptedDataKey;
+import software.amazon.encryption.s3.materials.MaterialsDescription;
 import software.amazon.encryption.s3.materials.S3Keyring;
 
 import java.io.ByteArrayOutputStream;
@@ -137,8 +138,8 @@ private ContentMetadata readFromMap(Map<String, String> metadata, GetObjectRespo
                 .keyProviderInfo(keyProviderInfo.getBytes(StandardCharsets.UTF_8))
                 .build();
 
-        // Get encrypted data key encryption context or materials description (depending on the keyring)
-        final Map<String, String> encryptionContextOrMatDesc = new HashMap<>();
+        // Parse the JSON materials description or encryption context
+        final Map<String, String> matDescMap = new HashMap<>();
         // The V2 client treats null value here as empty, do the same to avoid incompatibility
         String jsonEncryptionContext = metadata.getOrDefault(MetadataKeyConstants.ENCRYPTED_DATA_KEY_CONTEXT, "{}");
         // When the encryption context contains non-US-ASCII characters,
@@ -150,19 +151,37 @@ private ContentMetadata readFromMap(Map<String, String> metadata, GetObjectRespo
             JsonNode objectNode = parser.parse(decodedJsonEncryptionContext);
 
             for (Map.Entry<String, JsonNode> entry : objectNode.asObject().entrySet()) {
-                encryptionContextOrMatDesc.put(entry.getKey(), entry.getValue().asString());
+                matDescMap.put(entry.getKey(), entry.getValue().asString());
             }
         } catch (Exception e) {
             throw new RuntimeException(e);
         }
 
+        // By default, assume the context is a materials description unless it's a KMS keyring
+        Map<String, String> encryptionContext;
+        MaterialsDescription materialsDescription;
+
+        if (keyProviderInfo.contains("kms")) {
+            // For KMS keyrings, use the map as encryption context
+            encryptionContext = matDescMap;
+            materialsDescription = MaterialsDescription.builder().build();
+        } else {
+            // For all other keyrings (AES, RSA), use the map as materials description
+            materialsDescription = MaterialsDescription.builder()
+                .putAll(matDescMap)
+                .build();
+            // Set an empty encryption context
+            encryptionContext = new HashMap<>();
+        }
+
         // Get content iv
         byte[] iv = DECODER.decode(metadata.get(MetadataKeyConstants.CONTENT_IV));
 
         return ContentMetadata.builder()
                 .algorithmSuite(algorithmSuite)
                 .encryptedDataKey(edk)
-                .encryptionContextOrMatDesc(encryptionContextOrMatDesc)
+                .encryptionContext(encryptionContext)
+                .materialsDescription(materialsDescription)
                 .contentIv(iv)
                 .contentRange(contentRange)
                 .build();

src/main/java/software/amazon/encryption/s3/internal/GetEncryptedObjectPipeline.java

@@ -91,7 +91,8 @@ private DecryptionMaterials prepareMaterialsFromRequest(final GetObjectRequest g
                 .s3Request(getObjectRequest)
                 .algorithmSuite(algorithmSuite)
                 .encryptedDataKeys(encryptedDataKeys)
-                .encryptionContext(contentMetadata.encryptedDataKeyMatDescOrContext())
+                .encryptionContext(contentMetadata.encryptionContext())
+                .materialsDescription(contentMetadata.materialsDescription())
                 .ciphertextLength(getObjectResponse.contentLength())
                 .contentRange(getObjectRequest.range())
                 .build();

src/main/java/software/amazon/encryption/s3/materials/AesKeyMaterial.java

@@ -0,0 +1,67 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.encryption.s3.materials;
+
+import javax.crypto.SecretKey;
+
+/**
+ * A concrete implementation of RawKeyMaterial for AES keys.
+ * This class provides a more convenient way to create key material for AES keyrings
+ * without having to specify the generic type parameter.
+ */
+public class AesKeyMaterial extends RawKeyMaterial<SecretKey> {
+
+    /**
+     * Creates a new AesKeyMaterial with the specified materials description and key material.
+     *
+     * @param materialsDescription the materials description
+     * @param keyMaterial the AES key material
+     */
+    public AesKeyMaterial(MaterialsDescription materialsDescription, SecretKey keyMaterial) {
+        super(materialsDescription, keyMaterial);
+    }
+
+    /**
+     * @return a new builder instance for AesKeyMaterial
+     */
+    public static Builder aesBuilder() {
+        return new Builder();
+    }
+
+    /**
+     * Builder for AesKeyMaterial.
+     */
+    public static class Builder {
+        private MaterialsDescription _materialsDescription;
+        private SecretKey _keyMaterial;
+
+        /**
+         * Sets the materials description for this AES key material.
+         *
+         * @param materialsDescription the materials description
+         * @return a reference to this object so that method calls can be chained together.
+         */
+        public Builder materialsDescription(MaterialsDescription materialsDescription) {
+            this._materialsDescription = materialsDescription;
+            return this;
+        }
+
+        /**
+         * Sets the AES key material.
+         *
+         * @param keyMaterial the AES key material
+         * @return a reference to this object so that method calls can be chained together.
+         */
+        public Builder keyMaterial(SecretKey keyMaterial) {
+            this._keyMaterial = keyMaterial;
+            return this;
+        }
+
+        /**
+         * @return the built AesKeyMaterial
+         */
+        public AesKeyMaterial build() {
+            return new AesKeyMaterial(_materialsDescription, _keyMaterial);
+        }
+    }
+}

src/main/java/software/amazon/encryption/s3/materials/AesKeyring.java

@@ -20,7 +20,7 @@
  * This keyring can wrap keys with the active keywrap algorithm and
  * unwrap with the active and legacy algorithms for AES keys.
  */
-public class AesKeyring extends RawKeyring {
+public class AesKeyring extends RawKeyring<SecretKey> {
 
     private static final String KEY_ALGORITHM = "AES";
 
@@ -41,13 +41,16 @@ public String keyProviderInfo() {
             return KEY_PROVIDER_INFO;
         }
 
-        @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
-            final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
-            cipher.init(Cipher.DECRYPT_MODE, _wrappingKey);
+    @Override
+    public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
+        // Find the appropriate key material to use for decryption
+        SecretKey keyToUse = findKeyMaterialForDecryption(materials, _wrappingKey);
 
-            return cipher.doFinal(encryptedDataKey);
-        }
+        final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
+        cipher.init(Cipher.DECRYPT_MODE, keyToUse);
+
+        return cipher.doFinal(encryptedDataKey);
+    }
     };
 
     private final DecryptDataKeyStrategy _aesWrapStrategy = new DecryptDataKeyStrategy() {
@@ -65,14 +68,17 @@ public String keyProviderInfo() {
             return KEY_PROVIDER_INFO;
         }
 
-        @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
-            final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
-            cipher.init(Cipher.UNWRAP_MODE, _wrappingKey);
+    @Override
+    public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
+        // Find the appropriate key material to use for decryption
+        SecretKey keyToUse = findKeyMaterialForDecryption(materials, _wrappingKey);
 
-            Key plaintextKey = cipher.unwrap(encryptedDataKey, CIPHER_ALGORITHM, Cipher.SECRET_KEY);
-            return plaintextKey.getEncoded();
-        }
+        final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
+        cipher.init(Cipher.UNWRAP_MODE, keyToUse);
+
+        Key plaintextKey = cipher.unwrap(encryptedDataKey, CIPHER_ALGORITHM, Cipher.SECRET_KEY);
+        return plaintextKey.getEncoded();
+    }
     };
 
     private final DataKeyStrategy _aesGcmStrategy = new DataKeyStrategy() {
@@ -126,22 +132,25 @@ public byte[] encryptDataKey(SecureRandom secureRandom,
             return encodedBytes;
         }
 
-        @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
-            byte[] iv = new byte[IV_LENGTH_BYTES];
-            byte[] ciphertext = new byte[encryptedDataKey.length - iv.length];
+    @Override
+    public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
+        byte[] iv = new byte[IV_LENGTH_BYTES];
+        byte[] ciphertext = new byte[encryptedDataKey.length - iv.length];
 
-            System.arraycopy(encryptedDataKey, 0, iv, 0, iv.length);
-            System.arraycopy(encryptedDataKey, iv.length, ciphertext, 0, ciphertext.length);
+        System.arraycopy(encryptedDataKey, 0, iv, 0, iv.length);
+        System.arraycopy(encryptedDataKey, iv.length, ciphertext, 0, ciphertext.length);
 
-            GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_BITS, iv);
-            final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
-            cipher.init(Cipher.DECRYPT_MODE, _wrappingKey, gcmParameterSpec);
+        // Find the appropriate key material to use for decryption
+        SecretKey keyToUse = findKeyMaterialForDecryption(materials, _wrappingKey);
 
-            final byte[] aADBytes = AlgorithmSuite.ALG_AES_256_GCM_IV12_TAG16_NO_KDF.cipherName().getBytes(StandardCharsets.UTF_8);
-            cipher.updateAAD(aADBytes);
-            return cipher.doFinal(ciphertext);
-        }
+        GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_BITS, iv);
+        final Cipher cipher = CryptoFactory.createCipher(CIPHER_ALGORITHM, materials.cryptoProvider());
+        cipher.init(Cipher.DECRYPT_MODE, keyToUse, gcmParameterSpec);
+
+        final byte[] aADBytes = AlgorithmSuite.ALG_AES_256_GCM_IV12_TAG16_NO_KDF.cipherName().getBytes(StandardCharsets.UTF_8);
+        cipher.updateAAD(aADBytes);
+        return cipher.doFinal(ciphertext);
+    }
     };
 
     private final Map<String, DecryptDataKeyStrategy> decryptDataKeyStrategies = new HashMap<>();
@@ -175,7 +184,7 @@ protected Map<String, DecryptDataKeyStrategy> decryptDataKeyStrategies() {
         return decryptDataKeyStrategies;
     }
 
-  public static class Builder extends RawKeyring.Builder<AesKeyring, Builder> {
+  public static class Builder extends RawKeyring.Builder<AesKeyring, Builder, SecretKey> {
         private SecretKey _wrappingKey;
 
         private Builder() {