Skip to content

Commit 3da36d9

Browse files
SalusaSecondusSalusaSecondus
committed
Adding @HandleUnknownAttributes, @TableAadOverride and lots more tests
1 parent e7888cc commit 3da36d9

22 files changed

+3135
-59
lines changed

‎src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/AttributeEncryptor.java‎

Lines changed: 160 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -14,19 +14,22 @@
1414
*/
1515
package com.amazonaws.services.dynamodbv2.datamodeling;
1616

17-
import java.lang.reflect.Method;
1817
import java.util.Collections;
1918
import java.util.EnumSet;
2019
import java.util.HashMap;
2120
import java.util.Map;
2221
import java.util.Set;
2322
import java.util.concurrent.ConcurrentHashMap;
2423

24+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMappingsRegistry.Mapping;
25+
import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMappingsRegistry.Mappings;
2526
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DoNotEncrypt;
2627
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DoNotTouch;
2728
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DynamoDBEncryptor;
2829
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionContext;
2930
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.EncryptionFlags;
31+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.HandleUnknownAttributes;
32+
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.TableAadOverride;
3033
import com.amazonaws.services.dynamodbv2.datamodeling.encryption.providers.EncryptionMaterialsProvider;
3134
import com.amazonaws.services.dynamodbv2.model.AttributeValue;
3235

