SNOW-1640968 dispose disposable objects

Snowflake.Data.Tests/UnitTests/ConcatenatedStreamTest.cs

@@ -2,13 +2,14 @@
  * Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
  */
 
+using Snowflake.Data.Tests.Util;
+using NUnit.Framework;
+using System;
+using System.IO;
+using System.Text;
+
 namespace Snowflake.Data.Tests.UnitTests
 {
-    using NUnit.Framework;
-    using Snowflake.Data.Core;
-    using System;
-    using System.IO;
-    using System.Text;
 
     [TestFixture]
     class ConcatenatedStreamTest

Snowflake.Data.Tests/UnitTests/SFRemoteStorageClientTest.cs

@@ -71,7 +71,7 @@
         const bool IsAsync = true;
 
         [SetUp]
         public void BeforeTest()
         {
             t_realSourceFilePath = TestNameWithWorker + "_realSrcFilePath.txt";
             t_downloadFileName = TestNameWithWorker + "_mockFileName.txt";
@@ -122,7 +122,7 @@
         }
 
         [TearDown]
         public void AfterTest()
         {
             // Delete temporary files from upload
             if (File.Exists(_fileMetadata.realSrcFilePath))
@@ -527,14 +527,14 @@
 
             // Get encrypted stream from file
             SFEncryptionMetadata encryptionMetadata = new SFEncryptionMetadata();
-            Stream stream = EncryptionProvider.EncryptFile(
+            StreamPair streamPair = EncryptionProvider.EncryptFile(
                 t_downloadFileName,
                 _fileMetadata.encryptionMaterial,
                 encryptionMetadata,
                 FileTransferConfiguration.FromFileMetadata(_fileMetadata));
 
             // Set up the stream and metadata for decryption
-            MockRemoteStorageClient.SetEncryptionData(stream, encryptionMetadata.iv, encryptionMetadata.key);
+            MockRemoteStorageClient.SetEncryptionData(streamPair.MainStream, encryptionMetadata.iv, encryptionMetadata.key);
         }
 
         [Test]

Snowflake.Data/Core/ConcatenatedStream.cs → Snowflake.Data.Tests/Util/ConcatenatedStream.cs

@@ -1,12 +1,12 @@
 /*
- * Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
+ * Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
  */
 
 using System;
 using System.Collections.Generic;
 using System.IO;
 
-namespace Snowflake.Data.Core
+namespace Snowflake.Data.Tests.Util
 {
     /// <summary>
     ///     Used to concat multiple streams without copying. Since we need to preappend '[' and append ']'

Snowflake.Data/Core/Authenticator/KeyPairAuthenticator.cs

@@ -100,36 +100,29 @@
             {
                 try
                 {
-                    PemReader pr = null;
-                    if (null != pkPwd)
+                    using (PemReader pr = CreatePemReader(tr, pkPwd))
                     {
-                        IPasswordFinder ipwdf = new PasswordFinder(pkPwd);
-                        pr = new PemReader(tr, ipwdf);
-                    }
-                    else
-                    {
-                        pr = new PemReader(tr);
-                    }
-
-                    object key = pr.ReadObject();
-                    // Infer what the pem reader is sending back based on the object properties
-                    if (key.GetType().GetProperty("Private") != null)
-                    {
-                        // PKCS1 key
-                        keypair = (AsymmetricCipherKeyPair)key;
-                        rsaParams = DotNetUtilities.ToRSAParameters(
-                        keypair.Private as RsaPrivateCrtKeyParameters);
-                    }
-                    else
-                    {
-                        // PKCS8 key
-                        RsaPrivateCrtKeyParameters pk = (RsaPrivateCrtKeyParameters)key;
-                        rsaParams = DotNetUtilities.ToRSAParameters(pk);
-                        keypair = DotNetUtilities.GetRsaKeyPair(rsaParams);
-                    }
-                    if (keypair == null)
-                    {
-                        throw new Exception("Unknown error.");
+                        object key = pr.ReadObject();
+                        // Infer what the pem reader is sending back based on the object properties
+                        if (key.GetType().GetProperty("Private") != null)
+                        {
+                            // PKCS1 key
+                            keypair = (AsymmetricCipherKeyPair)key;
+                            rsaParams = DotNetUtilities.ToRSAParameters(
+                                keypair.Private as RsaPrivateCrtKeyParameters);
+                        }
+                        else
+                        {
+                            // PKCS8 key
+                            RsaPrivateCrtKeyParameters pk = (RsaPrivateCrtKeyParameters)key;
+                            rsaParams = DotNetUtilities.ToRSAParameters(pk);
+                            keypair = DotNetUtilities.GetRsaKeyPair(rsaParams);
+                        }
+
+                        if (keypair == null)
+                        {
+                            throw new Exception("Unknown error.");
+                        }
                     }
                 }
                 catch (Exception e)
@@ -207,6 +200,19 @@
             return jwtToken;
         }
 
+        private PemReader CreatePemReader(TextReader textReader, string privateKeyPassword)
+        {
+            if (null != privateKeyPassword)
+            {
+                IPasswordFinder ipwdf = new PasswordFinder(privateKeyPassword);
+                return new PemReader(textReader, ipwdf);
+            }
+            else
+            {
+                return new PemReader(textReader);
+            }
+        }
+
         /// <summary>
         /// Helper class to handle the password for the certificate if there is one.
         /// </summary>

Snowflake.Data/Core/FileTransfer/EncryptionProvider.cs

@@ -10,7 +10,7 @@
 namespace Snowflake.Data.Core.FileTransfer
 {
     /// <summary>
-    /// The encryption materials. 
+    /// The encryption materials.
     /// </summary>
     internal class MaterialDescriptor
     {
@@ -22,7 +22,7 @@ internal class MaterialDescriptor
     }
 
     /// <summary>
-    /// The encryptor/decryptor for PUT/GET files. 
+    /// The encryptor/decryptor for PUT/GET files.
     /// </summary>
     class EncryptionProvider
     {
@@ -43,7 +43,7 @@ class EncryptionProvider
         /// <param name="encryptionMetadata">Store the encryption metadata into</param>
         /// <param name="transferConfiguration">Contains parameters used during encryption process</param>
         /// <returns>The encrypted bytes of the file to upload.</returns>
-        public static Stream EncryptFile(
+        public static StreamPair EncryptFile(
             string inFile,
             PutGetEncryptionMaterial encryptionMaterial,
             SFEncryptionMetadata encryptionMetadata,
@@ -63,7 +63,7 @@ public static Stream EncryptFile(
         /// <param name="encryptionMetadata">Store the encryption metadata into</param>
         /// <param name="transferConfiguration">Contains parameters used during encryption process</param>
         /// <returns>The encrypted bytes of the file to upload.</returns>
-        public static Stream EncryptStream(
+        public static StreamPair EncryptStream(
             Stream inputStream,
             PutGetEncryptionMaterial encryptionMaterial,
             SFEncryptionMetadata encryptionMetadata,
@@ -89,7 +89,7 @@ public static Stream EncryptStream(
 
             // Encrypt file key
             byte[] encryptedFileKey = encryptFileKey(decodedMasterKey, keyData);
-            
+
             // Store encryption metadata information
             MaterialDescriptor matDesc = new MaterialDescriptor
             {
@@ -116,17 +116,21 @@ public static Stream EncryptStream(
         /// <returns>The encrypted key.</returns>
         private static byte[] encryptFileKey(byte[] masterKey, byte[] unencryptedFileKey)
         {
-            Aes aes = Aes.Create();            
-            aes.Key = masterKey;
-            aes.Mode = CipherMode.ECB;
-            aes.Padding = PaddingMode.PKCS7;
-
-            MemoryStream cipherStream = new MemoryStream();
-            CryptoStream cryptoStream = new CryptoStream(cipherStream, aes.CreateEncryptor(), CryptoStreamMode.Write);            
-            cryptoStream.Write(unencryptedFileKey, 0, unencryptedFileKey.Length);
-            cryptoStream.FlushFinalBlock();
+            using (Aes aes = Aes.Create())
+            {
+                aes.Key = masterKey;
+                aes.Mode = CipherMode.ECB;
+                aes.Padding = PaddingMode.PKCS7;
 
-            return cipherStream.ToArray();
+                using (MemoryStream cipherStream = new MemoryStream())
+                using (var encryptor = aes.CreateEncryptor())
+                using (CryptoStream cryptoStream = new CryptoStream(cipherStream, encryptor, CryptoStreamMode.Write))
+                {
+                    cryptoStream.Write(unencryptedFileKey, 0, unencryptedFileKey.Length);
+                    cryptoStream.FlushFinalBlock();
+                    return cipherStream.ToArray();
+                }
+            }
         }
 
         /// <summary>
@@ -137,32 +141,41 @@ private static byte[] encryptFileKey(byte[] masterKey, byte[] unencryptedFileKey
         /// <param name="iv">The encryption IV or null if it needs to be generated.</param>
         /// <param name="transferConfiguration">Contains parameters used during encryption process</param>
         /// <returns>The encrypted bytes.</returns>
-        private static Stream CreateEncryptedBytesStream(
+        private static StreamPair CreateEncryptedBytesStream(
             Stream inputStream,
             byte[] key,
             byte[] iv,
             FileTransferConfiguration transferConfiguration)
         {
-            Aes aes = Aes.Create();            
-            aes.Key = key;
-            aes.Mode = CipherMode.CBC;
-            aes.Padding = PaddingMode.PKCS7;
-            aes.IV = iv;
-            inputStream.Position = 0;
-
-            var targetStream = new FileBackedOutputStream(transferConfiguration.MaxBytesInMemory, transferConfiguration.TempDir);
-            CryptoStream cryptoStream = new CryptoStream(targetStream, aes.CreateEncryptor(), CryptoStreamMode.Write);
-            byte[] buffer = new byte[transferConfiguration.MaxBytesInMemory];
-            int bytesRead;
-            while ((bytesRead = inputStream.Read(buffer, 0, buffer.Length)) > 0)
+            using (Aes aes = Aes.Create())
             {
-                cryptoStream.Write(buffer, 0, bytesRead);
+                aes.Key = key;
+                aes.Mode = CipherMode.CBC;
+                aes.Padding = PaddingMode.PKCS7;
+                aes.IV = iv;
+                inputStream.Position = 0;
+
+                var targetStream = new FileBackedOutputStream(transferConfiguration.MaxBytesInMemory, transferConfiguration.TempDir);
+                using (var encryptor = aes.CreateEncryptor())
+                {
+                    CryptoStream cryptoStream = new CryptoStream(targetStream, encryptor, CryptoStreamMode.Write);
+                    byte[] buffer = new byte[transferConfiguration.MaxBytesInMemory];
+                    int bytesRead;
+                    while ((bytesRead = inputStream.Read(buffer, 0, buffer.Length)) > 0)
+                    {
+                        cryptoStream.Write(buffer, 0, bytesRead);
+                    }
+                    cryptoStream.FlushFinalBlock();
+
+                    return new StreamPair
+                    {
+                        MainStream = targetStream,
+                        HelperStream = cryptoStream // cryptoStream cannot be closed here because it would close target stream as well
+                    };
+                }
             }
-            cryptoStream.FlushFinalBlock();
-
-            return targetStream;
         }
-        
+
         /// <summary>
         /// Decrypt data and write to the outStream.
         /// </summary>
@@ -195,14 +208,15 @@ public static string DecryptFile(
             byte[] decryptedFileKey = decryptFileKey(decodedMasterKey, keyBytes);
 
             // Create key decipher with decoded key and AES ECB
-            using (var decryptedBytesStream = CreateDecryptedBytesStream(
+            using (var decryptedBytesStreamPair = CreateDecryptedBytesStream(
                        inFile,
                        decryptedFileKey,
                        ivBytes,
                        transferConfiguration))
             {
                 using (var decryptedFileStream = File.Create(tempFileName))
                 {
+                    var decryptedBytesStream = decryptedBytesStreamPair.MainStream;
                     decryptedBytesStream.Position = 0;
                     decryptedBytesStream.CopyTo(decryptedFileStream);
                 }
@@ -218,17 +232,22 @@ public static string DecryptFile(
         /// <returns>The encrypted key.</returns>
         private static byte[] decryptFileKey(byte[] masterKey, byte[] unencryptedFileKey)
         {
-            Aes aes = Aes.Create();
-            aes.Key = masterKey;
-            aes.Mode = CipherMode.ECB;
-            aes.Padding = PaddingMode.PKCS7;
+            using (Aes aes = Aes.Create())
+            {
+                aes.Key = masterKey;
+                aes.Mode = CipherMode.ECB;
+                aes.Padding = PaddingMode.PKCS7;
 
-            MemoryStream cipherStream = new MemoryStream();
-            CryptoStream cryptoStream = new CryptoStream(cipherStream, aes.CreateDecryptor(), CryptoStreamMode.Write);
-            cryptoStream.Write(unencryptedFileKey, 0, unencryptedFileKey.Length);
-            cryptoStream.FlushFinalBlock();
+                using (MemoryStream cipherStream = new MemoryStream())
+                using (var decryptor = aes.CreateDecryptor())
+                using (CryptoStream cryptoStream = new CryptoStream(cipherStream, decryptor, CryptoStreamMode.Write))
+                {
+                    cryptoStream.Write(unencryptedFileKey, 0, unencryptedFileKey.Length);
+                    cryptoStream.FlushFinalBlock();
 
-            return cipherStream.ToArray();
+                    return cipherStream.ToArray();
+                }
+            }
         }
 
         /// <summary>
@@ -238,28 +257,36 @@ private static byte[] decryptFileKey(byte[] masterKey, byte[] unencryptedFileKey
         /// <param name="key">The encryption key.</param>
         /// <param name="iv">The encryption IV or null if it needs to be generated.</param>
         /// <returns>The decrypted bytes stream</returns>
-        private static Stream CreateDecryptedBytesStream(
+        private static StreamPair CreateDecryptedBytesStream(
             string inFile,
             byte[] key,
             byte[] iv,
             FileTransferConfiguration transferConfiguration)
         {
-            Aes aes = Aes.Create();
-            aes.Key = key;
-            aes.Mode = CipherMode.CBC;
-            aes.Padding = PaddingMode.PKCS7;
-            aes.IV = iv;
-
-            var targetStream = new FileBackedOutputStream(transferConfiguration.MaxBytesInMemory, transferConfiguration.TempDir);
-            CryptoStream cryptoStream = new CryptoStream(targetStream, aes.CreateDecryptor(), CryptoStreamMode.Write);
-
-            using(Stream inStream = File.OpenRead(inFile))
+            using (Aes aes = Aes.Create())
             {
-                inStream.CopyTo(cryptoStream);
-            }
-            cryptoStream.FlushFinalBlock();
+                aes.Key = key;
+                aes.Mode = CipherMode.CBC;
+                aes.Padding = PaddingMode.PKCS7;
+                aes.IV = iv;
 
-            return targetStream;
+                var targetStream = new FileBackedOutputStream(transferConfiguration.MaxBytesInMemory, transferConfiguration.TempDir);
+                using (var decryptor = aes.CreateDecryptor())
+                {
+                    CryptoStream cryptoStream = new CryptoStream(targetStream, decryptor, CryptoStreamMode.Write);
+                    using (Stream inStream = File.OpenRead(inFile))
+                    {
+                        inStream.CopyTo(cryptoStream);
+                    }
+                    cryptoStream.FlushFinalBlock();
+
+                    return new StreamPair
+                    {
+                        MainStream = targetStream,
+                        HelperStream = cryptoStream // cryptoStream cannot be closed here because it would close target stream as well
+                    };
+                }
+            }
         }
     }
 }

Snowflake.Data/Core/FileTransfer/FileBackedOutputStream.cs

@@ -89,7 +89,7 @@ internal bool IsUsingFileOutputStream()
         {
             return _fileOutputStream != null;
         }
-        
+
         private void SwitchFromMemoryToTempFile()
         {
             _fileName = GenerateTempFilePath();

Snowflake.Data/Core/FileTransfer/StorageClient/SFLocalStorageUtil.cs

@@ -62,7 +62,7 @@ internal static void DownloadOneFile(SFFileMetadata fileMetadata)
             }
 
             // Create stream object for reader and writer
-            Stream stream = new MemoryStream(File.ReadAllBytes(realSrcFilePath));
+            using (Stream stream = new MemoryStream(File.ReadAllBytes(realSrcFilePath)))
             using (var fileStream = File.Create(output))
             {
                 // Write file