Merge pull request #9 from smswz/request-encrypt-context

Add request-scoped encryption context

Author: smswz

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

@@ -1,10 +1,13 @@
 package software.amazon.encryption.s3;
 
 import java.security.KeyPair;
-import java.security.SecureRandom;
+import java.util.Map;
+import java.util.function.Consumer;
 import javax.crypto.SecretKey;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.core.interceptor.ExecutionAttribute;
 import software.amazon.awssdk.core.sync.RequestBody;
 import software.amazon.awssdk.core.sync.ResponseTransformer;
 import software.amazon.awssdk.services.s3.S3Client;
@@ -16,28 +19,40 @@
 import software.amazon.encryption.s3.internal.PutEncryptedObjectPipeline;
 import software.amazon.encryption.s3.materials.AesKeyring;
 import software.amazon.encryption.s3.materials.CryptographicMaterialsManager;
-import software.amazon.encryption.s3.materials.DataKeyGenerator;
 import software.amazon.encryption.s3.materials.DefaultCryptoMaterialsManager;
-import software.amazon.encryption.s3.materials.DefaultDataKeyGenerator;
 import software.amazon.encryption.s3.materials.Keyring;
 import software.amazon.encryption.s3.materials.KmsKeyring;
 import software.amazon.encryption.s3.materials.RsaKeyring;
 