@@ -38,8 +41,7 @@
3841
public class AttributeEncryptor implements AttributeTransformer {
3942
private static final DynamoDBReflector reflector = new DynamoDBReflector();
4043
private final DynamoDBEncryptor encryptor;
41-
private final Map<Class<?>, Map<String, Set<EncryptionFlags>>> flagCache =
42-
new ConcurrentHashMap<Class<?>, Map<String, Set<EncryptionFlags>>>();
44+
private final Map<Class<?>, ModelClassMetadata> metadataCache = new ConcurrentHashMap<>();
4345

4446
public AttributeEncryptor(final DynamoDBEncryptor encryptor) {
4547
this.encryptor = encryptor;
@@ -56,11 +58,11 @@ public DynamoDBEncryptor getEncryptor() {
5658
@Override
5759
public Map<String, AttributeValue> transform(final Parameters<?> parameters) {
5860
// one map of attributeFlags per model class
59-
final Map<String, Set<EncryptionFlags>> attributeFlags = getAttributeFlags(parameters);
61+
final ModelClassMetadata metadata = getModelClassMetadata(parameters);
6062
try {
6163
return encryptor.encryptRecord(
6264
parameters.getAttributeValues(),
63-
attributeFlags,
65+
metadata.getEncryptionFlags(),
6466
paramsToContext(parameters));
6567
} catch (Exception ex) {
6668
throw new DynamoDBMappingException(ex);
@@ -69,7 +71,7 @@ public Map<String, AttributeValue> transform(final Parameters<?> parameters) {
6971

7072
@Override
7173
public Map<String, AttributeValue> untransform(final Parameters<?> parameters) {
72-
final Map<String, Set<EncryptionFlags>> attributeFlags = getAttributeFlags(parameters);
74+
final Map<String, Set<EncryptionFlags>> attributeFlags = getEncryptionFlags(parameters);
7375

7476
try {
7577
return encryptor.decryptRecord(
@@ -81,49 +83,177 @@ public Map<String, AttributeValue> untransform(final Parameters<?> parameters) {
8183
}
8284
}
8385

84-
private <T> Map<String, Set<EncryptionFlags>> getAttributeFlags(Parameters<T> parameters) {
86+
/*
87+
* For any attributes we see from DynamoDB that aren't modeled in the mapper class,
88+
* we either ignore them (the default behavior), or include them for encryption/signing
89+
* based on the presence of the @HandleUnknownAttributes annotation (unless the class
90+
* has @DoNotTouch, then we don't include them).
91+
*/
92+
private Map<String, Set<EncryptionFlags>> getEncryptionFlags(final Parameters<?> parameters) {
93+
final ModelClassMetadata metadata = getModelClassMetadata(parameters);
94+
95+
// If the class is annotated with @DoNotTouch, then none of the attributes are
96+
// encrypted or signed, so we don't need to bother looking for unknown attributes.
97+
if (metadata.getDoNotTouch()) {
98+
return metadata.getEncryptionFlags();
99+
}
100+
101+
final Set<EncryptionFlags> unknownAttributeBehavior = metadata.getUnknownAttributeBehavior();
102+
final Map<String, Set<EncryptionFlags>> attributeFlags = new HashMap<>();
103+
attributeFlags.putAll(metadata.getEncryptionFlags());
104+
105+
for (final String attributeName : parameters.getAttributeValues().keySet()) {
106+
if (!attributeFlags.containsKey(attributeName) &&
107+
!encryptor.getSignatureFieldName().equals(attributeName) &&
108+
!encryptor.getMaterialDescriptionFieldName().equals(attributeName)) {
109+
110+
attributeFlags.put(attributeName, unknownAttributeBehavior);
111+
}
112+
}
113+
114+
return attributeFlags;
115+
}
116+
117+
private <T> ModelClassMetadata getModelClassMetadata(Parameters<T> parameters) {
85118
// Due to the lack of explicit synchronization, it is possible that
86119
// elements in the cache will be added multiple times. Since they will
87120
// all be identical, this is okay. Avoiding explicit synchronization
88121
// means that in the general (retrieval) case, should never block and
89122
// should be extremely fast.
90123
final Class<T> clazz = parameters.getModelClass();
91-
Map<String, Set<EncryptionFlags>> attributeFlags = flagCache.get(clazz);
92-
if (attributeFlags == null) {
93-
attributeFlags = new HashMap<String, Set<EncryptionFlags>>();
124+
ModelClassMetadata metadata = metadataCache.get(clazz);
94125

95-
final boolean encryptionEnabled = !clazz.isAnnotationPresent(DoNotEncrypt.class);
96-
final boolean doNotTouch = clazz.isAnnotationPresent(DoNotTouch.class);
126+
if (metadata == null) {
127+
Map<String, Set<EncryptionFlags>> attributeFlags = new HashMap<>();
97128

98-
if (!doNotTouch) {
99-
final Method hashKeyGetter = reflector.getPrimaryHashKeyGetter(clazz);
100-
final Method rangeKeyGetter = reflector.getPrimaryRangeKeyGetter(clazz);
129+
final boolean handleUnknownAttributes = handleUnknownAttributes(clazz);
130+
final EnumSet<EncryptionFlags> unknownAttributeBehavior = EnumSet.noneOf(EncryptionFlags.class);
101131

102-
for (Method getter : reflector.getRelevantGetters(clazz)) {
132+
if (shouldTouch(clazz)) {
133+
Mappings mappings = DynamoDBMappingsRegistry.instance().mappingsOf(clazz);
134+
135+
for (Mapping mapping : mappings.getMappings()) {
103136
final EnumSet<EncryptionFlags> flags = EnumSet.noneOf(EncryptionFlags.class);
104-
if (!getter.isAnnotationPresent(DoNotTouch.class)) {
105-
if (encryptionEnabled && !getter.isAnnotationPresent(DoNotEncrypt.class)
106-
&& !getter.equals(hashKeyGetter) && !getter.equals(rangeKeyGetter)
107-
&& !reflector.isVersionAttributeGetter(getter)) {
137+
if (shouldTouch(mapping)) {
138+
if (shouldEncryptAttribute(clazz, mapping)) {
108139
flags.add(EncryptionFlags.ENCRYPT);
109140
}
110141
flags.add(EncryptionFlags.SIGN);
111142
}
112-
attributeFlags.put(reflector.getAttributeName(getter),
113-
Collections.unmodifiableSet(flags));
143+
attributeFlags.put(mapping.getAttributeName(), Collections.unmodifiableSet(flags));
144+
}
145+
146+
if (handleUnknownAttributes) {
147+
unknownAttributeBehavior.add(EncryptionFlags.SIGN);
148+
149+
if (shouldEncrypt(clazz)) {
150+
unknownAttributeBehavior.add(EncryptionFlags.ENCRYPT);
151+
}
114152
}
115153
}
116-
flagCache.put(clazz, Collections.unmodifiableMap(attributeFlags));
154+
155+
metadata = new ModelClassMetadata(Collections.unmodifiableMap(attributeFlags), doNotTouch(clazz),
156+
Collections.unmodifiableSet(unknownAttributeBehavior));
157+
metadataCache.put(clazz, metadata);
117158
}
118-
return attributeFlags;
159+
return metadata;
160+
}
161+
162+
/**
163+
* @return True if {@link DoNotTouch} is not present on the class level. False otherwise
164+
*/
165+
private boolean shouldTouch(Class<?> clazz) {
166+
return !doNotTouch(clazz);
167+
}
168+
169+
/**
170+
* @return True if {@link DoNotTouch} is not present on the getter level. False otherwise.
171+
*/
172+
private boolean shouldTouch(Mapping mapping) {
173+
return !doNotTouch(mapping);
174+
}
175+
176+
/**
177+
* @return True if {@link DoNotTouch} IS present on the class level. False otherwise.
178+
*/
179+
private boolean doNotTouch(Class<?> clazz) {
180+
return clazz.isAnnotationPresent(DoNotTouch.class);
181+
}
182+
183+
/**
184+
* @return True if {@link DoNotTouch} IS present on the getter level. False otherwise.
185+
*/
186+
private boolean doNotTouch(Mapping mapping) {
187+
return mapping.getter().isAnnotationPresent(DoNotTouch.class);
188+
}
189+
190+
/**
191+
* @return True if {@link DoNotEncrypt} is NOT present on the class level. False otherwise.
192+
*/
193+
private boolean shouldEncrypt(Class<?> clazz) {
194+
return !doNotEncrypt(clazz);
195+
}
196+
197+
/**
198+
* @return True if {@link DoNotEncrypt} IS present on the class level. False otherwise.
199+
*/
200+
private boolean doNotEncrypt(Class<?> clazz) {
201+
return clazz.isAnnotationPresent(DoNotEncrypt.class);
202+
}
203+
204+
/**
205+
* @return True if {@link DoNotEncrypt} IS present on the getter level. False otherwise.
206+
*/
207+
private boolean doNotEncrypt(Mapping mapping) {
208+
return mapping.getter().isAnnotationPresent(DoNotEncrypt.class);
209+
}
210+
211+
/**
212+
* @return True if the attribute should be encrypted, false otherwise.
213+
*/
214+
private boolean shouldEncryptAttribute(final Class<?> clazz, final Mapping mapping) {
215+
return !(doNotEncrypt(clazz) || doNotEncrypt(mapping) || mapping.isPrimaryKey() || mapping.isVersion());
119216
}
120-
217+
121218
private static EncryptionContext paramsToContext(Parameters<?> params) {
219+
final Class<?> clazz = params.getModelClass();
220+
final TableAadOverride override = clazz.getAnnotation(TableAadOverride.class);
221+
final String tableName = ((override == null) ? params.getTableName() : override.tableName());
222+
122223
return new EncryptionContext.Builder()
123-
.withHashKeyName(params.getHashKeyName())
124-
.withRangeKeyName(params.getRangeKeyName())
125-
.withTableName(params.getTableName())
126-
.withModeledClass(params.getModelClass())
127-
.withAttributeValues(params.getAttributeValues()).build();
224+
.withHashKeyName(params.getHashKeyName())
225+
.withRangeKeyName(params.getRangeKeyName())
226+
.withTableName(tableName)
227+
.withModeledClass(params.getModelClass())
228+
.withAttributeValues(params.getAttributeValues()).build();
229+
}
230+
231+
private boolean handleUnknownAttributes(Class<?> clazz) {
232+
return clazz.getAnnotation(HandleUnknownAttributes.class) != null;
233+
}
234+
235+
private static class ModelClassMetadata {
236+
private final Map<String, Set<EncryptionFlags>> encryptionFlags;
237+
private final boolean doNotTouch;
238+
private final Set<EncryptionFlags> unknownAttributeBehavior;
239+
240+
public ModelClassMetadata(Map<String, Set<EncryptionFlags>> encryptionFlags,
241+
boolean doNotTouch, Set<EncryptionFlags> unknownAttributeBehavior) {
242+
this.encryptionFlags = encryptionFlags;
243+
this.doNotTouch = doNotTouch;
244+
this.unknownAttributeBehavior = unknownAttributeBehavior;
245+
}
246+
247+
public Map<String, Set<EncryptionFlags>> getEncryptionFlags() {
248+
return encryptionFlags;
249+
}
250+
251+
public boolean getDoNotTouch() {
252+
return doNotTouch;
253+
}
254+
255+
public Set<EncryptionFlags> getUnknownAttributeBehavior() {
256+
return unknownAttributeBehavior;
257+
}
128258
}
129259
}

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -377,9 +377,12 @@ private void actualEncryption(Map<String, AttributeValue> itemAttributes,
377377
Map<String, Set<EncryptionFlags>> attributeFlags,
378378
Map<String, String> materialDescription,
379379
SecretKey encryptionKey) throws GeneralSecurityException {
380-
materialDescription.put(symmetricEncryptionModeHeader,
381-
SYMMETRIC_ENCRYPTION_MODE);
382-
final String encryptionMode = encryptionKey != null ? encryptionKey.getAlgorithm() + SYMMETRIC_ENCRYPTION_MODE : null;
380+
String encryptionMode = null;
381+
if (encryptionKey != null) {
382+
materialDescription.put(this.symmetricEncryptionModeHeader,
383+
SYMMETRIC_ENCRYPTION_MODE);
384+
encryptionMode = encryptionKey.getAlgorithm() + SYMMETRIC_ENCRYPTION_MODE;
385+
}
383386
Cipher cipher = null;
384387
int ivSize = -1;
385388

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
/*
2+
* Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
16+
17+
import java.lang.annotation.ElementType;
18+
import java.lang.annotation.Retention;
19+
import java.lang.annotation.RetentionPolicy;
20+
import java.lang.annotation.Target;
21+
22+
/**
23+
* Marker annotation that indicates that attributes found during unmarshalling
24+
* that are in the DynamoDB item but not modeled in the mapper model class
25+
* should be included in for decryption/signature verification. The default
26+
* behavior (without this annotation) is to ignore them, which can lead to
27+
* signature verification failures when attributes are removed from model classes.
28+
*
29+
* If this annotation is added to a class with @DoNotEncrypt, then the unknown
30+
* attributes will only be included in the signature calculation, and if it's
31+
* added to a class with default encryption behavior, the unknown attributes
32+
* will be signed and decrypted.
33+
*
34+
* @author Dan Cavallaro
35+
*/
36+
@Target(value = {ElementType.TYPE})
37+
@Retention(value = RetentionPolicy.RUNTIME)
38+
public @interface HandleUnknownAttributes {}

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
/*
2+
* Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
16+
17+
import java.lang.annotation.ElementType;
18+
import java.lang.annotation.Retention;
19+
import java.lang.annotation.RetentionPolicy;
20+
import java.lang.annotation.Target;
21+
22+
/**
23+
* Overrides the default tablename used as part of the data signature with
24+
* {@code tableName} instead. This can be useful when multiple tables are
25+
* used interchangably and data should be able to be copied or moved
26+
* between them without needing to be reencrypted.
27+
*
28+
* @author Greg Rubin
29+
*/
30+
@Target(value = {ElementType.TYPE})
31+
@Retention(value = RetentionPolicy.RUNTIME)
32+
public @interface TableAadOverride {
33+
String tableName();
34+
}

‎src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/providers/store/MetaStore.java‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
11
/*
22
* Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3-
*
3+
*
44
* Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except
55
* in compliance with the License. A copy of the License is located at
6-
*
6+
*
77
* http://aws.amazon.com/apache2.0
8-
*
8+
*
99
* or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS,
1010
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
1111
* specific language governing permissions and limitations under the License.
@@ -91,7 +91,7 @@ public MetaStore(final AmazonDynamoDB ddb, final String tableName,
9191
tmpExpected.put(DEFAULT_RANGE_KEY, new ExpectedAttributeValue().withExists(false));
9292
doesNotExist = Collections.unmodifiableMap(tmpExpected);
9393
}
94-
94+
9595
@Override
9696
public EncryptionMaterialsProvider getProvider(final String materialName, final long version) {
9797
final Map<String, AttributeValue> ddbKey = new HashMap<String, AttributeValue>();
@@ -134,15 +134,15 @@ public long getMaxVersion(final String materialName) {
134134
if (items.isEmpty()) {
135135
return -1L;
136136
} else {
137-
return Long.valueOf(items.get(0).get(DEFAULT_RANGE_KEY).getN());
137+
return Long.parseLong(items.get(0).get(DEFAULT_RANGE_KEY).getN());
138138
}
139139
}
140140

141141
@Override
142142
public long getVersionFromMaterialDescription(final Map<String, String> description) {
143143
final Matcher m = COMBINED_PATTERN.matcher(description.get(META_ID));
144144
if (m.matches()) {
145-
return Long.valueOf(m.group(2));
145+
return Long.parseLong(m.group(2));
146146
} else {
147147
throw new IllegalArgumentException("No meta id found");
148148
}

0 commit comments

Comments
 (0)