diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/DecodedResponse.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/DecodedResponse.java index 1f604f0d6cac..7e88de8cd37e 100644 --- a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/DecodedResponse.java +++ b/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. * *

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.

+ * {@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.

*/ class DecodedResponse extends HttpResponse { private final HttpResponse originalResponse; private final Flux 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. + *

{@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.

+ * + * @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 decodedBody) { + DecodedResponse(HttpResponse httpResponse, Flux 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 getBodyAsByteArray() { @Override public Mono getBodyAsString() { return getBodyAsByteArray() - .map(b -> CoreUtils.bomAwareToString(b, originalResponse.getHeaderValue(HttpHeaderName.CONTENT_TYPE))); + .map(b -> CoreUtils.bomAwareToString(b, adjustedHeaders.getValue(HttpHeaderName.CONTENT_TYPE))); } @Override diff --git a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java index 1b4c400d27d2..21150174a03a 100644 --- a/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java +++ b/sdk/storage/azure-storage-common/src/main/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicy.java @@ -70,14 +70,15 @@ public Mono 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 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)); + } } /** diff --git a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/DecodedResponseTests.java b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/DecodedResponseTests.java index 53f34be58f32..e99c169a1763 100644 --- a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/DecodedResponseTests.java +++ b/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]; diff --git a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicyTests.java b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicyTests.java index f9ff398a4611..6b0b8474ee5c 100644 --- a/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicyTests.java +++ b/sdk/storage/azure-storage-common/src/test/java/com/azure/storage/common/policy/StorageContentValidationDecoderPolicyTests.java @@ -18,6 +18,7 @@ import com.azure.storage.common.implementation.contentvalidation.StructuredMessageConstants; import com.azure.storage.common.implementation.contentvalidation.StructuredMessageEncoder; import com.azure.storage.common.implementation.contentvalidation.StructuredMessageFlags; +import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; @@ -83,13 +84,87 @@ public void decodesDynamicallySizedSegmentStructuredMessageThroughPipeline(int s sent.getHeaders().getValue(Constants.HeaderConstants.STRUCTURED_BODY_TYPE_HEADER_NAME)); } + @Test + public void contentLengthIsOverriddenToDecodedSizeWhenDecodingApplied() throws IOException { + byte[] payload = new byte[64]; + ThreadLocalRandom.current().nextBytes(payload); + + byte[] encoded = encodeStructuredMessageWireBytes(payload, 64, StructuredMessageFlags.STORAGE_CRC64); + long encodedLen = encoded.length; + long decodedLen = payload.length; + + HttpClient httpClient = request -> Mono.just( + new MockHttpResponse(request, 200, structuredDownloadResponseHeaders(encodedLen, decodedLen), encoded)); + + HttpPipeline pipeline = new HttpPipelineBuilder().policies((context, next) -> { + context.setData(StructuredMessageConstants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true); + return next.process(); + }, new StorageContentValidationDecoderPolicy()).httpClient(httpClient).build(); + + HttpRequest request = new HttpRequest(HttpMethod.GET, "https://example.blob.core.windows.net/c/b"); + try (HttpResponse response = pipeline.send(request, Context.NONE).block()) { + assertNotNull(response); + assertEquals(String.valueOf(decodedLen), response.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH)); + } + } + + @Test + public void contentLengthMatchesActualDecodedBodySize() throws IOException { + byte[] payload = new byte[128]; + ThreadLocalRandom.current().nextBytes(payload); + + byte[] encoded = encodeStructuredMessageWireBytes(payload, 64, StructuredMessageFlags.STORAGE_CRC64); + long encodedLen = encoded.length; + long decodedLen = payload.length; + + HttpClient httpClient = request -> Mono.just( + new MockHttpResponse(request, 200, structuredDownloadResponseHeaders(encodedLen, decodedLen), encoded)); + + HttpPipeline pipeline = new HttpPipelineBuilder().policies((context, next) -> { + context.setData(StructuredMessageConstants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true); + return next.process(); + }, new StorageContentValidationDecoderPolicy()).httpClient(httpClient).build(); + + HttpRequest request = new HttpRequest(HttpMethod.GET, "https://example.blob.core.windows.net/c/b"); + try (HttpResponse response = pipeline.send(request, Context.NONE).block()) { + assertNotNull(response); + assertEquals(String.valueOf(decodedLen), response.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH)); + + byte[] body = Objects.requireNonNull(FluxUtil.collectBytesInByteBufferStream(response.getBody()).block()); + assertEquals(decodedLen, body.length); + assertArrayEquals(payload, body); + } + } + + @Test + public void contentLengthIsUnchangedWhenDecodingNotApplied() throws IOException { + byte[] payload = new byte[64]; + ThreadLocalRandom.current().nextBytes(payload); + + byte[] encoded = encodeStructuredMessageWireBytes(payload, 64, StructuredMessageFlags.STORAGE_CRC64); + long encodedLen = encoded.length; + + HttpHeaders responseHeaders = new HttpHeaders().set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(encodedLen)); + HttpClient httpClient = request -> Mono.just(new MockHttpResponse(request, 200, responseHeaders, encoded)); + + HttpPipeline pipeline = new HttpPipelineBuilder().policies(new StorageContentValidationDecoderPolicy()) + .httpClient(httpClient) + .build(); + + HttpRequest request = new HttpRequest(HttpMethod.GET, "https://example.blob.core.windows.net/c/b"); + HttpResponse response = pipeline.send(request, Context.NONE).block(); + + assertNotNull(response); + assertEquals(String.valueOf(encodedLen), response.getHeaders().getValue(HttpHeaderName.CONTENT_LENGTH)); + } + private static Stream segmentPayloadSizeAndTotalPayloadSizeSupplier() { return Stream.of(Arguments.of(10 * 1024 * 1024, 10 * 1024 * 1024 + 1), // larger than 4 MiB Arguments.of(3 * 1024 * 1024, 3 * 1024 * 1024 + 1), // smaller than 4 MiB, but not KB Arguments.of(5 * 1024 * 1024 + 1, 15 * 1024 * 1024)); } - private static HttpHeaders structuredDownloadResponseHeaders(int contentLength, long structuredContentLength) { + private static HttpHeaders structuredDownloadResponseHeaders(long contentLength, long structuredContentLength) { HttpHeaders headers = new HttpHeaders(); headers.set(HttpHeaderName.CONTENT_LENGTH, String.valueOf(contentLength)); headers.set(Constants.HeaderConstants.STRUCTURED_BODY_TYPE_HEADER_NAME,