Add support for EncryptionContext overrides to the DynamoDBEncryptor

examples/com/amazonaws/examples/AwsKmsEncryptedObject.java

@@ -52,7 +52,8 @@ public static void encryptRecord(final String cmkArn, final String region) {
     // Encryptor creation
     final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);
     // Mapper Creation
-    // Please note the use of SaveBehavior.CLOBBER. Omitting this can result in data-corruption.
+    // Please note the use of SaveBehavior.CLOBBER (SaveBehavior.PUT works as well).
+    // Omitting this can result in data-corruption.
     DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder().withSaveBehavior(SaveBehavior.CLOBBER).build();
     DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
 

examples/com/amazonaws/examples/EncryptionContextOverridesWithDynamoDBMapper.java

@@ -0,0 +1,155 @@
+package com.amazonaws.examples;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
+import com.amazonaws.services.dynamodbv2.datamodeling.AttributeEncryptor;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBAttribute;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBHashKey;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapper;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapperConfig;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBRangeKey;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBTable;
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DynamoDBEncryptor;
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionContext;
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionFlags;
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.providers.DirectKmsMaterialProvider;
+import com.amazonaws.services.dynamodbv2.model.AttributeValue;
+import com.amazonaws.services.kms.AWSKMS;
+import com.amazonaws.services.kms.AWSKMSClientBuilder;
+
+import java.security.GeneralSecurityException;
+import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableNameUsingMap;
+
+public class EncryptionContextOverridesWithDynamoDBMapper {
+    public static void main(String[] args) throws GeneralSecurityException {
+        final String cmkArn = args[0];
+        final String region = args[1];
+        final String encryptionContextTableName = args[2];
+
+        AmazonDynamoDB ddb = null;
+        AWSKMS kms = null;
+        try {
+            ddb = AmazonDynamoDBClientBuilder.standard().withRegion(region).build();
+            kms = AWSKMSClientBuilder.standard().withRegion(region).build();
+            encryptRecord(cmkArn, encryptionContextTableName, ddb, kms);
+        } finally {
+            if (ddb != null) {
+                ddb.shutdown();
+            }
+            if (kms != null) {
+                kms.shutdown();
+            }
+        }
+    }
+
+    public static void encryptRecord(final String cmkArn,
+                                     final String newEncryptionContextTableName,
+                                     AmazonDynamoDB ddb,
+                                     AWSKMS kms) throws GeneralSecurityException {
+        // Sample object to be encrypted
+        ExampleItem record = new ExampleItem();
+        record.setPartitionAttribute("is this");
+        record.setSortAttribute(55);
+        record.setExample("my data");
+
+        // Set up our configuration and clients
+        final DirectKmsMaterialProvider cmp = new DirectKmsMaterialProvider(kms, cmkArn);
+        final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);
+
+        Map<String, String> tableNameEncryptionContextOverrides = new HashMap<>();
+        tableNameEncryptionContextOverrides.put("ExampleTableForEncryptionContextOverrides", newEncryptionContextTableName);
+        tableNameEncryptionContextOverrides.put("AnotherExampleTableForEncryptionContextOverrides", "this table doesn't exist");
+
+        // Supply an operator to override the table name used in the encryption context
+        encryptor.setEncryptionContextOverrideOperator(
+                overrideEncryptionContextTableNameUsingMap(tableNameEncryptionContextOverrides)
+        );
+
+        // Mapper Creation
+        // Please note the use of SaveBehavior.CLOBBER (SaveBehavior.PUT works as well).
+        // Omitting this can result in data-corruption.
+        DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder()
+                .withSaveBehavior(DynamoDBMapperConfig.SaveBehavior.CLOBBER).build();
+        DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
+
+        System.out.println("Plaintext Record: " + record.toString());
+        // Save the record to the DynamoDB table
+        mapper.save(record);
+
+        // Retrieve (and decrypt) it from DynamoDB
+        ExampleItem decrypted_record = mapper.load(ExampleItem.class, "is this", 55);
+        System.out.println("Decrypted Record: " + decrypted_record.toString());
+
+        // Setup new configuration to decrypt without using an overridden EncryptionContext
+        final Map<String, AttributeValue> itemKey = new HashMap<>();
+        itemKey.put("partition_attribute", new AttributeValue().withS("is this"));
+        itemKey.put("sort_attribute", new AttributeValue().withN("55"));
+
+        final EnumSet<EncryptionFlags> signOnly = EnumSet.of(EncryptionFlags.SIGN);
+        final EnumSet<EncryptionFlags> encryptAndSign = EnumSet.of(EncryptionFlags.ENCRYPT, EncryptionFlags.SIGN);
+        final Map<String, AttributeValue> encryptedItem = ddb.getItem("ExampleTableForEncryptionContextOverrides", itemKey)
+                .getItem();
+        System.out.println("Encrypted Record: " + encryptedItem);
+
+        Map<String, Set<EncryptionFlags>> encryptionFlags = new HashMap<>();
+        encryptionFlags.put("partition_attribute", signOnly);
+        encryptionFlags.put("sort_attribute", signOnly);
+        encryptionFlags.put("example", encryptAndSign);
+
+        final DynamoDBEncryptor encryptorWithoutOverrides = DynamoDBEncryptor.getInstance(cmp);
+
+        // Decrypt the record without using an overridden EncryptionContext
+        encryptorWithoutOverrides.decryptRecord(encryptedItem,
+                encryptionFlags,
+                new EncryptionContext.Builder().withHashKeyName("partition_attribute")
+                        .withRangeKeyName("sort_attribute")
+                        .withTableName(newEncryptionContextTableName)
+                        .build());
+        System.out.printf("The example item was encrypted using the table name '%s' in the EncryptionContext%n", newEncryptionContextTableName);
+    }
+
+    @DynamoDBTable(tableName = "ExampleTableForEncryptionContextOverrides")
+    public static final class ExampleItem {
+        private String partitionAttribute;
+        private int sortAttribute;
+        private String example;
+
+        @DynamoDBHashKey(attributeName = "partition_attribute")
+        public String getPartitionAttribute() {
+            return partitionAttribute;
+        }
+
+        public void setPartitionAttribute(String partitionAttribute) {
+            this.partitionAttribute = partitionAttribute;
+        }
+
+        @DynamoDBRangeKey(attributeName = "sort_attribute")
+        public int getSortAttribute() {
+            return sortAttribute;
+        }
+
+        public void setSortAttribute(int sortAttribute) {
+            this.sortAttribute = sortAttribute;
+        }
+
+        @DynamoDBAttribute(attributeName = "example")
+        public String getExample() {
+            return example;
+        }
+
+        public void setExample(String example) {
+            this.example = example;
+        }
+
+        public String toString() {
+            return String.format("{partition_attribute: %s, sort_attribute: %s, example: %s}",
+                    partitionAttribute, sortAttribute, example);
+        }
+    }
+
+}

