Skip to content

Commit f89ea4b

Browse files
SalusaSecondusSalusaSecondus
authored
Merge pull request #60 from johnwalker/pitr
Add support for EncryptionContext overrides to the DynamoDBEncryptor
2 parents cda6411 + 9cedf08 commit f89ea4b

File tree

7 files changed

+475
-7
lines changed

7 files changed

+475
-7
lines changed

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,9 @@ 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.
56-
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder().withSaveBehavior(SaveBehavior.CLOBBER).build();
55+
// Please note the use of SaveBehavior.PUT (SaveBehavior.CLOBBER works as well).
56+
// Omitting this can result in data-corruption.
57+
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder().withSaveBehavior(SaveBehavior.PUT).build();
5758
DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
5859

5960
System.out.println("Plaintext Record: " + record);

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

+155
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
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.EncryptionContext;
14+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionFlags;
15+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.providers.DirectKmsMaterialProvider;
16+
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
17+
import com.amazonaws.services.kms.AWSKMS;
18+
import com.amazonaws.services.kms.AWSKMSClientBuilder;
19+
20+
import java.security.GeneralSecurityException;
21+
import java.util.EnumSet;
22+
import java.util.HashMap;
23+
import java.util.Map;
24+
import java.util.Set;
25+
26+
import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableNameUsingMap;
27+
28+
public class EncryptionContextOverridesWithDynamoDBMapper {
29+
public static void main(String[] args) throws GeneralSecurityException {
30+
final String cmkArn = args[0];
31+
final String region = args[1];
32+
final String encryptionContextTableName = args[2];
33+
34+
AmazonDynamoDB ddb = null;
35+
AWSKMS kms = null;
36+
try {
37+
ddb = AmazonDynamoDBClientBuilder.standard().withRegion(region).build();
38+
kms = AWSKMSClientBuilder.standard().withRegion(region).build();
39+
encryptRecord(cmkArn, encryptionContextTableName, ddb, kms);
40+
} finally {
41+
if (ddb != null) {
42+
ddb.shutdown();
43+
}
44+
if (kms != null) {
45+
kms.shutdown();
46+
}
47+
}
48+
}
49+
50+
public static void encryptRecord(final String cmkArn,
51+
final String newEncryptionContextTableName,
52+
AmazonDynamoDB ddb,
53+
AWSKMS kms) throws GeneralSecurityException {
54+
// Sample object to be encrypted
55+
ExampleItem record = new ExampleItem();
56+
record.setPartitionAttribute("is this");
57+
record.setSortAttribute(55);
58+
record.setExample("my data");
59+
60+
// Set up our configuration and clients
61+
final DirectKmsMaterialProvider cmp = new DirectKmsMaterialProvider(kms, cmkArn);
62+
final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);
63+
64+
Map<String, String> tableNameEncryptionContextOverrides = new HashMap<>();
65+
tableNameEncryptionContextOverrides.put("ExampleTableForEncryptionContextOverrides", newEncryptionContextTableName);
66+
tableNameEncryptionContextOverrides.put("AnotherExampleTableForEncryptionContextOverrides", "this table doesn't exist");
67+
68+
// Supply an operator to override the table name used in the encryption context
69+
encryptor.setEncryptionContextOverrideOperator(
70+
overrideEncryptionContextTableNameUsingMap(tableNameEncryptionContextOverrides)
71+
);
72+
73+
// Mapper Creation
74+
// Please note the use of SaveBehavior.PUT (SaveBehavior.CLOBBER works as well).
75+
// Omitting this can result in data-corruption.
76+
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder()
77+
.withSaveBehavior(DynamoDBMapperConfig.SaveBehavior.PUT).build();
78+
DynamoDBMapper mapper = new DynamoDBMapper(ddb, mapperConfig, new AttributeEncryptor(encryptor));
79+
80+
System.out.println("Plaintext Record: " + record.toString());
81+
// Save the record to the DynamoDB table
82+
mapper.save(record);
83+
84+
// Retrieve (and decrypt) it from DynamoDB
85+
ExampleItem decrypted_record = mapper.load(ExampleItem.class, "is this", 55);
86+
System.out.println("Decrypted Record: " + decrypted_record.toString());
87+
88+
// Setup new configuration to decrypt without using an overridden EncryptionContext
89+
final Map<String, AttributeValue> itemKey = new HashMap<>();
90+
itemKey.put("partition_attribute", new AttributeValue().withS("is this"));
91+
itemKey.put("sort_attribute", new AttributeValue().withN("55"));
92+
93+
final EnumSet<EncryptionFlags> signOnly = EnumSet.of(EncryptionFlags.SIGN);
94+
final EnumSet<EncryptionFlags> encryptAndSign = EnumSet.of(EncryptionFlags.ENCRYPT, EncryptionFlags.SIGN);
95+
final Map<String, AttributeValue> encryptedItem = ddb.getItem("ExampleTableForEncryptionContextOverrides", itemKey)
96+
.getItem();
97+
System.out.println("Encrypted Record: " + encryptedItem);
98+
99+
Map<String, Set<EncryptionFlags>> encryptionFlags = new HashMap<>();
100+
encryptionFlags.put("partition_attribute", signOnly);
101+
encryptionFlags.put("sort_attribute", signOnly);
102+
encryptionFlags.put("example", encryptAndSign);
103+
104+
final DynamoDBEncryptor encryptorWithoutOverrides = DynamoDBEncryptor.getInstance(cmp);
105+
106+
// Decrypt the record without using an overridden EncryptionContext
107+
encryptorWithoutOverrides.decryptRecord(encryptedItem,
108+
encryptionFlags,
109+
new EncryptionContext.Builder().withHashKeyName("partition_attribute")
110+
.withRangeKeyName("sort_attribute")
111+
.withTableName(newEncryptionContextTableName)
112+
.build());
113+
System.out.printf("The example item was encrypted using the table name '%s' in the EncryptionContext%n", newEncryptionContextTableName);
114+
}
115+
116+
@DynamoDBTable(tableName = "ExampleTableForEncryptionContextOverrides")
117+
public static final class ExampleItem {
118+
private String partitionAttribute;
119+
private int sortAttribute;
120+
private String example;
121+
122+
@DynamoDBHashKey(attributeName = "partition_attribute")
123+
public String getPartitionAttribute() {
124+
return partitionAttribute;
125+
}
126+
127+
public void setPartitionAttribute(String partitionAttribute) {
128+
this.partitionAttribute = partitionAttribute;
129+
}
130+
131+
@DynamoDBRangeKey(attributeName = "sort_attribute")
132+
public int getSortAttribute() {
133+
return sortAttribute;
134+
}
135+
136+
public void setSortAttribute(int sortAttribute) {
137+
this.sortAttribute = sortAttribute;
138+
}
139+
140+
@DynamoDBAttribute(attributeName = "example")
141+
public String getExample() {
142+
return example;
143+
}
144+
145+
public void setExample(String example) {
146+
this.example = example;
147+
}
148+
149+
public String toString() {
150+
return String.format("{partition_attribute: %s, sort_attribute: %s, example: %s}",
151+
partitionAttribute, sortAttribute, example);
152+
}
153+
}
154+
155+
}

