Add StorageContentValidationDecoderPolicy for pipeline-based structured message decoding with smart retry support

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/implementation/util/BuilderHelper.java

@@ -39,6 +39,7 @@
 import com.azure.storage.common.policy.ResponseValidationPolicyBuilder;
 import com.azure.storage.common.policy.ScrubEtagPolicy;
 import com.azure.storage.common.policy.StorageBearerTokenChallengeAuthorizationPolicy;
+import com.azure.storage.common.policy.StorageContentValidationDecoderPolicy;
 import com.azure.storage.common.policy.StorageSharedKeyCredentialPolicy;
 
 import java.net.MalformedURLException;
@@ -140,6 +141,9 @@ public static HttpPipeline buildPipeline(StorageSharedKeyCredential storageShare
 
         HttpPolicyProviders.addAfterRetryPolicies(policies);
 
+        // Add structured message decoder policy to handle structured message decoding
+        policies.add(new StorageContentValidationDecoderPolicy());
+
         policies.add(getResponseValidationPolicy());
 
         policies.add(new HttpLoggingPolicy(logOptions));

sdk/storage/azure-storage-blob/src/main/java/com/azure/storage/blob/specialized/BlobAsyncClientBase.java

@@ -81,9 +81,9 @@
 import com.azure.storage.blob.sas.BlobServiceSasSignatureValues;
 import com.azure.storage.common.StorageSharedKeyCredential;
 import com.azure.storage.common.Utility;
+import com.azure.storage.common.implementation.Constants;
 import com.azure.storage.common.implementation.SasImplUtils;
 import com.azure.storage.common.implementation.StorageImplUtils;
-import com.azure.storage.common.implementation.structuredmessage.StructuredMessageDecodingStream;
 import com.azure.storage.common.DownloadContentValidationOptions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -1337,6 +1337,16 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
             ? new Context("azure-eagerly-convert-headers", true)
             : context.addData("azure-eagerly-convert-headers", true);
 
+        // Add structured message decoding context if enabled
+        if (contentValidationOptions != null
+            && contentValidationOptions.isStructuredMessageValidationEnabled()) {
+            firstRangeContext = firstRangeContext.addData(
+                Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true);
+            firstRangeContext = firstRangeContext.addData(
+                Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY,
+                contentValidationOptions);
+        }
+
         return downloadRange(finalRange, finalRequestConditions, finalRequestConditions.getIfMatch(), finalGetMD5,
             firstRangeContext).map(response -> {
                 BlobsDownloadHeaders blobsDownloadHeaders = new BlobsDownloadHeaders(response.getHeaders());
@@ -1357,16 +1367,6 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
                     finalCount = finalRange.getCount();
                 }
 
-                // Apply structured message decoding if enabled - this allows both MD5 and structured message to coexist
-                Flux<ByteBuffer> processedStream = response.getValue();
-                if (contentValidationOptions != null
-                    && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                    // Use the content length from headers to determine expected length for structured message decoding
-                    Long contentLength = blobDownloadHeaders.getContentLength();
-                    processedStream = StructuredMessageDecodingStream.wrapStreamIfNeeded(response.getValue(),
-                        contentLength, contentValidationOptions);
-                }
-
                 // The resume function takes throwable and offset at the destination.
                 // I.e. offset is relative to the starting point.
                 BiFunction<Throwable, Long, Mono<StreamResponse>> onDownloadErrorResume = (throwable, offset) -> {
@@ -1391,27 +1391,15 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down
 
                     try {
                         return downloadRange(new BlobRange(initialOffset + offset, newCount), finalRequestConditions,
-                            eTag, finalGetMD5, context);
+                            eTag, finalGetMD5, firstRangeContext);
                     } catch (Exception e) {
                         return Mono.error(e);
                     }
                 };
 
-                // If structured message decoding was applied, we need to create a new StreamResponse with the processed stream
-                if (contentValidationOptions != null
-                    && contentValidationOptions.isStructuredMessageValidationEnabled()) {
-                    // Create a new StreamResponse using the deprecated but available constructor
-                    @SuppressWarnings("deprecation")
-                    StreamResponse processedResponse = new StreamResponse(response.getRequest(),
-                        response.getStatusCode(), response.getHeaders(), processedStream);
-
-                    return BlobDownloadAsyncResponseConstructorProxy.create(processedResponse, onDownloadErrorResume,
-                        finalOptions);
-                } else {
-                    // No structured message processing needed, use original response
-                    return BlobDownloadAsyncResponseConstructorProxy.create(response, onDownloadErrorResume,
-                        finalOptions);
-                }
+                // Structured message decoding is now handled by StructuredMessageDecoderPolicy
+                return BlobDownloadAsyncResponseConstructorProxy.create(response, onDownloadErrorResume,
+                    finalOptions);
             });
     }
 

