Skip to content

feat: Update overrideEncryptionContextTableName #1112

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 6 commits into from
Dec 12, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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