feat(tvix/store): validate blob size in NARRenderer r/5917

Make sure the blob size in the current proto node matches what we get
back from the blob backend.

Change-Id: I939fa18f37c7bc86ada8a495c7be622e69ec47f8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8129
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>

3 files changed, 115 insertions, 0 deletions

diff --git a/tvix/store/src/nar/mod.rs b/tvix/store/src/nar/mod.rs
index d7d2cec4d8..a3a8677d92 100644
--- a/tvix/store/src/nar/mod.rs
+++ b/tvix/store/src/nar/mod.rs
@@ -17,6 +17,9 @@ pub enum RenderError {
     #[error("unable to find blob {}, referred from {}", BASE64.encode(.0), .1)]
     BlobNotFound(Vec<u8>, String),
 
+    #[error("unexpected size in metadata for blob {}, referred from {} returned, expected {}, got {}", BASE64.encode(.0), .1, .2, .3)]
+    UnexpectedBlobMeta(Vec<u8>, String, u32, u32),
+
     #[error("failure using the NAR writer: {0}")]
     NARWriterError(std::io::Error),
 }
diff --git a/tvix/store/src/nar/renderer.rs b/tvix/store/src/nar/renderer.rs
index d8d9886b31..94a392d361 100644
--- a/tvix/store/src/nar/renderer.rs
+++ b/tvix/store/src/nar/renderer.rs
@@ -76,6 +76,17 @@ impl<BS: BlobService, CS: ChunkService + Clone, DS: DirectoryService> NARRendere
                         return Err(RenderError::BlobNotFound(digest, proto_file_node.name));
                     }
                     Some(blob_meta) => {
+                        // make sure the blob_meta size matches what we expect from proto_file_node
+                        let blob_meta_size = blob_meta.chunks.iter().fold(0, |acc, e| acc + e.size);
+                        if blob_meta_size != proto_file_node.size {
+                            return Err(RenderError::UnexpectedBlobMeta(
+                                digest,
+                                proto_file_node.name,
+                                proto_file_node.size,
+                                blob_meta_size,
+                            ));
+                        }
+
                         let mut blob_reader = std::io::BufReader::new(BlobReader::open(
                             &self.chunk_service,
                             blob_meta,
diff --git a/tvix/store/src/tests/nar_renderer.rs b/tvix/store/src/tests/nar_renderer.rs
index e19dd88777..7271a45f59 100644
--- a/tvix/store/src/tests/nar_renderer.rs
+++ b/tvix/store/src/tests/nar_renderer.rs
@@ -100,6 +100,107 @@ fn single_symlink() {
     );
 }
 
+/// Make sure the NARRenderer fails if the blob size in the proto node doesn't
+/// match what's in the store.
+#[test]
+fn single_file_missing_blob() {
+    let tmpdir = TempDir::new().unwrap();
+
+    let blob_service = gen_blob_service(tmpdir.path());
+    let chunk_service = gen_chunk_service(tmpdir.path());
+
+    let renderer = NARRenderer::new(
+        blob_service,
+        chunk_service,
+        gen_directory_service(tmpdir.path()),
+    );
+    let mut buf: Vec<u8> = vec![];
+
+    let e = renderer
+        .write_nar(
+            &mut buf,
+            crate::proto::node::Node::File(FileNode {
+                name: "doesntmatter".to_string(),
+                digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
+                size: HELLOWORLD_BLOB_CONTENTS.len() as u32,
+                executable: false,
+            }),
+        )
+        .expect_err("must fail");
+
+    if let crate::nar::RenderError::BlobNotFound(actual_digest, _) = e {
+        assert_eq!(HELLOWORLD_BLOB_DIGEST.to_vec(), actual_digest);
+    } else {
+        panic!("unexpected error")
+    }
+}
+
+/// Make sure the NAR Renderer fails if the returned blob meta has another size
+/// than specified in the proto node.
+#[test]
+fn single_file_wrong_blob_size() {
+    let tmpdir = TempDir::new().unwrap();
+
+    let blob_service = gen_blob_service(tmpdir.path());
+    let chunk_service = gen_chunk_service(tmpdir.path());
+
+    // insert blob and chunk into the stores
+    chunk_service
+        .put(HELLOWORLD_BLOB_CONTENTS.to_vec())
+        .unwrap();
+
+    blob_service
+        .put(
+            &HELLOWORLD_BLOB_DIGEST,
+            proto::BlobMeta {
+                chunks: vec![proto::blob_meta::ChunkMeta {
+                    digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
+                    size: HELLOWORLD_BLOB_CONTENTS.len() as u32,
+                }],
+                ..Default::default()
+            },
+        )
+        .unwrap();
+
+    let renderer = NARRenderer::new(
+        blob_service,
+        chunk_service,
+        gen_directory_service(tmpdir.path()),
+    );
+    let mut buf: Vec<u8> = vec![];
+
+    let e = renderer
+        .write_nar(
+            &mut buf,
+            crate::proto::node::Node::File(FileNode {
+                name: "doesntmatter".to_string(),
+                digest: HELLOWORLD_BLOB_DIGEST.to_vec(),
+                size: 42, // <- note the wrong size here!
+                executable: false,
+            }),
+        )
+        .expect_err("must fail");
+
+    if let crate::nar::RenderError::UnexpectedBlobMeta(digest, _, expected_size, actual_size) = e {
+        assert_eq!(
+            digest,
+            HELLOWORLD_BLOB_DIGEST.to_vec(),
+            "expect digest to match"
+        );
+        assert_eq!(
+            expected_size, 42,
+            "expected expected size to be what's passed in the request"
+        );
+        assert_eq!(
+            actual_size,
+            HELLOWORLD_BLOB_CONTENTS.len() as u32,
+            "expected actual size to be correct"
+        );
+    } else {
+        panic!("unexpected error")
+    }
+}
+
 #[test]
 fn single_file() {
     let tmpdir = TempDir::new().unwrap();