Skip to content

Commit 37fa00b

Browse files
ajewellamzajewellamz
authored
feat: better support of single table design (#736)
* feat: better support of single table design
1 parent 69ec293 commit 37fa00b

File tree

10 files changed

+972
-48
lines changed

10 files changed

+972
-48
lines changed

‎DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedClientEncryption.java

+153-38
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import java.util.HashMap;
55
import java.util.HashSet;
66
import java.util.List;
7+
import java.util.ArrayList;
78
import java.util.Map;
89
import java.util.Objects;
910
import java.util.Optional;
@@ -40,6 +41,7 @@ public static DynamoDbEncryptionInterceptor CreateDynamoDbEncryptionInterceptor(
4041
.build();
4142
}
4243

44+
// return all attribute names that are keys in any index
4345
private static Set<String> attributeNamesUsedInIndices(
4446
final TableMetadata tableMetadata
4547
) {
@@ -59,6 +61,7 @@ private static Set<String> attributeNamesUsedInIndices(
5961
return allIndexAttributes;
6062
}
6163

64+
// return attributes used in the primary table index
6265
private static Set<String> attributeNamesUsedInPrimaryKey(
6366
final TableMetadata tableMetadata
6467
) {
@@ -69,62 +72,85 @@ private static Set<String> attributeNamesUsedInPrimaryKey(
6972
return keyAttributes;
7073
}
7174

72-
private static DynamoDbTableEncryptionConfig getTableConfig(
73-
final DynamoDbEnhancedTableEncryptionConfig configWithSchema,
74-
final String tableName
75-
) {
76-
Map<String, CryptoAction> actions = new HashMap<>();
75+
private static void throwUsageError(String tableName, String attributeName, String usage, String usage2)
76+
{
77+
throw DynamoDbEncryptionException.builder()
78+
.message(String.format(
79+
"Attribute %s of table %s is used as both %s and %s.",
80+
attributeName, tableName, usage, usage2))
81+
.build();
82+
83+
}
84+
85+
// Any given attribute MUST have one and only one Cryptographic Action.
86+
// i.e: It can't be both SignOnly and DoNothing.
87+
// validateAttributeUsage throws an error if
88+
// the given attribute is marked with `usage` and another Cryptographic Action.
89+
// For example, for a SignOnly, signOnly will be empty, and an error must be reported
90+
// if the attribute exists in any of the other sets.
91+
private static void validateAttributeUsage(
92+
String tableName,
93+
String attributeName,
94+
String usage,
95+
Optional<Set<String>> signOnly,
96+
Optional<Set<String>> signAndInclude,
97+
Optional<Set<String>> doNothing
98+
)
99+
{
100+
if (signOnly.isPresent()) {
101+
if (signOnly.get().contains(attributeName)) {
102+
throwUsageError(tableName, attributeName, usage, "@DynamoDbEncryptionSignOnly");
103+
}
104+
}
105+
if (signAndInclude.isPresent()) {
106+
if (signAndInclude.get().contains(attributeName)) {
107+
throwUsageError(tableName, attributeName, usage, "@DynamoDbEncryptionSignAndIncludeInEncryptionContext");
108+
}
109+
}
110+
if (doNothing.isPresent()) {
111+
if (doNothing.get().contains(attributeName)) {
112+
throwUsageError(tableName, attributeName, usage, "@DynamoDbEncryptionDoNothing");
113+
}
114+
}
115+
}
77116

78-
TableSchema<?> topTableSchema = configWithSchema.schemaOnEncrypt();
117+
// return a map containing all top level attributes in the schema
118+
// If an attribute is used in an index, it is SignOnly
119+
// Else if an attribute is tagged with a single action, it gets that action
120+
// Else if an attribute is tagged with a multiple actions, an error is thrown
121+
// Else if an attribute is not tagged, it is to be encrypted
122+
private static Map<String, CryptoAction> getActionsFromSchema(String tableName, TableSchema<?> topTableSchema)
123+
{
79124
Set<String> signOnlyAttributes = getSignOnlyAttributes(topTableSchema);
80125
Set<String> signAndIncludeAttributes = getSignAndIncludeInEncryptionContextAttributes(topTableSchema);
81126
Set<String> doNothingAttributes = getDoNothingAttributes(topTableSchema);
82127
Set<String> keyAttributes = attributeNamesUsedInIndices(topTableSchema.tableMetadata());
83128
Set<String> tableKeys = attributeNamesUsedInPrimaryKey(topTableSchema.tableMetadata());
84-
85-
if (!Collections.disjoint(keyAttributes, doNothingAttributes)) {
86-
throw DynamoDbEncryptionException.builder()
87-
.message(String.format(
88-
"Cannot use @DynamoDbEncryptionDoNothing on primary key attributes. Found on Table Name: %s",
89-
tableName))
90-
.build();
91-
} else if (!Collections.disjoint(signOnlyAttributes, doNothingAttributes)) {
92-
throw DynamoDbEncryptionException.builder()
93-
.message(String.format(
94-
"Cannot use @DynamoDbEncryptionDoNothing and @DynamoDbEncryptionSignOnly on same attribute. Found on Table Name: %s",
95-
tableName))
96-
.build();
97-
} else if (!Collections.disjoint(signOnlyAttributes, signAndIncludeAttributes)) {
98-
throw DynamoDbEncryptionException.builder()
99-
.message(String.format(
100-
"Cannot use @DynamoDbEncryptionSignAndIncludeInEncryptionContext and @DynamoDbEncryptionSignOnly on same attribute. Found on Table Name: %s",
101-
tableName))
102-
.build();
103-
} else if (!Collections.disjoint(doNothingAttributes, signAndIncludeAttributes)) {
104-
throw DynamoDbEncryptionException.builder()
105-
.message(String.format(
106-
"Cannot use @DynamoDbEncryptionSignAndIncludeInEncryptionContext and @DynamoDbEncryptionDoNothing on same attribute. Found on Table Name: %s",
107-
tableName))
108-
.build();
109-
}
110-
111129
List<String> attributeNames = topTableSchema.attributeNames();
130+
131+
Map<String, CryptoAction> actions = new HashMap<>();
112132
StringBuilder path = new StringBuilder();
113133
path.append(tableName).append(".");
114134
for (String attributeName : attributeNames) {
115135
if (tableKeys.contains(attributeName)) {
116-
if (signAndIncludeAttributes.isEmpty()) {
136+
if (signAndIncludeAttributes.isEmpty()) {
137+
validateAttributeUsage(tableName, attributeName, "a primary key", Optional.empty(), Optional.of(signAndIncludeAttributes), Optional.of(doNothingAttributes));
117138
actions.put(attributeName, CryptoAction.SIGN_ONLY);
118-
} else {
139+
} else {
140+
validateAttributeUsage(tableName, attributeName, "a primary key", Optional.of(signOnlyAttributes), Optional.empty(), Optional.of(doNothingAttributes));
119141
actions.put(attributeName, CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT);
120-
}
142+
}
121143
} else if (keyAttributes.contains(attributeName)) {
144+
validateAttributeUsage(tableName, attributeName, "an index key", Optional.empty(), Optional.of(signAndIncludeAttributes), Optional.of(doNothingAttributes));
122145
actions.put(attributeName, CryptoAction.SIGN_ONLY);
123146
} else if (signOnlyAttributes.contains(attributeName)) {
147+
validateAttributeUsage(tableName, attributeName, "@DynamoDbEncryptionSignOnly", Optional.empty(), Optional.of(signAndIncludeAttributes), Optional.of(doNothingAttributes));
124148
actions.put(attributeName, CryptoAction.SIGN_ONLY);
125149
} else if (signAndIncludeAttributes.contains(attributeName)) {
150+
validateAttributeUsage(tableName, attributeName, "@DynamoDbEncryptionSignAndIncludeInEncryptionContext", Optional.of(signOnlyAttributes), Optional.empty(), Optional.of(doNothingAttributes));
126151
actions.put(attributeName, CryptoAction.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT);
127152
} else if (doNothingAttributes.contains(attributeName)) {
153+
validateAttributeUsage(tableName, attributeName, "@DynamoDbEncryptionDoNothing", Optional.of(signOnlyAttributes), Optional.of(signAndIncludeAttributes), Optional.empty());
128154
actions.put(attributeName, CryptoAction.DO_NOTHING);
129155
} else {
130156
// non-key attributes are ENCRYPT_AND_SIGN unless otherwise annotated
@@ -134,12 +160,101 @@ private static DynamoDbTableEncryptionConfig getTableConfig(
134160
// Detect Encryption Flags that are Ignored b/c they are in a Nested Class
135161
scanForIgnoredEncryptionTags(topTableSchema, attributeName, path);
136162
}
163+
return actions;
164+
}
165+
166+
// given action maps from multiple tables, merge them into one
167+
// we throw an error if the one attribute is given two different actions
168+
private static Map<String, CryptoAction> mergeActions(List<Map<String, CryptoAction>> actionList)
169+
{
170+
// most common case
171+
if (actionList.size() == 1) {
172+
return actionList.get(0);
173+
}
174+
175+
// Gather set of all attributes
176+
HashSet<String> attributes = new HashSet<>();
177+
for (Map<String, CryptoAction> config : actionList) {
178+
attributes.addAll(config.keySet());
179+
}
180+
181+
// for each attribute, ensure that everyone agrees on its action
182+
Map<String, CryptoAction> actions = new HashMap<>();
183+
for (String attr : attributes) {
184+
Optional<CryptoAction> action = Optional.empty();
185+
for (Map<String, CryptoAction> config : actionList) {
186+
CryptoAction act = config.get(attr);
187+
if (act != null) {
188+
if (action.isPresent()) {
189+
if (!action.get().equals(act)) {
190+
throw DynamoDbEncryptionException.builder()
191+
.message(String.format(
192+
"Attribute %s set to %s in one table and %s in another.",
193+
attr, action.get(), act))
194+
.build();
195+
}
196+
} else {
197+
action = Optional.of(act);
198+
}
199+
}
200+
}
201+
actions.put(attr, action.get());
202+
}
203+
return actions;
204+
}
205+
206+
// return the partition key name
207+
// throw an error if two schemas disagree
208+
private static String getPartitionKeyName(List<TableSchema<?>> schemas)
209+
{
210+
String partitionName = schemas.get(0).tableMetadata().primaryPartitionKey();
211+
for (TableSchema<?> schema : schemas) {
212+
String part = schema.tableMetadata().primaryPartitionKey();
213+
if (!partitionName.equals(part)) {
214+
throw DynamoDbEncryptionException.builder()
215+
.message(String.format(
216+
"Primary Key set to %s in one table and %s in another.",
217+
partitionName, part))
218+
.build();
219+
}
220+
}
221+
return partitionName;
222+
}
223+
224+
// return the sort key name
225+
// throw an error if two schemas disagree
226+
private static Optional<String> getSortKeyName(List<TableSchema<?>> schemas)
227+
{
228+
Optional<String> sortName = schemas.get(0).tableMetadata().primarySortKey();
229+
for (TableSchema<?> schema : schemas) {
230+
Optional<String> sort = schema.tableMetadata().primarySortKey();
231+
if (!sortName.equals(sort)) {
232+
throw DynamoDbEncryptionException.builder()
233+
.message(String.format(
234+
"Primary Key set to %s in one table and %s in another.",
235+
sortName, sort))
236+
.build();
237+
}
238+
}
239+
return sortName;
240+
}
241+
242+
// Convert enhanced client config to regular config
243+
private static DynamoDbTableEncryptionConfig getTableConfig(
244+
final DynamoDbEnhancedTableEncryptionConfig configWithSchema,
245+
final String tableName
246+
) {
247+
List<Map<String, CryptoAction>> actionList = new ArrayList<>();
248+
for (TableSchema<?> schema : configWithSchema.schemaOnEncrypt()) {
249+
actionList.add(getActionsFromSchema(tableName, schema));
250+
}
251+
Map<String, CryptoAction> actions = mergeActions(actionList);
137252

138253
DynamoDbTableEncryptionConfig.Builder builder = DynamoDbTableEncryptionConfig.builder();
139-
String partitionName = topTableSchema.tableMetadata().primaryPartitionKey();
254+
String partitionName = getPartitionKeyName(configWithSchema.schemaOnEncrypt());
140255
builder = builder.partitionKeyName(partitionName);
141256

142-
Optional<String> sortName = topTableSchema.tableMetadata().primarySortKey();
257+
Optional<String> sortName = getSortKeyName(configWithSchema.schemaOnEncrypt());
143258
if (sortName.isPresent()) {
144259
builder = builder.sortKeyName(sortName.get());
145260
}

‎DynamoDbEncryption/runtimes/java/src/main/java/software/amazon/cryptography/dbencryptionsdk/dynamodb/enhancedclient/DynamoDbEnhancedTableEncryptionConfig.java

+10-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package software.amazon.cryptography.dbencryptionsdk.dynamodb.enhancedclient;
22

33
import java.util.List;
4+
import java.util.ArrayList;
45
import java.util.Objects;
56

67
import software.amazon.cryptography.dbencryptionsdk.dynamodb.model.LegacyOverride;
@@ -14,7 +15,7 @@
1415

1516
public class DynamoDbEnhancedTableEncryptionConfig {
1617
private final String logicalTableName;
17-
private final TableSchema<?> schemaOnEncrypt;
18+
private final List<TableSchema<?>> schemaOnEncrypt;
1819
private final List<String> allowedUnsignedAttributes;
1920
private final String allowedUnsignedAttributePrefix;
2021
private final Keyring keyring;
@@ -39,7 +40,7 @@ protected DynamoDbEnhancedTableEncryptionConfig(BuilderImpl builder) {
3940

4041
public String logicalTableName() { return this.logicalTableName; }
4142

42-
public TableSchema<?> schemaOnEncrypt() {
43+
public List<TableSchema<?>> schemaOnEncrypt() {
4344
return this.schemaOnEncrypt;
4445
}
4546

@@ -83,7 +84,7 @@ public interface Builder {
8384
String logicalTableName();
8485
Builder logicalTableName(String logicalTableName);
8586
Builder schemaOnEncrypt(TableSchema<?> schemaOnEncrypt);
86-
TableSchema<?> schemaOnEncrypt();
87+
List<TableSchema<?>> schemaOnEncrypt();
8788
Builder allowedUnsignedAttributes(List<String> allowedUnsignedAttributes);
8889
List<String> allowedUnsignedAttributes();
8990
Builder allowedUnsignedAttributePrefix(String allowedUnsignedAttributePrefix);
@@ -101,7 +102,7 @@ public interface Builder {
101102

102103
protected static class BuilderImpl implements Builder {
103104
protected String logicalTableName;
104-
protected TableSchema<?> schemaOnEncrypt;
105+
protected List<TableSchema<?>> schemaOnEncrypt;
105106
protected List<String> allowedUnsignedAttributes;
106107
protected String allowedUnsignedAttributePrefix;
107108
protected Keyring keyring;
@@ -132,11 +133,14 @@ public Builder logicalTableName(String logicalTableName) {
132133
public String logicalTableName() { return this.logicalTableName; }
133134

134135
public Builder schemaOnEncrypt(TableSchema<?> schemaOnEncrypt) {
135-
this.schemaOnEncrypt = schemaOnEncrypt;
136+
if (Objects.isNull(this.schemaOnEncrypt())) {
137+
this.schemaOnEncrypt = new ArrayList();
138+
}
139+
this.schemaOnEncrypt.add(schemaOnEncrypt);
136140
return this;
137141
}
138142

139-
public TableSchema<?> schemaOnEncrypt() {
143+
public List<TableSchema<?>> schemaOnEncrypt() {
140144
return this.schemaOnEncrypt;
141145
}
142146

0 commit comments

Comments
 (0)