‎pom.xml

+1-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
<dependency>
4444
<groupId>com.amazonaws</groupId>
4545
<artifactId>aws-java-sdk-bom</artifactId>
46-
<version>1.11.434</version>
46+
<version>1.11.460</version>
4747
<type>pom</type>
4848
<scope>import</scope>
4949
</dependency>

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

+32-2
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

+67
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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+
// Prevent instantiation
14+
private EncryptionContextOperators() {
15+
}
16+
17+
/**
18+
* An operator for overriding EncryptionContext's table name for a specific DynamoDBEncryptor. If any table names or
19+
* the encryption context itself is null, then it returns the original EncryptionContext.
20+
*
21+
* @param originalTableName the name of the table that should be overridden in the Encryption Context
22+
* @param newTableName the table name that should be used in the Encryption Context
23+
* @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
24+
*/
25+
public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableName(
26+
String originalTableName,
27+
String newTableName) {
28+
return encryptionContext -> {
29+
if (encryptionContext == null
30+
|| encryptionContext.getTableName() == null
31+
|| originalTableName == null
32+
|| newTableName == null) {
33+
return encryptionContext;
34+
}
35+
if (originalTableName.equals(encryptionContext.getTableName())) {
36+
return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
37+
} else {
38+
return encryptionContext;
39+
}
40+
};
41+
}
42+
43+
/**
44+
* An operator for mapping multiple table names in the Encryption Context to a new table name. If the table name for
45+
* a given EncryptionContext is missing, then it returns the original EncryptionContext. Similarly, it returns the
46+
* original EncryptionContext if the value it is overridden to is null, or if the original table name is null.
47+
*
48+
* @param tableNameOverrideMap a map specifying the names of tables that should be overridden,
49+
* and the values to which they should be overridden. If the given table name
50+
* corresponds to null, or isn't in the map, then the table name won't be overridden.
51+
* @return A UnaryOperator that produces a new EncryptionContext with the supplied table name
52+
*/
53+
public static UnaryOperator<EncryptionContext> overrideEncryptionContextTableNameUsingMap(
54+
Map<String, String> tableNameOverrideMap) {
55+
return encryptionContext -> {
56+
if (tableNameOverrideMap == null || encryptionContext == null || encryptionContext.getTableName() == null) {
57+
return encryptionContext;
58+
}
59+
String newTableName = tableNameOverrideMap.get(encryptionContext.getTableName());
60+
if (newTableName != null) {
61+
return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
62+
} else {
63+
return encryptionContext;
64+
}
65+
};
66+
}
67+
}

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

+67-2
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)