diff options
| -rw-r--r-- | tvix/Cargo.lock | 64 | ||||
| -rw-r--r-- | tvix/Cargo.nix | 234 | ||||
| -rw-r--r-- | tvix/Cargo.toml | 2 | ||||
| -rw-r--r-- | tvix/nar-bridge/Cargo.toml | 2 | ||||
| -rw-r--r-- | tvix/nar-bridge/src/nar.rs | 71 |
5 files changed, 352 insertions, 21 deletions
diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock index 22f4d194d43a..32674fbdfb91 100644 --- a/tvix/Cargo.lock +++ b/tvix/Cargo.lock @@ -362,6 +362,44 @@ dependencies = [ ] [[package]] +name = "axum-extra" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be6ea09c9b96cb5076af0de2e383bd2bc0c18f827cf1967bdd353e0b910d733" +dependencies = [ + "axum", + "axum-core", + "bytes", + "futures-util", + "headers", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "serde", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-range" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c30398a7f716ebdd7f3c8a4f7a7a6df48a30e002007fd57b2a7a00fac864bd" +dependencies = [ + "axum", + "axum-extra", + "bytes", + "futures", + "http-body", + "pin-project", + "tokio", +] + +[[package]] name = "backtrace" version = "0.3.69" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1501,6 +1539,30 @@ dependencies = [ ] [[package]] +name = "headers" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "322106e6bd0cba2d5ead589ddb8150a13d7c4217cf80d7c4f682ca994ccc6aa9" +dependencies = [ + "base64 0.21.7", + "bytes", + "headers-core", + "http", + "httpdate", + "mime", + "sha1", +] + +[[package]] +name = "headers-core" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4" +dependencies = [ + "http", +] + +[[package]] name = "heck" version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2133,6 +2195,8 @@ name = "nar-bridge" version = "0.1.0" dependencies = [ "axum", + "axum-extra", + "axum-range", "bytes", "clap", "data-encoding", diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix index 027aaba4845b..a1cb903b70a1 100644 --- a/tvix/Cargo.nix +++ b/tvix/Cargo.nix @@ -1338,6 +1338,173 @@ rec { }; resolvedDefaultFeatures = [ "tracing" ]; }; + "axum-extra" = rec { + crateName = "axum-extra"; + version = "0.9.3"; + edition = "2021"; + sha256 = "0cyp22wy0lykpmkikkr7z0c0rg6j7cw2xpphd83vav5rr44ymrhb"; + libName = "axum_extra"; + dependencies = [ + { + name = "axum"; + packageId = "axum"; + usesDefaultFeatures = false; + } + { + name = "axum-core"; + packageId = "axum-core"; + } + { + name = "bytes"; + packageId = "bytes"; + } + { + name = "futures-util"; + packageId = "futures-util"; + usesDefaultFeatures = false; + features = [ "alloc" ]; + } + { + name = "headers"; + packageId = "headers"; + optional = true; + } + { + name = "http"; + packageId = "http"; + } + { + name = "http-body"; + packageId = "http-body"; + } + { + name = "http-body-util"; + packageId = "http-body-util"; + } + { + name = "mime"; + packageId = "mime"; + } + { + name = "pin-project-lite"; + packageId = "pin-project-lite"; + } + { + name = "serde"; + packageId = "serde"; + } + { + name = "tower"; + packageId = "tower"; + usesDefaultFeatures = false; + features = [ "util" ]; + } + { + name = "tower-layer"; + packageId = "tower-layer"; + } + { + name = "tower-service"; + packageId = "tower-service"; + } + { + name = "tracing"; + packageId = "tracing"; + optional = true; + usesDefaultFeatures = false; + } + ]; + devDependencies = [ + { + name = "axum"; + packageId = "axum"; + } + { + name = "serde"; + packageId = "serde"; + features = [ "derive" ]; + } + { + name = "tower"; + packageId = "tower"; + features = [ "util" ]; + } + ]; + features = { + "async-read-body" = [ "dep:tokio-util" "tokio-util?/io" "dep:tokio" ]; + "cookie" = [ "dep:cookie" ]; + "cookie-key-expansion" = [ "cookie" "cookie?/key-expansion" ]; + "cookie-private" = [ "cookie" "cookie?/private" ]; + "cookie-signed" = [ "cookie" "cookie?/signed" ]; + "default" = [ "tracing" ]; + "erased-json" = [ "dep:serde_json" ]; + "form" = [ "dep:serde_html_form" ]; + "json-deserializer" = [ "dep:serde_json" "dep:serde_path_to_error" ]; + "json-lines" = [ "dep:serde_json" "dep:tokio-util" "dep:tokio-stream" "tokio-util?/io" "tokio-stream?/io-util" "dep:tokio" ]; + "multipart" = [ "dep:multer" ]; + "protobuf" = [ "dep:prost" ]; + "query" = [ "dep:serde_html_form" ]; + "tracing" = [ "dep:tracing" "axum-core/tracing" ]; + "typed-header" = [ "dep:headers" ]; + "typed-routing" = [ "dep:axum-macros" "dep:percent-encoding" "dep:serde_html_form" "dep:form_urlencoded" ]; + }; + resolvedDefaultFeatures = [ "default" "tracing" "typed-header" ]; + }; + "axum-range" = rec { + crateName = "axum-range"; + version = "0.4.0"; + edition = "2021"; + sha256 = "1gb4r3x00yiaggapy002w0q8mx3dg9x4z2iwgzfyn5pplyc07hxi"; + libName = "axum_range"; + dependencies = [ + { + name = "axum"; + packageId = "axum"; + usesDefaultFeatures = false; + } + { + name = "axum-extra"; + packageId = "axum-extra"; + features = [ "typed-header" ]; + } + { + name = "bytes"; + packageId = "bytes"; + } + { + name = "futures"; + packageId = "futures"; + usesDefaultFeatures = false; + features = [ "std" ]; + } + { + name = "http-body"; + packageId = "http-body"; + } + { + name = "pin-project"; + packageId = "pin-project"; + } + { + name = "tokio"; + packageId = "tokio"; + features = [ "fs" "io-util" ]; + } + ]; + devDependencies = [ + { + name = "axum"; + packageId = "axum"; + features = [ "macros" ]; + } + { + name = "tokio"; + packageId = "tokio"; + features = [ "rt" "rt-multi-thread" "macros" ]; + } + ]; + features = { }; + }; "backtrace" = rec { crateName = "backtrace"; version = "0.3.69"; @@ -1412,7 +1579,7 @@ rec { "default" = [ "std" ]; "std" = [ "alloc" ]; }; - resolvedDefaultFeatures = [ "alloc" "std" ]; + resolvedDefaultFeatures = [ "alloc" "default" "std" ]; }; "base64 0.22.1" = rec { crateName = "base64"; @@ -4734,6 +4901,63 @@ rec { }; resolvedDefaultFeatures = [ "ahash" "allocator-api2" "default" "inline-more" "raw" ]; }; + "headers" = rec { + crateName = "headers"; + version = "0.4.0"; + edition = "2015"; + sha256 = "1abari69kjl2yv2dg06g2x17qgd1a20xp7aqmmg2vfhcppk0c89j"; + authors = [ + "Sean McArthur <sean@seanmonstar.com>" + ]; + dependencies = [ + { + name = "base64"; + packageId = "base64 0.21.7"; + } + { + name = "bytes"; + packageId = "bytes"; + } + { + name = "headers-core"; + packageId = "headers-core"; + } + { + name = "http"; + packageId = "http"; + } + { + name = "httpdate"; + packageId = "httpdate"; + } + { + name = "mime"; + packageId = "mime"; + } + { + name = "sha1"; + packageId = "sha1"; + } + ]; + features = { }; + }; + "headers-core" = rec { + crateName = "headers-core"; + version = "0.3.0"; + edition = "2015"; + sha256 = "1r1w80i2bhmyh8s5mjr2dz6baqlrm6cak6yvzm4jq96lacjs5d2l"; + libName = "headers_core"; + authors = [ + "Sean McArthur <sean@seanmonstar.com>" + ]; + dependencies = [ + { + name = "http"; + packageId = "http"; + } + ]; + + }; "heck 0.4.1" = rec { crateName = "heck"; version = "0.4.1"; @@ -6616,6 +6840,14 @@ rec { features = [ "http2" ]; } { + name = "axum-extra"; + packageId = "axum-extra"; + } + { + name = "axum-range"; + packageId = "axum-range"; + } + { name = "bytes"; packageId = "bytes"; } diff --git a/tvix/Cargo.toml b/tvix/Cargo.toml index 16bf9c2dd99c..bcb8e6f7e78e 100644 --- a/tvix/Cargo.toml +++ b/tvix/Cargo.toml @@ -46,6 +46,8 @@ async-process = "2.2.4" async-stream = "0.3.5" async-tempfile = "0.4.0" axum = "0.7.5" +axum-extra = "0.9.3" +axum-range = "0.4.0" # https://github.com/liufuyang/bigtable_rs/pull/86 bigtable_rs = { git = "https://github.com/liufuyang/bigtable_rs", rev = "1818355a5373a5bc2c84287e3a4e3807154ac8ef" } bitflags = "2.6.0" diff --git a/tvix/nar-bridge/Cargo.toml b/tvix/nar-bridge/Cargo.toml index 6ca0479a9a81..ac23b597311f 100644 --- a/tvix/nar-bridge/Cargo.toml +++ b/tvix/nar-bridge/Cargo.toml @@ -5,6 +5,8 @@ edition = "2021" [dependencies] axum = { workspace = true, features = ["http2"] } +axum-extra = { workspace = true } +axum-range = { workspace = true } tower = { workspace = true } tower-http = { workspace = true, features = ["trace"] } bytes = { workspace = true } diff --git a/tvix/nar-bridge/src/nar.rs b/tvix/nar-bridge/src/nar.rs index 2f6c0cc5a0ef..707c01e4bcda 100644 --- a/tvix/nar-bridge/src/nar.rs +++ b/tvix/nar-bridge/src/nar.rs @@ -1,7 +1,9 @@ -use axum::body::Body; use axum::extract::Query; use axum::http::StatusCode; use axum::response::Response; +use axum::{body::Body, response::IntoResponse}; +use axum_extra::{headers::Range, TypedHeader}; +use axum_range::{KnownSize, Ranged}; use bytes::Bytes; use data_encoding::BASE64URL_NOPAD; use futures::TryStreamExt; @@ -22,6 +24,7 @@ pub(crate) struct GetNARParams { #[instrument(skip(blob_service, directory_service))] pub async fn get( + ranges: Option<TypedHeader<Range>>, axum::extract::Path(root_node_enc): axum::extract::Path<String>, axum::extract::Query(GetNARParams { nar_size }): Query<GetNARParams>, axum::extract::State(AppState { @@ -29,7 +32,7 @@ pub async fn get( directory_service, .. }): axum::extract::State<AppState>, -) -> Result<Response, StatusCode> { +) -> Result<impl axum::response::IntoResponse, StatusCode> { use prost::Message; // b64decode the root node passed *by the user* let root_node_proto = BASE64URL_NOPAD @@ -62,24 +65,52 @@ pub async fn get( return Err(StatusCode::BAD_REQUEST); } - 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) - .header("content-type", nix_http::MIME_TYPE_NAR) - .body(Body::from_stream(ReaderStream::new(r))) - .unwrap()) + Ok(( + // headers + [ + ("cache-control", "max-age=31536000, immutable"), + ("content-type", nix_http::MIME_TYPE_NAR), + ], + // If this is a range request, construct a seekable NAR reader + if let Some(TypedHeader(ranges)) = ranges { + let r = + tvix_store::nar::seekable::Reader::new(root_node, blob_service, directory_service) + .await + .map_err(|e| { + warn!(err=%e, "failed to construct seekable nar reader"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + // ensure the user-supplied nar size was correct, no point returning data otherwise. + if r.stream_len() != nar_size { + warn!( + actual_nar_size = r.stream_len(), + supplied_nar_size = nar_size, + "wrong nar size supplied" + ); + return Err(StatusCode::BAD_REQUEST); + } + Ranged::new(Some(ranges), KnownSize::sized(r, nar_size)).into_response() + } else { + // use the non-seekable codepath if there's no range(s) requested, + // as it uses less memory. + 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"); + } + }); + + Response::builder() + .header("content-length", nar_size) + .body(Body::from_stream(ReaderStream::new(r))) + .unwrap() + }, + )) } #[instrument(skip(blob_service, directory_service, request))] |