Skip to content

Commit

Permalink
feat: Update overrideEncryptionContextTableName (#1112)
Browse files Browse the repository at this point in the history
The names in `overrideEncryptionContextTableName`
are relative to the code and not the customer using the code.

This updates these names and docs
to reflect the intuition customers have
when dealing with two different table names.

---------

Co-authored-by: texastony <5892063+texastony@users.noreply.github.com>
  • Loading branch information
seebees and texastony authored Dec 12, 2023
1 parent 0dc4898 commit c7450c5
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 17 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,8 @@
* "partition_attribute" for Strings and a sort (range) key named "sort_attribute" for numbers.
*/
public class EncryptionContextOverridesWithDynamoDBMapper {
public static final String TABLE_NAME_TO_OVERRIDE = "ExampleTableForEncryptionContextOverrides";
public static final String ORIGINAL_TABLE_NAME_TO_OVERRIDE =
"ExampleTableForEncryptionContextOverrides";
public static final String PARTITION_ATTRIBUTE = "partition_attribute";
public static final String SORT_ATTRIBUTE = "sort_attribute";

Expand Down Expand Up @@ -78,7 +79,7 @@ public static void main(String[] args) throws GeneralSecurityException {

public static void encryptRecord(
final String cmkArn,
final String newEncryptionContextTableName,
final String currentTableName,
AmazonDynamoDB ddbClient,
AWSKMS kmsClient)
throws GeneralSecurityException {
Expand All @@ -95,7 +96,7 @@ public static void encryptRecord(
final DynamoDBEncryptor encryptor = DynamoDBEncryptor.getInstance(cmp);

Map<String, String> tableNameEncryptionContextOverrides = new HashMap<>();
tableNameEncryptionContextOverrides.put(TABLE_NAME_TO_OVERRIDE, newEncryptionContextTableName);
tableNameEncryptionContextOverrides.put(ORIGINAL_TABLE_NAME_TO_OVERRIDE, currentTableName);
tableNameEncryptionContextOverrides.put(
"AnotherExampleTableForEncryptionContextOverrides", "this table doesn't exist");

Expand Down Expand Up @@ -133,7 +134,7 @@ public static void encryptRecord(
final EnumSet<EncryptionFlags> encryptAndSign =
EnumSet.of(EncryptionFlags.ENCRYPT, EncryptionFlags.SIGN);
final Map<String, AttributeValue> encryptedItem =
ddbClient.getItem(TABLE_NAME_TO_OVERRIDE, itemKey).getItem();
ddbClient.getItem(ORIGINAL_TABLE_NAME_TO_OVERRIDE, itemKey).getItem();
System.out.println("Encrypted Record: " + encryptedItem);

Map<String, Set<EncryptionFlags>> encryptionFlags = new HashMap<>();
Expand All @@ -151,19 +152,19 @@ public static void encryptRecord(
new EncryptionContext.Builder()
.withHashKeyName(PARTITION_ATTRIBUTE)
.withRangeKeyName(SORT_ATTRIBUTE)
.withTableName(newEncryptionContextTableName)
.withTableName(currentTableName)
.build());
System.out.printf(
"The example item was encrypted using the table name '%s' in the EncryptionContext%n",
newEncryptionContextTableName);
currentTableName);

// The decrypted field matches the original field before encryption
assert record
.getExample()
.equals(decrypted_without_override_record.get(STRING_FIELD_NAME).getS());
}

@DynamoDBTable(tableName = TABLE_NAME_TO_OVERRIDE)
@DynamoDBTable(tableName = ORIGINAL_TABLE_NAME_TO_OVERRIDE)
public static final class ExampleItem {
private String partitionAttribute;
private int sortAttribute;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

package com.amazonaws.examples;

import static com.amazonaws.examples.EncryptionContextOverridesWithDynamoDBMapper.ORIGINAL_TABLE_NAME_TO_OVERRIDE;
import static com.amazonaws.examples.EncryptionContextOverridesWithDynamoDBMapper.PARTITION_ATTRIBUTE;
import static com.amazonaws.examples.EncryptionContextOverridesWithDynamoDBMapper.SORT_ATTRIBUTE;
import static com.amazonaws.examples.EncryptionContextOverridesWithDynamoDBMapper.TABLE_NAME_TO_OVERRIDE;
import static com.amazonaws.examples.TestUtils.US_WEST_2;
import static com.amazonaws.examples.TestUtils.US_WEST_2_KEY_ID;
import static com.amazonaws.examples.TestUtils.createDDBTable;
Expand All @@ -26,7 +26,7 @@ public void testEncryptAndDecrypt() throws GeneralSecurityException {
final AmazonDynamoDB ddb = DynamoDBEmbedded.create();

// Create the table under test
createDDBTable(ddb, TABLE_NAME_TO_OVERRIDE, PARTITION_ATTRIBUTE, SORT_ATTRIBUTE);
createDDBTable(ddb, ORIGINAL_TABLE_NAME_TO_OVERRIDE, PARTITION_ATTRIBUTE, SORT_ATTRIBUTE);

EncryptionContextOverridesWithDynamoDBMapper.encryptRecord(
US_WEST_2_KEY_ID, OVERRIDE_TABLE_NAME, ddb, kms);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,25 +26,32 @@ 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.
* any table names or the encryption context is null, 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
* <p>The client automatically adds the current table name to the encryption context so it's bound
* to the ciphertext. Use this method when the encryption context of encrypted table items
* includes a different table name, such as when a table is backed up, or table items are
* moved/copied to a different table. If you don't override the name of the current table with the
* table name in the encryption context, decrypt fails. This override affects the encryption
* context of all table items, including newly encrypted items.
*
* @param originalTableName Use this table name in the encryption context
* @param currentTableName Override this table name 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) {
String originalTableName, String currentTableName) {
return encryptionContext -> {
if (encryptionContext == null
|| encryptionContext.getTableName() == null
|| originalTableName == null
|| newTableName == null) {
|| currentTableName == null) {
return encryptionContext;
}
if (originalTableName.equals(encryptionContext.getTableName())) {
return new EncryptionContext.Builder(encryptionContext).withTableName(newTableName).build();
return new EncryptionContext.Builder(encryptionContext)
.withTableName(currentTableName)
.build();
} else {
return encryptionContext;
}
Expand Down

0 comments on commit c7450c5

Please sign in to comment.