Fix offset overflow and TOC size accounting for 0.5.3

.github/workflows/rust.yml

@@ -2,9 +2,9 @@ name: Rust
 
 on:
   push:
-    branches: [ "main" ]
+    branches: [ "master" ]
   pull_request:
-    branches: [ "main" ]
+    branches: [ "master" ]
 
 permissions:
   contents: read

Cargo.toml

@@ -3,7 +3,7 @@ members = ["codegen", "tools/create-data-file", "tools/dump-data-file"]
 
 [package]
 name = "data_bucket"
-version = "0.5.2"
+version = "0.5.3"
 edition = "2021"
 authors = ["Handy-caT"]
 license = "MIT"

src/lib.rs

@@ -17,11 +17,14 @@ pub use page::{
     get_index_page_size_from_data_length, map_data_pages_to_general, parse_data_page,
     parse_data_pages_batch, parse_general_header_by_index, parse_page, parse_pages_batch,
     persist_page, persist_pages_batch, seek_by_link, seek_to_page_start, update_at, DataPage,
-    GeneralHeader, GeneralPage, IndexPage, IndexPageUtility, IndexValue, Interval, PageType,
-    SpaceInfoPage, TableOfContentsPage, UnsizedIndexPage, UnsizedIndexPageUtility, DATA_VERSION,
+    GeneralHeader, GeneralPage, IndexPage, IndexPageUtility, IndexValue, Interval,
+    PageOverflowError, PageType, SpaceInfoPage, TableOfContentsOverflowError, TableOfContentsPage,
+    UnsizedIndexPage, UnsizedIndexPageUtility, DATA_VERSION, EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE,
     GENERAL_HEADER_SIZE, INNER_PAGE_SIZE, PAGE_SIZE,
 };
 pub use persistence::{PersistableIndex, PersistableTable};
 pub use space::Id as SpaceId;
 pub use util::access_archived;
-pub use util::{align, align8, align_vec, Persistable, SizeMeasurable, VariableSizeMeasurable};
+pub use util::{
+    align, align8, align_to, align_vec, Persistable, SizeMeasurable, VariableSizeMeasurable,
+};

src/page/data.rs

@@ -18,7 +18,11 @@ impl<const DATA_LENGTH: usize> DataPage<DATA_LENGTH> {
             ));
         }
 
