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

shell_tunnel/api/
websocket.rs
