V2 Encryption Bug 4MB file downloads

sdk/storage/azure-storage-blob-cryptography/CHANGELOG.md

@@ -7,6 +7,8 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+- Fixed a bug where downloadToFile and openInputStream was throwing an InvalidRange exception if the target file size was a multiple of the
+  authenticated region length.
 
 ### Other Changes
 

sdk/storage/azure-storage-blob-cryptography/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "java",
   "TagPrefix": "java/storage/azure-storage-blob-cryptography",
-  "Tag": "java/storage/azure-storage-blob-cryptography_e245db55c0"
+  "Tag": "java/storage/azure-storage-blob-cryptography_9fd41be58a"
 }

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

@@ -750,9 +750,14 @@ private <T> Mono<T> populateRequestConditionsAndContext(BlobRequestConditions re
 
                 String encryptionDataKey = StorageImplUtils.getEncryptionDataKey(response.getValue().getMetadata());
                 if (encryptionDataKey != null) {
-                    result = result.contextWrite(context -> context.put(ENCRYPTION_DATA_KEY,
-                        EncryptionData.getAndValidateEncryptionData(encryptionDataKey, requiresEncryption)));
+                    EncryptionData encryptionData = EncryptionData.getAndValidateEncryptionData(encryptionDataKey,
+                        requiresEncryption);
+
+                    result = result.contextWrite(context -> context.put(ENCRYPTION_DATA_KEY, encryptionData))
+                        .contextWrite(context -> context.put(Constants.ADJUSTED_BLOB_LENGTH_KEY,
+                            EncryptedBlobLength.computeAdjustedBlobLength(encryptionData, response.getValue().getBlobSize())));
                 }
+
                 return result;
             });
     }

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

@@ -489,8 +489,12 @@ private Context populateRequestConditionsAndContext(BlobRequestConditions reques
 
         String encryptionDataKey = StorageImplUtils.getEncryptionDataKey(initialProperties.getMetadata());
         if (encryptionDataKey != null) {
-            context = context.addData(ENCRYPTION_DATA_KEY, EncryptionData.getAndValidateEncryptionData(
-                encryptionDataKey, encryptedBlobAsyncClient.isEncryptionRequired()));
+            EncryptionData encryptionData = EncryptionData.getAndValidateEncryptionData(
+                encryptionDataKey, encryptedBlobAsyncClient.isEncryptionRequired());
+
+            context = context.addData(ENCRYPTION_DATA_KEY, encryptionData)
+                .addData(Constants.ADJUSTED_BLOB_LENGTH_KEY,
+                    EncryptedBlobLength.computeAdjustedBlobLength(encryptionData, initialProperties.getBlobSize()));
         }
         return context;
     }
@@ -572,5 +576,4 @@ public BlobQueryResponse queryWithResponse(BlobQueryOptions queryOptions,
     BlobClientSideEncryptionOptions getClientSideEncryptionOptions() {
         return encryptedBlobAsyncClient.getClientSideEncryptionOptions();
     }
-
 }

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

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.storage.blob.specialized.cryptography;
+
+import com.azure.core.util.logging.ClientLogger;
+
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.ENCRYPTION_PROTOCOL_V1;
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.ENCRYPTION_PROTOCOL_V2;
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.ENCRYPTION_PROTOCOL_V2_1;
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.NONCE_LENGTH;
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.TAG_LENGTH;
+
+/**
+ * This class provides helper methods for adjusting encrypted downloads.
+ */
+final class EncryptedBlobLength {
+    private static final ClientLogger LOGGER = new ClientLogger(EncryptedBlobLength.class);
+
+    static Long computeAdjustedBlobLength(EncryptionData encryptionData, Long encryptedLength) {
+        switch (encryptionData.getEncryptionAgent().getProtocol()) {
+            /*
+             Technically, the total unencrypted length may be different for v1,
+             but because this helper method is only used for partitioning ranged downloads,
+             the size does not need to be adjusted for v1.
+             */
+            case ENCRYPTION_PROTOCOL_V1:
+                return encryptedLength;
+            case ENCRYPTION_PROTOCOL_V2:
+            case ENCRYPTION_PROTOCOL_V2_1:
+                long regionLength = encryptionData.getEncryptedRegionInfo().getDataLength();
+                long region = (long) Math.ceil((double) encryptedLength / (double) (regionLength + NONCE_LENGTH + TAG_LENGTH));
+                long offset = (NONCE_LENGTH + TAG_LENGTH) * region;
+                return encryptedLength - offset;
+            default:
+                throw LOGGER.logExceptionAsError(new IllegalArgumentException("Unexpected protocol version"));
+        }
+    }
+}

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

