feat: support context_and_sign in enhanced client

DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEncryptionSignAndIncludeInEncryptionContext.java

@@ -0,0 +1,13 @@
+package software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag;
+
+@Target({ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+@BeanTableSchemaAttributeTag(EncryptionAttributeTags.class)
+public @interface DynamoDbEncryptionSignAndIncludeInEncryptionContext {
+}

DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java

@@ -24,6 +24,7 @@
 
 import static software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.DoNothingTag.CUSTOM_DDB_ENCRYPTION_DO_NOTHING_PREFIX;
 import static software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.SignOnlyTag.CUSTOM_DDB_ENCRYPTION_SIGN_ONLY_PREFIX;
+import static software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.SignAndIncludeInEncryptionContextTag.CUSTOM_DDB_ENCRYPTION_SIGN_AND_INCLUDE_PREFIX;
 
 public class DynamoDbEnhancedClientEncryption {
     public static DynamoDbEncryptionInterceptor CreateDynamoDbEncryptionInterceptor(
@@ -66,6 +67,7 @@ private static DynamoDbTableEncryptionConfig getTableConfig(
 
         TableSchema<?> topTableSchema = configWithSchema.schemaOnEncrypt();
         Set<String> signOnlyAttributes = getSignOnlyAttributes(topTableSchema);
+        Set<String> signAndIncludeAttributes = getSignAndIncludeInEncryptionContextAttributes(topTableSchema);
         Set<String> doNothingAttributes = getDoNothingAttributes(topTableSchema);
         Set<String> keyAttributes = attributeNamesUsedInIndices(topTableSchema.tableMetadata());
 
@@ -81,17 +83,34 @@ private static DynamoDbTableEncryptionConfig getTableConfig(
                         "Cannot use @DynamoDbEncryptionDoNothing and @DynamoDbEncryptionSignOnly on same attribute. Found on Table Name: %s",
                         tableName))
                     .build();
+        } else if (!Collections.disjoint(signOnlyAttributes, signAndIncludeAttributes)) {
+            throw DynamoDbEncryptionException.builder()
+                    .message(String.format(
+                        "Cannot use @DynamoDbEncryptionSignAndIncludeInEncryptionContext and @DynamoDbEncryptionSignOnly on same attribute. Found on Table Name: %s",
+                        tableName))
+                    .build();
+        } else if (!Collections.disjoint(doNothingAttributes, signAndIncludeAttributes)) {
+            throw DynamoDbEncryptionException.builder()
+                    .message(String.format(
+                        "Cannot use @DynamoDbEncryptionSignAndIncludeInEncryptionContext and @DynamoDbEncryptionDoNothing on same attribute. Found on Table Name: %s",
+                        tableName))
+                    .build();
         }
 
         List<String> attributeNames = topTableSchema.attributeNames();
         StringBuilder path = new StringBuilder();
         path.append(tableName).append(".");
         for (String attributeName : attributeNames) {
             if (keyAttributes.contains(attributeName)) {
-                // key attributes are always SIGN_ONLY
-                actions.put(attributeName, CryptoAction.SIGN_ONLY);
+		if (signAndIncludeAttributes.isEmpty()) {
+                    actions.put(attributeName, CryptoAction.SIGN_ONLY);
+		} else {
+                    actions.put(attributeName, CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT);
+		}
             } else if (signOnlyAttributes.contains(attributeName)) {
                 actions.put(attributeName, CryptoAction.SIGN_ONLY);
+            } else if (signAndIncludeAttributes.contains(attributeName)) {
+                actions.put(attributeName, CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT);
             } else if (doNothingAttributes.contains(attributeName)) {
                 actions.put(attributeName, CryptoAction.DO_NOTHING);
             } else {
@@ -142,6 +161,13 @@ private static Set<String> getSignOnlyAttributes(TableSchema<?> tableSchema) {
             .orElseGet(HashSet::new);
     }
 
+    @SuppressWarnings("unchecked")
+    private static Set<String> getSignAndIncludeInEncryptionContextAttributes(TableSchema<?> tableSchema) {
+        return tableSchema.tableMetadata()
+            .customMetadataObject(CUSTOM_DDB_ENCRYPTION_SIGN_AND_INCLUDE_PREFIX, Set.class)
+            .orElseGet(HashSet::new);
+    }
+
     @SuppressWarnings("unchecked")
     private static Set<String> getDoNothingAttributes(TableSchema<?> tableSchema) {
         return tableSchema.tableMetadata()
@@ -186,6 +212,16 @@ private static void scanForIgnoredEncryptionTags(
                         attributePath.append(signOnlyAttributes.toArray()[0])))
                     .build();
             }
+            Set<String> signAndIncludeAttributes = getSignAndIncludeInEncryptionContextAttributes(subTableSchema);
+            if (signAndIncludeAttributes.size() > 0) {
+                throw DynamoDbEncryptionException.builder()
+                    .message(String.format(
+                        "Detected DynamoDbEncryption Tag %s on a nested attribute with Path %s. " +
+                            "This is NOT Supported at this time!",
+                        CUSTOM_DDB_ENCRYPTION_SIGN_AND_INCLUDE_PREFIX,
+                        attributePath.append(signAndIncludeAttributes.toArray()[0])))
+                    .build();
+            }
             Set<String> doNothingAttributes = getDoNothingAttributes(subTableSchema);
             if (doNothingAttributes.size() > 0) {
                 throw DynamoDbEncryptionException.builder()

DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/EncryptionAttributeTags.java

@@ -10,6 +10,10 @@ public static StaticAttributeTag attributeTagFor(DynamoDbEncryptionSignOnly anno
         return new SignOnlyTag();
     }
 
+    public static StaticAttributeTag attributeTagFor(DynamoDbEncryptionSignAndIncludeInEncryptionContext annotation) {
+        return new SignAndIncludeInEncryptionContextTag();
+    }
+
     public static StaticAttributeTag attributeTagFor(DynamoDbEncryptionDoNothing annotation) {
         return new DoNothingTag();
     }

DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/SignAndIncludeInEncryptionContextTag.java

@@ -0,0 +1,19 @@
+package software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.function.Consumer;
+
+import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
+
+public class SignAndIncludeInEncryptionContextTag implements StaticAttributeTag {
+    public static final String CUSTOM_DDB_ENCRYPTION_SIGN_AND_INCLUDE_PREFIX = "DynamoDbEncryption:SignAndIncludeInEncryptionContext";
+
+    @Override
+    public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName, AttributeValueType attributeValueType) {
+        return metadata -> metadata
+                .addCustomMetadataObject(CUSTOM_DDB_ENCRYPTION_SIGN_AND_INCLUDE_PREFIX, Collections.singleton(attributeName));
+    }
+}

DynamoDbEncryption/runtimes/java/src/test/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryptionTest.java

@@ -32,6 +32,7 @@ public class DynamoDbEnhancedClientEncryptionTest {
     public void TestMultipleTables() {
         TableSchema<SimpleClass> simpleSchema = TableSchema.fromBean(SimpleClass.class);
         TableSchema<SignOnlyClass> signOnlySchema = TableSchema.fromBean(SignOnlyClass.class);
+        TableSchema<SignAndIncludeInEncryptionContextClass> signAndIncludeSchema = TableSchema.fromBean(SignAndIncludeInEncryptionContextClass.class);
         Map<String, DynamoDbEnhancedTableEncryptionConfig> tableConfigs = new HashMap<>();
         tableConfigs.put("SimpleClassTestTable",
                 DynamoDbEnhancedTableEncryptionConfig.builder()
@@ -46,13 +47,19 @@ public void TestMultipleTables() {
                         .keyring(createKmsKeyring())
                         .schemaOnEncrypt(signOnlySchema)
                         .build());
+        tableConfigs.put("SignAndIncludeInEncryptionContextClassTestTable",
+                DynamoDbEnhancedTableEncryptionConfig.builder()
+                        .logicalTableName("SignAndIncludeInEncryptionContextClassTestTable")
+                        .keyring(createKmsKeyring())
+                        .schemaOnEncrypt(signAndIncludeSchema)
+                        .build());
         DynamoDbEncryptionInterceptor interceptor =
                 DynamoDbEnhancedClientEncryption.CreateDynamoDbEncryptionInterceptor(
                         CreateDynamoDbEncryptionInterceptorInput.builder()
                                 .tableEncryptionConfigs(tableConfigs)
                                 .build()
                 );
-        assertEquals(2, interceptor.config().tableEncryptionConfigs().size());
+        assertEquals(3, interceptor.config().tableEncryptionConfigs().size());
 
         DynamoDbTableEncryptionConfig simpleConfig = interceptor.config().tableEncryptionConfigs().get("SimpleClassTestTable");
         assertEquals(CryptoAction.DO_NOTHING, simpleConfig.attributeActionsOnEncrypt().get("doNothing"));
@@ -66,6 +73,12 @@ public void TestMultipleTables() {
         assertEquals(CryptoAction.SIGN_ONLY, signOnlyConfig.attributeActionsOnEncrypt().get("sort_key"));
         assertEquals(CryptoAction.SIGN_ONLY, signOnlyConfig.attributeActionsOnEncrypt().get("attr1"));
         assertEquals(CryptoAction.SIGN_ONLY, signOnlyConfig.attributeActionsOnEncrypt().get("attr2"));
+
+        DynamoDbTableEncryptionConfig signAndIncludeConfig = interceptor.config().tableEncryptionConfigs().get("SignAndIncludeInEncryptionContextClassTestTable");
+        assertEquals(CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT, signAndIncludeConfig.attributeActionsOnEncrypt().get("partition_key"));
+        assertEquals(CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT, signAndIncludeConfig.attributeActionsOnEncrypt().get("sort_key"));
+        assertEquals(CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT, signAndIncludeConfig.attributeActionsOnEncrypt().get("attr1"));
+        assertEquals(CryptoAction.SIGN_ONLY, signAndIncludeConfig.attributeActionsOnEncrypt().get("attr2"));
     }
 
     @Test

DynamoDbEncryption/runtimes/java/src/test/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/validdatamodels/SignAndIncludeInEncryptionContextClass.java

@@ -0,0 +1,55 @@
+package software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.validdatamodels;
+
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbAttribute;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPartitionKey;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSortKey;
+import software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.DynamoDbEncryptionSignOnly;
+import software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient.DynamoDbEncryptionSignAndIncludeInEncryptionContext;
+
+@DynamoDbBean
+public class SignAndIncludeInEncryptionContextClass {
+
+    private String partitionKey;
+    private int sortKey;
+    private String attr1;
+    private String attr2;
+
+    @DynamoDbPartitionKey
+    @DynamoDbAttribute(value = "partition_key")
+    public String getPartitionKey() {
+        return this.partitionKey;
+    }
+
+    public void setPartitionKey(String partitionKey) {
+        this.partitionKey = partitionKey;
+    }
+
+    @DynamoDbSortKey
+    @DynamoDbAttribute(value = "sort_key")
+    public int getSortKey() {
+        return this.sortKey;
+    }
+
+    public void setSortKey(int sortKey) {
+        this.sortKey = sortKey;
+    }
+
+    @DynamoDbEncryptionSignAndIncludeInEncryptionContext
+    public String getAttr1() {
+        return this.attr1;
+    }
+
+    public void setAttr1(String attr1) {
+        this.attr1 = attr1;
+    }
+
+    @DynamoDbEncryptionSignOnly
+    public String getAttr2() {
+        return this.attr2;
+    }
+
+    public void setAttr2(String attr2) {
+        this.attr2 = attr2;
+    }
+}