Client-side Encryption Added Test Coverage (Azure#24797)

* initial extra test

* Forced-multipart upload/download encryption tests

* undid formatting change

* added stress test (ignored)

Co-authored-by: jschrepp-MSFT <41338290+jschrepp-MSFT@users.noreply.github.com>

Author: jaschrep-msft

sdk/storage/Azure.Storage.Blobs/tests/ClientSideEncryptionTests.cs

@@ -218,6 +218,51 @@ private static byte[] Xor(byte[] a, byte[] b)
             return result;
         }
 
+        /// <summary>
+        /// Get Content Encryption Key and IV from uploaded/encrypted blob to determine expected ciphertext.
+        /// </summary>
+        /// <param name="plaintext">
+        /// Data to encrypt.
+        /// </param>
+        /// <param name="properties">
+        /// BlobProperties containing the wrapped CEK and IV to use for replicating encryption steps.
+        /// </param>
+        /// <param name="keyEncryptionKey">
+        /// KEK used to unwrap the CEK in <paramref name="properties"/>.
+        /// </param>
+        /// <returns>Expected encrypted data with the given CEK and IV.</returns>
+        private async Task<byte[]> ReplicateEncryption(byte[] plaintext, BlobProperties properties, IKeyEncryptionKey keyEncryptionKey)
+        {
+            // encrypt original data manually for comparison
+            if (!properties.Metadata.TryGetValue(Constants.ClientSideEncryption.EncryptionDataKey, out string serialEncryptionData))
+            {
+                Assert.Fail("No encryption metadata present.");
+            }
+            EncryptionData encryptionMetadata = EncryptionDataSerializer.Deserialize(serialEncryptionData);
+            Assert.NotNull(encryptionMetadata, "Never encrypted data.");
+
+            var explicitlyUnwrappedKey = IsAsync // can't instrument this
+                ? await keyEncryptionKey.UnwrapKeyAsync(s_algorithmName, encryptionMetadata.WrappedContentKey.EncryptedKey, s_cancellationToken).ConfigureAwait(false)
+                : keyEncryptionKey.UnwrapKey(s_algorithmName, encryptionMetadata.WrappedContentKey.EncryptedKey, s_cancellationToken);
+
+            return EncryptData(
+                plaintext,
+                explicitlyUnwrappedKey,
+                encryptionMetadata.ContentEncryptionIV);
+        }
+
+        /// <summary>
+        /// Download a blob without decrypting it.
+        /// </summary>
+        /// <param name="blob">Encrypted blob to download.</param>
+        /// <returns>Ciphertext.</returns>
+        private async Task<byte[]> DownloadBypassDecryption(BlobClient blob)
+        {
+            var encryptedDataStream = new MemoryStream();
+            await InstrumentClient(new BlobClient(blob.Uri, Tenants.GetNewSharedKeyCredentials())).DownloadToAsync(encryptedDataStream, cancellationToken: s_cancellationToken);
+            return encryptedDataStream.ToArray();
+        }
+
         [Test]
         [LiveOnly]
         public void CanSwapKey()
