Add StorageContentValidationDecoderPolicy for pipeline-based structured message decoding #46951
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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.DownloadContentValidationOptions; | ||
| import reactor.core.publisher.Flux; | ||
| import reactor.core.publisher.Mono; | ||
|
|
@@ -1333,10 +1333,21 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down | |
| DownloadRetryOptions finalOptions = (options == null) ? new DownloadRetryOptions() : options; | ||
|
|
||
| // The first range should eagerly convert headers as they'll be used to create response types. | ||
| Context initialContext = context == null | ||
| ? new Context("azure-eagerly-convert-headers", true) | ||
| : context.addData("azure-eagerly-convert-headers", true); | ||
|
|
||
| // Add structured message decoding context if enabled | ||
| final Context firstRangeContext; | ||
| if (contentValidationOptions != null | ||
| && contentValidationOptions.isStructuredMessageValidationEnabled()) { | ||
| firstRangeContext = initialContext | ||
| .addData(Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEY, true) | ||
| .addData(Constants.STRUCTURED_MESSAGE_VALIDATION_OPTIONS_CONTEXT_KEY, contentValidationOptions); | ||
| } else { | ||
| firstRangeContext = initialContext; | ||
| } | ||
|
|
||
| return downloadRange(finalRange, finalRequestConditions, finalRequestConditions.getIfMatch(), finalGetMD5, | ||
| firstRangeContext).map(response -> { | ||
| BlobsDownloadHeaders blobsDownloadHeaders = new BlobsDownloadHeaders(response.getHeaders()); | ||
|
|
@@ -1357,16 +1368,6 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down | |
| finalCount = finalRange.getCount(); | ||
| } | ||
|
|
||
| // 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 +1392,15 @@ Mono<BlobDownloadAsyncResponse> downloadStreamWithResponse(BlobRange range, Down | |
|
|
||
| try { | ||
| return downloadRange(new BlobRange(initialOffset + offset, newCount), finalRequestConditions, | ||
| eTag, finalGetMD5, firstRangeContext); | ||
|
There was a problem hiding this comment. @copilot getting this compile error here : Variable used in lambda expression should be final or effectively final for firstRangeContext. so please fix it. There was a problem hiding this comment. Fixed in commit (in progress). The issue was that Solution: Introduced There was a problem hiding this comment. Fixed in commit f9d9113. The issue was that Solution: Introduced |
||
| } catch (Exception e) { | ||
| return Mono.error(e); | ||
| } | ||
| }; | ||
|
|
||
| // Structured message decoding is now handled by StructuredMessageDecoderPolicy | ||
| return BlobDownloadAsyncResponseConstructorProxy.create(response, onDownloadErrorResume, | ||
| finalOptions); | ||
| }); | ||
| } | ||
|
|
||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,227 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT License. | ||
|
|
||
| package com.azure.storage.blob.specialized; | ||
|
|
||
| import com.azure.core.test.utils.TestUtils; | ||
| import com.azure.core.util.FluxUtil; | ||
| import com.azure.storage.blob.BlobAsyncClient; | ||
| import com.azure.storage.blob.BlobTestBase; | ||
| import com.azure.storage.blob.models.BlobRange; | ||
| import com.azure.storage.blob.models.BlobRequestConditions; | ||
| import com.azure.storage.blob.models.DownloadRetryOptions; | ||
| import com.azure.storage.common.DownloadContentValidationOptions; | ||
| import com.azure.storage.common.implementation.Constants; | ||
| import com.azure.storage.common.implementation.structuredmessage.StructuredMessageEncoder; | ||
| import com.azure.storage.common.implementation.structuredmessage.StructuredMessageFlags; | ||
| import org.junit.jupiter.api.BeforeEach; | ||
| import org.junit.jupiter.api.Test; | ||
| import reactor.core.publisher.Flux; | ||
| import reactor.test.StepVerifier; | ||
|
|
||
| import java.io.IOException; | ||
| import java.nio.ByteBuffer; | ||
|
|
||
| import static org.junit.jupiter.api.Assertions.assertEquals; | ||
| import static org.junit.jupiter.api.Assertions.assertNotNull; | ||
| import static org.junit.jupiter.api.Assertions.assertTrue; | ||
|
|
||
| /** | ||
| * Tests for structured message decoding during blob downloads using StorageContentValidationDecoderPolicy. | ||
| * These tests verify that the pipeline policy correctly decodes structured messages when content validation is enabled. | ||
| */ | ||
| public class BlobMessageDecoderDownloadTests extends BlobTestBase { | ||
|
|
||
| private BlobAsyncClient bc; | ||
|
|
||
| @BeforeEach | ||
| public void setup() { | ||
| String blobName = generateBlobName(); | ||
| bc = ccAsync.getBlobAsyncClient(blobName); | ||
| bc.upload(Flux.just(ByteBuffer.wrap(new byte[0])), null).block(); | ||
| } | ||
|
|
||
| @Test | ||
| public void downloadStreamWithResponseContentValidation() throws IOException { | ||
| 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(true); | ||
|
|
||
| StepVerifier | ||
| .create(bc.upload(input, null, true) | ||
| .then(bc.downloadStreamWithResponse((BlobRange) null, (DownloadRetryOptions) null, | ||
gunjansingh-msft marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| (BlobRequestConditions) null, false, validationOptions)) | ||
| .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue()))) | ||
| .assertNext(r -> TestUtils.assertArraysEqual(r, randomData)) | ||
| .verifyComplete(); | ||
| } | ||
|
|
||
| @Test | ||
| public void downloadStreamWithResponseContentValidationRange() throws IOException { | ||
|
There was a problem hiding this comment. @copilot This test is failing due to this error : [ERROR] Failures: There was a problem hiding this comment. The test was incorrect. Range downloads are incompatible with structured message validation because the decoder needs the complete encoded message to validate, but range downloads only fetch partial bytes. Fixed in commit (in progress):
The original test tried to validate a 512-byte range of a 1081-byte encoded message, which correctly failed because the decoder cannot validate an incomplete structured message. There was a problem hiding this comment. The test was incorrect. Range downloads are incompatible with structured message validation because the decoder needs the complete encoded message to validate, but range downloads only fetch partial bytes. Fixed in commit 49b0a17:
The original test tried to validate a 512-byte range of a 1081-byte encoded message, which correctly failed because the decoder cannot validate an incomplete structured message. |
||
| // Note: Range downloads are not compatible with structured message validation | ||
| // because you need the complete encoded message for validation. | ||
| // This test verifies that range downloads work without validation. | ||
| 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); | ||
|
|
||
| // Range download without validation should work | ||
| BlobRange range = new BlobRange(0, 512L); | ||
|
|
||
| StepVerifier.create(bc.upload(input, null, true) | ||
| .then(bc.downloadStreamWithResponse(range, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) null, false)) | ||
| .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue()))).assertNext(r -> { | ||
| assertNotNull(r); | ||
| // Should get exactly 512 bytes of encoded data | ||
| assertEquals(512, r.length); | ||
| }).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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) 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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) 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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) 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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) 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(); | ||
| } | ||
|
|
||
| @Test | ||
| public void downloadStreamWithResponseContentValidationSmallSegment() throws IOException { | ||
| // Test with small segment size to ensure boundary conditions are handled | ||
| byte[] randomData = getRandomByteArray(256); | ||
| StructuredMessageEncoder encoder | ||
| = new StructuredMessageEncoder(randomData.length, 128, 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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) null, false, validationOptions)) | ||
| .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue()))) | ||
| .assertNext(r -> TestUtils.assertArraysEqual(r, randomData)) | ||
| .verifyComplete(); | ||
| } | ||
|
|
||
| @Test | ||
| public void downloadStreamWithResponseContentValidationVeryLargeBlob() throws IOException { | ||
| // Test with very large data to verify chunking and policy work correctly with large blobs | ||
| byte[] randomData = getRandomByteArray(10 * Constants.KB); | ||
| StructuredMessageEncoder encoder | ||
| = new StructuredMessageEncoder(randomData.length, 2048, 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((BlobRange) null, (DownloadRetryOptions) null, | ||
| (BlobRequestConditions) null, false, validationOptions)) | ||
| .flatMap(r -> FluxUtil.collectBytesInByteBufferStream(r.getValue()))) | ||
| .assertNext(r -> TestUtils.assertArraysEqual(r, randomData)) | ||
| .verifyComplete(); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@copilot do you think the way we are using the pipelinepolicy here is correct? will it only apply to the download methods affected due to conetent validation or the entire client? This is how it's being used in the custom encoder pipeline policy.
https://github.com/Azure/azure-sdk-for-java/pull/46002/files#diff-711d7fd895212bdcd4514bc8cce946c4ca9ebe39116262b70bec1496139825ac
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, the implementation is correct. The pipeline policy is added to the entire client (in
BuilderHelper.buildPipeline()), similar to howBlobDecryptionPolicyworks. However, it only activates for specific download requests based on the context key.The policy checks for
Constants.STRUCTURED_MESSAGE_DECODING_CONTEXT_KEYin the context, which is only set indownloadStreamWithResponse()whencontentValidationOptions.isStructuredMessageValidationEnabled()is true. For all other requests (uploads, metadata operations, etc.), the policy is present in the pipeline but immediately passes through without any processing.This context-based activation ensures:
This matches the pattern from
BlobDecryptionPolicywhich is also added to the entire client pipeline but only activates based on context data.