1 files changed, 66 insertions, 3 deletions
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"),
+    }
+}