@@ -259,7 +304,7 @@ public void CanSwapKey()
         [LiveOnly] // cannot seed content encryption key
         public async Task UploadAsync(long dataSize)
         {
-            var data = GetRandomBuffer(dataSize);
+            var plaintext = GetRandomBuffer(dataSize);
             var mockKey = GetIKeyEncryptionKey().Object;
             await using (var disposable = await GetTestContainerEncryptionAsync(
                 new ClientSideEncryptionOptions(ClientSideEncryptionVersion.V1_0)
@@ -272,28 +317,53 @@ public async Task UploadAsync(long dataSize)
                 var blob = InstrumentClient(disposable.Container.GetBlobClient(blobName));
 
                 // upload with encryption
-                await blob.UploadAsync(new MemoryStream(data), cancellationToken: s_cancellationToken);
+                await blob.UploadAsync(new MemoryStream(plaintext), cancellationToken: s_cancellationToken);
 
-                // download without decrypting
-                var encryptedDataStream = new MemoryStream();
-                await InstrumentClient(new BlobClient(blob.Uri, Tenants.GetNewSharedKeyCredentials())).DownloadToAsync(encryptedDataStream, cancellationToken: s_cancellationToken);
-                var encryptedData = encryptedDataStream.ToArray();
+                var encryptedData = await DownloadBypassDecryption(blob);
+                byte[] expectedEncryptedData = await ReplicateEncryption(plaintext, await blob.GetPropertiesAsync(), mockKey);
 
-                // encrypt original data manually for comparison
-                if (!(await blob.GetPropertiesAsync()).Value.Metadata.TryGetValue(Constants.ClientSideEncryption.EncryptionDataKey, out string serialEncryptionData))
+                // compare data
+                Assert.AreEqual(expectedEncryptedData, encryptedData);
+            }
+        }
+
+        [TestCase(1)]
+        [TestCase(2)]
+        [TestCase(4)]
+        [TestCase(8)]
+        [LiveOnly] // cannot seed content encryption key
+        public async Task UploadAsyncSplit(int concurrency)
+        {
+            int blockSize = Constants.KB;
+            int dataSize = 16 * Constants.KB;
+            var plaintext = GetRandomBuffer(dataSize);
+            var mockKey = GetIKeyEncryptionKey().Object;
+            await using (var disposable = await GetTestContainerEncryptionAsync(
+                new ClientSideEncryptionOptions(ClientSideEncryptionVersion.V1_0)
                 {
-                    Assert.Fail("No encryption metadata present.");
-                }
-                EncryptionData encryptionMetadata = EncryptionDataSerializer.Deserialize(serialEncryptionData);
-                Assert.NotNull(encryptionMetadata, "Never encrypted data.");
+                    KeyEncryptionKey = mockKey,
+                    KeyWrapAlgorithm = s_algorithmName
+                }))
+            {
+                var blobName = GetNewBlobName();
+                var blob = InstrumentClient(disposable.Container.GetBlobClient(blobName));
+
+                // upload with encryption
+                await blob.UploadAsync(
+                    new MemoryStream(plaintext),
+                    new BlobUploadOptions
+                    {
+                        TransferOptions = new StorageTransferOptions
+                        {
+                            InitialTransferSize = blockSize,
+                            MaximumTransferSize = blockSize,
+                            MaximumConcurrency = concurrency
+                        }
+                    },
+                    cancellationToken: s_cancellationToken);
 
-                var explicitlyUnwrappedKey = IsAsync // can't instrument this
-                    ? await mockKey.UnwrapKeyAsync(s_algorithmName, encryptionMetadata.WrappedContentKey.EncryptedKey, s_cancellationToken).ConfigureAwait(false)
-                    : mockKey.UnwrapKey(s_algorithmName, encryptionMetadata.WrappedContentKey.EncryptedKey, s_cancellationToken);
-                byte[] expectedEncryptedData = EncryptData(
-                    data,
-                    explicitlyUnwrappedKey,
-                    encryptionMetadata.ContentEncryptionIV);
+                var encryptedData = await DownloadBypassDecryption(blob);
+                byte[] expectedEncryptedData = await ReplicateEncryption(plaintext, await blob.GetPropertiesAsync(), mockKey);
 
                 // compare data
                 Assert.AreEqual(expectedEncryptedData, encryptedData);
@@ -340,6 +410,54 @@ await blob.DownloadToAsync(stream,
             }
         }
 
+        [TestCase(1)]
+        [TestCase(2)]
+        [TestCase(4)]
+        [TestCase(8)]
+        [LiveOnly] // cannot seed content encryption key
+        public async Task RoundtripSplitAsync(int concurrency)
+        {
+            int blockSize = Constants.KB;
+            int dataSize = 16 * Constants.KB;
+
+            var data = GetRandomBuffer(dataSize);
+            var mockKey = GetIKeyEncryptionKey();
+            var mockKeyResolver = GetIKeyEncryptionKeyResolver(mockKey.Object).Object;
+            var transferOptions = new StorageTransferOptions
+            {
+                InitialTransferSize = blockSize,
+                MaximumTransferSize = blockSize,
+                MaximumConcurrency = concurrency
+            };
+            await using (var disposable = await GetTestContainerEncryptionAsync(
+                new ClientSideEncryptionOptions(ClientSideEncryptionVersion.V1_0)
+                {
+                    KeyEncryptionKey = mockKey.Object,
+                    KeyResolver = mockKeyResolver,
+                    KeyWrapAlgorithm = s_algorithmName
+                }))
+            {
+                var blob = InstrumentClient(disposable.Container.GetBlobClient(GetNewBlobName()));
+
+                // upload with encryption
+                await blob.UploadAsync(new MemoryStream(data), transferOptions: transferOptions, cancellationToken: s_cancellationToken);
+
+                // download with decryption
+                byte[] downloadData;
+                using (var stream = new MemoryStream())
+                {
+                    await blob.DownloadToAsync(stream,
+                        transferOptions: transferOptions,
+                        cancellationToken: s_cancellationToken);
+                    downloadData = stream.ToArray();
+                }
+
+                // compare data
+                Assert.AreEqual(data, downloadData);
+                VerifyUnwrappedKeyWasCached(mockKey);
+            }
+        }
+
         [TestCase(Constants.MB, 64*Constants.KB)]
         [TestCase(Constants.MB, Constants.MB)]
         [TestCase(Constants.MB, 4*Constants.MB)]
@@ -711,7 +829,7 @@ public async Task AppropriateRangeDownloadOnPlaintext(int rangeOffset, int? rang
         [Test]
         [LiveOnly] // cannot seed content encryption key
         [Ignore("stress test")]
-        public async Task StressAsync()
+        public async Task StressManyBlobsAsync()
         {
             static async Task<byte[]> RoundTripData(BlobClient client, byte[] data)
             {
@@ -755,6 +873,48 @@ static async Task<byte[]> RoundTripData(BlobClient client, byte[] data)
             }
         }
 
+        [Test]
+        [LiveOnly] // cannot seed content encryption key
+        [Ignore("stress test")]
+        public async Task StressLargeBlobAsync()
+        {
+            const int dataSize = 100 * Constants.MB;
+            const int blockSize = 8 * Constants.MB;
+
+            var data = GetRandomBuffer(dataSize);
+            var mockKey = GetIKeyEncryptionKey().Object;
+            var mockKeyResolver = GetIKeyEncryptionKeyResolver(mockKey).Object;
+            var transferOptions = new StorageTransferOptions
+            {
+                InitialTransferSize = blockSize,
+                MaximumTransferSize = blockSize,
+                MaximumConcurrency = 100
+            };
+            await using (var disposable = await GetTestContainerEncryptionAsync(
+                new ClientSideEncryptionOptions(ClientSideEncryptionVersion.V1_0)
+                {
+                    KeyEncryptionKey = mockKey,
+                    KeyResolver = mockKeyResolver,
+                    KeyWrapAlgorithm = s_algorithmName
+                }))
+            {
+                var client = disposable.Container.GetBlobClient(GetNewBlobName());
+                using (var dataStream = new MemoryStream(data))
+                {
+                    await client.UploadAsync(dataStream, transferOptions: transferOptions, cancellationToken: s_cancellationToken);
+                }
+
+                byte[] downloadResult;
+                using (var downloadStream = new MemoryStream())
+                {
+                    await client.DownloadToAsync(downloadStream, transferOptions: transferOptions, cancellationToken: s_cancellationToken);
+                    downloadResult = downloadStream.ToArray();
+                }
+
+                Assert.AreEqual(data, downloadResult);
+            }
+        }
+
         [Test]
         [LiveOnly]
         public async Task EncryptedReuploadSuccess()