sdk/storage/azure-storage-blob/src/test/java/com/azure/storage/blob/specialized/BlobBaseAsyncApiTests.java

@@ -607,6 +607,96 @@ public void downloadStreamWithResponseContentValidationRange() throws IOExceptio
             }).verifyComplete();
     }
 
+    @Test
+    public void downloadStreamWithResponseContentValidationLargeBlob() throws IOException {
+        // Test with larger data to verify chunking works correctly
+        byte[] randomData = getRandomByteArray(5 * Constants.KB);
+        StructuredMessageEncoder encoder
+            = new StructuredMessageEncoder(randomData.length, 1024, StructuredMessageFlags.STORAGE_CRC64);
+        ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+        Flux<ByteBuffer> input = Flux.just(encodedData);
+
+        DownloadContentValidationOptions validationOptions
+            = new DownloadContentValidationOptions().setStructuredMessageValidationEnabled(true);
+
+        StepVerifier
+            .create(bc.upload(input, null, true)
+                .then(bc.downloadStreamWithResponse(null, null, null, false, validationOptions))
+                .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue())))
+            .assertNext(r -> TestUtils.assertArraysEqual(r, randomData))
+            .verifyComplete();
+    }
+
+    @Test
+    public void downloadStreamWithResponseContentValidationMultipleSegments() throws IOException {
+        // Test with multiple segments to ensure all segments are decoded correctly
+        byte[] randomData = getRandomByteArray(2 * Constants.KB);
+        StructuredMessageEncoder encoder
+            = new StructuredMessageEncoder(randomData.length, 512, StructuredMessageFlags.STORAGE_CRC64);
+        ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+        Flux<ByteBuffer> input = Flux.just(encodedData);
+
+        DownloadContentValidationOptions validationOptions
+            = new DownloadContentValidationOptions().setStructuredMessageValidationEnabled(true);
+
+        StepVerifier
+            .create(bc.upload(input, null, true)
+                .then(bc.downloadStreamWithResponse(null, null, null, false, validationOptions))
+                .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue())))
+            .assertNext(r -> TestUtils.assertArraysEqual(r, randomData))
+            .verifyComplete();
+    }
+
+    @Test
+    public void downloadStreamWithResponseNoValidation() throws IOException {
+        // Test that download works normally when validation is not enabled
+        byte[] randomData = getRandomByteArray(Constants.KB);
+        StructuredMessageEncoder encoder
+            = new StructuredMessageEncoder(randomData.length, 512, StructuredMessageFlags.STORAGE_CRC64);
+        ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+        Flux<ByteBuffer> input = Flux.just(encodedData);
+
+        // No validation options - should download encoded data as-is
+        StepVerifier
+            .create(bc.upload(input, null, true)
+                .then(bc.downloadStreamWithResponse(null, null, null, false))
+                .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue())))
+            .assertNext(r -> {
+                assertNotNull(r);
+                // Should get encoded data, not decoded
+                assertTrue(r.length > randomData.length); // Encoded data is larger
+            })
+            .verifyComplete();
+    }
+
+    @Test
+    public void downloadStreamWithResponseValidationDisabled() throws IOException {
+        // Test with validation options but validation disabled
+        byte[] randomData = getRandomByteArray(Constants.KB);
+        StructuredMessageEncoder encoder
+            = new StructuredMessageEncoder(randomData.length, 512, StructuredMessageFlags.STORAGE_CRC64);
+        ByteBuffer encodedData = encoder.encode(ByteBuffer.wrap(randomData));
+
+        Flux<ByteBuffer> input = Flux.just(encodedData);
+
+        DownloadContentValidationOptions validationOptions
+            = new DownloadContentValidationOptions().setStructuredMessageValidationEnabled(false);
+
+        StepVerifier
+            .create(bc.upload(input, null, true)
+                .then(bc.downloadStreamWithResponse(null, null, null, false, validationOptions))
+                .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue())))
+            .assertNext(r -> {
+                assertNotNull(r);
+                // Should get encoded data, not decoded
+                assertTrue(r.length > randomData.length); // Encoded data is larger
+            })
+            .verifyComplete();
+    }
+
     @RequiredServiceVersion(clazz = BlobServiceVersion.class, min = "2024-08-04")
     @Test
     public void copyFromURLSourceErrorAndStatusCode() {

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

@@ -94,6 +94,17 @@ public final class Constants {
 
     public static final String SKIP_ECHO_VALIDATION_KEY = "skipEchoValidation";
 
+    /**
+     * Context key used to signal that structured message decoding should be applied.
+     */
+    public static final String STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY = "azure-storage-structured-message-decoding";
+
+    /**
+     * Context key used to pass DownloadContentValidationOptions to the policy.
+     */
+    public static final String STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY =
+        "azure-storage-structured-message-validation-options";
+
     private Constants() {
     }
 

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

@@ -0,0 +1,164 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.storage.common.policy;
+
+import com.azure.core.http.HttpHeaderName;
+import com.azure.core.http.HttpHeaders;
+import com.azure.core.http.HttpMethod;
+import com.azure.core.http.HttpPipelineCallContext;
+import com.azure.core.http.HttpPipelineNextPolicy;
+import com.azure.core.http.HttpResponse;
+import com.azure.core.http.policy.HttpPipelinePolicy;
+import com.azure.core.util.FluxUtil;
+import com.azure.core.util.logging.ClientLogger;
+import com.azure.storage.common.DownloadContentValidationOptions;
+import com.azure.storage.common.implementation.Constants;
+import com.azure.storage.common.implementation.structuredmessage.StructuredMessageDecodingStream;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.Charset;
+
+/**
+ * This is a decoding policy in an {@link com.azure.core.http.HttpPipeline} to decode structured messages in
+ * storage download requests. The policy checks for a context value to determine when to apply structured message decoding.
+ */
+public class StorageContentValidationDecoderPolicy implements HttpPipelinePolicy {
+    private static final ClientLogger LOGGER = new ClientLogger(StorageContentValidationDecoderPolicy.class);
+
+    /**
+     * Creates a new instance of {@link StorageContentValidationDecoderPolicy}.
+     */
+    public StorageContentValidationDecoderPolicy() {
+    }
+
+    @Override
+    public Mono<HttpResponse> process(HttpPipelineCallContext context, HttpPipelineNextPolicy next) {
+        // Check if structured message decoding is enabled for this request
+        if (!shouldApplyDecoding(context)) {
+            return next.process();
+        }
+
+        return next.process().map(httpResponse -> {
+            // Only apply decoding to download responses (GET requests with body)
+            if (!isDownloadResponse(httpResponse)) {
+                return httpResponse;
+            }
+
+            DownloadContentValidationOptions validationOptions = getValidationOptions(context);
+            Long contentLength = getContentLength(httpResponse.getHeaders());
+
+            if (contentLength != null && contentLength > 0 && validationOptions != null) {
+                Flux<ByteBuffer> decodedStream = StructuredMessageDecodingStream.wrapStreamIfNeeded(
+                    httpResponse.getBody(), contentLength, validationOptions);
+                return new DecodedResponse(httpResponse, decodedStream);
+            }
+
+            return httpResponse;
+        });
+    }
+
+    /**
+     * Checks if structured message decoding should be applied based on context.
+     *
+     * @param context The pipeline call context.
+     * @return true if decoding should be applied, false otherwise.
+     */
+    private boolean shouldApplyDecoding(HttpPipelineCallContext context) {
+        return context.getData(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY)
+            .map(value -> value instanceof Boolean && (Boolean) value)
+            .orElse(false);
+    }
+
+    /**
+     * Gets the validation options from context.
+     *
+     * @param context The pipeline call context.
+     * @return The validation options or null if not present.
+     */
+    private DownloadContentValidationOptions getValidationOptions(HttpPipelineCallContext context) {
+        return context.getData(Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY)
+            .filter(value -> value instanceof DownloadContentValidationOptions)
+            .map(value -> (DownloadContentValidationOptions) value)
+            .orElse(null);
+    }
+
+    /**
+     * Gets the content length from response headers.
+     *
+     * @param headers The response headers.
+     * @return The content length or null if not present.
+     */
+    private Long getContentLength(HttpHeaders headers) {
+        String contentLengthStr = headers.getValue(HttpHeaderName.CONTENT_LENGTH);
+        if (contentLengthStr != null) {
+            try {
+                return Long.parseLong(contentLengthStr);
+            } catch (NumberFormatException e) {
+                LOGGER.warning("Invalid content length in response headers: " + contentLengthStr);
+            }
+        }
+        return null;
+    }
+
+    /**
+     * Checks if the response is a download response (GET request with body).
+     *
+     * @param httpResponse The HTTP response.
+     * @return true if it's a download response, false otherwise.
+     */
+    private boolean isDownloadResponse(HttpResponse httpResponse) {
+        return httpResponse.getRequest().getHttpMethod() == HttpMethod.GET && httpResponse.getBody() != null;
+    }
+
+    /**
+     * HTTP response wrapper that provides a decoded response body.
+     */
+    static class DecodedResponse extends HttpResponse {
+        private final Flux<ByteBuffer> decodedBody;
+        private final HttpResponse originalResponse;
+
+        DecodedResponse(HttpResponse httpResponse, Flux<ByteBuffer> decodedBody) {
+            super(httpResponse.getRequest());
+            this.originalResponse = httpResponse;
+            this.decodedBody = decodedBody;
+        }
+
+        @Override
+        public int getStatusCode() {
+            return originalResponse.getStatusCode();
+        }
+
+        @Override
+        public String getHeaderValue(String name) {
+            return originalResponse.getHeaderValue(name);
+        }
+
+        @Override
+        public HttpHeaders getHeaders() {
+            return originalResponse.getHeaders();
+        }
+
+        @Override
+        public Flux<ByteBuffer> getBody() {
+            return decodedBody;
+        }
+
+        @Override
+        public Mono<byte[]> getBodyAsByteArray() {
+            return FluxUtil.collectBytesInByteBufferStream(decodedBody);
+        }
+
+        @Override
+        public Mono<String> getBodyAsString() {
+            return getBodyAsByteArray().map(String::new);
+        }
+
+        @Override
+        public Mono<String> getBodyAsString(Charset charset) {
+            return getBodyAsByteArray().map(bytes -> new String(bytes, charset));
+        }
+    }
+}

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

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.storage.common.policy;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+/**
+ * Unit tests for {@link StorageContentValidationDecoderPolicy}.
+ * 
+ * Note: The policy behavior is primarily validated through integration tests in BlobBaseAsyncApiTests
+ * which test the end-to-end download scenarios with structured message validation.
+ */
+public class StorageContentValidationDecoderPolicyTest {
+
+    @Test
+    public void policyCanBeInstantiated() {
+        // Verify the policy can be constructed successfully
+        StorageContentValidationDecoderPolicy policy = new StorageContentValidationDecoderPolicy();
+        assertNotNull(policy);
+    }
+}