feat(tvix/nar-bridge): support uploading NAR files r/8377

This ingests NAR files into the {Blob,Directory}Service, which are
already part of the AppState.

As we then need to correlate the root node to the uploaded PathInfo, we
need to keep a (short-lived) lookup table from NARHash to root node
around. We insert it into a `LruCache` after the NAR is uploaded, and
use `peek()` to do the lookup, which doesn't update the LRU list.

Change-Id: I48a4c6246bacf76559c5a4ccad2a0bc25c1b7900
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11986
Tested-by: BuildkiteCI
Reviewed-by: Brian Olsen <me@griff.name>

1 files changed, 109 insertions, 1 deletions

diff --git a/tvix/nar-bridge/src/nar.rs b/tvix/nar-bridge/src/nar.rs
index 63cd86e49969..5bce0c973ef3 100644
--- a/tvix/nar-bridge/src/nar.rs
+++ b/tvix/nar-bridge/src/nar.rs
@@ -4,9 +4,13 @@ use axum::http::StatusCode;
 use axum::response::Response;
 use bytes::Bytes;
 use data_encoding::BASE64URL_NOPAD;
+use futures::TryStreamExt;
+use nix_compat::nixbase32;
 use serde::Deserialize;
+use std::io;
 use tokio_util::io::ReaderStream;
-use tracing::{instrument, warn};
+use tracing::{instrument, warn, Span};
+use tvix_store::nar::ingest_nar_and_hash;
 
 use crate::AppState;
 
@@ -75,3 +79,107 @@ pub async fn get(
         .body(Body::from_stream(ReaderStream::new(r)))
         .unwrap())
 }
+
+#[instrument(skip(blob_service, directory_service, request))]
+pub async fn put(
+    axum::extract::Path(nar_str): axum::extract::Path<String>,
+    axum::extract::State(AppState {
+        blob_service,
+        directory_service,
+        root_nodes,
+        ..
+    }): axum::extract::State<AppState>,
+    request: axum::extract::Request,
+) -> Result<&'static str, StatusCode> {
+    let nar_hash_expected = parse_nar_str(&nar_str)?;
+
+    let s = request.into_body().into_data_stream();
+
+    let mut r = tokio_util::io::StreamReader::new(s.map_err(|e| {
+        warn!(err=%e, "failed to read request body");
+        io::Error::new(io::ErrorKind::BrokenPipe, e.to_string())
+    }));
+
+    // ingest the NAR
+    let (root_node, nar_hash_actual, nar_size) =
+        ingest_nar_and_hash(blob_service.clone(), directory_service.clone(), &mut r)
+            .await
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
+            .map_err(|e| {
+                warn!(err=%e, "failed to ingest nar");
+                StatusCode::INTERNAL_SERVER_ERROR
+            })?;
+
+    let s = Span::current();
+    s.record("nar_hash.expected", nixbase32::encode(&nar_hash_expected));
+    s.record("nar_size", nar_size);
+
+    if nar_hash_expected != nar_hash_actual {
+        warn!(
+            nar_hash.expected = nixbase32::encode(&nar_hash_expected),
+            nar_hash.actual = nixbase32::encode(&nar_hash_actual),
+            "nar hash mismatch"
+        );
+        return Err(StatusCode::BAD_REQUEST);
+    }
+
+    // store mapping of narhash to root node into root_nodes.
+    // we need it later to populate the root node when accepting the PathInfo.
+    root_nodes.write().put(nar_hash_actual, root_node);
+
+    Ok("")
+}
+
+// FUTUREWORK: maybe head by narhash. Though not too critical, as we do
+// implement HEAD for .narinfo.
+
+/// Parses a `14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0k.nar`
+/// string and returns the nixbase32-decoded digest.
+/// No compression is supported.
+fn parse_nar_str(s: &str) -> Result<[u8; 32], StatusCode> {
+    if !s.is_char_boundary(52) {
+        warn!("invalid string, no char boundary at 32");
+        return Err(StatusCode::NOT_FOUND);
+    }
+
+    Ok(match s.split_at(52) {
+        (hash_str, ".nar") => {
+            // we know this is 52 bytes
+            let hash_str_fixed: [u8; 52] = hash_str.as_bytes().try_into().unwrap();
+            nixbase32::decode_fixed(hash_str_fixed).map_err(|e| {
+                warn!(err=%e, "invalid digest");
+                StatusCode::NOT_FOUND
+            })?
+        }
+        _ => {
+            warn!("invalid string");
+            return Err(StatusCode::BAD_REQUEST);
+        }
+    })
+}
+
+#[cfg(test)]
+mod test {
+    use super::parse_nar_str;
+    use hex_literal::hex;
+
+    #[test]
+    fn success() {
+        assert_eq!(
+            hex!("13a8cf7ca57f68a9f1752acee36a72a55187d3a954443c112818926f26109d91"),
+            parse_nar_str("14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0k.nar").unwrap()
+        )
+    }
+
+    #[test]
+    fn failure() {
+        assert!(
+            parse_nar_str("14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0k.nar.x").is_err()
+        );
+        assert!(
+            parse_nar_str("14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0k.nar.xz").is_err()
+        );
+        assert!(parse_nar_str("14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0").is_err());
+        assert!(parse_nar_str("14cx20k6z4hq508kqi2lm79qfld5f9mf7kiafpqsjs3zlmycza0🦊.nar").is_err())
+    }
+}