Eagerly Read in LengthValidatingInputStream to Fix Flaky JDK HttpClient Test (Azure#31020)

Author: alzimmermsft

sdk/core/azure-core-http-jdk-httpclient/src/main/java/com/azure/core/http/jdk/httpclient/BodyPublisherUtils.java

@@ -13,7 +13,6 @@
 import com.azure.core.util.BinaryData;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.ProgressReporter;
-import com.azure.core.util.logging.ClientLogger;
 import reactor.adapter.JdkFlowAdapter;
 
 import java.net.http.HttpRequest;
@@ -27,8 +26,6 @@
 import static java.net.http.HttpRequest.BodyPublishers.ofString;
 
 final class BodyPublisherUtils {
-    private static final ClientLogger LOGGER = new ClientLogger(BodyPublisherUtils.class);
-
     private BodyPublisherUtils() {
     }
 
@@ -37,7 +34,7 @@ private BodyPublisherUtils() {
      * If progress reporter is not null, configures it to track request body upload.
      *
      * @param request {@link com.azure.core.http.HttpRequest} instance
-     * @progressReporter optional progress reporter.
+     * @param progressReporter optional progress reporter.
      * @return the request BodyPublisher
      */
     public static HttpRequest.BodyPublisher toBodyPublisher(com.azure.core.http.HttpRequest request, ProgressReporter progressReporter) {
@@ -46,7 +43,7 @@ public static HttpRequest.BodyPublisher toBodyPublisher(com.azure.core.http.Http
             return noBody();
         }
 
-        HttpRequest.BodyPublisher publisher = null;
+        HttpRequest.BodyPublisher publisher;
         BinaryDataContent bodyContent = BinaryDataHelper.getContent(body);
         if (bodyContent instanceof ByteArrayContent) {
             publisher = ofByteArray(bodyContent.toBytes());

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

@@ -684,15 +684,12 @@ public void asyncPutRequestWithBinaryDataBodyAndMoreThanContentLength() {
      */
     @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)))
+        BinaryData body = BinaryData.fromStream(new ByteArrayInputStream("test".getBytes(StandardCharsets.UTF_8)));
+        StepVerifier.create(createService(Service9.class).putAsyncBodyAndContentLength(body, 3L))
             .verifyErrorSatisfies(exception -> {
                 assertTrue(exception instanceof UnexpectedLengthException
                     || (exception.getSuppressed().length > 0
-                    && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
+                        && exception.getSuppressed()[0] instanceof UnexpectedLengthException));
                 assertTrue(exception.getMessage().contains("more than"));
             });
     }

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

@@ -46,10 +46,25 @@ final class LengthValidatingInputStream extends InputStream {
 
     @Override
     public synchronized int read(byte[] b, int off, int len) throws IOException {
-        int readSize = inner.read(b, off, len);
-        validateLength(readSize);
+        int totalRead = 0;
+        int readSize;
+        do {
+            // Attempt to read until the byte array is filled or the inner stream ends.
+            // This results in finishing validation faster and prevents scenarios such as with JDK HttpClient where
+            // a large buffer will be requested and it validates read length before we can finish validation.
+            readSize = inner.read(b, off + totalRead, len - totalRead);
+            validateLength(readSize);
+
+            if (readSize != -1) {
+                totalRead += readSize;
+            } else if (totalRead == 0) {
+                // If the inner stream was already read to completion the first read will return -1, need to set
+                // total read to -1 to prevent any infinite read loops by returning 0.
+                totalRead = -1;
+            }
+        } while (readSize != -1 && totalRead != len);
 
-        return readSize;
+        return totalRead;
     }
 
     @Override