[BUG FIX] Client-side encryption respects StorageTransferOptions.InitialTransferSize (Azure#25392)

* encryption respects storage transfer options

* more tests

* resolved internal access compilation conflict

* fixed test

* changelog

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

Author: jaschrep-msft

sdk/storage/Azure.Storage.Blobs/CHANGELOG.md

@@ -6,6 +6,7 @@
 - Added support for blob names container invalid XML characters.
 - Added support for updating the key encryption key on a client-side encrypted blob.
 - Fixed a bug where BlobClient.Upload() and UploadAsync() when using client-side encryption would modify the Dictionary instance passed by the caller for blob metadata.
+- Fixed a bug where BlobClient.Upload() and UploadAsync() when using client-side encryption would not respect StorageTransferOptions.InitialTransferSize.
 
 ## 12.11.0-beta.1 (2021-11-03)
 - Added support for service version 2020-12-06.

sdk/storage/Azure.Storage.Blobs/src/BlobClient.cs

@@ -1622,14 +1622,20 @@ internal async Task<Response<BlobContentInfo>> StagedUploadInternal(
             bool async = true,
             CancellationToken cancellationToken = default)
         {
+            long? expectedContentLength = null;
             if (UsingClientSideEncryption)
             {
                 if (UsingClientSideEncryption && options.TransactionalHashingOptions != default)
                 {
                     throw Errors.TransactionalHashingNotSupportedWithClientSideEncryption();
                 }
 
-                // content is now unseekable, so PartitionedUploader will be forced to do a buffered multipart upload
+                // if content length was known, we retain that for dividing REST requests appropriately
+                expectedContentLength = content.GetLengthOrDefault();
+                if (expectedContentLength.HasValue)
+                {
+                    expectedContentLength = ClientSideEncryptor.ExpectedCiphertextLength(expectedContentLength.Value);
+                }
                 (content, options.Metadata) = await new BlobClientSideEncryptor(new ClientSideEncryptor(ClientSideEncryption))
                     .ClientSideEncryptInternal(content, options.Metadata, async, cancellationToken).ConfigureAwait(false);
             }
@@ -1641,6 +1647,7 @@ internal async Task<Response<BlobContentInfo>> StagedUploadInternal(
 
             return await uploader.UploadInternal(
                 content,
+                expectedContentLength,
                 options,
                 options?.ProgressHandler,
                 async,

sdk/storage/Azure.Storage.Blobs/src/BlockBlobClient.cs

@@ -533,6 +533,7 @@ public virtual Response<BlobContentInfo> Upload(
 
             return uploader.UploadInternal(
                 content,
+                expectedContentLength: default,
                 options,
                 options.ProgressHandler,
                 async: false,
@@ -589,6 +590,7 @@ public virtual async Task<Response<BlobContentInfo>> UploadAsync(
 
             return await uploader.UploadInternal(
                 content,
+                expectedContentLength: default,
                 options,
                 options.ProgressHandler,
                 async: true,

sdk/storage/Azure.Storage.Blobs/tests/ClientBuilderExtensions.InternalOnly.cs

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Storage.Blobs.Specialized;
+using BlobsClientBuilder = Azure.Storage.Test.Shared.ClientBuilder<
+    Azure.Storage.Blobs.BlobServiceClient,
+    Azure.Storage.Blobs.BlobClientOptions>;
+
+/**
+ * Azure.Storage.Blobs.Batch (and any other sub-package of blobs) finds itself in the situation
+ * where it cannot have access to Azure.Storage.Blobs internals. It requires access to
+ * Azure.Storage.Blobs.Batch internals, which shares compile-includes with Azure.Storage.Blobs.
+ * Some client builder extensions require access to the internals of blobs, but batch needs
+ * general access to client builder extensions. This file contains extensions that require access
+ * to blobs internals, and would cause compile conflicts were other packages to have access to those
+ * internals.
+ */
+namespace Azure.Storage.Blobs.Tests
+{
+    public static partial class ClientBuilderExtensions
+    {
+        public static BlockBlobClient ToBlockBlobClient(
+            this BlobsClientBuilder clientBuilder,
+            BlobBaseClient client)
+            => clientBuilder.AzureCoreRecordedTestBase.InstrumentClient(new BlockBlobClient(client.Uri, client.ClientConfiguration));
+    }
+}

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

@@ -5,14 +5,15 @@
 using System.Collections.Generic;
 using System.Threading.Tasks;
 using Azure.Storage.Blobs.Models;
+using Azure.Storage.Blobs.Specialized;
 using Azure.Storage.Test.Shared;
 using BlobsClientBuilder = Azure.Storage.Test.Shared.ClientBuilder<
     Azure.Storage.Blobs.BlobServiceClient,
     Azure.Storage.Blobs.BlobClientOptions>;
 
 namespace Azure.Storage.Blobs.Tests
 {
-    public static class ClientBuilderExtensions
+    public static partial class ClientBuilderExtensions
     {
         /// <summary>
         /// Creates a new <see cref="ClientBuilder{TServiceClient, TServiceClientOptions}"/>

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

@@ -1058,6 +1058,61 @@ await encryptedBlobClient.UploadAsync(
             }
         }
 
+        [Test]
+        [LiveOnly]
+        /// <summary>
+        /// Crypto transform streams are unseekable and have no <see cref="Stream.Length"/>.
+        /// When length is unknown, <see cref="PartitionedUploader{TServiceSpecificArgs, TCompleteUploadReturn}"/>
+        /// doesn't even attempt a one-shot upload.
+        /// This tests if we correctly inform the uploader of an expected stream length so it
+        /// can respect the given <see cref="StorageTransferOptions"/>.
+        /// </summary>
+        public async Task PutBlobPutBlockSwitch([Values(true, false)] bool oneshot)
+        {
+            const int dataSize = 1 * Constants.KB;
+
+            // Arrange
+            byte[] data = GetRandomBuffer(dataSize);
+            int transferSize = oneshot
+                    ? 2 * dataSize // big enough for put blob even after AES-CBC PKCS7 padding
+                    : dataSize / 2;
+            StorageTransferOptions transferOptions = new StorageTransferOptions
+            {
+                InitialTransferSize = transferSize,
+                MaximumTransferSize = transferSize
+            };
+
+            IKeyEncryptionKey key = GetIKeyEncryptionKey().Object;
+            await using var disposable = await GetTestContainerEncryptionAsync(
+                new ClientSideEncryptionOptions(ClientSideEncryptionVersion.V1_0)
+                {
+                    KeyEncryptionKey = key,
+                    KeyWrapAlgorithm = s_algorithmName
+                });
+            var blob = disposable.Container.GetBlobClient(GetNewBlobName());
+
+            // Act
+            await blob.UploadAsync(
+                new MemoryStream(data),
+                new BlobUploadOptions { TransferOptions = transferOptions },
+                cancellationToken: s_cancellationToken);
+
+            // Assert
+            Assert.IsTrue(await blob.ExistsAsync());
+            Assert.Greater((await blob.GetPropertiesAsync()).Value.ContentLength, 0);
+            // block list will return empty when putblob was used
+            BlockList blockList = await BlobsClientBuilder.ToBlockBlobClient(blob).GetBlockListAsync();
+            Assert.IsEmpty(blockList.UncommittedBlocks);
+            if (oneshot)
+            {
+                Assert.IsEmpty(blockList.CommittedBlocks);
+            }
+            else
+            {
+                Assert.IsNotEmpty(blockList.CommittedBlocks);
+            }
+        }
+
         [RecordedTest]
         public void CanGenerateSas_WithClientSideEncryptionOptions_True()
         {

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

@@ -318,6 +318,7 @@ private async Task<Response<BlobContentInfo>> InvokeUploadAsync(PartitionedUploa
         {
             return await uploader.UploadInternal(
                 content,
+                expectedContentLength: default,
                 new BlobUploadOptions
                 {
                     HttpHeaders = s_blobHttpHeaders,

sdk/storage/Azure.Storage.Common/src/Shared/ClientsideEncryption/ClientSideEncryptor.cs

@@ -21,6 +21,14 @@ public ClientSideEncryptor(ClientSideEncryptionOptions options)
             _keyWrapAlgorithm = options.KeyWrapAlgorithm;
         }
 
+        public static long ExpectedCiphertextLength(long plaintextLength)
+        {
+            const int aesBlockSizeBytes = 16;
+
+            // pkcs7 padding output length algorithm
+            return plaintextLength + (aesBlockSizeBytes - (plaintextLength % aesBlockSizeBytes));
+        }
+
         /// <summary>
         /// Wraps the given read-stream in a CryptoStream and provides the metadata used to create
         /// that stream.

sdk/storage/Azure.Storage.Common/src/Shared/PartitionedUploader.cs

@@ -11,6 +11,8 @@
 using Azure.Core.Pipeline;
 using Azure.Storage.Shared;
 
+#pragma warning disable SA1402  // File may only contain a single type
+
 namespace Azure.Storage
 {
     internal class PartitionedUploader<TServiceSpecificArgs, TCompleteUploadReturn>
@@ -162,6 +164,7 @@ public PartitionedUploader(
 
         public async Task<Response<TCompleteUploadReturn>> UploadInternal(
             Stream content,
+            long? expectedContentLength,
             TServiceSpecificArgs args,
             IProgress<long> progressHandler,
             bool async,
@@ -184,11 +187,27 @@ public async Task<Response<TCompleteUploadReturn>> UploadInternal(
             // some strategies are unavailable if we don't know the stream length, and some can still work
             // we may introduce separately provided stream lengths in the future for unseekable streams with
             // an expected length
-            long? length = GetLengthOrDefault(content);
+            long? length = expectedContentLength ?? content.GetLengthOrDefault();
 
             // If we know the length and it's small enough
             if (length < _singleUploadThreshold)
             {
+                // may not be seekable. buffer if that's the case
+                if (!content.CanSeek)
+                {
+                    content = await PooledMemoryStream.BufferStreamPartitionInternal(
+                        content,
+                        // we've passed a comparison on length; we know there is a value
+                        length.Value,
+                        length.Value,
+                        // for the purposes of a one-shot, absolutePosition is always zero
+                        absolutePosition: 0,
+                        _arrayPool,
+                        maxArrayPoolRentalSize: default,
+                        async,
+                        cancellationToken).ConfigureAwait(false);
+                }
+
                 // Upload it in a single request
                 return await _singleUploadInternal(
                     content,
@@ -472,25 +491,6 @@ await _uploadPartitionInternal(
             }
         }
 
-        /// <summary>
-        /// Some streams will throw if you try to access their length so we wrap
-        /// the check in a TryGet helper.
-        /// </summary>
-        private static long? GetLengthOrDefault(Stream content)
-        {
-            try
-            {
-                if (content.CanSeek)
-                {
-                    return content.Length - content.Position;
-                }
-            }
-            catch (NotSupportedException)
-            {
-            }
-            return default;
-        }
-
         #region Stream Splitters
         /// <summary>
         /// Partition a stream into a series of blocks buffered as needed by an array pool.
@@ -630,4 +630,26 @@ private static Task<SlicedStream> GetStreamedPartitionInternal(
             => Task.FromResult((SlicedStream)WindowStream.GetWindow(stream, maxCount, absolutePosition));
         #endregion
     }
+
+    internal static partial class StreamExtensions
+    {
+        /// <summary>
+        /// Some streams will throw if you try to access their length so we wrap
+        /// the check in a TryGet helper.
+        /// </summary>
+        public static long? GetLengthOrDefault(this Stream content)
+        {
+            try
+            {
+                if (content.CanSeek)
+                {
+                    return content.Length - content.Position;
+                }
+            }
+            catch (NotSupportedException)
+            {
+            }
+            return default;
+        }
+    }
 }

sdk/storage/Azure.Storage.Common/tests/PartitionedUploaderTests.cs

@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core.Pipeline;
+using Azure.Core.TestFramework;
 using Azure.Storage.Test;
 using Moq;
 using NUnit.Framework;
@@ -134,10 +135,75 @@ public async Task AlwaysUploadSeekableStreamsSequentialUploads(int streamSize, i
                 s_hashingOptions,
                 operationName: s_operationName);
 
-            Response<object> result = await partitionedUploader.UploadInternal(stream.Object, s_objectArgs, s_progress, IsAsync, s_cancellation).ConfigureAwait(false);
+            Response<object> result = await partitionedUploader.UploadInternal(stream.Object, default, s_objectArgs, s_progress, IsAsync, s_cancellation).ConfigureAwait(false);
 
             // assert streams were actually sent to delegates; the delegates themselves threw if conditions weren't met
             Assert.Greater(singleUpload.Invocations.Count + uploadPartition.Invocations.Count, 0);
         }
+
+        [Test]
+        public async Task InterpretsLengthNonSeekableStream([Values(true, false)] bool oneshot)
+        {
+            // Arrange
+            const int dataSize = Constants.KB;
+            const int numPartitions = 2;
+            var data = TestHelper.GetRandomBuffer(dataSize);
+            int blockSize = oneshot ? dataSize * 2 : dataSize / numPartitions;
+
+            var stream = new Mock<MemoryStream>(MockBehavior.Loose, data);
+            stream.CallBase = true;
+
+            // make stream unseekable (can't get length from stream)
+            stream.SetupGet(s => s.CanSeek).Returns(false);
+            stream.SetupGet(s => s.Position).Throws(new NotSupportedException());
+            stream.SetupSet(s => s.Position = default).Throws(new NotSupportedException());
+            stream.Setup(s => s.Seek(It.IsAny<long>(), It.IsAny<SeekOrigin>())).Throws(new NotSupportedException());
+            stream.SetupGet(s => s.Length).Throws(new NotSupportedException());
+            stream.Setup(s => s.SetLength(It.IsAny<long>())).Throws(new NotSupportedException());
+
+            // confirm our stream cannot give a length
+            Assert.Throws<NotSupportedException>(() => _ = stream.Object.Length);
+
+            var createScope = GetMockCreateScope();
+            var initializeDestination = GetMockInitializeDestinationInternal();
+            var singleUpload = GetMockSingleUploadInternal(dataSize);
+            var uploadPartition = GetMockUploadPartitionInternal(blockSize);
+            var commitPartitions = GetMockCommitPartitionedUploadInternal();
+            var partitionedUploader = new PartitionedUploader<object, object>(
+                new PartitionedUploader<object, object>.Behaviors
+                {
+                    Scope = createScope.Object,
+                    InitializeDestination = initializeDestination.Object,
+                    SingleUpload = singleUpload.Object,
+                    UploadPartition = uploadPartition.Object,
+                    CommitPartitionedUpload = commitPartitions.Object
+                },
+                new StorageTransferOptions()
+                {
+                    InitialTransferSize = blockSize,
+                    MaximumTransferSize = blockSize,
+                    MaximumConcurrency = 1
+                },
+                s_hashingOptions,
+                operationName: s_operationName);
+
+            // Act
+            // give uploader an expected content length for unseekable stream
+            Response<object> result = await partitionedUploader.UploadInternal(stream.Object, dataSize, s_objectArgs, s_progress, IsAsync, s_cancellation);
+
+            // Assert
+            if (oneshot)
+            {
+                Assert.AreEqual(1, singleUpload.Invocations.Count);
+                Assert.IsEmpty(uploadPartition.Invocations);
+                Assert.IsEmpty(commitPartitions.Invocations);
+            }
+            else
+            {
+                Assert.IsEmpty(singleUpload.Invocations);
+                Assert.AreEqual(numPartitions, uploadPartition.Invocations.Count);
+                Assert.AreEqual(1, commitPartitions.Invocations.Count);
+            }
+        }
     }
 }