refactor(tvix/store/directorysvc): move from Vec<u8> to B3Digest r/6177

This introduces a new struct, B3Digest, which internally holds a
Vec<u8>, but only allows construction with 32 bytes.

It also implements display, which will print the base64 representation.
This should reduce some boilerplate when parsing Vec<u8>.

Change-Id: Ia91aa40cb691916773abc8f93e6ed79a5fd34863
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8592
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

5 files changed, 35 insertions, 29 deletions

diff --git a/tvix/store/src/proto/grpc_directoryservice_wrapper.rs b/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
index 35af132215e9..6d2df310137f 100644
--- a/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
+++ b/tvix/store/src/proto/grpc_directoryservice_wrapper.rs
@@ -1,6 +1,5 @@
-use crate::directoryservice::DirectoryService;
 use crate::proto;
-use data_encoding::BASE64;
+use crate::{directoryservice::DirectoryService, B3Digest};
 use std::collections::HashMap;
 use tokio::{sync::mpsc::channel, task};
 use tokio_stream::wrappers::ReceiverStream;
@@ -41,13 +40,9 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
             match &req_inner.by_what {
                 None => return Err(Status::invalid_argument("by_what needs to be specified")),
                 Some(proto::get_directory_request::ByWhat::Digest(digest)) => {
-                    let digest: [u8; 32] = digest
-                        .as_slice()
-                        .try_into()
+                    let digest = B3Digest::from_vec(digest.to_vec())
                         .map_err(|_e| Status::invalid_argument("invalid digest length"))?;
 
-                    let digest_b64: String = BASE64.encode(&digest);
-
                     task::spawn(async move {
                         if !req_inner.recursive {
                             let e: Result<proto::Directory, Status> =
@@ -55,7 +50,7 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
                                     Ok(Some(directory)) => Ok(directory),
                                     Ok(None) => Err(Status::not_found(format!(
                                         "directory {} not found",
-                                        digest_b64
+                                        digest
                                     ))),
                                     Err(e) => Err(e.into()),
                                 };
@@ -97,8 +92,8 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
         // This keeps track of the seen directory keys, and their size.
         // This is used to validate the size field of a reference to a previously sent directory.
         // We don't need to keep the contents around, they're stored in the DB.
-        let mut seen_directories_sizes: HashMap<[u8; 32], u32> = HashMap::new();
-        let mut last_directory_dgst: Option<[u8; 32]> = None;
+        let mut seen_directories_sizes: HashMap<B3Digest, u32> = HashMap::new();
+        let mut last_directory_dgst: Option<B3Digest> = None;
 
         // Consume directories, and insert them into the store.
         // Reject directory messages that refer to Directories not sent in the same stream.
@@ -107,7 +102,7 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
             if let Err(e) = directory.validate() {
                 return Err(Status::invalid_argument(format!(
                     "directory {} failed validation: {}",
-                    BASE64.encode(&directory.digest()),
+                    directory.digest(),
                     e,
                 )));
             }
@@ -116,10 +111,7 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
             // to ensure it has been seen already in this stream, and that the size
             // matches what we recorded.
             for child_directory in &directory.directories {
-                let child_directory_digest: [u8; 32] = child_directory
-                    .digest
-                    .clone()
-                    .try_into()
+                let child_directory_digest = B3Digest::from_vec(child_directory.digest.to_vec())
                     .map_err(|_e| Status::internal("invalid child directory digest len"))?;
 
                 match seen_directories_sizes.get(&child_directory_digest) {
@@ -127,8 +119,8 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
                         return Err(Status::invalid_argument(format!(
                             "child directory '{}' ({}) in directory '{}' not seen yet",
                             child_directory.name,
-                            BASE64.encode(&child_directory_digest),
-                            BASE64.encode(&directory.digest()),
+                            &child_directory_digest,
+                            &directory.digest(),
                         )));
                     }
                     Some(seen_child_directory_size) => {
@@ -136,11 +128,11 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
                             return Err(Status::invalid_argument(format!(
                                     "child directory '{}' ({}) in directory '{}' referred with wrong size, expected {}, actual {}",
                                     child_directory.name,
-                                    BASE64.encode(&child_directory_digest),
-                                    BASE64.encode(&directory.digest()),
+                                    &child_directory_digest,
+                                    &directory.digest(),
                                     seen_child_directory_size,
                                     child_directory.size,
-                                )));
+                            )));
                         }
                     }
                 }
@@ -154,8 +146,8 @@ impl<DS: DirectoryService + Send + Sync + Clone + 'static>
             // reachable from that (root) node.
 
             let dgst = directory.digest();
