added support for custom instruction file suffix in getObject

Author: Anirav Kareddy

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

@@ -6,7 +6,6 @@
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.awscore.exception.AwsServiceException;
-import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
@@ -68,13 +67,11 @@
 import software.amazon.encryption.s3.materials.EncryptionMaterials;
 import software.amazon.encryption.s3.materials.Keyring;
 import software.amazon.encryption.s3.materials.KmsKeyring;
-import software.amazon.encryption.s3.materials.MaterialsDescription;
 import software.amazon.encryption.s3.materials.MultipartConfiguration;
 import software.amazon.encryption.s3.materials.PartialRsaKeyPair;
 import software.amazon.encryption.s3.materials.RawKeyring;
 import software.amazon.encryption.s3.materials.RsaKeyring;
 
-import javax.crypto.DecapsulateException;
 import javax.crypto.SecretKey;
 import java.io.IOException;
 import java.net.URI;
@@ -85,6 +82,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -110,6 +108,7 @@ public class S3EncryptionClient extends DelegatingS3Client {
     // Used for request-scoped encryption contexts for supporting keys
     public static final ExecutionAttribute<Map<String, String>> ENCRYPTION_CONTEXT = new ExecutionAttribute<>("EncryptionContext");
     public static final ExecutionAttribute<MultipartConfiguration> CONFIGURATION = new ExecutionAttribute<>("MultipartConfiguration");
+    public static final ExecutionAttribute<String> CUSTOM_INSTRUCTION_FILE_SUFFIX = new ExecutionAttribute<>("CustomInstructionFileSuffix");
 
     private final S3Client _wrappedClient;
     private final S3AsyncClient _wrappedAsyncClient;
@@ -157,6 +156,11 @@ public static Consumer<AwsRequestOverrideConfiguration.Builder> withAdditionalCo
                 builder.putExecutionAttribute(S3EncryptionClient.ENCRYPTION_CONTEXT, encryptionContext);
     }
 
+    public static Consumer<AwsRequestOverrideConfiguration.Builder> withCustomInstructionFileSuffix(String customInstructionFileSuffix) {
+        return builder ->
+                builder.putExecutionAttribute(S3EncryptionClient.CUSTOM_INSTRUCTION_FILE_SUFFIX, customInstructionFileSuffix);
+    }
+
     /**
      * Attaches multipart configuration to a request. Must be used as a parameter to
      * {@link S3Request#overrideConfiguration()} in the request.
@@ -204,13 +208,15 @@ public ReEncryptInstructionFileResponse reEncryptInstructionFile(ReEncryptInstru
           DecryptMaterialsRequest.builder()
             .algorithmSuite(algorithmSuite)
             .encryptedDataKeys(Collections.singletonList(encryptedDataKey))
+            .s3Request(request)
             .build()
         );
         byte[] plaintextDataKey = decryptedMaterials.plaintextDataKey();
 
         EncryptionMaterials encryptionMaterials = EncryptionMaterials.builder()
           .algorithmSuite(algorithmSuite)
           .plaintextDataKey(plaintextDataKey)
+          .s3Request(request)
           .build();
 
         RawKeyring newKeyring = reEncryptInstructionFileRequest.newKeyring();
@@ -221,11 +227,18 @@ public ReEncryptInstructionFileResponse reEncryptInstructionFile(ReEncryptInstru
         }
 
         ContentMetadataEncodingStrategy encodeStrategy = new ContentMetadataEncodingStrategy(_instructionFileConfig);
-        encodeStrategy.encodeMetadata(encryptedMaterials, iv, PutObjectRequest.builder()
-          .bucket(reEncryptInstructionFileRequest.bucket())
-          .key(reEncryptInstructionFileRequest.key())
-          .build());
 
+        if (reEncryptInstructionFileRequest.instructionFileSuffix().equals(INSTRUCTION_FILE_SUFFIX)) {
+            encodeStrategy.encodeMetadata(encryptedMaterials, iv, PutObjectRequest.builder()
+              .bucket(reEncryptInstructionFileRequest.bucket())
+              .key(reEncryptInstructionFileRequest.key())
+              .build());
+        } else {
+            encodeStrategy.encodeMetadata(encryptedMaterials, iv, PutObjectRequest.builder()
+              .bucket(reEncryptInstructionFileRequest.bucket())
+              .key(reEncryptInstructionFileRequest.key())
+              .build(), reEncryptInstructionFileRequest.instructionFileSuffix());
+        }
         return new ReEncryptInstructionFileResponse(reEncryptInstructionFileRequest.bucket(),
             reEncryptInstructionFileRequest.key(), reEncryptInstructionFileRequest.instructionFileSuffix());
 

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

@@ -9,6 +9,7 @@
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
+import software.amazon.encryption.s3.S3EncryptionClient;
 import software.amazon.encryption.s3.S3EncryptionClientException;
 import software.amazon.encryption.s3.algorithms.AlgorithmSuite;
 import software.amazon.encryption.s3.materials.EncryptedDataKey;
@@ -22,6 +23,7 @@
 import java.util.Base64;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.CompletionException;
 
 import static software.amazon.encryption.s3.S3EncryptionClientUtilities.INSTRUCTION_FILE_SUFFIX;
@@ -224,9 +226,13 @@ private ContentMetadata decodeFromObjectMetadata(GetObjectRequest request, GetOb
     }
 
     private ContentMetadata decodeFromInstructionFile(GetObjectRequest request, GetObjectResponse response) {
+       String instructionFileSuffix = request.overrideConfiguration()
+         .flatMap(config -> config.executionAttributes().getOptionalAttribute(S3EncryptionClient.CUSTOM_INSTRUCTION_FILE_SUFFIX))
+         .orElse(INSTRUCTION_FILE_SUFFIX);
+
         GetObjectRequest instructionGetObjectRequest = GetObjectRequest.builder()
                 .bucket(request.bucket())
-                .key(request.key() + INSTRUCTION_FILE_SUFFIX)
+                .key(request.key() + instructionFileSuffix)
                 .build();
 
         ResponseInputStream<GetObjectResponse> instruction;