From 2e5ceeaf5c16e4eba0e8266379b4bd1b6cf2379d Mon Sep 17 00:00:00 2001 From: Juergen Kunz Date: Tue, 17 Mar 2026 11:15:18 +0000 Subject: [PATCH] fix(protocol,edge): optimize tunnel frame handling and zero-copy uploads in edge I/O --- changelog.md | 7 + rust/crates/remoteingress-core/src/edge.rs | 241 +++-- rust/crates/remoteingress-core/src/hub.rs | 888 +++++++----------- rust/crates/remoteingress-protocol/src/lib.rs | 149 ++- ts/00_commitinfo_data.ts | 2 +- 5 files changed, 555 insertions(+), 732 deletions(-) diff --git a/changelog.md b/changelog.md index 904f711..6cdbdcd 100644 --- a/changelog.md +++ b/changelog.md @@ -1,5 +1,12 @@ # Changelog +## 2026-03-17 - 4.8.3 - fix(protocol,edge) +optimize tunnel frame handling and zero-copy uploads in edge I/O + +- extract hub frame processing into a shared edge handler to remove duplicated tunnel logic +- add zero-copy frame header encoding and read payloads directly into framed buffers for client-to-hub uploads +- refactor TunnelIo read/write state to avoid unsafe queue access and reduce buffer churn with incremental parsing + ## 2026-03-17 - 4.8.2 - fix(rust-edge) refactor tunnel I/O to preserve TLS state and prioritize control frames diff --git a/rust/crates/remoteingress-core/src/edge.rs b/rust/crates/remoteingress-core/src/edge.rs index ceb41a6..4438ae4 100644 --- a/rust/crates/remoteingress-core/src/edge.rs +++ b/rust/crates/remoteingress-core/src/edge.rs @@ -13,6 +13,15 @@ use serde::{Deserialize, Serialize}; use remoteingress_protocol::*; +type EdgeTlsStream = tokio_rustls::client::TlsStream; + +/// Result of processing a frame (shared with hub.rs pattern). +#[allow(dead_code)] +enum EdgeFrameAction { + Continue, + Disconnect(String), +} + /// Per-stream state tracked in the edge's client_writers map. struct EdgeStreamState { /// Channel to deliver FRAME_DATA_BACK payloads to the hub_to_client task. @@ -272,6 +281,83 @@ enum EdgeLoopResult { Reconnect(String), // reason for disconnection } +/// Process a single frame received from the hub side of the tunnel. +/// Handles FRAME_DATA_BACK, FRAME_WINDOW_UPDATE_BACK, FRAME_CLOSE_BACK, FRAME_CONFIG, FRAME_PING. +async fn handle_edge_frame( + frame: Frame, + tunnel_io: &mut remoteingress_protocol::TunnelIo, + client_writers: &Arc>>, + listen_ports: &Arc>>, + event_tx: &mpsc::Sender, + tunnel_writer_tx: &mpsc::Sender>, + tunnel_data_tx: &mpsc::Sender>, + port_listeners: &mut HashMap>, + active_streams: &Arc, + next_stream_id: &Arc, + edge_id: &str, + connection_token: &CancellationToken, + bind_address: &str, +) -> EdgeFrameAction { + match frame.frame_type { + FRAME_DATA_BACK => { + let mut writers = client_writers.lock().await; + if let Some(state) = writers.get(&frame.stream_id) { + if state.back_tx.try_send(frame.payload).is_err() { + log::warn!("Stream {} back-channel full, closing", frame.stream_id); + writers.remove(&frame.stream_id); + } + } + } + FRAME_WINDOW_UPDATE_BACK => { + if let Some(increment) = decode_window_update(&frame.payload) { + if increment > 0 { + let writers = client_writers.lock().await; + if let Some(state) = writers.get(&frame.stream_id) { + let prev = state.send_window.fetch_add(increment, Ordering::Release); + if prev + increment > MAX_WINDOW_SIZE { + state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); + } + state.window_notify.notify_one(); + } + } + } + } + FRAME_CLOSE_BACK => { + let mut writers = client_writers.lock().await; + writers.remove(&frame.stream_id); + } + FRAME_CONFIG => { + if let Ok(update) = serde_json::from_slice::(&frame.payload) { + log::info!("Config update from hub: ports {:?}", update.listen_ports); + *listen_ports.write().await = update.listen_ports.clone(); + let _ = event_tx.try_send(EdgeEvent::PortsUpdated { + listen_ports: update.listen_ports.clone(), + }); + apply_port_config( + &update.listen_ports, + port_listeners, + tunnel_writer_tx, + tunnel_data_tx, + client_writers, + active_streams, + next_stream_id, + edge_id, + connection_token, + bind_address, + ); + } + } + FRAME_PING => { + // Queue PONG directly — no channel round-trip, guaranteed delivery + tunnel_io.queue_ctrl(encode_frame(0, FRAME_PONG, &[])); + } + _ => { + log::warn!("Unexpected frame type {} from hub", frame.frame_type); + } + } + EdgeFrameAction::Continue +} + async fn connect_to_hub_and_run( config: &EdgeConfig, connected: &Arc>, @@ -436,73 +522,22 @@ async fn connect_to_hub_and_run( let result = 'io_loop: loop { // Drain any buffered frames loop { - match tunnel_io.try_parse_frame() { - Some(Ok(frame)) => { - last_activity = Instant::now(); - liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); - match frame.frame_type { - FRAME_DATA_BACK => { - let mut writers = client_writers.lock().await; - if let Some(state) = writers.get(&frame.stream_id) { - if state.back_tx.try_send(frame.payload).is_err() { - log::warn!("Stream {} back-channel full, closing", frame.stream_id); - writers.remove(&frame.stream_id); - } - } - } - FRAME_WINDOW_UPDATE_BACK => { - if let Some(increment) = decode_window_update(&frame.payload) { - if increment > 0 { - let writers = client_writers.lock().await; - if let Some(state) = writers.get(&frame.stream_id) { - let prev = state.send_window.fetch_add(increment, Ordering::Release); - if prev + increment > MAX_WINDOW_SIZE { - state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); - } - state.window_notify.notify_one(); - } - } - } - } - FRAME_CLOSE_BACK => { - let mut writers = client_writers.lock().await; - writers.remove(&frame.stream_id); - } - FRAME_CONFIG => { - if let Ok(update) = serde_json::from_slice::(&frame.payload) { - log::info!("Config update from hub: ports {:?}", update.listen_ports); - *listen_ports.write().await = update.listen_ports.clone(); - let _ = event_tx.try_send(EdgeEvent::PortsUpdated { - listen_ports: update.listen_ports.clone(), - }); - apply_port_config( - &update.listen_ports, - &mut port_listeners, - &tunnel_writer_tx, - &tunnel_data_tx, - &client_writers, - active_streams, - next_stream_id, - &config.edge_id, - connection_token, - bind_address, - ); - } - } - FRAME_PING => { - // Queue PONG directly — no channel round-trip, guaranteed delivery - tunnel_io.queue_ctrl(encode_frame(0, FRAME_PONG, &[])); - } - _ => { - log::warn!("Unexpected frame type {} from hub", frame.frame_type); - } - } - } + let frame = match tunnel_io.try_parse_frame() { + Some(Ok(f)) => f, Some(Err(e)) => { log::error!("Hub frame error: {}", e); break 'io_loop EdgeLoopResult::Reconnect(format!("hub_frame_error: {}", e)); } None => break, + }; + last_activity = Instant::now(); + liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); + if let EdgeFrameAction::Disconnect(reason) = handle_edge_frame( + frame, &mut tunnel_io, &client_writers, listen_ports, event_tx, + &tunnel_writer_tx, &tunnel_data_tx, &mut port_listeners, + active_streams, next_stream_id, &config.edge_id, connection_token, bind_address, + ).await { + break 'io_loop EdgeLoopResult::Reconnect(reason); } } @@ -515,61 +550,12 @@ async fn connect_to_hub_and_run( remoteingress_protocol::TunnelEvent::Frame(frame) => { last_activity = Instant::now(); liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); - match frame.frame_type { - FRAME_DATA_BACK => { - let mut writers = client_writers.lock().await; - if let Some(state) = writers.get(&frame.stream_id) { - if state.back_tx.try_send(frame.payload).is_err() { - log::warn!("Stream {} back-channel full, closing", frame.stream_id); - writers.remove(&frame.stream_id); - } - } - } - FRAME_WINDOW_UPDATE_BACK => { - if let Some(increment) = decode_window_update(&frame.payload) { - if increment > 0 { - let writers = client_writers.lock().await; - if let Some(state) = writers.get(&frame.stream_id) { - let prev = state.send_window.fetch_add(increment, Ordering::Release); - if prev + increment > MAX_WINDOW_SIZE { - state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); - } - state.window_notify.notify_one(); - } - } - } - } - FRAME_CLOSE_BACK => { - let mut writers = client_writers.lock().await; - writers.remove(&frame.stream_id); - } - FRAME_CONFIG => { - if let Ok(update) = serde_json::from_slice::(&frame.payload) { - log::info!("Config update from hub: ports {:?}", update.listen_ports); - *listen_ports.write().await = update.listen_ports.clone(); - let _ = event_tx.try_send(EdgeEvent::PortsUpdated { - listen_ports: update.listen_ports.clone(), - }); - apply_port_config( - &update.listen_ports, - &mut port_listeners, - &tunnel_writer_tx, - &tunnel_data_tx, - &client_writers, - active_streams, - next_stream_id, - &config.edge_id, - connection_token, - bind_address, - ); - } - } - FRAME_PING => { - tunnel_io.queue_ctrl(encode_frame(0, FRAME_PONG, &[])); - } - _ => { - log::warn!("Unexpected frame type {} from hub", frame.frame_type); - } + if let EdgeFrameAction::Disconnect(reason) = handle_edge_frame( + frame, &mut tunnel_io, &client_writers, listen_ports, event_tx, + &tunnel_writer_tx, &tunnel_data_tx, &mut port_listeners, + active_streams, next_stream_id, &config.edge_id, connection_token, bind_address, + ).await { + break EdgeLoopResult::Reconnect(reason); } } remoteingress_protocol::TunnelEvent::Eof => { @@ -813,15 +799,21 @@ async fn handle_client_connection( let _ = client_write.shutdown().await; }); - // Task: client -> hub (upload direction) with per-stream flow control - let mut buf = vec![0u8; 32768]; + // Task: client -> hub (upload direction) with per-stream flow control. + // Zero-copy: read payload directly after the header, then prepend header. + let mut buf = vec![0u8; FRAME_HEADER_SIZE + 32768]; loop { - // Wait for send window to have capacity (with stall timeout) + // Wait for send window to have capacity (with stall timeout). + // Safe pattern: register notified BEFORE checking the condition + // to avoid missing a notify_one that fires between load and select. loop { + let notified = window_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); let w = send_window.load(Ordering::Acquire); if w > 0 { break; } tokio::select! { - _ = window_notify.notified() => continue, + _ = notified => continue, _ = client_token.cancelled() => break, _ = tokio::time::sleep(Duration::from_secs(120)) => { log::warn!("Stream {} upload stalled (window empty for 120s)", stream_id); @@ -844,15 +836,16 @@ async fn handle_client_connection( let adaptive_cap = remoteingress_protocol::compute_window_for_stream_count( active_streams.load(Ordering::Relaxed), ) as usize; - let max_read = w.min(buf.len()).min(adaptive_cap); + let max_read = w.min(32768).min(adaptive_cap); tokio::select! { - read_result = client_read.read(&mut buf[..max_read]) => { + read_result = client_read.read(&mut buf[FRAME_HEADER_SIZE..FRAME_HEADER_SIZE + max_read]) => { match read_result { Ok(0) => break, Ok(n) => { send_window.fetch_sub(n as u32, Ordering::Release); - let data_frame = encode_frame(stream_id, FRAME_DATA, &buf[..n]); + encode_frame_header(&mut buf, stream_id, FRAME_DATA, n); + let data_frame = buf[..FRAME_HEADER_SIZE + n].to_vec(); if tunnel_data_tx.send(data_frame).await.is_err() { log::warn!("Stream {} data channel closed, closing", stream_id); break; diff --git a/rust/crates/remoteingress-core/src/hub.rs b/rust/crates/remoteingress-core/src/hub.rs index 57028eb..8986761 100644 --- a/rust/crates/remoteingress-core/src/hub.rs +++ b/rust/crates/remoteingress-core/src/hub.rs @@ -12,6 +12,19 @@ use serde::{Deserialize, Serialize}; use remoteingress_protocol::*; +type HubTlsStream = tokio_rustls::server::TlsStream; + +/// Per-stream data channel capacity. With 4MB window and 32KB frames, +/// at most ~128 frames are in-flight. 256 provides comfortable headroom. +const PER_STREAM_DATA_CAPACITY: usize = 256; + +/// Result of processing a frame. +#[allow(dead_code)] +enum FrameAction { + Continue, + Disconnect(String), +} + /// Per-stream state tracked in the hub's stream map. struct HubStreamState { /// Channel to deliver FRAME_DATA payloads to the upstream writer task. @@ -123,7 +136,7 @@ pub struct TunnelHub { struct ConnectedEdgeInfo { connected_at: u64, peer_addr: String, - active_streams: Arc>>, + edge_stream_count: Arc, config_tx: mpsc::Sender, #[allow(dead_code)] // kept alive for Drop — cancels child tokens when edge is removed cancel_token: CancellationToken, @@ -189,11 +202,10 @@ impl TunnelHub { let mut connected = Vec::new(); for (id, info) in edges.iter() { - let streams = info.active_streams.lock().await; connected.push(ConnectedEdgeStatus { edge_id: id.clone(), connected_at: info.connected_at, - active_streams: streams.len(), + active_streams: info.edge_stream_count.load(Ordering::Relaxed) as usize, peer_addr: info.peer_addr.clone(), }); } @@ -287,6 +299,285 @@ impl Drop for TunnelHub { /// Maximum concurrent streams per edge connection. const MAX_STREAMS_PER_EDGE: usize = 1024; +/// Process a single frame received from the edge side of the tunnel. +/// Handles FRAME_OPEN, FRAME_DATA, FRAME_WINDOW_UPDATE, FRAME_CLOSE, and FRAME_PONG. +async fn handle_hub_frame( + frame: Frame, + tunnel_io: &mut remoteingress_protocol::TunnelIo, + streams: &mut HashMap, + stream_semaphore: &Arc, + edge_stream_count: &Arc, + edge_id: &str, + event_tx: &mpsc::Sender, + ctrl_tx: &mpsc::Sender>, + data_tx: &mpsc::Sender>, + target_host: &str, + edge_token: &CancellationToken, + cleanup_tx: &mpsc::Sender, +) -> FrameAction { + match frame.frame_type { + FRAME_OPEN => { + // A4: Check stream limit before processing + let permit = match stream_semaphore.clone().try_acquire_owned() { + Ok(p) => p, + Err(_) => { + log::warn!("Edge {} exceeded max streams ({}), rejecting stream {}", + edge_id, MAX_STREAMS_PER_EDGE, frame.stream_id); + let close_frame = encode_frame(frame.stream_id, FRAME_CLOSE_BACK, &[]); + tunnel_io.queue_ctrl(close_frame); + return FrameAction::Continue; + } + }; + + // Payload is PROXY v1 header line + let proxy_header = String::from_utf8_lossy(&frame.payload).to_string(); + + // Parse destination port from PROXY header + let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443); + + let stream_id = frame.stream_id; + let cleanup = cleanup_tx.clone(); + let writer_tx = ctrl_tx.clone(); // control: CLOSE_BACK, WINDOW_UPDATE_BACK + let data_writer_tx = data_tx.clone(); // data: DATA_BACK + let target = target_host.to_string(); + let stream_token = edge_token.child_token(); + + let _ = event_tx.try_send(HubEvent::StreamOpened { + edge_id: edge_id.to_string(), + stream_id, + }); + + // Create channel for data from edge to this stream + let (stream_data_tx, mut stream_data_rx) = mpsc::channel::>(PER_STREAM_DATA_CAPACITY); + // Adaptive initial window: scale with current stream count + // to keep total in-flight data within the 32MB budget. + let initial_window = compute_window_for_stream_count( + edge_stream_count.load(Ordering::Relaxed), + ); + let send_window = Arc::new(AtomicU32::new(initial_window)); + let window_notify = Arc::new(Notify::new()); + streams.insert(stream_id, HubStreamState { + data_tx: stream_data_tx, + cancel_token: stream_token.clone(), + send_window: Arc::clone(&send_window), + window_notify: Arc::clone(&window_notify), + }); + + // Spawn task: connect to SmartProxy, send PROXY header, pipe data + let stream_counter = Arc::clone(edge_stream_count); + tokio::spawn(async move { + let _permit = permit; // hold semaphore permit until stream completes + stream_counter.fetch_add(1, Ordering::Relaxed); + + let result = async { + // A2: Connect to SmartProxy with timeout + let mut upstream = tokio::time::timeout( + Duration::from_secs(10), + TcpStream::connect((target.as_str(), dest_port)), + ) + .await + .map_err(|_| -> Box { + format!("connect to SmartProxy {}:{} timed out (10s)", target, dest_port).into() + })??; + + upstream.set_nodelay(true)?; + upstream.write_all(proxy_header.as_bytes()).await?; + + let (mut up_read, mut up_write) = + upstream.into_split(); + + // Forward data from edge (via channel) to SmartProxy + // After writing to upstream, send WINDOW_UPDATE_BACK to edge + let writer_token = stream_token.clone(); + let wub_tx = writer_tx.clone(); + let stream_counter_w = Arc::clone(&stream_counter); + let writer_for_edge_data = tokio::spawn(async move { + let mut consumed_since_update: u32 = 0; + loop { + tokio::select! { + data = stream_data_rx.recv() => { + match data { + Some(data) => { + let len = data.len() as u32; + // Check cancellation alongside the write so we respond + // promptly to FRAME_CLOSE instead of blocking up to 60s. + let write_result = tokio::select! { + r = tokio::time::timeout( + Duration::from_secs(60), + up_write.write_all(&data), + ) => r, + _ = writer_token.cancelled() => break, + }; + match write_result { + Ok(Ok(())) => {} + Ok(Err(_)) => break, + Err(_) => { + log::warn!("Stream {} write to upstream timed out (60s)", stream_id); + break; + } + } + // Track consumption for adaptive flow control. + // Increment capped to adaptive window to limit per-stream in-flight data. + consumed_since_update += len; + let adaptive_window = remoteingress_protocol::compute_window_for_stream_count( + stream_counter_w.load(Ordering::Relaxed), + ); + let threshold = adaptive_window / 2; + if consumed_since_update >= threshold { + let increment = consumed_since_update.min(adaptive_window); + let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, increment); + if wub_tx.try_send(frame).is_ok() { + consumed_since_update -= increment; + } + // If try_send fails, keep accumulating — retry on next threshold + } + } + None => break, + } + } + _ = writer_token.cancelled() => break, + } + } + // Send final window update for remaining consumed bytes + if consumed_since_update > 0 { + let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, consumed_since_update); + let _ = wub_tx.try_send(frame); + } + let _ = up_write.shutdown().await; + }); + + // Forward data from SmartProxy back to edge via writer channel + // with per-stream flow control (check send_window before reading). + // Zero-copy: read payload directly after the header, then prepend header. + let mut buf = vec![0u8; FRAME_HEADER_SIZE + 32768]; + loop { + // Wait for send window to have capacity (with stall timeout). + // Safe pattern: register notified BEFORE checking the condition + // to avoid missing a notify_one that fires between load and select. + loop { + let notified = window_notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + let w = send_window.load(Ordering::Acquire); + if w > 0 { break; } + tokio::select! { + _ = notified => continue, + _ = stream_token.cancelled() => break, + _ = tokio::time::sleep(Duration::from_secs(120)) => { + log::warn!("Stream {} download stalled (window empty for 120s)", stream_id); + break; + } + } + } + if stream_token.is_cancelled() { break; } + + // Limit read size to available window. + // IMPORTANT: if window is 0 (stall timeout fired), we must NOT + // read into an empty buffer — read(&mut buf[..0]) returns Ok(0) + // which would be falsely interpreted as EOF. + let w = send_window.load(Ordering::Acquire) as usize; + if w == 0 { + log::warn!("Stream {} download: window still 0 after stall timeout, closing", stream_id); + break; + } + // Adaptive: cap read to current per-stream target window + let adaptive_cap = remoteingress_protocol::compute_window_for_stream_count( + stream_counter.load(Ordering::Relaxed), + ) as usize; + let max_read = w.min(32768).min(adaptive_cap); + + tokio::select! { + read_result = up_read.read(&mut buf[FRAME_HEADER_SIZE..FRAME_HEADER_SIZE + max_read]) => { + match read_result { + Ok(0) => break, + Ok(n) => { + send_window.fetch_sub(n as u32, Ordering::Release); + encode_frame_header(&mut buf, stream_id, FRAME_DATA_BACK, n); + let frame = buf[..FRAME_HEADER_SIZE + n].to_vec(); + if data_writer_tx.send(frame).await.is_err() { + log::warn!("Stream {} data channel closed, closing", stream_id); + break; + } + } + Err(_) => break, + } + } + _ = stream_token.cancelled() => break, + } + } + + // Send CLOSE_BACK via DATA channel (must arrive AFTER last DATA_BACK). + // Use send().await to guarantee delivery (try_send silently drops if full). + if !stream_token.is_cancelled() { + let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); + let _ = data_writer_tx.send(close_frame).await; + } + + writer_for_edge_data.abort(); + Ok::<(), Box>(()) + } + .await; + + if let Err(e) = result { + log::error!("Stream {} error: {}", stream_id, e); + // Send CLOSE_BACK via DATA channel on error (must arrive after any DATA_BACK). + // Use send().await to guarantee delivery. + if !stream_token.is_cancelled() { + let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); + let _ = data_writer_tx.send(close_frame).await; + } + } + + // Signal main loop to remove stream from the map + let _ = cleanup.send(stream_id).await; + stream_counter.fetch_sub(1, Ordering::Relaxed); + }); + } + FRAME_DATA => { + // Non-blocking dispatch to per-stream channel. + // With flow control, the sender should rarely exceed the channel capacity. + if let Some(state) = streams.get(&frame.stream_id) { + if state.data_tx.try_send(frame.payload).is_err() { + log::warn!("Stream {} data channel full, closing stream", frame.stream_id); + if let Some(state) = streams.remove(&frame.stream_id) { + state.cancel_token.cancel(); + } + } + } + } + FRAME_WINDOW_UPDATE => { + // Edge consumed data — increase our send window for this stream + if let Some(increment) = decode_window_update(&frame.payload) { + if increment > 0 { + if let Some(state) = streams.get(&frame.stream_id) { + let prev = state.send_window.fetch_add(increment, Ordering::Release); + if prev + increment > MAX_WINDOW_SIZE { + state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); + } + state.window_notify.notify_one(); + } + } + } + } + FRAME_CLOSE => { + if let Some(state) = streams.remove(&frame.stream_id) { + state.cancel_token.cancel(); + let _ = event_tx.try_send(HubEvent::StreamClosed { + edge_id: edge_id.to_string(), + stream_id: frame.stream_id, + }); + } + } + FRAME_PONG => { + log::debug!("Received PONG from edge {}", edge_id); + } + _ => { + log::warn!("Unexpected frame type {} from edge", frame.frame_type); + } + } + FrameAction::Continue +} + /// Handle a single edge connection: authenticate, then enter frame loop. async fn handle_edge_connection( stream: TcpStream, @@ -368,8 +659,11 @@ async fn handle_edge_connection( tls_stream.flush().await?; // Track this edge - let streams: Arc>> = - Arc::new(Mutex::new(HashMap::new())); + let mut streams: HashMap = HashMap::new(); + // Per-edge active stream counter for adaptive flow control + let edge_stream_count = Arc::new(AtomicU32::new(0)); + // Cleanup channel: spawned stream tasks send stream_id here when done + let (cleanup_tx, mut cleanup_rx) = mpsc::channel::(256); let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() @@ -385,16 +679,13 @@ async fn handle_edge_connection( ConnectedEdgeInfo { connected_at: now, peer_addr, - active_streams: streams.clone(), + edge_stream_count: edge_stream_count.clone(), config_tx, cancel_token: edge_token.clone(), }, ); } - // Per-edge active stream counter for adaptive flow control - let edge_stream_count = Arc::new(AtomicU32::new(0)); - // QoS dual-channel: ctrl frames have priority over data frames. // Stream handlers send through these channels -> TunnelIo drains them. let (ctrl_tx, mut ctrl_rx) = mpsc::channel::>(256); @@ -444,292 +735,36 @@ async fn handle_edge_connection( let mut disconnect_reason = "unknown".to_string(); 'hub_loop: loop { + // Drain completed stream cleanups from spawned tasks + while let Ok(stream_id) = cleanup_rx.try_recv() { + if streams.remove(&stream_id).is_some() { + let _ = event_tx.try_send(HubEvent::StreamClosed { + edge_id: edge_id.clone(), + stream_id, + }); + } + } + // Drain any buffered frames loop { - match tunnel_io.try_parse_frame() { - Some(Ok(frame)) => { - // Reset liveness on any received frame - last_activity = Instant::now(); - liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); - - match frame.frame_type { - FRAME_OPEN => { - // A4: Check stream limit before processing - let permit = match stream_semaphore.clone().try_acquire_owned() { - Ok(p) => p, - Err(_) => { - log::warn!("Edge {} exceeded max streams ({}), rejecting stream {}", - edge_id, MAX_STREAMS_PER_EDGE, frame.stream_id); - let close_frame = encode_frame(frame.stream_id, FRAME_CLOSE_BACK, &[]); - tunnel_io.queue_ctrl(close_frame); - continue; - } - }; - - // Payload is PROXY v1 header line - let proxy_header = String::from_utf8_lossy(&frame.payload).to_string(); - - // Parse destination port from PROXY header - let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443); - - let stream_id = frame.stream_id; - let edge_id_clone = edge_id.clone(); - let event_tx_clone = event_tx.clone(); - let streams_clone = streams.clone(); - let writer_tx = ctrl_tx.clone(); // control: CLOSE_BACK, WINDOW_UPDATE_BACK - let data_writer_tx = data_tx.clone(); // data: DATA_BACK - let target = target_host.clone(); - let stream_token = edge_token.child_token(); - - let _ = event_tx.try_send(HubEvent::StreamOpened { - edge_id: edge_id.clone(), - stream_id, - }); - - // Create channel for data from edge to this stream (capacity 16 is sufficient with flow control) - let (data_tx, mut data_rx) = mpsc::channel::>(1024); - // Adaptive initial window: scale with current stream count - // to keep total in-flight data within the 32MB budget. - let initial_window = compute_window_for_stream_count( - edge_stream_count.load(Ordering::Relaxed), - ); - let send_window = Arc::new(AtomicU32::new(initial_window)); - let window_notify = Arc::new(Notify::new()); - { - let mut s = streams.lock().await; - s.insert(stream_id, HubStreamState { - data_tx, - cancel_token: stream_token.clone(), - send_window: Arc::clone(&send_window), - window_notify: Arc::clone(&window_notify), - }); - } - - // Spawn task: connect to SmartProxy, send PROXY header, pipe data - let stream_counter = Arc::clone(&edge_stream_count); - tokio::spawn(async move { - let _permit = permit; // hold semaphore permit until stream completes - stream_counter.fetch_add(1, Ordering::Relaxed); - - let result = async { - // A2: Connect to SmartProxy with timeout - let mut upstream = tokio::time::timeout( - Duration::from_secs(10), - TcpStream::connect((target.as_str(), dest_port)), - ) - .await - .map_err(|_| -> Box { - format!("connect to SmartProxy {}:{} timed out (10s)", target, dest_port).into() - })??; - - upstream.set_nodelay(true)?; - upstream.write_all(proxy_header.as_bytes()).await?; - - let (mut up_read, mut up_write) = - upstream.into_split(); - - // Forward data from edge (via channel) to SmartProxy - // After writing to upstream, send WINDOW_UPDATE_BACK to edge - let writer_token = stream_token.clone(); - let wub_tx = writer_tx.clone(); - let stream_counter_w = Arc::clone(&stream_counter); - let writer_for_edge_data = tokio::spawn(async move { - let mut consumed_since_update: u32 = 0; - loop { - tokio::select! { - data = data_rx.recv() => { - match data { - Some(data) => { - let len = data.len() as u32; - // Check cancellation alongside the write so we respond - // promptly to FRAME_CLOSE instead of blocking up to 60s. - let write_result = tokio::select! { - r = tokio::time::timeout( - Duration::from_secs(60), - up_write.write_all(&data), - ) => r, - _ = writer_token.cancelled() => break, - }; - match write_result { - Ok(Ok(())) => {} - Ok(Err(_)) => break, - Err(_) => { - log::warn!("Stream {} write to upstream timed out (60s)", stream_id); - break; - } - } - // Track consumption for adaptive flow control. - // Increment capped to adaptive window to limit per-stream in-flight data. - consumed_since_update += len; - let adaptive_window = remoteingress_protocol::compute_window_for_stream_count( - stream_counter_w.load(Ordering::Relaxed), - ); - let threshold = adaptive_window / 2; - if consumed_since_update >= threshold { - let increment = consumed_since_update.min(adaptive_window); - let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, increment); - if wub_tx.try_send(frame).is_ok() { - consumed_since_update -= increment; - } - // If try_send fails, keep accumulating — retry on next threshold - } - } - None => break, - } - } - _ = writer_token.cancelled() => break, - } - } - // Send final window update for remaining consumed bytes - if consumed_since_update > 0 { - let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, consumed_since_update); - let _ = wub_tx.try_send(frame); - } - let _ = up_write.shutdown().await; - }); - - // Forward data from SmartProxy back to edge via writer channel - // with per-stream flow control (check send_window before reading) - let mut buf = vec![0u8; 32768]; - loop { - // Wait for send window to have capacity (with stall timeout) - loop { - let w = send_window.load(Ordering::Acquire); - if w > 0 { break; } - tokio::select! { - _ = window_notify.notified() => continue, - _ = stream_token.cancelled() => break, - _ = tokio::time::sleep(Duration::from_secs(120)) => { - log::warn!("Stream {} download stalled (window empty for 120s)", stream_id); - break; - } - } - } - if stream_token.is_cancelled() { break; } - - // Limit read size to available window. - // IMPORTANT: if window is 0 (stall timeout fired), we must NOT - // read into an empty buffer — read(&mut buf[..0]) returns Ok(0) - // which would be falsely interpreted as EOF. - let w = send_window.load(Ordering::Acquire) as usize; - if w == 0 { - log::warn!("Stream {} download: window still 0 after stall timeout, closing", stream_id); - break; - } - // Adaptive: cap read to current per-stream target window - let adaptive_cap = remoteingress_protocol::compute_window_for_stream_count( - stream_counter.load(Ordering::Relaxed), - ) as usize; - let max_read = w.min(buf.len()).min(adaptive_cap); - - tokio::select! { - read_result = up_read.read(&mut buf[..max_read]) => { - match read_result { - Ok(0) => break, - Ok(n) => { - send_window.fetch_sub(n as u32, Ordering::Release); - let frame = - encode_frame(stream_id, FRAME_DATA_BACK, &buf[..n]); - if data_writer_tx.send(frame).await.is_err() { - log::warn!("Stream {} data channel closed, closing", stream_id); - break; - } - } - Err(_) => break, - } - } - _ = stream_token.cancelled() => break, - } - } - - // Send CLOSE_BACK via DATA channel (must arrive AFTER last DATA_BACK). - // Use send().await to guarantee delivery (try_send silently drops if full). - if !stream_token.is_cancelled() { - let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); - let _ = data_writer_tx.send(close_frame).await; - } - - writer_for_edge_data.abort(); - Ok::<(), Box>(()) - } - .await; - - if let Err(e) = result { - log::error!("Stream {} error: {}", stream_id, e); - // Send CLOSE_BACK via DATA channel on error (must arrive after any DATA_BACK). - // Use send().await to guarantee delivery. - if !stream_token.is_cancelled() { - let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); - let _ = data_writer_tx.send(close_frame).await; - } - } - - // Clean up stream (guard against duplicate if FRAME_CLOSE already removed it) - let was_present = { - let mut s = streams_clone.lock().await; - s.remove(&stream_id).is_some() - }; - if was_present { - let _ = event_tx_clone.try_send(HubEvent::StreamClosed { - edge_id: edge_id_clone, - stream_id, - }); - } - stream_counter.fetch_sub(1, Ordering::Relaxed); - }); - } - FRAME_DATA => { - // Non-blocking dispatch to per-stream channel. - // With flow control, the sender should rarely exceed the channel capacity. - let mut s = streams.lock().await; - if let Some(state) = s.get(&frame.stream_id) { - if state.data_tx.try_send(frame.payload).is_err() { - log::warn!("Stream {} data channel full, closing stream", frame.stream_id); - if let Some(state) = s.remove(&frame.stream_id) { - state.cancel_token.cancel(); - } - } - } - } - FRAME_WINDOW_UPDATE => { - // Edge consumed data — increase our send window for this stream - if let Some(increment) = decode_window_update(&frame.payload) { - if increment > 0 { - let s = streams.lock().await; - if let Some(state) = s.get(&frame.stream_id) { - let prev = state.send_window.fetch_add(increment, Ordering::Release); - if prev + increment > MAX_WINDOW_SIZE { - state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); - } - state.window_notify.notify_one(); - } - } - } - } - FRAME_CLOSE => { - let mut s = streams.lock().await; - if let Some(state) = s.remove(&frame.stream_id) { - state.cancel_token.cancel(); - let _ = event_tx.try_send(HubEvent::StreamClosed { - edge_id: edge_id.clone(), - stream_id: frame.stream_id, - }); - } - } - FRAME_PONG => { - log::debug!("Received PONG from edge {}", edge_id); - } - _ => { - log::warn!("Unexpected frame type {} from edge", frame.frame_type); - } - } - } + let frame = match tunnel_io.try_parse_frame() { + Some(Ok(f)) => f, Some(Err(e)) => { log::error!("Edge {} frame error: {}", edge_id, e); disconnect_reason = format!("edge_frame_error: {}", e); break 'hub_loop; } None => break, + }; + last_activity = Instant::now(); + liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); + if let FrameAction::Disconnect(reason) = handle_hub_frame( + frame, &mut tunnel_io, &mut streams, &stream_semaphore, &edge_stream_count, + &edge_id, &event_tx, &ctrl_tx, &data_tx, &target_host, &edge_token, + &cleanup_tx, + ).await { + disconnect_reason = reason; + break 'hub_loop; } } @@ -744,280 +779,15 @@ async fn handle_edge_connection( match event { remoteingress_protocol::TunnelEvent::Frame(frame) => { - // Reset liveness on any received frame last_activity = Instant::now(); liveness_deadline.as_mut().reset(last_activity + liveness_timeout_dur); - - match frame.frame_type { - FRAME_OPEN => { - // A4: Check stream limit before processing - let permit = match stream_semaphore.clone().try_acquire_owned() { - Ok(p) => p, - Err(_) => { - log::warn!("Edge {} exceeded max streams ({}), rejecting stream {}", - edge_id, MAX_STREAMS_PER_EDGE, frame.stream_id); - let close_frame = encode_frame(frame.stream_id, FRAME_CLOSE_BACK, &[]); - tunnel_io.queue_ctrl(close_frame); - continue; - } - }; - - // Payload is PROXY v1 header line - let proxy_header = String::from_utf8_lossy(&frame.payload).to_string(); - - // Parse destination port from PROXY header - let dest_port = parse_dest_port_from_proxy(&proxy_header).unwrap_or(443); - - let stream_id = frame.stream_id; - let edge_id_clone = edge_id.clone(); - let event_tx_clone = event_tx.clone(); - let streams_clone = streams.clone(); - let writer_tx = ctrl_tx.clone(); // control: CLOSE_BACK, WINDOW_UPDATE_BACK - let data_writer_tx = data_tx.clone(); // data: DATA_BACK - let target = target_host.clone(); - let stream_token = edge_token.child_token(); - - let _ = event_tx.try_send(HubEvent::StreamOpened { - edge_id: edge_id.clone(), - stream_id, - }); - - // Create channel for data from edge to this stream (capacity 16 is sufficient with flow control) - let (data_tx, mut data_rx) = mpsc::channel::>(256); - // Adaptive initial window: scale with current stream count - // to keep total in-flight data within the 32MB budget. - let initial_window = compute_window_for_stream_count( - edge_stream_count.load(Ordering::Relaxed), - ); - let send_window = Arc::new(AtomicU32::new(initial_window)); - let window_notify = Arc::new(Notify::new()); - { - let mut s = streams.lock().await; - s.insert(stream_id, HubStreamState { - data_tx, - cancel_token: stream_token.clone(), - send_window: Arc::clone(&send_window), - window_notify: Arc::clone(&window_notify), - }); - } - - // Spawn task: connect to SmartProxy, send PROXY header, pipe data - let stream_counter = Arc::clone(&edge_stream_count); - tokio::spawn(async move { - let _permit = permit; // hold semaphore permit until stream completes - stream_counter.fetch_add(1, Ordering::Relaxed); - - let result = async { - // A2: Connect to SmartProxy with timeout - let mut upstream = tokio::time::timeout( - Duration::from_secs(10), - TcpStream::connect((target.as_str(), dest_port)), - ) - .await - .map_err(|_| -> Box { - format!("connect to SmartProxy {}:{} timed out (10s)", target, dest_port).into() - })??; - - upstream.set_nodelay(true)?; - upstream.write_all(proxy_header.as_bytes()).await?; - - let (mut up_read, mut up_write) = - upstream.into_split(); - - // Forward data from edge (via channel) to SmartProxy - // After writing to upstream, send WINDOW_UPDATE_BACK to edge - let writer_token = stream_token.clone(); - let wub_tx = writer_tx.clone(); - let stream_counter_w = Arc::clone(&stream_counter); - let writer_for_edge_data = tokio::spawn(async move { - let mut consumed_since_update: u32 = 0; - loop { - tokio::select! { - data = data_rx.recv() => { - match data { - Some(data) => { - let len = data.len() as u32; - // Check cancellation alongside the write so we respond - // promptly to FRAME_CLOSE instead of blocking up to 60s. - let write_result = tokio::select! { - r = tokio::time::timeout( - Duration::from_secs(60), - up_write.write_all(&data), - ) => r, - _ = writer_token.cancelled() => break, - }; - match write_result { - Ok(Ok(())) => {} - Ok(Err(_)) => break, - Err(_) => { - log::warn!("Stream {} write to upstream timed out (60s)", stream_id); - break; - } - } - // Track consumption for adaptive flow control. - // Increment capped to adaptive window to limit per-stream in-flight data. - consumed_since_update += len; - let adaptive_window = remoteingress_protocol::compute_window_for_stream_count( - stream_counter_w.load(Ordering::Relaxed), - ); - let threshold = adaptive_window / 2; - if consumed_since_update >= threshold { - let increment = consumed_since_update.min(adaptive_window); - let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, increment); - if wub_tx.try_send(frame).is_ok() { - consumed_since_update -= increment; - } - // If try_send fails, keep accumulating — retry on next threshold - } - } - None => break, - } - } - _ = writer_token.cancelled() => break, - } - } - // Send final window update for remaining consumed bytes - if consumed_since_update > 0 { - let frame = encode_window_update(stream_id, FRAME_WINDOW_UPDATE_BACK, consumed_since_update); - let _ = wub_tx.try_send(frame); - } - let _ = up_write.shutdown().await; - }); - - // Forward data from SmartProxy back to edge via writer channel - // with per-stream flow control (check send_window before reading) - let mut buf = vec![0u8; 32768]; - loop { - // Wait for send window to have capacity (with stall timeout) - loop { - let w = send_window.load(Ordering::Acquire); - if w > 0 { break; } - tokio::select! { - _ = window_notify.notified() => continue, - _ = stream_token.cancelled() => break, - _ = tokio::time::sleep(Duration::from_secs(120)) => { - log::warn!("Stream {} download stalled (window empty for 120s)", stream_id); - break; - } - } - } - if stream_token.is_cancelled() { break; } - - // Limit read size to available window. - // IMPORTANT: if window is 0 (stall timeout fired), we must NOT - // read into an empty buffer — read(&mut buf[..0]) returns Ok(0) - // which would be falsely interpreted as EOF. - let w = send_window.load(Ordering::Acquire) as usize; - if w == 0 { - log::warn!("Stream {} download: window still 0 after stall timeout, closing", stream_id); - break; - } - // Adaptive: cap read to current per-stream target window - let adaptive_cap = remoteingress_protocol::compute_window_for_stream_count( - stream_counter.load(Ordering::Relaxed), - ) as usize; - let max_read = w.min(buf.len()).min(adaptive_cap); - - tokio::select! { - read_result = up_read.read(&mut buf[..max_read]) => { - match read_result { - Ok(0) => break, - Ok(n) => { - send_window.fetch_sub(n as u32, Ordering::Release); - let frame = - encode_frame(stream_id, FRAME_DATA_BACK, &buf[..n]); - if data_writer_tx.send(frame).await.is_err() { - log::warn!("Stream {} data channel closed, closing", stream_id); - break; - } - } - Err(_) => break, - } - } - _ = stream_token.cancelled() => break, - } - } - - // Send CLOSE_BACK via DATA channel (must arrive AFTER last DATA_BACK). - // Use send().await to guarantee delivery (try_send silently drops if full). - if !stream_token.is_cancelled() { - let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); - let _ = data_writer_tx.send(close_frame).await; - } - - writer_for_edge_data.abort(); - Ok::<(), Box>(()) - } - .await; - - if let Err(e) = result { - log::error!("Stream {} error: {}", stream_id, e); - // Send CLOSE_BACK via DATA channel on error (must arrive after any DATA_BACK). - // Use send().await to guarantee delivery. - if !stream_token.is_cancelled() { - let close_frame = encode_frame(stream_id, FRAME_CLOSE_BACK, &[]); - let _ = data_writer_tx.send(close_frame).await; - } - } - - // Clean up stream (guard against duplicate if FRAME_CLOSE already removed it) - let was_present = { - let mut s = streams_clone.lock().await; - s.remove(&stream_id).is_some() - }; - if was_present { - let _ = event_tx_clone.try_send(HubEvent::StreamClosed { - edge_id: edge_id_clone, - stream_id, - }); - } - stream_counter.fetch_sub(1, Ordering::Relaxed); - }); - } - FRAME_DATA => { - // Non-blocking dispatch to per-stream channel. - // With flow control, the sender should rarely exceed the channel capacity. - let mut s = streams.lock().await; - if let Some(state) = s.get(&frame.stream_id) { - if state.data_tx.try_send(frame.payload).is_err() { - log::warn!("Stream {} data channel full, closing stream", frame.stream_id); - if let Some(state) = s.remove(&frame.stream_id) { - state.cancel_token.cancel(); - } - } - } - } - FRAME_WINDOW_UPDATE => { - // Edge consumed data — increase our send window for this stream - if let Some(increment) = decode_window_update(&frame.payload) { - if increment > 0 { - let s = streams.lock().await; - if let Some(state) = s.get(&frame.stream_id) { - let prev = state.send_window.fetch_add(increment, Ordering::Release); - if prev + increment > MAX_WINDOW_SIZE { - state.send_window.store(MAX_WINDOW_SIZE, Ordering::Release); - } - state.window_notify.notify_one(); - } - } - } - } - FRAME_CLOSE => { - let mut s = streams.lock().await; - if let Some(state) = s.remove(&frame.stream_id) { - state.cancel_token.cancel(); - let _ = event_tx.try_send(HubEvent::StreamClosed { - edge_id: edge_id.clone(), - stream_id: frame.stream_id, - }); - } - } - FRAME_PONG => { - log::debug!("Received PONG from edge {}", edge_id); - } - _ => { - log::warn!("Unexpected frame type {} from edge", frame.frame_type); - } + if let FrameAction::Disconnect(reason) = handle_hub_frame( + frame, &mut tunnel_io, &mut streams, &stream_semaphore, &edge_stream_count, + &edge_id, &event_tx, &ctrl_tx, &data_tx, &target_host, &edge_token, + &cleanup_tx, + ).await { + disconnect_reason = reason; + break; } } remoteingress_protocol::TunnelEvent::Eof => { diff --git a/rust/crates/remoteingress-protocol/src/lib.rs b/rust/crates/remoteingress-protocol/src/lib.rs index c9323fe..b4d14dc 100644 --- a/rust/crates/remoteingress-protocol/src/lib.rs +++ b/rust/crates/remoteingress-protocol/src/lib.rs @@ -72,6 +72,16 @@ pub fn encode_frame(stream_id: u32, frame_type: u8, payload: &[u8]) -> Vec { buf } +/// Write a frame header into `buf[0..FRAME_HEADER_SIZE]`. +/// The caller must ensure payload is already at `buf[FRAME_HEADER_SIZE..FRAME_HEADER_SIZE + payload_len]`. +/// This enables zero-copy encoding: read directly into `buf[FRAME_HEADER_SIZE..]`, then +/// prepend the header without copying the payload. +pub fn encode_frame_header(buf: &mut [u8], stream_id: u32, frame_type: u8, payload_len: usize) { + buf[0..4].copy_from_slice(&stream_id.to_be_bytes()); + buf[4] = frame_type; + buf[5..9].copy_from_slice(&(payload_len as u32).to_be_bytes()); +} + /// Build a PROXY protocol v1 header line. /// Format: `PROXY TCP4 \r\n` pub fn build_proxy_v1_header( @@ -173,6 +183,21 @@ pub enum TunnelEvent { Cancelled, } +/// Write state extracted into a sub-struct so the borrow checker can see +/// disjoint field access between `self.write` and `self.stream`. +struct WriteState { + ctrl_queue: VecDeque>, // PONG, WINDOW_UPDATE, CLOSE, OPEN — always first + data_queue: VecDeque>, // DATA, DATA_BACK — only when ctrl is empty + offset: usize, // progress within current frame being written + flush_needed: bool, +} + +impl WriteState { + fn has_work(&self) -> bool { + !self.ctrl_queue.is_empty() || !self.data_queue.is_empty() + } +} + /// Single-owner I/O engine for the tunnel TLS connection. /// /// Owns the TLS stream directly — no `tokio::io::split()`, no mutex. @@ -184,11 +209,9 @@ pub struct TunnelIo { // Read state: accumulate bytes, parse frames incrementally read_buf: Vec, read_pos: usize, - // Write state: dual priority queues - ctrl_queue: VecDeque>, // PONG, WINDOW_UPDATE, CLOSE, OPEN — always first - data_queue: VecDeque>, // DATA, DATA_BACK — only when ctrl is empty - write_offset: usize, // progress within current frame being written - flush_needed: bool, + parse_pos: usize, + // Write state: extracted sub-struct for safe disjoint borrows + write: WriteState, } impl TunnelIo { @@ -202,42 +225,52 @@ impl TunnelIo { stream, read_buf, read_pos, - ctrl_queue: VecDeque::new(), - data_queue: VecDeque::new(), - write_offset: 0, - flush_needed: false, + parse_pos: 0, + write: WriteState { + ctrl_queue: VecDeque::new(), + data_queue: VecDeque::new(), + offset: 0, + flush_needed: false, + }, } } /// Queue a high-priority control frame (PONG, WINDOW_UPDATE, CLOSE, OPEN). pub fn queue_ctrl(&mut self, frame: Vec) { - self.ctrl_queue.push_back(frame); + self.write.ctrl_queue.push_back(frame); } /// Queue a lower-priority data frame (DATA, DATA_BACK). pub fn queue_data(&mut self, frame: Vec) { - self.data_queue.push_back(frame); + self.write.data_queue.push_back(frame); } /// Try to parse a complete frame from the read buffer. + /// Uses a parse_pos cursor to avoid drain() on every frame. pub fn try_parse_frame(&mut self) -> Option> { - if self.read_pos < FRAME_HEADER_SIZE { + let available = self.read_pos - self.parse_pos; + if available < FRAME_HEADER_SIZE { return None; } + let base = self.parse_pos; let stream_id = u32::from_be_bytes([ - self.read_buf[0], self.read_buf[1], self.read_buf[2], self.read_buf[3], + self.read_buf[base], self.read_buf[base + 1], + self.read_buf[base + 2], self.read_buf[base + 3], ]); - let frame_type = self.read_buf[4]; + let frame_type = self.read_buf[base + 4]; let length = u32::from_be_bytes([ - self.read_buf[5], self.read_buf[6], self.read_buf[7], self.read_buf[8], + self.read_buf[base + 5], self.read_buf[base + 6], + self.read_buf[base + 7], self.read_buf[base + 8], ]); if length > MAX_PAYLOAD_SIZE { let header = [ - self.read_buf[0], self.read_buf[1], self.read_buf[2], self.read_buf[3], - self.read_buf[4], self.read_buf[5], self.read_buf[6], self.read_buf[7], - self.read_buf[8], + self.read_buf[base], self.read_buf[base + 1], + self.read_buf[base + 2], self.read_buf[base + 3], + self.read_buf[base + 4], self.read_buf[base + 5], + self.read_buf[base + 6], self.read_buf[base + 7], + self.read_buf[base + 8], ]; log::error!( "CORRUPT FRAME HEADER: raw={:02x?} stream_id={} type=0x{:02x} length={}", @@ -250,21 +283,23 @@ impl TunnelIo { } let total_frame_size = FRAME_HEADER_SIZE + length as usize; - if self.read_pos < total_frame_size { + if available < total_frame_size { return None; } - let payload = self.read_buf[FRAME_HEADER_SIZE..total_frame_size].to_vec(); - self.read_buf.drain(..total_frame_size); - self.read_pos -= total_frame_size; + let payload = self.read_buf[base + FRAME_HEADER_SIZE..base + total_frame_size].to_vec(); + self.parse_pos += total_frame_size; + + // Compact when parse_pos > half the data to reclaim memory + if self.parse_pos > self.read_pos / 2 && self.parse_pos > 0 { + self.read_buf.drain(..self.parse_pos); + self.read_pos -= self.parse_pos; + self.parse_pos = 0; + } Some(Ok(Frame { stream_id, frame_type, payload })) } - fn has_write_work(&self) -> bool { - !self.ctrl_queue.is_empty() || !self.data_queue.is_empty() - } - /// Poll-based I/O step. Returns Ready on events, Pending when idle. /// /// Order: write(ctrl→data) → flush → read → channels → timers @@ -279,20 +314,16 @@ impl TunnelIo { // 1. WRITE: drain ctrl queue first, then data queue. // TLS poll_write writes plaintext to session buffer (always Ready). // Batch up to 16 frames per poll cycle. + // Safe: `self.write` and `self.stream` are disjoint fields. let mut writes = 0; - while self.has_write_work() && writes < 16 { - // Determine which queue to write from and the frame data. - // We access the queues via raw pointers to avoid borrow conflicts with self.stream. - let from_ctrl = !self.ctrl_queue.is_empty(); - let frame_ptr: *const Vec = if from_ctrl { - self.ctrl_queue.front().unwrap() + while self.write.has_work() && writes < 16 { + let from_ctrl = !self.write.ctrl_queue.is_empty(); + let frame = if from_ctrl { + self.write.ctrl_queue.front().unwrap() } else { - self.data_queue.front().unwrap() + self.write.data_queue.front().unwrap() }; - // SAFETY: the frame is not modified while we hold the pointer — poll_write - // only writes to self.stream, and advance_write only runs after poll_write returns. - let frame = unsafe { &*frame_ptr }; - let remaining = &frame[self.write_offset..]; + let remaining = &frame[self.write.offset..]; match Pin::new(&mut self.stream).poll_write(cx, remaining) { Poll::Ready(Ok(0)) => { @@ -301,12 +332,12 @@ impl TunnelIo { )); } Poll::Ready(Ok(n)) => { - self.write_offset += n; - self.flush_needed = true; - if self.write_offset >= frame.len() { - if from_ctrl { self.ctrl_queue.pop_front(); } - else { self.data_queue.pop_front(); } - self.write_offset = 0; + self.write.offset += n; + self.write.flush_needed = true; + if self.write.offset >= frame.len() { + if from_ctrl { self.write.ctrl_queue.pop_front(); } + else { self.write.data_queue.pop_front(); } + self.write.offset = 0; writes += 1; } } @@ -316,9 +347,9 @@ impl TunnelIo { } // 2. FLUSH: push encrypted data from TLS session to TCP. - if self.flush_needed { + if self.write.flush_needed { match Pin::new(&mut self.stream).poll_flush(cx) { - Poll::Ready(Ok(())) => self.flush_needed = false, + Poll::Ready(Ok(())) => self.write.flush_needed = false, Poll::Ready(Err(e)) => return Poll::Ready(TunnelEvent::WriteError(e)), Poll::Pending => {} // TCP waker will notify us } @@ -329,6 +360,12 @@ impl TunnelIo { // the waker without re-registering it, causing the task to sleep until a // timer or channel wakes it (potentially 15+ seconds of lost reads). loop { + // Compact if needed to make room for reads + if self.parse_pos > 0 && self.read_buf.len() - self.read_pos < 32768 { + self.read_buf.drain(..self.parse_pos); + self.read_pos -= self.parse_pos; + self.parse_pos = 0; + } if self.read_buf.len() < self.read_pos + 32768 { self.read_buf.resize(self.read_pos + 32768, 0); } @@ -358,7 +395,7 @@ impl TunnelIo { let mut got_new = false; loop { match ctrl_rx.poll_recv(cx) { - Poll::Ready(Some(frame)) => { self.ctrl_queue.push_back(frame); got_new = true; } + Poll::Ready(Some(frame)) => { self.write.ctrl_queue.push_back(frame); got_new = true; } Poll::Ready(None) => { return Poll::Ready(TunnelEvent::WriteError( std::io::Error::new(std::io::ErrorKind::BrokenPipe, "ctrl channel closed"), @@ -369,7 +406,7 @@ impl TunnelIo { } loop { match data_rx.poll_recv(cx) { - Poll::Ready(Some(frame)) => { self.data_queue.push_back(frame); got_new = true; } + Poll::Ready(Some(frame)) => { self.write.data_queue.push_back(frame); got_new = true; } Poll::Ready(None) => { return Poll::Ready(TunnelEvent::WriteError( std::io::Error::new(std::io::ErrorKind::BrokenPipe, "data channel closed"), @@ -390,7 +427,7 @@ impl TunnelIo { // 6. SELF-WAKE: only when we have frames AND flush is done. // If flush is pending, the TCP write-readiness waker will notify us. // If we got new channel frames, wake to write them. - if got_new || (!self.flush_needed && self.has_write_work()) { + if got_new || (!self.write.flush_needed && self.write.has_work()) { cx.waker().wake_by_ref(); } @@ -406,6 +443,22 @@ impl TunnelIo { mod tests { use super::*; + #[test] + fn test_encode_frame_header() { + let payload = b"hello"; + let mut buf = vec![0u8; FRAME_HEADER_SIZE + payload.len()]; + buf[FRAME_HEADER_SIZE..].copy_from_slice(payload); + encode_frame_header(&mut buf, 42, FRAME_DATA, payload.len()); + assert_eq!(buf, encode_frame(42, FRAME_DATA, payload)); + } + + #[test] + fn test_encode_frame_header_empty_payload() { + let mut buf = vec![0u8; FRAME_HEADER_SIZE]; + encode_frame_header(&mut buf, 99, FRAME_CLOSE, 0); + assert_eq!(buf, encode_frame(99, FRAME_CLOSE, &[])); + } + #[test] fn test_encode_frame() { let data = b"hello"; diff --git a/ts/00_commitinfo_data.ts b/ts/00_commitinfo_data.ts index 9eb320c..993b859 100644 --- a/ts/00_commitinfo_data.ts +++ b/ts/00_commitinfo_data.ts @@ -3,6 +3,6 @@ */ export const commitinfo = { name: '@serve.zone/remoteingress', - version: '4.8.2', + version: '4.8.3', description: 'Edge ingress tunnel for DcRouter - accepts incoming TCP connections at network edge and tunnels them to DcRouter SmartProxy preserving client IP via PROXY protocol v1.' }