use crate::control::Hydrant; use axum::Router; use axum::routing::get; use axum::{ Json, extract::{Path, Query, State}, http::{HeaderValue, StatusCode, header}, response::{IntoResponse, Response}, }; use serde::{Deserialize, Serialize}; use tower_http::services::ServeFile; use super::ws::{ Message, WebSocket, WebSocketUpgrade, WsAction, control_frame_only_limits, run_socket, }; pub fn router() -> Router { Router::new() .route("/", get(handle_stream)) .route("/head", get(handle_head)) .route("/manifest", get(handle_manifest)) .route("/chunk/{sha256}", get(handle_chunk)) } /// the newest stored event id. identity and account events take ids too but /// are never replayed, so this is the furthest a cursor replay can reach. #[derive(Serialize)] struct StreamHead<'a> { id: Option, stream: &'a str, } pub async fn handle_head(State(hydrant): State) -> Response { match hydrant.stream_head() { Ok(id) => Json(StreamHead { id, stream: hydrant.stream_id(), }) .into_response(), Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), } } pub async fn handle_manifest(State(hydrant): State) -> Response { match hydrant.stream_chunks() { Some(chunks) => Json(chunks.manifest()).into_response(), None => (StatusCode::NOT_FOUND, "stream chunks are off").into_response(), } } pub async fn handle_chunk( State(hydrant): State, Path(sha256): Path, request: axum::extract::Request, ) -> Response { let Some(path) = hydrant.stream_chunks().and_then(|c| c.servable(&sha256)) else { return StatusCode::NOT_FOUND.into_response(); }; let served = ServeFile::new(path) .with_buf_chunk_size(256 * 1024) .try_call(request) .await; let mut response = match served { Ok(response) => response.map(axum::body::Body::new), Err(e) => return (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), }; if !response.status().is_success() { return response; } let headers = response.headers_mut(); headers.insert( header::CONTENT_TYPE, HeaderValue::from_static("application/zstd"), ); headers.insert( header::CACHE_CONTROL, HeaderValue::from_static("public, max-age=31536000, immutable"), ); if let Ok(etag) = HeaderValue::try_from(format!("\"{sha256}\"")) { headers.insert(header::ETAG, etag); } response } #[derive(Deserialize)] pub struct StreamQuery { pub cursor: Option, /// leave it out and the cursor is not checked against this stream pub stream: Option, } pub async fn handle_stream( State(hydrant): State, Query(query): Query, ws: WebSocketUpgrade, ) -> impl IntoResponse { ws.limits(control_frame_only_limits()) .on_upgrade(move |socket| handle_socket(socket, hydrant, query)) } async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { let send_timeout = hydrant.stream_send_timeout(); let draining = hydrant.state.drain.subscribe(); if let Some(stream) = query.stream.filter(|s| s != hydrant.stream_id()) { let json = serde_json::json!({ "type": "error", "error": "StreamMismatch", "message": format!("cursor is from stream {stream}, this is {}", hydrant.stream_id()), }); let refusal = Message::text(json.to_string()); run_socket( socket, futures::stream::iter([refusal]), |refusal| WsAction::Close(Some(refusal)), send_timeout, |_| None, draining, ) .await; return; } let frames = hydrant.subscribe_frames(query.cursor); run_socket( socket, frames, |item| match item { Ok(frame) => WsAction::Send(Message::text(frame)), Err(err) => { let json = serde_json::json!({ "type": "error", "error": err.code(), "message": err.to_string(), }); WsAction::Close(Some(Message::text(json.to_string()))) } }, send_timeout, |timeout_dur| { let json = serde_json::json!({ "type": "error", "error": "ConsumerTooSlow", "message": format!( "stream socket send blocked for at least {} seconds", timeout_dur.as_secs() ), }); Some(Message::text(json.to_string())) }, draining, ) .await; } #[cfg(test)] mod tests { use super::*; use jacquard_common::types::did::Did; use miette::IntoDiagnostic; async fn head(hydrant: &Hydrant) -> miette::Result { let response = handle_head(State(hydrant.clone())).await; assert_eq!(response.status(), StatusCode::OK); let body = axum::body::to_bytes(response.into_body(), 1024) .await .into_diagnostic()?; serde_json::from_slice(&body).into_diagnostic() } #[tokio::test] async fn head_tracks_durable_events_not_the_id_counter() -> miette::Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), ..Default::default() }) .await?; let stream = hydrant.stream_id().to_owned(); assert_eq!( head(&hydrant).await?, serde_json::json!({ "id": null, "stream": stream }) ); let ids = &hydrant.state.db.stream.ids; hydrant.seed_events_for_bench(3); let last_seeded = ids.head().expect("seeded events took ids"); assert_eq!( head(&hydrant).await?, serde_json::json!({ "id": last_seeded, "stream": stream }) ); let did = Did::new_static("did:plc:aaaaaaaaaaaaaaaaaaaaaaaa").into_diagnostic()?; crate::ops::announce_identity(&hydrant.state.db, &did, None).await; assert_eq!( ids.head(), Some(last_seeded + 1), "the identity event must have taken an id" ); assert_eq!( head(&hydrant).await?, serde_json::json!({ "id": last_seeded, "stream": stream }) ); Ok(()) } #[tokio::test] async fn the_manifest_lists_the_chunks_and_is_404_while_they_are_off() -> miette::Result<()> { let manifest = |hydrant: &Hydrant| handle_manifest(State(hydrant.clone())); let off = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(crate::config::Config { database_path: off.path().to_path_buf(), ..Default::default() }) .await?; assert_eq!(manifest(&hydrant).await.status(), StatusCode::NOT_FOUND); let on = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(crate::config::Config { database_path: on.path().to_path_buf(), stream_chunk_span: 4, ..Default::default() }) .await?; hydrant.seed_events_for_bench(8); crate::control::stream::chunks::build(&hydrant.state)?; let listed = manifest(&hydrant).await; assert_eq!(listed.status(), StatusCode::OK); let body = axum::body::to_bytes(listed.into_body(), 1 << 20) .await .into_diagnostic()?; let chunks = hydrant.stream_chunks().expect("chunks are on").manifest(); assert!(!chunks.chunks.is_empty()); assert_eq!( serde_json::from_slice::(&body).into_diagnostic()?, serde_json::to_value(chunks).into_diagnostic()? ); Ok(()) } #[tokio::test] async fn a_chunk_is_served_whole_or_by_range() -> miette::Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; let hydrant = Hydrant::new(crate::config::Config { database_path: tmp.path().to_path_buf(), stream_chunk_span: 4, ..Default::default() }) .await?; hydrant.seed_events_for_bench(8); crate::control::stream::chunks::build(&hydrant.state)?; let chunks = hydrant.stream_chunks().expect("chunks are on"); let sha256 = chunks.manifest().chunks[0].sha256.clone(); let file = std::fs::read(chunks.path(&sha256)).into_diagnostic()?; let get = |sha256: &str, range: Option<&str>| { let mut request = axum::extract::Request::builder().uri("/"); if let Some(range) = range { request = request.header(header::RANGE, range); } handle_chunk( State(hydrant.clone()), Path(sha256.to_owned()), request.body(axum::body::Body::empty()).unwrap(), ) }; let body = |response: Response| axum::body::to_bytes(response.into_body(), 1 << 20); let whole = get(&sha256, None).await; assert_eq!(whole.status(), StatusCode::OK); let headers = whole.headers(); assert_eq!(headers[header::CONTENT_TYPE], "application/zstd"); assert_eq!( headers[header::CACHE_CONTROL], "public, max-age=31536000, immutable" ); assert_eq!(headers[header::ETAG], format!("\"{sha256}\"").as_str()); assert_eq!(body(whole).await.into_diagnostic()?, file); let part = get(&sha256, Some("bytes=4-9")).await; assert_eq!(part.status(), StatusCode::PARTIAL_CONTENT); assert_eq!(body(part).await.into_diagnostic()?, file[4..10]); assert_eq!(get("00", None).await.status(), StatusCode::NOT_FOUND); Ok(()) } }