path: root/tvix/store/src/nar/renderer.rs

use std::io::{self, BufReader};

use crate::{
    blobservice::BlobService,
    directoryservice::DirectoryService,
    proto::{self, NamedNode},
};
use data_encoding::BASE64;
use nix_compat::nar;

use super::RenderError;

/// A NAR renderer, using a blob_service, chunk_service and directory_service
/// to render a NAR to a writer.
#[derive(Clone)]
pub struct NARRenderer<BS: BlobService, DS: DirectoryService> {
    blob_service: BS,
    directory_service: DS,
}

impl<BS: BlobService, DS: DirectoryService> NARRenderer<BS, DS> {
    pub fn new(blob_service: BS, directory_service: DS) -> Self {
        Self {
            blob_service,
            directory_service,
        }
    }

    /// 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)?;

        self.walk_node(nar_root_node, proto_root_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: [u8; 32] =
                    proto_file_node.digest.to_owned().try_into().map_err(|_e| {
                        RenderError::StoreError(crate::Error::StorageError(
                            "invalid digest len in file node".to_string(),
                        ))
                    })?;

                // TODO: handle error
                let mut blob_reader = match self.blob_service.open_read(&digest).unwrap() {
                    Some(blob_reader) => Ok(BufReader::new(blob_reader)),
                    None => Err(RenderError::NARWriterError(io::Error::new(
                        io::ErrorKind::NotFound,
                        format!("blob with digest {} not found", BASE64.encode(&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: [u8; 32] =
                    proto_directory_node
                        .digest
                        .to_owned()
                        .try_into()
                        .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)?;

                match resp {
                    // if it's None, that's an error!
                    None => {
                        return Err(RenderError::DirectoryNotFound(
                            digest.to_vec(),
                            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)?;
                        }

                        // close the directory
                        nar_node_directory
                            .close()
                            .map_err(RenderError::NARWriterError)?;
                    }
                }
            }
        }
        Ok(())
    }
}