test(tvix/nar-bridge): start testing handlers

We currently only had some integration tests (as part of tvix-boot)
testing nar-bridge functionality as a smoketest, but with axum-test we
can test individual handlers and peek at the store afterwards, which is
much more granular.

This adds tests for the nar-specific request handlers.

Change-Id: I7f2345df89ac43b9b372ecc66f696e95e2fcad18
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12916
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>
This commit is contained in:
Florian Klink 2024-11-24 18:37:03 +02:00 committed by clbot
parent 4f9112f1cd
commit 30b631ea72
5 changed files with 870 additions and 49 deletions

View file

@ -18,7 +18,7 @@ use crate::AppState;
#[derive(Debug, Deserialize)]
pub(crate) struct GetNARParams {
#[serde(rename = "narsize")]
nar_size: u64,
nar_size: Option<u64>,
}
#[instrument(skip_all)]
@ -34,6 +34,13 @@ pub async fn get_head(
}): axum::extract::State<AppState>,
) -> Result<impl axum::response::IntoResponse, StatusCode> {
use prost::Message;
// We insist on the nar_size field being set.
// If it's not present, the client is misbehaving somehow.
let nar_size = nar_size.ok_or_else(|| {
warn!("no nar_size parameter set");
StatusCode::BAD_REQUEST
})?;
// b64decode the root node passed *by the user*
let root_node_proto = BASE64URL_NOPAD
.decode(root_node_enc.as_bytes())
@ -57,7 +64,7 @@ pub async fn get_head(
let root_node = root_node.try_into_anonymous_node().map_err(|e| {
warn!(err=%e, "root node validation failed");
StatusCode::BAD_REQUEST
StatusCode::NOT_FOUND
})?;
Ok((
@ -116,7 +123,7 @@ 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
/// Nix probes at {filehash}.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
@ -205,3 +212,264 @@ pub async fn put(
Ok("")
}
#[cfg(test)]
mod tests {
use std::{
num::NonZero,
sync::{Arc, LazyLock},
};
use axum::{http::Method, Router};
use bytes::Bytes;
use data_encoding::BASE64URL_NOPAD;
use nix_compat::nixbase32;
use sha2::Digest;
use tracing_test::traced_test;
use tvix_castore::{
blobservice::{BlobService, MemoryBlobService},
directoryservice::{DirectoryService, MemoryDirectoryService},
fixtures::HELLOWORLD_BLOB_DIGEST,
};
use tvix_store::{
fixtures::{
CASTORE_NODE_COMPLICATED, CASTORE_NODE_SYMLINK, NAR_CONTENTS_COMPLICATED,
NAR_CONTENTS_HELLOWORLD, NAR_CONTENTS_SYMLINK,
},
pathinfoservice::{MemoryPathInfoService, PathInfoService},
};
use crate::AppState;
pub static NAR_STR_SYMLINK: LazyLock<String> = LazyLock::new(|| {
use prost::Message;
BASE64URL_NOPAD.encode(
&tvix_castore::proto::Node::from_name_and_node("".into(), CASTORE_NODE_SYMLINK.clone())
.encode_to_vec(),
)
});
/// Accepts a router without state, and returns a [axum_test::TestServer].
fn gen_server(
router: axum::Router<AppState>,
) -> (
axum_test::TestServer,
impl BlobService,
impl DirectoryService,
impl PathInfoService,
) {
let blob_service = Arc::new(MemoryBlobService::default());
let directory_service = Arc::new(MemoryDirectoryService::default());
let path_info_service = Arc::new(MemoryPathInfoService::default());
let app = router.with_state(AppState::new(
blob_service.clone(),
directory_service.clone(),
path_info_service.clone(),
NonZero::new(100).unwrap(),
));
(
axum_test::TestServer::new(app).unwrap(),
blob_service,
directory_service,
path_info_service,
)
}
#[traced_test]
#[tokio::test]
async fn test_get_head() {
let (server, _blob_service, _directory_service, _path_info_service) =
gen_server(Router::new().route(
"/nar/tvix-castore/:root_node_enc",
axum::routing::get(super::get_head),
));
// Empty nar_str should be NotFound
server
.method(Method::HEAD, "/nar/tvix-castore/")
.expect_failure()
.await
.assert_status_not_found();
let valid_url = &format!("/nar/tvix-castore/{}", &*NAR_STR_SYMLINK);
let qps = &[("narsize", &NAR_CONTENTS_SYMLINK.len().to_string())];
// Missing narsize should be BadRequest
server
.method(Method::HEAD, valid_url)
.expect_failure()
.await
.assert_status_bad_request();
let invalid_url = {
use prost::Message;
let n = tvix_castore::proto::Node {
node: Some(tvix_castore::proto::node::Node::Directory(
tvix_castore::proto::DirectoryNode {
name: "".into(),
digest: "invalid b64".into(),
size: 1,
},
)),
};
&format!(
"/nar/tvix-castore/{}",
BASE64URL_NOPAD.encode(&n.encode_to_vec())
)
};
// Invalid node proto should return NotFound
server
.method(Method::HEAD, invalid_url)
.add_query_params(qps)
.expect_failure()
.await
.assert_status_not_found();
// success, HEAD
server
.method(Method::HEAD, valid_url)
.add_query_params(qps)
.expect_success()
.await;
// success, GET
assert_eq!(
NAR_CONTENTS_SYMLINK.as_slice(),
server
.get(valid_url)
.add_query_params(qps)
.expect_success()
.await
.into_bytes(),
"Expected to get back NAR_CONTENTS_SYMLINK"
)
}
/// Uploading a NAR with a different file hash than what's specified in the URL
/// is considered an error.
#[traced_test]
#[tokio::test]
async fn test_put_wrong_narhash() {
let (server, _blob_service, _directory_service, _path_info_service) =
gen_server(Router::new().route("/nar/:nar_str", axum::routing::put(super::put)));
server
.put("/nar/0000000000000000000000000000000000000000000000000000.nar")
.bytes(Bytes::from_static(&NAR_CONTENTS_SYMLINK))
.expect_failure()
.await;
}
/// Uploading a NAR with compression is not supported.
#[traced_test]
#[tokio::test]
async fn test_put_with_compression_fail() {
let (server, _blob_service, _directory_service, _path_info_service) =
gen_server(Router::new().route("/nar/:nar_str", axum::routing::put(super::put)));
let nar_sha256: [u8; 32] = sha2::Sha256::new_with_prefix(NAR_CONTENTS_SYMLINK.as_slice())
.finalize()
.into();
let nar_url = format!("/nar/{}.nar.zst", nixbase32::encode(&nar_sha256));
server
.put(&nar_url)
.bytes(Bytes::from_static(&NAR_CONTENTS_SYMLINK))
.expect_failure()
.await
.assert_status_unauthorized();
}
/// Upload a NAR with a single file, ensure the blob exists later on.
#[traced_test]
#[tokio::test]
async fn test_put_success() {
let (server, blob_service, _directory_service, _path_info_service) =
gen_server(Router::new().route("/nar/:nar_str", axum::routing::put(super::put)));
let nar_sha256: [u8; 32] =
sha2::Sha256::new_with_prefix(NAR_CONTENTS_HELLOWORLD.as_slice())
.finalize()
.into();
let nar_url = format!("/nar/{}.nar", nixbase32::encode(&nar_sha256));
server
.put(&nar_url)
.bytes(Bytes::from_static(&NAR_CONTENTS_HELLOWORLD))
.expect_success()
.await;
assert!(blob_service
.has(&HELLOWORLD_BLOB_DIGEST)
.await
.expect("blobservice"))
}
// Upload a NAR with blobs and directories, ensure blobs and directories
// were uploaded, by rendering the NAR stream from the root node we know
// describes these contents.
#[traced_test]
#[tokio::test]
async fn test_put_success2() {
let (server, blob_service, directory_service, _path_info_service) =
gen_server(Router::new().route("/nar/:nar_str", axum::routing::put(super::put)));
let nar_sha256: [u8; 32] =
sha2::Sha256::new_with_prefix(NAR_CONTENTS_COMPLICATED.as_slice())
.finalize()
.into();
let nar_url = format!("/nar/{}.nar", nixbase32::encode(&nar_sha256));
server
.put(&nar_url)
.bytes(Bytes::from_static(&NAR_CONTENTS_COMPLICATED))
.expect_success()
.await;
let mut buf = Vec::new();
tvix_store::nar::write_nar(
&mut buf,
&CASTORE_NODE_COMPLICATED,
blob_service,
directory_service,
)
.await
.expect("write nar");
assert_eq!(NAR_CONTENTS_COMPLICATED, buf[..]);
}
/// Upload a NAR, ensure a HEAD by NarHash returns a 2xx code.
#[traced_test]
#[tokio::test]
async fn test_put_root_nodes() {
let (server, _blob_service, _directory_servicee, _path_info_service) = gen_server(
Router::new()
.route("/nar/:nar_str", axum::routing::put(super::put))
.route("/nar/:nar_str", axum::routing::get(super::head_root_nodes)),
);
let nar_sha256: [u8; 32] =
sha2::Sha256::new_with_prefix(NAR_CONTENTS_COMPLICATED.as_slice())
.finalize()
.into();
let nar_url = format!("/nar/{}.nar", nixbase32::encode(&nar_sha256));
// upload NAR
server
.put(&nar_url)
.bytes(Bytes::from_static(&NAR_CONTENTS_COMPLICATED))
.expect_success()
.await;
// check HEAD by NarHash
server.method(Method::HEAD, &nar_url).expect_success().await;
}
}