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<Hydrant> {
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<u64>,
stream: &'a str,
}
pub async fn handle_head(State(hydrant): State<Hydrant>) -> 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<Hydrant>) -> 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<Hydrant>,
Path(sha256): Path<String>,
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<u64>,
/// leave it out and the cursor is not checked against this stream
pub stream: Option<String>,
}
pub async fn handle_stream(
State(hydrant): State<Hydrant>,
Query(query): Query<StreamQuery>,
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<serde_json::Value> {
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::<serde_json::Value>(&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(())
}
}