feat(tvix/store): add write_nar function r/5794

This adds a function that 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].

We need this in various places:

 - tvix-store's calculate_nar() RPC method needs to render a NAR stream
   to get the nar hash, which is necessary to give things imported in
   the store a "NAR-based" store path.

 - communication with (remote) Nix (via daemon protocol) needs a NAR
   representation.

 - Things like nar-bridge, exposing a NAR/NARInfo HTTP interface need a
   NAR representation.

Change-Id: I7fb2e0bf01814a1c09094c0e35394d9d6b3e43b6
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7956
Reviewed-by: tazjin <tazjin@tvl.su>
Tested-by: BuildkiteCI

1 files changed, 70 insertions, 0 deletions

diff --git a/tvix/store/src/nar.rs b/tvix/store/src/nar.rs
new file mode 100644
index 000000000000..efcaf652e138
--- /dev/null
+++ b/tvix/store/src/nar.rs
@@ -0,0 +1,70 @@
+//! This provides some common "client-side" libraries to interact with a tvix-
+//! store, in this case to render NAR.
+use crate::{
+    client::StoreClient,
+    proto::{self, NamedNode},
+};
+use anyhow::Result;
+use nix_compat::nar;
+
+/// 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 a [StoreClient] to do the necessary lookups as it traverses the
+/// structure.
+pub fn write_nar<W: std::io::Write, SC: StoreClient>(
+    w: &mut W,
+    proto_root_node: proto::node::Node,
+    store_client: &mut SC,
+) -> Result<()> {
+    // Initialize NAR writer
+    let nar_root_node = nar::writer::open(w)?;
+
+    walk_node(nar_root_node, proto_root_node, store_client)
+}
+
+/// Process an intermediate node in the structure.
+/// This consumes the node.
+fn walk_node<SC: StoreClient>(
+    nar_node: nar::writer::Node,
+    proto_node: proto::node::Node,
+    store_client: &mut SC,
+) -> Result<()> {
+    match proto_node {
+        proto::node::Node::Symlink(proto_symlink_node) => {
+            nar_node.symlink(&proto_symlink_node.target)?;
+        }
+        proto::node::Node::File(proto_file_node) => {
+            nar_node.file(
+                proto_file_node.executable,
+                proto_file_node.size.into(),
+                &mut store_client.open_blob(proto_file_node.digest)?,
+            )?;
+        }
+        proto::node::Node::Directory(proto_directory_node) => {
+            // look up that node from the store client
+            let proto_directory = store_client.get_directory(proto_directory_node.digest)?;
+
+            // if it's None, that's an error!
+            if proto_directory.is_none() {
+                // TODO: proper error handling
+                panic!("not found!")
+            }
+
+            // start a directory node
+            let mut nar_node_directory = nar_node.directory()?;
+
+            // 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.unwrap().nodes() {
+                let child_node = nar_node_directory.entry(proto_node.get_name())?;
+                walk_node(child_node, proto_node, store_client)?;
+            }
+
+            // close the directory
+            nar_node_directory.close()?;
+        }
+    }
+    Ok(())
+}