Pass binary data through rest proxy. (#29084)

* wip

* more tests.

* pass binary data.

* wip

* wip.

* more testing.

* stream v1

* tests.

* more testing.

* more testing.

* modify this test.

* check if http clients hit eof

* static logger.

Author: kasobol-msft

sdk/core/azure-core-test/src/main/java/com/azure/core/test/implementation/RestProxyTests.java

@@ -49,6 +49,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.io.ByteArrayInputStream;
 import java.nio.ByteBuffer;
 import java.nio.channels.AsynchronousFileChannel;
 import java.nio.charset.StandardCharsets;
@@ -446,6 +447,13 @@ HttpBinJSON putBodyAndContentLength(@BodyParam(ContentType.APPLICATION_OCTET_STR
         Mono<HttpBinJSON> putAsyncBodyAndContentLength(@BodyParam(ContentType.APPLICATION_OCTET_STREAM) Flux<ByteBuffer> body,
             @HeaderParam("Content-Length") long contentLength);
 
+        @Put("put")
+        @ExpectedResponses({200})
+        @UnexpectedResponseExceptionType(MyRestException.class)
+        Mono<HttpBinJSON> putAsyncBodyAndContentLength(
+            @BodyParam(ContentType.APPLICATION_OCTET_STREAM) BinaryData body,
+            @HeaderParam("Content-Length") long contentLength);
+
         @Put("put")
         @ExpectedResponses({201})
         HttpBinJSON putWithUnexpectedResponse(@BodyParam(ContentType.APPLICATION_OCTET_STREAM) String putBody);
@@ -592,6 +600,88 @@ public void asyncPutRequestWithBodyAndMoreThanContentLength() {
             });
     }
 
+    @Test
+    public void asyncPutRequestWithBinaryDataBodyAndEqualContentLength() {
+        Mono<BinaryData> bodyMono = BinaryData.fromFlux(
+            Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8))));
+        StepVerifier.create(
+                bodyMono.flatMap(body ->
+                    createService(Service9.class).putAsyncBodyAndContentLength(body, 4L)))
+            .assertNext(json -> {
+                assertEquals("test", json.data());
+                assertEquals(ContentType.APPLICATION_OCTET_STREAM, json.getHeaderValue("Content-Type"));
+                assertEquals("4", json.getHeaderValue("Content-Length"));
+            }).verifyComplete();
+    }
+
+    @Test
+    public void asyncPutRequestWithBinaryDataBodyAndLessThanContentLength() {
+        Mono<BinaryData> bodyMono = BinaryData.fromFlux(
+            Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8))));
+        StepVerifier.create(
+                bodyMono.flatMap(body ->
+                    createService(Service9.class).putAsyncBodyAndContentLength(body, 5L)))
+            .verifyErrorSatisfies(exception -> {
+                assertTrue(exception instanceof UnexpectedLengthException
+                    || (exception.getSuppressed().length > 0
+                    && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
+                assertTrue(exception.getMessage().contains("less than"));
+            });
+    }
+
+    /**
+     * LengthValidatingInputStream in rest proxy relies on reader
+     * reaching EOF. This test specifically targets InputStream to assert this behavior.
+     */
+    @Test
+    public void asyncPutRequestWithStreamBinaryDataBodyAndLessThanContentLength() {
+        Mono<BinaryData> bodyMono = Mono.just(BinaryData.fromStream(
+            new ByteArrayInputStream("test".getBytes(StandardCharsets.UTF_8))));
+        StepVerifier.create(
+                bodyMono.flatMap(body ->
+                    createService(Service9.class).putAsyncBodyAndContentLength(body, 5L)))
+            .verifyErrorSatisfies(exception -> {
+                assertTrue(exception instanceof UnexpectedLengthException
+                    || (exception.getSuppressed().length > 0
+                    && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
+                assertTrue(exception.getMessage().contains("less than"));
+            });
+    }
+
+    @Test
+    public void asyncPutRequestWithBinaryDataBodyAndMoreThanContentLength() {
+        Mono<BinaryData> bodyMono = BinaryData.fromFlux(
+            Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8))));
+        StepVerifier.create(
+                bodyMono.flatMap(body ->
+                    createService(Service9.class).putAsyncBodyAndContentLength(body, 3L)))
+            .verifyErrorSatisfies(exception -> {
+                assertTrue(exception instanceof UnexpectedLengthException
+                    || (exception.getSuppressed().length > 0
+                    && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
+                assertTrue(exception.getMessage().contains("more than"));
+            });
+    }
+
+    /**
+     * LengthValidatingInputStream in rest proxy relies on reader
+     * reaching EOF. This test specifically targets InputStream to assert this behavior.
+     */
+    @Test
+    public void asyncPutRequestWithStreamBinaryDataBodyAndMoreThanContentLength() {
+        Mono<BinaryData> bodyMono = Mono.just(BinaryData.fromStream(
+            new ByteArrayInputStream("test".getBytes(StandardCharsets.UTF_8))));
+        StepVerifier.create(
+                bodyMono.flatMap(body ->
+                    createService(Service9.class).putAsyncBodyAndContentLength(body, 3L)))
+            .verifyErrorSatisfies(exception -> {
+                assertTrue(exception instanceof UnexpectedLengthException
+                    || (exception.getSuppressed().length > 0
+                    && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
+                assertTrue(exception.getMessage().contains("more than"));
+            });
+    }
+
     @Test
     public void syncPutRequestWithUnexpectedResponse() {
         HttpResponseException e = assertThrows(HttpResponseException.class,
@@ -1593,6 +1683,36 @@ public void segmentUploadTest() throws Exception {
         assertEquals("quick brown fox", response.getValue().data());
     }
 
+    @Host("http://localhost")
+    @ServiceInterface(name = "FluxUploadService")
+    interface BinaryDataUploadService {
+        @Put("/put")
+        Response<HttpBinJSON> put(@BodyParam("text/plain") BinaryData content,
+                                  @HeaderParam("Content-Length") long contentLength);
+    }
+
+    @Test
+    public void binaryDataUploadTest() throws Exception {
+        Path filePath = Paths.get(getClass().getClassLoader().getResource("upload.txt").toURI());
+        BinaryData data = BinaryData.fromFile(filePath);
+
+        final HttpClient httpClient = createHttpClient();
+        // Scenario: Log the body so that body buffering/replay behavior is exercised.
+        //
+        // Order in which policies applied will be the order in which they added to builder
+        //
+        final HttpPipeline httpPipeline = new HttpPipelineBuilder()
+            .httpClient(httpClient)
+            .policies(new PortPolicy(getWireMockPort(), true),
+                new HttpLoggingPolicy(new HttpLogOptions().setLogLevel(HttpLogDetailLevel.BODY_AND_HEADERS)))
+            .build();
+        //
+        Response<HttpBinJSON> response = RestProxy
+            .create(BinaryDataUploadService.class, httpPipeline).put(data, Files.size(filePath));
+
+        assertEquals("The quick brown fox jumps over the lazy dog", response.getValue().data());
+    }
+
     @Host("{url}")
     @ServiceInterface(name = "Service22")
     interface Service22 {

sdk/core/azure-core/pom.xml

@@ -49,6 +49,7 @@
     <jacoco.min.branchcoverage>0.60</jacoco.min.branchcoverage>
     <javaModulesSurefireArgLine>
       --add-exports com.azure.core/com.azure.core.implementation.http=ALL-UNNAMED
+      --add-exports com.azure.core/com.azure.core.implementation.http.rest=ALL-UNNAMED
       --add-exports com.azure.core/com.azure.core.implementation.serializer=ALL-UNNAMED
       --add-exports com.azure.core/com.azure.core.implementation.jackson=ALL-UNNAMED
 

sdk/core/azure-core/src/main/java/com/azure/core/http/rest/RequestOptions.java

@@ -237,7 +237,7 @@ public RequestOptions addRequestCallback(Consumer<HttpRequest> requestCallback)
      */
     public RequestOptions setBody(BinaryData requestBody) {
         Objects.requireNonNull(requestBody, "'requestBody' cannot be null.");
-        this.requestCallback = this.requestCallback.andThen(request -> request.setBody(requestBody.toBytes()));
+        this.requestCallback = this.requestCallback.andThen(request -> request.setBody(requestBody));
         return this;
     }
 

sdk/core/azure-core/src/main/java/com/azure/core/http/rest/RestProxy.java

@@ -10,7 +10,6 @@
 import com.azure.core.exception.ResourceModifiedException;
 import com.azure.core.exception.ResourceNotFoundException;
 import com.azure.core.exception.TooManyRedirectsException;
-import com.azure.core.exception.UnexpectedLengthException;
 import com.azure.core.http.ContentType;
 import com.azure.core.http.HttpHeaders;
 import com.azure.core.http.HttpMethod;
@@ -23,6 +22,7 @@
 import com.azure.core.http.policy.UserAgentPolicy;
 import com.azure.core.implementation.TypeUtil;
 import com.azure.core.implementation.http.UnexpectedExceptionInformation;
+import com.azure.core.implementation.http.rest.RestProxyUtils;
 import com.azure.core.implementation.serializer.HttpResponseDecoder;
 import com.azure.core.implementation.serializer.HttpResponseDecoder.HttpDecodedResponse;
 import com.azure.core.util.Base64Url;
@@ -63,9 +63,6 @@
  * as asynchronous Single objects that resolve to a deserialized Java object.
  */
 public final class RestProxy implements InvocationHandler {
-    private static final ByteBuffer VALIDATION_BUFFER = ByteBuffer.allocate(0);
-    private static final String BODY_TOO_LARGE = "Request body emitted %d bytes, more than the expected %d bytes.";
-    private static final String BODY_TOO_SMALL = "Request body emitted %d bytes, less than the expected %d bytes.";
     private static final String MUST_IMPLEMENT_PAGE_ERROR =
         "Unable to create PagedResponse<T>. Body must be of a type that implements: " + Page.class;
 
@@ -140,11 +137,9 @@ public Object invoke(Object proxy, final Method method, Object[] args) {
                 options.getRequestCallback().accept(request);
             }
 
-            if (request.getBody() != null) {
-                request.setBody(validateLength(request));
-            }
-
-            final Mono<HttpResponse> asyncResponse = send(request, context);
+            Context finalContext = context;
+            final Mono<HttpResponse> asyncResponse = RestProxyUtils.validateLengthAsync(request)
+                .flatMap(r -> send(r, finalContext));
 
             Mono<HttpDecodedResponse> asyncDecodedResponse = this.decoder.decode(asyncResponse, methodParser);
 
@@ -177,43 +172,6 @@ static Context mergeRequestOptionsContext(Context context, RequestOptions option
         return context;
     }
 
-    static Flux<ByteBuffer> validateLength(final HttpRequest request) {
-        final Flux<ByteBuffer> bbFlux = request.getBody();
-        if (bbFlux == null) {
-            return Flux.empty();
-        }
-
-        final long expectedLength = Long.parseLong(request.getHeaders().getValue("Content-Length"));
-
-        return Flux.defer(() -> {
-            final long[] currentTotalLength = new long[1];
-            return Flux.concat(bbFlux, Flux.just(VALIDATION_BUFFER)).handle((buffer, sink) -> {
-                if (buffer == null) {
-                    return;
-                }
-
-                if (buffer == VALIDATION_BUFFER) {
-                    if (expectedLength != currentTotalLength[0]) {
-                        sink.error(new UnexpectedLengthException(String.format(BODY_TOO_SMALL,
-                            currentTotalLength[0], expectedLength), currentTotalLength[0], expectedLength));
-                    } else {
-                        sink.complete();
-                    }
-                    return;
-                }
-
-                currentTotalLength[0] += buffer.remaining();
-                if (currentTotalLength[0] > expectedLength) {
-                    sink.error(new UnexpectedLengthException(String.format(BODY_TOO_LARGE,
-                        currentTotalLength[0], expectedLength), currentTotalLength[0], expectedLength));
-                    return;
-                }
-
-                sink.next(buffer);
-            });
-        });
-    }
-
     /**
      * Starts the tracing span for the current service call, additionally set metadata attributes on the span by passing
      * additional context information.
@@ -321,7 +279,7 @@ private HttpRequest configRequest(final HttpRequest request, final SwaggerMethod
                 // sending the request to the service. There is no memory copy that happens here. Sources like
                 // InputStream, File and Flux<ByteBuffer> will not be eagerly copied into memory until it's required
                 // by the HttpClient implementations.
-                request.setBody(binaryData.toFluxByteBuffer());
+                request.setBody(binaryData);
                 return request;
             }
 

sdk/core/azure-core/src/main/java/com/azure/core/implementation/http/rest/LengthValidatingInputStream.java

@@ -0,0 +1,111 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+package com.azure.core.implementation.http.rest;
+
+import com.azure.core.exception.UnexpectedLengthException;
+import com.azure.core.util.logging.ClientLogger;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Objects;
+
+import static com.azure.core.implementation.http.rest.RestProxyUtils.BODY_TOO_LARGE;
+import static com.azure.core.implementation.http.rest.RestProxyUtils.BODY_TOO_SMALL;
+
+/**
+ * An {@link InputStream} decorator that tracks the number of bytes read from an inner {@link InputStream} and throws
+ * an exception if the number of bytes read doesn't match what was expected.
+ *
+ * This implementation assumes that reader is going to read until EOF.
+ */
+final class LengthValidatingInputStream extends InputStream {
+
+    private static final ClientLogger LOGGER = new ClientLogger(LengthValidatingInputStream.class);
+
+    private final InputStream inner;
+    private final long expectedReadSize;
+
+    private long position;
+    private long mark = -1;
+
+    /**
+     * Creates a new {@link LengthValidatingInputStream}.
+     *
+     * @param inputStream The {@link InputStream} being decorated.
+     * @param expectedReadSize The expected number of bytes to be read from the inner {@code inputStream}.
+     */
+    LengthValidatingInputStream(InputStream inputStream, long expectedReadSize) {
+        this.inner = Objects.requireNonNull(inputStream, "'inputStream' cannot be null.");
+
+        if (expectedReadSize < 0) {
+            throw LOGGER.logExceptionAsError(new IllegalArgumentException("'expectedReadSize' cannot be less than 0."));
+        }
+
+        this.expectedReadSize = expectedReadSize;
+    }
+
+    @Override
+    public synchronized int read(byte[] b, int off, int len) throws IOException {
+        int readSize = inner.read(b, off, len);
+        validateLength(readSize);
+
+        return readSize;
+    }
+
+    @Override
+    public synchronized long skip(long n) throws IOException {
+        long skipped = inner.skip(n);
+        position += skipped;
+        return skipped;
+    }
+
+    @Override
+    public int available() throws IOException {
+        return inner.available();
+    }
+
+    @Override
+    public void close() throws IOException {
+        inner.close();
+    }
+
+    @Override
+    public synchronized void mark(int readlimit) {
+        inner.mark(readlimit);
+        mark = position;
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+        inner.reset();
+        position = mark;
+    }
+
+    @Override
+    public boolean markSupported() {
+        return inner.markSupported();
+    }
+
+    @Override
+    public synchronized int read() throws IOException {
+        int read = inner.read();
+        validateLength(read == -1 ? -1 : 1);
+
+        return read;
+    }
+
+    private void validateLength(int readSize) {
+        if (readSize == -1) {
+            // If the inner InputStream has reached termination validate that the read bytes matches what was expected.
+            if (position > expectedReadSize) {
+                throw new UnexpectedLengthException(String.format(BODY_TOO_LARGE,
+                    position, expectedReadSize), position, expectedReadSize);
+            } else if (position < expectedReadSize) {
+                throw new UnexpectedLengthException(String.format(BODY_TOO_SMALL,
+                    position, expectedReadSize), position, expectedReadSize);
+            }
+        } else {
+            position += readSize;
+        }
+    }
+}