Merge pull request aws#3747 from aws/zoewang/OutputStreamFix

zoewangg · web-flow · commit bd1a49c4010c · 2025-03-17T11:27:32.000-07:00
Copy byte array in OutputStreamPublisher
diff --git a/.changes/2.31.2.json b/.changes/2.31.2.json
@@ -37,6 +37,12 @@
             "category": "AWS SDK for Java v2",
             "contributor": "",
             "description": "Updated endpoint and partition metadata."
+        },
+        {
+            "type": "bugfix",
+            "category": "AWS SDK for Java v2",
+            "contributor": "",
+            "description": "Copy bytes written to OutputStream of BlockingOutputStreamAsyncRequestBody"
         }
     ]
 }
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -3,6 +3,8 @@
 ## __AWS SDK for Java v2__
   - ### Features
     - Updated endpoint and partition metadata.
+  - ### Bugfixes
+    - Copy bytes written to OutputStream of BlockingOutputStreamAsyncRequestBody.
 
 ## __AWS WAFV2__
   - ### Features
diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/BlockingOutputStreamAsyncRequestBodyTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/async/BlockingOutputStreamAsyncRequestBodyTest.java
@@ -24,15 +24,20 @@
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Random;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import org.apache.commons.lang.math.RandomUtils;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 import software.amazon.awssdk.utils.CancellableOutputStream;
 import software.amazon.awssdk.utils.async.ByteBufferStoringSubscriber;
 import software.amazon.awssdk.utils.async.StoringSubscriber;
 
 class BlockingOutputStreamAsyncRequestBodyTest {
+    private Random random = new Random(3470);
     @Test
     public void outputStream_waitsForSubscription() throws IOException {
         ScheduledExecutorService executor = Executors.newScheduledThreadPool(1);
@@ -77,4 +82,55 @@ public void outputStream_writesToSubscriber() throws IOException {
         assertThat(out.get()).isEqualTo((byte) 1);
     }
 
+    @Test
+    public void outputStream_writesArrayWithOffsetUseSameArray_shouldNotOverride() throws IOException {
+        BlockingOutputStreamAsyncRequestBody requestBody =
+            AsyncRequestBody.forBlockingOutputStream(320L);
+        int totalLength = 320;
+        ByteBuffer expected = ByteBuffer.allocate(totalLength);
+        ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(totalLength);
+        requestBody.subscribe(subscriber);
+
+        int bytesToWrite = 32;
+        CancellableOutputStream outputStream = requestBody.outputStream();
+        byte[] bytes = new byte[512];
+        for (int i = 0; i < 10; i++) {
+            random.nextBytes(bytes);
+            expected.put(bytes, 10, bytesToWrite);
+            outputStream.write(bytes, 10, bytesToWrite);
+        }
+
+        outputStream.close();
+        ByteBuffer out = ByteBuffer.allocate(totalLength);
+        assertThat(subscriber.transferTo(out)).isEqualTo(END_OF_STREAM);
+        out.flip();
+
+        assertThat(out.array()).containsExactly(expected.array());
+    }
+
+    @Test
+    public void outputStream_writesArrayUseSameArray_shouldNotOverride() throws IOException {
+        BlockingOutputStreamAsyncRequestBody requestBody =
+            AsyncRequestBody.forBlockingOutputStream(320L);
+        int totalLength = 320;
+        ByteBuffer expected = ByteBuffer.allocate(totalLength);
+        ByteBufferStoringSubscriber subscriber = new ByteBufferStoringSubscriber(totalLength);
+        requestBody.subscribe(subscriber);
+
+        CancellableOutputStream outputStream = requestBody.outputStream();
+        byte[] bytes = new byte[32];
+        for (int i = 0; i < 10; i++) {
+            random.nextBytes(bytes);
+            expected.put(bytes);
+            outputStream.write(bytes);
+        }
+
+        outputStream.close();
+        ByteBuffer out = ByteBuffer.allocate(totalLength);
+        assertThat(subscriber.transferTo(out)).isEqualTo(END_OF_STREAM);
+        out.flip();
+
+        assertThat(out.array()).containsExactly(expected.array());
+    }
+
 }
diff --git a/services/s3/src/it/java/software/amazon/awssdk/services/s3/PutObjectIntegrationTest.java b/services/s3/src/it/java/software/amazon/awssdk/services/s3/PutObjectIntegrationTest.java
@@ -16,6 +16,7 @@
 
 package software.amazon.awssdk.services.s3;
 
+import static java.util.Base64.getEncoder;
 import static org.assertj.core.api.Assertions.assertThat;
 import static software.amazon.awssdk.services.s3.internal.checksums.ChecksumsEnabledValidator.CHECKSUM;
 import static software.amazon.awssdk.testutils.service.S3BucketUtils.temporaryBucketName;
@@ -25,45 +26,69 @@
 import java.io.FilterInputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.OutputStream;
 import java.net.URI;
 import java.nio.charset.StandardCharsets;
+import java.security.DigestInputStream;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.CompletableFuture;
-import org.junit.AfterClass;
-import org.junit.BeforeClass;
-import org.junit.Test;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.awssdk.core.ResponseInputStream;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.BlockingInputStreamAsyncRequestBody;
+import software.amazon.awssdk.core.async.BlockingOutputStreamAsyncRequestBody;
+import software.amazon.awssdk.core.checksums.RequestChecksumCalculation;
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.sync.RequestBody;
 import software.amazon.awssdk.http.ContentStreamProvider;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.utils.IoUtils;
 
 /**
  * Integration tests for {@code PutObject}.
  */
 public class PutObjectIntegrationTest extends S3IntegrationTestBase {
     private static final String BUCKET = temporaryBucketName(PutObjectIntegrationTest.class);
-    private static final String ASYNC_KEY = "async-key";
-    private static final String SYNC_KEY = "sync-key";
-    private static final String TEXT_CONTENT_TYPE = "text/plain";
+    private static final String ASYNC_KEY = "async-key" ;
+    private static final String SYNC_KEY = "sync-key" ;
+    private static final String TEXT_CONTENT_TYPE = "text/plain" ;
     private static final byte[] CONTENT = "Hello".getBytes(StandardCharsets.UTF_8);
+    private static final Random RANDOM = new Random(3470);
 
-    @BeforeClass
+    @BeforeAll
     public static void setUp() throws Exception {
         S3IntegrationTestBase.setUp();
         createBucket(BUCKET);
     }
 
-    @AfterClass
+    @AfterAll
     public static void tearDown() {
         deleteBucketAndAllContents(BUCKET);
     }
 
+    public static Stream<Arguments> s3Clients() {
+        return Stream.of(
+            Arguments.of(s3AsyncClientBuilder().requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED).build()),
+            Arguments.of(s3AsyncClientBuilder().build()),
+            Arguments.of(crtClientBuilder().build()));
+    }
+
     @Test
     public void objectInputStreamsAreClosed() {
         TestContentProvider provider = new TestContentProvider(CONTENT);
@@ -97,6 +122,88 @@ public void blockingInputStreamAsyncRequestBody_withContentType_isHonored() {
         assertThat(response.contentType()).isEqualTo(TEXT_CONTENT_TYPE);
     }
 
+    @ParameterizedTest
+    @MethodSource("s3Clients")
+    public void blockingOutputStreamAsyncRequestBody_writeArrayWithOffset_shouldSucceed(S3AsyncClient client) throws Exception {
+        BlockingOutputStreamAsyncRequestBody body = AsyncRequestBody.forBlockingOutputStream(320L);
+
+        PutObjectRequest.Builder request = PutObjectRequest.builder()
+                                                           .bucket(BUCKET)
+                                                           .key(ASYNC_KEY);
+        CompletableFuture<PutObjectResponse> responseFuture = client.putObject(request.build(), body);
+        MessageDigest expectedDigest = writeArrayWithOffset(body, 32);
+        responseFuture.join();
+
+        ResponseInputStream<GetObjectResponse> responseInputStream =
+            client.getObject(b -> b.bucket(BUCKET).key(ASYNC_KEY),
+                              AsyncResponseTransformer.toBlockingInputStream()).join();
+
+        DigestInputStream digestInputStream = new DigestInputStream(
+            responseInputStream, MessageDigest.getInstance("SHA-256"));
+        IoUtils.drainInputStream(digestInputStream);
+        MessageDigest actual = digestInputStream.getMessageDigest();
+
+        assertThat(getEncoder().encodeToString(actual.digest()))
+            .isEqualTo(getEncoder().encodeToString(expectedDigest.digest()));
+    }
+
+    @ParameterizedTest
+    @MethodSource("s3Clients")
+    public void blockingOutputStreamAsyncRequestBody_writeArray_shouldSucceed(S3AsyncClient client) throws Exception {
+        BlockingOutputStreamAsyncRequestBody body = AsyncRequestBody.forBlockingOutputStream(320L);
+
+        PutObjectRequest.Builder request = PutObjectRequest.builder()
+                                                           .bucket(BUCKET)
+                                                           .key(ASYNC_KEY);
+        CompletableFuture<PutObjectResponse> responseFuture = client.putObject(request.build(), body);
+        MessageDigest expectedDigest = writeArray(body);
+        responseFuture.join();
+
+        ResponseInputStream<GetObjectResponse> responseInputStream =
+            client.getObject(b -> b.bucket(BUCKET).key(ASYNC_KEY),
+                              AsyncResponseTransformer.toBlockingInputStream()).join();
+
+        DigestInputStream digestInputStream = new DigestInputStream(
+            responseInputStream, MessageDigest.getInstance("SHA-256"));
+        IoUtils.drainInputStream(digestInputStream);
+        MessageDigest actual = digestInputStream.getMessageDigest();
+
+        assertThat(getEncoder().encodeToString(actual.digest()))
+            .isEqualTo(getEncoder().encodeToString(expectedDigest.digest()));
+    }
+
+    private static MessageDigest writeArray(BlockingOutputStreamAsyncRequestBody body)
+        throws NoSuchAlgorithmException {
+        MessageDigest digest = MessageDigest.getInstance("SHA-256");
+        try (OutputStream outputStream = body.outputStream()) {
+            byte[] buffer = new byte[32];
+            for (int i = 0; i < 10; i++) {
+                RANDOM.nextBytes(buffer);
+                digest.update(buffer);
+                outputStream.write(buffer);
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        return digest;
+    }
+
+    private static MessageDigest writeArrayWithOffset(BlockingOutputStreamAsyncRequestBody body, int chunkSize)
+        throws NoSuchAlgorithmException {
+        MessageDigest digest = MessageDigest.getInstance("SHA-256");
+        try (OutputStream outputStream = body.outputStream()) {
+            byte[] buffer = new byte[1024];
+            for (int i = 0; i < 10; i++) {
+                RANDOM.nextBytes(buffer);
+                digest.update(buffer, 10, chunkSize);
+                outputStream.write(buffer, 10, chunkSize);
+            }
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+        return digest;
+    }
+
     @Test
     public void s3Client_usingHttpAndDisableChunkedEncoding() {
         try (S3Client s3Client = s3ClientBuilder()
diff --git a/utils/src/main/java/software/amazon/awssdk/utils/async/OutputStreamPublisher.java b/utils/src/main/java/software/amazon/awssdk/utils/async/OutputStreamPublisher.java
@@ -62,13 +62,19 @@ public void write(int b) {
     @Override
     public void write(byte[] b) {
         flush();
-        send(ByteBuffer.wrap(b));
+        ByteBuffer buffer = ByteBuffer.allocate(b.length);
+        buffer.put(b);
+        buffer.flip();
+        send(buffer.asReadOnlyBuffer());
     }
 
     @Override
     public void write(byte[] b, int off, int len) {
         flush();
-        send(ByteBuffer.wrap(b, off, len));
+        ByteBuffer buffer = ByteBuffer.allocate(len);
+        buffer.put(b, off, len);
+        buffer.flip();
+        send(buffer.asReadOnlyBuffer());
     }
 
     @Override
diff --git a/utils/src/test/java/software/amazon/awssdk/utils/async/OutputStreamPublisherTest.java b/utils/src/test/java/software/amazon/awssdk/utils/async/OutputStreamPublisherTest.java
@@ -19,6 +19,10 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.Random;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -27,13 +31,15 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
+import software.amazon.awssdk.utils.BinaryUtils;
 import software.amazon.awssdk.utils.ThreadFactoryBuilder;
 import software.amazon.awssdk.utils.async.ByteBufferStoringSubscriber.TransferResult;
 
 public class OutputStreamPublisherTest {
     private StoringSubscriber<ByteBuffer> storingSubscriber;
     private ByteBufferStoringSubscriber byteStoringSubscriber;
     private OutputStreamPublisher publisher;
+    private Random random = new Random(3470);
 
     @BeforeEach
     public void setup() {
@@ -60,6 +66,55 @@ public void oneByteWritesAreFlushedEventually() {
         });
     }
 
+    @Test
+    public void writeArrayReuseSameArray_shouldNotOverride() {
+        publisher.subscribe(storingSubscriber);
+        List<ByteBuffer> expectedBuffers = new ArrayList<>();
+
+        byte[] bytes = new byte[512];
+        for (int i = 0; i < 10; i++) {
+            random.nextBytes(bytes);
+            ByteBuffer byteBuffer = ByteBuffer.allocate(512);
+            byteBuffer.put(bytes);
+            byteBuffer.flip();
+
+            expectedBuffers.add(byteBuffer);
+            publisher.write(bytes);
+        }
+
+        List<ByteBuffer> actualBuffers = new ArrayList<>();
+        while(storingSubscriber.peek().isPresent()) {
+            StoringSubscriber.Event<ByteBuffer> event = storingSubscriber.poll().get();
+            actualBuffers.add(event.value());
+        }
+        assertThat(actualBuffers).hasSameElementsAs(expectedBuffers);
+    }
+
+    @Test
+    public void writeArrayWithOffsetReuseSameArray_shouldNotOverride() {
+        publisher.subscribe(storingSubscriber);
+        int bytesToWrite = 32;
+        List<ByteBuffer> expectedBuffers = new ArrayList<>();
+
+        byte[] bytes = new byte[512];
+        for (int i = 0; i < 10; i++) {
+            random.nextBytes(bytes);
+            ByteBuffer byteBuffer = ByteBuffer.allocate(bytesToWrite);
+            byteBuffer.put(bytes, 10, bytesToWrite);
+            byteBuffer.flip();
+
+            expectedBuffers.add(byteBuffer);
+            publisher.write(bytes, 10, bytesToWrite);
+        }
+
+        List<ByteBuffer> actualBuffers = new ArrayList<>();
+        while(storingSubscriber.peek().isPresent()) {
+            StoringSubscriber.Event<ByteBuffer> event = storingSubscriber.poll().get();
+            actualBuffers.add(event.value());
+        }
+        assertThat(actualBuffers).hasSameElementsAs(expectedBuffers);
+    }
+
     @Test
     public void flushDrainsBufferedBytes() {
         publisher.subscribe(storingSubscriber);

Original file line number	Diff line number	Diff line change
`@@ -37,6 +37,12 @@`
`37`	`37`	`"category": "AWS SDK for Java v2",`
`38`	`38`	`"contributor": "",`
`39`	`39`	`"description": "Updated endpoint and partition metadata."`
	`40`	`+ },`
	`41`	`+ {`
	`42`	`+ "type": "bugfix",`
	`43`	`+ "category": "AWS SDK for Java v2",`
	`44`	`+ "contributor": "",`
	`45`	`+ "description": "Copy bytes written to OutputStream of BlockingOutputStreamAsyncRequestBody"`
`40`	`46`	`}`
`41`	`47`	`]`
`42`	`48`	`}`