From 5afa2543eabecf40cdb514fc6b935d0fc6b4812d Mon Sep 17 00:00:00 2001
From: Anurakt Ghosh <ganurakt@vmware.com>
Date: Wed, 19 Jul 2023 03:12:36 +0530
Subject: [PATCH] fix Invalid Password exception caused due to corruption of
 hidden page in encrypted files.

---
 LiteDB.Tests/Internals/Aes_Tests.cs     | 83 ++++++++++++++++++++++++-
 LiteDB/Engine/Disk/Streams/AesStream.cs | 14 ++++-
 2 files changed, 93 insertions(+), 4 deletions(-)

diff --git a/LiteDB.Tests/Internals/Aes_Tests.cs b/LiteDB.Tests/Internals/Aes_Tests.cs
index fe5489118..91c7031d3 100644
--- a/LiteDB.Tests/Internals/Aes_Tests.cs
+++ b/LiteDB.Tests/Internals/Aes_Tests.cs
@@ -1,5 +1,8 @@
-using System.IO;
+using System;
+using System.IO;
 using System.Linq;
+using System.Reflection;
+using System.Security.Cryptography;
 using FluentAssertions;
 using LiteDB.Engine;
 using Xunit;
@@ -61,5 +64,83 @@ public void Encrypt_Decrypt_Stream()
                 output2.All(x => x == 102).Should().BeTrue();
             }
         }
+
+        /// <summary>
+        /// Test whether AesStream can handle stream that has invalid page size.
+        /// </summary>
+        [Fact]
+        public void AesStream_Invalid_Page_Size()
+        {
+            var fakeContent = new byte[] {
+                1,22,222,184,3,227,126,129,205,182,182,143,201,181,242,107,36,
+                0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
+                20,88,18,70,65,77,202,50,184,177,167,59,80,255,67,66,20,
+                88,18,70,65,77,202,50,184,177,167,59,80,255,67,66
+            };
+
+            // stream of 64bytes (invalid page size)
+            using (var memoryStream = new MemoryStream())
+            {
+                memoryStream.Write(fakeContent, 0, fakeContent.Length);
+                memoryStream.Position = 0;
+
+                using (var crypto = new AesStream("password", memoryStream))
+                {
+                    // 1st page is hidden, so AesStream.Length returns (stream.Length - PAGE_SIZE)
+                    Assert.Equal(0, crypto.Length);
+
+                    // AesStream should add padding to the underlying stream to make its size equivalent to PAGE_SIZE
+                    Assert.Equal(8192, memoryStream.Length);
+                }
+            }
+        }
+
+        /// <summary>
+        /// Test whether AesStream can handle stream where bytes 32-64 are empty.
+        /// </summary>
+        [Fact]
+        public void AesStream_Invalid_Password()
+        {
+            // stream of 8192 bytes where bytes 32 to 64 is empty. 
+            using (var memoryStream = new MemoryStream())
+            {
+                // 1st byte to indicate the stream is encrypted
+                memoryStream.WriteByte(1);
+
+                // next 16 bytes contain salt
+                var salt = new byte[16];
+                using (var rng = RandomNumberGenerator.Create())
+                {
+                    rng.GetBytes(salt);
+                }
+                memoryStream.Write(salt, 0, salt.Length);
+
+                // remaining (8192 - 17) bytes are empty
+                var emptyContent = new byte[8175];
+                memoryStream.Write(emptyContent, 0, emptyContent.Length);
+
+                // reset the stream position to 0
+                memoryStream.Position = 0;
+
+                using (var crypto = new AesStream("password", memoryStream))
+                {
+                    // 1st page is hidden, so AesStream.Length returns (stream.Length - PAGE_SIZE)
+                    Assert.Equal(0, crypto.Length);
+
+                    // AesStream should add padding to the underlying stream to make its size equivalent to PAGE_SIZE
+                    Assert.Equal(8192, memoryStream.Length);
+
+                    // AesStream should fill bytes 32-64 with encrypted 1s
+                    var checkBytes = new byte[32];
+                    var cryptoReader = typeof(AesStream)
+                        .GetField("_reader", BindingFlags.Instance | BindingFlags.NonPublic)
+                        .GetValue(crypto) as CryptoStream;
+                    
+                    memoryStream.Position = 32;
+                    cryptoReader.Read(checkBytes, 0, checkBytes.Length);
+                    Assert.All(checkBytes, b => Assert.Equal(1, b));
+                }
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/LiteDB/Engine/Disk/Streams/AesStream.cs b/LiteDB/Engine/Disk/Streams/AesStream.cs
index 516cb752a..d0eb3f48f 100644
--- a/LiteDB/Engine/Disk/Streams/AesStream.cs
+++ b/LiteDB/Engine/Disk/Streams/AesStream.cs
@@ -62,9 +62,6 @@ public AesStream(string password, Stream stream)
                     // first byte =1 means this datafile is encrypted
                     _stream.WriteByte(1);
                     _stream.Write(this.Salt, 0, ENCRYPTION_SALT_SIZE);
-
-                    // fill with 0 full PAGE_SIZE
-                    _stream.Write(_emptyContent, 0, _emptyContent.Length);
                 }
                 else
                 {
@@ -109,6 +106,17 @@ public AesStream(string password, Stream stream)
 
                 var checkBuffer = new byte[32];
 
+                if (!isNew)
+                {
+                    // check whether bytes 32 to 64 is empty. This indicates LiteDb was unable to write encrypted 1s during last attempt.
+                    _stream.Read(checkBuffer, 0, checkBuffer.Length);
+                    isNew = checkBuffer.All(x => x == 0);
+                    
+                    // reset checkBuffer and stream position
+                    Array.Clear(checkBuffer, 0, checkBuffer.Length);
+                    _stream.Position = 32;
+                }
+
                 // fill checkBuffer with encrypted 1 to check when open
                 if (isNew)
                 {