Add CryptoStream.FlushFinalBlockAsync

dotnet · Jul 21, 2020 · d129af6 · d129af6
1 parent bf0076c
commit d129af6
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 6 deletions.
diff --git a/...es/System.Security.Cryptography.Primitives/ref/System.Security.Cryptography.Primitives.cs b/...es/System.Security.Cryptography.Primitives/ref/System.Security.Cryptography.Primitives.cs
@@ -82,6 +82,7 @@ public override void EndWrite(System.IAsyncResult asyncResult) { }
         public override void Flush() { }
         public override System.Threading.Tasks.Task FlushAsync(System.Threading.CancellationToken cancellationToken) { throw null; }
         public void FlushFinalBlock() { }
+        public System.Threading.Tasks.ValueTask FlushFinalBlockAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
         public override int Read(byte[] buffer, int offset, int count) { throw null; }
         public override System.Threading.Tasks.Task<int> ReadAsync(byte[] buffer, int offset, int count, System.Threading.CancellationToken cancellationToken) { throw null; }
         public override int ReadByte() { throw null; }

diff --git a/.../System.Security.Cryptography.Primitives/src/System/Security/Cryptography/CryptoStream.cs b/.../System.Security.Cryptography.Primitives/src/System/Security/Cryptography/CryptoStream.cs
@@ -100,9 +100,23 @@ public bool HasFlushedFinalBlock
         // byte[] ciphertext = ms.ToArray();
         // cs.Close();
         public void FlushFinalBlock() =>
-            FlushFinalBlockAsync(useAsync: false).AsTask().GetAwaiter().GetResult();
+            FlushFinalBlockAsync(useAsync: false, default).AsTask().GetAwaiter().GetResult();
+
+        /// <summary>
+        /// Asynchronously updates the underlying data source or repository with the
+        /// current state of the buffer, then clears the buffer.
+        /// </summary>
+        /// <param name="cancellationToken">The token to monitor for cancellation requests. The default value is <see cref="CancellationToken.None"/>.</param>
+        /// <returns>A task that represents the asynchronous flush operation.</returns>
+        public ValueTask FlushFinalBlockAsync(CancellationToken cancellationToken = default)
+        {
+            if (cancellationToken.IsCancellationRequested)
+                return ValueTask.FromCanceled(cancellationToken);
+
+            return FlushFinalBlockAsync(useAsync: true, cancellationToken);
+        }
 
-        private async ValueTask FlushFinalBlockAsync(bool useAsync)
+        private async ValueTask FlushFinalBlockAsync(bool useAsync, CancellationToken cancellationToken)
         {
             if (_finalBlockTransformed)
                 throw new NotSupportedException(SR.Cryptography_CryptoStream_FlushFinalBlockTwice);
@@ -116,7 +130,7 @@ private async ValueTask FlushFinalBlockAsync(bool useAsync)
                 byte[] finalBytes = _transform.TransformFinalBlock(_inputBuffer!, 0, _inputBufferIndex);
                 if (useAsync)
                 {
-                    await _stream.WriteAsync(new ReadOnlyMemory<byte>(finalBytes)).ConfigureAwait(false);
+                    await _stream.WriteAsync(new ReadOnlyMemory<byte>(finalBytes), cancellationToken).ConfigureAwait(false);
                 }
                 else
                 {
@@ -129,14 +143,14 @@ private async ValueTask FlushFinalBlockAsync(bool useAsync)
             {
                 if (!innerCryptoStream.HasFlushedFinalBlock)
                 {
-                    await innerCryptoStream.FlushFinalBlockAsync(useAsync).ConfigureAwait(false);
+                    await innerCryptoStream.FlushFinalBlockAsync(useAsync, cancellationToken).ConfigureAwait(false);
                 }
             }
             else
             {
                 if (useAsync)
                 {
-                    await _stream.FlushAsync().ConfigureAwait(false);
+                    await _stream.FlushAsync(cancellationToken).ConfigureAwait(false);
                 }
                 else
                 {
@@ -691,7 +705,7 @@ private async ValueTask DisposeAsyncCore()
             {
                 if (!_finalBlockTransformed)
                 {
-                    await FlushFinalBlockAsync(useAsync: true).ConfigureAwait(false);
+                    await FlushFinalBlockAsync(useAsync: true, default).ConfigureAwait(false);
                 }
 
                 if (!_leaveOpen)

diff --git a/src/libraries/System.Security.Cryptography.Primitives/tests/CryptoStream.cs b/src/libraries/System.Security.Cryptography.Primitives/tests/CryptoStream.cs
@@ -187,6 +187,34 @@ public static void FlushAsync()
             }
         }
 
+        [Fact]
+        public static async Task FlushFinalBlockAsync()
+        {
+            ICryptoTransform encryptor = new IdentityTransform(1, 1, true);
+            using (MemoryStream output = new MemoryStream())
+            using (CryptoStream encryptStream = new CryptoStream(output, encryptor, CryptoStreamMode.Write))
+            {
+                await encryptStream.WriteAsync(new byte[] { 1, 2, 3, 4, 5 }, 0, 5);
+                await encryptStream.FlushFinalBlockAsync();
+                Assert.True(encryptStream.HasFlushedFinalBlock);
+                Assert.Equal(5, output.ToArray().Length);
+            }
+        }
+
+        [Fact]
+        public static async Task FlushFinalBlockAsync_Cancelled()
+        {
+            ICryptoTransform encryptor = new IdentityTransform(1, 1, true);
+            using (MemoryStream output = new MemoryStream())
+            using (CryptoStream encryptStream = new CryptoStream(output, encryptor, CryptoStreamMode.Write))
+            {
+                await encryptStream.WriteAsync(new byte[] { 1, 2, 3, 4, 5 }, 0, 5);
+                ValueTask waitable = encryptStream.FlushFinalBlockAsync(new Threading.CancellationToken(canceled: true));
+                Assert.True(waitable.IsCanceled);
+                Assert.False(encryptStream.HasFlushedFinalBlock);
+            }
+        }
+
         [Fact]
         public static void FlushCalledOnFlushAsync_DeriveClass()
         {