From 643ea0bdb698a4d72bf0eaa2b68524d39659d74f Mon Sep 17 00:00:00 2001 From: Michael Neale Date: Thu, 6 Nov 2025 09:42:18 +1100 Subject: [PATCH] refactoring for smaller functions --- crates/goose-server/src/tunnel/lapstone.rs | 452 ++++++++++++--------- 1 file changed, 260 insertions(+), 192 deletions(-) diff --git a/crates/goose-server/src/tunnel/lapstone.rs b/crates/goose-server/src/tunnel/lapstone.rs index 8fcdec454f..f60025d76c 100644 --- a/crates/goose-server/src/tunnel/lapstone.rs +++ b/crates/goose-server/src/tunnel/lapstone.rs @@ -77,25 +77,19 @@ struct TunnelResponse { is_last_chunk: Option, } -async fn handle_request( - message: TunnelMessage, - port: u16, - ws_tx: WebSocketSender, - server_secret: String, -) -> Result<()> { - let request_id = message.request_id.clone(); - info!("→ {} {} [{}]", message.method, message.path, request_id); - - let client = reqwest::Client::new(); - let url = format!("http://127.0.0.1:{}{}", port, message.path); - +fn build_request( + client: &reqwest::Client, + url: &str, + message: &TunnelMessage, + server_secret: &str, +) -> reqwest::RequestBuilder { let mut request_builder = match message.method.as_str() { - "GET" => client.get(&url), - "POST" => client.post(&url), - "PUT" => client.put(&url), - "DELETE" => client.delete(&url), - "PATCH" => client.patch(&url), - _ => client.get(&url), + "GET" => client.get(url), + "POST" => client.post(url), + "PUT" => client.put(url), + "DELETE" => client.delete(url), + "PATCH" => client.patch(url), + _ => client.get(url), }; if let Some(headers) = &message.headers { @@ -107,7 +101,7 @@ async fn handle_request( } } - request_builder = request_builder.header("X-Secret-Key", &server_secret); + request_builder = request_builder.header("X-Secret-Key", server_secret); if let Some(body) = &message.body { if message.method != "GET" && message.method != "HEAD" { @@ -115,6 +109,132 @@ async fn handle_request( } } + request_builder +} + +async fn handle_streaming_response( + response: reqwest::Response, + status: u16, + headers_map: HashMap, + request_id: String, + message_path: String, + ws_tx: WebSocketSender, +) -> Result<()> { + info!("← {} {} [{}] (streaming)", status, message_path, request_id); + + let mut stream = response.bytes_stream(); + let mut chunk_index = 0; + let mut is_first_chunk = true; + + while let Some(chunk_result) = stream.next().await { + match chunk_result { + Ok(chunk) => { + let chunk_str = String::from_utf8_lossy(&chunk).to_string(); + let tunnel_response = TunnelResponse { + request_id: request_id.clone(), + status, + headers: if is_first_chunk { + Some(headers_map.clone()) + } else { + None + }, + body: Some(chunk_str), + error: None, + chunk_index: Some(chunk_index), + total_chunks: None, + is_chunked: None, + is_streaming: Some(true), + is_first_chunk: Some(is_first_chunk), + is_last_chunk: Some(false), + }; + send_response(ws_tx.clone(), tunnel_response).await?; + chunk_index += 1; + is_first_chunk = false; + } + Err(e) => { + error!("Error reading stream chunk: {}", e); + break; + } + } + } + + let tunnel_response = TunnelResponse { + request_id: request_id.clone(), + status, + headers: None, + body: Some(String::new()), + error: None, + chunk_index: Some(chunk_index), + total_chunks: None, + is_chunked: None, + is_streaming: Some(true), + is_first_chunk: Some(false), + is_last_chunk: Some(true), + }; + send_response(ws_tx, tunnel_response).await?; + info!( + "← {} {} [{}] (complete, {} chunks)", + status, message_path, request_id, chunk_index + ); + Ok(()) +} + +async fn handle_chunked_response( + body: String, + status: u16, + headers_map: HashMap, + request_id: String, + message_path: String, + ws_tx: WebSocketSender, +) -> Result<()> { + const MAX_WS_SIZE: usize = 900_000; + let total_chunks = body.len().div_ceil(MAX_WS_SIZE); + info!( + "← {} {} [{}] ({} bytes, {} chunks)", + status, + message_path, + request_id, + body.len(), + total_chunks + ); + + for (i, chunk) in body.as_bytes().chunks(MAX_WS_SIZE).enumerate() { + let chunk_str = String::from_utf8_lossy(chunk).to_string(); + let tunnel_response = TunnelResponse { + request_id: request_id.clone(), + status, + headers: if i == 0 { + Some(headers_map.clone()) + } else { + None + }, + body: Some(chunk_str), + error: None, + chunk_index: Some(i), + total_chunks: Some(total_chunks), + is_chunked: Some(true), + is_streaming: None, + is_first_chunk: None, + is_last_chunk: None, + }; + send_response(ws_tx.clone(), tunnel_response).await?; + } + Ok(()) +} + +async fn handle_request( + message: TunnelMessage, + port: u16, + ws_tx: WebSocketSender, + server_secret: String, +) -> Result<()> { + let request_id = message.request_id.clone(); + info!("→ {} {} [{}]", message.method, message.path, request_id); + + let client = reqwest::Client::new(); + let url = format!("http://127.0.0.1:{}{}", port, message.path); + let request_builder = build_request(&client, &url, &message, &server_secret); + let response = match request_builder.send().await { Ok(resp) => resp, Err(e) => { @@ -150,100 +270,22 @@ async fn handle_request( .unwrap_or(false); if is_streaming { - info!("← {} {} [{}] (streaming)", status, message.path, request_id); - - // Stream the response chunk by chunk - let mut stream = response.bytes_stream(); - let mut chunk_index = 0; - let mut is_first_chunk = true; - - while let Some(chunk_result) = stream.next().await { - match chunk_result { - Ok(chunk) => { - let chunk_str = String::from_utf8_lossy(&chunk).to_string(); - let tunnel_response = TunnelResponse { - request_id: request_id.clone(), - status, - headers: if is_first_chunk { - Some(headers_map.clone()) - } else { - None - }, - body: Some(chunk_str), - error: None, - chunk_index: Some(chunk_index), - total_chunks: None, - is_chunked: None, - is_streaming: Some(true), - is_first_chunk: Some(is_first_chunk), - is_last_chunk: Some(false), - }; - send_response(ws_tx.clone(), tunnel_response).await?; - chunk_index += 1; - is_first_chunk = false; - } - Err(e) => { - error!("Error reading stream chunk: {}", e); - break; - } - } - } - - // Send final chunk to indicate end of stream - let tunnel_response = TunnelResponse { - request_id: request_id.clone(), + handle_streaming_response( + response, status, - headers: None, - body: Some(String::new()), - error: None, - chunk_index: Some(chunk_index), - total_chunks: None, - is_chunked: None, - is_streaming: Some(true), - is_first_chunk: Some(false), - is_last_chunk: Some(true), - }; - send_response(ws_tx, tunnel_response).await?; - info!( - "← {} {} [{}] (complete, {} chunks)", - status, message.path, request_id, chunk_index - ); + headers_map, + request_id, + message.path, + ws_tx, + ) + .await?; } else { let body = response.text().await.unwrap_or_default(); - const MAX_WS_SIZE: usize = 900_000; // 900KB + const MAX_WS_SIZE: usize = 900_000; if body.len() > MAX_WS_SIZE { - let total_chunks = body.len().div_ceil(MAX_WS_SIZE); - info!( - "← {} {} [{}] ({} bytes, {} chunks)", - status, - message.path, - request_id, - body.len(), - total_chunks - ); - - for (i, chunk) in body.as_bytes().chunks(MAX_WS_SIZE).enumerate() { - let chunk_str = String::from_utf8_lossy(chunk).to_string(); - let tunnel_response = TunnelResponse { - request_id: request_id.clone(), - status, - headers: if i == 0 { - Some(headers_map.clone()) - } else { - None - }, - body: Some(chunk_str), - error: None, - chunk_index: Some(i), - total_chunks: Some(total_chunks), - is_chunked: Some(true), - is_streaming: None, - is_first_chunk: None, - is_last_chunk: None, - }; - send_response(ws_tx.clone(), tunnel_response).await?; - } + handle_chunked_response(body, status, headers_map, request_id, message.path, ws_tx) + .await?; } else { let tunnel_response = TunnelResponse { request_id: request_id.clone(), @@ -276,6 +318,105 @@ async fn send_response(ws_tx: WebSocketSender, response: TunnelResponse) -> Resu Ok(()) } +fn configure_tcp_keepalive( + stream: &tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, +) { + let tcp_stream = stream.get_ref().get_ref(); + + #[cfg(unix)] + let socket: Socket = { + let fd = tcp_stream.as_raw_fd(); + unsafe { Socket::from_raw_fd(fd) } + }; + + #[cfg(windows)] + let socket: Socket = { + let sock = tcp_stream.as_raw_socket(); + unsafe { Socket::from_raw_socket(sock) } + }; + + let keepalive = TcpKeepalive::new() + .with_time(Duration::from_secs(30)) + .with_interval(Duration::from_secs(30)); + + if let Err(e) = socket.set_tcp_keepalive(&keepalive) { + warn!("Failed to set TCP keep-alive: {}", e); + } else { + info!("✓ TCP keep-alive enabled (30s interval)"); + } + + std::mem::forget(socket); +} + +async fn handle_websocket_messages( + mut read: futures::stream::SplitStream< + tokio_tungstenite::WebSocketStream< + tokio_tungstenite::MaybeTlsStream, + >, + >, + ws_tx: WebSocketSender, + port: u16, + server_secret: String, + last_activity: Arc>, + active_tasks: Arc>>>, +) { + while let Some(msg) = read.next().await { + match msg { + Ok(Message::Text(text)) => { + *last_activity.write().await = Instant::now(); + + match serde_json::from_str::(&text) { + Ok(tunnel_msg) => { + let ws_tx_clone = ws_tx.clone(); + let server_secret_clone = server_secret.clone(); + let task = tokio::spawn(async move { + if let Err(e) = + handle_request(tunnel_msg, port, ws_tx_clone, server_secret_clone) + .await + { + error!("Error handling request: {}", e); + } + }); + active_tasks.write().await.push(task); + } + Err(e) => { + error!("Error parsing tunnel message: {}", e); + } + } + } + Ok(Message::Close(_)) => { + info!("✗ Connection closed by server"); + break; + } + Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => { + *last_activity.write().await = Instant::now(); + } + Err(e) => { + error!("✗ WebSocket error: {}", e); + break; + } + _ => {} + } + } +} + +async fn cleanup_connection( + ws_tx: WebSocketSender, + active_tasks: Arc>>>, +) { + if let Some(mut tx) = ws_tx.write().await.take() { + let _ = tx.close().await; + } + + let tasks = active_tasks.write().await.drain(..).collect::>(); + info!("Aborting {} active request tasks", tasks.len()); + for task in tasks { + task.abort(); + } +} + async fn run_tunnel_loop( port: u16, agent_id: String, @@ -295,33 +436,7 @@ async fn run_tunnel_loop( let ws_stream = match connect_async(url.clone()).await { Ok((stream, _)) => { - // Enable TCP keep-alive to detect dead connections faster - let tcp_stream = stream.get_ref().get_ref(); - - #[cfg(unix)] - let socket: Socket = { - let fd = tcp_stream.as_raw_fd(); - unsafe { Socket::from_raw_fd(fd) } - }; - - #[cfg(windows)] - let socket: Socket = { - let sock = tcp_stream.as_raw_socket(); - unsafe { Socket::from_raw_socket(sock) } - }; - - let keepalive = TcpKeepalive::new() - .with_time(Duration::from_secs(30)) - .with_interval(Duration::from_secs(30)); - - if let Err(e) = socket.set_tcp_keepalive(&keepalive) { - warn!("Failed to set TCP keep-alive: {}", e); - } else { - info!("✓ TCP keep-alive enabled (30s interval)"); - } - - std::mem::forget(socket); - + configure_tcp_keepalive(&stream); stream } Err(e) => { @@ -336,11 +451,9 @@ async fn run_tunnel_loop( let public_url = format!("{}/tunnel/{}", worker_url, agent_id); info!("✓ Public URL: {}", public_url); - let (write, mut read) = ws_stream.split(); + let (write, read) = ws_stream.split(); let ws_tx: WebSocketSender = Arc::new(RwLock::new(Some(write))); let last_activity = Arc::new(RwLock::new(Instant::now())); - - // we need to track active request tasks so we can abort them on disconnect let active_tasks: Arc>>> = Arc::new(RwLock::new(Vec::new())); let last_activity_clone = last_activity.clone(); @@ -362,64 +475,19 @@ async fn run_tunnel_loop( _ = idle_task => { info!("✗ Idle timeout triggered, reconnecting..."); } - _ = async { - // Main message loop - while let Some(msg) = read.next().await { - match msg { - Ok(Message::Text(text)) => { - *last_activity.write().await = Instant::now(); - - match serde_json::from_str::(&text) { - Ok(tunnel_msg) => { - let ws_tx_clone = ws_tx.clone(); - let server_secret_clone = server_secret.clone(); - let task = tokio::spawn(async move { - if let Err(e) = handle_request( - tunnel_msg, - port, - ws_tx_clone, - server_secret_clone, - ) - .await - { - error!("Error handling request: {}", e); - } - }); - active_tasks.write().await.push(task); - } - Err(e) => { - error!("Error parsing tunnel message: {}", e); - } - } - } - Ok(Message::Close(_)) => { - info!("✗ Connection closed by server"); - break; - } - Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => { - *last_activity.write().await = Instant::now(); - } - Err(e) => { - error!("✗ WebSocket error: {}", e); - break; - } - _ => {} - } - } - } => { + _ = handle_websocket_messages( + read, + ws_tx.clone(), + port, + server_secret.clone(), + last_activity, + active_tasks.clone() + ) => { info!("✗ Connection ended"); } } - if let Some(mut tx) = ws_tx.write().await.take() { - let _ = tx.close().await; - } - - let tasks = active_tasks.write().await.drain(..).collect::>(); - info!("Aborting {} active request tasks", tasks.len()); - for task in tasks { - task.abort(); - } + cleanup_connection(ws_tx, active_tasks).await; if handle.read().await.is_none() { info!("Tunnel stopped, not reconnecting");