+/**
+ * This client is a drop-in replacement for the S3 client. It will automatically encrypt objects
+ * on putObject and decrypt objects on getObject using the provided encryption key(s).
+ */
 public class S3EncryptionClient implements S3Client {
 
+    // Used for request-scoped encryption contexts for supporting keys
+    public static final ExecutionAttribute<Map<String,String>> ENCRYPTION_CONTEXT = new ExecutionAttribute<>("EncryptionContext");
+
     private final S3Client _wrappedClient;
     private final CryptographicMaterialsManager _cryptoMaterialsManager;
+    private final boolean _enableLegacyModes;
 
     private S3EncryptionClient(Builder builder) {
         _wrappedClient = builder._wrappedClient;
         _cryptoMaterialsManager = builder._cryptoMaterialsManager;
-        // TODO: store _enableLegacyModes and pass onto pipeline
+        _enableLegacyModes = builder._enableLegacyModes;
     }
 
     public static Builder builder() {
         return new Builder();
     }
 
+    // Helper function to attach encryption contexts to a request
+    public static Consumer<AwsRequestOverrideConfiguration.Builder> withAdditionalEncryptionContext(Map<String, String> encryptionContext) {
+        return builder ->
+                builder.putExecutionAttribute(S3EncryptionClient.ENCRYPTION_CONTEXT, encryptionContext);
+    }
+
     @Override
     public PutObjectResponse putObject(PutObjectRequest putObjectRequest, RequestBody requestBody)
             throws AwsServiceException, SdkClientException {
@@ -58,6 +73,7 @@ public <T> T getObject(GetObjectRequest getObjectRequest,
         GetEncryptedObjectPipeline pipeline = GetEncryptedObjectPipeline.builder()
                 .s3Client(_wrappedClient)
                 .cryptoMaterialsManager(_cryptoMaterialsManager)
+                .enableLegacyModes(_enableLegacyModes)
                 .build();
 
         return pipeline.getObject(getObjectRequest, responseTransformer);

src/main/java/software/amazon/encryption/s3/internal/ContentMetadata.java

@@ -7,18 +7,6 @@
 
 public class ContentMetadata {
 
-    /*
-    public static final String ENCRYPTED_DATA_KEY = "x-amz-key-v2";
-    // This is the name of the keyring/algorithm e.g. AES/GCM or kms+context
-    public static final String ENCRYPTED_DATA_KEY_ALGORITHM = "x-amz-wrap-alg";
-    public static final String ENCRYPTED_DATA_KEY_CONTEXT = "x-amz-matdesc";
-
-    public static final String CONTENT_NONCE = "x-amz-iv";
-    // This is usually an actual Java cipher e.g. AES/GCM/NoPadding
-    public static final String CONTENT_CIPHER = "x-amz-cek-alg";
-    public static final String CONTENT_CIPHER_TAG_LENGTH = "x-amz-tag-len";
-     */
-
     private final AlgorithmSuite _algorithmSuite;
 
     private final EncryptedDataKey _encryptedDataKey;

src/main/java/software/amazon/encryption/s3/internal/CryptoMaterialsManagerKeyStrategy.java

@@ -20,16 +20,24 @@
 import software.amazon.encryption.s3.materials.DecryptionMaterials;
 import software.amazon.encryption.s3.materials.EncryptedDataKey;
 
+/**
+ * This class will determine the necessary mechanisms to decrypt objects returned from S3.
+ * Due to supporting various legacy modes, this is not a predefined pipeline like
+ * PutEncryptedObjectPipeline. There are several branches in this graph that are determined as more
+ * information is available from the returned object.
+ */
 public class GetEncryptedObjectPipeline {
 
-    final private S3Client _s3Client;
-    final private CryptographicMaterialsManager _cryptoMaterialsManager;
+    private final S3Client _s3Client;
+    private final CryptographicMaterialsManager _cryptoMaterialsManager;
+    private final boolean _enableLegacyModes;
 
     public static Builder builder() { return new Builder(); }
 
     private GetEncryptedObjectPipeline(Builder builder) {
         this._s3Client = builder._s3Client;
         this._cryptoMaterialsManager = builder._cryptoMaterialsManager;
+        this._enableLegacyModes = builder._enableLegacyModes;
     }
 
     public <T> T getObject(GetObjectRequest getObjectRequest,
@@ -47,15 +55,20 @@ public <T> T getObject(GetObjectRequest getObjectRequest,
         ContentMetadata contentMetadata = ContentMetadataStrategy.decode(_s3Client, getObjectRequest, getObjectResponse);
 
         AlgorithmSuite algorithmSuite = contentMetadata.algorithmSuite();
+        if (!_enableLegacyModes && algorithmSuite.isLegacy()) {
+            throw new S3EncryptionClientException("Enable legacy modes to use legacy content encryption: " + algorithmSuite.cipherName());
+        }
+
         List<EncryptedDataKey> encryptedDataKeys = Collections.singletonList(contentMetadata.encryptedDataKey());
 
-        DecryptMaterialsRequest request = DecryptMaterialsRequest.builder()
+        DecryptMaterialsRequest materialsRequest = DecryptMaterialsRequest.builder()
+                .s3Request(getObjectRequest)
                 .algorithmSuite(algorithmSuite)
                 .encryptedDataKeys(encryptedDataKeys)
                 .encryptionContext(contentMetadata.encryptedDataKeyContext())
                 .build();
 
-        DecryptionMaterials materials = _cryptoMaterialsManager.decryptMaterials(request);
+        DecryptionMaterials materials = _cryptoMaterialsManager.decryptMaterials(materialsRequest);
 
         ContentDecryptionStrategy contentDecryptionStrategy = null;
         switch (algorithmSuite) {
@@ -79,6 +92,7 @@ public <T> T getObject(GetObjectRequest getObjectRequest,
     public static class Builder {
         private S3Client _s3Client;
         private CryptographicMaterialsManager _cryptoMaterialsManager;
+        private boolean _enableLegacyModes;
 
         private Builder() {}
 
@@ -92,6 +106,11 @@ public Builder cryptoMaterialsManager(CryptographicMaterialsManager cryptoMateri
             return this;
         }
 
+        public Builder enableLegacyModes(boolean enableLegacyModes) {
+            this._enableLegacyModes = enableLegacyModes;
+            return this;
+        }
+
         public GetEncryptedObjectPipeline build() {
             return new GetEncryptedObjectPipeline(this);
         }

src/main/java/software/amazon/encryption/s3/internal/GetEncryptedObjectPipeline.java

@@ -29,12 +29,14 @@ private PutEncryptedObjectPipeline(Builder builder) {
     }
 
     public PutObjectResponse putObject(PutObjectRequest request, RequestBody requestBody) {
-        EncryptionMaterials materials = _cryptoMaterialsManager.getEncryptionMaterials(
-                EncryptionMaterialsRequest.builder()
-                        .build());
+        EncryptionMaterialsRequest.Builder requestBuilder = EncryptionMaterialsRequest.builder()
+                .s3Request(request);
+
+        EncryptionMaterials materials = _cryptoMaterialsManager.getEncryptionMaterials(requestBuilder.build());
 
         byte[] input;
         try {
+            // TODO: this needs to be a stream and not a byte[]
             input = IoUtils.toByteArray(requestBody.contentStreamProvider().newStream());
         } catch (IOException e) {
             throw new S3EncryptionClientException("Cannot read input.", e);
@@ -49,6 +51,7 @@ public PutObjectResponse putObject(PutObjectRequest request, RequestBody request
     public static class Builder {
         private S3Client _s3Client;
         private CryptographicMaterialsManager _cryptoMaterialsManager;
+        // Default to AesGcm since it is the only active (non-legacy) content encryption strategy
         private ContentEncryptionStrategy _contentEncryptionStrategy =
                 AesGcmContentStrategy
                         .builder()

src/main/java/software/amazon/encryption/s3/internal/KeyUnwrapStrategy.java

@@ -3,7 +3,6 @@
 import java.security.InvalidAlgorithmParameterException;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
-import java.security.SecureRandom;
 import javax.crypto.BadPaddingException;
 import javax.crypto.Cipher;
 import javax.crypto.IllegalBlockSizeException;

src/main/java/software/amazon/encryption/s3/internal/PutEncryptedObjectPipeline.java

@@ -38,11 +38,11 @@ public String keyProviderInfo() {
         }
 
         @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, EncryptedDataKey encryptedDataKey) throws GeneralSecurityException {
+        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
             Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
             cipher.init(Cipher.DECRYPT_MODE, _wrappingKey);
 
-            return cipher.doFinal(encryptedDataKey.ciphertext());
+            return cipher.doFinal(encryptedDataKey);
         }
     };
 
@@ -62,11 +62,11 @@ public String keyProviderInfo() {
         }
 
         @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, EncryptedDataKey encryptedDataKey) throws GeneralSecurityException {
+        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
             final Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
             cipher.init(Cipher.UNWRAP_MODE, _wrappingKey);
 
-            Key plaintextKey = cipher.unwrap(encryptedDataKey.ciphertext(), CIPHER_ALGORITHM, Cipher.SECRET_KEY);
+            Key plaintextKey = cipher.unwrap(encryptedDataKey, CIPHER_ALGORITHM, Cipher.SECRET_KEY);
             return plaintextKey.getEncoded();
         }
     };
@@ -113,16 +113,16 @@ public byte[] encryptDataKey(SecureRandom secureRandom,
         }
 
         @Override
-        public byte[] decryptDataKey(DecryptionMaterials materials, EncryptedDataKey encryptedDataKey) throws GeneralSecurityException {
-            byte[] encodedBytes = encryptedDataKey.ciphertext();
+        public byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey) throws GeneralSecurityException {
+            byte[] encodedBytes = encryptedDataKey;
             byte[] nonce = new byte[NONCE_LENGTH_BYTES];
             byte[] ciphertext = new byte[encodedBytes.length - nonce.length];
 
             System.arraycopy(encodedBytes, 0, nonce, 0, nonce.length);
             System.arraycopy(encodedBytes, nonce.length, ciphertext, 0, ciphertext.length);
 
             GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_BITS, nonce);
-            final Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
+            final Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);
             cipher.init(Cipher.DECRYPT_MODE, _wrappingKey, gcmParameterSpec);
 
             AlgorithmSuite algorithmSuite = materials.algorithmSuite();

src/main/java/software/amazon/encryption/s3/legacy/internal/AesCbcContentStrategy.java

@@ -7,6 +7,6 @@ public interface DecryptDataKeyStrategy {
 
     String keyProviderInfo();
 
-    byte[] decryptDataKey(DecryptionMaterials materials, EncryptedDataKey encryptedDataKey)
+    byte[] decryptDataKey(DecryptionMaterials materials, byte[] encryptedDataKey)
             throws GeneralSecurityException;
 }

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

@@ -3,15 +3,18 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
 
 public class DecryptMaterialsRequest {
 
+    private final GetObjectRequest _s3Request;
     private final AlgorithmSuite _algorithmSuite;
     private final List<EncryptedDataKey> _encryptedDataKeys;
     private final Map<String, String> _encryptionContext;
 
     private DecryptMaterialsRequest(Builder builder) {
+        this._s3Request = builder._s3Request;
         this._algorithmSuite = builder._algorithmSuite;
         this._encryptedDataKeys = builder._encryptedDataKeys;
         this._encryptionContext = builder._encryptionContext;
@@ -21,6 +24,10 @@ static public Builder builder() {
         return new Builder();
     }
 
+    public GetObjectRequest s3Request() {
+        return _s3Request;
+    }
+
     public AlgorithmSuite algorithmSuite() {
         return _algorithmSuite;
     }
@@ -35,13 +42,19 @@ public Map<String, String> encryptionContext() {
 
     static public class Builder {
 
+        public GetObjectRequest _s3Request = null;
         private AlgorithmSuite _algorithmSuite = AlgorithmSuite.ALG_AES_256_GCM_IV12_TAG16_NO_KDF;
         private Map<String, String> _encryptionContext = Collections.emptyMap();
         private List<EncryptedDataKey> _encryptedDataKeys = Collections.emptyList();
 
         private Builder() {
         }
 
+        public Builder s3Request(GetObjectRequest s3Request) {
+            _s3Request = s3Request;
+            return this;
+        }
+
         public Builder algorithmSuite(AlgorithmSuite algorithmSuite) {
             _algorithmSuite = algorithmSuite;
             return this;