fix: handle contentLength in request object (#106)

Author: kessplas

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

@@ -5,6 +5,7 @@
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.encryption.s3.S3EncryptionClientException;
 import software.amazon.encryption.s3.materials.CryptographicMaterialsManager;
 import software.amazon.encryption.s3.materials.EncryptionMaterials;
 import software.amazon.encryption.s3.materials.EncryptionMaterialsRequest;
@@ -33,17 +34,38 @@ private PutEncryptedObjectPipeline(Builder builder) {
     }
 
     public CompletableFuture<PutObjectResponse> putObject(PutObjectRequest request, AsyncRequestBody requestBody) {
-        EncryptionMaterialsRequest.Builder requestBuilder = EncryptionMaterialsRequest.builder()
+        final Long contentLength;
+        if (request.contentLength() != null) {
+            if (requestBody.contentLength().isPresent() && !request.contentLength().equals(requestBody.contentLength().get())) {
+                // if the contentLength values do not match, throw an exception, since we don't know which is correct
+                throw new S3EncryptionClientException("The contentLength provided in the request object MUST match the " +
+                        "contentLength in the request body");
+            } else if (!requestBody.contentLength().isPresent()) {
+                // no contentLength in request body, use the one in request
+                contentLength = request.contentLength();
+            } else {
+                // only remaining case is when the values match, so either works here
+                contentLength = request.contentLength();
+            }
+        } else {
+            contentLength = requestBody.contentLength().orElse(-1L);
+        }
+
+        EncryptionMaterialsRequest encryptionMaterialsRequest = EncryptionMaterialsRequest.builder()
                 .s3Request(request)
-                .plaintextLength(requestBody.contentLength().orElse(-1L));
+                .plaintextLength(contentLength)
+                .build();
 
-        EncryptionMaterials materials = _cryptoMaterialsManager.getEncryptionMaterials(requestBuilder.build());
+        EncryptionMaterials materials = _cryptoMaterialsManager.getEncryptionMaterials(encryptionMaterialsRequest);
 
         EncryptedContent encryptedContent = _asyncContentEncryptionStrategy.encryptContent(materials, requestBody);
 
         Map<String, String> metadata = new HashMap<>(request.metadata());
         metadata = _contentMetadataEncodingStrategy.encodeMetadata(materials, encryptedContent.getNonce(), metadata);
-        PutObjectRequest encryptedPutRequest = request.toBuilder().metadata(metadata).build();
+        PutObjectRequest encryptedPutRequest = request.toBuilder()
+                .contentLength(encryptedContent.getCiphertextLength())
+                .metadata(metadata)
+                .build();
         return _s3AsyncClient.putObject(encryptedPutRequest, encryptedContent.getAsyncCiphertext());
     }
 

src/test/java/software/amazon/encryption/s3/S3EncryptionClientTest.java

@@ -511,29 +511,6 @@ public void s3EncryptionClientFromAESKeyringUsesDifferentSecureRandomThanKeyring
         v3Client.close();
     }
 
-    /**
-     * A simple, reusable round-trip (encryption + decryption) using a given
-     * S3Client. Useful for testing client configuration.
-     *
-     * @param v3Client the client under test
-     */
-    private void simpleV3RoundTrip(final S3Client v3Client, final String objectKey) {
-        final String input = "SimpleTestOfV3EncryptionClient";
-
-        v3Client.putObject(builder -> builder
-                        .bucket(BUCKET)
-                        .key(objectKey)
-                        .build(),
-                RequestBody.fromString(input));
-
-        ResponseBytes<GetObjectResponse> objectResponse = v3Client.getObjectAsBytes(builder -> builder
-                .bucket(BUCKET)
-                .key(objectKey)
-                .build());
-        String output = objectResponse.asUtf8String();
-        assertEquals(input, output);
-    }
-
     @Test
     public void cryptoProviderV3toV3Enabled() {
         final String objectKey = appendTestSuffix("crypto-provider-enabled-v3-to-v3");
@@ -599,4 +576,58 @@ public void cryptoProviderV2toV3Enabled() {
         deleteObject(BUCKET, objectKey, v3Client);
         v3Client.close();
     }
+
+    @Test
+    public void contentLengthRequest() {
+        final String objectKey = appendTestSuffix("content-length");
+
+        S3Client s3EncryptionClient = S3EncryptionClient.builder()
+                .kmsKeyId(KMS_KEY_ID)
+                .build();
+
+        final String input = "SimpleTestOfV3EncryptionClientCopyObject";
+        final int inputLength = input.length();
+
+        s3EncryptionClient.putObject(builder -> builder
+                        .bucket(BUCKET)
+                        .key(objectKey)
+                        .contentLength((long) inputLength)
+                        .build(),
+                RequestBody.fromString(input));
+
+        ResponseBytes<GetObjectResponse> objectResponse = s3EncryptionClient.getObjectAsBytes(builder -> builder
+                .bucket(BUCKET)
+                .key(objectKey)
+                .build());
+        String output = objectResponse.asUtf8String();
+        assertEquals(input, output);
+
+        // Cleanup
+        deleteObject(BUCKET, objectKey, s3EncryptionClient);
+        s3EncryptionClient.close();
+    }
+
+    /**
+     * A simple, reusable round-trip (encryption + decryption) using a given
+     * S3Client. Useful for testing client configuration.
+     *
+     * @param v3Client the client under test
+     */
+    private void simpleV3RoundTrip(final S3Client v3Client, final String objectKey) {
+        final String input = "SimpleTestOfV3EncryptionClient";
+
+        v3Client.putObject(builder -> builder
+                        .bucket(BUCKET)
+                        .key(objectKey)
+                        .build(),
+                RequestBody.fromString(input));
+
+        ResponseBytes<GetObjectResponse> objectResponse = v3Client.getObjectAsBytes(builder -> builder
+                .bucket(BUCKET)
+                .key(objectKey)
+                .build());
+        String output = objectResponse.asUtf8String();
+        assertEquals(input, output);
+    }
+
 }