Add some contextual comments

Author: smswz

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

@@ -24,8 +24,13 @@
 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;
@@ -42,6 +47,7 @@ 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);

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

@@ -20,6 +20,12 @@
 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 {
 
     private final S3Client _s3Client;

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

@@ -36,6 +36,7 @@ public PutObjectResponse putObject(PutObjectRequest request, RequestBody request
 
         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);
@@ -50,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/materials/KmsKeyring.java

@@ -4,7 +4,6 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
-import java.util.TreeMap;
 import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.ApiName;
 import software.amazon.awssdk.core.SdkBytes;
@@ -100,8 +99,7 @@ public EncryptionMaterials modifyMaterials(EncryptionMaterials materials) {
 
         @Override
         public byte[] encryptDataKey(SecureRandom secureRandom, EncryptionMaterials materials) {
-            // Convert to TreeMap for sorting of keys
-            TreeMap<String, String> encryptionContext = new TreeMap<>(materials.encryptionContext());
+            HashMap<String, String> encryptionContext = new HashMap<>(materials.encryptionContext());
             EncryptRequest request = EncryptRequest.builder()
                     .keyId(_wrappingKeyId)
                     .encryptionContext(encryptionContext)

src/test/java/S3EncryptionClientTest.java

@@ -536,11 +536,7 @@ public void KmsContextV3toV2() throws IOException {
                         .overrideConfiguration(withAdditionalEncryptionContext(encryptionContext)),
                 RequestBody.fromString(input));
 
-        EncryptedGetObjectRequest getObjectRequest = new EncryptedGetObjectRequest(
-                BUCKET,
-                BUCKET_KEY
-        ).withExtraMaterialsDescription(encryptionContext);
-        String output = IOUtils.toString(v2Client.getObject(getObjectRequest).getObjectContent());
+        String output = v2Client.getObjectAsString(BUCKET, BUCKET_KEY);
         assertEquals(input, output);
     }