content length override changes

Author: Local Merge

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/implementation/models/BlobsDownloadHeaders.java

@@ -301,6 +301,11 @@ public final class BlobsDownloadHeaders {
     @Generated
     private byte[] xMsContentCrc64;
 
+    /*
+     * The x-ms-original-content-length property.
+     */
+    private Long xMsOriginalContentLength;
+
     private static final HttpHeaderName X_MS_CREATION_TIME = HttpHeaderName.fromString("x-ms-creation-time");
 
     private static final HttpHeaderName X_MS_OR_POLICY_ID = HttpHeaderName.fromString("x-ms-or-policy-id");
@@ -369,6 +374,9 @@ public final class BlobsDownloadHeaders {
 
     private static final HttpHeaderName X_MS_CONTENT_CRC64 = HttpHeaderName.fromString("x-ms-content-crc64");
 
+    private static final HttpHeaderName X_MS_ORIGINAL_CONTENT_LENGTH
+        = HttpHeaderName.fromString("x-ms-original-content-length");
+
     // HttpHeaders containing the raw property values.
     /**
      * Creates an instance of BlobsDownloadHeaders class.
@@ -535,6 +543,12 @@ public BlobsDownloadHeaders(HttpHeaders rawHeaders) {
         } else {
             this.xMsContentCrc64 = null;
         }
+        String xMsOriginalContentLength = rawHeaders.getValue(X_MS_ORIGINAL_CONTENT_LENGTH);
+        if (xMsOriginalContentLength != null) {
+            this.xMsOriginalContentLength = Long.parseLong(xMsOriginalContentLength);
+        } else {
+            this.xMsOriginalContentLength = null;
+        }
         Map<String, String> xMsMetaHeaderCollection = new LinkedHashMap<>();
         Map<String, String> xMsOrHeaderCollection = new LinkedHashMap<>();
 
@@ -1596,7 +1610,7 @@ public byte[] getXMsContentCrc64() {
 
     /**
      * Set the xMsContentCrc64 property: The x-ms-content-crc64 property.
-     * 
+     *
      * @param xMsContentCrc64 the xMsContentCrc64 value to set.
      * @return the BlobsDownloadHeaders object itself.
      */
@@ -1605,4 +1619,24 @@ public BlobsDownloadHeaders setXMsContentCrc64(byte[] xMsContentCrc64) {
         this.xMsContentCrc64 = CoreUtils.clone(xMsContentCrc64);
         return this;
     }
+
+    /**
+     * Get the xMsOriginalContentLength property: The wire size of the encoded structured message body before decoding.
+     *
+     * @return the xMsOriginalContentLength value.
+     */
+    public Long getXMsOriginalContentLength() {
+        return this.xMsOriginalContentLength;
+    }
+
+    /**
+     * Set the xMsOriginalContentLength property: The wire size of the encoded structured message body before decoding.
+     *
+     * @param xMsOriginalContentLength the xMsOriginalContentLength value to set.
+     * @return the BlobsDownloadHeaders object itself.
+     */
+    public BlobsDownloadHeaders setXMsOriginalContentLength(Long xMsOriginalContentLength) {
+        this.xMsOriginalContentLength = xMsOriginalContentLength;
+        return this;
+    }
 }

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/models/BlobDownloadHeaders.java

@@ -1044,4 +1044,25 @@ public BlobDownloadHeaders setCreationTime(OffsetDateTime creationTime) {
         internalHeaders.setXMsCreationTime(creationTime);
         return this;
     }
+
+    /**
+     * Get the originalContentLength property: The wire size of the encoded structured message body before decoding.
+     * Only present when content validation is active and the response was decoded transparently.
+     *
+     * @return the originalContentLength value.
+     */
+    public Long getOriginalContentLength() {
+        return internalHeaders.getXMsOriginalContentLength();
+    }
+
+    /**
+     * Set the originalContentLength property: The wire size of the encoded structured message body before decoding.
+     *
+     * @param originalContentLength the originalContentLength value to set.
+     * @return the BlobDownloadHeaders object itself.
+     */
+    public BlobDownloadHeaders setOriginalContentLength(Long originalContentLength) {
+        internalHeaders.setXMsOriginalContentLength(originalContentLength);
+        return this;
+    }
 }

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/implementation/Constants.java

