Validate page 1 header

ddwalias · ddwalias · commit caf528e6ba89 · 2025-11-08T21:41:21.000+07:00
diff --git a/core/lib.rs b/core/lib.rs
@@ -90,7 +90,7 @@ pub use storage::database::IOContext;
 pub use storage::encryption::{CipherMode, EncryptionContext, EncryptionKey};
 use storage::page_cache::PageCache;
 use storage::pager::{AtomicDbState, DbState};
-use storage::sqlite3_ondisk::PageSize;
+use storage::sqlite3_ondisk::{DatabaseHeader, PageSize};
 pub use storage::{
     buffer_pool::BufferPool,
     database::DatabaseStorage,
@@ -662,41 +662,36 @@ impl Database {
         self.open_flags.get().contains(OpenFlags::ReadOnly)
     }
 
-    /// If we do not have a physical WAL file, but we know the database file is initialized on disk,
-    /// we need to read the page_size from the database header.
-    fn read_page_size_from_db_header(&self) -> Result<PageSize> {
+    /// Read the first page-size block from disk and verify it looks like a database header.
+    fn read_db_header_block(&self) -> Result<Arc<Buffer>> {
         turso_assert!(
             self.db_state.get().is_initialized(),
-            "read_page_size_from_db_header called on uninitialized database"
+            "read_db_header_block called on uninitialized database"
         );
         turso_assert!(
-            PageSize::MIN % 512 == 0,
+            PageSize::MIN.is_multiple_of(512),
             "header read must be a multiple of 512 for O_DIRECT"
         );
         let buf = Arc::new(Buffer::new_temporary(PageSize::MIN as usize));
         let c = Completion::new_read(buf.clone(), move |_res| {});
         let c = self.db_file.read_header(c)?;
         self.io.wait_for_completion(c)?;
+        DatabaseHeader::validate_bytes(buf.as_slice())?;
+        Ok(buf)
+    }
+
+    /// If we do not have a physical WAL file, but we know the database file is initialized on disk,
+    /// we need to read the page_size from the database header.
+    fn read_page_size_from_db_header(&self) -> Result<PageSize> {
+        let buf = self.read_db_header_block()?;
         let page_size = u16::from_be_bytes(buf.as_slice()[16..18].try_into().unwrap());
         let page_size = PageSize::new_from_header_u16(page_size)?;
         Ok(page_size)
     }
 
     fn read_reserved_space_bytes_from_db_header(&self) -> Result<u8> {
-        turso_assert!(
-            self.db_state.get().is_initialized(),
-            "read_reserved_space_bytes_from_db_header called on uninitialized database"
-        );
-        turso_assert!(
-            PageSize::MIN % 512 == 0,
-            "header read must be a multiple of 512 for O_DIRECT"
-        );
-        let buf = Arc::new(Buffer::new_temporary(PageSize::MIN as usize));
-        let c = Completion::new_read(buf.clone(), move |_res| {});
-        let c = self.db_file.read_header(c)?;
-        self.io.wait_for_completion(c)?;
-        let reserved_bytes = u8::from_be_bytes(buf.as_slice()[20..21].try_into().unwrap());
-        Ok(reserved_bytes)
+        let buf = self.read_db_header_block()?;
+        Ok(buf.as_slice()[20])
     }
 
     /// Read the page size in order of preference:
diff --git a/core/storage/encryption.rs b/core/storage/encryption.rs
@@ -65,7 +65,7 @@ use turso_macros::{match_ignore_ascii_case, AtomicEnum};
 /// ```
 ///
 /// constants used for the Turso page header in the encrypted dbs.
-const TURSO_HEADER_PREFIX: &[u8] = b"Turso";
+pub const TURSO_HEADER_PREFIX: &[u8] = b"Turso";
 const TURSO_VERSION: u8 = 0x00;
 const VERSION_OFFSET: usize = 5;
 const CIPHER_OFFSET: usize = 6;
diff --git a/core/storage/pager.rs b/core/storage/pager.rs
@@ -54,6 +54,9 @@ pub struct HeaderRef(PageRef);
 impl HeaderRef {
     pub fn from_pager(pager: &Pager) -> Result<IOResult<Self>> {
         loop {
+            if let Some(page) = pager.cached_header_page() {
+                return Ok(IOResult::Done(Self(page)));
+            }
             let state = pager.header_ref_state.read().clone();
             tracing::trace!("HeaderRef::from_pager - {:?}", state);
             match state {
@@ -74,6 +77,8 @@ impl HeaderRef {
                         page.get().id == DatabaseHeader::PAGE_ID,
                         "incorrect header page id"
                     );
+                    Pager::validate_header_page(&page)?;
+                    pager.set_cached_header_page(&page);
                     *pager.header_ref_state.write() = HeaderRefState::Start;
                     break Ok(IOResult::Done(Self(page)));
                 }
@@ -94,6 +99,10 @@ pub struct HeaderRefMut(PageRef);
 impl HeaderRefMut {
     pub fn from_pager(pager: &Pager) -> Result<IOResult<Self>> {
         loop {
+            if let Some(page) = pager.cached_header_page() {
+                pager.add_dirty(page.as_ref())?;
+                return Ok(IOResult::Done(Self(page)));
+            }
             let state = pager.header_ref_state.read().clone();
             tracing::trace!(?state);
             match state {
@@ -114,6 +123,8 @@ impl HeaderRefMut {
                         page.get().id == DatabaseHeader::PAGE_ID,
                         "incorrect header page id"
                     );
+                    Pager::validate_header_page(&page)?;
+                    pager.set_cached_header_page(&page);
 
                     pager.add_dirty(&page)?;
                     *pager.header_ref_state.write() = HeaderRefState::Start;
@@ -546,6 +557,7 @@ pub struct Pager {
     /// Maximum number of pages allowed in the database. Default is 1073741823 (SQLite default).
     max_page_count: AtomicU32,
     header_ref_state: RwLock<HeaderRefState>,
+    header_page: RwLock<Option<PageRef>>,
     #[cfg(not(feature = "omit_autovacuum"))]
     vacuum_state: RwLock<VacuumState>,
     pub(crate) io_ctx: RwLock<IOContext>,
@@ -660,6 +672,7 @@ impl Pager {
             allocate_page_state: RwLock::new(AllocatePageState::Start),
             max_page_count: AtomicU32::new(DEFAULT_MAX_PAGE_COUNT),
             header_ref_state: RwLock::new(HeaderRefState::Start),
+            header_page: RwLock::new(None),
             #[cfg(not(feature = "omit_autovacuum"))]
             vacuum_state: RwLock::new(VacuumState {
                 ptrmap_get_state: PtrMapGetState::Start,
@@ -677,6 +690,53 @@ impl Pager {
         Ok(())
     }
 
+    fn cached_header_page(&self) -> Option<PageRef> {
+        let maybe_page = self.header_page.read().clone();
+        if let Some(page) = maybe_page {
+            if page.get().contents.is_some() {
+                Some(page)
+            } else {
+                drop(page);
+                self.clear_cached_header_page();
+                None
+            }
+        } else {
+            None
+        }
+    }
+
+    fn set_cached_header_page(&self, page: &PageRef) {
+        let mut guard = self.header_page.write();
+        if guard
+            .as_ref()
+            .map(|cached| Arc::ptr_eq(cached, page))
+            .unwrap_or(false)
+        {
+            return;
+        }
+        if let Some(old) = guard.take() {
+            if old.is_pinned() {
+                old.unpin();
+            }
+        }
+        page.pin();
+        *guard = Some(page.clone());
+    }
+
+    fn clear_cached_header_page(&self) {
+        let mut guard = self.header_page.write();
+        if let Some(page) = guard.take() {
+            if page.is_pinned() {
+                page.unpin();
+            }
+        }
+    }
+
+    fn validate_header_page(page: &PageRef) -> Result<()> {
+        let content: &PageContent = page.get_contents();
+        DatabaseHeader::validate_bytes(&content.buffer.as_slice()[0..DatabaseHeader::SIZE])
+    }
+
     /// Open the subjournal if not yet open.
     /// The subjournal is a file that is used to store the "before images" of pages for the
     /// current savepoint. If the savepoint is rolled back, the pages can be restored from the subjournal.
@@ -2120,6 +2180,7 @@ impl Pager {
     /// of a rollback or in case we want to invalidate page cache after starting a read transaction
     /// right after new writes happened which would invalidate current page cache.
     pub fn clear_page_cache(&self, clear_dirty: bool) {
+        self.clear_cached_header_page();
         let dirty_pages = self.dirty_pages.read();
         let mut cache = self.page_cache.write();
         for page_id in dirty_pages.iter() {
@@ -2713,6 +2774,9 @@ impl Pager {
         })?;
         page.set_loaded();
         page.clear_wal_tag();
+        if id == DatabaseHeader::PAGE_ID {
+            self.set_cached_header_page(&page);
+        }
         Ok(())
     }
 
diff --git a/core/storage/sqlite3_ondisk.rs b/core/storage/sqlite3_ondisk.rs
@@ -58,7 +58,8 @@ use crate::storage::btree::offset::{
 };
 use crate::storage::btree::{payload_overflow_threshold_max, payload_overflow_threshold_min};
 use crate::storage::buffer_pool::BufferPool;
-use crate::storage::database::{DatabaseStorage, EncryptionOrChecksum};
+use crate::storage::database::{DatabaseFile, DatabaseStorage, EncryptionOrChecksum};
+use crate::storage::encryption::TURSO_HEADER_PREFIX;
 use crate::storage::pager::Pager;
 use crate::storage::wal::READMARK_NOT_USED;
 use crate::types::{SerialType, SerialTypeKind, TextSubtype, ValueRef};
@@ -295,6 +296,7 @@ pub struct DatabaseHeader {
 }
 
 impl DatabaseHeader {
+    pub const MAGIC: [u8; 16] = *b"SQLite format 3\0";
     pub const PAGE_ID: usize = 1;
     pub const SIZE: usize = size_of::<Self>();
 
@@ -305,12 +307,51 @@ impl DatabaseHeader {
     pub fn usable_space(self) -> usize {
         (self.page_size.get() as usize) - (self.reserved_space as usize)
     }
+
+    pub fn validate_bytes(header: &[u8]) -> Result<()> {
+        if header.len() < Self::SIZE {
+            return Err(LimboError::NotADB);
+        }
+
+        let has_sqlite_magic = header.starts_with(&Self::MAGIC);
+        let has_turso_magic = header.starts_with(TURSO_HEADER_PREFIX);
+        if !has_sqlite_magic && !has_turso_magic {
+            return Err(LimboError::NotADB);
+        }
+
+        let write_version = header[18];
+        let read_version = header[19];
+        if !matches!(write_version, 1 | 2) || !matches!(read_version, 1 | 2) {
+            return Err(LimboError::NotADB);
+        }
+
+        if header[21] != 64 || header[22] != 32 || header[23] != 32 {
+            return Err(LimboError::NotADB);
+        }
+
+        let page_size_raw = u16::from_be_bytes([header[16], header[17]]);
+        let page_size = match PageSize::new_from_header_u16(page_size_raw) {
+            Ok(size) => size.get() as usize,
+            Err(_) => return Err(LimboError::NotADB),
+        };
+
+        let reserved_bytes = header[20] as usize;
+        if reserved_bytes >= page_size {
+            return Err(LimboError::NotADB);
+        }
+
+        if page_size - reserved_bytes < 480 {
+            return Err(LimboError::NotADB);
+        }
+
+        Ok(())
+    }
 }
 
 impl Default for DatabaseHeader {
     fn default() -> Self {
         Self {
-            magic: *b"SQLite format 3\0",
+            magic: Self::MAGIC,
             page_size: Default::default(),
             write_version: Version::Wal,
             read_version: Version::Wal,
