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                let _busy = BusySession::begin(&state.store, &id).ok();
103
104                // Execute with streaming
105                match state.executor.execute_async(&cmd).await {
106                    Ok((mut rx, handle)) => {
107                        // Stream output chunks
108                        while let Some(chunk) = rx.recv().await {
109                            let output = WsServerMessage::Output {
110                                data: String::from_utf8_lossy(&chunk.raw).to_string(),
111                                is_final: false,
112                            };
113                            if let Ok(json) = serde_json::to_string(&output) {
114                                if sink.send(Message::Text(json.into())).await.is_err() {
115                                    break;
116                                }
117                            }
118                        }
119
120                        // Let the command go once nobody is reading it. Holding
121                        // the receiver across the await below would stall the
122                        // loop that enforces the timeout — see `execute_async`.
123                        drop(rx);
124
125                        // Wait for completion and send result
126                        match handle.await {
127                            Ok(Ok(result)) => {
128                                state
129                                    .audit
130                                    .record_async(
131                                        crate::audit::AuditEvent::new("execute")
132                                            .with_identity(identity.clone())
133                                            .with_route("WS /api/v1/sessions/{id}/ws")
134                                            .with_command(&command)
135                                            .with_session(session_id)
136                                            .with_outcome(
137                                                result.exit_code,
138                                                result.timed_out,
139                                                result.duration.as_millis() as u64,
140                                            ),
141                                    )
142                                    .await;
143
144                                // Update session context
145                                state
146                                    .store
147                                    .update(&id, |s| {
148                                        s.context.record_execution(&command, result.exit_code);
149                                    })
150                                    .ok();
151
152                                let result_msg = WsServerMessage::Result {
153                                    success: result.exit_code.map(|c| c == 0).unwrap_or(false)
154                                        && !result.timed_out,
155                                    exit_code: result.exit_code,
156                                    duration_ms: result.duration.as_millis() as u64,
157                                    timed_out: result.timed_out,
158                                    total_bytes: result.total_bytes,
159                                };
160                                if let Ok(json) = serde_json::to_string(&result_msg) {
161                                    let _ = sink.send(Message::Text(json.into())).await;
162                                }
163                            }
164                            Ok(Err(e)) => {
165                                let err = WsServerMessage::Error {
166                                    code: "EXECUTION_ERROR".to_string(),
167                                    message: e.to_string(),
168                                };
169                                if let Ok(json) = serde_json::to_string(&err) {
170                                    let _ = sink.send(Message::Text(json.into())).await;
171                                }
172                            }
173                            Err(e) => {
174                                let err = WsServerMessage::Error {
175                                    code: "TASK_ERROR".to_string(),
176                                    message: e.to_string(),
177                                };
178                                if let Ok(json) = serde_json::to_string(&err) {
179                                    let _ = sink.send(Message::Text(json.into())).await;
180                                }
181                            }
182                        }
183                    }
184                    Err(e) => {
185                        let err = WsServerMessage::Error {
186                            code: "EXECUTION_ERROR".to_string(),
187                            message: e.to_string(),
188                        };
189                        if let Ok(json) = serde_json::to_string(&err) {
190                            let _ = sink.send(Message::Text(json.into())).await;
191                        }
192                    }
193                }
194                // `_busy` drops here: idle again on every way out, including the
195                // ones that never reached the executor and the ones no branch
196                // here expresses.
197            }
198            WsClientMessage::Ping => {
199                let pong = WsServerMessage::Pong;
200                if let Ok(json) = serde_json::to_string(&pong) {
201                    let _ = sink.send(Message::Text(json.into())).await;
202                }
203            }
204            _ => {
205                // Ignore other message types from client
206            }
207        }
208    }
209}
210
211/// One-shot WebSocket execution (no session required).
212pub async fn ws_oneshot_handler(
213    ws: WebSocketUpgrade,
214    State(state): State<AppState>,
215    identity: Option<axum::Extension<crate::audit::Identity>>,
216) -> impl IntoResponse {
217    let identity = identity.map(|axum::Extension(id)| id);
218    ws.on_upgrade(move |socket| handle_oneshot_socket(socket, state, identity))
219}
220
221/// Handle one-shot WebSocket connection.
222async fn handle_oneshot_socket(
223    socket: WebSocket,
224    state: AppState,
225    identity: Option<crate::audit::Identity>,
226) {
227    let (mut sink, mut stream) = socket.split();
228
229    while let Some(msg) = stream.next().await {
230        let msg = match msg {
231            Ok(Message::Text(text)) => text.to_string(),
232            Ok(Message::Close(_)) => break,
233            Ok(Message::Ping(data)) => {
234                let _ = sink.send(Message::Pong(data)).await;
235                continue;
236            }
237            Ok(_) => continue,
238            Err(_) => break,
239        };
240
241        let ws_msg: WsClientMessage = match serde_json::from_str(&msg) {
242            Ok(m) => m,
243            Err(e) => {
244                let err = WsServerMessage::Error {
245                    code: "PARSE_ERROR".to_string(),
246                    message: e.to_string(),
247                };
248                if let Ok(json) = serde_json::to_string(&err) {
249                    let _ = sink.send(Message::Text(json.into())).await;
250                }
251                continue;
252            }
253        };
254
255        match ws_msg {
256            WsClientMessage::Execute {
257                command,
258                timeout_secs,
259            } => {
260                let mut cmd = Command::new(&command);
261                if let Some(secs) = timeout_secs {
262                    cmd = cmd.timeout(Duration::from_secs(secs));
263                }
264
265                match state.executor.execute_async(&cmd).await {
266                    Ok((mut rx, handle)) => {
267                        while let Some(chunk) = rx.recv().await {
268                            let output = WsServerMessage::Output {
269                                data: String::from_utf8_lossy(&chunk.raw).to_string(),
270                                is_final: false,
271                            };
272                            if let Ok(json) = serde_json::to_string(&output) {
273                                if sink.send(Message::Text(json.into())).await.is_err() {
274                                    break;
275                                }
276                            }
277                        }
278
279                        // Nobody is reading any more: release the command so its
280                        // timeout can still be enforced (see `execute_async`).
281                        // This path has no session, so a stalled command here
282                        // showed up only as a child that never died.
283                        drop(rx);
284
285                        match handle.await {
286                            Ok(Ok(result)) => {
287                                state
288                                    .audit
289                                    .record_async(
290                                        crate::audit::AuditEvent::new("execute")
291                                            .with_identity(identity.clone())
292                                            .with_route("WS /api/v1/ws")
293                                            .with_command(&command)
294                                            .with_outcome(
295                                                result.exit_code,
296                                                result.timed_out,
297                                                result.duration.as_millis() as u64,
298                                            ),
299                                    )
300                                    .await;
301
302                                let result_msg = WsServerMessage::Result {
303                                    success: result.exit_code.map(|c| c == 0).unwrap_or(false)
304                                        && !result.timed_out,
305                                    exit_code: result.exit_code,
306                                    duration_ms: result.duration.as_millis() as u64,
307                                    timed_out: result.timed_out,
308                                    total_bytes: result.total_bytes,
309                                };
310                                if let Ok(json) = serde_json::to_string(&result_msg) {
311                                    let _ = sink.send(Message::Text(json.into())).await;
312                                }
313                            }
314                            Ok(Err(e)) => {
315                                let err = WsServerMessage::Error {
316                                    code: "EXECUTION_ERROR".to_string(),
317                                    message: e.to_string(),
318                                };
319                                if let Ok(json) = serde_json::to_string(&err) {
320                                    let _ = sink.send(Message::Text(json.into())).await;
321                                }
322                            }
323                            Err(e) => {
324                                let err = WsServerMessage::Error {
325                                    code: "TASK_ERROR".to_string(),
326                                    message: e.to_string(),
327                                };
328                                if let Ok(json) = serde_json::to_string(&err) {
329                                    let _ = sink.send(Message::Text(json.into())).await;
330                                }
331                            }
332                        }
333                    }
334                    Err(e) => {
335                        let err = WsServerMessage::Error {
336                            code: "EXECUTION_ERROR".to_string(),
337                            message: e.to_string(),
338                        };
339                        if let Ok(json) = serde_json::to_string(&err) {
340                            let _ = sink.send(Message::Text(json.into())).await;
341                        }
342                    }
343                }
344            }
345            WsClientMessage::Ping => {
346                let pong = WsServerMessage::Pong;
347                if let Ok(json) = serde_json::to_string(&pong) {
348                    let _ = sink.send(Message::Text(json.into())).await;
349                }
350            }
351            _ => {}
352        }
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359
360    #[test]
361    fn test_ws_message_execute_parse() {
362        let json = r#"{"type": "execute", "command": "echo hello"}"#;
363        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
364        match msg {
365            WsClientMessage::Execute { command, .. } => assert_eq!(command, "echo hello"),
366            _ => panic!("Expected Execute message"),
367        }
368    }
369
370    #[test]
371    fn test_ws_message_ping_parse() {
372        let json = r#"{"type": "ping"}"#;
373        let msg: WsClientMessage = serde_json::from_str(json).unwrap();
374        assert!(matches!(msg, WsClientMessage::Ping));
375    }
376}
shell_tunnel/api/websocket.rs

shell_tunnel/api/
websocket.rs
