Storage - Content validation content-length override

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

@@ -15,30 +15,36 @@
 import java.nio.charset.Charset;
 
 /**
- * {@link HttpResponse} wrapper that exposes a decoded body stream while preserving the request, status code, and
- * headers of the original response.
+ * {@link HttpResponse} wrapper that exposes a decoded body stream while preserving the request and status code of
+ * the original response.
  *
  * <p>The policy hands this class a Flux that already represents validated, framing-stripped bytes (produced by the
  * decoder pipeline). This class's only job is to make that Flux look like the body of the original
- * {@link HttpResponse}. Status code, headers, and request remain identical to the underlying response so callers
- * cannot distinguish a validated download from a normal one – the validation is transparent.</p>
+ * {@link HttpResponse}. {@code Content-Length} is overridden to the decoded payload size so it matches what callers
+ * will actually read; all other headers are forwarded verbatim. The validation is transparent to callers.</p>
  */
 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 from the decoded payload, not the larger wire size of the structured message.</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 decodedContentLength The size of the decoded payload that callers will consume.
      */
-    DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody) {
+    DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody, long decodedContentLength) {
         super(httpResponse.getRequest());
         this.originalResponse = httpResponse;
         this.decodedBody = decodedBody;
+        HttpHeaders headers = new HttpHeaders(httpResponse.getHeaders());
+        headers.set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(decodedContentLength));
+        this.adjustedHeaders = headers;
     }
 
     @Override
@@ -49,12 +55,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
@@ -70,7 +76,7 @@ public Mono<byte[]> getBodyAsByteArray() {
     @Override
     public Mono<String> getBodyAsString() {
         return getBodyAsByteArray()
-            .map(b -> CoreUtils.bomAwareToString(b, originalResponse.getHeaderValue(HttpHeaderName.CONTENT_TYPE)));
+            .map(b -> CoreUtils.bomAwareToString(b, adjustedHeaders.getValue(HttpHeaderName.CONTENT_TYPE)));
     }
 
     @Override

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

@@ -70,14 +70,15 @@ public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineN
                 return httpResponse;
             }
 
-            // Confirm the service actually honored our structured-body request before we hand the body to the decoder.
-            validateStructuredMessageHeaders(httpResponse);
+            // Confirm the service honored our structured-body request and parse the decoded length in one step,
+            // failing fast with a consistent error if either header is absent or not parseable as a long.
+            long decodedContentLength = validateAndGetDecodedContentLength(httpResponse);
 
             // 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);
+            return new DecodedResponse(httpResponse, decodedStream, decodedContentLength);
         });
     }
 
@@ -92,11 +93,14 @@ private boolean shouldApplyDecoding(HttpPipelineCallContext context) {
     }
 
     /**
-     * Verifies the response acknowledges the structured-body request: presence of the
-     * {@code x-ms-structured-body} header and the {@code x-ms-structured-content-length} 
-     * header. If either is missing, the service is sending us a normal body and we must not run the decoder over it.
+     * Verifies the response acknowledges the structured-body request ({@code x-ms-structured-body} present) and
+     * parses {@code x-ms-structured-content-length} in one step. Throws with a consistent error message if either
+     * header is absent or the length value is not parseable as a {@code long}, matching the fail-fast behaviour
+     * used for other validation failures.
+     *
+     * @return the decoded payload size declared by the service.
      */
-    private void validateStructuredMessageHeaders(HttpResponse httpResponse) {
+    private long validateAndGetDecodedContentLength(HttpResponse httpResponse) {
         String structuredBody
             = httpResponse.getHeaders().getValue(Constants.HeaderConstants.STRUCTURED_BODY_TYPE_HEADER_NAME);
         String structuredContentLength
@@ -105,6 +109,12 @@ private void validateStructuredMessageHeaders(HttpResponse httpResponse) {
             throw LOGGER.logExceptionAsError(
                 new IllegalStateException("Structured message was requested but the response did not acknowledge it."));
         }
+        try {
+            return Long.parseLong(structuredContentLength);
+        } catch (NumberFormatException e) {
+            throw LOGGER.logExceptionAsError(new IllegalStateException("x-ms-structured-content-length header value '"
+                + structuredContentLength + "' could not be parsed as a long.", e));
+        }
     }
 
     /**

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

@@ -64,17 +64,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")), 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);
 
         assertEquals("value", wrapper.getHeaderValue(CUSTOM_HEADER.getCaseInsensitiveName()));
         assertNull(wrapper.getHeaderValue("nonexistent"));
@@ -84,7 +86,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);
 
         StepVerifier.create(wrapper.getBody().reduce(new ByteArrayOutputStream(), (sink, buf) -> {
             byte[] copy = new byte[buf.remaining()];
@@ -98,7 +100,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);
 
         StepVerifier.create(wrapper.getBodyAsByteArray())
             .expectNextMatches(b -> Arrays.equals(decoded, b))
@@ -111,7 +113,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);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -124,7 +126,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);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -140,7 +142,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);
 
         StepVerifier.create(wrapper.getBodyAsString()).expectNext(text).verifyComplete();
     }
@@ -150,7 +152,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);
 
         StepVerifier.create(wrapper.getBodyAsString(StandardCharsets.ISO_8859_1)).expectNext(text).verifyComplete();
     }
@@ -160,7 +162,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);
 
         try (InputStream stream = wrapper.getBodyAsInputStream().block()) {
             assertNotNull(stream);
@@ -172,7 +174,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);
 
         ByteArrayOutputStream sink = new ByteArrayOutputStream();
         try (WritableByteChannel channel = Channels.newChannel(sink)) {
@@ -186,7 +188,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);
 
         HttpResponse buffered = wrapper.buffer();
         assertNotNull(buffered);
@@ -201,15 +203,32 @@ 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));
+        long decodedSize = decoded.length;
+        HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(decodedSize + 32));
         DecodedResponse wrapper
-            = new DecodedResponse(mockResponse(200, h, bytes("encoded wire body")), fluxOf(decoded));
+            = new DecodedResponse(mockResponse(200, h, bytes("encoded wire body")), fluxOf(decoded), decodedSize);
 
         BinaryData data = wrapper.getBodyAsBinaryData();
         assertNotNull(data);
         assertArrayEquals(decoded, data.toBytes());
     }
 
+    @Test
+    public void contentLengthIsOverriddenToDecodedSize() {
+        long wireSize = 500L;
+        long decodedSize = 300L;
+        HttpHeaders h = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(wireSize))
+            .set(CUSTOM_HEADER, "preserve-me");
+        DecodedResponse wrapper
+            = new DecodedResponse(mockResponse(200, h, new byte[0]), fluxOf(new byte[0]), decodedSize);
+
+        assertEquals(String.valueOf(decodedSize), wrapper.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH));
+        assertEquals("preserve-me", wrapper.getHeaders().getValue(CUSTOM_HEADER));
+        // Deprecated getHeaderValue must reflect the same override.
+        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];