diff --git a/tvix/nar-bridge/src/lib.rs b/tvix/nar-bridge/src/lib.rs index 46390e865..2f3dd8243 100644 --- a/tvix/nar-bridge/src/lib.rs +++ b/tvix/nar-bridge/src/lib.rs @@ -59,6 +59,7 @@ pub fn gen_router(priority: u64) -> Router { .route("/nar/tvix-castore/:root_node_enc", get(nar::get)) .route("/:narinfo_str", get(narinfo::get)) .route("/:narinfo_str", head(narinfo::head)) + .route("/:narinfo_str", put(narinfo::put)) .route("/nix-cache-info", get(move || nix_cache_info(priority))) } diff --git a/tvix/nar-bridge/src/narinfo.rs b/tvix/nar-bridge/src/narinfo.rs index 7b6c1bdfd..b985a37c9 100644 --- a/tvix/nar-bridge/src/narinfo.rs +++ b/tvix/nar-bridge/src/narinfo.rs @@ -1,10 +1,15 @@ use axum::http::StatusCode; -use nix_compat::nixbase32; -use tracing::{instrument, warn, Span}; -use tvix_castore::proto::node::Node; +use bytes::Bytes; +use nix_compat::{narinfo::NarInfo, nixbase32}; +use tracing::{info, instrument, warn, Span}; +use tvix_castore::proto::{self as castorepb, node::Node}; +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, @@ -84,6 +89,76 @@ pub async fn get( Ok(narinfo.to_string()) } +#[instrument(skip(path_info_service, root_nodes, request))] +pub async fn put( + axum::extract::Path(narinfo_str): axum::extract::Path, + axum::extract::State(AppState { + path_info_service, + root_nodes, + .. + }): axum::extract::State, + request: axum::extract::Request, +) -> Result<&'static str, StatusCode> { + let _narinfo_digest = parse_narinfo_str(&narinfo_str)?; + 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 = root_nodes + .read() + .peek(&narinfo.nar_hash) + .map(|v| v.to_owned()); + + match maybe_root_node { + Some(root_node) => { + info!(narinfo.store_path=%narinfo.store_path, narinfo.store_path=?narinfo.store_path, "NARINFO STORE PATH"); + + // 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 { + node: Some(root_node.rename(narinfo.store_path.to_string().into())), + }); + + // 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) + } + } +} + /// Parses a `3mzh8lvgbynm9daj7c82k2sfsfhrsfsy.narinfo` string and returns the /// nixbase32-decoded digest. fn parse_narinfo_str(s: &str) -> Result<[u8; 20], StatusCode> {