shell_tunnel/api/websocket.rs
1//! WebSocket handler for real-time command streaming.
2
3use std::time::Duration;
4
5use axum::{
6 extract::{
7 ws::{Message, WebSocket, WebSocketUpgrade},
8 Path, State,
9 },
10 response::IntoResponse,
11};
12use futures_util::{SinkExt, StreamExt};
13
14use super::handlers::AppState;
15use super::types::{WsClientMessage, WsServerMessage};
16use crate::execution::Command;
17use crate::session::{BusySession, SessionId};
18
19/// WebSocket upgrade handler.
20pub async fn ws_handler(
21 ws: WebSocketUpgrade,
22 State(state): State<AppState>,
23 Path(session_id): Path<u64>,
24 identity: Option<axum::Extension<crate::audit::Identity>>,
25) -> impl IntoResponse {
26 // Taken here because extensions belong to the upgrade request, not to the
27 // socket that outlives it.
28 let identity = identity.map(|axum::Extension(id)| id);
29 ws.on_upgrade(move |socket| handle_socket(socket, state, session_id, identity))
30}
31
32/// Handle WebSocket connection.
33async fn handle_socket(
34 socket: WebSocket,
35 state: AppState,
36 session_id: u64,
37 identity: Option<crate::audit::Identity>,
38) {
39 let id = SessionId::from_raw(session_id);
40
41 // Verify session exists
42 if state.store.get(&id).ok().flatten().is_none() {
43 let (mut sink, _) = socket.split();
44 let err = WsServerMessage::Error {
45 code: "SESSION_NOT_FOUND".to_string(),
46 message: format!("Session {} not found", session_id),
47 };
48 if let Ok(json) = serde_json::to_string(&err) {
49 let _ = sink.send(Message::Text(json.into())).await;
50 }
51 return;
52 }
53
54 let (mut sink, mut stream) = socket.split();
55
56 // Process incoming messages
57 while let Some(msg) = stream.next().await {
58 let msg = match msg {
59 Ok(Message::Text(text)) => text.to_string(),
60 Ok(Message::Close(_)) => break,
61 Ok(Message::Ping(data)) => {
62 let _ = sink.send(Message::Pong(data)).await;
63 continue;
64 }
65 Ok(_) => continue,
66 Err(_) => break,
67 };
68
69 // Parse WebSocket message
70 let ws_msg: WsClientMessage = match serde_json::from_str(&msg) {
71 Ok(m) => m,
72 Err(e) => {
73 let err = WsServerMessage::Error {
74 code: "PARSE_ERROR".to_string(),
75 message: e.to_string(),
76 };
77 if let Ok(json) = serde_json::to_string(&err) {
78 let _ = sink.send(Message::Text(json.into())).await;
79 }
80 continue;
81 }
82 };
83
84 match ws_msg {
85 WsClientMessage::Execute {
86 command,
87 timeout_secs,
88 } => {
89 // Build command
90 let mut cmd = Command::new(&command);
91 if let Some(secs) = timeout_secs {
92 cmd = cmd.timeout(Duration::from_secs(secs));
93 }
94
95 // Busy for the whole command, exactly as the REST path is. This
96 // handler streams through `execute_async` rather than
97 // `execute_in_session`, so nothing else here touches the session
98 // — without it a command driven over the socket left the session
99 // reporting `running: false` and its idle clock running while a
100 // build was under way. The guard also covers the ways out that
101 // no branch below expresses.
102 //
103 // Held in an `Option` so the delivery loop can put it down the
104 // moment the command itself is over: see `while_running`.
105 let mut busy = BusySession::begin(&state.store, &id).ok();
106
107 // Execute with streaming
108 match state.executor.execute_async(&cmd).await {
109 Ok((mut rx, mut handle)) => {
110 // The command's own outcome, if it finished while output
111 // was still being delivered — which is the normal case
112 // for a consumer that reads slowly or not at all.
113 let mut finished = None;
114
115 // Stream output chunks
116 while let Some(chunk) = rx.recv().await {
117 let output = WsServerMessage::Output {
118 data: String::from_utf8_lossy(&chunk.raw).to_string(),
119 is_final: false,
120 };
121 if let Ok(json) = serde_json::to_string(&output) {
122 let sent = while_running(
123 sink.send(Message::Text(json.into())),
124 &mut handle,
125 &mut finished,
126 &mut busy,
127 )
128 .await;
129 if sent.is_err() {
130 break;
131 }
132 }
133 }
134
135 // Let the command go once nobody is reading it. Holding
136 // the receiver across the await below would stall the
137 // loop that enforces the timeout — see `execute_async`.
138 drop(rx);
139
140 // Wait for completion and send result. `while_running`
141 // has usually taken it already — awaiting a handle it
142 // has consumed would panic, which is why it is stored
143 // rather than re-awaited.
144 let outcome = match finished {
145 Some(res) => res,
146 None => handle.await,
147 };
148 match outcome {
149 Ok(Ok(result)) => {
150 state
151 .audit
152 .record_async(
153 crate::audit::AuditEvent::new("execute")
154 .with_identity(identity.clone())
155 .with_route("WS /api/v1/sessions/{id}/ws")
156 .with_command(&command)
157 .with_session(session_id)
158 .with_outcome(
159 result.exit_code,
160 result.timed_out,
161 result.duration.as_millis() as u64,
162 ),
163 )
164 .await;
165
166 // Update session context
167 state
168 .store
169 .update(&id, |s| {
170 s.context.record_execution(&command, result.exit_code);
171 })
172 .ok();
173
174 let result_msg = WsServerMessage::Result {
175 success: result.exit_code.map(|c| c == 0).unwrap_or(false)
176 && !result.timed_out,
177 exit_code: result.exit_code,
178 duration_ms: result.duration.as_millis() as u64,
179 timed_out: result.timed_out,
180 total_bytes: result.total_bytes,
181 };
182 if let Ok(json) = serde_json::to_string(&result_msg) {
183 let _ = sink.send(Message::Text(json.into())).await;
184 }
185 }
186 Ok(Err(e)) => {
187 let err = WsServerMessage::Error {
188 code: "EXECUTION_ERROR".to_string(),
189 message: e.to_string(),
190 };
191 if let Ok(json) = serde_json::to_string(&err) {
192 let _ = sink.send(Message::Text(json.into())).await;
193 }
194 }
195 Err(e) => {
196 let err = WsServerMessage::Error {
197 code: "TASK_ERROR".to_string(),
198 message: e.to_string(),
199 };
200 if let Ok(json) = serde_json::to_string(&err) {
201 let _ = sink.send(Message::Text(json.into())).await;
202 }
203 }
204 }
205 }
206 Err(e) => {
207 let err = WsServerMessage::Error {
208 code: "EXECUTION_ERROR".to_string(),
209 message: e.to_string(),
210 };
211 if let Ok(json) = serde_json::to_string(&err) {
212 let _ = sink.send(Message::Text(json.into())).await;
213 }
214 }
215 }
216 // `busy` drops here if `while_running` did not already put it
217 // down: idle again on every way out, including the ones that
218 // never reached the executor and the ones no branch here
219 // expresses.
220 }
221 WsClientMessage::Ping => {
222 let pong = WsServerMessage::Pong;
223 if let Ok(json) = serde_json::to_string(&pong) {
224 let _ = sink.send(Message::Text(json.into())).await;
225 }
226 }
227 _ => {
228 // Ignore other message types from client
229 }
230 }
231 }
232}
233
234/// The task running one streamed command, as `execute_async` hands it over.
235type CommandHandle = tokio::task::JoinHandle<crate::Result<crate::execution::ExecutionResult>>;
236
237/// What that task produced: the command's result, or the join error if the task
238/// itself came apart.
239type CommandOutcome =
240 std::result::Result<crate::Result<crate::execution::ExecutionResult>, tokio::task::JoinError>;
241
242/// Await `fut`, releasing the session's busy guard the moment the *command*
243/// ends rather than when delivery does.
244///
245/// The two are not the same wait, and the gap between them is unbounded. A
246/// consumer that stops reading its socket parks `sink.send` for as long as it
247/// likes; the command still dies at its own deadline (`forward_chunk` stops
248/// waiting on a stalled receiver there), but the handler is still inside a send
249/// and cannot reach the end of the arm where the guard would drop. Measured on
250/// 0.21.0: a command with `timeout_secs: 5` died at 5.005 s and its session went
251/// on reporting `running: true` for **75 seconds**, ending not at any deadline
252/// but at the instant the consumer resumed reading. A silent command released at
253/// 5.3 s, which is what says the delay belongs to undelivered output rather than
254/// to the deadline. The direct and relayed paths behaved identically.
255///
256/// That mattered beyond the field's own honesty: the idle sweep skips sessions
257/// in [`SessionState::Active`](crate::session::SessionState::Active), on the
258/// stated grounds that a command's deadline bounds how long that can last. A
259/// session pinned by an unread socket is therefore never swept at all.
260///
261/// `fut` is pinned once and re-polled rather than re-created, so a send that was
262/// half-written when the command finished is resumed, not cancelled — cancelling
263/// mid-frame would corrupt the stream. The handle branch is disabled after it
264/// fires, since polling a completed `JoinHandle` panics.
265///
266/// This bounds the *session*, not the socket: a consumer that never reads again
267/// still holds the connection, which is the separate open item (there is no
268/// WebSocket idle timeout).
269async fn while_running<F: std::future::Future>(
270 fut: F,
271 handle: &mut CommandHandle,
272 finished: &mut Option<CommandOutcome>,
273 busy: &mut Option<BusySession>,
274) -> F::Output {
275 tokio::pin!(fut);
276 loop {
277 if finished.is_some() {
278 return fut.await;
279 }
280 tokio::select! {
281 out = &mut fut => return out,
282 res = &mut *handle => {
283 *finished = Some(res);
284 // Dropping the guard returns the session to idle and restarts
285 // its clock from the command's end, which is the instant the
286 // clock is meant to measure from.
287 *busy = None;
288 }
289 }
290 }
291}
292
293/// One-shot WebSocket execution (no session required).
294pub async fn ws_oneshot_handler(
295 ws: WebSocketUpgrade,
296 State(state): State<AppState>,
297 identity: Option<axum::Extension<crate::audit::Identity>>,
298) -> impl IntoResponse {
299 let identity = identity.map(|axum::Extension(id)| id);
300 ws.on_upgrade(move |socket| handle_oneshot_socket(socket, state, identity))
301}
302
303/// Handle one-shot WebSocket connection.
304async fn handle_oneshot_socket(
305 socket: WebSocket,
306 state: AppState,
307 identity: Option<crate::audit::Identity>,
308) {
309 let (mut sink, mut stream) = socket.split();
310
311 while let Some(msg) = stream.next().await {
312 let msg = match msg {
313 Ok(Message::Text(text)) => text.to_string(),
314 Ok(Message::Close(_)) => break,
315 Ok(Message::Ping(data)) => {
316 let _ = sink.send(Message::Pong(data)).await;
317 continue;
318 }
319 Ok(_) => continue,
320 Err(_) => break,
321 };
322
323 let ws_msg: WsClientMessage = match serde_json::from_str(&msg) {
324 Ok(m) => m,
325 Err(e) => {
326 let err = WsServerMessage::Error {
327 code: "PARSE_ERROR".to_string(),
328 message: e.to_string(),
329 };
330 if let Ok(json) = serde_json::to_string(&err) {
331 let _ = sink.send(Message::Text(json.into())).await;
332 }
333 continue;
334 }
335 };
336
337 match ws_msg {
338 WsClientMessage::Execute {
339 command,
340 timeout_secs,
341 } => {
342 let mut cmd = Command::new(&command);
343 if let Some(secs) = timeout_secs {
344 cmd = cmd.timeout(Duration::from_secs(secs));
345 }
346
347 match state.executor.execute_async(&cmd).await {
348 Ok((mut rx, handle)) => {
349 while let Some(chunk) = rx.recv().await {
350 let output = WsServerMessage::Output {
351 data: String::from_utf8_lossy(&chunk.raw).to_string(),
352 is_final: false,
353 };
354 if let Ok(json) = serde_json::to_string(&output) {
355 if sink.send(Message::Text(json.into())).await.is_err() {
356 break;
357 }
358 }
359 }
360
361 // Nobody is reading any more: release the command so its
362 // timeout can still be enforced (see `execute_async`).
363 // This path has no session, so a stalled command here
364 // showed up only as a child that never died.
365 drop(rx);
366
367 match handle.await {
368 Ok(Ok(result)) => {
369 state
370 .audit
371 .record_async(
372 crate::audit::AuditEvent::new("execute")
373 .with_identity(identity.clone())
374 .with_route("WS /api/v1/ws")
375 .with_command(&command)
376 .with_outcome(
377 result.exit_code,
378 result.timed_out,
379 result.duration.as_millis() as u64,
380 ),
381 )
382 .await;
383
384 let result_msg = WsServerMessage::Result {
385 success: result.exit_code.map(|c| c == 0).unwrap_or(false)
386 && !result.timed_out,
387 exit_code: result.exit_code,
388 duration_ms: result.duration.as_millis() as u64,
389 timed_out: result.timed_out,
390 total_bytes: result.total_bytes,
391 };
392 if let Ok(json) = serde_json::to_string(&result_msg) {
393 let _ = sink.send(Message::Text(json.into())).await;
394 }
395 }
396 Ok(Err(e)) => {
397 let err = WsServerMessage::Error {
398 code: "EXECUTION_ERROR".to_string(),
399 message: e.to_string(),
400 };
401 if let Ok(json) = serde_json::to_string(&err) {
402 let _ = sink.send(Message::Text(json.into())).await;
403 }
404 }
405 Err(e) => {
406 let err = WsServerMessage::Error {
407 code: "TASK_ERROR".to_string(),
408 message: e.to_string(),
409 };
410 if let Ok(json) = serde_json::to_string(&err) {
411 let _ = sink.send(Message::Text(json.into())).await;
412 }
413 }
414 }
415 }
416 Err(e) => {
417 let err = WsServerMessage::Error {
418 code: "EXECUTION_ERROR".to_string(),
419 message: e.to_string(),
420 };
421 if let Ok(json) = serde_json::to_string(&err) {
422 let _ = sink.send(Message::Text(json.into())).await;
423 }
424 }
425 }
426 }
427 WsClientMessage::Ping => {
428 let pong = WsServerMessage::Pong;
429 if let Ok(json) = serde_json::to_string(&pong) {
430 let _ = sink.send(Message::Text(json.into())).await;
431 }
432 }
433 _ => {}
434 }
435 }
436}
437
438#[cfg(test)]
439mod tests {
440 use super::*;
441
442 #[test]
443 fn test_ws_message_execute_parse() {
444 let json = r#"{"type": "execute", "command": "echo hello"}"#;
445 let msg: WsClientMessage = serde_json::from_str(json).unwrap();
446 match msg {
447 WsClientMessage::Execute { command, .. } => assert_eq!(command, "echo hello"),
448 _ => panic!("Expected Execute message"),
449 }
450 }
451
452 #[test]
453 fn test_ws_message_ping_parse() {
454 let json = r#"{"type": "ping"}"#;
455 let msg: WsClientMessage = serde_json::from_str(json).unwrap();
456 assert!(matches!(msg, WsClientMessage::Ping));
457 }
458}