diff --git a/src/api/code_runner/mod.rs b/src/api/code_runner/mod.rs index e66e3bd..fe6ba4c 100644 --- a/src/api/code_runner/mod.rs +++ b/src/api/code_runner/mod.rs @@ -77,4 +77,6 @@ pub struct ExecTask { pub mod languages; #[cfg(feature = "server")] pub mod progress; +#[cfg(feature = "server")] +pub mod sse; pub mod execute; diff --git a/src/api/code_runner/sse.rs b/src/api/code_runner/sse.rs new file mode 100644 index 0000000..5be6001 --- /dev/null +++ b/src/api/code_runner/sse.rs @@ -0,0 +1,71 @@ +//! SSE 流式输出端点。`GET /api/exec/stream?task_id=X`。 +//! +//! 鉴权 + 限流已在 [`super::execute::start_exec_stream`] 完成(校验链),此处只 +//! 校验 task 存在。task_id 是 UUID,不可枚举;只有过了校验的调用者才知道 task_id。 +//! +//! 前端用原生 EventSource 连接本端点(`web_sys::EventSource`),按 event 类型 +//! 分发:`stdout` → 终端 writeStdout,`stderr` → writeStderr,`done` → 终态收尾。 +//! keep-alive comment 每 15s 一次,防反向代理超时关闭空闲连接。 + +use std::convert::Infallible; +use std::time::Duration; + +use axum::extract::Query; +use axum::http::StatusCode; +use axum::response::sse::{Event, KeepAlive, Sse}; +use futures::StreamExt; +use tokio_stream::wrappers::ReceiverStream; + +use crate::api::code_runner::progress::EXEC_STREAMS; +use crate::infra::docker::OutputChunk; + +/// SSE 查询参数:`?task_id=X`。 +#[derive(serde::Deserialize)] +pub struct StreamQuery { + pub task_id: String, +} + +/// done 事件的 JSON payload。 +#[derive(serde::Serialize)] +struct DonePayload { + exit_code: Option, + oom_killed: bool, + timed_out: bool, +} + +/// SSE handler:从 EXEC_STREAMS 取出 receiver(取出即移除,防重复连接), +/// 包成 ReceiverStream,映射成 SSE Event 流。 +pub async fn exec_stream( + Query(q): Query, +) -> Result>>, (StatusCode, String)> { + // 取出 rx(取出即移除):同一 task_id 只能连一次 SSE, + // 防止多客户端或重连导致 receiver 被多次消费。 + let (_, entry) = EXEC_STREAMS + .remove(&q.task_id) + .ok_or((StatusCode::NOT_FOUND, "任务不存在或已结束".to_string()))?; + + let stream = ReceiverStream::new(entry.rx).map(|chunk| { + Ok::<_, Infallible>(match chunk { + OutputChunk::Stdout(s) => Event::default().event("stdout").data(s), + OutputChunk::Stderr(s) => Event::default().event("stderr").data(s), + OutputChunk::Done { + exit_code, + oom_killed, + timed_out, + } => Event::default() + .event("done") + .json_data(DonePayload { + exit_code, + oom_killed, + timed_out, + }) + .unwrap_or_else(|_| Event::default().event("done").data("{}")), + }) + }); + + Ok(Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(15)) + .text("keep-alive"), + )) +} diff --git a/src/main.rs b/src/main.rs index b4f92e1..e1f959b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -474,6 +474,19 @@ fn main() { )) .layer(axum::middleware::from_fn(crate::api::csrf::csrf_middleware)); + // SSE 流式输出端点:GET /api/exec/stream?task_id=X + // 不挂 TimeoutLayer!SSE 是长连接,30s timeout 会杀掉流。 + // 鉴权 + 限流已在 start_exec_stream server function 完成(校验链), + // 此处只校验 task_id 存在;CSRF 对 GET 放行(is_write_method 返回 false)。 + // CompressionLayer 跳过 text/event-stream(见 compression_layer_from_env 注释), + // 但 sse_route 本身不挂 compression,更安全。 + let sse_route = axum::Router::new() + .route( + "/api/exec/stream", + axum::routing::get(crate::api::code_runner::sse::exec_stream), + ) + .layer(axum::middleware::from_fn(crate::api::csrf::csrf_middleware)); + // Dioxus 应用路由:自动挂载所有 server function 并渲染前端组件 let dioxus_app = axum::Router::new().serve_dioxus_application(config, router::AppRouter); @@ -515,8 +528,12 @@ fn main() { axum::routing::get(|| async { StatusCode::NOT_FOUND }), ); - // 合并:upload 路由保持自己独立的 300s 超时;export 路由 120s;app routes 加可选压缩/30s;static routes 无任何中间件 - let router = upload_route.merge(export_route).merge(app_routes).merge(static_routes); + // 合并:upload 路由 300s 超时;export 路由 120s;sse 路由无超时(长连接);app routes 加可选压缩/30s;static routes 无任何中间件 + let router = upload_route + .merge(export_route) + .merge(sse_route) + .merge(app_routes) + .merge(static_routes); Ok(router) });