about summary refs log tree commit diff
path: root/tvix/store/src/nar/renderer.rs
diff options
context:
space:
mode:
Diffstat (limited to 'tvix/store/src/nar/renderer.rs')
-rw-r--r--tvix/store/src/nar/renderer.rs217
1 files changed, 113 insertions, 104 deletions
diff --git a/tvix/store/src/nar/renderer.rs b/tvix/store/src/nar/renderer.rs
index a9a6d989e1b8..6ea76e1429be 100644
--- a/tvix/store/src/nar/renderer.rs
+++ b/tvix/store/src/nar/renderer.rs
@@ -5,131 +5,140 @@ use crate::{
     proto::{self, NamedNode},
     B3Digest,
 };
+use count_write::CountWrite;
 use nix_compat::nar;
+use sha2::{Digest, Sha256};
 use std::io::{self, BufReader};
 use tracing::warn;
 
-/// A NAR renderer, using a blob_service, chunk_service and directory_service
-/// to render a NAR to a writer.
-pub struct NARRenderer<DS: DirectoryService> {
-    blob_service: Box<dyn BlobService>,
+/// Invoke [render_nar], and return the size and sha256 digest of the produced
+/// NAR output.
+pub fn calculate_size_and_sha256<DS: DirectoryService + Clone>(
+    root_node: &proto::node::Node,
+    blob_service: &Box<dyn BlobService>,
     directory_service: DS,
+) -> Result<(u64, [u8; 32]), RenderError> {
+    let h = Sha256::new();
+    let mut cw = CountWrite::from(h);
+
+    writer_nar(&mut cw, root_node, blob_service, directory_service)?;
+
+    Ok((cw.count(), cw.into_inner().finalize().into()))
 }
 
-impl<DS: DirectoryService> NARRenderer<DS> {
-    pub fn new(blob_service: Box<dyn BlobService>, directory_service: DS) -> Self {
-        Self {
-            blob_service,
-            directory_service,
-        }
-    }
+/// Accepts a [proto::node::Node] pointing to the root of a (store) path,
+/// and uses the passed blob_service and directory_service to
+/// perform the necessary lookups as it traverses the structure.
+/// The contents in NAR serialization are writen to the passed [std::io::Write].
+pub fn writer_nar<W: std::io::Write, DS: DirectoryService + Clone>(
+    w: &mut W,
+    proto_root_node: &proto::node::Node,
+    blob_service: &Box<dyn BlobService>,
+    directory_service: DS,
+) -> Result<(), RenderError> {
+    // Initialize NAR writer
+    let nar_root_node = nar::writer::open(w).map_err(RenderError::NARWriterError)?;
 
-    /// Consumes a [proto::node::Node] pointing to the root of a (store) path,
-    /// and writes the contents in NAR serialization to the passed
-    /// [std::io::Write].
-    ///
-    /// It uses the different clients in the struct to perform the necessary
-    /// lookups as it traverses the structure.
-    pub fn write_nar<W: std::io::Write>(
-        &self,
-        w: &mut W,
-        proto_root_node: &proto::node::Node,
-    ) -> Result<(), RenderError> {
-        // Initialize NAR writer
-        let nar_root_node = nar::writer::open(w).map_err(RenderError::NARWriterError)?;
+    walk_node(
+        nar_root_node,
+        proto_root_node,
+        blob_service,
+        directory_service,
+    )
+}
 
-        self.walk_node(nar_root_node, proto_root_node)
-    }
+/// Process an intermediate node in the structure.
+/// This consumes the node.
+fn walk_node<DS: DirectoryService + Clone>(
+    nar_node: nar::writer::Node,
+    proto_node: &proto::node::Node,
+    blob_service: &Box<dyn BlobService>,
+    directory_service: DS,
+) -> Result<(), RenderError> {
+    match proto_node {
+        proto::node::Node::Symlink(proto_symlink_node) => {
+            nar_node
+                .symlink(&proto_symlink_node.target)
+                .map_err(RenderError::NARWriterError)?;
+        }
+        proto::node::Node::File(proto_file_node) => {
+            let digest = B3Digest::from_vec(proto_file_node.digest.clone()).map_err(|_e| {
+                warn!(
+                    file_node = ?proto_file_node,
+                    "invalid digest length in file node",
+                );
 
-    /// Process an intermediate node in the structure.
-    /// This consumes the node.
-    fn walk_node(
-        &self,
-        nar_node: nar::writer::Node,
-        proto_node: &proto::node::Node,
-    ) -> Result<(), RenderError> {
-        match proto_node {
-            proto::node::Node::Symlink(proto_symlink_node) => {
-                nar_node
-                    .symlink(&proto_symlink_node.target)
-                    .map_err(RenderError::NARWriterError)?;
-            }
-            proto::node::Node::File(proto_file_node) => {
-                let digest = B3Digest::from_vec(proto_file_node.digest.clone()).map_err(|_e| {
-                    warn!(
-                        file_node = ?proto_file_node,
-                        "invalid digest length in file node",
-                    );
+                RenderError::StoreError(crate::Error::StorageError(
+                    "invalid digest len in file node".to_string(),
+                ))
+            })?;
 
+            let mut blob_reader = match blob_service
+                .open_read(&digest)
+                .map_err(RenderError::StoreError)?
+            {
+                Some(blob_reader) => Ok(BufReader::new(blob_reader)),
+                None => Err(RenderError::NARWriterError(io::Error::new(
+                    io::ErrorKind::NotFound,
+                    format!("blob with digest {} not found", &digest),
+                ))),
+            }?;
+
+            nar_node
+                .file(
+                    proto_file_node.executable,
+                    proto_file_node.size.into(),
+                    &mut blob_reader,
+                )
+                .map_err(RenderError::NARWriterError)?;
+        }
+        proto::node::Node::Directory(proto_directory_node) => {
+            let digest =
+                B3Digest::from_vec(proto_directory_node.digest.to_vec()).map_err(|_e| {
                     RenderError::StoreError(crate::Error::StorageError(
-                        "invalid digest len in file node".to_string(),
+                        "invalid digest len in directory node".to_string(),
                     ))
                 })?;
 
-                let mut blob_reader = match self
-                    .blob_service
-                    .open_read(&digest)
-                    .map_err(RenderError::StoreError)?
-                {
-                    Some(blob_reader) => Ok(BufReader::new(blob_reader)),
-                    None => Err(RenderError::NARWriterError(io::Error::new(
-                        io::ErrorKind::NotFound,
-                        format!("blob with digest {} not found", &digest),
-                    ))),
-                }?;
-
-                nar_node
-                    .file(
-                        proto_file_node.executable,
-                        proto_file_node.size.into(),
-                        &mut blob_reader,
-                    )
-                    .map_err(RenderError::NARWriterError)?;
-            }
-            proto::node::Node::Directory(proto_directory_node) => {
-                let digest =
-                    B3Digest::from_vec(proto_directory_node.digest.to_vec()).map_err(|_e| {
-                        RenderError::StoreError(crate::Error::StorageError(
-                            "invalid digest len in directory node".to_string(),
-                        ))
-                    })?;
-
-                // look it up with the directory service
-                let resp = self
-                    .directory_service
-                    .get(&digest)
-                    .map_err(RenderError::StoreError)?;
+            // look it up with the directory service
+            let resp = directory_service
+                .get(&digest)
+                .map_err(RenderError::StoreError)?;
 
-                match resp {
-                    // if it's None, that's an error!
-                    None => {
-                        return Err(RenderError::DirectoryNotFound(
-                            digest,
-                            proto_directory_node.name.to_owned(),
-                        ))
-                    }
-                    Some(proto_directory) => {
-                        // start a directory node
-                        let mut nar_node_directory =
-                            nar_node.directory().map_err(RenderError::NARWriterError)?;
-
-                        // for each node in the directory, create a new entry with its name,
-                        // and then invoke walk_node on that entry.
-                        for proto_node in proto_directory.nodes() {
-                            let child_node = nar_node_directory
-                                .entry(proto_node.get_name())
-                                .map_err(RenderError::NARWriterError)?;
-                            self.walk_node(child_node, &proto_node)?;
-                        }
+            match resp {
+                // if it's None, that's an error!
+                None => {
+                    return Err(RenderError::DirectoryNotFound(
+                        digest,
+                        proto_directory_node.name.to_owned(),
+                    ))
+                }
+                Some(proto_directory) => {
+                    // start a directory node
+                    let mut nar_node_directory =
+                        nar_node.directory().map_err(RenderError::NARWriterError)?;
 
-                        // close the directory
-                        nar_node_directory
-                            .close()
+                    // for each node in the directory, create a new entry with its name,
+                    // and then invoke walk_node on that entry.
+                    for proto_node in proto_directory.nodes() {
+                        let child_node = nar_node_directory
+                            .entry(proto_node.get_name())
                             .map_err(RenderError::NARWriterError)?;
+                        walk_node(
+                            child_node,
+                            &proto_node,
+                            blob_service,
+                            directory_service.clone(),
+                        )?;
                     }
+
+                    // close the directory
+                    nar_node_directory
+                        .close()
+                        .map_err(RenderError::NARWriterError)?;
                 }
             }
         }
-        Ok(())
     }
+    Ok(())
 }