From 55d84e20b5b27f129cb613a1e43a7a71cf17d99e Mon Sep 17 00:00:00 2001 From: Bikram Sharma Date: Thu, 24 Sep 2026 14:30:51 -0700 Subject: [PATCH 1/3] fix: close getObject response stream for buffering transformers S3EncryptionClient.getObject never closed the ResponseInputStream from the toBlockingInputStream() pipeline, leaking Netty direct memory on every call with a buffering transformer (getObjectAsBytes, toFile). Close the stream when the transformer does not need the connection left open; streaming transformers (toInputStream) still return the stream for the caller to close. Fixes #518. --- .../encryption/s3/S3EncryptionClient.java | 23 ++++++++++-- ...ryptionClientGetObjectStreamCloseTest.java | 36 +++++++++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) create mode 100644 src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java diff --git a/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java b/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java index 3138b2b60..d526a2567 100644 --- a/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java +++ b/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java @@ -457,13 +457,32 @@ public T getObject(GetObjectRequest getObjectRequest, .commitmentPolicy(_commitmentPolicy) .build(); + ResponseInputStream joinFutureGet = null; + boolean callerOwnsStream = false; try { - ResponseInputStream joinFutureGet = pipeline.getObject(getObjectRequest, AsyncResponseTransformer.toBlockingInputStream()).join(); - return responseTransformer.transform(joinFutureGet.response(), AbortableInputStream.create(joinFutureGet)); + joinFutureGet = pipeline.getObject(getObjectRequest, AsyncResponseTransformer.toBlockingInputStream()).join(); + T result = responseTransformer.transform(joinFutureGet.response(), AbortableInputStream.create(joinFutureGet)); + // Streaming transformers (e.g. toInputStream) hand the stream back to the caller, who is + // then responsible for closing it. Buffering transformers (e.g. toBytes, toFile) fully + // consume the stream and return a materialized result, so ownership is not transferred and + // the stream is closed in the finally block below to release its buffers. + callerOwnsStream = responseTransformer.needsConnectionLeftOpen(); + return result; } catch (CompletionException e) { throw new S3EncryptionClientException(e.getCause().getMessage(), e.getCause()); } catch (Exception e) { throw new S3EncryptionClientException("Unable to transform response.", e); + } finally { + // Close the stream unless ownership was successfully handed to the caller. This covers the + // buffering case (leak fix) and also the case where transform threw before returning, so a + // streaming transformer's stream does not leak when the caller never received it. + if (joinFutureGet != null && !callerOwnsStream) { + try { + joinFutureGet.close(); + } catch (IOException e) { + throw new S3EncryptionClientException("Unable to close response stream.", e); + } + } } } diff --git a/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java new file mode 100644 index 000000000..959f2f529 --- /dev/null +++ b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java @@ -0,0 +1,36 @@ +// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.encryption.s3; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.File; +import java.io.IOException; + +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.core.sync.ResponseTransformer; +import software.amazon.awssdk.services.s3.model.GetObjectResponse; + +/** + * Pins the contract that {@link S3EncryptionClient#getObject} relies on to decide whether to close + * the response stream: buffering transformers do not need the connection left open (so the client + * closes the stream and avoids the leak), while streaming transformers do (the caller closes it). + */ +public class S3EncryptionClientGetObjectStreamCloseTest { + + @Test + public void bufferingTransformersDoNotNeedConnectionLeftOpen() throws IOException { + assertFalse(ResponseTransformer.toBytes().needsConnectionLeftOpen()); + + File tempFile = File.createTempFile("s3ec-close-test", ".tmp"); + tempFile.delete(); // toFile requires the file to not already exist + tempFile.deleteOnExit(); + assertFalse(ResponseTransformer.toFile(tempFile.toPath()).needsConnectionLeftOpen()); + } + + @Test + public void streamingTransformerNeedsConnectionLeftOpen() { + assertTrue(ResponseTransformer.toInputStream().needsConnectionLeftOpen()); + } +} From a2aa8c115099d06ad83acfc59c4bb3027f7f06ce Mon Sep 17 00:00:00 2001 From: Bikram Sharma Date: Fri, 2 Oct 2026 14:48:40 -0700 Subject: [PATCH 2/3] test: exercise getObject stream release end-to-end Replace the SDK-contract test with one that drives S3EncryptionClient.getObject through a real S3AsyncClient backed by an in-memory transport, so it fails without the stream close: a transformer that stops reading early or throws must release the response stream, while toInputStream leaves it open for the caller. Log instead of throw when closing the stream fails, so a close failure does not mask the transformed result or the original exception. --- .../encryption/s3/S3EncryptionClient.java | 14 +- ...ryptionClientGetObjectStreamCloseTest.java | 308 +++++++++++++++++- 2 files changed, 301 insertions(+), 21 deletions(-) diff --git a/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java b/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java index d526a2567..e91a8b894 100644 --- a/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java +++ b/src/main/java/software/amazon/encryption/s3/S3EncryptionClient.java @@ -462,10 +462,8 @@ public T getObject(GetObjectRequest getObjectRequest, try { joinFutureGet = pipeline.getObject(getObjectRequest, AsyncResponseTransformer.toBlockingInputStream()).join(); T result = responseTransformer.transform(joinFutureGet.response(), AbortableInputStream.create(joinFutureGet)); - // Streaming transformers (e.g. toInputStream) hand the stream back to the caller, who is - // then responsible for closing it. Buffering transformers (e.g. toBytes, toFile) fully - // consume the stream and return a materialized result, so ownership is not transferred and - // the stream is closed in the finally block below to release its buffers. + // Transformers that need the connection left open (e.g. toInputStream) return the stream to + // the caller, who becomes responsible for closing it. callerOwnsStream = responseTransformer.needsConnectionLeftOpen(); return result; } catch (CompletionException e) { @@ -473,14 +471,14 @@ public T getObject(GetObjectRequest getObjectRequest, } catch (Exception e) { throw new S3EncryptionClientException("Unable to transform response.", e); } finally { - // Close the stream unless ownership was successfully handed to the caller. This covers the - // buffering case (leak fix) and also the case where transform threw before returning, so a - // streaming transformer's stream does not leak when the caller never received it. + // Unless the caller now owns the stream, close it so the underlying connection is released + // even if the transformer did not read the response to the end or threw. if (joinFutureGet != null && !callerOwnsStream) { try { joinFutureGet.close(); } catch (IOException e) { - throw new S3EncryptionClientException("Unable to close response stream.", e); + // Don't let a close failure mask the transformed result or the original exception. + LogFactory.getLog(getClass()).debug("Unable to close GetObject response stream.", e); } } } diff --git a/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java index 959f2f529..1f6b31d78 100644 --- a/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java +++ b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java @@ -2,35 +2,317 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.encryption.s3; +import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; -import java.io.File; -import java.io.IOException; +import java.io.ByteArrayOutputStream; +import java.nio.ByteBuffer; +import java.security.SecureRandom; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import javax.crypto.SecretKey; +import javax.crypto.spec.SecretKeySpec; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.reactivestreams.Subscriber; +import org.reactivestreams.Subscription; + +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.ResponseInputStream; +import software.amazon.awssdk.core.async.SdkPublisher; +import software.amazon.awssdk.core.checksums.RequestChecksumCalculation; +import software.amazon.awssdk.core.sync.RequestBody; import software.amazon.awssdk.core.sync.ResponseTransformer; +import software.amazon.awssdk.http.SdkHttpFullResponse; +import software.amazon.awssdk.http.SdkHttpMethod; +import software.amazon.awssdk.http.async.AsyncExecuteRequest; +import software.amazon.awssdk.http.async.SdkAsyncHttpClient; +import software.amazon.awssdk.http.async.SdkAsyncHttpResponseHandler; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3AsyncClient; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; import software.amazon.awssdk.services.s3.model.GetObjectResponse; +import software.amazon.awssdk.services.s3.model.PutObjectRequest; /** - * Pins the contract that {@link S3EncryptionClient#getObject} relies on to decide whether to close - * the response stream: buffering transformers do not need the connection left open (so the client - * closes the stream and avoids the leak), while streaming transformers do (the caller closes it). + * Verifies that {@link S3EncryptionClient#getObject} releases the response stream it opens on the + * wrapped async client, so the underlying HTTP connection is returned to the pool. + *

