Implement unexpectlengthexception in azure core

eng/code-quality-reports/src/main/resources/checkstyle/checkstyle-suppressions.xml

@@ -155,6 +155,9 @@
 
   <!-- ClientLogger class suppression -->
   <suppress checks="com.azure.tools.checkstyle.checks.GoodLoggingCheck" files="ClientLogger.java"/>
+  <!-- Use the logger in Utility static method. -->
+  <suppress checks="com.azure.tools.checkstyle.checks.GoodLoggingCheck" files="com.azure.storage.common.Utility.java"/>
+
   <!-- Azure Core Test -->
   <suppress checks="com.azure.tools.checkstyle.checks.GoodLoggingCheck" files=".*[/\\]core[/\\]test[/\\].*\.java"/>
   <!-- Class has static methods which using static logger instance, issue link: https://github.com/Azure/azure-sdk-for-java/issues/5137-->

sdk/core/azure-core-test/src/main/java/com/azure/core/test/models/NetworkCallError.java

@@ -3,6 +3,7 @@
 
 package com.azure.core.test.models;
 
+import com.azure.core.exception.UnexpectedLengthException;
 import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.net.UnknownHostException;
@@ -12,12 +13,14 @@
  * during the pipeline and deserialize them back into their actual throwable class when running in playback mode.
  */
 public class NetworkCallError {
-    @JsonProperty("Throwable")
-    private Throwable throwable;
-
     @JsonProperty("ClassName")
     private String className;
 
+    @JsonProperty("ErrorMessage")
+    private String errorMessage;
+
+    private Throwable throwable;
+
     /**
      * Empty constructor used by deserialization.
      */
@@ -32,6 +35,7 @@ public NetworkCallError() {
     public NetworkCallError(Throwable throwable) {
         this.throwable = throwable;
         this.className = throwable.getClass().getName();
+        this.errorMessage = throwable.getMessage();
     }
 
     /**
@@ -40,13 +44,16 @@ public NetworkCallError(Throwable throwable) {
     public Throwable get() {
         switch (className) {
             case "java.lang.NullPointerException":
-                return new NullPointerException(throwable.getMessage());
+                return new NullPointerException(this.errorMessage);
 
             case "java.lang.IndexOutOfBoundsException":
-                return new IndexOutOfBoundsException(throwable.getMessage());
+                return new IndexOutOfBoundsException(this.errorMessage);
 
             case "java.net.UnknownHostException":
-                return new UnknownHostException(throwable.getMessage());
+                return new UnknownHostException(this.errorMessage);
+
+            case "com.azure.core.exception.UnexpectedLengthException":
+                return new UnexpectedLengthException(this.errorMessage, 0L, 0L);
 
             default:
                 return throwable;
@@ -71,4 +78,14 @@ public void throwable(Throwable throwable) {
     public void className(String className) {
         this.className = className;
     }
+
+    /**
+     * Sets the error message of the class of the throwable. This is used during deserialization the construct the
+     * throwable as the actual class that was thrown.
+     *
+     * @param errorMessage Error msg from the exception.
+     */
+    public void errorMessage(String errorMessage) {
+        this.errorMessage = errorMessage;
+    }
 }

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/UnexpectedLengthException.java → sdk/core/azure-core/src/main/java/com/azure/core/exception/UnexpectedLengthException.java

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-package com.azure.storage.blob;
+package com.azure.core.exception;
 
 /**
  * This exception class represents an error when the specified input length doesn't match the data length.
@@ -10,7 +10,13 @@ public final class UnexpectedLengthException extends IllegalStateException {
     private final long bytesRead;
     private final long bytesExpected;
 
-    UnexpectedLengthException(String message, long bytesRead, long bytesExpected) {
+    /**
+     * Constructor of the UnexpectedLengthException.
+     * @param message The message for the exception.
+     * @param bytesRead The number of bytes read from resource.
+     * @param bytesExpected The number of bytes expected from the receiver.
+     */
+    public UnexpectedLengthException(String message, long bytesRead, long bytesExpected) {
         super(message);
         this.bytesRead = bytesRead;
         this.bytesExpected = bytesExpected;

sdk/core/azure-core/src/main/java/com/azure/core/implementation/RestProxy.java

@@ -22,6 +22,7 @@
 import com.azure.core.http.rest.PagedResponse;
 import com.azure.core.http.rest.Response;
 import com.azure.core.http.rest.ResponseBase;
+import com.azure.core.exception.UnexpectedLengthException;
 import com.azure.core.implementation.http.ContentType;
 import com.azure.core.implementation.http.PagedResponseBase;
 import com.azure.core.implementation.http.UrlBuilder;
@@ -142,10 +143,14 @@ public Object invoke(Object proxy, final Method method, Object[] args) {
                 Context context = methodParser.context(args).addData("caller-method", methodParser.fullyQualifiedMethodName());
                 context = startTracingSpan(method, context);
 
+                if (request.body() != null) {
+                    request.body(validateLength(request));
+                }
+
                 final Mono<HttpResponse> asyncResponse = send(request, context);
-                //
+
                 Mono<HttpDecodedResponse> asyncDecodedResponse = this.decoder.decode(asyncResponse, methodParser);
-                //
+
                 return handleHttpResponse(request, asyncDecodedResponse, methodParser, methodParser.returnType(), context);
             }
 
@@ -154,6 +159,36 @@ public Object invoke(Object proxy, final Method method, Object[] args) {
         }
     }
 
+    private Flux<ByteBuffer> validateLength(final HttpRequest request) {
+        final Flux<ByteBuffer> bbFlux = request.body();
+        if (bbFlux == null) {
+            return Flux.empty();
+        }
+
+        return Flux.defer(() -> {
+            Long expectedLength = Long.valueOf(request.headers().value("Content-Length"));
+            final long[] currentTotalLength = new long[1];
+            return bbFlux.doOnEach(s -> {
+                if (s.isOnNext()) {
+                    ByteBuffer byteBuffer = s.get();
+                    int currentLength = (byteBuffer == null) ? 0 : byteBuffer.remaining();
+                    currentTotalLength[0] += currentLength;
+                    if (currentTotalLength[0] > expectedLength) {
+                        throw logger.logExceptionAsError(new UnexpectedLengthException(
+                            String.format("Request body emitted %d bytes more than the expected %d bytes.",
+                                currentTotalLength[0], expectedLength), currentTotalLength[0], expectedLength));
+                    }
+                } else if (s.isOnComplete()) {
+                    if (expectedLength.compareTo(currentTotalLength[0]) != 0) {
+                        throw logger.logExceptionAsError(new UnexpectedLengthException(
+                            String.format("Request body emitted %d bytes less than the expected %d bytes.",
+                                currentTotalLength[0], expectedLength), currentTotalLength[0], expectedLength));
+                    }
+                }
+            });
+        });
+    }
+
     private Method determineResumeMethod(Method method, String resumeMethodName) {
         for (Method potentialResumeMethod : method.getDeclaringClass().getMethods()) {
             if (potentialResumeMethod.getName().equals(resumeMethodName)) {
@@ -290,6 +325,8 @@ private HttpRequest configRequest(HttpRequest request, SwaggerMethodParser metho
                 if (!bodyContentString.isEmpty()) {
                     request.body(bodyContentString);
                 }
+            } else if (bodyContentObject instanceof ByteBuffer) {
+                request.body(Flux.just((ByteBuffer) bodyContentObject));
             } else {
                 final String bodyContentString = serializer.serialize(bodyContentObject, SerializerEncoding.fromHeaders(request.headers()));
                 request.body(bodyContentString);

sdk/core/azure-core/src/main/java/com/azure/core/implementation/util/FluxUtil.java

@@ -78,7 +78,7 @@ public static byte[] byteBufferToArray(ByteBuffer byteBuffer) {
         byteBuffer.get(byteArray);
         return byteArray;
     }
-
+    
     /**
      * This method converts the incoming {@code subscriberContext} from {@link reactor.util.context.Context Reactor
      * Context} to {@link Context Azure Context} and calls the given lambda function with this context and returns a

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

@@ -36,13 +36,17 @@
 import com.azure.core.http.rest.ResponseBase;
 import com.azure.core.http.rest.StreamResponse;
 import com.azure.core.http.rest.VoidResponse;
+import com.azure.core.exception.UnexpectedLengthException;
 import com.azure.core.implementation.http.ContentType;
 import com.azure.core.implementation.serializer.SerializerAdapter;
 import com.azure.core.implementation.serializer.jackson.JacksonAdapter;
 import com.azure.core.implementation.util.FluxUtil;
+import java.nio.charset.StandardCharsets;
 import org.junit.Assert;
 import org.junit.Ignore;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -57,6 +61,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import reactor.test.StepVerifier;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
@@ -74,6 +79,9 @@ public abstract class RestProxyTests {
      */
     protected abstract HttpClient createHttpClient();
 
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+
     @Host("http://httpbin.org")
     @ServiceInterface(name = "Service1")
     private interface Service1 {
@@ -430,6 +438,16 @@ private interface Service9 {
         @ExpectedResponses({200})
         Mono<HttpBinJSON> putAsync(@BodyParam(ContentType.APPLICATION_OCTET_STREAM) int putBody);
 
+        @Put("put")
+        @ExpectedResponses({200})
+        @UnexpectedResponseExceptionType(MyRestException.class)
+        HttpBinJSON putBodyAndContentLength(@BodyParam("application/octet-stream") ByteBuffer body, @HeaderParam("Content-Length") long contentLength);
+
+        @Put("put")
+        @ExpectedResponses({200})
+        @UnexpectedResponseExceptionType(MyRestException.class)
+        Mono<HttpBinJSON> putAsyncBodyAndContentLength(@BodyParam("application/octet-stream") Flux<ByteBuffer> body, @HeaderParam("Content-Length") long contentLength);
+
         @Put("put")
         @ExpectedResponses({201})
         HttpBinJSON putWithUnexpectedResponse(@BodyParam(ContentType.APPLICATION_OCTET_STREAM) String putBody);
@@ -500,6 +518,79 @@ public void asyncPutRequestWithIntBody() {
         assertEquals("42", json.data());
     }
 
+    // Test all scenarios for the body length and content length comparison for sync API
+    @Test
+    public void syncPutRequestWithBodyAndEqualContentLength() {
+        ByteBuffer body = ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8));
+        final HttpBinJSON json = createService(Service9.class)
+            .putBodyAndContentLength(body, 4L);
+        assertEquals("test", json.data());
+        assertEquals("application/octet-stream", json.headers().get(("Content-Type")));
+        assertEquals("4", json.headers().get(("Content-Length")));
+        body.clear();
+    }
+
+    @Test
+    public void syncPutRequestWithBodyLessThanContentLength() {
+        ByteBuffer body = ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8));
+        thrown.expect(UnexpectedLengthException.class);
+        thrown.expectMessage("less than");
+        createService(Service9.class)
+            .putBodyAndContentLength(body, 5L);
+        body.clear();
+    }
+
+    @Test
+    public void syncPutRequestWithBodyMoreThanContentLength() {
+        ByteBuffer body = ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8));
+        thrown.expect(UnexpectedLengthException.class);
+        thrown.expectMessage("more than");
+        createService(Service9.class)
+            .putBodyAndContentLength(body, 3L);
+        body.clear();
+    }
+
+    // Test all scenarios for the body length and content length comparison for Async API
+    @Test
+    public void asyncPutRequestWithBodyAndEqualContentLength() {
+        Flux<ByteBuffer> body = Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8)));
+        StepVerifier.create(createService(Service9.class)
+            .putAsyncBodyAndContentLength(body, 4L))
+            .assertNext(
+                json -> {
+                    assertEquals("test", json.data());
+                    assertEquals("application/octet-stream", json.headers().get(("Content-Type")));
+                    assertEquals("4", json.headers().get(("Content-Length")));
+                }
+            ).verifyComplete();
+    }
+
+    @Test
+    public void asyncPutRequestWithBodyAndLessThanContentLength() {
+        Flux<ByteBuffer> body = Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8)));
+        StepVerifier.create(createService(Service9.class)
+            .putAsyncBodyAndContentLength(body, 5L))
+            .verifyErrorSatisfies(
+                exception -> {
+                    assertTrue(exception instanceof UnexpectedLengthException);
+                    assertTrue(exception.getMessage().contains("less than"));
+                }
+            );
+    }
+
+    @Test
+    public void asyncPutRequestWithBodyAndMoreThanContentLength() {
+        Flux<ByteBuffer> body = Flux.just(ByteBuffer.wrap("test".getBytes(StandardCharsets.UTF_8)));
+        StepVerifier.create(createService(Service9.class)
+            .putAsyncBodyAndContentLength(body, 3L))
+            .verifyErrorSatisfies(
+                exception -> {
+                    assertTrue(exception instanceof UnexpectedLengthException);
+                    assertTrue(exception.getMessage().contains("more than"));
+                }
+            );
+    }
+
     @Test
     public void syncPutRequestWithUnexpectedResponse() {
         try {

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/AppendBlobClient.java

@@ -4,6 +4,7 @@
 package com.azure.storage.blob;
 
 import com.azure.core.http.rest.Response;
+import com.azure.core.exception.UnexpectedLengthException;
 import com.azure.core.util.Context;
 import com.azure.storage.blob.models.AppendBlobAccessConditions;
 import com.azure.storage.blob.models.AppendBlobItem;
@@ -14,6 +15,7 @@
 import com.azure.storage.blob.models.SourceModifiedAccessConditions;
 import com.azure.storage.blob.models.StorageException;
 import com.azure.storage.common.Utility;
+import java.util.Objects;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
@@ -159,24 +161,14 @@ public AppendBlobItem appendBlock(InputStream data, long length) {
      * @param context Additional context that is passed through the Http pipeline during the service call.
      *
      * @return A {@link Response} whose {@link Response#value() value} contains the append blob operation.
+     * @throws UnexpectedLengthException when the length of data does not match the input {@code length}.
+     * @throws NullPointerException if the input data is null.
      */
     public Response<AppendBlobItem> appendBlockWithResponse(InputStream data, long length,
         AppendBlobAccessConditions appendBlobAccessConditions, Duration timeout, Context context) {
-        Flux<ByteBuffer> fbb = Flux.range(0, (int) Math.ceil((double) length / (double) MAX_APPEND_BLOCK_BYTES))
-            .map(i -> i * MAX_APPEND_BLOCK_BYTES)
-            .concatMap(pos -> Mono.fromCallable(() -> {
-                long count = pos + MAX_APPEND_BLOCK_BYTES > length ? length - pos : MAX_APPEND_BLOCK_BYTES;
-                byte[] cache = new byte[(int) count];
-                int read = 0;
-                while (read < count) {
-                    read += data.read(cache, read, (int) count - read);
-                }
-
-                return ByteBuffer.wrap(cache);
-            }));
-
-        Mono<Response<AppendBlobItem>> response = appendBlobAsyncClient.appendBlockWithResponse(
-            fbb.subscribeOn(Schedulers.elastic()), length, appendBlobAccessConditions, context);
+        Objects.requireNonNull(data);
+        Flux<ByteBuffer> fbb = Utility.convertStreamToByteBuffer(data, length, MAX_APPEND_BLOCK_BYTES);
+        Mono<Response<AppendBlobItem>> response = appendBlobAsyncClient.appendBlockWithResponse(fbb.subscribeOn(Schedulers.elastic()), length, appendBlobAccessConditions, context);
         return Utility.blockWithOptionalTimeout(response, timeout);
     }