Merge pull request #76 from johnwalker/metastoredirectkms

Add ExtraDataSupplier to Metastore, and support for overriding KMS requests in DirectKMSMaterialProvider

Author: johnwalker

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

@@ -126,7 +126,7 @@ public DecryptionMaterials getDecryptionMaterials(EncryptionContext context) {
         DecryptRequest request = appendUserAgent(new DecryptRequest());
         request.setCiphertextBlob(ByteBuffer.wrap(Base64.decode(materialDescription.get(ENVELOPE_KEY))));
         request.setEncryptionContext(ec);
-        final DecryptResult decryptResult = kms.decrypt(request);
+        final DecryptResult decryptResult = decrypt(request, context);
         validateEncryptionKeyId(decryptResult.getKeyId(), context);
 
         final Hkdf kdf;
@@ -167,7 +167,7 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
         req.setNumberOfBytes(256 / 8);
         req.setEncryptionContext(ec);
 
-        final GenerateDataKeyResult dataKeyResult = kms.generateDataKey(req);
+        final GenerateDataKeyResult dataKeyResult = generateDataKey(req, context);
 
         final Map<String, String> materialDescription = new HashMap<>();
         materialDescription.putAll(description);
@@ -192,7 +192,8 @@ public EncryptionMaterials getEncryptionMaterials(EncryptionContext context) {
     }
 
     /**
-     * Get encryption key id.
+     * Get encryption key id that is used to create the {@link EncryptionMaterials}.
+     *
      * @return encryption key id.
      */
     protected String getEncryptionKeyId() {
@@ -202,6 +203,7 @@ protected String getEncryptionKeyId() {
     /**
      * Select encryption key id to be used to generate data key. The default implementation of this method returns
      * {@link DirectKmsMaterialProvider#encryptionKeyId}.
+     *
      * @param context encryption context.
      * @return the encryptionKeyId.
      * @throws DynamoDBMappingException when we fails to select a valid encryption key id.
@@ -211,7 +213,9 @@ protected String selectEncryptionKeyId(EncryptionContext context) throws DynamoD
     }
 
     /**
-     * Validate the encryption key id. The default implementation of this method does nothing.
+     * Validate the encryption key id. The default implementation of this method does not validate
+     * encryption key id.
+     *
      * @param encryptionKeyId encryption key id from {@link DecryptResult}.
      * @param context encryption context.
      * @throws DynamoDBMappingException when encryptionKeyId is invalid.
@@ -221,6 +225,34 @@ protected void validateEncryptionKeyId(String encryptionKeyId, EncryptionContext
         // No action taken.
     }
 
+    /**
+     * Decrypts ciphertext. The default implementation calls KMS to decrypt the ciphertext using the parameters
+     * provided in the {@link DecryptRequest}. Subclass can override the default implementation to provide
+     * additional request parameters using attributes within the {@link EncryptionContext}.
+     *
+     * @param request request parameters to decrypt the given ciphertext.
+     * @param context additional useful data to decrypt the ciphertext.
+     * @return the decrypted plaintext for the given ciphertext.
+     */
+    protected DecryptResult decrypt(final DecryptRequest request, final EncryptionContext context) {
+        return kms.decrypt(request);
+    }
+
+    /**
+     * Returns a data encryption key that you can use in your application to encrypt data locally. The default
+     * implementation calls KMS to generate the data key using the parameters provided in the
+     * {@link GenerateDataKeyRequest}. Subclass can override the default implementation to provide additional
+     * request parameters using attributes within the {@link EncryptionContext}.
+     *
+     * @param request request parameters to generate the data key.
+     * @param context additional useful data to generate the data key.
+     * @return the newly generated data key which includes both the plaintext and ciphertext.
+     */
+    protected GenerateDataKeyResult generateDataKey(final GenerateDataKeyRequest request,
+            final EncryptionContext context) {
+        return kms.generateDataKey(request);
+    }
+
     /**
      * Extracts relevant information from {@code context} and uses it to populate fields in
      * {@code kmsEc}. Currently, these fields are:

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

@@ -17,8 +17,10 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
@@ -70,45 +72,100 @@ public class MetaStore extends ProviderStore {
     private static final String DEFAULT_HASH_KEY = "N";
     private static final String DEFAULT_RANGE_KEY = "V";
 
+    /** Default no-op implementation of {@link ExtraDataSupplier}. */
+    private static final EmptyExtraDataSupplier EMPTY_EXTRA_DATA_SUPPLIER
+            = new EmptyExtraDataSupplier();
+
+    /**  DDB fields that must be encrypted. */
+    private static final Set<String> ENCRYPTED_FIELDS;
+    static {
+        final Set<String> tempEncryptedFields = new HashSet<>();
+        tempEncryptedFields.add(MATERIAL_TYPE_VERSION);
+        tempEncryptedFields.add(ENCRYPTION_KEY_FIELD);
+        tempEncryptedFields.add(ENCRYPTION_ALGORITHM_FIELD);
+        tempEncryptedFields.add(INTEGRITY_KEY_FIELD);
+        tempEncryptedFields.add(INTEGRITY_ALGORITHM_FIELD);
+        ENCRYPTED_FIELDS = tempEncryptedFields;
+    }
+
     private final Map<String, ExpectedAttributeValue> doesNotExist;
+    private final Set<String> doNotEncrypt;
     private final String tableName;
     private final AmazonDynamoDB ddb;
     private final DynamoDBEncryptor encryptor;
     private final EncryptionContext ddbCtx;
+    private final ExtraDataSupplier extraDataSupplier;
+
+    /**
+     * Provides extra data that should be persisted along with the standard material data.
+     */
+    public interface ExtraDataSupplier {
 
+        /**
+         * Gets the extra data attributes for the specified material name.
+         *
+         * @param materialName material name.
+         * @param version version number.
+         * @return plain text of the extra data.
+         */
+        Map<String, AttributeValue> getAttributes(final String materialName, final long version);
+
+        /**
+         * Gets the extra data field names that should be signed only but not encrypted.
+         *
+         * @return signed only fields.
+         */
+        Set<String> getSignedOnlyFieldNames();
+    }
+
+    /**
+     * Create a new MetaStore with specified table name.
+     *
+     * @param ddb Interface for accessing DynamoDB.
+     * @param tableName DynamoDB table name for this {@link MetaStore}.
+     * @param encryptor used to perform crypto operations on the record attributes.
+     */
     public MetaStore(final AmazonDynamoDB ddb, final String tableName,
             final DynamoDBEncryptor encryptor) {
+        this(ddb, tableName, encryptor, EMPTY_EXTRA_DATA_SUPPLIER);
+    }
+
+    /**
+     * Create a new MetaStore with specified table name and extra data supplier.
+     *
+     * @param ddb Interface for accessing DynamoDB.
+     * @param tableName DynamoDB table name for this {@link MetaStore}.
+     * @param encryptor used to perform crypto operations on the record attributes
+     * @param extraDataSupplier provides extra data that should be stored along with the material.
+     */
+    public MetaStore(final AmazonDynamoDB ddb, final String tableName,
+            final DynamoDBEncryptor encryptor, final ExtraDataSupplier extraDataSupplier) {
         this.ddb = checkNotNull(ddb, "ddb must not be null");
         this.tableName = checkNotNull(tableName, "tableName must not be null");
         this.encryptor = checkNotNull(encryptor, "encryptor must not be null");
+        this.extraDataSupplier = checkNotNull(extraDataSupplier, "extraDataSupplier must not be null");
 
-        ddbCtx = new EncryptionContext.Builder().withTableName(this.tableName)
+        this.ddbCtx = new EncryptionContext.Builder().withTableName(this.tableName)
                 .withHashKeyName(DEFAULT_HASH_KEY).withRangeKeyName(DEFAULT_RANGE_KEY).build();
 
-        final Map<String, ExpectedAttributeValue> tmpExpected = new HashMap<String, ExpectedAttributeValue>();
+        final Map<String, ExpectedAttributeValue> tmpExpected = new HashMap<>();
         tmpExpected.put(DEFAULT_HASH_KEY, new ExpectedAttributeValue().withExists(false));
         tmpExpected.put(DEFAULT_RANGE_KEY, new ExpectedAttributeValue().withExists(false));
         doesNotExist = Collections.unmodifiableMap(tmpExpected);
+
+        this.doNotEncrypt = getSignedOnlyFields(extraDataSupplier);
     }
 
     @Override
     public EncryptionMaterialsProvider getProvider(final String materialName, final long version) {
-        final Map<String, AttributeValue> ddbKey = new HashMap<String, AttributeValue>();
-        ddbKey.put(DEFAULT_HASH_KEY, new AttributeValue().withS(materialName));
-        ddbKey.put(DEFAULT_RANGE_KEY, new AttributeValue().withN(Long.toString(version)));
-        final Map<String, AttributeValue> item = ddbGet(ddbKey);
-        if (item == null || item.isEmpty()) {
-            throw new IndexOutOfBoundsException("No material found: " + materialName + "#" + version);
-        }
+        Map<String, AttributeValue> item = getMaterialItem(materialName, version);
         return decryptProvider(item);
     }
 
     @Override
     public EncryptionMaterialsProvider getOrCreate(final String materialName, final long nextId) {
-        final SecretKeySpec encryptionKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_ENCRYPTION);
-        final SecretKeySpec integrityKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_INTEGRITY);
-        final Map<String, AttributeValue> ciphertext = conditionalPut(encryptKeys(materialName,
-                nextId, encryptionKey, integrityKey));
+        final Map<String, AttributeValue> plaintext = createMaterialItem(materialName, nextId);
+        final Map<String, AttributeValue> ciphertext = conditionalPut(getEncryptedText(plaintext));
         return decryptProvider(ciphertext);
     }
 
@@ -145,20 +202,14 @@ public long getVersionFromMaterialDescription(final Map<String, String> descript
 
     /**
      * This API retrieves the intermediate keys from the source region and replicates it in the target region.
-     * @param materialName
-     * @param version
-     * @param targetMetaStore
+     *
+     * @param materialName material name of the encryption material.
+     * @param version version of the encryption material.
+     * @param targetMetaStore target MetaStore where the encryption material to be stored.
      */
     public void replicate(final String materialName, final long version, final MetaStore targetMetaStore) {
         try {
-            final Map<String, AttributeValue> ddbKey = new HashMap<String, AttributeValue>();
-            ddbKey.put(DEFAULT_HASH_KEY, new AttributeValue().withS(materialName));
-            ddbKey.put(DEFAULT_RANGE_KEY, new AttributeValue().withN(Long.toString(version)));
-            final Map<String, AttributeValue> item = ddbGet(ddbKey);
-            if (item == null || item.isEmpty()) {
-                throw new IndexOutOfBoundsException("No material found: " + materialName + "#" + version);
-            }
-
+            Map<String, AttributeValue> item = getMaterialItem(materialName, version);
             final Map<String, AttributeValue> plainText = getPlainText(item);
             final Map<String, AttributeValue> encryptedText = targetMetaStore.getEncryptedText(plainText);
             final PutItemRequest put = new PutItemRequest().withTableName(targetMetaStore.tableName).withItem(encryptedText)
@@ -168,8 +219,25 @@ public void replicate(final String materialName, final long version, final MetaS
             //Item already present.
         }
     }
+
+    private Map<String, AttributeValue> getMaterialItem(final String materialName, final long version) {
+        final Map<String, AttributeValue> ddbKey = new HashMap<>();
+        ddbKey.put(DEFAULT_HASH_KEY, new AttributeValue().withS(materialName));
+        ddbKey.put(DEFAULT_RANGE_KEY, new AttributeValue().withN(Long.toString(version)));
+        final Map<String, AttributeValue> item = ddbGet(ddbKey);
+        if (item == null || item.isEmpty()) {
+            throw new IndexOutOfBoundsException("No material found: " + materialName + "#" + version);
+        }
+        return item;
+    }
+
     /**
      * Creates a DynamoDB Table with the correct properties to be used with a ProviderStore.
+     *
+     * @param ddb interface for accessing DynamoDB
+     * @param tableName name of table that stores the meta data of the material.
+     * @param provisionedThroughput required provisioned throughput of the this table.
+     * @return result of create table request.
      */
     public static CreateTableResult createTable(final AmazonDynamoDB ddb, final String tableName,
             final ProvisionedThroughput provisionedThroughput) {
@@ -181,14 +249,52 @@ public static CreateTableResult createTable(final AmazonDynamoDB ddb, final Stri
 
     }
 
+    /**
+     * Empty extra data supplier. This default class is intended to simplify the default
+     * implementation of {@link MetaStore}.
+     */
+    private static class EmptyExtraDataSupplier implements ExtraDataSupplier {
+        @Override
+        public Map<String, AttributeValue> getAttributes(String materialName, long version) {
+            return Collections.emptyMap();
+        }
+
+        @Override
+        public Set<String> getSignedOnlyFieldNames() {
+            return Collections.emptySet();
+        }
+    }
+
+    /**
+     * Get a set of fields that must be signed but not encrypted.
+     *
+     * @param extraDataSupplier extra data supplier that is used to return sign only field names.
+     * @return fields that must be signed.
+     */
+    private static Set<String> getSignedOnlyFields(final ExtraDataSupplier extraDataSupplier) {
+        final Set<String> signedOnlyFields = extraDataSupplier.getSignedOnlyFieldNames();
+        for (final String signedOnlyField : signedOnlyFields) {
+            if (ENCRYPTED_FIELDS.contains(signedOnlyField)) {
+                throw new IllegalArgumentException(signedOnlyField + " must be encrypted");
+            }
+        }
+
+        // fields that should not be encrypted
+        final Set<String> doNotEncryptFields = new HashSet<>();
+        doNotEncryptFields.add(DEFAULT_HASH_KEY);
+        doNotEncryptFields.add(DEFAULT_RANGE_KEY);
+        doNotEncryptFields.addAll(signedOnlyFields);
+        return Collections.unmodifiableSet(doNotEncryptFields);
+    }
+
     private Map<String, AttributeValue> conditionalPut(final Map<String, AttributeValue> item) {
         try {
             final PutItemRequest put = new PutItemRequest().withTableName(tableName).withItem(item)
                     .withExpected(doesNotExist);
             ddb.putItem(put);
             return item;
         } catch (final ConditionalCheckFailedException ex) {
-            final Map<String, AttributeValue> ddbKey = new HashMap<String, AttributeValue>();
+            final Map<String, AttributeValue> ddbKey = new HashMap<>();
             ddbKey.put(DEFAULT_HASH_KEY, item.get(DEFAULT_HASH_KEY));
             ddbKey.put(DEFAULT_RANGE_KEY, item.get(DEFAULT_RANGE_KEY));
             return ddbGet(ddbKey);
@@ -201,19 +307,29 @@ private Map<String, AttributeValue> ddbGet(final Map<String, AttributeValue> ddb
                 .withKey(ddbKey)).getItem();
     }
 
-    private Map<String, AttributeValue> encryptKeys(final String name, final long version,
-            final SecretKeySpec encryptionKey, final SecretKeySpec integrityKey) {
+    /**
+     * Build an material item for a given material name and version with newly generated
+     * encryption and integrity keys.
+     *
+     * @param materialName material name.
+     * @param version version of the material.
+     * @return newly generated plaintext material item.
+     */
+    private Map<String, AttributeValue> createMaterialItem(final String materialName, final long version) {
+        final SecretKeySpec encryptionKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_ENCRYPTION);
+        final SecretKeySpec integrityKey = new SecretKeySpec(Utils.getRandom(32), DEFAULT_INTEGRITY);
+
         final Map<String, AttributeValue> plaintext = new HashMap<String, AttributeValue>();
-        plaintext.put(DEFAULT_HASH_KEY, new AttributeValue().withS(name));
+        plaintext.put(DEFAULT_HASH_KEY, new AttributeValue().withS(materialName));
         plaintext.put(DEFAULT_RANGE_KEY, new AttributeValue().withN(Long.toString(version)));
         plaintext.put(MATERIAL_TYPE_VERSION, new AttributeValue().withS("0"));
-        plaintext.put(ENCRYPTION_KEY_FIELD,
-                new AttributeValue().withB(ByteBuffer.wrap(encryptionKey.getEncoded())));
+        plaintext.put(ENCRYPTION_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(encryptionKey.getEncoded())));
         plaintext.put(ENCRYPTION_ALGORITHM_FIELD, new AttributeValue().withS(encryptionKey.getAlgorithm()));
-        plaintext
-        .put(INTEGRITY_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(integrityKey.getEncoded())));
+        plaintext.put(INTEGRITY_KEY_FIELD, new AttributeValue().withB(ByteBuffer.wrap(integrityKey.getEncoded())));
         plaintext.put(INTEGRITY_ALGORITHM_FIELD, new AttributeValue().withS(integrityKey.getAlgorithm()));
-        return getEncryptedText(plaintext);
+        plaintext.putAll(extraDataSupplier.getAttributes(materialName, version));
+
+        return plaintext;
     }
 
     private EncryptionMaterialsProvider decryptProvider(final Map<String, AttributeValue> item) {
@@ -237,19 +353,31 @@ private EncryptionMaterialsProvider decryptProvider(final Map<String, AttributeV
                 buildDescription(plaintext));
     }
 
-    private Map<String, AttributeValue> getPlainText(Map<String, AttributeValue> item) {
+    /**
+     * Decrypts attributes in the ciphertext item using {@link DynamoDBEncryptor}.
+     * except the attribute names specified in doNotEncrypt.
+     * @param ciphertext the ciphertext to be decrypted.
+     * @throws AmazonClientException when failed to decrypt material item.
+     * @return decrypted item.
+     */
+    private Map<String, AttributeValue> getPlainText(final Map<String, AttributeValue> ciphertext) {
         try {
-            return encryptor.decryptAllFieldsExcept(item,
-                    ddbCtx, DEFAULT_HASH_KEY, DEFAULT_RANGE_KEY);
+            return encryptor.decryptAllFieldsExcept(ciphertext, ddbCtx, doNotEncrypt);
         } catch (final GeneralSecurityException e) {
             throw new AmazonClientException(e);
         }
     }
 
+    /**
+     * Encrypts attributes in the plaintext item using {@link DynamoDBEncryptor}.
+     * except the attribute names specified in doNotEncrypt.
+     *
+     * @throws AmazonClientException when failed to encrypt material item.
+     * @param plaintext plaintext to be encrypted.
+     */
     private Map<String, AttributeValue> getEncryptedText(Map<String, AttributeValue> plaintext) {
         try {
-            return encryptor.encryptAllFieldsExcept(plaintext, ddbCtx, DEFAULT_HASH_KEY,
-                    DEFAULT_RANGE_KEY);
+            return encryptor.encryptAllFieldsExcept(plaintext, ddbCtx, doNotEncrypt);
         } catch (final GeneralSecurityException e) {
             throw new AmazonClientException(e);
         }