From baae5ce473ed83f35f343656eedb14bb60fbecc7 Mon Sep 17 00:00:00 2001 From: edef Date: Tue, 10 Oct 2023 08:55:00 +0000 Subject: fix(tvix/castore): handle Directory::size overflow explicitly We use checked arithmetic for computing the total size, and verify that size is in-bounds in Directory::validate. If an out-of-bounds size makes it to the "unchecked" size method, we either panic (in debug mode), or silently saturate to u32::MAX. No new panic sites are added, since overflows in debug mode already panic at the language level. Change-Id: I95b8c066a42614fa447f08b4f8fe74e16fbe8bf9 Reviewed-on: https://cl.tvl.fyi/c/depot/+/9616 Reviewed-by: flokli Tested-by: BuildkiteCI --- tvix/castore/src/proto/mod.rs | 30 +++++++++++--- tvix/castore/src/proto/tests/directory.rs | 69 +++++++++++++++++++++++++++++-- 2 files changed, 90 insertions(+), 9 deletions(-) (limited to 'tvix') diff --git a/tvix/castore/src/proto/mod.rs b/tvix/castore/src/proto/mod.rs index 0737f3f08443..ba3fcbceb14d 100644 --- a/tvix/castore/src/proto/mod.rs +++ b/tvix/castore/src/proto/mod.rs @@ -40,6 +40,8 @@ pub enum ValidateDirectoryError { /// Invalid digest length encountered #[error("Invalid Digest length: {0}")] InvalidDigestLen(usize), + #[error("Total size exceeds u32::MAX")] + SizeOverflow, } /// Checks a Node name for validity as an intermediate node, and returns an @@ -130,16 +132,29 @@ fn insert_once<'n>( Ok(()) } +fn checked_sum(iter: impl IntoIterator) -> Option { + iter.into_iter().try_fold(0u32, |acc, i| acc.checked_add(i)) +} + impl Directory { /// The size of a directory is the number of all regular and symlink elements, /// the number of directory elements, and their size fields. pub fn size(&self) -> u32 { - self.files.len() as u32 - + self.symlinks.len() as u32 - + self - .directories - .iter() - .fold(0, |acc: u32, e| (acc + 1 + e.size)) + if cfg!(debug_assertions) { + self.size_checked() + .expect("Directory::size exceeds u32::MAX") + } else { + self.size_checked().unwrap_or(u32::MAX) + } + } + + fn size_checked(&self) -> Option { + checked_sum([ + self.files.len().try_into().ok()?, + self.symlinks.len().try_into().ok()?, + self.directories.len().try_into().ok()?, + checked_sum(self.directories.iter().map(|e| e.size))?, + ]) } /// Calculates the digest of a Directory, which is the blake3 hash of a @@ -201,6 +216,9 @@ impl Directory { insert_once(&mut seen_names, &symlink_node.name)?; } + self.size_checked() + .ok_or(ValidateDirectoryError::SizeOverflow)?; + Ok(()) } diff --git a/tvix/castore/src/proto/tests/directory.rs b/tvix/castore/src/proto/tests/directory.rs index 688a50f5285e..d46e8cb6c3c4 100644 --- a/tvix/castore/src/proto/tests/directory.rs +++ b/tvix/castore/src/proto/tests/directory.rs @@ -62,7 +62,7 @@ fn size() { #[test] #[cfg_attr(not(debug_assertions), ignore)] -#[should_panic] +#[should_panic = "Directory::size exceeds u32::MAX"] fn size_unchecked_panic() { let d = Directory { directories: vec![DirectoryNode { @@ -78,7 +78,7 @@ fn size_unchecked_panic() { #[test] #[cfg_attr(debug_assertions, ignore)] -fn size_unchecked_wrap() { +fn size_unchecked_saturate() { let d = Directory { directories: vec![DirectoryNode { name: "foo".into(), @@ -88,7 +88,53 @@ fn size_unchecked_wrap() { ..Default::default() }; - assert_eq!(d.size(), 0); + assert_eq!(d.size(), u32::MAX); +} + +#[test] +fn size_checked() { + // We don't test the overflow cases that rely purely on immediate + // child count, since that would take an absurd amount of memory. + { + let d = Directory { + directories: vec![DirectoryNode { + name: "foo".into(), + digest: DUMMY_DIGEST.to_vec().into(), + size: u32::MAX - 1, + }], + ..Default::default() + }; + assert_eq!(d.size_checked(), Some(u32::MAX)); + } + { + let d = Directory { + directories: vec![DirectoryNode { + name: "foo".into(), + digest: DUMMY_DIGEST.to_vec().into(), + size: u32::MAX, + }], + ..Default::default() + }; + assert_eq!(d.size_checked(), None); + } + { + let d = Directory { + directories: vec![ + DirectoryNode { + name: "foo".into(), + digest: DUMMY_DIGEST.to_vec().into(), + size: u32::MAX / 2, + }, + DirectoryNode { + name: "foo".into(), + digest: DUMMY_DIGEST.to_vec().into(), + size: u32::MAX / 2, + }, + ], + ..Default::default() + }; + assert_eq!(d.size_checked(), None); + } } #[test] @@ -316,3 +362,20 @@ fn validate_sorting() { d.validate().expect("validate shouldn't error"); } } + +#[test] +fn validate_overflow() { + let d = Directory { + directories: vec![DirectoryNode { + name: "foo".into(), + digest: DUMMY_DIGEST.to_vec().into(), + size: u32::MAX, + }], + ..Default::default() + }; + + match d.validate().expect_err("must fail") { + ValidateDirectoryError::SizeOverflow => {} + _ => panic!("unexpected error"), + } +} -- cgit 1.4.1