awesome-bookmarks/goose
3
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-17 12:56:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactoring for smaller functions

Browse Source
This commit is contained in:
Michael Neale
2025-11-06 09:42:18 +11:00
parent 1611779d3a
commit 643ea0bdb6
1 changed files with 260 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/tunnel/lapstone.rs
+260 -192
View File
@@ -77,25 +77,19 @@ struct TunnelResponse {
is_last_chunk: Option<bool>,
}
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<String, String>,
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<String, String>,
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<tokio::net::TcpStream>,
>,
) {
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<tokio::net::TcpStream>,
>,
>,
ws_tx: WebSocketSender,
port: u16,
server_secret: String,
last_activity: Arc<RwLock<Instant>>,
active_tasks: Arc<RwLock<Vec<JoinHandle<()>>>>,
) {
while let Some(msg) = read.next().await {
match msg {
Ok(Message::Text(text)) => {
*last_activity.write().await = Instant::now();
match serde_json::from_str::<TunnelMessage>(&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<RwLock<Vec<JoinHandle<()>>>>,
) {
if let Some(mut tx) = ws_tx.write().await.take() {
let _ = tx.close().await;
}
let tasks = active_tasks.write().await.drain(..).collect::<Vec<_>>();
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<RwLock<Vec<JoinHandle<()>>>> = 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::<TunnelMessage>(&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::<Vec<_>>();
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");