Fix issue #54 with stream being unintentionally closed before all par…

…ts are uploaded. (#56)

src/Amazon.Extensions.S3.Encryption.csproj

@@ -2,7 +2,7 @@
 
     <PropertyGroup>
         <TargetFrameworks>net35;net45;netstandard2.0;netcoreapp3.1</TargetFrameworks>
-        <Version>2.1.1</Version>
+        <Version>2.1.2</Version>
         <GeneratePackageOnBuild>true</GeneratePackageOnBuild>
         <PackageId>Amazon.Extensions.S3.Encryption</PackageId>
         <Title>Amazon S3 Encryption Client for .NET</Title>
@@ -15,8 +15,8 @@
         <PackageIcon>icon.png</PackageIcon>
         <RepositoryUrl>https://github.com/aws/amazon-s3-encryption-client-dotnet/</RepositoryUrl>
         <Company>Amazon Web Services</Company>
-        <AssemblyVersion>2.1.1</AssemblyVersion>
-        <FileVersion>2.1.1</FileVersion>
+        <AssemblyVersion>2.1.2</AssemblyVersion>
+        <FileVersion>2.1.2</FileVersion>
 
         <SignAssembly>true</SignAssembly>
         <AssemblyOriginatorKeyFile>..\public.snk</AssemblyOriginatorKeyFile>

src/Internal/SetupDecryptionHandlerV1.cs

@@ -22,6 +22,7 @@
 using Amazon.Runtime.Internal.Util;
 using Amazon.S3;
 using ThirdParty.Json.LitJson;
+using Amazon.Extensions.S3.Encryption.Util;
 
 namespace Amazon.Extensions.S3.Encryption.Internal
 {
@@ -168,7 +169,7 @@ protected override void UpdateMultipartUploadEncryptionContext(UploadPartRequest
             {
                 object stream = null;
 
-                if (!((Amazon.Runtime.Internal.IAmazonWebServiceRequest) uploadPartRequest).RequestState.TryGetValue(AmazonS3EncryptionClient.S3CryptoStream, out stream))
+                if (!((Amazon.Runtime.Internal.IAmazonWebServiceRequest) uploadPartRequest).RequestState.TryGetValue(Constants.S3CryptoStreamRequestState, out stream))
                     throw new AmazonS3Exception("Cannot retrieve S3 crypto stream from request state, hence cannot get Initialization vector for next uploadPart ");
 
                 var encryptionStream = stream as AESEncryptionUploadPartStream;

src/Internal/SetupDecryptionHandlerV2.cs

@@ -135,7 +135,7 @@ protected override void UpdateMultipartUploadEncryptionContext(UploadPartRequest
             {
                 object stream = null;
 
-                if (!((IAmazonWebServiceRequest) uploadPartRequest).RequestState.TryGetValue(AmazonS3EncryptionClient.S3CryptoStream, out stream))
+                if (!((IAmazonWebServiceRequest) uploadPartRequest).RequestState.TryGetValue(Constants.S3CryptoStreamRequestState, out stream))
                     throw new AmazonS3Exception("Cannot retrieve S3 crypto stream from request state, hence cannot get Initialization vector for next uploadPart ");
 
                 var encryptionStream = stream as AESEncryptionUploadPartStream;

src/Internal/SetupEncryptionHandlerV1.cs

@@ -15,6 +15,7 @@
 
  using Amazon.Runtime;
  using Amazon.S3.Model;
+using Amazon.Extensions.S3.Encryption.Util;
 
  namespace Amazon.Extensions.S3.Encryption.Internal
 {
@@ -153,7 +154,7 @@ protected override void GenerateEncryptedUploadPartRequest(UploadPartRequest req
                 request.InputStream = EncryptionUtils.EncryptRequestUsingInstruction(request.InputStream, instructions);
                 contextForEncryption.IsFinalPart = true;
             }
-            ((Amazon.Runtime.Internal.IAmazonWebServiceRequest)request).RequestState.Add(AmazonS3EncryptionClient.S3CryptoStream, request.InputStream);
+            ((Amazon.Runtime.Internal.IAmazonWebServiceRequest)request).RequestState.Add(Constants.S3CryptoStreamRequestState, request.InputStream);
         }
     }
 }

src/Internal/SetupEncryptionHandlerV2.cs

@@ -13,11 +13,12 @@
  * permissions and limitations under the License.
  */
 
- using Amazon.Extensions.S3.Encryption.Util;
+using Amazon.Extensions.S3.Encryption.Util;
 using Amazon.Runtime;
- using Amazon.S3.Model;
+using Amazon.S3.Model;
+using System;
 
- namespace Amazon.Extensions.S3.Encryption.Internal
+namespace Amazon.Extensions.S3.Encryption.Internal
 {
     /// <summary>
     /// Custom pipeline handler to encrypt the data as it is being uploaded to S3 for AmazonS3EncryptionClientV2.
@@ -37,6 +38,64 @@ public SetupEncryptionHandlerV2(AmazonS3EncryptionClientBase encryptionClient) :
         {
         }
 
+        /// <inheritdoc/>
+        public override void InvokeSync(IExecutionContext executionContext)
+        {
+            try
+            {
+                base.InvokeSync(executionContext);
+            }
+            catch (Exception)
+            {
+                HandleException(executionContext);
+                throw;
+            }
+        }
+
+#if AWS_ASYNC_API
+        /// <inheritdoc/>
+        public override async System.Threading.Tasks.Task<T> InvokeAsync<T>(IExecutionContext executionContext)
+        {
+            try
+            {
+                return await base.InvokeAsync<T>(executionContext);
+            }
+            catch (Exception)
+            {
+                HandleException(executionContext);
+                throw;
+            }
+        }
+#endif
+
+        /// <summary>
+        /// If the crypto stream that is reused for each part has its disposed disabled then the SDK 
+        /// did not close the stream after the exception occurred. This method is called after a exception
+        /// has ocurred and force the crypto stream to be closed.
+        /// </summary>
+        /// <param name="executionContext"></param>
+        private void HandleException(IExecutionContext executionContext)
+        {
+            var request = executionContext.RequestContext.OriginalRequest;
+            var uploadPartRequest = request as UploadPartRequest;
+            if (uploadPartRequest != null)
+            {
+                var contextForEncryption = this.EncryptionClient.CurrentMultiPartUploadKeys[uploadPartRequest.UploadId];
+                if (contextForEncryption == null)
+                    return;
+
+                var aesGcmEncryptStream = contextForEncryption.CryptoStream as AesGcmEncryptStream;
+                if (aesGcmEncryptStream == null)
+                    return;
+
+                if (aesGcmEncryptStream.DisableDispose)
+                {
+                    aesGcmEncryptStream.DisableDispose = false;
+                    aesGcmEncryptStream.Dispose();
+                }
+            }
+        }
+
         /// <inheritdoc/>
         protected override EncryptionInstructions GenerateInstructions(IExecutionContext executionContext)
         {
@@ -149,7 +208,7 @@ protected override void GenerateEncryptedUploadPartRequest(UploadPartRequest req
                 UpdateRequestInputStream(request, contextForEncryption, instructions);
                 contextForEncryption.IsFinalPart = true;
             }
-            ((Amazon.Runtime.Internal.IAmazonWebServiceRequest)request).RequestState.Add(AmazonS3EncryptionClient.S3CryptoStream, request.InputStream);
+            ((Amazon.Runtime.Internal.IAmazonWebServiceRequest)request).RequestState.Add(Constants.S3CryptoStreamRequestState, request.InputStream);
 
         }
 
@@ -167,9 +226,19 @@ private static void UpdateRequestInputStream(UploadPartRequest request, UploadPa
             // Clear the buffer filled for retry request
             var aesGcmEncryptCachingStream = request.InputStream as AesGcmEncryptCachingStream;
             if (aesGcmEncryptCachingStream != null)
-            {
+            {                
                 aesGcmEncryptCachingStream.ClearReadBufferToPosition();
             }
+
+            var aesGcmEncryptStream = request.InputStream as AesGcmEncryptStream;
+            if (aesGcmEncryptStream != null)
+            {
+                // The stream is reused across multi part uploads to maintain the encryption state.
+                // The SDK will attempt to close the stream after the part is upload but setting
+                // DisableDispose to true for anything besides the last part will make the 
+                // disable a noop.
+                aesGcmEncryptStream.DisableDispose = !request.IsLastPart;
+            }
         }
     }
 }

src/Util/AesGcmEncryptStream.cs

@@ -148,5 +148,30 @@ public override async System.Threading.Tasks.Task<int> ReadAsync(byte[] buffer,
             }
         }
 #endif
+
+        /// <summary>
+        /// If set to true the Close and Dispose methods will be a noop. This is necessary in multipart
+        /// upload scenarios when we want the SDK to only dispose the stream on the last part.
+        /// </summary>
+        internal bool DisableDispose { get; set; }
+
+#if !NETSTANDARD
+        /// <inheritdoc/>
+        public override void Close()
+        {
+            if (!DisableDispose)
+            {
+                base.Close();
+            }
+        }
+#else
+        protected override void Dispose(bool disposing)
+        {
+            if (!DisableDispose)
+            {
+                base.Dispose(disposing);
+            }
+        }
+#endif
     }
 }

src/Util/Constants.cs

@@ -0,0 +1,21 @@
+/*
+* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+* 
+* Licensed under the Apache License, Version 2.0 (the "License").
+* You may not use this file except in compliance with the License.
+* A copy of the License is located at
+* 
+*  http://aws.amazon.com/apache2.0
+* 
+* or in the "license" file accompanying this file. This file is distributed
+* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+* express or implied. See the License for the specific language governing
+* permissions and limitations under the License.
+*/
+namespace Amazon.Extensions.S3.Encryption.Util
+{
+    internal class Constants
+    {
+        internal const string S3CryptoStreamRequestState = "S3-Crypto-Stream";
+    }
+}

test/IntegrationTests/_bcl/Utilities/EncryptionTestsUtils.cs

@@ -42,7 +42,7 @@ public static void TestTransferUtility(IAmazonS3 s3EncryptionClient, string buck
 
         public static void TestTransferUtility(IAmazonS3 s3EncryptionClient, IAmazonS3 s3DecryptionClient, string bucketName)
         {
-            var directory = TransferUtilityTests.CreateTestDirectory(10 * TransferUtilityTests.KILO_SIZE);
+            var directory = TransferUtilityTests.CreateTestDirectory(30 * MegaByteSize);
             var keyPrefix = directory.Name;
             var directoryPath = directory.FullName;