feat(tvix/nar-bridge): init r/8376

This adds an implementation of nar-bridge in Rust.
Currently, only the GET parts are implemented.

Contrary to the Go variant, this doesn't try to keep a mapping from nar
hashes to root node in memory, it simply encodes the root node itself
(stripped by its basename) into the URL.

This pulls in a more recent version of axum than what we use in
tonic, causing two versions of http and hyper, however dealing with
`Body::from_stream` in axum 0.6 is much more annoying, and
https://github.com/hyperium/tonic/pull/1740 suggests this will be fixed
soon.

Change-Id: Ia4c2dbda7cd3fdbe47a75f3e33544d19eac6e44e
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11898
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: Brian Olsen <me@griff.name>
Tested-by: BuildkiteCI

1 files changed, 77 insertions, 0 deletions

diff --git a/tvix/nar-bridge/src/nar.rs b/tvix/nar-bridge/src/nar.rs
new file mode 100644
index 000000000000..63cd86e49969
--- /dev/null
+++ b/tvix/nar-bridge/src/nar.rs
@@ -0,0 +1,77 @@
+use axum::body::Body;
+use axum::extract::Query;
+use axum::http::StatusCode;
+use axum::response::Response;
+use bytes::Bytes;
+use data_encoding::BASE64URL_NOPAD;
+use serde::Deserialize;
+use tokio_util::io::ReaderStream;
+use tracing::{instrument, warn};
+
+use crate::AppState;
+
+#[derive(Debug, Deserialize)]
+pub(crate) struct GetNARParams {
+    #[serde(rename = "narsize")]
+    nar_size: u64,
+}
+
+#[instrument(skip(blob_service, directory_service))]
+pub async fn get(
+    axum::extract::Path(root_node_enc): axum::extract::Path<String>,
+    axum::extract::Query(GetNARParams { nar_size }): Query<GetNARParams>,
+    axum::extract::State(AppState {
+        blob_service,
+        directory_service,
+        ..
+    }): axum::extract::State<AppState>,
+) -> Result<Response, StatusCode> {
+    use prost::Message;
+    // b64decode the root node passed *by the user*
+    let root_node_proto = BASE64URL_NOPAD
+        .decode(root_node_enc.as_bytes())
+        .map_err(|e| {
+            warn!(err=%e, "unable to decode root node b64");
+            StatusCode::NOT_FOUND
+        })?;
+
+    // check the proto size to be somewhat reasonable before parsing it.
+    if root_node_proto.len() > 4096 {
+        warn!("rejected too large root node");
+        return Err(StatusCode::BAD_REQUEST);
+    }
+
+    // parse the proto
+    let root_node: tvix_castore::proto::Node = Message::decode(Bytes::from(root_node_enc))
+        .map_err(|e| {
+            warn!(err=%e, "unable to decode root node proto");
+            StatusCode::NOT_FOUND
+        })?;
+
+    // validate it.
+    let root_node = root_node
+        .validate()
+        .map_err(|e| {
+            warn!(err=%e, "root node validation failed");
+            StatusCode::BAD_REQUEST
+        })?
+        .to_owned();
+
+    let (w, r) = tokio::io::duplex(1024 * 8);
+
+    // spawn a task rendering the NAR to the client
+    tokio::spawn(async move {
+        if let Err(e) =
+            tvix_store::nar::write_nar(w, &root_node, blob_service, directory_service).await
+        {
+            warn!(err=%e, "failed to write out NAR");
+        }
+    });
+
+    Ok(Response::builder()
+        .status(StatusCode::OK)
+        .header("cache-control", "max-age=31536000, immutable")
+        .header("content-length", nar_size)
+        .body(Body::from_stream(ReaderStream::new(r)))
+        .unwrap())
+}