@@ -264,6 +264,11 @@ public static final class HeaderConstants {
         public static final HttpHeaderName STRUCTURED_CONTENT_LENGTH_HEADER_NAME
             = HttpHeaderName.fromString(STRUCTURED_CONTENT_LENGTH);
 
+        public static final String ORIGINAL_CONTENT_LENGTH = "x-ms-original-content-length";
+
+        public static final HttpHeaderName ORIGINAL_CONTENT_LENGTH_HEADER_NAME
+            = HttpHeaderName.fromString(ORIGINAL_CONTENT_LENGTH);
+
         /**
          * Metadata key ("hdi_isfolder") used to mark virtual directories in Azure Blob Storage.
          *

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/DecodedResponse.java

@@ -8,6 +8,7 @@
 import com.azure.core.http.HttpResponse;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.FluxUtil;
+import com.azure.storage.common.implementation.Constants;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
@@ -26,19 +27,30 @@
 class DecodedResponse extends HttpResponse {
     private final HttpResponse originalResponse;
     private final Flux<ByteBuffer> decodedBody;
+    private final HttpHeaders adjustedHeaders;
 
     /**
      * Wraps {@code httpResponse} with a body backed by {@code decodedBody}.
      *
-     * @param httpResponse The original response from the storage service. Its request, status code, and headers
-     * are preserved verbatim.
-     * @param decodedBody The Flux of CRC-validated, framing-stripped payload bytes produced by the decoder
-     * pipeline.
+     * <p>{@code Content-Length} is overridden to {@code decodedContentLength} so callers see the size of the bytes
+     * they will actually read. The original wire size is preserved in {@code x-ms-original-content-length}.</p>
+     *
+     * @param httpResponse The original response from the storage service.
+     * @param decodedBody The Flux of CRC-validated, framing-stripped payload bytes produced by the decoder pipeline.
+     * @param originalContentLength The wire size of the encoded structured message body.
+     * @param decodedContentLength The size of the decoded payload that callers will consume.
      */
-    DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody) {
+    DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody, long originalContentLength,
+        long decodedContentLength) {
         super(httpResponse.getRequest());
         this.originalResponse = httpResponse;
         this.decodedBody = decodedBody;
+        HttpHeaders headers = new HttpHeaders();
+        httpResponse.getHeaders().stream().forEach(h -> headers.set(h.getName(), h.getValue()));
+        headers.set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(decodedContentLength));
+        headers.set(Constants.HeaderConstants.ORIGINAL_CONTENT_LENGTH_HEADER_NAME,
+            String.valueOf(originalContentLength));
+        this.adjustedHeaders = headers;
     }
 
     @Override
@@ -49,12 +61,12 @@ public int getStatusCode() {
     @Override
     @SuppressWarnings("deprecation")
     public String getHeaderValue(String name) {
-        return originalResponse.getHeaderValue(name);
+        return adjustedHeaders.getValue(name);
     }
 
     @Override
     public HttpHeaders getHeaders() {
-        return originalResponse.getHeaders();
+        return adjustedHeaders;
     }
 
     @Override

sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java

@@ -73,11 +73,14 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
             // Confirm the service actually honored our structured-body request before we hand the body to the decoder.
             validateStructuredMessageHeaders(httpResponse);
 
+            Long decodedContentLength = getStructuredContentLength(httpResponse.getHeaders());
+
             // Fresh decoder per response so retries each get a clean state machine.
             StructuredMessageDecoder decoder = new StructuredMessageDecoder(contentLength);
 
             Flux<ByteBuffer> decodedStream = decodeStream(httpResponse.getBody(), decoder);
-            return new DecodedResponse(httpResponse, decodedStream);
+            // decodedContentLength is guaranteed non-null here: validateStructuredMessageHeaders confirmed its presence.
+            return new DecodedResponse(httpResponse, decodedStream, contentLength, decodedContentLength);
         });
     }
 
@@ -123,6 +126,22 @@ private static Long getContentLength(HttpHeaders headers) {
         return null;
     }
 
