Merge pull request #3647 from aws/zoewang/lengthaware-enforce-length-merge

Validate sent body len, remove buffering for known content length cases

Author: zoewangg

.changes/next-release/feature-AWSSDKforJavav2-02883e2.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "The SDK now throws exception for input streaming operation if the stream has fewer bytes (i.e. reaches EOF) before the expected length is reached."
+}

.changes/next-release/feature-AWSSDKforJavav2-9dffbaf.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "The SDK now does not buffer input data from `RequestBody#fromInputStream` in cases where the InputStream does not support mark and reset."
+}

.changes/next-release/feature-AWSSDKforJavav2-ee5927f.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "The SDK now does not buffer input data from ContentStreamProvider in cases where content length is known."
+}

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

@@ -15,23 +15,29 @@
 
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
+import static software.amazon.awssdk.http.Header.CONTENT_LENGTH;
+
 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 +46,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 +72,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 +103,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.debug(() -> 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();
+    }
 }

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,13 @@ 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 IllegalStateException("The request content has fewer bytes than the "
+                                            + "specified "
+                                            + "content-length: " + length + " bytes.");
         }
         return read;
     }
@@ -61,12 +67,20 @@ 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 IllegalStateException("The request content has fewer bytes than the "
+                                            + "specified "
+                                            + "content-length: " + length + " bytes.");
+        }
+
         return read;
     }
 

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);
     }
 
     /**

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;
+    }
 }