1 files changed, 162 insertions, 0 deletions
diff --git a/tvix/nar-bridge/src/narinfo.rs b/tvix/nar-bridge/src/narinfo.rs
new file mode 100644
index 000000000000..fc90f0b86629
--- /dev/null
+++ b/tvix/nar-bridge/src/narinfo.rs
@@ -0,0 +1,162 @@
+use axum::{http::StatusCode, response::IntoResponse};
+use bytes::Bytes;
+use nix_compat::{narinfo::NarInfo, nix_http, nixbase32};
+use prost::Message;
+use tracing::{instrument, warn, Span};
+use tvix_castore::proto::{self as castorepb};
+use tvix_store::proto::PathInfo;
+
+use crate::AppState;
+
+/// The size limit for NARInfo uploads nar-bridge receives
+const NARINFO_LIMIT: usize = 2 * 1024 * 1024;
+
+#[instrument(skip(path_info_service))]
+pub async fn head(
+    axum::extract::Path(narinfo_str): axum::extract::Path<String>,
+    axum::extract::State(AppState {
+        path_info_service, ..
+    }): axum::extract::State<AppState>,
+) -> Result<impl IntoResponse, StatusCode> {
+    let digest = nix_http::parse_narinfo_str(&narinfo_str).ok_or(StatusCode::NOT_FOUND)?;
+    Span::current().record("path_info.digest", &narinfo_str[0..32]);
+
+    if path_info_service
+        .get(digest)
+        .await
+        .map_err(|e| {
+            warn!(err=%e, "failed to get PathInfo");
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?
+        .is_some()
+    {
+        Ok(([("content-type", nix_http::MIME_TYPE_NARINFO)], ""))
+    } else {
+        warn!("PathInfo not found");
+        Err(StatusCode::NOT_FOUND)
+    }
+}
+
+#[instrument(skip(path_info_service))]
+pub async fn get(
+    axum::extract::Path(narinfo_str): axum::extract::Path<String>,
+    axum::extract::State(AppState {
+        path_info_service, ..
+    }): axum::extract::State<AppState>,
+) -> Result<impl IntoResponse, StatusCode> {
+    let digest = nix_http::parse_narinfo_str(&narinfo_str).ok_or(StatusCode::NOT_FOUND)?;
+    Span::current().record("path_info.digest", &narinfo_str[0..32]);
+
+    // fetch the PathInfo
+    let path_info = path_info_service
+        .get(digest)
+        .await
+        .map_err(|e| {
+            warn!(err=%e, "failed to get PathInfo");
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?
+        .ok_or(StatusCode::NOT_FOUND)?;
+
+    let store_path = path_info.validate().map_err(|e| {
+        warn!(err=%e, "invalid PathInfo");
+        StatusCode::INTERNAL_SERVER_ERROR
+    })?;
+
+    let mut narinfo = path_info.to_narinfo(store_path.as_ref()).ok_or_else(|| {
+        warn!(path_info=?path_info, "PathInfo contained no NAR data");
+        StatusCode::INTERNAL_SERVER_ERROR
+    })?;
+
+    // encode the (unnamed) root node in the NAR url itself.
+    // We strip the name from the proto node before sending it out.
+    // It's not needed to render the NAR, it'll make the URL shorter, and it
+    // will make caching these requests easier.
+    let (_, root_node) = path_info
+        .node
+        .as_ref()
+        .expect("invalid pathinfo")
+        .to_owned()
+        .into_name_and_node()
+        .expect("invalid pathinfo");
+
+    let url = format!(
+        "nar/tvix-castore/{}?narsize={}",
+        data_encoding::BASE64URL_NOPAD
+            .encode(&castorepb::Node::from_name_and_node("".into(), root_node).encode_to_vec()),
+        narinfo.nar_size,
+    );
+
+    narinfo.url = &url;
+
+    Ok((
+        [("content-type", nix_http::MIME_TYPE_NARINFO)],
+        narinfo.to_string(),
+    ))
+}
+
+#[instrument(skip(path_info_service, root_nodes, request))]
+pub async fn put(
+    axum::extract::Path(narinfo_str): axum::extract::Path<String>,
+    axum::extract::State(AppState {
+        path_info_service,
+        root_nodes,
+        ..
+    }): axum::extract::State<AppState>,
+    request: axum::extract::Request,
+) -> Result<&'static str, StatusCode> {
+    let _narinfo_digest = nix_http::parse_narinfo_str(&narinfo_str).ok_or(StatusCode::UNAUTHORIZED);
+    Span::current().record("path_info.digest", &narinfo_str[0..32]);
+
+    let narinfo_bytes: Bytes = axum::body::to_bytes(request.into_body(), NARINFO_LIMIT)
+        .await
+        .map_err(|e| {
+            warn!(err=%e, "unable to fetch body");
+            StatusCode::BAD_REQUEST
+        })?;
+
+    // Parse the narinfo from the body.
+    let narinfo_str = std::str::from_utf8(narinfo_bytes.as_ref()).map_err(|e| {
+        warn!(err=%e, "unable decode body as string");
+        StatusCode::BAD_REQUEST
+    })?;
+
+    let narinfo = NarInfo::parse(narinfo_str).map_err(|e| {
+        warn!(err=%e, "unable to parse narinfo");
+        StatusCode::BAD_REQUEST
+    })?;
+
+    // Extract the NARHash from the PathInfo.
+    Span::current().record("path_info.nar_info", nixbase32::encode(&narinfo.nar_hash));
+
+    // populate the pathinfo.
+    let mut pathinfo = PathInfo::from(&narinfo);
+
+    // Lookup root node with peek, as we don't want to update the LRU list.
+    // We need to be careful to not hold the RwLock across the await point.
+    let maybe_root_node: Option<tvix_castore::Node> =
+        root_nodes.read().peek(&narinfo.nar_hash).cloned();
+
+    match maybe_root_node {
+        Some(root_node) => {
+            // Set the root node from the lookup.
+            // We need to rename the node to the narinfo storepath basename, as
+            // that's where it's stored in PathInfo.
+            pathinfo.node = Some(castorepb::Node::from_name_and_node(
+                narinfo.store_path.to_string().into(),
+                root_node,
+            ));
+
+            // Persist the PathInfo.
+            path_info_service.put(pathinfo).await.map_err(|e| {
+                warn!(err=%e, "failed to persist the PathInfo");
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+            Ok("")
+        }
+        None => {
+            warn!("received narinfo with unknown NARHash");
+            Err(StatusCode::BAD_REQUEST)
+        }
+    }
+}