-        if (link.offset + link.length) as usize > DATA_LENGTH {
+        // Sum in usize: `offset + length` in u32 can wrap past 4 GiB and
+        // slip under the bound with a range that is actually out of page.
+        let start = link.offset as usize;
+        let end = link.offset as usize + link.length as usize;
+        if end > DATA_LENGTH {
             return Err(eyre!(
                 "Link range (offset: {}, length: {}) exceeds data bounds ({})",
                 link.offset,
@@ -27,16 +31,17 @@ impl<const DATA_LENGTH: usize> DataPage<DATA_LENGTH> {
             ));
         }
 
-        let start = link.offset as usize;
-        let end = (link.offset + link.length) as usize;
         self.data[start..end].copy_from_slice(new_data);
 
-        self.length = self.length.max(link.offset + link.length);
+        self.length = self.length.max(end as u32);
         Ok(())
     }
 
     pub fn get_at(&self, link: Link) -> Result<&[u8]> {
-        if (link.offset + link.length) as usize > DATA_LENGTH {
+        // Sum in usize, see `update_at`.
+        let start = link.offset as usize;
+        let end = link.offset as usize + link.length as usize;
+        if end > DATA_LENGTH {
             return Err(eyre!(
                 "Link range (offset: {}, length: {}) exceeds data bounds ({})",
                 link.offset,
@@ -45,8 +50,6 @@ impl<const DATA_LENGTH: usize> DataPage<DATA_LENGTH> {
             ));
         }
 
-        let start = link.offset as usize;
-        let end = (link.offset + link.length) as usize;
         Ok(&self.data[start..end])
     }
 }
@@ -126,6 +129,42 @@ mod tests {
             .contains("Link range (offset: 98, length: 3) exceeds data bounds (100)"));
     }
 
+    #[test]
+    fn test_update_at_offset_plus_length_wrapping_u32() {
+        let mut data = DataPage {
+            length: 0,
+            data: [0; 100],
+        };
+
+        // In u32, offset + length wraps to 5 and used to pass the bounds
+        // check, panicking on the slice instead of returning an error.
+        let link = Link {
+            page_id: 1.into(),
+            offset: u32::MAX - 2,
+            length: 8,
+        };
+
+        let err = data.update_at(link, &[1, 2, 3, 4, 5, 6, 7, 8]).unwrap_err();
+        assert!(err.to_string().contains("exceeds data bounds"));
+    }
+
+    #[test]
+    fn test_get_at_offset_plus_length_wrapping_u32() {
+        let data = DataPage {
+            length: 0,
+            data: [0; 100],
+        };
+
+        let link = Link {
+            page_id: 1.into(),
+            offset: u32::MAX - 2,
+            length: 8,
+        };
+
+        let err = data.get_at(link).unwrap_err();
+        assert!(err.to_string().contains("exceeds data bounds"));
+    }
+
     #[test]
     fn test_get_at_out_of_bounds() {
         let data = DataPage {

src/page/index/mod.rs

@@ -7,9 +7,10 @@ use rkyv::{Archive, Deserialize, Serialize};
 use tokio::fs::File;
 use tokio::io::{AsyncSeekExt, AsyncWriteExt};
 
+use crate::page::PageOverflowError;
 use crate::{
-    align, align8, seek_to_page_start, Link, Persistable, SizeMeasurable, VariableSizeMeasurable,
-    GENERAL_HEADER_SIZE,
+    align, align_to, seek_to_page_start, Link, Persistable, SizeMeasurable, VariableSizeMeasurable,
+    GENERAL_HEADER_SIZE, INNER_PAGE_SIZE,
 };
 
 mod page;
@@ -22,7 +23,9 @@ use crate::page::PageId;
 
 pub use page::{get_index_page_size_from_data_length, IndexPage};
 pub use page_for_unsized::{UnsizedIndexPage, UnsizedIndexPageUtility};
-pub use table_of_contents_page::TableOfContentsPage;
+pub use table_of_contents_page::{
+    TableOfContentsOverflowError, TableOfContentsPage, EMPTY_TABLE_OF_CONTENTS_PAGE_SIZE,
+};
 
 pub trait IndexPageUtility<T> {
     type Utility: Persistable + Send + Sync;
@@ -38,10 +41,21 @@ pub trait IndexPageUtility<T> {
         utility: Self::Utility,
     ) -> impl std::future::Future<Output = eyre::Result<()>> + Send {
         async move {
+            let bytes = utility.as_bytes();
+            let utility_length = bytes.as_ref().len();
+            // An oversized utility must fail here, in its own persist,
+            // instead of writing past the page slot into the neighbor page.
+            if utility_length > INNER_PAGE_SIZE {
+                return Err(eyre::Report::new(PageOverflowError {
+                    page_id,
+                    data_length: utility_length,
+                    capacity: INNER_PAGE_SIZE,
+                }));
+            }
             seek_to_page_start(file, page_id.0).await?;
             file.seek(SeekFrom::Current(GENERAL_HEADER_SIZE as i64))
                 .await?;
-            file.write_all(utility.as_bytes().as_ref()).await?;
+            file.write_all(bytes.as_ref()).await?;
             Ok(())
         }
     }
@@ -62,12 +76,15 @@ where
     T: SizeMeasurable,
 {
     fn aligned_size(&self) -> usize {
-        if let Some(align) = T::align() {
-            if align % 8 == 0 {
-                return align8(self.key.aligned_size() + self.link.aligned_size());
+        let len = self.key.aligned_size() + self.link.aligned_size();
+        if let Some(key_align) = T::align() {
+            if key_align % 8 == 0 {
+                // rkyv pads the archived value out to the key's real
+                // alignment (16 for u128-likes), so round to it, not to 8.
+                return align_to(len, key_align);
             }
         }
-        align(self.key.aligned_size() + self.link.aligned_size())
+        align(len)
     }
 }
 

src/page/index/page.rs

@@ -18,9 +18,10 @@ use tokio::fs::File;
 use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
 
 use crate::page::index::IndexPageUtility;
-use crate::page::{IndexValue, PageId};
+use crate::page::{IndexValue, PageId, PageOverflowError};
 use crate::{
     align, align8, seek_to_page_start, Link, Persistable, SizeMeasurable, GENERAL_HEADER_SIZE,
+    INNER_PAGE_SIZE,
 };
 
 pub fn get_index_page_size_from_data_length<T>(length: usize) -> usize
@@ -202,6 +203,27 @@ impl<T: Default + SizeMeasurable> IndexPage<T> {
         Self::read_value(file).await
     }
 
+    /// Rejects a slot write whose byte range would leave the page, so a bad
+    /// `value_index` (or an oversized serialized value) corrupts nothing.
+    ///
+    /// `offset` is relative to the page start and already includes the
+    /// general header.
+    fn check_value_write_bounds(
+        page_id: PageId,
+        offset: usize,
+        value_length: usize,
+    ) -> Result<(), PageOverflowError> {
+        let write_end_in_slot = offset + value_length - GENERAL_HEADER_SIZE;
+        if write_end_in_slot > INNER_PAGE_SIZE {
+            return Err(PageOverflowError {
+                page_id,
+                data_length: write_end_in_slot,
+                capacity: INNER_PAGE_SIZE,
+            });
+        }
+        Ok(())
+    }
+
     fn get_value_offset(size: usize, value_index: usize) -> usize
     where
         T: Default + SizeMeasurable,
@@ -238,11 +260,11 @@ impl<T: Default + SizeMeasurable> IndexPage<T> {
             rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
         >,
     {
-        seek_to_page_start(file, page_id.0).await?;
-
         let offset = Self::get_value_offset(size, value_index as usize);
-        file.seek(SeekFrom::Current(offset as i64)).await?;
         let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&value)?;
+        Self::check_value_write_bounds(page_id, offset, bytes.len())?;
+        seek_to_page_start(file, page_id.0).await?;
+        file.seek(SeekFrom::Current(offset as i64)).await?;
         file.write_all(bytes.as_slice()).await?;
 
         if value_index != size as u16 - 1 {
@@ -278,12 +300,12 @@ impl<T: Default + SizeMeasurable> IndexPage<T> {
             rkyv::api::high::HighValidator<'a, rkyv::rancor::Error>,
         >,
     {
-        seek_to_page_start(file, page_id.0).await?;
-
         let offset = Self::get_value_offset(size, value_index as usize);
-        file.seek(SeekFrom::Current(offset as i64)).await?;
         let value = IndexValue::<T>::default();
         let bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&value)?;
+        Self::check_value_write_bounds(page_id, offset, bytes.len())?;
+        seek_to_page_start(file, page_id.0).await?;
+        file.seek(SeekFrom::Current(offset as i64)).await?;
         file.write_all(bytes.as_slice()).await?;
 
         Ok(())
@@ -391,6 +413,84 @@ mod tests {
         assert_eq!(new_page.index_values, page.index_values);
     }
 
+    #[tokio::test]
+    async fn persist_and_remove_value_reject_writes_past_the_page_slot() {
+        use super::{IndexPageUtility, PageOverflowError, SizedIndexPageUtility};
+
+        let path = std::env::temp_dir().join(format!(
+            "data_bucket_slot_write_bounds_{}.wt",
+            std::process::id()
+        ));
+        let mut file = tokio::fs::OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .truncate(true)
+            .open(&path)
+            .await
+            .unwrap();
+
+        let value = IndexValue::<u64> {
+            key: 7,
+            link: Default::default(),
+        };
+
+        // An in-bounds slot write works.
+        IndexPage::<u64>::persist_value(&mut file, 1.into(), 4, value.clone(), 3)
+            .await
+            .unwrap();
+        // tokio's File buffers writes; flush so metadata() sees them.
+        tokio::io::AsyncWriteExt::flush(&mut file).await.unwrap();
+        let length_after_valid_write = file.metadata().await.unwrap().len();
+
+        // A value index whose slot lies past the page must be rejected
+        // before anything is written.
+        let err = IndexPage::<u64>::persist_value(&mut file, 1.into(), 4, value, 2000)
+            .await
+            .unwrap_err();
+        assert!(
+            err.downcast_ref::<PageOverflowError>().is_some(),
+            "expected PageOverflowError, got: {err}"
+        );
+
+        let err = IndexPage::<u64>::remove_value(&mut file, 1.into(), 4, 2000)
+            .await
+            .unwrap_err();
+        assert!(
+            err.downcast_ref::<PageOverflowError>().is_some(),
+            "expected PageOverflowError, got: {err}"
+        );
+
+        // A utility larger than the page slot must be rejected too.
+        let utility = SizedIndexPageUtility::<u64> {
+            size: 0,
+            node_id: IndexValue::default(),
+            current_index: 0,
+            current_length: 0,
+            slots: vec![0u16; 20_000],
+        };
+        let err = <IndexPage<u64> as IndexPageUtility<u64>>::persist_index_page_utility(
+            &mut file,
+            1.into(),
+            utility,
+        )
+        .await
+        .unwrap_err();
+        assert!(
+            err.downcast_ref::<PageOverflowError>().is_some(),
+            "expected PageOverflowError, got: {err}"
+        );
+
+        // Nothing was written by the rejected operations.
+        assert_eq!(
+            file.metadata().await.unwrap().len(),
+            length_after_valid_write
+        );
+
+        drop(file);
+        std::fs::remove_file(&path).unwrap();
+    }
+
     #[test]
     fn test_split() {
         let mut page = IndexPage::<u64>::new(