use prost::Message;
use std::cmp::Ordering;
mod grpc_blobservice_wrapper;
mod grpc_directoryservice_wrapper;
use crate::{path::PathComponent, B3Digest, DirectoryError};
pub use grpc_blobservice_wrapper::GRPCBlobServiceWrapper;
pub use grpc_directoryservice_wrapper::GRPCDirectoryServiceWrapper;
tonic::include_proto!("tvix.castore.v1");
#[cfg(feature = "tonic-reflection")]
/// Compiled file descriptors for implementing [gRPC
/// reflection](https://github.com/grpc/grpc/blob/master/doc/server-reflection.md) with e.g.
/// [`tonic_reflection`](https://docs.rs/tonic-reflection).
pub const FILE_DESCRIPTOR_SET: &[u8] = tonic::include_file_descriptor_set!("tvix.castore.v1");
#[cfg(test)]
mod tests;
/// Errors that occur during StatBlobResponse validation
#[derive(Debug, PartialEq, Eq, thiserror::Error)]
pub enum ValidateStatBlobResponseError {
/// Invalid digest length encountered
#[error("Invalid digest length {0} for chunk #{1}")]
InvalidDigestLen(usize, usize),
}
fn checked_sum(iter: impl IntoIterator<Item = u64>) -> Option<u64> {
iter.into_iter().try_fold(0u64, |acc, i| acc.checked_add(i))
}
impl Directory {
/// The size of a directory is the number of all regular and symlink elements,
/// the number of directory elements, and their size fields.
pub fn size(&self) -> u64 {
if cfg!(debug_assertions) {
self.size_checked()
.expect("Directory::size exceeds u64::MAX")
} else {
self.size_checked().unwrap_or(u64::MAX)
}
}
fn size_checked(&self) -> Option<u64> {
checked_sum([
self.files.len().try_into().ok()?,
self.symlinks.len().try_into().ok()?,
self.directories.len().try_into().ok()?,
checked_sum(self.directories.iter().map(|e| e.size))?,
])
}
/// 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) -> B3Digest {
let mut hasher = blake3::Hasher::new();
hasher
.update(&self.encode_to_vec())
.finalize()
.as_bytes()
.into()
}
}
impl TryFrom<Directory> for crate::Directory {
type Error = DirectoryError;
fn try_from(value: Directory) -> Result<Self, Self::Error> {
// Check directories, files and symlinks are sorted
// We'll notice duplicates across all three fields when constructing the Directory.
// FUTUREWORK: use is_sorted() once stable, and/or implement the producer for
// [crate::Directory::try_from_iter] iterating over all three and doing all checks inline.
value
.directories
.iter()
.try_fold(&b""[..], |prev_name, e| {
match e.name.as_ref().cmp(prev_name) {
Ordering::Less => Err(DirectoryError::WrongSorting(e.name.to_owned())),
Ordering::Equal => Err(DirectoryError::DuplicateName(
e.name
.to_owned()
.try_into()
.map_err(DirectoryError::InvalidName)?,
)),
Ordering::Greater => Ok(e.name.as_ref()),
}
})?;
value.files.iter().try_fold(&b""[..], |prev_name, e| {
match e.name.as_ref().cmp(prev_name) {
Ordering::Less => Err(DirectoryError::WrongSorting(e.name.to_owned())),
Ordering::Equal => Err(DirectoryError::DuplicateName(
e.name
.to_owned()
.try_into()
.map_err(DirectoryError::InvalidName)?,
)),
Ordering::Greater => Ok(e.name.as_ref()),
}
})?;
value.symlinks.iter().try_fold(&b""[..], |prev_name, e| {
match e.name.as_ref().cmp(prev_name) {
Ordering::Less => Err(DirectoryError::WrongSorting(e.name.to_owned())),
Ordering::Equal => Err(DirectoryError::DuplicateName(
e.name
.to_owned()
.try_into()
.map_err(DirectoryError::InvalidName)?,
)),
Ordering::Greater => Ok(e.name.as_ref()),
}
})?;
// FUTUREWORK: use is_sorted() once stable, and/or implement the producer for
// [crate::Directory::try_from_iter] iterating over all three and doing all checks inline.
let mut elems: Vec<(PathComponent, crate::Node)> =
Vec::with_capacity(value.directories.len() + value.files.len() + value.symlinks.len());
for e in value.directories {
elems.push(
Node {
node: Some(node::Node::Directory(e)),
}
.try_into_name_and_node()?,
);
}
for e in value.files {
elems.push(
Node {
node: Some(node::Node::File(e)),
}
.try_into_name_and_node()?,
)
}
for e in value.symlinks {
elems.push(
Node {
node: Some(node::Node::Symlink(e)),
}
.try_into_name_and_node()?,
)
}
crate::Directory::try_from_iter(elems)
}
}
impl From<crate::Directory> for Directory {
fn from(value: crate::Directory) -> Self {
let mut directories = vec![];
let mut files = vec![];
let mut symlinks = vec![];
for (name, node) in value.into_nodes() {
match node {
crate::Node::File {
digest,
size,
executable,
} => files.push(FileNode {
name: name.into(),
digest: digest.into(),
size,
executable,
}),
crate::Node::Directory { digest, size } => directories.push(DirectoryNode {
name: name.into(),
digest: digest.into(),
size,
}),
crate::Node::Symlink { target } => {
symlinks.push(SymlinkNode {
name: name.into(),
target: target.into(),
});
}
}
}
Directory {
directories,
files,
symlinks,
}
}
}
impl Node {
/// Converts a proto [Node] to a [crate::Node], and splits off the name as a [PathComponent].
pub fn try_into_name_and_node(self) -> Result<(PathComponent, crate::Node), DirectoryError> {
let (name_bytes, node) = self.try_into_unchecked_name_and_checked_node()?;
Ok((
name_bytes.try_into().map_err(DirectoryError::InvalidName)?,
node,
))
}
/// Converts a proto [Node] to a [crate::Node], and splits off the name as a
/// [bytes::Bytes] without doing any checking of it.
fn try_into_unchecked_name_and_checked_node(
self,
) -> Result<(bytes::Bytes, crate::Node), DirectoryError> {
match self.node.ok_or_else(|| DirectoryError::NoNodeSet)? {
node::Node::Directory(n) => {
let digest = B3Digest::try_from(n.digest)
.map_err(|e| DirectoryError::InvalidNode(n.name.clone(), e.into()))?;
let node = crate::Node::Directory {
digest,
size: n.size,
};
Ok((n.name, node))
}
node::Node::File(n) => {
let digest = B3Digest::try_from(n.digest)
.map_err(|e| DirectoryError::InvalidNode(n.name.clone(), e.into()))?;
let node = crate::Node::File {
digest,
size: n.size,
executable: n.executable,
};
Ok((n.name, node))
}
node::Node::Symlink(n) => {
let node = crate::Node::Symlink {
target: n.target.try_into().map_err(|e| {
DirectoryError::InvalidNode(
n.name.clone(),
crate::ValidateNodeError::InvalidSymlinkTarget(e),
)
})?,
};
Ok((n.name, node))
}
}
}
/// Converts a proto [Node] to a [crate::Node], and splits off the name and returns it as a
/// [bytes::Bytes].
///
/// The name must be empty.
pub fn try_into_anonymous_node(self) -> Result<crate::Node, DirectoryError> {
let (name, node) = Self::try_into_unchecked_name_and_checked_node(self)?;
if !name.is_empty() {
return Err(DirectoryError::NameInAnonymousNode);
}
Ok(node)
}
/// Constructs a [Node] from a name and [crate::Node].
/// The name is a [bytes::Bytes], not a [PathComponent], as we have use an
/// empty name in some places.
pub fn from_name_and_node(name: bytes::Bytes, n: crate::Node) -> Self {
match n {
crate::Node::Directory { digest, size } => Self {
node: Some(node::Node::Directory(DirectoryNode {
name,
digest: digest.into(),
size,
})),
},
crate::Node::File {
digest,
size,
executable,
} => Self {
node: Some(node::Node::File(FileNode {
name,
digest: digest.into(),
size,
executable,
})),
},
crate::Node::Symlink { target } => Self {
node: Some(node::Node::Symlink(SymlinkNode {
name,
target: target.into(),
})),
},
}
}
}
impl StatBlobResponse {
/// Validates a StatBlobResponse. All chunks must have valid blake3 digests.
/// It is allowed to send an empty list, if no more granular chunking is
/// available.
pub fn validate(&self) -> Result<(), ValidateStatBlobResponseError> {
for (i, chunk) in self.chunks.iter().enumerate() {
if chunk.digest.len() != blake3::KEY_LEN {
return Err(ValidateStatBlobResponseError::InvalidDigestLen(
chunk.digest.len(),
i,
));
}
}
Ok(())
}
}
