-
Notifications
You must be signed in to change notification settings - Fork 71
Add support for EncryptionContext overrides to the DynamoDBEncryptor #60
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
8000
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
155 changes: 155 additions & 0 deletions
155
examples/com/amazonaws/examples/EncryptionContextOverridesWithDynamoDBMapper.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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.amazo 8000 naws.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.PUT (SaveBehavior.CLOBBER works as well). | ||
// Omitting this can result in data-corruption. | ||
DynamoDBMapperConfig mapperConfig = DynamoDBMapperConfig.builder() | ||
.withSaveBehavior(DynamoDBMapperConfig.SaveBehavior.PUT).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); | ||
} | ||
} | ||
|
||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
67 changes: 67 additions & 0 deletions
67
...azonaws/services/dynamodbv2/datamodeling/encryption/utils/EncryptionContextOperators.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,67 @@ | ||
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 { | ||
johnwalker marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
// Prevent instantiation | ||
private 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; | ||
} | ||
}; | ||
} | ||
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should be a
UnaryOperator
here and throughout.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'd have liked
UnaryOperator
as well, but the chaining methods inherited fromFunction<T,T>
returns aFunction<T,T>
instead of an UnaryOperator. Ex. calling.andThen
on anUnaryOperator<T>
returns aFunction<T, T>
.I'm tempted to drop
UnaryOperator
completely and instead useFunction<T,T>
throughout, but also think that returning a more specific object from a helper class is acceptable. What do you think?