Skip to content

Commit 4e682ea

Browse files
johnwalkerjohnwalker
committed
Direct copies of DirectKMS and MetaStore
1 parent b5fce22 commit 4e682ea

File tree

2 files changed

+187
-25
lines changed

2 files changed

+187
-25
lines changed

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

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ public DecryptionMaterials getDecryptionMaterials(EncryptionContext context) {
126126
DecryptRequest request = appendUserAgent(new DecryptRequest());
127127
request.setCiphertextBlob(ByteBuffer.wrap(Base64.decode(materialDescription.get(ENVELOPE_KEY))));
128128
request.setEncryptionContext(ec);
129-
final DecryptResult decryptResult = kms.decrypt(request);
129+
final DecryptResult decryptResult = decrypt(request, context);
130130
validateEncryptionKeyId(decryptResult.getKeyId(), context);
131131

132132
final Hkdf kdf;
@@ -167,7 +167,7 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
167167
req.setNumberOfBytes(256 / 8);
168168
req.setEncryptionContext(ec);
169169

170-
final GenerateDataKeyResult dataKeyResult = kms.generateDataKey(req);
170+
final GenerateDataKeyResult dataKeyResult = generateDataKey(req, context);
171171

172172
final Map<String, String> materialDescription = new HashMap<>();
173173
materialDescription.putAll(description);
@@ -192,7 +192,8 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
192192
}
193193

