diff options
Diffstat (limited to 'tvix/store/src/proto')
| -rw-r--r-- | tvix/store/src/proto/grpc_directoryservice_wrapper.rs | 36 | ||||
| -rw-r--r-- | tvix/store/src/proto/grpc_pathinfoservice_wrapper.rs | 5 | ||||
| -rw-r--r-- | tvix/store/src/proto/mod.rs | 11 | ||||
| -rw-r--r-- | tvix/store/src/proto/tests/directory.rs | 10 | ||||
| -rw-r--r-- | tvix/store/src/proto/tests/grpc_directoryservice.rs | 2 |
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( |