@@ -35,7 +35,7 @@ final class EncryptedBlobRange {
      * 0-31 for v1
      * 0-(4mb+12-1)
      */
-    private final int offsetAdjustment;
+    private final Long offsetAdjustment;
 
     /**
      * How many bytes to download, including the adjustments for encryption block boundaries and the IV.
@@ -78,7 +78,7 @@ static EncryptedBlobRange getEncryptedBlobRangeFromHeader(String stringRange, En
     EncryptedBlobRange(BlobRange originalRange, EncryptionData encryptionData) {
         if (originalRange == null) {
             this.originalRange = new BlobRange(0);
-            this.offsetAdjustment = 0;
+            this.offsetAdjustment = 0L;
             this.amountPlaintextToSkip = 0; // In cases where this block is executed, this value does not matter
             return;
         }
@@ -111,7 +111,7 @@ static EncryptedBlobRange getEncryptedBlobRangeFromHeader(String stringRange, En
                     }
                 }
 
-                this.offsetAdjustment = tempOffsetAdjustment;
+                this.offsetAdjustment = (long) tempOffsetAdjustment;
 
                 /*
                 Align adjustedDownloadCount with encryption block boundary at the end of the range. Note that it is impossible
@@ -153,7 +153,7 @@ static EncryptedBlobRange getEncryptedBlobRangeFromHeader(String stringRange, En
                 }
 
                 // Offset adjustment is difference in two starting values
-                this.offsetAdjustment = (int) (originalRange.getOffset() - regionStartOffset);
+                this.offsetAdjustment = (originalRange.getOffset() - regionStartOffset);
 
                 break;
             default:
@@ -171,7 +171,7 @@ BlobRange getOriginalRange() {
     /**
      * @return Offset from beginning of BlobRange, 0-31.
      */
-    int getOffsetAdjustment() {
+    Long getOffsetAdjustment() {
         return this.offsetAdjustment;
     }
 

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

@@ -0,0 +1,99 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.storage.blob.specialized.cryptography;
+
+import com.azure.storage.common.implementation.Constants;
+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;
+
+import java.util.stream.Stream;
+
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.GCM_ENCRYPTION_REGION_LENGTH;
+import static com.azure.storage.blob.specialized.cryptography.CryptographyConstants.NONCE_LENGTH;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class EncryptedBlobLengthTests extends BlobCryptographyTestBase {
+
+    //to prevent having to cast int -> long everywhere
+    private static final long FOUR_MB = 4 * Constants.MB;
+    private static final long SIXTEEN_MB = 16 * Constants.MB;
+    private static final long ONE_MB = Constants.MB;
+    private static final long ONE_KB = Constants.KB;
+
+    @ParameterizedTest
+    @MethodSource("correctAdjustedLengthV1Supplier")
+    public void correctAdjustedLengthV1(Long encryptedLength) {
+        EncryptionData encryptionData = new EncryptionData();
+        encryptionData.setEncryptionAgent(new EncryptionAgent("1.0", null));
+
+        Long newLength = EncryptedBlobLength.computeAdjustedBlobLength(encryptionData, encryptedLength);
+        assertEquals(encryptedLength, newLength);
+    }
+
+    private static Stream<Arguments> correctAdjustedLengthV1Supplier() {
+        return Stream.of(
+            Arguments.of(FOUR_MB),
+            Arguments.of(SIXTEEN_MB)
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("correctAdjustedLengthV2Supplier")
+    public void correctAdjustedLengthV2(Long encryptedLength, Long expectedLength) {
+        EncryptionData encryptionData = new EncryptionData();
+        encryptionData.setEncryptionAgent(new EncryptionAgent("2.0", null));
+        encryptionData.setEncryptedRegionInfo(new EncryptedRegionInfo(GCM_ENCRYPTION_REGION_LENGTH, NONCE_LENGTH));
+
+        Long newLength = EncryptedBlobLength.computeAdjustedBlobLength(encryptionData, encryptedLength);
+        assertEquals(expectedLength, newLength);
+    }
+
+    private static Stream<Arguments> correctAdjustedLengthV2Supplier() {
+        return Stream.of(
+            Arguments.of(FOUR_MB + 28, FOUR_MB),
+            Arguments.of(FOUR_MB + 57, FOUR_MB + 1),
+            Arguments.of(SIXTEEN_MB + 112, SIXTEEN_MB),
+            Arguments.of(28L, 0L),
+            Arguments.of(0L, 0L),
+            Arguments.of(ONE_MB + 28, ONE_MB)
+        );
+    }
+
+    @ParameterizedTest
+    @MethodSource("correctAdjustedLengthVariableRegionSupplier")
+    public void correctAdjustedLengthVariableRegion(Long encryptedLength, Long expectedLength, Long regionLength) {
+        EncryptionData encryptionData = new EncryptionData();
+        encryptionData.setEncryptionAgent(new EncryptionAgent("2.1", null));
+        encryptionData.setEncryptedRegionInfo(new EncryptedRegionInfo(regionLength, NONCE_LENGTH));
+
+        Long newLength = EncryptedBlobLength.computeAdjustedBlobLength(encryptionData, encryptedLength);
+        assertEquals(expectedLength, newLength);
+    }
+
+    private static Stream<Arguments> correctAdjustedLengthVariableRegionSupplier() {
+        return Stream.of(
+            Arguments.of(FOUR_MB + 112, FOUR_MB, ONE_MB),
+            Arguments.of(SIXTEEN_MB + 448, SIXTEEN_MB, ONE_MB),
+            Arguments.of(FOUR_MB + 114688, FOUR_MB, ONE_KB),
+            Arguments.of(SIXTEEN_MB + 458752, SIXTEEN_MB, ONE_KB),
+            Arguments.of(FOUR_MB + 448, FOUR_MB, 256 * ONE_KB),
+            Arguments.of(SIXTEEN_MB + 1792, SIXTEEN_MB, 256 * ONE_KB),
+            Arguments.of(ONE_MB + 28672, ONE_MB, ONE_KB),
+            Arguments.of(ONE_MB + 28, ONE_MB, ONE_MB)
+        );
+    }
+
+    @Test
+    public void badProtocol() {
+        EncryptionData encryptionData = new EncryptionData();
+        encryptionData.setEncryptionAgent(new EncryptionAgent("garbage", null));
+
+        assertThrows(IllegalArgumentException.class, () -> EncryptedBlobLength.computeAdjustedBlobLength(encryptionData,
+                null));
+    }
+
+}

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

@@ -114,8 +114,7 @@ public void encryptedBlobRangeFromEncryptionData(int regionLength) {
 
     @ParameterizedTest
     @MethodSource("provideRanges")
-    public void testAdjustedBlobRange(int originalOffset, long originalCount, int expectedNewOffset, long expectedNewCount) {
-        int regionLength = 16;
+    public void testAdjustedBlobRange(long originalOffset, long originalCount, long expectedNewOffset, long expectedNewCount, int regionLength) {
         EncryptionData encryptionData = new EncryptionData()
             .setEncryptionAgent(new EncryptionAgent(ENCRYPTION_PROTOCOL_V2, EncryptionAlgorithm.AES_GCM_256))
             .setEncryptedRegionInfo(new EncryptedRegionInfo(regionLength, NONCE_LENGTH));
@@ -131,10 +130,11 @@ public void testAdjustedBlobRange(int originalOffset, long originalCount, int ex
 
     private static Stream<Arguments> provideRanges() {
         return Stream.of(
-            Arguments.of(5, 10, 0, 44), // Entirely within a single region, adjustment includes nonce and tag
-            Arguments.of(16, 16, 44, 44), // Exactly one region
-            Arguments.of(15, 35, 0, 176), // Straddles across four regions
-            Arguments.of(32, 15, 88, 44)  // Starts exactly at the second region
+            Arguments.of(5L, 10L, 0L, 44L, 16), // Entirely within a single region, adjustment includes nonce and tag
+            Arguments.of(16L, 16L, 44L, 44L, 16), // Exactly one region
+            Arguments.of(15L, 35L, 0L, 176L, 16), // Straddles across four regions
+            Arguments.of(32L, 15L, 88L, 44L, 16), // Starts exactly at the second region
+            Arguments.of(3674210304L, 10066367L, 8818104720L, 24159312L, 20) // Tests for integer overflow with a very large blob
         );
     }
 }