fix bug in decoding, add tests for empty aad

snowflakedb · Oct 22, 2024 · 1aebc4b · 1aebc4b
1 parent 4f4bc6b
commit 1aebc4b
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 4 deletions.
diff --git 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
@@ -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)