use std::path::{Path, PathBuf};
use axum::body::{Body, Bytes};
use axum::extract::{Path as AxumPath, Request, State};
use axum::http::{StatusCode, header};
use axum::response::{IntoResponse, Response};
use futures_util::StreamExt as _;
use serde_json::{Value, json};
use tokio::sync::mpsc;
use tracing::{debug, warn};
use trusty_common::uds::UdsRpcError;
use trusty_common::uds::stream_client::FramedStream;
use super::map::{Call, map_request};
use super::{
CALL_TIMEOUT, MAX_FRAME_BYTES, STREAM_OPEN_TIMEOUT, SearchRpcError, call, json_response,
open_stream,
};
use crate::server::AppState;
const SSE_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(20);
const SSE_BUFFER: usize = 64;
pub async fn search_api_handler(
State(state): State<AppState>,
AxumPath(path): AxumPath<String>,
req: Request,
) -> Response {
let (parts, body) = req.into_parts();
let query = parts.uri.query().map(str::to_owned);
let body_limit = usize::try_from(MAX_FRAME_BYTES).unwrap_or(usize::MAX);
let body_bytes: Bytes = match axum::body::to_bytes(body, body_limit).await {
Ok(b) => b,
Err(e) => {
warn!("search_uds: could not read the request body: {e}");
return (
StatusCode::PAYLOAD_TOO_LARGE,
format!("request body exceeds the {MAX_FRAME_BYTES}-byte frame budget"),
)
.into_response();
}
};
let mapped = match map_request(&parts.method, &path, query.as_deref(), &body_bytes) {
Ok(call) => call,
Err(reason) => {
warn!(
"search_uds: {} /{path} is not mapped: {reason}",
parts.method
);
return (
StatusCode::NOT_IMPLEMENTED,
axum::Json(json!({ "error": reason, "service": super::SEARCH_SERVICE })),
)
.into_response();
}
};
let socket = match state.search_socket_path() {
Ok(p) => p,
Err(reason) => return SearchRpcError::Unresolved(reason).into_response(),
};
match mapped {
Call::Unary { method, params } => {
debug!("search_uds: {} /{path} → {method}", parts.method);
match call(&socket, method, params, CALL_TIMEOUT).await {
Ok(result) => json_response(&result),
Err(e) => e.into_response(),
}
}
Call::Stream { method, params } => {
debug!("search_uds: {} /{path} → {method} (stream)", parts.method);
stream_response(&socket, method, params, STREAM_OPEN_TIMEOUT).await
}
}
}
pub async fn deprecated_search_api_handler(
state: State<AppState>,
path: AxumPath<String>,
req: Request,
) -> Response {
tracing::trace!("search_uds: DEPRECATED /proxy/search/… — use /api/search/… instead (#1849)");
search_api_handler(state, path, req).await
}
async fn stream_response(
socket: &Path,
method: &'static str,
params: Value,
open_timeout: std::time::Duration,
) -> Response {
let deadline = tokio::time::Instant::now() + open_timeout;
let mut stream = match open_stream(socket, method, params, open_timeout).await {
Ok(s) => s,
Err(e) => return e.into_response(),
};
let peeked = match tokio::time::timeout_at(deadline, stream.next_frame()).await {
Ok(frame) => frame,
Err(_) => {
return SearchRpcError::Unreachable(format!(
"{} did not answer {method} within {}s",
super::SEARCH_SERVICE,
open_timeout.as_secs_f32()
))
.into_response();
}
};
let first = match peeked {
Some(Ok(item)) => Some(item),
Some(Err(e)) => return stream_error(method, e).into_response(),
None => None,
};
let head = futures_util::stream::iter(
first
.into_iter()
.map(|item| Ok::<Bytes, std::convert::Infallible>(sse_data(&item))),
);
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/event-stream")
.header(header::CACHE_CONTROL, "no-cache")
.header("X-Accel-Buffering", "no")
.body(Body::from_stream(head.chain(sse_tail(stream, method))))
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
}
fn sse_tail(
mut stream: FramedStream<Value>,
method: &'static str,
) -> impl futures_util::Stream<Item = Result<Bytes, std::convert::Infallible>> {
let (tx, rx) = mpsc::channel::<Result<Value, UdsRpcError>>(SSE_BUFFER);
tokio::spawn(async move {
loop {
let item = tokio::select! {
biased;
() = tx.closed() => return,
item = stream.next_frame() => item,
};
let Some(item) = item else { return };
let terminal = item.is_err();
if tx.send(item).await.is_err() || terminal {
return;
}
}
});
let heartbeat = tokio::time::interval_at(
tokio::time::Instant::now() + SSE_HEARTBEAT_INTERVAL,
SSE_HEARTBEAT_INTERVAL,
);
futures_util::stream::unfold(Some((rx, heartbeat)), move |state| async move {
let (mut rx, mut heartbeat) = state?;
tokio::select! {
biased;
item = rx.recv() => match item {
Some(Ok(value)) => Some((Ok(sse_data(&value)), Some((rx, heartbeat)))),
Some(Err(e)) => {
warn!("search_uds: {method} failed mid-stream: {e}");
let event = json!({ "type": "error", "message": e.to_string() });
Some((Ok(sse_data(&event)), None))
}
None => None,
},
_ = heartbeat.tick() => Some((
Ok(Bytes::from_static(b": heartbeat\n\n")),
Some((rx, heartbeat)),
)),
}
})
}
fn sse_data(value: &Value) -> Bytes {
Bytes::from(format!("data: {value}\n\n"))
}
fn stream_error(method: &str, e: UdsRpcError) -> SearchRpcError {
match e {
UdsRpcError::Stream { error, .. } => SearchRpcError::Refused {
code: error.code,
message: error.message,
},
other => SearchRpcError::Unreachable(format!(
"{} did not stream {method}: {other}",
super::SEARCH_SERVICE
)),
}
}
impl AppState {
pub(crate) fn search_socket_path(&self) -> Result<PathBuf, String> {
match &self.search_socket {
Some(p) => Ok(p.as_ref().clone()),
None => super::socket_path(),
}
}
#[must_use]
pub fn with_search_socket(mut self, socket: PathBuf) -> Self {
self.search_socket = Some(std::sync::Arc::new(socket));
self
}
}
#[cfg(test)]
mod tests {
use super::*;
use trusty_common::uds::server::RpcError;
#[test]
fn sse_data_is_one_line_per_event() {
let framed = sse_data(&json!({ "message": "a\nb" }));
let text = String::from_utf8(framed.to_vec()).expect("utf-8");
assert_eq!(text, "data: {\"message\":\"a\\nb\"}\n\n");
assert_eq!(text.matches("\n\n").count(), 1);
}
#[test]
fn stream_error_carries_the_daemon_code() {
let err = stream_error(
"search.index.reindex.stream",
UdsRpcError::Stream {
path: PathBuf::from("/tmp/x.sock"),
error: RpcError::new(-32004, "no reindex in progress for 'ghost'"),
},
);
assert_eq!(err.status(), StatusCode::NOT_FOUND);
assert!(err.message().contains("ghost"), "{err:?}");
}
#[tokio::test(flavor = "multi_thread")]
async fn a_socket_that_never_answers_is_a_prompt_bad_gateway() {
let tmp = tempfile::TempDir::new().expect("tempdir");
let socket = tmp.path().join("silent.sock");
let listener = trusty_common::uds::bind_hardened(&socket).expect("bind");
let _accepting = tokio::spawn(async move {
let Ok((mut conn, _)) = listener.accept().await else {
return;
};
let mut sink = Vec::new();
let _ = tokio::io::AsyncReadExt::read_to_end(&mut conn, &mut sink).await;
std::future::pending::<()>().await;
});
let started = std::time::Instant::now();
let response = stream_response(
&socket,
super::super::METHOD_STATUS_STREAM,
json!({}),
std::time::Duration::from_millis(300),
)
.await;
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
assert!(
started.elapsed() < std::time::Duration::from_secs(5),
"the open must be bounded, not left to the per-frame budget: {:?}",
started.elapsed()
);
}
fn stalls_then_drains(socket: PathBuf, drain_after: std::time::Duration) -> PathBuf {
let listener = trusty_common::uds::bind_hardened(&socket).expect("bind");
tokio::spawn(async move {
let Ok((mut conn, _)) = listener.accept().await else {
return;
};
tokio::time::sleep(drain_after).await;
let mut sink = Vec::new();
let _ = tokio::io::AsyncReadExt::read_to_end(&mut conn, &mut sink).await;
std::future::pending::<()>().await;
});
socket
}
fn bulky_params() -> Value {
json!({ "blob": "x".repeat(512 * 1024) })
}
#[tokio::test(flavor = "multi_thread")]
async fn a_slow_open_and_a_silent_first_frame_share_one_budget() {
const BUDGET: std::time::Duration = std::time::Duration::from_millis(1000);
const DRAIN_AFTER: std::time::Duration = std::time::Duration::from_millis(600);
const CEILING: std::time::Duration = std::time::Duration::from_millis(1300);
let tmp = tempfile::TempDir::new().expect("tempdir");
let slow = stalls_then_drains(tmp.path().join("slow-open.sock"), DRAIN_AFTER);
let started = std::time::Instant::now();
let opened = open_stream(
&slow,
super::super::METHOD_STATUS_STREAM,
bulky_params(),
BUDGET,
)
.await;
let open_took = started.elapsed();
assert!(opened.is_ok(), "the open must succeed, slowly: {opened:?}");
assert!(
open_took >= DRAIN_AFTER,
"the write must park until the peer drains, or this test proves nothing: {open_took:?}"
);
drop(opened);
let silent = stalls_then_drains(tmp.path().join("silent-frame.sock"), DRAIN_AFTER);
let started = std::time::Instant::now();
let response = stream_response(
&silent,
super::super::METHOD_STATUS_STREAM,
bulky_params(),
BUDGET,
)
.await;
let elapsed = started.elapsed();
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
assert!(
elapsed >= BUDGET,
"a shared deadline still spends the whole budget: {elapsed:?}"
);
assert!(
elapsed < CEILING,
"the dial, the write and the first frame read must share ONE {BUDGET:?} deadline, \
not take one each: {elapsed:?}"
);
}
#[test]
fn stream_error_reports_a_transport_failure_as_unreachable() {
let err = stream_error(
"search.status.stream",
UdsRpcError::NoResponse {
path: PathBuf::from("/tmp/x.sock"),
},
);
assert_eq!(err.status(), StatusCode::BAD_GATEWAY);
}
}