From 4095ffe9989dc41ac31e6a2862277a8e7be95e8a Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Sat, 27 Jul 2024 16:33:33 -0700
Subject: [PATCH] Fix potential corruption after restoring a savepoint

If the file had grown enough that the region tracker had to be
reallocated the restoration would free the tracker. This could lead to
corruption, if that page was allocated for other use, when the tracker
then got flushed to the same location

This would only occur in real world usage on very large database files, since the default region
tracker is sized for 4TiB
---
 src/transactions.rs                       | 8 +++++---
 src/tree_store/page_store/page_manager.rs | 4 ++++
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/transactions.rs b/src/transactions.rs
index 94c79f17..fa474e22 100644
--- a/src/transactions.rs
+++ b/src/transactions.rs
@@ -709,12 +709,14 @@ impl WriteTransaction {
             .pages_allocated_since_raw_state(savepoint.get_regional_allocator_states());
 
         // We don't want to rollback the system tree, so keep any pages it references
-        let referenced_by_system_tree = self
+        let mut whitelist = self
             .system_tables
             .lock()
             .unwrap()
             .table_tree
             .all_referenced_pages()?;
+        // The tracker page could have changed too. Don't erase it.
+        whitelist.insert(self.mem.region_tracker_page());
 
         // Find the oldest transaction in the current freed tree, for use below. We do this before
         // freeing pages to ensure that this tree is still valid
@@ -732,7 +734,7 @@ impl WriteTransaction {
 
         let mut freed_pages = vec![];
         for page in allocated_since_savepoint {
-            if referenced_by_system_tree.contains(&page) {
+            if whitelist.contains(&page) {
                 continue;
             }
             freed_pages.push(page);
@@ -773,7 +775,7 @@ impl WriteTransaction {
                 // Re-process the freed pages, but ignore any that would be double-frees
                 if self.mem.is_allocated(page)
                     && !freed_pages_hash.contains(&page)
-                    && !referenced_by_system_tree.contains(&page)
+                    && !whitelist.contains(&page)
                 {
                     freed_pages.push(page);
                     freed_pages_hash.insert(page);
diff --git a/src/tree_store/page_store/page_manager.rs b/src/tree_store/page_store/page_manager.rs
index b670b0d7..05ccc80c 100644
--- a/src/tree_store/page_store/page_manager.rs
+++ b/src/tree_store/page_store/page_manager.rs
@@ -423,6 +423,10 @@ impl TransactionalMemory {
         result
     }
 
+    pub(crate) fn region_tracker_page(&self) -> PageNumber {
+        self.state.lock().unwrap().header.region_tracker()
+    }
+
     // Relocates the region tracker to a lower page, if possible
     // Returns true if the page was moved
     pub(crate) fn relocate_region_tracker(&self) -> Result<bool> {