diff options
Diffstat (limited to 'tvix/store/src/directoryservice/grpc.rs')
| -rw-r--r-- | tvix/store/src/directoryservice/grpc.rs | 140 |
1 files changed, 135 insertions, 5 deletions
diff --git a/tvix/store/src/directoryservice/grpc.rs b/tvix/store/src/directoryservice/grpc.rs index b036f16aca..e44106b329 100644 --- a/tvix/store/src/directoryservice/grpc.rs +++ b/tvix/store/src/directoryservice/grpc.rs @@ -1,7 +1,11 @@ +use std::collections::HashSet; + use super::DirectoryService; use crate::proto::{self, get_directory_request::ByWhat}; -use tonic::transport::Channel; -use tonic::Code; +use data_encoding::BASE64; +use tonic::{transport::Channel, Status}; +use tonic::{Code, Streaming}; +use tracing::{instrument, warn}; /// Connects to a (remote) tvix-store DirectoryService over gRPC. #[derive(Clone)] @@ -29,13 +33,13 @@ impl GRPCDirectoryService { } impl DirectoryService for GRPCDirectoryService { + type DirectoriesIterator = StreamIterator; + fn get(&self, digest: &[u8; 32]) -> Result<Option<crate::proto::Directory>, crate::Error> { // Get a new handle to the gRPC client, and copy the digest. let mut grpc_client = self.grpc_client.clone(); let digest = digest.to_owned(); - // TODO: do requests recursively, populate a backing other - // [DirectoryService] as cache, and ask it first. let task = self.tokio_handle.spawn(async move { let mut s = grpc_client .get(proto::GetDirectoryRequest { @@ -50,7 +54,29 @@ impl DirectoryService for GRPCDirectoryService { }); match self.tokio_handle.block_on(task)? { - Ok(resp) => Ok(resp), + Ok(Some(directory)) => { + // Validate the retrieved Directory indeed has the + // digest we expect it to have, to detect corruptions. + let actual_digest = directory.digest(); + if actual_digest != digest { + Err(crate::Error::StorageError(format!( + "requested directory with digest {}, but got {}", + BASE64.encode(&digest), + BASE64.encode(&actual_digest) + ))) + } else if let Err(e) = directory.validate() { + // Validate the Directory itself is valid. + warn!("directory failed validation: {}", e.to_string()); + Err(crate::Error::StorageError(format!( + "directory {} failed validation: {}", + BASE64.encode(&digest), + e, + ))) + } else { + Ok(Some(directory)) + } + } + Ok(None) => Ok(None), Err(e) if e.code() == Code::NotFound => Ok(None), Err(e) => Err(crate::Error::StorageError(e.to_string())), } @@ -76,6 +102,110 @@ impl DirectoryService for GRPCDirectoryService { Err(e) => Err(crate::Error::StorageError(e.to_string())), } } + + #[instrument(skip_all, fields(directory.digest = BASE64.encode(root_directory_digest)))] + fn get_recursive(&self, root_directory_digest: &[u8; 32]) -> Self::DirectoriesIterator { + let mut grpc_client = self.grpc_client.clone(); + let root_directory_digest = root_directory_digest.to_owned(); + + let task: tokio::task::JoinHandle<Result<Streaming<proto::Directory>, Status>> = + self.tokio_handle.spawn(async move { + let s = grpc_client + .get(proto::GetDirectoryRequest { + recursive: true, + by_what: Some(ByWhat::Digest(root_directory_digest.to_vec())), + }) + .await? + .into_inner(); + + Ok(s) + }); + + let stream = self.tokio_handle.block_on(task).unwrap().unwrap(); + + StreamIterator::new(self.tokio_handle.clone(), &root_directory_digest, stream) + } +} + +pub struct StreamIterator { + /// A handle into the active tokio runtime. Necessary to run futures to completion. + tokio_handle: tokio::runtime::Handle, + // A stream of [proto::Directory] + stream: Streaming<proto::Directory>, + // The Directory digests we received so far + received_directory_digests: HashSet<[u8; 32]>, + // The Directory digests we're still expecting to get sent. + expected_directory_digests: HashSet<[u8; 32]>, +} + +impl StreamIterator { + pub fn new( + tokio_handle: tokio::runtime::Handle, + root_digest: &[u8; 32], + stream: Streaming<proto::Directory>, + ) -> Self { + Self { + tokio_handle, + stream, + received_directory_digests: HashSet::new(), + expected_directory_digests: HashSet::from([*root_digest]), + } + } +} + +impl Iterator for StreamIterator { + type Item = Result<proto::Directory, crate::Error>; + + fn next(&mut self) -> Option<Self::Item> { + match self.tokio_handle.block_on(self.stream.message()) { + Ok(ok) => match ok { + Some(directory) => { + // validate the directory itself. + if let Err(e) = directory.validate() { + return Some(Err(crate::Error::StorageError(format!( + "directory {} failed validation: {}", + BASE64.encode(&directory.digest()), + e, + )))); + } + // validate we actually expected that directory, and move it from expected to received. + let directory_digest = directory.digest(); + let was_expected = self.expected_directory_digests.remove(&directory_digest); + if !was_expected { + // FUTUREWORK: dumb clients might send the same stuff twice. + // as a fallback, we might want to tolerate receiving + // it if it's in received_directory_digests (as that + // means it once was in expected_directory_digests) + return Some(Err(crate::Error::StorageError(format!( + "received unexpected directory {}", + BASE64.encode(&directory_digest) + )))); + } + self.received_directory_digests.insert(directory_digest); + + // register all children in expected_directory_digests. + for child_directories in &directory.directories { + self.expected_directory_digests + .insert(child_directories.digest.clone().try_into().unwrap()); + } + + Some(Ok(directory)) + } + None => { + // If we were still expecting something, that's an error. + if !self.expected_directory_digests.is_empty() { + Some(Err(crate::Error::StorageError(format!( + "still expected {} directories, but got premature end of stream", + self.expected_directory_digests.len(), + )))) + } else { + None + } + } + }, + Err(e) => Some(Err(crate::Error::StorageError(e.to_string()))), + } + } } #[cfg(test)] |