Skip to content

Commit cdec144

Browse files
davidh44davidh44
authored
Preserve contentLength for InMemoryPublisher (aws#7141)
* Preserve contentLength for InMemoryPublisher * Use imports instead of FQCNs * Add unit tests for InMemoryPublisher contentLength * Use imports instead of FQCNs * Remove changelog
1 parent 5ca6ef3 commit cdec144

9 files changed

Lines changed: 152 additions & 13 deletions

File tree

‎core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/FlexibleChecksummer.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -81,9 +81,10 @@ public CompletableFuture<Publisher<ByteBuffer>> checksum(Publisher<ByteBuffer> p
8181
}
8282

8383
payload.subscribe(checksumSubscriber);
84-
CompletableFuture<Publisher<ByteBuffer>> result = checksumSubscriber.completeFuture();
85-
result.thenRun(() -> addChecksums(request));
86-
return result;
84+
return checksumSubscriber.completeFuture().thenApply(checksummedPayload -> {
85+
addChecksums(request);
86+
return checksummedPayload;
87+
});
8788
}
8889

8990
private void addChecksums(SdkHttpRequest.Builder request) {

‎core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/ChecksumSubscriber.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ public void onError(Throwable throwable) {
9595

9696
@Override
9797
public void onComplete() {
98-
checksumming.complete(new InMemoryPublisher(bufferedPayload));
98+
long totalBytes = bufferedPayload.stream().mapToLong(ByteBuffer::remaining).sum();
99+
checksumming.complete(new InMemoryPublisher(bufferedPayload, totalBytes));
99100
}
100101

101102
public CompletableFuture<Publisher<ByteBuffer>> completeFuture() {

‎core/http-auth-aws/src/main/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisher.java‎

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,24 +18,33 @@
1818
import java.nio.ByteBuffer;
1919
import java.util.ArrayList;
2020
import java.util.List;
21+
import java.util.Optional;
2122
import java.util.concurrent.atomic.AtomicBoolean;
2223
import java.util.concurrent.atomic.AtomicLong;
23-
import org.reactivestreams.Publisher;
2424
import org.reactivestreams.Subscriber;
2525
import org.reactivestreams.Subscription;
2626
import software.amazon.awssdk.annotations.SdkInternalApi;
27+
import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
2728
import software.amazon.awssdk.utils.Validate;
2829

2930
/**
30-
* Temporarily used for buffering all data into memory.
31+
* A content-length-aware publisher that replays buffered data. Used by {@link ChecksumSubscriber} to replay the payload after
32+
* checksumming.
3133
*/
3234
@SdkInternalApi
33-
public class InMemoryPublisher implements Publisher<ByteBuffer> {
35+
public class InMemoryPublisher implements SdkHttpContentPublisher {
3436
private final AtomicBoolean subscribed = new AtomicBoolean(false);
3537
private final List<ByteBuffer> data;
38+
private final long length;
3639

37-
public InMemoryPublisher(List<ByteBuffer> data) {
40+
public InMemoryPublisher(List<ByteBuffer> data, long length) {
3841
this.data = new ArrayList<>(Validate.noNullElements(data, "Data must not contain null elements."));
42+
this.length = length;
43+
}
44+
45+
@Override
46+
public Optional<Long> contentLength() {
47+
return Optional.of(length);
3948
}
4049

4150
@Override

‎core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisherTckTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ public Publisher<ByteBuffer> createPublisher(long elements) {
3737
for (long i = 0; i < elements; i++) {
3838
data.add(ByteBuffer.wrap(new byte[]{(byte) (i % 127)}));
3939
}
40-
return new InMemoryPublisher(data);
40+
return new InMemoryPublisher(data, elements);
4141
}
4242

4343
@Override

‎core/http-auth-aws/src/test/java/software/amazon/awssdk/http/auth/aws/internal/signer/io/InMemoryPublisherTest.java‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,11 +31,27 @@
3131

3232
public class InMemoryPublisherTest {
3333

34+
@Test
35+
public void contentLength_returnsProvidedLength() {
36+
List<ByteBuffer> data = Arrays.asList(
37+
ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8)),
38+
ByteBuffer.wrap(" world".getBytes(StandardCharsets.UTF_8))
39+
);
40+
InMemoryPublisher publisher = new InMemoryPublisher(data, 11);
41+
assertThat(publisher.contentLength()).hasValue(11L);
42+
}
43+
44+
@Test
45+
public void contentLength_emptyPayload_returnsZero() {
46+
InMemoryPublisher publisher = new InMemoryPublisher(new ArrayList<>(), 0);
47+
assertThat(publisher.contentLength()).hasValue(0L);
48+
}
49+
3450
@Test
3551
public void subscribe_deliversAllData() throws Exception {
3652
byte[] bytes = "test data".getBytes(StandardCharsets.UTF_8);
3753
List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap(bytes));
38-
InMemoryPublisher publisher = new InMemoryPublisher(data);
54+
InMemoryPublisher publisher = new InMemoryPublisher(data, bytes.length);
3955

4056
List<ByteBuffer> received = new ArrayList<>();
4157
CountDownLatch completed = new CountDownLatch(1);
@@ -72,7 +88,7 @@ public void subscribe_reEntrantRequestFromOnNext_doesNotDeadlock() throws Except
7288
ByteBuffer.wrap("b".getBytes(StandardCharsets.UTF_8)),
7389
ByteBuffer.wrap("c".getBytes(StandardCharsets.UTF_8))
7490
);
75-
InMemoryPublisher publisher = new InMemoryPublisher(data);
91+
InMemoryPublisher publisher = new InMemoryPublisher(data, 3);
7692

7793
List<ByteBuffer> received = new ArrayList<>();
7894
CountDownLatch completed = new CountDownLatch(1);
@@ -116,7 +132,7 @@ public void onComplete() {
116132
@Test
117133
public void subscribe_secondSubscription_getsError() {
118134
List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap("x".getBytes(StandardCharsets.UTF_8)));
119-
InMemoryPublisher publisher = new InMemoryPublisher(data);
135+
InMemoryPublisher publisher = new InMemoryPublisher(data, 1);
120136

121137
publisher.subscribe(new Subscriber<ByteBuffer>() {
122138
@Override public void onSubscribe(Subscription s) { s.request(1); }
@@ -139,7 +155,7 @@ public void subscribe_secondSubscription_getsError() {
139155
@Test
140156
public void subscribe_requestNonPositive_signalsError() throws Exception {
141157
List<ByteBuffer> data = Arrays.asList(ByteBuffer.wrap("x".getBytes(StandardCharsets.UTF_8)));
142-
InMemoryPublisher publisher = new InMemoryPublisher(data);
158+
InMemoryPublisher publisher = new InMemoryPublisher(data, 1);
143159

144160
AtomicBoolean gotError = new AtomicBoolean(false);
145161
CountDownLatch completed = new CountDownLatch(1);

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncSigningStage.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Optional;
2222
import java.util.concurrent.CompletableFuture;
2323
import org.reactivestreams.Publisher;
24+
import org.reactivestreams.Subscriber;
2425
import software.amazon.awssdk.annotations.SdkInternalApi;
2526
import software.amazon.awssdk.core.SelectedAuthScheme;
2627
import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -38,6 +39,7 @@
3839
import software.amazon.awssdk.core.signer.Signer;
3940
import software.amazon.awssdk.http.SdkHttpFullRequest;
4041
import software.amazon.awssdk.http.SdkHttpRequest;
42+
import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
4143
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
4244
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
4345
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
@@ -154,6 +156,19 @@ private static void updateAsyncRequestBodyInContexts(RequestExecutionContext con
154156
Publisher<ByteBuffer> signedPayload = optionalPayload.get();
155157
if (signedPayload instanceof AsyncRequestBody) {
156158
newAsyncRequestBody = (AsyncRequestBody) signedPayload;
159+
} else if (signedPayload instanceof SdkHttpContentPublisher) {
160+
SdkHttpContentPublisher contentPublisher = (SdkHttpContentPublisher) signedPayload;
161+
newAsyncRequestBody = new AsyncRequestBody() {
162+
@Override
163+
public Optional<Long> contentLength() {
164+
return contentPublisher.contentLength();
165+
}
166+
167+
@Override
168+
public void subscribe(Subscriber<? super ByteBuffer> s) {
169+
contentPublisher.subscribe(s);
170+
}
171+
};
157172
} else {
158173
newAsyncRequestBody = AsyncRequestBody.fromPublisher(signedPayload);
159174
}

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeAsyncHttpRequestStage.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import static software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute.SDK_HTTP_EXECUTION_ATTRIBUTES;
1919
import static software.amazon.awssdk.core.internal.http.timers.TimerUtils.resolveTimeoutInMillis;
2020
import static software.amazon.awssdk.http.Header.CONTENT_LENGTH;
21+
import static software.amazon.awssdk.http.Header.TRANSFER_ENCODING;
2122

2223
import java.nio.ByteBuffer;
2324
import java.time.Duration;
@@ -246,6 +247,10 @@ private boolean shouldSetContentLength(SdkHttpFullRequest request, SdkHttpConten
246247
return false;
247248
}
248249

250+
if (request.firstMatchingHeader(TRANSFER_ENCODING).isPresent()) {
251+
return false;
252+
}
253+
249254
return Optional.ofNullable(requestProvider).flatMap(SdkHttpContentPublisher::contentLength).isPresent();
250255
}
251256

‎core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/AsyncSigningStageTest.java‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import java.time.Instant;
3434
import java.time.temporal.ChronoUnit;
3535
import java.util.HashMap;
36+
import java.util.Optional;
3637
import java.util.concurrent.CompletableFuture;
3738
import org.junit.Before;
3839
import org.junit.Test;
@@ -43,6 +44,8 @@
4344
import org.mockito.Mock;
4445
import org.mockito.junit.MockitoJUnitRunner;
4546
import org.reactivestreams.Publisher;
47+
import org.reactivestreams.Subscriber;
48+
import org.reactivestreams.Subscription;
4649
import software.amazon.awssdk.core.SdkRequest;
4750
import software.amazon.awssdk.core.SelectedAuthScheme;
4851
import software.amazon.awssdk.core.async.AsyncRequestBody;
@@ -58,6 +61,7 @@
5861
import software.amazon.awssdk.core.signer.Signer;
5962
import software.amazon.awssdk.http.SdkHttpFullRequest;
6063
import software.amazon.awssdk.http.SdkHttpRequest;
64+
import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
6165
import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
6266
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignRequest;
6367
import software.amazon.awssdk.http.auth.spi.signer.AsyncSignedRequest;
@@ -596,6 +600,50 @@ private RequestExecutionContext createContext(SelectedAuthScheme<Identity> selec
596600
return context;
597601
}
598602

603+
@Test
604+
public void execute_signerReturnsSdkHttpContentPublisher_preservesContentLength() throws Exception {
605+
AsyncRequestBody asyncPayload = AsyncRequestBody.fromString("async request body");
606+
607+
SelectedAuthScheme<Identity> selectedAuthScheme = new SelectedAuthScheme<>(
608+
CompletableFuture.completedFuture(identity),
609+
httpSigner,
610+
AuthSchemeOption.builder()
611+
.schemeId("my.auth#myAuth")
612+
.build());
613+
RequestExecutionContext context = createContext(selectedAuthScheme, asyncPayload, null);
614+
615+
SdkHttpRequest signedRequest = ValidSdkObjects.sdkHttpFullRequest().build();
616+
617+
SdkHttpContentPublisher contentPublisher =
618+
new SdkHttpContentPublisher() {
619+
@Override
620+
public Optional<Long> contentLength() {
621+
return Optional.of(42L);
622+
}
623+
624+
@Override
625+
public void subscribe(Subscriber<? super ByteBuffer> s) {
626+
s.onSubscribe(new Subscription() {
627+
@Override public void request(long n) { s.onComplete(); }
628+
@Override public void cancel() { }
629+
});
630+
}
631+
};
632+
633+
when(httpSigner.signAsync(ArgumentMatchers.<AsyncSignRequest<? extends Identity>>any()))
634+
.thenReturn(
635+
CompletableFuture.completedFuture(AsyncSignedRequest.builder()
636+
.request(signedRequest)
637+
.payload(contentPublisher)
638+
.build()));
639+
640+
SdkHttpFullRequest request = ValidSdkObjects.sdkHttpFullRequest().build();
641+
stage.execute(request, context).join();
642+
643+
assertThat(context.requestProvider()).isNotNull();
644+
assertThat(context.requestProvider().contentLength()).hasValue(42L);
645+
}
646+
599647
private interface TestAsyncRequestBodySigner extends Signer, AsyncRequestBodySigner {
600648
}
601649

‎core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeAsyncHttpRequestStageTest.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import org.mockito.ArgumentCaptor;
4747
import org.mockito.Mock;
4848
import org.mockito.junit.MockitoJUnitRunner;
49+
import software.amazon.awssdk.core.async.AsyncRequestBody;
4950
import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
5051
import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
5152
import software.amazon.awssdk.core.http.ExecutionContext;
@@ -56,6 +57,7 @@
5657
import software.amazon.awssdk.core.internal.http.TransformingAsyncResponseHandler;
5758
import software.amazon.awssdk.core.internal.http.timers.ClientExecutionAndRequestTimerTestUtils;
5859
import software.amazon.awssdk.core.internal.util.AsyncResponseHandlerTestUtils;
60+
import software.amazon.awssdk.http.Header;
5961
import software.amazon.awssdk.http.SdkHttpFullRequest;
6062
import software.amazon.awssdk.http.SdkHttpMethod;
6163
import software.amazon.awssdk.http.async.AsyncExecuteRequest;
@@ -237,6 +239,48 @@ public void execute_handlerFutureCompletedExceptionally_doesNotAttemptSynchronou
237239
}
238240
}
239241

242+
@Test
243+
public void execute_requestHasTransferEncodingHeader_doesNotAddContentLength() {
244+
stage = new MakeAsyncHttpRequestStage<>(
245+
combinedAsyncResponseHandler(AsyncResponseHandlerTestUtils.noOpResponseHandler(),
246+
AsyncResponseHandlerTestUtils.noOpResponseHandler()),
247+
clientDependencies(null));
248+
249+
SdkHttpFullRequest sdkHttpRequest = SdkHttpFullRequest.builder()
250+
.method(SdkHttpMethod.POST)
251+
.host("service.us-east-1.amazonaws.com")
252+
.protocol("https")
253+
.putHeader(Header.TRANSFER_ENCODING, "chunked")
254+
.build();
255+
256+
AsyncRequestBody requestBody = AsyncRequestBody.fromString("hello world");
257+
258+
ExecutionContext executionContext = ExecutionContext.builder()
259+
.executionAttributes(new ExecutionAttributes())
260+
.build();
261+
262+
RequestExecutionContext context = RequestExecutionContext.builder()
263+
.originalRequest(ValidSdkObjects.sdkRequest())
264+
.executionContext(executionContext)
265+
.requestProvider(requestBody)
266+
.build();
267+
268+
CompletableFuture<SdkHttpFullRequest> requestFuture = CompletableFuture.completedFuture(sdkHttpRequest);
269+
270+
try {
271+
stage.execute(requestFuture, context);
272+
} catch (Exception e) {
273+
// ignored, we only care about the request sent to the HTTP client
274+
} finally {
275+
ArgumentCaptor<AsyncExecuteRequest> httpRequestCaptor = ArgumentCaptor.forClass(AsyncExecuteRequest.class);
276+
verify(sdkAsyncHttpClient).execute(httpRequestCaptor.capture());
277+
278+
SdkHttpFullRequest capturedRequest = (SdkHttpFullRequest) httpRequestCaptor.getValue().request();
279+
assertThat(capturedRequest.firstMatchingHeader(Header.CONTENT_LENGTH)).isNotPresent();
280+
assertThat(capturedRequest.firstMatchingHeader(Header.TRANSFER_ENCODING)).hasValue("chunked");
281+
}
282+
}
283+
240284
private HttpClientDependencies clientDependencies(Duration timeout) {
241285
SdkClientConfiguration configuration = SdkClientConfiguration.builder()
242286
.option(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, Runnable::run)

0 commit comments

Comments
 (0)