src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/DynamoDBEncryptor.java

@@ -81,7 +81,8 @@ public class DynamoDBEncryptor {
     private final String signingAlgorithmHeader;
 
     public static final String DEFAULT_SIGNING_ALGORITHM_HEADER = DEFAULT_DESCRIPTION_BASE + "signingAlg";
-
+    private Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator;
+
     protected DynamoDBEncryptor(EncryptionMaterialsProvider provider, String descriptionBase) {
         this.encryptionMaterialsProvider = provider;
         this.descriptionBase = descriptionBase;
@@ -254,6 +255,11 @@ public Map<String, AttributeValue> decryptRecord(
             .withAttributeValues(itemAttributes)
             .build();
 
+        Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator = getEncryptionContextOverrideOperator();
+        if (encryptionContextOverrideOperator != null) {
+            context = encryptionContextOverrideOperator.apply(context);
+        }
+
         materials = encryptionMaterialsProvider.getDecryptionMaterials(context);
         decryptionKey = materials.getDecryptionKey();
         if (materialDescription.containsKey(signingAlgorithmHeader)) {
@@ -307,7 +313,13 @@ public Map<String, AttributeValue> encryptRecord(
         context = new EncryptionContext.Builder(context)
             .withAttributeValues(itemAttributes)
             .build();
-
+
+        Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator =
+                getEncryptionContextOverrideOperator();
+        if (encryptionContextOverrideOperator != null) {
+            context = encryptionContextOverrideOperator.apply(context);
+        }
+
         EncryptionMaterials materials = encryptionMaterialsProvider.getEncryptionMaterials(context);
         // We need to copy this because we modify it to record other encryption details
         Map<String, String> materialDescription = new HashMap<String, String>(
@@ -559,6 +571,24 @@ protected static Map<String, String> unmarshallDescription(AttributeValue attrib
         }
     }
 
+    /**
+     * @param encryptionContextOverrideOperator  the nullable operator which will be used to override
+     *                                           the EncryptionContext.
+     * @see com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators
+     */
+    public final void setEncryptionContextOverrideOperator(
+            Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator) {
+        this.encryptionContextOverrideOperator = encryptionContextOverrideOperator;
+    }
+
+    /**
+     * @return the operator used to override the EncryptionContext
+     * @see #setEncryptionContextOverrideOperator(Function)
+     */
+    public final Function<EncryptionContext, EncryptionContext> getEncryptionContextOverrideOperator() {
+        return encryptionContextOverrideOperator;
+    }
+
     private static byte[] toByteArray(ByteBuffer buffer) {
         buffer = buffer.duplicate();
         // We can only return the array directly if:

src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/utils/EncryptionContextOperators.java

@@ -0,0 +1,65 @@
+package com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils;
+
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionContext;
+
+import java.util.Map;
+import java.util.function.UnaryOperator;
+
+/**
+ * Implementations of common operators for overriding the EncryptionContext
+ */
+public class EncryptionContextOperators {
+    /**
+     * An operator for overriding EncryptionContext's table name for a specific DynamoDBEncryptor. If any table names or
+     * the encryption context itself is null, then it returns the original EncryptionContext.
+     *
+     * @param originalTableName the name of the table that should be overridden in the Encryption Context
+     * @param newTableName the table name that should be used in the Encryption Context
+     * @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
+     */
+    public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableName(
+            String originalTableName,
+            String newTableName) {
+        return encryptionContext -> {
+            if (encryptionContext == null
+                    || encryptionContext.getTableName() == null
+                    || originalTableName == null
+                    || newTableName == null) {
+                return encryptionContext;
+            }
+            if (originalTableName.equals(encryptionContext.getTableName())) {
+                return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
+            } else {
+                return encryptionContext;
+            }
+        };
+    }
+
+
+    /**
+     * An operator for mapping multiple table names in the Encryption Context to a new table name. If the table name for
+     * a given EncryptionContext is missing, then it returns the original EncryptionContext. Similarly, it returns the
+     * original EncryptionContext if the value it is overridden to is null, or if the original table name is null.
+     *
+     * @param tableNameOverrideMap a map specifying the names of tables that should be overridden,
+     *                             and the values to which they should be overridden. If the given table name
+     *                             corresponds to null, or isn't in the map, then the table name won't be overridden.
+     * @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
+     */
+    public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableNameUsingMap(
+            Map<String, String> tableNameOverrideMap) {
+        return encryptionContext -> {
+            if (tableNameOverrideMap == null || encryptionContext == null || encryptionContext.getTableName() == null) {
+                return encryptionContext;
+            }
+            String newTableName = tableNameOverrideMap.get(encryptionContext.getTableName());
+            if (newTableName != null) {
+                return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
+            } else {
+                return encryptionContext;
+            }
+        };
+    }
+
+
+}

src/test/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/DynamoDBEncryptorTest.java

@@ -14,8 +14,10 @@
  */
 package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
 
+import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableName;
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
@@ -31,7 +33,6 @@
 import java.security.NoSuchProviderException;
 import java.security.Security;
 import java.security.SignatureException;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -296,7 +297,71 @@ public void RsaSignedOnlyBadSignature() throws GeneralSecurityException {
         encryptedAttributes.get("hashKey").setN("666");
         encryptor.decryptAllFieldsExcept(encryptedAttributes, context, attribs.keySet().toArray(new String[0]));
     }
-
+
+    /**
+     * Tests that no exception is thrown when the encryption context override operator is null
+     * @throws GeneralSecurityException
+     */
+    @Test
+    public void testNullEncryptionContextOperator() throws GeneralSecurityException {
+        DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
+        encryptor.setEncryptionContextOverrideOperator(null);
+        encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
+    }
+
+    /**
+     * Tests decrypt and encrypt with an encryption context override operator
+     * @throws GeneralSecurityException
+     */
+    @Test
+    public void testTableNameOverriddenEncryptionContextOperator() throws GeneralSecurityException {
+        // Ensure that the table name is different from what we override the table to.
+        assertNotEquals(context.getTableName(), "TheBestTableName");
+        DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
+        encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
+        Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
+        Map<String, AttributeValue> decryptedItems = encryptor.decryptAllFieldsExcept(encryptedItems, context, Collections.emptyList());
+        assertThat(decryptedItems, AttrMatcher.match(attribs));
+    }
+
+
+    /**
+     * Tests encrypt with an encryption context override operator, and a second encryptor without an override
+     * @throws GeneralSecurityException
+     */
+    @Test
+    public void testTableNameOverriddenEncryptionContextOperatorWithSecondEncryptor() throws GeneralSecurityException {
+        // Ensure that the table name is different from what we override the table to.
+        assertNotEquals(context.getTableName(), "TheBestTableName");
+        DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
+        DynamoDBEncryptor encryptorWithoutOverride = DynamoDBEncryptor.getInstance(prov);
+        encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
+        Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
+
+        EncryptionContext expectedOverriddenContext = new EncryptionContext.Builder(context).withTableName("TheBestTableName").build();
+        Map<String, AttributeValue> decryptedItems = encryptorWithoutOverride.decryptAllFieldsExcept(encryptedItems,
+                expectedOverriddenContext, Collections.emptyList());
+        assertThat(decryptedItems, AttrMatcher.match(attribs));
+    }
+
+    /**
+     * Tests encrypt with an encryption context override operator, and a second encryptor without an override
+     * @throws GeneralSecurityException
+     */
+    @Test(expected = SignatureException.class)
+    public void testTableNameOverriddenEncryptionContextOperatorWithSecondEncryptorButTheOriginalEncryptionContext() throws GeneralSecurityException {
+        // Ensure that the table name is different from what we override the table to.
+        assertNotEquals(context.getTableName(), "TheBestTableName");
+        DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
+        DynamoDBEncryptor encryptorWithoutOverride = DynamoDBEncryptor.getInstance(prov);
+        encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
+        Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
+
+        // Use the original encryption context, and expect a signature failure
+        Map<String, AttributeValue> decryptedItems = encryptorWithoutOverride.decryptAllFieldsExcept(encryptedItems,
+                context, Collections.emptyList());
+    }
+
     @Test
     public void EcdsaSignedOnly() throws GeneralSecurityException {