Skip to content

Commit cfb0fa0

Browse files
johnwalkerjohnwalker
committed
Add support for EncryptionContext overrides to the DynamoDBEncryptor
There are people asking for overrides, and it would be better if they didn't need to wait for longer term, internal refactors before they are able to use overrides. This adds optional operators that give the client the last say on the EncryptionContext's value. Note: I haven't tested the example, since I didn't figure out a way to run them. How are we running example code? If we don't have a suggested way yet, I can figure something out and document it.
1 parent cda6411 commit cfb0fa0

File tree

6 files changed

+433
-5
lines changed

6 files changed

+433
-5
lines changed

‎examples/com/amazonaws/examples/AwsKmsEncryptedObject.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,8 @@ public static void encryptRecord(final String cmkArn, final String region) {
5252
// Encryptor creation
5353
final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);
5454
// Mapper Creation
55-
// Please note the use of SaveBehavior.CLOBBER. Omitting this can result in data-corruption.
55+
// Please note the use of SaveBehavior.CLOBBER (SaveBehavior.PUT works as well).
56+
// Omitting this can result in data-corruption.
5657
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder().withSaveBehavior(SaveBehavior.CLOBBER).build();
5758
DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
5859

‎examples/com/amazonaws/examples/EncryptionContextOverridesWithDynamoDBMapper.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package com.amazonaws.examples;
2+
3+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
4+
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClientBuilder;
5+
import com.amazonaws.services.dynamodbv2.datamodeling.AttributeEncryptor;
6+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBAttribute;
7+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBHashKey;
8+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapper;
9+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapperConfig;
10+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBRangeKey;
11+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBTable;
12+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DynamoDBEncryptor;
13+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.providers.DirectKmsMaterialProvider;
14+
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
15+
import com.amazonaws.services.kms.AWSKMS;
16+
import com.amazonaws.services.kms.AWSKMSClientBuilder;
17+
18+
import java.security.GeneralSecurityException;
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
22+
import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableNameUsingMap;
23+
24+
public class EncryptionContextOverridesWithDynamoDBMapper {
25+
public static void main(String[] args) throws GeneralSecurityException {
26+
final String cmkArn = args[0];
27+
final String region = args[1];
28+
final String encryptionContextTableName = args[2];
29+
30+
encryptRecord(cmkArn, region, encryptionContextTableName);
31+
}
32+
33+
public static void encryptRecord(final String cmkArn,
34+
final String region,
35+
final String newEncryptionContextTableName) {
36+
// Sample object to be encrypted
37+
ExampleItem record = new ExampleItem();
38+
record.setPartitionAttribute("is this");
39+
record.setSortAttribute(55);
40+
record.setExample("my data");
41+
42+
// Set up our configuration and clients
43+
final AmazonDynamoDB ddb = AmazonDynamoDBClientBuilder.standard().withRegion(region).build();
44+
final AWSKMS kms = AWSKMSClientBuilder.standard().withRegion(region).build();
45+
final DirectKmsMaterialProvider cmp = new DirectKmsMaterialProvider(kms, cmkArn);
46+
// Encryptor creation
47+
final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);
48+
49+
Map<String, String> tableNameEncryptionContextOverrides = new HashMap<>();
50+
tableNameEncryptionContextOverrides.put("ExampleTableForEncryptionContextOverrides", newEncryptionContextTableName);
51+
tableNameEncryptionContextOverrides.put("AnotherExampleTableForEncryptionContextOverrides", "this table doesn't exist");
52+
53+
// Here we supply an operator to override the table name used in the encryption context
54+
encryptor.setEncryptionContextOverrideOperator(
55+
overrideEncryptionContextTableNameUsingMap(tableNameEncryptionContextOverrides)
56+
);
57+
58+
// Mapper Creation
59+
// Please note the use of SaveBehavior.CLOBBER (SaveBehavior.PUT works as well).
60+
// Omitting this can result in data-corruption.
61+
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder()
62+
.withSaveBehavior(DynamoDBMapperConfig.SaveBehavior.CLOBBER).build();
63+
DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
64+
65+
System.out.println("Plaintext Record: " + record);
66+
// Save the record to the DynamoDB table
67+
mapper.save(record);
68+
69+
// Retrieve the encrypted record (directly without decrypting) from Dynamo so we can see it in our example
70+
final Map<String, AttributeValue> itemKey = new HashMap<>();
71+
itemKey.put("partition_attribute", new AttributeValue().withS("is this"));
72+
itemKey.put("sort_attribute", new AttributeValue().withN("55"));
73+
System.out.println("Encrypted Record: " + ddb.getItem("ExampleTableForEncryptionContextOverrides",
74+
itemKey).getItem());
75+
76+
// Retrieve (and decrypt) it from DynamoDB
77+
ExampleItem decrypted_record = mapper.load(ExampleItem.class, "is this", 55);
78+
System.out.println("Decrypted Record: " + decrypted_record);
79+
}
80+
81+
@DynamoDBTable(tableName = "ExampleTableForEncryptionContextOverrides")
82+
public static final class ExampleItem {
83+
private String partitionAttribute;
84+
private int sortAttribute;
85+
private String example;
86+
87+
@DynamoDBHashKey(attributeName = "partition_attribute")
88+
public String getPartitionAttribute() {
89+
return partitionAttribute;
90+
}
91+
92+
public void setPartitionAttribute(String partitionAttribute) {
93+
this.partitionAttribute = partitionAttribute;
94+
}
95+
96+
@DynamoDBRangeKey(attributeName = "sort_attribute")
97+
public int getSortAttribute() {
98+
return sortAttribute;
99+
}
100+
101+
public void setSortAttribute(int sortAttribute) {
102+
this.sortAttribute = sortAttribute;
103+
}
104+
105+
@DynamoDBAttribute(attributeName = "example")
106+
public String getExample() {
107+
return example;
108+
}
109+
110+
public void setExample(String example) {
111+
this.example = example;
112+
}
113+
}
114+
115+
}

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

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,8 @@ public class DynamoDBEncryptor {
8181
private final String signingAlgorithmHeader;
8282

8383
public static final String DEFAULT_SIGNING_ALGORITHM_HEADER = DEFAULT_DESCRIPTION_BASE + "signingAlg";
84-
84+
private Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator;
85+
8586
protected DynamoDBEncryptor(EncryptionMaterialsProvider provider, String descriptionBase) {
8687
this.encryptionMaterialsProvider = provider;
8788
this.descriptionBase = descriptionBase;
@@ -254,6 +255,11 @@ public Map<String, AttributeValue> decryptRecord(
254255
.withAttributeValues(itemAttributes)
255256
.build();
256257

258+
Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator = getEncryptionContextOverrideOperator();
259+
if (encryptionContextOverrideOperator != null) {
260+
context = encryptionContextOverrideOperator.apply(context);
261+
}
262+
257263
materials = encryptionMaterialsProvider.getDecryptionMaterials(context);
258264
decryptionKey = materials.getDecryptionKey();
259265
if (materialDescription.containsKey(signingAlgorithmHeader)) {
@@ -307,7 +313,13 @@ public Map<String, AttributeValue> encryptRecord(
307313
context = new EncryptionContext.Builder(context)
308314
.withAttributeValues(itemAttributes)
309315
.build();
310-
316+
317+
Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator =
318+
getEncryptionContextOverrideOperator();
319+
if (encryptionContextOverrideOperator != null) {
320+
context = encryptionContextOverrideOperator.apply(context);
321+
}
322+
311323
EncryptionMaterials materials = encryptionMaterialsProvider.getEncryptionMaterials(context);
312324
// We need to copy this because we modify it to record other encryption details
313325
Map<String, String> materialDescription = new HashMap<String, String>(
@@ -559,6 +571,24 @@ protected static Map<String, String> unmarshallDescription(AttributeValue attrib
559571
}
560572
}
561573

574+
/**
575+
* @param encryptionContextOverrideOperator the nullable operator which will be used to override
576+
* the EncryptionContext.
577+
* @see com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators
578+
*/
579+
public final void setEncryptionContextOverrideOperator(
580+
Function<EncryptionContext, EncryptionContext> encryptionContextOverrideOperator) {
581+
this.encryptionContextOverrideOperator = encryptionContextOverrideOperator;
582+
}
583+
584+
/**
585+
* @return the operator used to override the EncryptionContext
586+
* @see #setEncryptionContextOverrideOperator(Function)
587+
*/
588+
public final Function<EncryptionContext, EncryptionContext> getEncryptionContextOverrideOperator() {
589+
return encryptionContextOverrideOperator;
590+
}
591+
562592
private static byte[] toByteArray(ByteBuffer buffer) {
563593
buffer = buffer.duplicate();
564594
// We can only return the array directly if:

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils;
2+
3+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionContext;
4+
5+
import java.util.Map;
6+
import java.util.function.UnaryOperator;
7+
8+
/**
9+
* Implementations of common operators for overriding the EncryptionContext
10+
*/
11+
public class EncryptionContextOperators {
12+
/**
13+
* An operator for overriding EncryptionContext's table name for a specific DynamoDBEncryptor. If any table names or
14+
* the encryption context itself is null, then it returns the original EncryptionContext.
15+
*
16+
* @param originalTableName the name of the table that should be overridden in the Encryption Context
17+
* @param newTableName the table name that should be used in the Encryption Context
18+
* @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
19+
*/
20+
public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableName(
21+
String originalTableName,
22+
String newTableName) {
23+
return encryptionContext -> {
24+
if (encryptionContext == null
25+
|| encryptionContext.getTableName() == null
26+
|| originalTableName == null
27+
|| newTableName == null) {
28+
return encryptionContext;
29+
}
30+
if (originalTableName.equals(encryptionContext.getTableName())) {
31+
return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
32+
} else {
33+
return encryptionContext;
34+
}
35+
};
36+
}
37+
38+
39+
/**
40+
* An operator for mapping multiple table names in the Encryption Context to a new table name. If the table name for
41+
* a given EncryptionContext is missing, then it returns the original EncryptionContext. Similarly, it returns the
42+
* original EncryptionContext if the value it is overridden to is null, or if the original table name is null.
43+
*
44+
* @param tableNameOverrideMap a map specifying the names of tables that should be overridden,
45+
* and the values to which they should be overridden. If the given table name
46+
* corresponds to null, or isn't in the map, then the table name won't be overridden.
47+
* @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
48+
*/
49+
public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableNameUsingMap(
50+
Map<String, String> tableNameOverrideMap) {
51+
return encryptionContext -> {
52+
if (tableNameOverrideMap == null || encryptionContext == null || encryptionContext.getTableName() == null) {
53+
return encryptionContext;
54+
}
55+
String newTableName = tableNameOverrideMap.get(encryptionContext.getTableName());
56+
if (newTableName != null) {
57+
return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
58+
} else {
59+
return encryptionContext;
60+
}
61+
};
62+
}
63+
64+
65+
}

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

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
*/
1515
package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
1616

17+
import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableName;
1718
import static org.junit.Assert.assertArrayEquals;
1819
import static org.junit.Assert.assertEquals;
20+
import static org.junit.Assert.assertNotEquals;
1921
import static org.junit.Assert.assertNotNull;
2022
import static org.junit.Assert.assertNull;
2123
import static org.junit.Assert.assertThat;
@@ -31,7 +33,6 @@
3133
import java.security.NoSuchProviderException;
3234
import java.security.Security;
3335
import java.security.SignatureException;
34-
import java.util.Arrays;
3536
import java.util.Collection;
3637
import java.util.Collections;
3738
import java.util.HashMap;
@@ -296,7 +297,71 @@ public void RsaSignedOnlyBadSignature() throws GeneralSecurityException {
296297
encryptedAttributes.get("hashKey").setN("666");
297298
encryptor.decryptAllFieldsExcept(encryptedAttributes, context, attribs.keySet().toArray(new String[0]));
298299
}
299-
300+
301+
/**
302+
* Tests that no exception is thrown when the encryption context override operator is null
303+
* @throws GeneralSecurityException
304+
*/
305+
@Test
306+
public void testNullEncryptionContextOperator() throws GeneralSecurityException {
307+
DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
308+
encryptor.setEncryptionContextOverrideOperator(null);
309+
encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
310+
}
311+
312+
/**
313+
* Tests decrypt and encrypt with an encryption context override operator
314+
* @throws GeneralSecurityException
315+
*/
316+
@Test
317+
public void testTableNameOverriddenEncryptionContextOperator() throws GeneralSecurityException {
318+
// Ensure that the table name is different from what we override the table to.
319+
assertNotEquals(context.getTableName(), "TheBestTableName");
320+
DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
321+
encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
322+
Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
323+
Map<String, AttributeValue> decryptedItems = encryptor.decryptAllFieldsExcept(encryptedItems, context, Collections.emptyList());
324+
assertThat(decryptedItems, AttrMatcher.match(attribs));
325+
}
326+
327+
328+
/**
329+
* Tests encrypt with an encryption context override operator, and a second encryptor without an override
330+
* @throws GeneralSecurityException
331+
*/
332+
@Test
333+
public void testTableNameOverriddenEncryptionContextOperatorWithSecondEncryptor() throws GeneralSecurityException {
334+
// Ensure that the table name is different from what we override the table to.
335+
assertNotEquals(context.getTableName(), "TheBestTableName");
336+
DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
337+
DynamoDBEncryptor encryptorWithoutOverride = DynamoDBEncryptor.getInstance(prov);
338+
encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
339+
Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
340+
341+
EncryptionContext expectedOverriddenContext = new EncryptionContext.Builder(context).withTableName("TheBestTableName").build();
342+
Map<String, AttributeValue> decryptedItems = encryptorWithoutOverride.decryptAllFieldsExcept(encryptedItems,
343+
expectedOverriddenContext, Collections.emptyList());
344+
assertThat(decryptedItems, AttrMatcher.match(attribs));
345+
}
346+
347+
/**
348+
* Tests encrypt with an encryption context override operator, and a second encryptor without an override
349+
* @throws GeneralSecurityException
350+
*/
351+
@Test(expected = SignatureException.class)
352+
public void testTableNameOverriddenEncryptionContextOperatorWithSecondEncryptorButTheOriginalEncryptionContext() throws GeneralSecurityException {
353+
// Ensure that the table name is different from what we override the table to.
354+
assertNotEquals(context.getTableName(), "TheBestTableName");
355+
DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(prov);
356+
DynamoDBEncryptor encryptorWithoutOverride = DynamoDBEncryptor.getInstance(prov);
357+
encryptor.setEncryptionContextOverrideOperator(overrideEncryptionContextTableName(context.getTableName(), "TheBestTableName"));
358+
Map<String, AttributeValue> encryptedItems = encryptor.encryptAllFieldsExcept(attribs, context, Collections.emptyList());
359+
360+
// Use the original encryption context, and expect a signature failure
361+
Map<String, AttributeValue> decryptedItems = encryptorWithoutOverride.decryptAllFieldsExcept(encryptedItems,
362+
context, Collections.emptyList());
363+
}
364+
300365
@Test
301366
public void EcdsaSignedOnly() throws GeneralSecurityException {
302367

0 commit comments

Comments
 (0)