//! Terminal proxy handler - reverse proxy to ttyd via Unix socket use axum::{ body::Body, extract::{ ws::{Message as AxumMessage, WebSocket, WebSocketUpgrade}, OriginalUri, Path, State, }, http::{Request, StatusCode}, response::Response, }; use futures::{SinkExt, StreamExt}; use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::UnixStream; use tokio_tungstenite::tungstenite::{ client::IntoClientRequest, http::HeaderValue, Message as TungsteniteMessage, }; use crate::error::AppError; use crate::extensions::TTYD_SOCKET_PATH; use crate::state::AppState; /// Handle WebSocket upgrade for terminal pub async fn terminal_ws( State(_state): State>, OriginalUri(original_uri): OriginalUri, ws: WebSocketUpgrade, ) -> Response { let query_string = original_uri .query() .map(|q| format!("?{}", q)) .unwrap_or_default(); // Use the tty subprotocol that ttyd expects ws.protocols(["tty"]) .on_upgrade(move |socket| handle_terminal_websocket(socket, query_string)) } /// Handle terminal WebSocket connection - bridge browser and ttyd async fn handle_terminal_websocket(client_ws: WebSocket, query_string: String) { // Connect to ttyd Unix socket let unix_stream = match UnixStream::connect(TTYD_SOCKET_PATH).await { Ok(s) => s, Err(e) => { tracing::error!("Failed to connect to ttyd socket: {}", e); return; } }; // Build WebSocket request for ttyd with tty subprotocol let uri_str = format!("ws://localhost/api/terminal/ws{}", query_string); let mut request = match uri_str.into_client_request() { Ok(r) => r, Err(e) => { tracing::error!("Failed to create WebSocket request: {}", e); return; } }; request .headers_mut() .insert("Sec-WebSocket-Protocol", HeaderValue::from_static("tty")); // Create WebSocket connection to ttyd let ws_stream = match tokio_tungstenite::client_async(request, unix_stream).await { Ok((ws, _)) => ws, Err(e) => { tracing::error!("Failed to establish WebSocket with ttyd: {}", e); return; } }; // Split both WebSocket connections let (mut client_tx, mut client_rx) = client_ws.split(); let (mut ttyd_tx, mut ttyd_rx) = ws_stream.split(); // Forward messages from browser to ttyd let client_to_ttyd = tokio::spawn(async move { while let Some(msg) = client_rx.next().await { let ttyd_msg = match msg { Ok(AxumMessage::Text(text)) => TungsteniteMessage::Text(text.to_string().into()), Ok(AxumMessage::Binary(data)) => TungsteniteMessage::Binary(data), Ok(AxumMessage::Ping(data)) => TungsteniteMessage::Ping(data), Ok(AxumMessage::Pong(data)) => TungsteniteMessage::Pong(data), Ok(AxumMessage::Close(_)) => { let _ = ttyd_tx.send(TungsteniteMessage::Close(None)).await; break; } Err(_) => break, }; if ttyd_tx.send(ttyd_msg).await.is_err() { break; } } }); // Forward messages from ttyd to browser let ttyd_to_client = tokio::spawn(async move { while let Some(msg) = ttyd_rx.next().await { let client_msg = match msg { Ok(TungsteniteMessage::Text(text)) => AxumMessage::Text(text.to_string().into()), Ok(TungsteniteMessage::Binary(data)) => AxumMessage::Binary(data), Ok(TungsteniteMessage::Ping(data)) => AxumMessage::Ping(data), Ok(TungsteniteMessage::Pong(data)) => AxumMessage::Pong(data), Ok(TungsteniteMessage::Close(_)) => { let _ = client_tx.send(AxumMessage::Close(None)).await; break; } Ok(TungsteniteMessage::Frame(_)) => continue, Err(_) => break, }; if client_tx.send(client_msg).await.is_err() { break; } } }); // Wait for either direction to complete tokio::select! { _ = client_to_ttyd => {} _ = ttyd_to_client => {} } } /// Proxy HTTP requests to ttyd pub async fn terminal_proxy( State(_state): State>, path: Option>, req: Request, ) -> Result { let path_str = path.map(|p| p.0).unwrap_or_default(); // Connect to ttyd Unix socket let mut unix_stream = UnixStream::connect(TTYD_SOCKET_PATH) .await .map_err(|e| AppError::ServiceUnavailable(format!("ttyd not running: {}", e)))?; // Build HTTP request to forward let method = req.method().as_str(); let query = req .uri() .query() .map(|q| format!("?{}", q)) .unwrap_or_default(); let uri_path = if path_str.is_empty() { format!("/api/terminal/{}", query) } else { format!("/api/terminal/{}{}", path_str, query) }; // Forward relevant headers let mut headers_str = String::new(); for (name, value) in req.headers() { if let Ok(v) = value.to_str() { let name_lower = name.as_str().to_lowercase(); if !matches!( name_lower.as_str(), "connection" | "keep-alive" | "transfer-encoding" | "upgrade" ) { headers_str.push_str(&format!("{}: {}\r\n", name, v)); } } } let http_request = format!( "{} {} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n{}\r\n", method, uri_path, headers_str ); // Send request unix_stream .write_all(http_request.as_bytes()) .await .map_err(|e| AppError::Internal(format!("Failed to send request: {}", e)))?; // Read response let mut response_buf = Vec::new(); unix_stream .read_to_end(&mut response_buf) .await .map_err(|e| AppError::Internal(format!("Failed to read response: {}", e)))?; // Parse HTTP response let response_str = String::from_utf8_lossy(&response_buf); let header_end = response_str .find("\r\n\r\n") .ok_or_else(|| AppError::Internal("Invalid HTTP response".to_string()))?; let headers_part = &response_str[..header_end]; let body_start = header_end + 4; // Parse status line let status_line = headers_part .lines() .next() .ok_or_else(|| AppError::Internal("Missing status line".to_string()))?; let status_code: u16 = status_line .split_whitespace() .nth(1) .and_then(|s| s.parse().ok()) .unwrap_or(200); // Build response let mut builder = Response::builder().status(StatusCode::from_u16(status_code).unwrap_or(StatusCode::OK)); // Forward response headers for line in headers_part.lines().skip(1) { if let Some((name, value)) = line.split_once(':') { let name = name.trim(); let value = value.trim(); if !matches!( name.to_lowercase().as_str(), "connection" | "keep-alive" | "transfer-encoding" ) { builder = builder.header(name, value); } } } let body = if body_start < response_buf.len() { Body::from(response_buf[body_start..].to_vec()) } else { Body::empty() }; builder .body(body) .map_err(|e| AppError::Internal(format!("Failed to build response: {}", e))) } /// Terminal index page pub async fn terminal_index( State(state): State>, req: Request, ) -> Result { terminal_proxy(State(state), None, req).await }