2 files changed, 32 insertions, 8 deletions
diff --git a/tvix/nar-bridge/src/lib.rs b/tvix/nar-bridge/src/lib.rs
index c4a6c8d5f2dc..9246a334c7f6 100644
--- a/tvix/nar-bridge/src/lib.rs
+++ b/tvix/nar-bridge/src/lib.rs
@@ -53,10 +53,8 @@ impl AppState {
 pub fn gen_router(priority: u64) -> Router<AppState> {
     Router::new()
         .route("/", get(root))
-        // FUTUREWORK: respond for NARs that we still have in root_nodes (at least HEAD)
-        // This avoids some unnecessary NAR uploading from multiple concurrent clients, and is cheap.
         .route("/nar/:nar_str", get(four_o_four))
-        .route("/nar/:nar_str", head(four_o_four))
+        .route("/nar/:nar_str", head(nar::head_root_nodes))
         .route("/nar/:nar_str", put(nar::put))
         .route("/nar/tvix-castore/:root_node_enc", get(nar::get_head))
         .route("/nar/tvix-castore/:root_node_enc", head(nar::get_head))
diff --git a/tvix/nar-bridge/src/nar.rs b/tvix/nar-bridge/src/nar.rs
index abc0d854d7c7..292be2b1c5ec 100644
--- a/tvix/nar-bridge/src/nar.rs
+++ b/tvix/nar-bridge/src/nar.rs
@@ -1,6 +1,5 @@
 use axum::extract::Query;
-use axum::http::StatusCode;
-use axum::response::Response;
+use axum::http::{Response, StatusCode};
 use axum::{body::Body, response::IntoResponse};
 use axum_extra::{headers::Range, TypedHeader};
 use axum_range::{KnownSize, Ranged};
@@ -116,6 +115,36 @@ pub async fn get_head(
     ))
 }
 
+/// Handler to respond to GET/HEAD requests for recently uploaded NAR files.
+/// Nix probes at {narhash}.nar[.compression_suffix] to determine whether a NAR
+/// has already been uploaded, by responding to (some of) these requests we
+/// avoid it unnecessarily uploading.
+/// We don't keep a full K/V from NAR hash to root note around, only the
+/// in-memory cache used to connect to the castore node when processing a PUT
+/// for the NARInfo.
+#[instrument(skip_all, fields(nar_str))]
+pub async fn head_root_nodes(
+    axum::extract::Path(nar_str): axum::extract::Path<String>,
+    axum::extract::State(AppState { root_nodes, .. }): axum::extract::State<AppState>,
+) -> Result<impl axum::response::IntoResponse, StatusCode> {
+    let (nar_hash, compression_suffix) =
+        nix_http::parse_nar_str(&nar_str).ok_or(StatusCode::UNAUTHORIZED)?;
+
+    // No paths with compression suffix are supported.
+    if !compression_suffix.is_empty() {
+        warn!(%compression_suffix, "invalid compression suffix requested");
+        return Err(StatusCode::UNAUTHORIZED);
+    }
+
+    // Check root_nodes, updating the moving it to the most recently used,
+    // as it might be referred in a subsequent NARInfo upload.
+    if root_nodes.write().get(&nar_hash).is_some() {
+        Ok("")
+    } else {
+        Err(StatusCode::NOT_FOUND)
+    }
+}
+
 #[instrument(skip(blob_service, directory_service, request))]
 pub async fn put(
     axum::extract::Path(nar_str): axum::extract::Path<String>,
@@ -172,6 +201,3 @@ pub async fn put(
 
     Ok("")
 }
-
-// FUTUREWORK: maybe head by narhash. Though not too critical, as we do
-// implement HEAD for .narinfo.