fix(storage): harden B-tree and chain integrity

Author: presempathy-awb

src/btree/internal.rs

@@ -23,6 +23,26 @@ pub struct Internal {
     pub entries: Vec<InternalEntry>,
 }
 
+/// Encoded byte cost of one internal separator plus its slot-directory entry.
+pub(crate) fn separator_entry_size(key_len: usize) -> Result<usize> {
+    if key_len > u16::MAX as usize {
+        return Err(PagedbError::PayloadTooLarge);
+    }
+    2usize
+        .checked_add(key_len)
+        .and_then(|size| size.checked_add(8))
+        .and_then(|size| size.checked_add(2))
+        .ok_or(PagedbError::PayloadTooLarge)
+}
+
+#[must_use]
+pub(crate) fn separator_fits(key_len: usize, page_size: usize) -> bool {
+    separator_entry_size(key_len)
+        .ok()
+        .and_then(|entry_size| HEADER_LEN.checked_add(entry_size))
+        .is_some_and(|needed| needed <= body_capacity(page_size))
+}
+
 impl Internal {
     pub fn decode(body: &[u8]) -> Result<Self> {
         let h: NodeHeader = validate_node_body(body)?;

src/btree/leaf.rs

@@ -237,6 +237,24 @@ impl Leaf {
         let slot_dir_bytes = records.len() * 2;
         HEADER_LEN + prefix_len + slot_dir_bytes + record_bytes <= cap
     }
+
+    /// Whether one key/value record fits in an otherwise-empty leaf.
+    #[must_use]
+    pub(crate) fn single_record_fits_encoded(
+        key_len: usize,
+        value_encoded_size: usize,
+        page_size: usize,
+    ) -> bool {
+        if key_len > u16::MAX as usize {
+            return false;
+        }
+        HEADER_LEN
+            .checked_add(key_len)
+            .and_then(|size| size.checked_add(2))
+            .and_then(|size| size.checked_add(2))
+            .and_then(|size| size.checked_add(value_encoded_size))
+            .is_some_and(|needed| needed <= body_capacity(page_size))
+    }
 }
 
 impl Default for Leaf {

src/btree/overflow.rs

@@ -13,6 +13,8 @@
 //! `next` pointer sits at byte 4 in the root and byte 0 in a chain page — any
 //! chain walk must account for that offset.
 
+use std::collections::BTreeSet;
+
 use crate::errors::PagedbError;
 use crate::pager::Pager;
 use crate::pager::format::data_page::ENVELOPE_OVERHEAD;
@@ -272,8 +274,12 @@ pub async fn release<V: Vfs>(
     if info.refcount <= 1 {
         // Collect entire chain.
         let mut freed = vec![root_page_id];
+        let mut seen = BTreeSet::from([root_page_id]);
         let mut cur = info.next;
         while cur != 0 {
+            if !seen.insert(cur) {
+                return Err(PagedbError::overflow_chain_cycle(root_page_id, cur));
+            }
             let guard = pager
                 .read_main_page(cur, realm_id, PageKind::Overflow)
                 .await?;
@@ -303,22 +309,46 @@ pub async fn read_chain<V: Vfs>(
     root_page_id: u64,
     total_len: u64,
 ) -> Result<Vec<u8>> {
-    let mut out: Vec<u8> = Vec::with_capacity(usize::try_from(total_len).unwrap_or(0));
+    let total_len = usize::try_from(total_len)
+        .ok()
+        .filter(|len| isize::try_from(*len).is_ok())
+        .ok_or_else(|| PagedbError::overflow_body_malformed("chain.total_length"))?;
+    // Durable metadata must not choose an arbitrarily large allocation before
+    // any chain page has been authenticated. Grow only as bytes are verified.
+    let mut out: Vec<u8> = Vec::with_capacity(total_len.min(pager.page_size()));
 
     let info = read_root_page(pager, realm_id, root_page_id).await?;
+    if info.root_data.len() > total_len {
+        return Err(PagedbError::overflow_body_malformed(
+            "chain.assembled_length",
+        ));
+    }
     out.extend_from_slice(&info.root_data);
 
+    let mut seen = BTreeSet::from([root_page_id]);
     let mut next = info.next;
     while next != 0 {
+        if !seen.insert(next) {
+            return Err(PagedbError::overflow_chain_cycle(root_page_id, next));
+        }
         let guard = pager
             .read_main_page(next, realm_id, PageKind::Overflow)
             .await?;
         let body = guard.body();
         let (n, data) = decode_overflow(&body)?;
+        if out
+            .len()
+            .checked_add(data.len())
+            .is_none_or(|assembled| assembled > total_len)
+        {
+            return Err(PagedbError::overflow_body_malformed(
+                "chain.assembled_length",
+            ));
+        }
         out.extend_from_slice(data);
         next = n;
     }
-    if out.len() as u64 != total_len {
+    if out.len() != total_len {
         return Err(PagedbError::overflow_body_malformed(
             "chain.assembled_length",
         ));
@@ -334,9 +364,13 @@ pub async fn collect_chain<V: Vfs>(
     root_page_id: u64,
 ) -> Result<Vec<u64>> {
     let mut out = vec![root_page_id];
+    let mut seen = BTreeSet::from([root_page_id]);
     let info = read_root_page(pager, realm_id, root_page_id).await?;
     let mut next = info.next;
     while next != 0 {
+        if !seen.insert(next) {
+            return Err(PagedbError::overflow_chain_cycle(root_page_id, next));
+        }
         let guard = pager
             .read_main_page(next, realm_id, PageKind::Overflow)
             .await?;
@@ -350,8 +384,38 @@ pub async fn collect_chain<V: Vfs>(
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+    use std::time::Duration;
+
+    use crate::crypto::CipherId;
+    use crate::crypto::kdf::derive_mk;
+    use crate::errors::CorruptionDetail;
+    use crate::pager::PagerConfig;
+    use crate::vfs::memory::MemVfs;
+
     use super::*;
 
+    const TEST_PAGE_SIZE: usize = 4096;
+    const TEST_REALM: RealmId = RealmId::new([0xA4; 16]);
+
+    async fn test_pager() -> Arc<Pager<MemVfs>> {
+        let mk = derive_mk(&[0xA5; 32], &[0u8; 16], 0).unwrap();
+        let cfg = PagerConfig {
+            page_size: TEST_PAGE_SIZE,
+            buffer_pool_pages: 16,
+            segment_cache_pages: 16,
+            cipher_id: CipherId::Aes256Gcm,
+            mk_epoch: 0,
+            main_db_file_id: [0xB4; 16],
+            main_db_path: "/main.db".into(),
+            anchor_budget: 1_000_000,
+            dek_lru_capacity: 16,
+            observer_retry_count: 0,
+            metrics_enabled: true,
+        };
+        Arc::new(Pager::open(MemVfs::new(), mk, cfg).await.unwrap())
+    }
+
     #[test]
     fn round_trip_chain_page() {
         let mut body = vec![0u8; 4096 - ENVELOPE_OVERHEAD];
@@ -378,4 +442,111 @@ mod tests {
         // root: 4096 - 40 - 16 = 4040
         assert_eq!(overflow_root_capacity(4096), 4040);
     }
+
+    async fn cyclic_chain(pager: &Pager<MemVfs>, root_page_id: u64, chain_page_id: u64) {
+        let mut root_body = vec![0u8; TEST_PAGE_SIZE - ENVELOPE_OVERHEAD];
+        encode_overflow_root(&mut root_body, 1, chain_page_id, b"").unwrap();
+        pager
+            .write_main_page(root_page_id, TEST_REALM, PageKind::OverflowRoot, &root_body)
+            .await
+            .unwrap();
+
+        let mut chain_body = vec![0u8; TEST_PAGE_SIZE - ENVELOPE_OVERHEAD];
+        encode_overflow(&mut chain_body, chain_page_id, b"").unwrap();
+        pager
+            .write_main_page(chain_page_id, TEST_REALM, PageKind::Overflow, &chain_body)
+            .await
+            .unwrap();
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn read_chain_rejects_absurd_total_len_without_allocation_panic() {
+        let pager = test_pager().await;
+        let root_page_id = 42;
+        let mut body = vec![0u8; TEST_PAGE_SIZE - ENVELOPE_OVERHEAD];
+        encode_overflow_root(&mut body, 1, 0, b"small").unwrap();
+        pager
+            .write_main_page(root_page_id, TEST_REALM, PageKind::OverflowRoot, &body)
+            .await
+            .unwrap();
+
+        let error = read_chain(&pager, TEST_REALM, root_page_id, u64::MAX)
+            .await
+            .unwrap_err();
+        assert!(matches!(
+            error,
+            PagedbError::Corruption(CorruptionDetail::OverflowBodyMalformed { .. })
+        ));
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn read_chain_rejects_more_data_than_declared() {
+        let pager = test_pager().await;
+        let root_page_id = 43;
+        let mut body = vec![0u8; TEST_PAGE_SIZE - ENVELOPE_OVERHEAD];
+        encode_overflow_root(&mut body, 1, 0, b"too-long").unwrap();
+        pager
+            .write_main_page(root_page_id, TEST_REALM, PageKind::OverflowRoot, &body)
+            .await
+            .unwrap();
+
+        let error = read_chain(&pager, TEST_REALM, root_page_id, 1)
+            .await
+            .unwrap_err();
+        assert!(matches!(
+            error,
+            PagedbError::Corruption(CorruptionDetail::OverflowBodyMalformed { .. })
+        ));
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn read_chain_rejects_cycle_without_hanging() {
+        let pager = test_pager().await;
+        cyclic_chain(&pager, 44, 45).await;
+        let error = tokio::time::timeout(
+            Duration::from_secs(1),
+            read_chain(&pager, TEST_REALM, 44, 0),
+        )
+        .await
+        .expect("cycle detection should return before the timeout")
+        .unwrap_err();
+        assert!(matches!(
+            error,
+            PagedbError::Corruption(CorruptionDetail::OverflowChainCycle { .. })
+        ));
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn release_rejects_cycle_without_hanging() {
+        let pager = test_pager().await;
+        cyclic_chain(&pager, 46, 47).await;
+        let result =
+            tokio::time::timeout(Duration::from_secs(1), release(&pager, TEST_REALM, 46, 48))
+                .await
+                .expect("cycle detection should return before the timeout");
+        let Err(error) = result else {
+            panic!("overflow release cycles must not be accepted");
+        };
+        assert!(matches!(
+            error,
+            PagedbError::Corruption(CorruptionDetail::OverflowChainCycle { .. })
+        ));
+    }
+
+    #[tokio::test(flavor = "current_thread")]
+    async fn collect_chain_rejects_cycle_without_hanging() {
+        let pager = test_pager().await;
+        cyclic_chain(&pager, 49, 50).await;
+        let error = tokio::time::timeout(
+            Duration::from_secs(1),
+            collect_chain(&pager, TEST_REALM, 49),
+        )
+        .await
+        .expect("cycle detection should return before the timeout")
+        .expect_err("overflow collect cycles must not be accepted");
+        assert!(matches!(
+            error,
+            PagedbError::Corruption(CorruptionDetail::OverflowChainCycle { .. })
+        ));
+    }
 }

src/btree/tree/bulk.rs

@@ -6,6 +6,7 @@ use crate::Result;
 use crate::errors::PagedbError;
 use crate::vfs::Vfs;
 
+use crate::btree::internal;
 use crate::btree::leaf::{Leaf, LeafValue};
 use crate::btree::node;
 use crate::btree::overflow;
@@ -34,6 +35,18 @@ impl<V: Vfs> BTree<V> {
             return Ok(());
         }
 
+        for pair in pairs.windows(2) {
+            if pair[0].0 >= pair[1].0 {
+                return Err(PagedbError::Io(std::io::Error::new(
+                    std::io::ErrorKind::InvalidInput,
+                    "bulk_load input keys must be strictly increasing",
+                )));
+            }
+        }
+        for (key, value) in &pairs {
+            self.validate_insert_record_fits(key, value)?;
+        }
+
         // Resolve each value to its stored form, spilling values past the inline
         // threshold to overflow chains (same threshold as `put`). Inlining an
         // oversized value would exceed leaf capacity and fail the encode.
@@ -130,13 +143,19 @@ impl<V: Vfs> BTree<V> {
                 let mut used = node::HEADER_LEN;
                 for item in remaining.iter().skip(1) {
                     let sep_key = &item.1;
-                    let entry_bytes = 2 + sep_key.len() + 8 + 2; // record + slot
-                    if used + entry_bytes > body_cap {
+                    let entry_bytes = internal::separator_entry_size(sep_key.len())?;
+                    if used
+                        .checked_add(entry_bytes)
+                        .is_none_or(|needed| needed > body_cap)
+                    {
                         break;
                     }
                     used += entry_bytes;
                     count += 1;
                 }
+                if count == 1 && remaining.len() > 1 {
+                    return Err(PagedbError::PayloadTooLarge);
+                }
                 let chunk = &remaining[..count];
                 remaining = &remaining[count..];