2 files changed, 90 insertions, 9 deletions

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<Item = u32>) -> Option<u32> {
+    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<u32> {
+        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"),
+    }
+}