Validate sent body len, remove buffering

dagnir · zoewangg · commit 5f2af7768e7f · 2025-02-10T11:29:43.000-08:00
This commit removes the buffering of
`RequestBody.fromContentProvider(ContentStreamProvider,long,String)`,
which was used to ensure that the provider always provided the expected
number of bytes given by `length`.

Instead of doing the buffering, this commit updates
`SdkLengthAwareInputstream` to validate the number of bytes read is
correct, and throws an error if EOF is reached prematurely. We use this
stream to ensure that we sent the correct amount of bytes to the server
when sending the HTTP request.

Note that this a weaker check because we're only testing for the
*number* of bytes sent rather than their contents but has the benefit of
sidestepping buffering.
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeHttpRequestStage.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeHttpRequestStage.java
@@ -16,22 +16,26 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import java.time.Duration;
+import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
 import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
 import software.amazon.awssdk.core.internal.http.HttpClientDependencies;
 import software.amazon.awssdk.core.internal.http.InterruptMonitor;
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.core.internal.io.SdkLengthAwareInputStream;
 import software.amazon.awssdk.core.internal.util.MetricUtils;
 import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.ExecutableHttpRequest;
 import software.amazon.awssdk.http.HttpExecuteRequest;
 import software.amazon.awssdk.http.HttpExecuteResponse;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.metrics.MetricCollector;
+import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Pair;
 
 /**
@@ -40,6 +44,7 @@
 @SdkInternalApi
 public class MakeHttpRequestStage
     implements RequestPipeline<SdkHttpFullRequest, Pair<SdkHttpFullRequest, SdkHttpFullResponse>> {
+    private static final Logger LOG = Logger.loggerFor(MakeHttpRequestStage.class);
 
     private final SdkHttpClient sdkHttpClient;
 
@@ -65,6 +70,8 @@ private HttpExecuteResponse executeHttpRequest(SdkHttpFullRequest request, Reque
 
         MetricCollector httpMetricCollector = MetricUtils.createHttpMetricsCollector(context);
 
+        request = enforceContentLengthIfPresent(request);
+
         ExecutableHttpRequest requestCallable = sdkHttpClient
             .prepareRequest(HttpExecuteRequest.builder()
                                               .request(request)
@@ -94,4 +101,39 @@ private static long updateMetricCollectionAttributes(RequestExecutionContext con
                                                    now);
         return now;
     }
+
+    private static SdkHttpFullRequest enforceContentLengthIfPresent(SdkHttpFullRequest request) {
+        Optional<ContentStreamProvider> requestContentStreamProviderOptional = request.contentStreamProvider();
+
+        if (!requestContentStreamProviderOptional.isPresent()) {
+            return request;
+        }
+
+        Optional<Long> contentLength = contentLength(request);
+        if (!contentLength.isPresent()) {
+            LOG.warn(() -> String.format("Request contains a body but does not have a Content-Length header. Not validating "
+                                         + "the amount of data sent to the service: %s", request));
+            return request;
+        }
+
+        ContentStreamProvider requestContentProvider = requestContentStreamProviderOptional.get();
+        ContentStreamProvider lengthVerifyingProvider = () -> new SdkLengthAwareInputStream(requestContentProvider.newStream(),
+                                                                                            contentLength.get());
+        return request.toBuilder()
+                      .contentStreamProvider(lengthVerifyingProvider)
+                      .build();
+    }
+
+    private static Optional<Long> contentLength(SdkHttpFullRequest request) {
+        Optional<String> contentLengthHeader = request.firstMatchingHeader("Content-Length");
+
+        if (contentLengthHeader.isPresent()) {
+            try {
+                return Optional.of(Long.parseLong(contentLengthHeader.get()));
+            } catch (NumberFormatException e) {
+                LOG.warn(() -> "Unable to parse 'Content-Length' header. Treating it as non existent.");
+            }
+        }
+        return Optional.empty();
+    }
 }
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/SdkLengthAwareInputStream.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/io/SdkLengthAwareInputStream.java
@@ -25,8 +25,9 @@
 import software.amazon.awssdk.utils.Validate;
 
 /**
- * An {@code InputStream} that is aware of its length. The main purpose of this class is to support truncating streams to a
- * length that is shorter than the total length of the stream.
+ * An {@code InputStream} that is aware of its length. This class enforces that we sent exactly the number of bytes equal to
+ * the input length. If the wrapped stream has more bytes than the expected length, it will be truncated to length. If the stream
+ * has less bytes (i.e. reaches EOF) before the expected length is reached, it will throw {@code IOException}.
  */
 @SdkInternalApi
 public class SdkLengthAwareInputStream extends FilterInputStream {
@@ -48,8 +49,11 @@ public int read() throws IOException {
         }
 
         int read = super.read();
+
         if (read != -1) {
             remaining--;
+        } else if (remaining != 0) { // EOF, ensure we've read the number of expected bytes
+            throw new IOException("Reached EOF before reading entire expected content");
         }
         return read;
     }
