SNOW-1640968 dispose disposable objects (#1014)

Author: sfc-gh-knozderko

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

@@ -527,14 +527,14 @@ private void SetUpMockEncryptedFileForDownload()
 
             // 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.Tests/UnitTests/StreamPairTest.cs

@@ -0,0 +1,39 @@
+using System.IO;
+using NUnit.Framework;
+using Moq;
+using Snowflake.Data.Core.FileTransfer;
+
+namespace Snowflake.Data.Tests.UnitTests
+{
+    [TestFixture]
+    public class StreamPairTest
+    {
+        [Test]
+        public void TestCloseBothStreams()
+        {
+            // arrange
+            var mockedMainStream = new Mock<Stream>();
+            var mockedHelperStream = new Mock<Stream>();
+
+            // act
+            using (new StreamPair { MainStream = mockedMainStream.Object, HelperStream = mockedHelperStream.Object }) {}
+
+            // assert
+            mockedMainStream.Verify(stream => stream.Close());
+            mockedHelperStream.Verify(stream => stream.Close());
+        }
+
+        [Test]
+        public void TestCloseMainStreamOnlyWhenHelperStreamNotGiven()
+        {
+            // arrange
+            var mockedMainStream = new Mock<Stream>();
+
+            // act
+            using (new StreamPair { MainStream = mockedMainStream.Object }) {}
+
+            // assert
+            mockedMainStream.Verify(stream => stream.Close());
+        }
+    }
+}

Snowflake.Data/Core/ConcatenatedStream.cs renamed to 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 @@ private string GenerateJwtToken()
             {
                 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 @@ private string GenerateJwtToken()
             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>