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

shell_tunnel/api/
websocket.rs