@@ -61,12 +65,18 @@ public int read(byte[] b, int off, int len) throws IOException {
             return -1;
         }
 
-        len = Math.min(len, saturatedCast(remaining));
-        int read = super.read(b, off, len);
-        if (read > 0) {
+        int readLen = Math.min(len, saturatedCast(remaining));
+
+        int read = super.read(b, off, readLen);
+        if (read != -1) {
             remaining -= read;
         }
 
+        // EOF, ensure we've read the number of expected bytes
+        if (read == -1 && remaining != 0) {
+            throw new IOException("Reached EOF before reading entire expected content");
+        }
+
         return read;
     }
 
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/sync/RequestBody.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/sync/RequestBody.java
@@ -136,21 +136,14 @@ public static RequestBody fromFile(File file) {
      * @return RequestBody instance.
      */
     public static RequestBody fromInputStream(InputStream inputStream, long contentLength) {
-        // NOTE: does not have an effect if mark not supported
         IoUtils.markStreamWithMaxReadLimit(inputStream);
         InputStream nonCloseable = nonCloseableInputStream(inputStream);
-        ContentStreamProvider provider;
-        if (nonCloseable.markSupported()) {
-            // stream supports mark + reset
-            provider = () -> {
+        return fromContentProvider(() -> {
+            if (nonCloseable.markSupported()) {
                 invokeSafely(nonCloseable::reset);
-                return nonCloseable;
-            };
-        } else {
-            // stream doesn't support mark + reset, make sure to buffer it
-            provider = new BufferingContentStreamProvider(() -> nonCloseable, contentLength);
-        }
-        return new RequestBody(provider, contentLength, Mimetype.MIMETYPE_OCTET_STREAM);
+            }
+            return nonCloseable;
+        }, contentLength, Mimetype.MIMETYPE_OCTET_STREAM);
     }
 
     /**
@@ -224,9 +217,6 @@ public static RequestBody empty() {
     /**
      * Creates a {@link RequestBody} from the given {@link ContentStreamProvider}.
      * <p>
-     * Important: Be aware that this implementation requires buffering the contents for {@code ContentStreamProvider}, which can
-     * cause increased memory usage.
-     * <p>
      * If you are using this in conjunction with S3 and want to upload a stream with an unknown content length, you can refer
      * S3's documentation for
      * <a href="https://docs.aws.amazon.com/AmazonS3/latest/API/s3_example_s3_Scenario_UploadStream_section.html">alternative
@@ -239,7 +229,7 @@ public static RequestBody empty() {
      * @return The created {@code RequestBody}.
      */
     public static RequestBody fromContentProvider(ContentStreamProvider provider, long contentLength, String mimeType) {
-        return new RequestBody(new BufferingContentStreamProvider(provider, contentLength), contentLength, mimeType);
+        return new RequestBody(provider, contentLength, mimeType);
     }
 
     /**
diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeHttpRequestStageTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeHttpRequestStageTest.java
@@ -16,42 +16,48 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static software.amazon.awssdk.core.client.config.SdkClientOption.SYNC_HTTP_CLIENT;
+
+import java.io.ByteArrayInputStream;
 import java.io.IOException;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
+import java.io.InputStream;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.BeforeEach;
+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 org.mockito.ArgumentCaptor;
-import org.mockito.Mock;
-import org.mockito.junit.MockitoJUnitRunner;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.http.ExecutionContext;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.internal.http.HttpClientDependencies;
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.internal.http.timers.TimeoutTracker;
+import software.amazon.awssdk.core.internal.io.SdkLengthAwareInputStream;
+import software.amazon.awssdk.http.ContentStreamProvider;
 import software.amazon.awssdk.http.HttpExecuteRequest;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpMethod;
 import software.amazon.awssdk.metrics.MetricCollector;
 import utils.ValidSdkObjects;
 
-@RunWith(MockitoJUnitRunner.class)
 public class MakeHttpRequestStageTest {
 
-    @Mock
     private SdkHttpClient mockClient;
 
     private MakeHttpRequestStage stage;
 
-    @Before
+    @BeforeEach
     public void setup() throws IOException {
+        mockClient = mock(SdkHttpClient.class);
         SdkClientConfiguration config = SdkClientConfiguration.builder().option(SYNC_HTTP_CLIENT, mockClient).build();
         stage = new MakeHttpRequestStage(HttpClientDependencies.builder().clientConfiguration(config).build());
     }
@@ -94,4 +100,92 @@ public void testExecute_contextContainsMetricCollector_addsChildToExecuteRequest
             assertThat(httpRequestCaptor.getValue().metricCollector()).contains(childCollector);
         }
     }
+
+    @ParameterizedTest
+    @MethodSource("contentLengthVerificationInputs")
+    public void execute_testLengthChecking(String description,
+                                           ContentStreamProvider provider,
+                                           Long contentLength,
+                                           boolean expectLengthAware) {
+        SdkHttpFullRequest.Builder requestBuilder = SdkHttpFullRequest.builder()
+                                                                      .method(SdkHttpMethod.PUT)
+                                                          .host("mybucket.s3.us-west-2.amazonaws.com")
+                                                          .protocol("https");
+
+        if (provider != null) {
+            requestBuilder.contentStreamProvider(provider);
+        }
+
+        if (contentLength != null) {
+            requestBuilder.putHeader("Content-Length", String.valueOf(contentLength));
+        }
+
+        when(mockClient.prepareRequest(any()))
+            .thenThrow(new RuntimeException("BOOM"));
+
+        assertThatThrownBy(() -> stage.execute(requestBuilder.build(), createContext())).hasMessage("BOOM");
+
+        ArgumentCaptor<HttpExecuteRequest> requestCaptor = ArgumentCaptor.forClass(HttpExecuteRequest.class);
+
+        verify(mockClient).prepareRequest(requestCaptor.capture());
+
+        HttpExecuteRequest capturedRequest = requestCaptor.getValue();
+
+        if (provider != null) {
+            InputStream requestContentStream = capturedRequest.contentStreamProvider().get().newStream();
+
+            if (expectLengthAware) {
+                assertThat(requestContentStream).isInstanceOf(SdkLengthAwareInputStream.class);
+            } else {
+                assertThat(requestContentStream).isNotInstanceOf(SdkLengthAwareInputStream.class);
+            }
+        } else {
+            assertThat(capturedRequest.contentStreamProvider()).isEmpty();
+        }
+    }
+
+    private static Stream<Arguments> contentLengthVerificationInputs() {
+        return Stream.of(
+            Arguments.of(
+                "Provider present, ContentLength present",
+                (ContentStreamProvider) () -> new ByteArrayInputStream(new byte[16]),
+                16L,
+                true
+            ),
+            Arguments.of(
+                "Provider present, ContentLength not present",
+                (ContentStreamProvider) () -> new ByteArrayInputStream(new byte[16]),
+                null,
+                false
+            ),
+            Arguments.of(
+                "Provider not present, ContentLength present",
+                null,
+                16L,
+                false
+            ),
+            Arguments.of(
+                "Provider not present, ContentLength not present",
+                null,
+                null,
+                false
+            )
+        );
+    }
+
+    private static RequestExecutionContext createContext() {
+        ExecutionContext executionContext = ExecutionContext.builder()
+                                                            .executionAttributes(new ExecutionAttributes())
+                                                            .build();
+
+        RequestExecutionContext context = RequestExecutionContext.builder()
+                                                                 .originalRequest(ValidSdkObjects.sdkRequest())
+                                                                 .executionContext(executionContext)
+                                                                 .build();
+
+        context.apiCallAttemptTimeoutTracker(mock(TimeoutTracker.class));
+        context.apiCallTimeoutTracker(mock(TimeoutTracker.class));
+
+        return context;
+    }
 }
diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/io/SdkLengthAwareInputStreamTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/io/SdkLengthAwareInputStreamTest.java
diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/sync/RequestBodyTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/sync/RequestBodyTest.java
