Generate data keys from keyring.

Author: smswz

src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java

@@ -62,6 +62,8 @@ public S3EncryptionClient(S3Client client, MaterialsManager materialsManager) {
     public PutObjectResponse putObject(PutObjectRequest putObjectRequest, RequestBody requestBody)
             throws AwsServiceException, SdkClientException, S3Exception {
 
+        // TODO: This is proof-of-concept code and needs to be refactored
+
         // Get content encryption key
         EncryptionMaterials materials = _materialsManager.getEncryptionMaterials(EncryptionMaterialsRequest.builder()
                 .build());
@@ -129,6 +131,9 @@ public PutObjectResponse putObject(PutObjectRequest putObjectRequest, RequestBod
     @Override
     public <T> T getObject(GetObjectRequest getObjectRequest, ResponseTransformer<GetObjectResponse, T> responseTransformer)
             throws NoSuchKeyException, InvalidObjectStateException, AwsServiceException, SdkClientException, S3Exception {
+
+        // TODO: This is proof-of-concept code and needs to be refactored
+
         ResponseInputStream<GetObjectResponse> objectStream =  _wrappedClient.getObject(getObjectRequest);
         byte[] output;
         try {

src/main/java/software/amazon/encryption/s3/algorithms/AlgorithmSuite.java

@@ -39,6 +39,14 @@ public enum AlgorithmSuite {
         this._cipherMaxContentLengthBits = cipherMaxContentLengthBits;
     }
 
+    public String dataKeyAlgorithm() {
+        return _dataKeyAlgorithm;
+    }
+
+    public int dataKeyLengthBits() {
+        return _dataKeyLengthBits;
+    }
+
     public String cipherName() {
         return _cipherName;
     }

src/main/java/software/amazon/encryption/s3/materials/AESKeyring.java

@@ -20,20 +20,26 @@ public class AESKeyring implements Keyring {
     private static final int NONCE_LENGTH_BYTES = 12;
     private static final int TAG_LENGTH_BYTES = 16;
     private static final int TAG_LENGTH_BITS = TAG_LENGTH_BYTES * 8;
+
+    private final DataKeyGenerator _dataKeyGenerator;
     private final SecretKey _wrappingKey;
 
-    public AESKeyring(SecretKey wrappingKey) {
-        if (!wrappingKey.getAlgorithm().equals(KEY_ALGORITHM)) {
-            throw new IllegalArgumentException("Invalid algorithm '" + wrappingKey.getAlgorithm() + "', expecting " + KEY_ALGORITHM);
-        }
+    private AESKeyring(Builder builder) {
+        _dataKeyGenerator = builder._dataKeyGenerator;
+        _wrappingKey = builder._wrappingKey;
+    }
 
-        _wrappingKey = wrappingKey;
+    public static Builder builder() {
+        return new Builder();
     }
 
     @Override
     public EncryptionMaterials onEncrypt(EncryptionMaterials materials) {
         if (materials.plaintextDataKey() == null) {
-            throw new IllegalStateException("Missing data key to wrap");
+            SecretKey dataKey = _dataKeyGenerator.generateDataKey(materials.algorithmSuite());
+            materials = materials.toBuilder()
+                    .plaintextDataKey(dataKey.getEncoded())
+                    .build();
         }
 
         try {
@@ -108,4 +114,28 @@ public DecryptionMaterials onDecrypt(final DecryptionMaterials materials, List<E
 
         return materials;
     }
+
+    public static class Builder {
+        private DataKeyGenerator _dataKeyGenerator = new DefaultDataKeyGenerator();
+        private SecretKey _wrappingKey;
+
+        private Builder() {}
+
+        public Builder wrappingKey(SecretKey wrappingKey) {
+            if (!wrappingKey.getAlgorithm().equals(KEY_ALGORITHM)) {
+                throw new IllegalArgumentException("Invalid algorithm '" + wrappingKey.getAlgorithm() + "', expecting " + KEY_ALGORITHM);
+            }
+            _wrappingKey = wrappingKey;
+            return this;
+        }
+
+        public Builder dataKeyGenerator(DataKeyGenerator dataKeyGenerator) {
+            _dataKeyGenerator = dataKeyGenerator;
+            return this;
+        }
+
+        public AESKeyring build() {
+            return new AESKeyring(this);
+        }
+    }
 }

src/main/java/software/amazon/encryption/s3/materials/DataKeyGenerator.java

@@ -0,0 +1,8 @@
+package software.amazon.encryption.s3.materials;
+
+import javax.crypto.SecretKey;
+import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
+
+public interface DataKeyGenerator {
+    SecretKey generateDataKey(AlgorithmSuite algorithmSuite);
+}

src/main/java/software/amazon/encryption/s3/materials/DefaultDataKeyGenerator.java

@@ -0,0 +1,21 @@
+package software.amazon.encryption.s3.materials;
+
+import java.security.NoSuchAlgorithmException;
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
+
+public class DefaultDataKeyGenerator implements DataKeyGenerator {
+
+    public SecretKey generateDataKey(AlgorithmSuite algorithmSuite) {
+        KeyGenerator generator;
+        try {
+            generator = KeyGenerator.getInstance(algorithmSuite.dataKeyAlgorithm());
+        } catch (NoSuchAlgorithmException e) {
+            throw new UnsupportedOperationException("Unable to generate a(n) " + algorithmSuite.dataKeyAlgorithm() + " data key", e);
+        }
+
+        generator.init(algorithmSuite.dataKeyLengthBits());
+        return generator.generateKey();
+    }
+}

src/main/java/software/amazon/encryption/s3/materials/DefaultMaterialsManager.java

@@ -7,8 +7,6 @@
 import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
 
 public class DefaultMaterialsManager implements MaterialsManager {
-    // TODO: allow this to be configurable?
-    private final SecureRandom _secureRandom = new SecureRandom();
     private final Keyring _keyring;
 
 
@@ -17,29 +15,14 @@ public DefaultMaterialsManager(Keyring keyring) {
     }
 
     public EncryptionMaterials getEncryptionMaterials(EncryptionMaterialsRequest request) {
-        SecretKey key = generateDataKey();
-
         EncryptionMaterials materials = EncryptionMaterials.builder()
                 .algorithmSuite(AlgorithmSuite.ALG_AES_256_GCM_NO_KDF)
                 .encryptionContext(request.encryptionContext())
-                .plaintextDataKey(key.getEncoded())
                 .build();
 
         return _keyring.onEncrypt(materials);
     }
 
-    private SecretKey generateDataKey() {
-        KeyGenerator generator;
-        try {
-            generator = KeyGenerator.getInstance("AES");
-        } catch (NoSuchAlgorithmException e) {
-            throw new UnsupportedOperationException("Unable to generate an AES key", e);
-        }
-
-        generator.init(256 , _secureRandom);
-        return generator.generateKey();
-    }
-
     public DecryptionMaterials getDecryptionMaterials(DecryptionMaterialsRequest request) {
         DecryptionMaterials materials = DecryptionMaterials.builder()
                 .algorithmSuite(request.algorithmSuite())