+    /**
+     * Reads {@code x-ms-structured-content-length} as a {@code Long}, returning {@code null} if missing or
+     * unparseable.
+     */
+    private static Long getStructuredContentLength(HttpHeaders headers) {
+        String value = headers.getValue(Constants.HeaderConstants.STRUCTURED_CONTENT_LENGTH_HEADER_NAME);
+        if (value != null) {
+            try {
+                return Long.parseLong(value);
+            } catch (NumberFormatException e) {
+                // Malformed header; fall through to null.
+            }
+        }
+        return null;
+    }
+
     /**
      * @return true for a 2xx response to a GET request, the only response shape that carries a body we
      * can decode. 206 (Partial Content) on retried range downloads is included.

sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/DecodedResponseTests.java

@@ -10,6 +10,7 @@
 import com.azure.core.http.HttpResponse;
 import com.azure.core.test.http.MockHttpResponse;
 import com.azure.core.util.BinaryData;
+import com.azure.storage.common.implementation.Constants;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 import reactor.test.StepVerifier;
@@ -64,17 +65,19 @@ public void preservesRequestStatusCodeAndHeaders() {
         HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, "100").set(CUSTOM_HEADER, "value");
         MockHttpResponse original = mockResponse(206, h, bytes("encoded"));
 
-        DecodedResponse wrapper = new DecodedResponse(original, fluxOf(bytes("decoded")));
+        DecodedResponse wrapper = new DecodedResponse(original, fluxOf(bytes("decoded")), 100L, 80L);
 
         assertSame(original.getRequest(), wrapper.getRequest());
         assertEquals(206, wrapper.getStatusCode());
-        assertSame(h, wrapper.getHeaders());
+        // Content-Length is overridden to decoded size; other headers are preserved.
+        assertEquals("80", wrapper.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH));
+        assertEquals("value", wrapper.getHeaders().getValue(CUSTOM_HEADER));
     }
 
     @Test
     public void getHeaderValueByStringReturnsHeaderValue() {
         HttpHeaders h = headers(CUSTOM_HEADER, "value");
-        DecodedResponse wrapper = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(new byte[0]));
+        DecodedResponse wrapper = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(new byte[0]), 0L, 0L);
 
         assertEquals("value", wrapper.getHeaderValue(CUSTOM_HEADER.getCaseInsensitiveName()));
         assertNull(wrapper.getHeaderValue("nonexistent"));
@@ -84,7 +87,7 @@ public void getHeaderValueByStringReturnsHeaderValue() {
     public void getBodyReturnsDecodedFlux() {
         byte[] decoded = bytes("decoded body");
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded), 0L, 0L);
 
         StepVerifier.create(wrapper.getBody().reduce(new ByteArrayOutputStream(), (sink, buf) -> {
             byte[] copy = new byte[buf.remaining()];
@@ -98,7 +101,7 @@ public void getBodyReturnsDecodedFlux() {
     public void getBodyAsByteArrayReturnsDecodedBytes() {
         byte[] decoded = bytes("decoded body");
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded), 0L, 0L);
 
         StepVerifier.create(wrapper.getBodyAsByteArray())
             .expectNextMatches(b -> Arrays.equals(decoded, b))
@@ -111,7 +114,7 @@ public void getBodyAsStringDefaultsToUtf8WhenNoCharsetSpecified() {
         // BOM nor a Content-Type charset parameter is present. This test pins the "no headers, no BOM" path.
         String text = "héllo wörld – ✓";
         DecodedResponse wrapper = new DecodedResponse(mockResponse(200, new HttpHeaders(), new byte[0]),
-            fluxOf(text.getBytes(StandardCharsets.UTF_8)));
+            fluxOf(text.getBytes(StandardCharsets.UTF_8)), 0L, 0L);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -124,7 +127,7 @@ public void getBodyAsStringHonorsCharsetFromContentTypeHeader() {
         String text = "ümlaut";
         byte[] iso = text.getBytes(StandardCharsets.ISO_8859_1);
         HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_TYPE, "text/plain; charset=ISO-8859-1");
-        DecodedResponse wrapper = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(iso));
+        DecodedResponse wrapper = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(iso), 0L, 0L);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -140,7 +143,7 @@ public void getBodyAsStringDetectsUtf8BomAndStripsIt() {
         System.arraycopy(bom, 0, withBom, 0, bom.length);
         System.arraycopy(payload, 0, withBom, bom.length, payload.length);
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), new byte[0]), fluxOf(withBom));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), new byte[0]), fluxOf(withBom), 0L, 0L);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -150,7 +153,7 @@ public void getBodyAsStringDecodesUsingProvidedCharset() {
         String text = "ümlaut";
         byte[] latin1 = text.getBytes(StandardCharsets.ISO_8859_1);
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), new byte[0]), fluxOf(latin1));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), new byte[0]), fluxOf(latin1), 0L, 0L);
 
         StepVerifier.create(wrapper.getBodyAsString(StandardCharsets.ISO_8859_1)).expectNext(text).verifyComplete();
     }
@@ -160,7 +163,7 @@ public void inheritedGetBodyAsInputStreamUsesDecodedBytes() throws IOException {
         // Base getBodyAsInputStream() routes through getBodyAsByteArray(), so the override is exercised end-to-end.
         byte[] decoded = bytes("decoded stream");
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded), 0L, 0L);
 
         try (InputStream stream = wrapper.getBodyAsInputStream().block()) {
             assertNotNull(stream);
@@ -172,7 +175,7 @@ public void inheritedGetBodyAsInputStreamUsesDecodedBytes() throws IOException {
     public void inheritedWriteBodyToWritesDecodedBytes() throws IOException {
         byte[] decoded = bytes("write me");
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded), 0L, 0L);
 
         ByteArrayOutputStream sink = new ByteArrayOutputStream();
         try (WritableByteChannel channel = Channels.newChannel(sink)) {
@@ -186,7 +189,7 @@ public void inheritedWriteBodyToWritesDecodedBytes() throws IOException {
     public void inheritedBufferReturnsResponseBackedByDecodedBytes() {
         byte[] decoded = bytes("buffered");
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, new HttpHeaders(), bytes("encoded")), fluxOf(decoded), 0L, 0L);
 
         HttpResponse buffered = wrapper.buffer();
         assertNotNull(buffered);
@@ -201,15 +204,34 @@ public void inheritedGetBodyAsBinaryDataReturnsDecodedBytes() {
         // must contain the decoded payload, not the original wire body. A divergent Content-Length header is set
         // to make the wire vs decoded distinction explicit and guard against regressions in header forwarding.
         byte[] decoded = bytes("decoded payload");
-        HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(decoded.length + 32));
-        DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, h, bytes("encoded wire body")), fluxOf(decoded));
+        long originalSize = decoded.length + 32;
+        HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(originalSize));
+        DecodedResponse wrapper = new DecodedResponse(mockResponse(200, h, bytes("encoded wire body")), fluxOf(decoded),
+            originalSize, decoded.length);
 
         BinaryData data = wrapper.getBodyAsBinaryData();
         assertNotNull(data);
         assertArrayEquals(decoded, data.toBytes());
     }
 
+    @Test
+    public void contentLengthIsOverriddenToDecodedSizeAndOriginalIsPreserved() {
+        long originalSize = 500L;
+        long decodedSize = 300L;
+        HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(originalSize))
+            .set(CUSTOM_HEADER, "preserve-me");
+        DecodedResponse wrapper
+            = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(new byte[0]), originalSize, decodedSize);
+
+        assertEquals(String.valueOf(decodedSize), wrapper.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH));
+        assertEquals(String.valueOf(originalSize),
+            wrapper.getHeaders().getValue(Constants.HeaderConstants.ORIGINAL_CONTENT_LENGTH_HEADER_NAME));
+        assertEquals("preserve-me", wrapper.getHeaders().getValue(CUSTOM_HEADER));
+        // Deprecated getHeaderValue must reflect the same overrides.
+        assertEquals(String.valueOf(decodedSize),
+            wrapper.getHeaderValue(HttpHeaderName.CONTENT_LENGTH.getCaseInsensitiveName()));
+    }
+
     private static byte[] readAll(InputStream stream) throws IOException {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         byte[] buf = new byte[1024];