From 1aebc4b565f002b6ecf81609c087afe39cf991e9 Mon Sep 17 00:00:00 2001
From: Krzysztof Nozderko <krzysztof.nozderko@snowflake.com>
Date: Tue, 22 Oct 2024 13:58:46 +0200
Subject: [PATCH] fix bug in decoding, add tests for empty aad

---
 .../UnitTests/GcmEncryptionProviderTest.cs    | 41 ++++++++++++++++++-
 .../FileTransfer/GcmEncryptionProvider.cs     |  7 +++-
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/Snowflake.Data.Tests/UnitTests/GcmEncryptionProviderTest.cs b/Snowflake.Data.Tests/UnitTests/GcmEncryptionProviderTest.cs
index a8df88ff8..60c0c2059 100644
--- a/Snowflake.Data.Tests/UnitTests/GcmEncryptionProviderTest.cs
+++ b/Snowflake.Data.Tests/UnitTests/GcmEncryptionProviderTest.cs
@@ -27,6 +27,9 @@ public class GcmEncryptionProviderTest
         private const string InvalidAad = "invalid additional information";
         private static readonly byte[] s_invalidAadBytes = Encoding.UTF8.GetBytes(InvalidAad);
         private static readonly string s_invalidAadBase64 = Convert.ToBase64String(s_invalidAadBytes);
+        private static readonly string s_emptyAad = string.Empty;
+        private static readonly byte[] s_emptyAadBytes = Encoding.UTF8.GetBytes(s_emptyAad);
+        private static readonly string s_emptyAadBase64 = Convert.ToBase64String(s_emptyAadBytes);
         private static readonly PutGetEncryptionMaterial s_encryptionMaterial = new PutGetEncryptionMaterial
         {
             queryStageMasterKey = s_qsmk,
@@ -48,8 +51,8 @@ public void TestEncryptAndDecryptWithoutAad()
             // act
             using (var encryptedStream = GcmEncryptionProvider.Encrypt(
                        s_encryptionMaterial,
-                       encryptionMetadata,
-                       s_fileTransferConfiguration,// this is output parameter
+                       encryptionMetadata, // this is output parameter
+                       s_fileTransferConfiguration,
                        new MemoryStream(s_plainTextBytes),
                        null,
                        null))
@@ -73,6 +76,40 @@ public void TestEncryptAndDecryptWithoutAad()
             }
         }
 
+        [Test]
+        public void TestEncryptAndDecryptWithEmptyAad()
+        {
+            // arrange
+            SFEncryptionMetadata encryptionMetadata = new SFEncryptionMetadata();
+
+            // act
+            using (var encryptedStream = GcmEncryptionProvider.Encrypt(
+                       s_encryptionMaterial,
+                       encryptionMetadata, // this is output parameter
+                       s_fileTransferConfiguration,
+                       new MemoryStream(s_plainTextBytes),
+                       s_emptyAadBytes,
+                       s_emptyAadBytes))
+            {
+                var encryptedContent = ExtractContentBytes(encryptedStream);
+
+                // assert
+                Assert.NotNull(encryptionMetadata.key);
+                Assert.NotNull(encryptionMetadata.iv);
+                Assert.NotNull(encryptionMetadata.matDesc);
+                Assert.AreEqual(s_emptyAadBase64, encryptionMetadata.keyAad);
+                Assert.AreEqual(s_emptyAadBase64, encryptionMetadata.aad);
+
+                // act
+                using (var decryptedStream = GcmEncryptionProvider.Decrypt(new MemoryStream(encryptedContent), s_encryptionMaterial, encryptionMetadata, s_fileTransferConfiguration))
+                {
+                    // assert
+                    var decryptedText = ExtractContent(decryptedStream);
+                    CollectionAssert.AreEqual(s_plainTextBytes, decryptedText);
+                }
+            }
+        }
+
         [Test]
         public void TestEncryptAndDecryptWithAad()
         {
diff --git a/Snowflake.Data/Core/FileTransfer/GcmEncryptionProvider.cs b/Snowflake.Data/Core/FileTransfer/GcmEncryptionProvider.cs
index e60eccba1..50b80dd05 100644
--- a/Snowflake.Data/Core/FileTransfer/GcmEncryptionProvider.cs
+++ b/Snowflake.Data/Core/FileTransfer/GcmEncryptionProvider.cs
@@ -159,9 +159,12 @@ private static Stream DecryptContent(Stream inputStream, byte[] fileKeyBytes, by
                 {
                     cipherStream.Write(buffer, 0, bytesRead);
                 }
-
                 cipherStream.Flush(); // we cannot close or dispose cipherStream because closing cipherStream would close target stream
-                contentCipher.DoFinal(); // in case of decrypting we ignore the result which has to be empty
+                var lastBytes = contentCipher.DoFinal();
+                if (lastBytes != null && lastBytes.Length > 0)
+                {
+                    targetStream.Write(lastBytes, 0, lastBytes.Length);
+                }
                 return targetStream;
             }
             catch (Exception)