feat(tvix/store/directory): validate Directory and sizes r/5724

This calls out to Directory::validate() for all received Directory
messages, and also makes sure the sizes we refer a Directory message as
matches the sizes that have been calculated.

Change-Id: I316f9191d5872ee4ba6d78b9a4326f069b22fa63
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7882
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>

1 files changed, 61 insertions, 1 deletions

diff --git a/tvix/store/src/tests/directory_service.rs b/tvix/store/src/tests/directory_service.rs
index 38cc0897c2..8c66ccb538 100644
--- a/tvix/store/src/tests/directory_service.rs
+++ b/tvix/store/src/tests/directory_service.rs
@@ -5,7 +5,7 @@ use tonic::Status;
 use crate::proto::directory_service_server::DirectoryService;
 use crate::proto::get_directory_request::ByWhat;
 use crate::proto::GetDirectoryRequest;
-use crate::proto::{Directory, DirectoryNode};
+use crate::proto::{Directory, DirectoryNode, SymlinkNode};
 use crate::sled_directory_service::SledDirectoryService;
 use lazy_static::lazy_static;
 
@@ -206,3 +206,63 @@ async fn put_get_dedup() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+/// Trying to upload a Directory failing validation should fail.
+#[tokio::test]
+async fn put_reject_failed_validation() -> anyhow::Result<()> {
+    let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?;
+
+    // construct a broken Directory message that fails validation
+    let broken_directory = Directory {
+        symlinks: vec![SymlinkNode {
+            name: "".to_string(),
+            target: "doesntmatter".to_string(),
+        }],
+        ..Default::default()
+    };
+    assert!(broken_directory.validate().is_err());
+
+    // send it over, it must fail
+    let put_resp = service
+        .put(tonic_mock::streaming_request(vec![broken_directory]))
+        .await
+        .expect_err("must fail");
+
+    assert_eq!(put_resp.code(), tonic::Code::InvalidArgument);
+
+    Ok(())
+}
+
+/// Trying to upload a Directory with wrong size should fail.
+#[tokio::test]
+async fn put_reject_wrong_size() -> anyhow::Result<()> {
+    let service = SledDirectoryService::new(TempDir::new()?.path().to_path_buf())?;
+
+    // Construct a directory referring to DIRECTORY_A, but with wrong size.
+    let broken_parent_directory = Directory {
+        directories: vec![DirectoryNode {
+            name: "foo".to_string(),
+            digest: DIRECTORY_A.digest(),
+            size: 42,
+        }],
+        ..Default::default()
+    };
+    // Make sure we got the size wrong.
+    assert_ne!(
+        broken_parent_directory.directories[0].size,
+        DIRECTORY_A.size()
+    );
+
+    // now upload both (first A, then the broken parent). This must fail.
+    let put_resp = service
+        .put(tonic_mock::streaming_request(vec![
+            DIRECTORY_A.clone(),
+            broken_parent_directory,
+        ]))
+        .await
+        .expect_err("must fail");
+
+    assert_eq!(put_resp.code(), tonic::Code::InvalidArgument);
+
+    Ok(())
+}