diff --git a/src/main/java/software/amazon/encryption/s3/internal/InstructionFileConfig.java b/src/main/java/software/amazon/encryption/s3/internal/InstructionFileConfig.java index e5b4ab015..42de593a7 100644 --- a/src/main/java/software/amazon/encryption/s3/internal/InstructionFileConfig.java +++ b/src/main/java/software/amazon/encryption/s3/internal/InstructionFileConfig.java @@ -65,10 +65,12 @@ PutObjectResponse putInstructionFile(PutObjectRequest request, String instructio instFileMetadata.put(INSTRUCTION_FILE, ""); // Use toBuilder to keep all other fields the same as the actual request + // but set the content length, key, and metadata appropriately for the instruction file final PutObjectRequest instPutRequest = request.toBuilder() - .key(request.key() + instructionFileSuffix) - .metadata(instFileMetadata) - .build(); + .key(request.key() + instructionFileSuffix) + .contentLength((long) instructionFileContent.getBytes().length) + .metadata(instFileMetadata) + .build(); switch (_clientType) { case SYNCHRONOUS: return _s3Client.putObject(instPutRequest, RequestBody.fromString(instructionFileContent)); diff --git a/src/test/java/software/amazon/encryption/s3/internal/InstructionFileConfigUploadTest.java b/src/test/java/software/amazon/encryption/s3/internal/InstructionFileConfigUploadTest.java new file mode 100644 index 000000000..41cc5f287 --- /dev/null +++ b/src/test/java/software/amazon/encryption/s3/internal/InstructionFileConfigUploadTest.java @@ -0,0 +1,73 @@ +package software.amazon.encryption.s3.internal; + +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import software.amazon.awssdk.core.async.AsyncRequestBody; +import software.amazon.awssdk.core.sync.RequestBody; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; + +import java.util.concurrent.CompletableFuture; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.*; + +class InstructionFileConfigUploadTest { + + @Test + void uploadInstructionFileWithSetContentLengthSyncClient() { + // Create a mock for the S3 client + S3Client mockedS3Client = mock(S3Client.class); + // The argument captor is used to capture the PutObjectRequest passed to the putObject method + ArgumentCaptor instructionFilePutCaptor = ArgumentCaptor.forClass(PutObjectRequest.class); + + // Create the InstructionFileConfig with the mocked S3 client + InstructionFileConfig instructionFileConfig = InstructionFileConfig.builder() + .instructionFileClient(mockedS3Client) + .enableInstructionFilePutObject(true) + .build(); + + // Build some data for the test + PutObjectRequest putObjectRequest = PutObjectRequest.builder() + .key("someKey").build(); + String instructionFileContent = "some content that fakes an instruction file"; + + // call the actual method under test + instructionFileConfig.putInstructionFile(putObjectRequest, instructionFileContent); + + // Verify that the putObject method was called and the captured request has the correct content length + verify(mockedS3Client).putObject(instructionFilePutCaptor.capture(), any(RequestBody.class)); + assertEquals(instructionFileContent.getBytes().length, instructionFilePutCaptor.getValue().contentLength()); + } + + @Test + void uploadInstructionFileWithSetContentLengthAsyncClient() { + // Create a mock for the S3 client + S3AsyncClient mockedS3Client = mock(S3AsyncClient.class); + // The async putObject method returns a CompletableFuture, so we need to mock that behavior + when(mockedS3Client.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))) + .thenReturn(CompletableFuture.completedFuture(null)); + // The argument captor is used to capture the PutObjectRequest passed to the putObject method + ArgumentCaptor instructionFilePutCaptor = ArgumentCaptor.forClass(PutObjectRequest.class); + + // Create the InstructionFileConfig with the mocked S3 async client + InstructionFileConfig instructionFileConfig = InstructionFileConfig.builder() + .instructionFileAsyncClient(mockedS3Client) + .enableInstructionFilePutObject(true) + .build(); + + // Build some data for the test + PutObjectRequest putObjectRequest = PutObjectRequest.builder() + .key("someKey").build(); + String instructionFileContent = "some content that fakes an instruction file"; + + // call the actual method under test + instructionFileConfig.putInstructionFile(putObjectRequest, instructionFileContent); + + // Verify that the putObject method was called and the captured request has the correct content length + verify(mockedS3Client).putObject(instructionFilePutCaptor.capture(), any(AsyncRequestBody.class)); + assertEquals(instructionFileContent.getBytes().length, instructionFilePutCaptor.getValue().contentLength()); + } +} \ No newline at end of file