+ * The wrapped async client is a real {@link S3AsyncClient} backed by an in-memory transport, so the + * full encrypt / decrypt path runs. The object is larger than what the blocking input stream buffers, + * and delayed authentication is enabled so plaintext is streamed rather than buffered in full; a + * transformer that stops reading early therefore leaves the response body unconsumed, and the only + * thing that releases it is the client closing (cancelling) the stream. */ public class S3EncryptionClientGetObjectStreamCloseTest { + private static final String BUCKET = "test-bucket"; + private static final String KEY = "test-key"; + private static final int OBJECT_SIZE = 8 * 1024 * 1024; + private static final int CHUNK_SIZE = 64 * 1024; + + private InMemoryTransport transport; + private S3EncryptionClient client; + private byte[] plaintext; + + @BeforeEach + public void setUp() { + transport = new InMemoryTransport(); + StaticCredentialsProvider creds = StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")); + S3AsyncClient wrappedAsyncClient = S3AsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(creds) + .requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED) + .httpClient(transport) + .build(); + S3Client wrappedClient = S3Client.builder() + .region(Region.US_WEST_2) + .credentialsProvider(creds) + .build(); + + byte[] keyBytes = new byte[32]; + new SecureRandom().nextBytes(keyBytes); + SecretKey aesKey = new SecretKeySpec(keyBytes, "AES"); + client = S3EncryptionClient.builderV4() + .wrappedClient(wrappedClient) + .wrappedAsyncClient(wrappedAsyncClient) + .aesKey(aesKey) + .enableDelayedAuthenticationMode(true) + .build(); + + plaintext = new byte[OBJECT_SIZE]; + new SecureRandom().nextBytes(plaintext); + client.putObject(PutObjectRequest.builder().bucket(BUCKET).key(KEY).build(), RequestBody.fromBytes(plaintext)); + } + + @AfterEach + public void tearDown() { + client.close(); + transport.shutdown(); + } + @Test - public void bufferingTransformersDoNotNeedConnectionLeftOpen() throws IOException { - assertFalse(ResponseTransformer.toBytes().needsConnectionLeftOpen()); + public void bufferingTransformerThatStopsEarlyReleasesResponseStream() { + // needsConnectionLeftOpen() is false, so the client owns the stream and must close it. + int firstByte = client.getObject(getRequest(), (response, inputStream) -> inputStream.read()); - File tempFile = File.createTempFile("s3ec-close-test", ".tmp"); - tempFile.delete(); // toFile requires the file to not already exist - tempFile.deleteOnExit(); - assertFalse(ResponseTransformer.toFile(tempFile.toPath()).needsConnectionLeftOpen()); + assertEquals(plaintext[0] & 0xFF, firstByte); + assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after getObject returned"); } @Test - public void streamingTransformerNeedsConnectionLeftOpen() { - assertTrue(ResponseTransformer.toInputStream().needsConnectionLeftOpen()); + public void transformerThatThrowsReleasesResponseStream() { + assertThrows(S3EncryptionClientException.class, () -> client.getObject(getRequest(), (response, inputStream) -> { + inputStream.read(); + throw new IllegalStateException("transform failed"); + })); + + assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after transform threw"); + } + + @Test + public void streamingTransformerLeavesResponseStreamOpenForCaller() throws Exception { + try (ResponseInputStream stream = client.getObject(getRequest(), ResponseTransformer.toInputStream())) { + assertFalse(transport.lastBody.cancelled.get(), "caller-owned stream was closed by getObject"); + assertEquals(plaintext[0] & 0xFF, stream.read()); + } + + assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after the caller closed it"); + } + + @Test + public void fullyConsumedResponseStillDecryptsCorrectly() { + byte[] result = client.getObjectAsBytes(getRequest()).asByteArray(); + + assertEquals(OBJECT_SIZE, result.length); + assertTrue(java.util.Arrays.equals(plaintext, result)); + } + + private static GetObjectRequest getRequest() { + return GetObjectRequest.builder().bucket(BUCKET).key(KEY).build(); + } + + /** Stores a single object on PUT and serves it on GET, recording whether the GET body was cancelled. */ + private static final class InMemoryTransport implements SdkAsyncHttpClient { + private final ExecutorService executor = Executors.newCachedThreadPool(r -> { + Thread t = new Thread(r, "in-memory-transport"); + t.setDaemon(true); + return t; + }); + private final Map storedMetadata = new HashMap<>(); + private byte[] storedBody; + volatile RecordingBodyPublisher lastBody; + + @Override + public CompletableFuture execute(AsyncExecuteRequest request) { + return request.request().method() == SdkHttpMethod.PUT ? put(request) : get(request); + } + + private CompletableFuture put(AsyncExecuteRequest request) { + CompletableFuture done = new CompletableFuture<>(); + ByteArrayOutputStream body = new ByteArrayOutputStream(); + request.requestContentPublisher().subscribe(new Subscriber() { + @Override + public void onSubscribe(Subscription subscription) { + subscription.request(Long.MAX_VALUE); + } + + @Override + public void onNext(ByteBuffer byteBuffer) { + byte[] bytes = new byte[byteBuffer.remaining()]; + byteBuffer.get(bytes); + body.write(bytes, 0, bytes.length); + } + + @Override + public void onError(Throwable t) { + request.responseHandler().onError(t); + done.completeExceptionally(t); + } + + @Override + public void onComplete() { + storedBody = body.toByteArray(); + request.request().headers().forEach((name, values) -> { + if (name.toLowerCase().startsWith("x-amz-meta-")) { + storedMetadata.put(name, values.get(0)); + } + }); + SdkAsyncHttpResponseHandler handler = request.responseHandler(); + handler.onHeaders(SdkHttpFullResponse.builder().statusCode(200).putHeader("ETag", "\"etag\"").build()); + handler.onStream(new RecordingBodyPublisher(new byte[0], CHUNK_SIZE, executor)); + done.complete(null); + } + }); + return done; + } + + private CompletableFuture get(AsyncExecuteRequest request) { + SdkHttpFullResponse.Builder response = SdkHttpFullResponse.builder() + .statusCode(200) + .putHeader("ETag", "\"etag\"") + .putHeader("Content-Length", String.valueOf(storedBody.length)); + storedMetadata.forEach(response::putHeader); + lastBody = new RecordingBodyPublisher(storedBody, CHUNK_SIZE, executor); + + SdkAsyncHttpResponseHandler handler = request.responseHandler(); + handler.onHeaders(response.build()); + handler.onStream(lastBody); + return CompletableFuture.completedFuture(null); + } + + void shutdown() { + executor.shutdownNow(); + } + + @Override + public String clientName() { + return "in-memory-transport"; + } + + @Override + public void close() { + } + } + + /** + * Emits the body in chunks only as they are requested, like a real HTTP client applying + * backpressure, and records whether the subscriber cancelled. + */ + private static final class RecordingBodyPublisher implements SdkPublisher { + private final byte[] body; + private final int chunkSize; + private final ExecutorService executor; + final AtomicBoolean cancelled = new AtomicBoolean(false); + private final CompletableFuture cancelledFuture = new CompletableFuture<>(); + + private RecordingBodyPublisher(byte[] body, int chunkSize, ExecutorService executor) { + this.body = body; + this.chunkSize = chunkSize; + this.executor = executor; + } + + boolean awaitCancelled() { + try { + cancelledFuture.get(5, TimeUnit.SECONDS); + return true; + } catch (Exception e) { + return false; + } + } + + @Override + public void subscribe(Subscriber subscriber) { + subscriber.onSubscribe(new Subscription() { + private long demand; + private int position; + private boolean delivering; + private boolean terminated; + + @Override + public void request(long n) { + synchronized (this) { + if (terminated) { + return; + } + demand = demand + n < 0 ? Long.MAX_VALUE : demand + n; + if (delivering) { + return; + } + delivering = true; + } + executor.execute(this::deliver); + } + + private void deliver() { + while (true) { + ByteBuffer chunk; + boolean complete = false; + synchronized (this) { + if (terminated || demand == 0) { + delivering = false; + return; + } + if (position >= body.length) { + terminated = true; + complete = true; + chunk = null; + } else { + int length = Math.min(chunkSize, body.length - position); + chunk = ByteBuffer.wrap(body, position, length).slice(); + position += length; + demand--; + } + } + if (complete) { + subscriber.onComplete(); + return; + } + subscriber.onNext(chunk); + } + } + + @Override + public void cancel() { + synchronized (this) { + terminated = true; + } + cancelled.set(true); + cancelledFuture.complete(null); + } + }); + } } } From 837f0331817dbe1a60405f40cdd3bc7f13c45e06 Mon Sep 17 00:00:00 2001 From: Bikram Sharma Date: Fri, 2 Oct 2026 15:22:48 -0700 Subject: [PATCH 3/3] test: simplify getObject stream release test --- ...ryptionClientGetObjectStreamCloseTest.java | 274 ++++++------------ 1 file changed, 94 insertions(+), 180 deletions(-) diff --git a/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java index 1f6b31d78..463c628ac 100644 --- a/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java +++ b/src/test/java/software/amazon/encryption/s3/S3EncryptionClientGetObjectStreamCloseTest.java @@ -16,9 +16,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; -import javax.crypto.SecretKey; import javax.crypto.spec.SecretKeySpec; import org.junit.jupiter.api.AfterEach; @@ -30,6 +28,7 @@ import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.core.ResponseInputStream; +import software.amazon.awssdk.core.async.AsyncRequestBody; import software.amazon.awssdk.core.async.SdkPublisher; import software.amazon.awssdk.core.checksums.RequestChecksumCalculation; import software.amazon.awssdk.core.sync.RequestBody; @@ -38,7 +37,6 @@ import software.amazon.awssdk.http.SdkHttpMethod; import software.amazon.awssdk.http.async.AsyncExecuteRequest; import software.amazon.awssdk.http.async.SdkAsyncHttpClient; -import software.amazon.awssdk.http.async.SdkAsyncHttpResponseHandler; import software.amazon.awssdk.regions.Region; import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.S3Client; @@ -47,176 +45,133 @@ import software.amazon.awssdk.services.s3.model.PutObjectRequest; /** - * Verifies that {@link S3EncryptionClient#getObject} releases the response stream it opens on the - * wrapped async client, so the underlying HTTP connection is returned to the pool. + * Verifies that {@link S3EncryptionClient#getObject} closes the response stream unless the caller + * owns it, so the underlying HTTP connection is released. *

- * The wrapped async client is a real {@link S3AsyncClient} backed by an in-memory transport, so the - * full encrypt / decrypt path runs. The object is larger than what the blocking input stream buffers, - * and delayed authentication is enabled so plaintext is streamed rather than buffered in full; a - * transformer that stops reading early therefore leaves the response body unconsumed, and the only - * thing that releases it is the client closing (cancelling) the stream. + * The wrapped async client runs over an in-memory transport whose response body is delivered in + * chunks on demand, like a real connection. The object is larger than the blocking input stream + * buffers and delayed authentication streams plaintext, so a transformer that stops reading early + * leaves the body unconsumed; only closing the stream cancels it. */ public class S3EncryptionClientGetObjectStreamCloseTest { - private static final String BUCKET = "test-bucket"; - private static final String KEY = "test-key"; - private static final int OBJECT_SIZE = 8 * 1024 * 1024; - private static final int CHUNK_SIZE = 64 * 1024; + private static final GetObjectRequest GET_REQUEST = GetObjectRequest.builder().bucket("bucket").key("key").build(); - private InMemoryTransport transport; + private final ExecutorService executor = Executors.newCachedThreadPool(); + private final InMemoryTransport transport = new InMemoryTransport(executor); private S3EncryptionClient client; private byte[] plaintext; @BeforeEach public void setUp() { - transport = new InMemoryTransport(); StaticCredentialsProvider creds = StaticCredentialsProvider.create(AwsBasicCredentials.create("akid", "skid")); - S3AsyncClient wrappedAsyncClient = S3AsyncClient.builder() - .region(Region.US_WEST_2) - .credentialsProvider(creds) - .requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED) - .httpClient(transport) - .build(); - S3Client wrappedClient = S3Client.builder() - .region(Region.US_WEST_2) - .credentialsProvider(creds) - .build(); - - byte[] keyBytes = new byte[32]; - new SecureRandom().nextBytes(keyBytes); - SecretKey aesKey = new SecretKeySpec(keyBytes, "AES"); + byte[] key = new byte[32]; + new SecureRandom().nextBytes(key); client = S3EncryptionClient.builderV4() - .wrappedClient(wrappedClient) - .wrappedAsyncClient(wrappedAsyncClient) - .aesKey(aesKey) + .wrappedClient(S3Client.builder().region(Region.US_WEST_2).credentialsProvider(creds).build()) + .wrappedAsyncClient(S3AsyncClient.builder() + .region(Region.US_WEST_2) + .credentialsProvider(creds) + // Keeps the PUT body unencoded so the transport can store it as-is. + .requestChecksumCalculation(RequestChecksumCalculation.WHEN_REQUIRED) + .httpClient(transport) + .build()) + .aesKey(new SecretKeySpec(key, "AES")) .enableDelayedAuthenticationMode(true) .build(); - plaintext = new byte[OBJECT_SIZE]; + plaintext = new byte[8 * 1024 * 1024]; new SecureRandom().nextBytes(plaintext); - client.putObject(PutObjectRequest.builder().bucket(BUCKET).key(KEY).build(), RequestBody.fromBytes(plaintext)); + client.putObject(PutObjectRequest.builder().bucket("bucket").key("key").build(), RequestBody.fromBytes(plaintext)); } @AfterEach public void tearDown() { client.close(); - transport.shutdown(); + executor.shutdownNow(); } @Test - public void bufferingTransformerThatStopsEarlyReleasesResponseStream() { - // needsConnectionLeftOpen() is false, so the client owns the stream and must close it. - int firstByte = client.getObject(getRequest(), (response, inputStream) -> inputStream.read()); + public void transformerThatStopsEarlyReleasesResponseStream() { + int firstByte = client.getObject(GET_REQUEST, (response, inputStream) -> inputStream.read()); assertEquals(plaintext[0] & 0xFF, firstByte); - assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after getObject returned"); + assertTrue(transport.awaitBodyCancelled(), "response stream was not released after getObject returned"); } @Test public void transformerThatThrowsReleasesResponseStream() { - assertThrows(S3EncryptionClientException.class, () -> client.getObject(getRequest(), (response, inputStream) -> { + assertThrows(S3EncryptionClientException.class, () -> client.getObject(GET_REQUEST, (response, inputStream) -> { inputStream.read(); throw new IllegalStateException("transform failed"); })); - assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after transform threw"); + assertTrue(transport.awaitBodyCancelled(), "response stream was not released after transform threw"); } @Test public void streamingTransformerLeavesResponseStreamOpenForCaller() throws Exception { - try (ResponseInputStream stream = client.getObject(getRequest(), ResponseTransformer.toInputStream())) { - assertFalse(transport.lastBody.cancelled.get(), "caller-owned stream was closed by getObject"); + try (ResponseInputStream stream = client.getObject(GET_REQUEST, ResponseTransformer.toInputStream())) { assertEquals(plaintext[0] & 0xFF, stream.read()); + assertFalse(transport.bodyCancelled.isDone(), "caller-owned stream was closed by getObject"); } - assertTrue(transport.lastBody.awaitCancelled(), "response stream was not released after the caller closed it"); - } - - @Test - public void fullyConsumedResponseStillDecryptsCorrectly() { - byte[] result = client.getObjectAsBytes(getRequest()).asByteArray(); - - assertEquals(OBJECT_SIZE, result.length); - assertTrue(java.util.Arrays.equals(plaintext, result)); + assertTrue(transport.awaitBodyCancelled(), "response stream was not released after the caller closed it"); } - private static GetObjectRequest getRequest() { - return GetObjectRequest.builder().bucket(BUCKET).key(KEY).build(); - } - - /** Stores a single object on PUT and serves it on GET, recording whether the GET body was cancelled. */ + /** Stores the object on PUT and serves it on GET, recording whether the GET body was cancelled. */ private static final class InMemoryTransport implements SdkAsyncHttpClient { - private final ExecutorService executor = Executors.newCachedThreadPool(r -> { - Thread t = new Thread(r, "in-memory-transport"); - t.setDaemon(true); - return t; - }); - private final Map storedMetadata = new HashMap<>(); - private byte[] storedBody; - volatile RecordingBodyPublisher lastBody; + private final ExecutorService executor; + private final Map metadata = new HashMap<>(); + private byte[] body; + private volatile CompletableFuture bodyCancelled = new CompletableFuture<>(); - @Override - public CompletableFuture execute(AsyncExecuteRequest request) { - return request.request().method() == SdkHttpMethod.PUT ? put(request) : get(request); + InMemoryTransport(ExecutorService executor) { + this.executor = executor; } - private CompletableFuture put(AsyncExecuteRequest request) { - CompletableFuture done = new CompletableFuture<>(); - ByteArrayOutputStream body = new ByteArrayOutputStream(); - request.requestContentPublisher().subscribe(new Subscriber() { - @Override - public void onSubscribe(Subscription subscription) { - subscription.request(Long.MAX_VALUE); - } - - @Override - public void onNext(ByteBuffer byteBuffer) { - byte[] bytes = new byte[byteBuffer.remaining()]; - byteBuffer.get(bytes); - body.write(bytes, 0, bytes.length); - } - - @Override - public void onError(Throwable t) { - request.responseHandler().onError(t); - done.completeExceptionally(t); - } - - @Override - public void onComplete() { - storedBody = body.toByteArray(); - request.request().headers().forEach((name, values) -> { - if (name.toLowerCase().startsWith("x-amz-meta-")) { - storedMetadata.put(name, values.get(0)); - } - }); - SdkAsyncHttpResponseHandler handler = request.responseHandler(); - handler.onHeaders(SdkHttpFullResponse.builder().statusCode(200).putHeader("ETag", "\"etag\"").build()); - handler.onStream(new RecordingBodyPublisher(new byte[0], CHUNK_SIZE, executor)); - done.complete(null); - } - }); - return done; + boolean awaitBodyCancelled() { + try { + bodyCancelled.get(5, TimeUnit.SECONDS); + return true; + } catch (Exception e) { + return false; + } } - private CompletableFuture get(AsyncExecuteRequest request) { - SdkHttpFullResponse.Builder response = SdkHttpFullResponse.builder() - .statusCode(200) - .putHeader("ETag", "\"etag\"") - .putHeader("Content-Length", String.valueOf(storedBody.length)); - storedMetadata.forEach(response::putHeader); - lastBody = new RecordingBodyPublisher(storedBody, CHUNK_SIZE, executor); + @Override + public CompletableFuture execute(AsyncExecuteRequest request) { + SdkHttpFullResponse.Builder response = SdkHttpFullResponse.builder().statusCode(200).putHeader("ETag", "\"etag\""); + if (request.request().method() == SdkHttpMethod.PUT) { + ByteArrayOutputStream received = new ByteArrayOutputStream(); + return SdkPublisher.adapt(request.requestContentPublisher()) + .subscribe(buffer -> { + byte[] bytes = new byte[buffer.remaining()]; + buffer.get(bytes); + received.write(bytes, 0, bytes.length); + }) + .thenRun(() -> { + body = received.toByteArray(); + // S3EC stores its encryption metadata in user metadata headers. + request.request().headers().forEach((name, values) -> { + if (name.toLowerCase().startsWith("x-amz-meta-")) { + metadata.put(name, values.get(0)); + } + }); + request.responseHandler().onHeaders(response.build()); + request.responseHandler().onStream(AsyncRequestBody.empty()); + }); + } - SdkAsyncHttpResponseHandler handler = request.responseHandler(); - handler.onHeaders(response.build()); - handler.onStream(lastBody); + metadata.forEach(response::putHeader); + response.putHeader("Content-Length", String.valueOf(body.length)); + bodyCancelled = new CompletableFuture<>(); + CompletableFuture cancelled = bodyCancelled; + request.responseHandler().onHeaders(response.build()); + request.responseHandler().onStream(new OnDemandBody(body, executor, cancelled)); return CompletableFuture.completedFuture(null); } - void shutdown() { - executor.shutdownNow(); - } - @Override public String clientName() { return "in-memory-transport"; @@ -228,89 +183,48 @@ public void close() { } /** - * Emits the body in chunks only as they are requested, like a real HTTP client applying - * backpressure, and records whether the subscriber cancelled. + * Emits the body in 64 KiB chunks only as they are requested, like a real connection applying + * backpressure, and completes {@code cancelled} if the subscriber cancels. */ - private static final class RecordingBodyPublisher implements SdkPublisher { + private static final class OnDemandBody implements SdkPublisher { private final byte[] body; - private final int chunkSize; private final ExecutorService executor; - final AtomicBoolean cancelled = new AtomicBoolean(false); - private final CompletableFuture cancelledFuture = new CompletableFuture<>(); + private final CompletableFuture cancelled; - private RecordingBodyPublisher(byte[] body, int chunkSize, ExecutorService executor) { + OnDemandBody(byte[] body, ExecutorService executor, CompletableFuture cancelled) { this.body = body; - this.chunkSize = chunkSize; this.executor = executor; - } - - boolean awaitCancelled() { - try { - cancelledFuture.get(5, TimeUnit.SECONDS); - return true; - } catch (Exception e) { - return false; - } + this.cancelled = cancelled; } @Override public void subscribe(Subscriber subscriber) { subscriber.onSubscribe(new Subscription() { - private long demand; private int position; - private boolean delivering; - private boolean terminated; @Override public void request(long n) { - synchronized (this) { - if (terminated) { - return; - } - demand = demand + n < 0 ? Long.MAX_VALUE : demand + n; - if (delivering) { - return; - } - delivering = true; - } - executor.execute(this::deliver); - } - - private void deliver() { - while (true) { - ByteBuffer chunk; - boolean complete = false; + // Deliver on another thread so onNext never re-enters the caller of request(). + executor.execute(() -> { synchronized (this) { - if (terminated || demand == 0) { - delivering = false; - return; - } - if (position >= body.length) { - terminated = true; - complete = true; - chunk = null; - } else { - int length = Math.min(chunkSize, body.length - position); - chunk = ByteBuffer.wrap(body, position, length).slice(); + for (long i = 0; i < n && !cancelled.isDone(); i++) { + if (position >= body.length) { + if (position++ == body.length) { + subscriber.onComplete(); + } + return; + } + int length = Math.min(64 * 1024, body.length - position); + subscriber.onNext(ByteBuffer.wrap(body, position, length).slice()); position += length; - demand--; } } - if (complete) { - subscriber.onComplete(); - return; - } - subscriber.onNext(chunk); - } + }); } @Override public void cancel() { - synchronized (this) { - terminated = true; - } - cancelled.set(true); - cancelledFuture.complete(null); + cancelled.complete(null); } }); }