Add tunnel overload protection

apps/aether-gateway/src/tunnel/embedded/mod.rs

@@ -13,6 +13,7 @@ use aether_runtime::{
 use aether_runtime_state::{RuntimeSemaphore, RuntimeSemaphoreError, RuntimeSemaphoreSnapshot};
 use axum::extract::ws::WebSocketUpgrade;
 use axum::extract::State;
+use axum::http::HeaderMap;
 use axum::response::{IntoResponse, Json};
 use axum::routing::{get, post};
 use axum::Router;
@@ -193,7 +194,7 @@ async fn metrics(State(state): State<AppState>) -> impl IntoResponse {
 pub async fn ws_proxy(
     ws: WebSocketUpgrade,
     State(state): State<AppState>,
-    headers: axum::http::HeaderMap,
+    headers: HeaderMap,
 ) -> impl IntoResponse {
     let node_id = headers
         .get("x-node-id")
@@ -209,12 +210,7 @@ pub async fn ws_proxy(
         .trim()
         .to_string();
 
-    let max_streams: usize = headers
+    let max_streams = resolve_proxy_max_streams(&headers, state.max_streams);
-        .get("x-tunnel-max-streams")
-        .and_then(|v| v.to_str().ok())
-        .and_then(|v| v.parse().ok())
-        .unwrap_or(state.max_streams)
-        .clamp(64, 2048);
 
     if node_id.is_empty() {
         warn!("proxy connection rejected: missing X-Node-ID header");
@@ -262,3 +258,35 @@ pub async fn ws_proxy(
         })
         .into_response()
 }
+
+fn resolve_proxy_max_streams(headers: &HeaderMap, fallback: usize) -> usize {
+    headers
+        .get("x-tunnel-max-streams")
+        .and_then(|value| value.to_str().ok())
+        .and_then(|value| value.parse::<usize>().ok())
+        .unwrap_or(fallback)
+        .clamp(1, 2048)
+}
+
+#[cfg(test)]
+mod tests {
+    use axum::http::{HeaderMap, HeaderValue};
+
+    use super::resolve_proxy_max_streams;
+
+    #[test]
+    fn proxy_max_streams_honors_small_advertised_capacity() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-tunnel-max-streams", HeaderValue::from_static("8"));
+
+        assert_eq!(resolve_proxy_max_streams(&headers, 128), 8);
+    }
+
+    #[test]
+    fn proxy_max_streams_caps_unreasonably_large_capacity() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-tunnel-max-streams", HeaderValue::from_static("9999"));
+
+        assert_eq!(resolve_proxy_max_streams(&headers, 128), 2048);
+    }
+}

apps/aether-proxy/src/app.rs

@@ -23,6 +23,12 @@ use crate::{hardware, target_filter, tunnel};
 
 type TaskHandles = Arc<Mutex<Vec<JoinHandle<()>>>>;
 
+const AUTO_STREAM_LIMIT_MIN: usize = 16;
+const AUTO_STREAM_LIMIT_MAX: usize = 512;
+const AUTO_STREAM_LIMIT_PER_CPU: u64 = 64;
+const AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM: u64 = 4;
+const AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR: u64 = 40;
+
 #[derive(Debug, Clone, Copy)]
 struct TunnelPoolPolicy {
     min_connections: usize,
@@ -120,13 +126,25 @@ pub async fn run(mut config: Config, servers: Vec<ServerEntry>) -> anyhow::Resul
     // Collect hardware info (once at startup, sent during registration)
     let hw_info = hardware::collect();
 
+    if config.max_in_flight_streams.is_none() {
+        let auto = auto_max_in_flight_streams(&hw_info);
+        config.max_in_flight_streams = Some(auto);
+        info!(
+            max_in_flight_streams = auto,
+            "auto-detected max_in_flight_streams from hardware"
+        );
+    }
+
     // Auto-detect tunnel_max_streams from hardware if not explicitly set
     if config.tunnel_max_streams.is_none() {
-        let auto = (hw_info.estimated_max_concurrency / 10).clamp(64, 1024) as u32;
+        let stream_limit = config
+            .max_in_flight_streams
+            .unwrap_or_else(|| auto_max_in_flight_streams(&hw_info));
+        let auto = stream_limit.clamp(1, 1024) as u32;
         config.tunnel_max_streams = Some(auto);
         info!(
             tunnel_max_streams = auto,
-            "auto-detected tunnel_max_streams from hardware"
+            "auto-detected tunnel_max_streams from stream admission limit"
         );
     }
     let tunnel_pool_sizing = config.resolve_tunnel_pool_sizing(&hw_info)?;
@@ -447,6 +465,22 @@ fn registration_retry_policy(config: &Config) -> HttpRetryConfig {
     .normalized()
 }
 
+fn auto_max_in_flight_streams(hw_info: &crate::hardware::HardwareInfo) -> usize {
+    let by_cpu = u64::from(hw_info.cpu_cores.max(1)).saturating_mul(AUTO_STREAM_LIMIT_PER_CPU);
+    let by_memory = hw_info
+        .total_memory_mb
+        .max(AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM)
+        / AUTO_STREAM_LIMIT_MEMORY_MB_PER_STREAM;
+    let by_estimate = hw_info
+        .estimated_max_concurrency
+        .max(AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR)
+        / AUTO_STREAM_LIMIT_ESTIMATED_DIVISOR;
+    let raw = by_cpu.min(by_memory).min(by_estimate).max(1);
+    usize::try_from(raw)
+        .unwrap_or(AUTO_STREAM_LIMIT_MAX)
+        .clamp(AUTO_STREAM_LIMIT_MIN, AUTO_STREAM_LIMIT_MAX)
+}
+
 fn build_server_context(
     config: &Config,
     label: &str,
@@ -845,6 +879,32 @@ mod tests {
         assert!(!should_scale_down(200, 1, &policy));
     }
 
+    #[test]
+    fn auto_stream_limit_is_conservative_for_tiny_nodes() {
+        let hw = HardwareInfo {
+            cpu_cores: 1,
+            total_memory_mb: 183,
+            os_info: "test".to_string(),
+            fd_limit: 65_535,
+            estimated_max_concurrency: 2_000,
+        };
+
+        assert_eq!(auto_max_in_flight_streams(&hw), 45);
+    }
+
+    #[test]
+    fn auto_stream_limit_caps_large_nodes() {
+        let hw = HardwareInfo {
+            cpu_cores: 64,
+            total_memory_mb: 262_144,
+            os_info: "test".to_string(),
+            fd_limit: 1_048_576,
+            estimated_max_concurrency: 500_000,
+        };
+
+        assert_eq!(auto_max_in_flight_streams(&hw), AUTO_STREAM_LIMIT_MAX);
+    }
+
     async fn wait_for_registered_server(
         server_contexts: &Arc<Mutex<Vec<Arc<ServerContext>>>>,
     ) -> Arc<ServerContext> {

apps/aether-proxy/src/config.rs

@@ -756,7 +756,11 @@ impl Config {
         hw_info: &HardwareInfo,
     ) -> anyhow::Result<TunnelPoolSizing> {
         let per_tunnel_capacity = u64::from(self.tunnel_max_streams.unwrap_or(128).max(1));
-        let estimated = hw_info.estimated_max_concurrency.max(per_tunnel_capacity);
+        let estimated = self
+            .max_in_flight_streams
+            .and_then(|limit| u64::try_from(limit).ok())
+            .unwrap_or(hw_info.estimated_max_concurrency)
+            .max(per_tunnel_capacity);
         let cpu_soft_cap = u64::from(hw_info.cpu_cores.max(1))
             .saturating_mul(AUTO_TUNNEL_CONNECTIONS_PER_CPU_CAP)
             .clamp(
@@ -1613,6 +1617,36 @@ node_name = "proxy-test"
         assert_eq!(sizing.max_connections, 4);
     }
 
+    #[test]
+    fn auto_tunnel_pool_sizing_respects_stream_admission_limit() {
+        let config = Config::parse_from([
+            "aether-proxy",
+            "--aether-url",
+            "https://example.com",
+            "--management-token",
+            "ae_test",
+            "--node-name",
+            "proxy-test",
+            "--tunnel-max-streams",
+            "45",
+            "--max-in-flight-streams",
+            "45",
+        ]);
+        let hw = HardwareInfo {
+            cpu_cores: 1,
+            total_memory_mb: 183,
+            os_info: "test".to_string(),
+            fd_limit: 65_535,
+            estimated_max_concurrency: 2_000,
+        };
+
+        let sizing = config
+            .resolve_tunnel_pool_sizing(&hw)
+            .expect("sizing should resolve");
+        assert_eq!(sizing.initial_connections, 2);
+        assert_eq!(sizing.max_connections, 4);
+    }
+
     #[test]
     fn explicit_tunnel_connections_keep_fixed_pool_without_max_override() {
         let config = Config::parse_from([

apps/aether-proxy/src/tunnel/stream_handler.rs

@@ -838,20 +838,25 @@ async fn execute_upstream_request(
 }
 
 #[allow(clippy::too_many_arguments)]
-async fn relay_upstream_response(
+async fn relay_upstream_response<B>(
     server: &ServerContext,
     stream_id: u32,
     method: &hyper::Method,
     request_url: &url::Url,
     frame_tx: &FrameSender,
-    response: hyper::Response<hyper::body::Incoming>,
+    response: hyper::Response<B>,
     total_dns_ms: u64,
     total_elapsed: Duration,
     request_timing: upstream_client::RequestTiming,
     request_body_size: &AtomicUsize,
     redirect_count: usize,
     request_body_mode: &'static str,
-) -> Option<Duration> {
+    deadline: Instant,
+) -> Option<Duration>
+where
+    B: hyper::body::Body<Data = Bytes> + Send + Unpin + 'static,
+    B::Error: std::fmt::Display,
+{
     let status = response.status().as_u16();
     let ttfb_ms = total_elapsed.as_millis() as u64;
     let mut resp_headers: Vec<(String, String)> = Vec::with_capacity(response.headers().len() + 1);
@@ -911,7 +916,52 @@ async fn relay_upstream_response(
     }
 
     let mut stream = response.into_body().into_data_stream();
-    while let Some(chunk_result) = stream.next().await {
+    loop {
+        let Some(remaining) = remaining_timeout(deadline) else {
+            server.metrics.stream_errors.fetch_add(1, Ordering::Release);
+            let error_message = "upstream response body timeout".to_string();
+            log_stream_failure(
+                stream_log_context(
+                    server,
+                    stream_id,
+                    method,
+                    Some(request_url),
+                    redirect_count,
+                    request_body_size.load(Ordering::Relaxed),
+                ),
+                &error_message,
+                total_elapsed,
+            );
+            send_error(frame_tx, stream_id, &error_message).await;
+            return Some(total_elapsed);
+        };
+
+        let chunk_result = match tokio::time::timeout(remaining, stream.next()).await {
+            Ok(chunk_result) => chunk_result,
+            Err(_) => {
+                server.metrics.stream_errors.fetch_add(1, Ordering::Release);
+                let error_message = "upstream response body timeout".to_string();
+                log_stream_failure(
+                    stream_log_context(
+                        server,
+                        stream_id,
+                        method,
+                        Some(request_url),
+                        redirect_count,
+                        request_body_size.load(Ordering::Relaxed),
+                    ),
+                    &error_message,
+                    total_elapsed,
+                );
+                send_error(frame_tx, stream_id, &error_message).await;
+                return Some(total_elapsed);
+            }
+        };
+
+        let Some(chunk_result) = chunk_result else {
+            break;
+        };
+
         match chunk_result {
             Ok(chunk) => {
                 if chunk.len() <= MAX_CHUNK_SIZE {
@@ -1304,6 +1354,7 @@ async fn handle_stream_inner(
                         request_body_size.as_ref(),
                         redirects_followed,
                         request_body_mode,
+                        deadline,
                     )
                     .await;
                 }
@@ -1341,6 +1392,7 @@ async fn handle_stream_inner(
                             request_body_size.as_ref(),
                             redirects_followed,
                             request_body_mode,
+                            deadline,
                         )
                         .await;
                     }
@@ -1394,6 +1446,7 @@ async fn handle_stream_inner(
             request_body_size.as_ref(),
             redirects_followed,
             request_body_mode,
+            deadline,
         )
         .await;
     }
@@ -1922,6 +1975,47 @@ mod tests {
                 && value.starts_with("text/plain")));
     }
 
+    #[tokio::test]
+    async fn response_body_timeout_emits_stream_error() {
+        let state = sample_state(None, None);
+        let server = sample_server(&state);
+        let (frame_tx, sent, writer_handle) = spawn_test_writer();
+        let request_url = url::Url::parse("https://example.com/slow").expect("url");
+        let request_body_size = AtomicUsize::new(0);
+        let body = Body::from_stream(futures_util::stream::pending::<
+            Result<Bytes, std::convert::Infallible>,
+        >());
+        let response = Response::builder()
+            .status(StatusCode::OK)
+            .body(body)
+            .expect("response");
+
+        relay_upstream_response(
+            &server,
+            13,
+            &hyper::Method::GET,
+            &request_url,
+            &frame_tx,
+            response,
+            0,
+            Duration::ZERO,
+            upstream_client::RequestTiming::default(),
+            &request_body_size,
+            0,
+            "empty",
+            Instant::now(),
+        )
+        .await;
+
+        let result = collect_stream_result(frame_tx, sent, writer_handle).await;
+        assert_eq!(result.response.expect("response metadata").status, 200);
+        assert_eq!(
+            result.error.as_deref(),
+            Some("upstream response body timeout")
+        );
+        assert_eq!(server.metrics.stream_errors.load(Ordering::Acquire), 1);
+    }
+
     #[tokio::test]
     async fn follows_redirects_when_explicitly_enabled_for_replayable_post_requests() {
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0")