194194
/**
195-
* Get encryption key id.
195+
* Get encryption key id that is used to create the {@link EncryptionMaterials}.
196+
*
196197
* @return encryption key id.
197198
*/
198199
protected String getEncryptionKeyId() {
@@ -202,6 +203,7 @@ protected String getEncryptionKeyId() {
202203
/**
203204
* Select encryption key id to be used to generate data key. The default implementation of this method returns
204205
* {@link DirectKmsMaterialProvider#encryptionKeyId}.
206+
*
205207
* @param context encryption context.
206208
* @return the encryptionKeyId.
207209
* @throws DynamoDBMappingException when we fails to select a valid encryption key id.
@@ -211,7 +213,9 @@ protected String selectEncryptionKeyId(EncryptionContext context) throws DynamoD
211213
}
212214

213215
/**
214-
* Validate the encryption key id. The default implementation of this method does nothing.
216+
* Validate the encryption key id. The default implementation of this method does not validate
217+
* encryption key id.
218+
*
215219
* @param encryptionKeyId encryption key id from {@link DecryptResult}.
216220
* @param context encryption context.
217221
* @throws DynamoDBMappingException when encryptionKeyId is invalid.
@@ -221,6 +225,34 @@ protected void validateEncryptionKeyId(String encryptionKeyId, EncryptionContext
221225
// No action taken.
222226
}
223227

228+
/**
229+
* Decrypts ciphertext. The default implementation calls KMS to decrypt the ciphertext using the parameters
230+
* provided in the {@link DecryptRequest}. Subclass can override the default implementation to provide
231+
* additional request parameters using attributes within the {@link EncryptionContext}.
232+
*
233+
* @param request request parameters to decrypt the given ciphertext.
234+
* @param context additional useful data to decrypt the ciphertext.
235+
* @return the decrypted plaintext for the given ciphertext.
236+
*/
237+
protected DecryptResult decrypt(final DecryptRequest request, final EncryptionContext context) {
238+
return kms.decrypt(request);
239+
}
240+
241+
/**
242+
* Returns a data encryption key that you can use in your application to encrypt data locally. The default
243+
* implementation calls KMS to generate the data key using the parameters provided in the
244+
* {@link GenerateDataKeyRequest}. Subclass can override the default implementation to provide additional
245+
* request parameters using attributes within the {@link EncryptionContext}.
246+
*
247+
* @param request request parameters to generate the data key.
248+
* @param context additional useful data to generate the data key.
249+
* @return the newly generated data key which includes both the plaintext and ciphertext.
250+
*/
251+
protected GenerateDataKeyResult generateDataKey(final GenerateDataKeyRequest request,
252+
final EncryptionContext context) {
253+
return kms.generateDataKey(request);
254+
}
255+
224256
/**
225257
* Extracts relevant information from {@code context} and uses it to populate fields in
226258
* {@code kmsEc}. Currently, these fields are:

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

Lines changed: 151 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@
1717
import java.util.Arrays;
1818
import java.util.Collections;
1919
import java.util.HashMap;
20+
import java.util.HashSet;
2021
import java.util.List;
2122
import java.util.Map;
23+
import java.util.Set;
2224
import java.util.regex.Matcher;
2325
import java.util.regex.Pattern;
2426

@@ -70,25 +72,88 @@ public class MetaStore extends ProviderStore {
7072
private static final String DEFAULT_HASH_KEY = "N";
7173
private static final String DEFAULT_RANGE_KEY = "V";
7274

75+
/** Default no-op implementation of {@link ExtraDataSupplier}. */
76+
private static final EmptyExtraDataSupplier EMPTY_EXTRA_DATA_SUPPLIER
77+
= new EmptyExtraDataSupplier();
78+
79+
/** DDB fields that must be encrypted. */
80+
private static final Set<String> ENCRYPTED_FIELDS;
81+
static {
82+
final Set<String> tempEncryptedFields = new HashSet<>();
83+
tempEncryptedFields.add(MATERIAL_TYPE_VERSION);
84+
tempEncryptedFields.add(ENCRYPTION_KEY_FIELD);
85+
tempEncryptedFields.add(ENCRYPTION_ALGORITHM_FIELD);
86+
tempEncryptedFields.add(INTEGRITY_KEY_FIELD);
87+
tempEncryptedFields.add(INTEGRITY_ALGORITHM_FIELD);
88+
ENCRYPTED_FIELDS = tempEncryptedFields;
89+
}
90+
7391
private final Map<String, ExpectedAttributeValue> doesNotExist;
92+
private final Set<String> doNotEncrypt;
7493
private final String tableName;
7594
private final AmazonDynamoDB ddb;
7695
private final DynamoDBEncryptor encryptor;
7796
private final EncryptionContext ddbCtx;
97+
private final ExtraDataSupplier extraDataSupplier;
98+
99+
/**
100+
* Provides extra data that should be persisted along with the standard material data.
101+
*/
102+
public interface ExtraDataSupplier {
78103

104+
/**
105+
* Gets the extra data attributes for the specified material name.
106+
*
107+
* @param materialName material name.
108+
* @param version version number.
109+
* @return plain text of the extra data.
110+
*/
111+
Map<String, AttributeValue> getAttributes(final String materialName, final long version);
112+
113+
/**
114+
* Gets the extra data field names that should be signed only but not encrypted.
115+
*
116+
* @return signed only fields.
117+
*/
118+
Set<String> getSignedOnlyFieldNames();
119+
}
120+
121+
/**
122+
* Create a new MetaStore with specified table name.
123+
*
124+
* @param ddb Interface for accessing DynamoDB.
125+
* @param tableName DynamoDB table name for this {@link MetaStore}.
126+
* @param encryptor used to perform crypto operations on the record attributes.
127+
*/
79128
public MetaStore(final AmazonDynamoDB ddb, final String tableName,
80129
final DynamoDBEncryptor encryptor) {
130+
this(ddb, tableName, encryptor, EMPTY_EXTRA_DATA_SUPPLIER);
131+
}
132+
133+
/**
134+
* Create a new MetaStore with specified table name and extra data supplier.
135+
*
136+
* @param ddb Interface for accessing DynamoDB.
137+
* @param tableName DynamoDB table name for this {@link MetaStore}.
138+
* @param encryptor used to perform crypto operations on the record attributes
139+
* @param extraDataSupplier provides extra data that should be stored along with the material.
140+
*/
141+
public MetaStore(final AmazonDynamoDB ddb, final String tableName,
142+
final DynamoDBEncryptor encryptor, final ExtraDataSupplier extraDataSupplier) {
81143
this.ddb = checkNotNull(ddb, "ddb must not be null");
82144
this.tableName = checkNotNull(tableName, "tableName must not be null");
83145
this.encryptor = checkNotNull(encryptor, "encryptor must not be null");
146+
this.extraDataSupplier = checkNotNull(extraDataSupplier, "extraDataSupplier must not be null");
84147

85-
ddbCtx = new EncryptionContext.Builder().withTableName(this.tableName)
148+
this.ddbCtx = new EncryptionContext.Builder().withTableName(this.tableName)
86149
.withHashKeyName(DEFAULT_HASH_KEY).withRangeKeyName(DEFAULT_RANGE_KEY).build();
87150

88151
final Map<String, ExpectedAttributeValue> tmpExpected = new HashMap<String, ExpectedAttributeValue>();
89152
tmpExpected.put(DEFAULT_HASH_KEY, new ExpectedAttributeValue().withExists(false));
90153
tmpExpected.put(DEFAULT_RANGE_KEY, new ExpectedAttributeValue().withExists(false));
91154
doesNotExist = Collections.unmodifiableMap(tmpExpected);
155+
156+
this.doNotEncrypt = getSignedOnlyFields(extraDataSupplier);
92157
}
93158

94159
@Override
@@ -105,10 +170,8 @@ public EncryptionMaterialsProvider getProvider(final String materialName, final
105170

106171
@Override
107172
public EncryptionMaterialsProvider getOrCreate(final String materialName, final long nextId) {
108-
final SecretKeySpec encryptionKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_ENCRYPTION);
109-
final SecretKeySpec integrityKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_INTEGRITY);
110-
final Map<String, AttributeValue> ciphertext = conditionalPut(encryptKeys(materialName,
111-
nextId, encryptionKey, integrityKey));
173+
final Map<String, AttributeValue> plaintext = createMaterialItem(materialName, nextId);
174+
final Map<String, AttributeValue> ciphertext = conditionalPut(getEncryptedText(plaintext));
112175
return decryptProvider(ciphertext);
113176
}
114177

@@ -145,9 +208,10 @@ public long getVersionFromMaterialDescription(final Map<String, String> descript
145208

146209
/**
147210
* This API retrieves the intermediate keys from the source region and replicates it in the target region.
148-
* @param materialName
149-
* @param version
150-
* @param targetMetaStore
211+
*
212+
* @param materialName material name of the encryption material.
213+
* @param version version of the encryption material.
214+
* @param targetMetaStore target MetaStore where the encryption material to be stored.
151215
*/
152216
public void replicate(final String materialName, final long version, final MetaStore targetMetaStore) {
153217
try {
@@ -168,8 +232,14 @@ public void replicate(final String materialName, final long version, final MetaS
168232
//Item already present.
169233
}
170234
}
235+
171236
/**
172237
* Creates a DynamoDB Table with the correct properties to be used with a ProviderStore.
238+
*
239+
* @param ddb interface for accessing DynamoDB
240+
* @param tableName name of table that stores the meta data of the material.
241+
* @param provisionedThroughput required provisioned throughput of the this table.
242+
* @return result of create table request.
173243
*/
174244
public static CreateTableResult createTable(final AmazonDynamoDB ddb, final String tableName,
175245
final ProvisionedThroughput provisionedThroughput) {
@@ -181,6 +251,44 @@ public static CreateTableResult createTable(final AmazonDynamoDB ddb, final Stri
181251

182252
}
183253

254+
/**
255+
* Empty extra data supplier. This default class is intended to simplify the default
256+
* implementation of {@link MetaStore}.
257+
*/
258+
private static class EmptyExtraDataSupplier implements ExtraDataSupplier {
259+
@Override
260+
public Map<String, AttributeValue> getAttributes(String materialName, long version) {
261+
return Collections.emptyMap();
262+
}
263+
264+
@Override
265+
public Set<String> getSignedOnlyFieldNames() {
266+
return Collections.emptySet();
267+
}
268+
}
269+
270+
/**
271+
* Get a set of fields that must be signed but not encrypted.
272+
*
273+
* @param extraDataSupplier extra data supplier that is used to return sign only field names.
274+
* @return fields that must be signed.
275+
*/
276+
private static Set<String> getSignedOnlyFields(final ExtraDataSupplier extraDataSupplier) {
277+
final Set<String> signedOnlyFields = extraDataSupplier.getSignedOnlyFieldNames();
278+
for (final String signedOnlyField : signedOnlyFields) {
279+
if (ENCRYPTED_FIELDS.contains(signedOnlyField)) {
280+
throw new IllegalArgumentException(signedOnlyField + " must be encrypted");
281+
}
282+
}
283+
284+
// fields that should not be encrypted
285+
final Set<String> doNotEncryptFields = new HashSet<>();
286+
doNotEncryptFields.add(DEFAULT_HASH_KEY);
287+
doNotEncryptFields.add(DEFAULT_RANGE_KEY);
288+
doNotEncryptFields.addAll(signedOnlyFields);
289+
return Collections.unmodifiableSet(doNotEncryptFields);
290+
}
291+
184292
private Map<String, AttributeValue> conditionalPut(final Map<String, AttributeValue> item) {
185293
try {
186294
final PutItemRequest put = new PutItemRequest().withTableName(tableName).withItem(item)
@@ -201,19 +309,29 @@ private Map<String, AttributeValue> ddbGet(final Map<String, AttributeValue> ddb
201309
.withKey(ddbKey)).getItem();
202310
}
203311

204-
private Map<String, AttributeValue> encryptKeys(final String name, final long version,
205-
final SecretKeySpec encryptionKey, final SecretKeySpec integrityKey) {
312+
/**
313+
* Build an material item for a given material name and version with newly generated
314+
* encryption and integrity keys.
315+
*
316+
* @param materialName material name.
317+
* @param version version of the material.
318+
* @return newly generated plaintext material item.
319+
*/
320+
private Map<String, AttributeValue> createMaterialItem(final String materialName, final long version) {
321+
final SecretKeySpec encryptionKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_ENCRYPTION);
322+
final SecretKeySpec integrityKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_INTEGRITY);
323+
206324
final Map<String, AttributeValue> plaintext = new HashMap<String, AttributeValue>();
207-
plaintext.put(DEFAULT_HASH_KEY, new AttributeValue().withS(name));
325+
plaintext.put(DEFAULT_HASH_KEY, new AttributeValue().withS(materialName));
208326
plaintext.put(DEFAULT_RANGE_KEY, new AttributeValue().withN(Long.toString(version)));
209327
plaintext.put(MATERIAL_TYPE_VERSION, new AttributeValue().withS("0"));
210-
plaintext.put(ENCRYPTION_KEY_FIELD,
211-
new AttributeValue().withB(ByteBuffer.wrap(encryptionKey.getEncoded())));
328+
plaintext.put(ENCRYPTION_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(encryptionKey.getEncoded())));
212329
plaintext.put(ENCRYPTION_ALGORITHM_FIELD, new AttributeValue().withS(encryptionKey.getAlgorithm()));
213-
plaintext
214-
.put(INTEGRITY_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(integrityKey.getEncoded())));
330+
plaintext.put(INTEGRITY_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(integrityKey.getEncoded())));
215331
plaintext.put(INTEGRITY_ALGORITHM_FIELD, new AttributeValue().withS(integrityKey.getAlgorithm()));
216-
return getEncryptedText(plaintext);
332+
plaintext.putAll(extraDataSupplier.getAttributes(materialName, version));
333+
334+
return plaintext;
217335
}
218336

219337
private EncryptionMaterialsProvider decryptProvider(final Map<String, AttributeValue> item) {
@@ -237,19 +355,31 @@ private EncryptionMaterialsProvider decryptProvider(final Map<String, AttributeV
237355
buildDescription(plaintext));
238356
}
239357

240-
private Map<String, AttributeValue> getPlainText(Map<String, AttributeValue> item) {
358+
/**
359+
* Decrypts attributes in the ciphertext item using {@link DynamoDBEncryptor}.
360+
* except the attribute names specified in doNotEncrypt.
361+
* @param ciphertext the ciphertext to be decrypted.
362+
* @throws AmazonClientException when failed to decrypt material item.
363+
* @return decrypted item.
364+
*/
365+
private Map<String, AttributeValue> getPlainText(final Map<String, AttributeValue> ciphertext) {
241366
try {
242-
return encryptor.decryptAllFieldsExcept(item,
243-
ddbCtx, DEFAULT_HASH_KEY, DEFAULT_RANGE_KEY);
367+
return encryptor.decryptAllFieldsExcept(ciphertext, ddbCtx, doNotEncrypt);
244368
} catch (final GeneralSecurityException e) {
245369
throw new AmazonClientException(e);
246370
}
247371
}
248372

373+
/**
374+
* Encrypts attributes in the plaintext item using {@link DynamoDBEncryptor}.
375+
* except the attribute names specified in doNotEncrypt.
376+
*
377+
* @throws AmazonClientException when failed to encrypt material item.
378+
* @param plaintext plaintext to be encrypted.
379+
*/
249380
private Map<String, AttributeValue> getEncryptedText(Map<String, AttributeValue> plaintext) {
250381
try {
251-
return encryptor.encryptAllFieldsExcept(plaintext, ddbCtx, DEFAULT_HASH_KEY,
252-
DEFAULT_RANGE_KEY);
382+
return encryptor.encryptAllFieldsExcept(plaintext, ddbCtx, doNotEncrypt);
253383
} catch (final GeneralSecurityException e) {
254384
throw new AmazonClientException(e);
255385
}

0 commit comments

Comments
 (0)