-            seen_directories_sizes.insert(dgst, directory.size());
-            last_directory_dgst = Some(dgst);
+            seen_directories_sizes.insert(dgst.clone(), directory.size());
+            last_directory_dgst = Some(dgst.clone());
 
             // check if the directory already exists in the database. We can skip
             // inserting if it's already there, as that'd be a no-op.
diff --git a/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs b/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
index 8050ce10cc54..e82557b3a06c 100644
--- a/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
+++ b/tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs
@@ -70,7 +70,10 @@ impl<
         match request.into_inner().node {
             None => Err(Status::invalid_argument("no root node sent")),
             Some(root_node) => match self.nar_calculation_service.calculate_nar(&root_node) {
-                Ok(resp) => Ok(Response::new(resp)),
+                Ok((nar_size, nar_sha256)) => Ok(Response::new(proto::CalculateNarResponse {
+                    nar_size,
+                    nar_sha256: nar_sha256.to_vec(),
+                })),
                 Err(e) => Err(e.into()),
             },
         }
diff --git a/tvix/store/src/proto/mod.rs b/tvix/store/src/proto/mod.rs
index 528d0fb061fd..4db0b9731edc 100644
--- a/tvix/store/src/proto/mod.rs
+++ b/tvix/store/src/proto/mod.rs
@@ -17,6 +17,8 @@ pub use grpc_blobservice_wrapper::GRPCBlobServiceWrapper;
 pub use grpc_directoryservice_wrapper::GRPCDirectoryServiceWrapper;
 pub use grpc_pathinfoservice_wrapper::GRPCPathInfoServiceWrapper;
 
+use crate::B3Digest;
+
 tonic::include_proto!("tvix.store.v1");
 
 #[cfg(feature = "reflection")]
@@ -238,10 +240,15 @@ impl Directory {
 
     /// Calculates the digest of a Directory, which is the blake3 hash of a
     /// Directory protobuf message, serialized in protobuf canonical form.
-    pub fn digest(&self) -> [u8; 32] {
+    pub fn digest(&self) -> B3Digest {
         let mut hasher = blake3::Hasher::new();
 
-        *hasher.update(&self.encode_to_vec()).finalize().as_bytes()
+        let vec = hasher
+            .update(&self.encode_to_vec())
+            .finalize()
+            .as_bytes()
+            .to_vec();
+        B3Digest::from_vec(vec).unwrap()
     }
 
     /// validate checks the directory for invalid data, such as:
diff --git a/tvix/store/src/proto/tests/directory.rs b/tvix/store/src/proto/tests/directory.rs
index 03572629e606..8d6ca7241d7a 100644
--- a/tvix/store/src/proto/tests/directory.rs
+++ b/tvix/store/src/proto/tests/directory.rs
@@ -1,4 +1,7 @@
-use crate::proto::{Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError};
+use crate::{
+    proto::{Directory, DirectoryNode, FileNode, SymlinkNode, ValidateDirectoryError},
+    B3Digest,
+};
 use lazy_static::lazy_static;
 
 lazy_static! {
@@ -66,11 +69,12 @@ fn digest() {
 
     assert_eq!(
         d.digest(),
-        [
+        B3Digest::from_vec(vec![
             0xaf, 0x13, 0x49, 0xb9, 0xf5, 0xf9, 0xa1, 0xa6, 0xa0, 0x40, 0x4d, 0xea, 0x36, 0xdc,
             0xc9, 0x49, 0x9b, 0xcb, 0x25, 0xc9, 0xad, 0xc1, 0x12, 0xb7, 0xcc, 0x9a, 0x93, 0xca,
             0xe4, 0x1f, 0x32, 0x62
-        ]
+        ])
+        .unwrap()
     )
 }
 
diff --git a/tvix/store/src/proto/tests/grpc_directoryservice.rs b/tvix/store/src/proto/tests/grpc_directoryservice.rs
index 37428fbfa339..069e82f6463e 100644
--- a/tvix/store/src/proto/tests/grpc_directoryservice.rs
+++ b/tvix/store/src/proto/tests/grpc_directoryservice.rs
@@ -76,7 +76,7 @@ async fn put_get() {
         .into_inner();
 
     // the sent root_digest should match the calculated digest
-    assert_eq!(put_resp.root_digest, DIRECTORY_A.digest());
+    assert_eq!(put_resp.root_digest, DIRECTORY_A.digest().to_vec());
 
     // get it back
     let items = get_directories(