diff --git a/Cargo.lock b/Cargo.lock index 54c03e0b4..c1f29830b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4095,6 +4095,7 @@ name = "lpm-tunnel" version = "0.80.0" dependencies = [ "base64", + "bytes", "chrono", "futures-util", "hex", diff --git a/crates/lpm-cli/src/commands/tunnel.rs b/crates/lpm-cli/src/commands/tunnel.rs index e0fff6142..620385ae2 100644 --- a/crates/lpm-cli/src/commands/tunnel.rs +++ b/crates/lpm-cli/src/commands/tunnel.rs @@ -398,6 +398,9 @@ fn tunnel_ready_json( auto_ack: bool, ) -> serde_json::Value { serde_json::json!({ + "schema_version": 2, + "event": "ready", + "protocol": 4, "success": true, "tunnel_url": session.tunnel_url, "domain": session.domain, @@ -648,7 +651,7 @@ pub(crate) async fn run_start( |msg| { if json_output { let event = serde_json::json!({ - "schema_version": 1, + "schema_version": 2, "success": false, "event": "retry", "error_code": "tunnel_retry", @@ -664,11 +667,13 @@ pub(crate) async fn run_start( *usage_for_notices .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) = Some(usage.clone()); - if !initial - && !json_output - && let Some(summary) = tunnel_usage_summary(Some(usage)) - { - install_ui::warn_untrusted(&format!("Tunnel usage: {summary}")); + if !initial { + if json_output { + let event = serde_json::json!({ "schema_version": 2, "success": true, "event": "usage_notice", "quota_version": usage.quota_version, "usage": usage }); + println!("{event:#}"); + } else if let Some(summary) = tunnel_usage_summary(Some(usage)) { + install_ui::warn_untrusted(&format!("Tunnel usage: {summary}")); + } } }, ); @@ -1393,11 +1398,39 @@ pub(crate) fn tunnel_usage_summary( usage.overage_requests.unwrap_or_default() )); } - summary.push_str(if usage.overage_enabled.unwrap_or(false) { - " · overage on" + if let (Some(used), Some(included)) = (usage.busy_object_ms, usage.included_busy_object_ms) { + summary.push_str(&format!( + " · {:.1} / {:.1} busy h", + used as f64 / 3_600_000.0, + included as f64 / 3_600_000.0 + )); + } + if let (Some(used), Some(included)) = (usage.relay_messages, usage.included_relay_messages) { + summary.push_str(&format!(" · {used} / {included} relay messages")); + } + if let (Some(used), Some(included)) = (usage.transfer_bytes, usage.included_transfer_bytes) { + summary.push_str(&format!( + " · {:.1} / {:.1} GiB", + used as f64 / 1024_f64.powi(3), + included as f64 / 1024_f64.powi(3) + )); + } + if let (Some(used), Some(included)) = ( + usage.connection_attempts, + usage.included_connection_attempts, + ) { + summary.push_str(&format!(" · {used} / {included} connections")); + } + if usage.overage_enabled.unwrap_or(false) && usage.authorized_bundles.unwrap_or_default() > 0 { + summary.push_str(&format!( + " · {} / {} bundles (${:.2} cap)", + usage.billable_bundles.unwrap_or_default(), + usage.authorized_bundles.unwrap_or_default(), + usage.overage_spend_cap_cents.unwrap_or_default() as f64 / 100.0 + )); } else { - " · hard stop" - }); + summary.push_str(" · hard stop"); + } Some(summary) } @@ -2136,6 +2169,9 @@ mod tests { assert_eq!(output["local_url"], "http://[::1]:5173/app/"); insta::assert_json_snapshot!(output, @r###" { + "schema_version": 2, + "event": "ready", + "protocol": 4, "success": true, "tunnel_url": "https://orange-moon.lpm.fyi", "domain": "orange-moon.lpm.fyi", @@ -2282,6 +2318,8 @@ mod tests { #[test] fn tunnel_limit_summary_includes_account_and_relay_caps() { let limits = lpm_tunnel::TunnelLimitMetadata { + quota_version: Some(2), + public_websockets_available: Some(false), max_concurrent: Some(1), request_rate_limit_per_minute: Some(4_000), per_ip_rate_limit_per_minute: Some(600), @@ -2299,6 +2337,8 @@ mod tests { #[test] fn tunnel_limit_summary_displays_unlimited_request_rate() { let limits = lpm_tunnel::TunnelLimitMetadata { + quota_version: Some(2), + public_websockets_available: Some(false), max_concurrent: Some(3), request_rate_limit_per_minute: Some(20_000), per_ip_rate_limit_per_minute: Some(600), @@ -2324,7 +2364,34 @@ mod tests { }; assert_eq!( tunnel_usage_summary(Some(&usage)).as_deref(), - Some("100250 / 100000 requests · 250 overage · overage on") + Some("100250 / 100000 requests · 250 overage · hard stop") + ); + } + + #[test] + fn tunnel_usage_summary_shows_all_resources_connections_and_authorized_spend() { + let usage = lpm_tunnel::TunnelUsageMetadata { + accepted_requests: Some(25_000), + included_requests: Some(50_000), + busy_object_ms: Some(90_000_000), + included_busy_object_ms: Some(180_000_000), + relay_messages: Some(500_000), + included_relay_messages: Some(1_000_000), + transfer_bytes: Some(5 * 1024 * 1024 * 1024), + included_transfer_bytes: Some(10 * 1024 * 1024 * 1024), + connection_attempts: Some(10), + included_connection_attempts: Some(1_000), + overage_enabled: Some(true), + authorized_bundles: Some(2), + billable_bundles: Some(1), + overage_spend_cap_cents: Some(1_200), + ..Default::default() + }; + assert_eq!( + tunnel_usage_summary(Some(&usage)).as_deref(), + Some( + "25000 / 50000 requests · 25.0 / 50.0 busy h · 500000 / 1000000 relay messages · 5.0 / 10.0 GiB · 10 / 1000 connections · 1 / 2 bundles ($12.00 cap)" + ) ); } diff --git a/crates/lpm-inspect/src/db.rs b/crates/lpm-inspect/src/db.rs index 59b21867f..dc14afdee 100644 --- a/crates/lpm-inspect/src/db.rs +++ b/crates/lpm-inspect/src/db.rs @@ -977,13 +977,8 @@ fn insert_request_row( let headers_json = serde_json::to_string(&webhook.request_headers).unwrap_or_default(); let resp_headers_json = serde_json::to_string(&webhook.response_headers).unwrap_or_default(); - // Extract text for FTS indexing (only index UTF-8 content) - let req_body_text = std::str::from_utf8(&webhook.request_body) - .unwrap_or("") - .to_string(); - let res_body_text = std::str::from_utf8(&webhook.response_body) - .unwrap_or("") - .to_string(); + let req_body_text = std::str::from_utf8(&webhook.request_body).unwrap_or(""); + let res_body_text = std::str::from_utf8(&webhook.response_body).unwrap_or(""); // Before INSERT OR REPLACE: if this ID already exists, delete its FTS entry. // INSERT OR REPLACE deletes the old row (and its rowid) then inserts a new one diff --git a/crates/lpm-runner/src/lpm_json.rs b/crates/lpm-runner/src/lpm_json.rs index 3d7696fae..1fe0b8317 100644 --- a/crates/lpm-runner/src/lpm_json.rs +++ b/crates/lpm-runner/src/lpm_json.rs @@ -331,7 +331,7 @@ pub struct CertBlock { #[schemars(deny_unknown_fields)] pub struct TunnelConfig { /// Full tunnel domain (e.g., "acme-api.lpm.llc"). - /// Pro/Org only — free users get ephemeral random domains. + /// Public tunnels require a paid Pro or Org account. #[serde(default)] pub domain: Option, } diff --git a/crates/lpm-tunnel/Cargo.toml b/crates/lpm-tunnel/Cargo.toml index 0818e0fe4..942190c14 100644 --- a/crates/lpm-tunnel/Cargo.toml +++ b/crates/lpm-tunnel/Cargo.toml @@ -23,6 +23,7 @@ tracing = { workspace = true } thiserror = { workspace = true } rand = { workspace = true } base64 = { workspace = true } +bytes = "1" urlencoding = "2" chrono = { workspace = true } hex = { workspace = true } diff --git a/crates/lpm-tunnel/src/client.rs b/crates/lpm-tunnel/src/client.rs index 94b4752c7..5825f8690 100644 --- a/crates/lpm-tunnel/src/client.rs +++ b/crates/lpm-tunnel/src/client.rs @@ -15,14 +15,14 @@ use lpm_common::LpmError; use std::collections::HashMap; use std::future::Future; use std::pin::Pin; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::protocol::CloseFrame; use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; -/// Maximum WebSocket message size from relay (50 MB). -pub(crate) const MAX_WS_MESSAGE_SIZE: usize = 50 * 1024 * 1024; +/// Maximum WebSocket message size from relay (16 MiB). +pub(crate) const MAX_WS_MESSAGE_SIZE: usize = 16 * 1024 * 1024; /// Maximum WebSocket frame size from relay (16 MB). pub(crate) const MAX_WS_FRAME_SIZE: usize = 16 * 1024 * 1024; @@ -38,11 +38,11 @@ const PONG_TIMEOUT_SECS: u64 = 90; /// Maximum time to wait for in-flight tasks during graceful shutdown. const SHUTDOWN_TIMEOUT_SECS: u64 = 5; -const MAX_CONCURRENT_HTTP_FORWARDS: usize = 4; +const MAX_CONCURRENT_HTTP_FORWARDS: usize = 3; const HTTP_RESPONSE_MEMORY_PERMITS: usize = 64; const HTTP_RESPONSE_MEMORY_UNIT_BYTES: usize = 1024 * 1024; const HTTP_RESPONSE_ESTIMATED_OVERHEAD_MULTIPLIER: usize = 4; -const HTTP_REQUEST_MEMORY_PERMITS: usize = 256; +const HTTP_REQUEST_MEMORY_PERMITS: usize = 64; const HTTP_REQUEST_ESTIMATED_OVERHEAD_MULTIPLIER: usize = 4; const WEBSOCKET_MEMORY_PERMITS: usize = 128; const MAX_CONCURRENT_WEBSOCKETS: usize = 64; @@ -51,7 +51,7 @@ const WEBSOCKET_LOCAL_ENQUEUE_TIMEOUT_MILLIS: u64 = 100; const WEBSOCKET_SEND_TIMEOUT_SECS: u64 = 5; const WEBSOCKET_CAPTURE_HEADER_BYTES: usize = 64 * 1024; const WEBSOCKET_WRITE_BUFFER_BYTES: usize = 128 * 1024; -const WEBSOCKET_MAX_WRITE_BUFFER_BYTES: usize = 72 * 1024 * 1024; +const WEBSOCKET_MAX_WRITE_BUFFER_BYTES: usize = 2 * 1024 * 1024; const MAX_WEBSOCKET_CONNECTION_ID_BYTES: usize = 256; const MAX_WEBSOCKET_LOCAL_URL_BYTES: usize = 8 * 1024; const MAX_WEBSOCKET_CLOSE_REASON_BYTES: usize = 123; @@ -92,6 +92,15 @@ enum RelayHandshakeMessage { #[derive(serde::Serialize)] #[serde(tag = "type")] enum ClientExtensionMessage { + #[serde(rename = "transport_ack")] + TransportAck { + transport_seq: u64, + transport_ack_nonce: String, + }, + #[serde(rename = "client_ready")] + ClientReady, + #[serde(rename = "http_cancelled")] + HttpCancelled { id: String }, #[serde(rename = "ws_ready")] WebSocketReady { id: String }, #[serde(rename = "ws_reject")] @@ -107,7 +116,7 @@ enum RetryClass { #[derive(Debug)] struct TunnelConnectError { - error: LpmError, + error: Box, retry_class: RetryClass, capacity_limited: bool, } @@ -203,21 +212,151 @@ impl CapturedWebhookEvent { } struct ActiveHttpForward { - credits: tokio::sync::mpsc::Sender<()>, + credits: tokio::sync::mpsc::Sender, cancel: tokio_util::sync::CancellationToken, + next_credit: u64, } -struct CompletedHttpForward { +pub(crate) struct CompletedHttpForward { id: String, json: String, permit: tokio::sync::OwnedSemaphorePermit, memory_permit: Option>, + request_memory_permit: Arc, +} + +pub(crate) enum HttpRelayFrame { + Stream(ClientMessage), + Completed(CompletedHttpForward), +} + +#[derive(serde::Deserialize)] +struct RelayReceipt<'a> { + #[serde(rename = "type", borrow)] + kind: &'a str, + transport_seq: u64, + #[serde(borrow)] + transport_ack_nonce: &'a str, +} + +#[derive(serde::Deserialize)] +struct BorrowedRelayRequest<'a> { + #[serde(borrow)] + id: &'a str, + #[serde(borrow)] + method: &'a str, + #[serde(borrow)] + url: &'a str, + #[serde(borrow)] + headers: &'a serde_json::value::RawValue, + #[serde(borrow)] + body: &'a str, +} + +fn validate_relay_request(text: &str) -> Result<(), TunnelConnectError> { + let request: BorrowedRelayRequest<'_> = serde_json::from_str(text) + .map_err(|_| TunnelConnectError::permanent("Invalid relay request"))?; + const MAXIMUM_BODY_BYTES: usize = 10 * 1024 * 1024; + let padding = request + .body + .as_bytes() + .iter() + .rev() + .take_while(|byte| **byte == b'=') + .take(3) + .count(); + let decoded_bytes = request + .body + .len() + .div_ceil(4) + .checked_mul(3) + .and_then(|bytes| bytes.checked_sub(padding)); + if request.id.len() > 256 + || request.method.len() > 32 + || request.url.len() > 8192 + || request.headers.get().len() > 32 * 1024 + || request.body.len() > MAXIMUM_BODY_BYTES.div_ceil(3) * 4 + || padding > 2 + || decoded_bytes.is_none_or(|bytes| bytes > MAXIMUM_BODY_BYTES) + { + return Err(TunnelConnectError::permanent( + "Relay request metadata or body exceeds the limit", + )); + } + struct HeaderLimit; + impl<'de> serde::de::Visitor<'de> for HeaderLimit { + type Value = (); + fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("at most 128 headers") + } + fn visit_map>(self, mut map: M) -> Result<(), M::Error> { + let mut count = 0; + while map + .next_entry::()? + .is_some() + { + count += 1; + if count > 128 { + return Err(serde::de::Error::custom("too many headers")); + } + } + Ok(()) + } + } + let mut headers = serde_json::Deserializer::from_str(request.headers.get()); + serde::de::Deserializer::deserialize_map(&mut headers, HeaderLimit) + .map_err(|_| TunnelConnectError::permanent("Relay request headers exceed the limit"))?; + Ok(()) +} + +fn relay_receipt(text: &str, previous: u64) -> Result, TunnelConnectError> { + let receipt: RelayReceipt<'_> = serde_json::from_str(text) + .map_err(|_| TunnelConnectError::permanent("Invalid relay transport receipt"))?; + let nonce = receipt.transport_ack_nonce.as_bytes(); + if previous.checked_add(1) != Some(receipt.transport_seq) + || nonce.len() != 36 + || !nonce.iter().enumerate().all(|(index, byte)| { + if [8, 13, 18, 23].contains(&index) { + *byte == b'-' + } else { + byte.is_ascii_hexdigit() + } + }) + { + return Err(TunnelConnectError::permanent( + "Invalid relay transport receipt", + )); + } + Ok(receipt) +} + +fn encode_relay_acknowledgment(receipt: &RelayReceipt<'_>) -> Result { + serde_json::to_string(&ClientExtensionMessage::TransportAck { + transport_seq: receipt.transport_seq, + transport_ack_nonce: receipt.transport_ack_nonce.to_owned(), + }) + .map_err(|_| TunnelConnectError::permanent("Cannot encode relay receipt")) +} + +async fn acknowledge_relay( + write: &mut S, + receipt: &RelayReceipt<'_>, +) -> Result<(), TunnelConnectError> +where + S: futures_util::Sink + Unpin, +{ + send_to_relay( + write, + Message::Text(encode_relay_acknowledgment(receipt)?), + "acknowledge relay delivery", + ) + .await } impl TunnelConnectError { fn permanent(message: impl Into) -> Self { Self { - error: LpmError::Tunnel(message.into()), + error: Box::new(LpmError::Tunnel(message.into())), retry_class: RetryClass::Permanent, capacity_limited: false, } @@ -225,7 +364,7 @@ impl TunnelConnectError { fn transient(message: impl Into) -> Self { Self { - error: LpmError::Tunnel(message.into()), + error: Box::new(LpmError::Tunnel(message.into())), retry_class: RetryClass::Transient, capacity_limited: false, } @@ -233,7 +372,7 @@ impl TunnelConnectError { fn auth_rejected(message: impl Into) -> Self { Self { - error: LpmError::Tunnel(message.into()), + error: Box::new(LpmError::Tunnel(message.into())), retry_class: RetryClass::AuthRejected, capacity_limited: false, } @@ -246,7 +385,7 @@ impl TunnelConnectError { RetryClass::Transient }; Self { - error, + error: Box::new(error), retry_class, capacity_limited: false, } @@ -262,7 +401,7 @@ impl std::fmt::Display for TunnelConnectError { impl From for TunnelConnectError { fn from(error: LpmError) -> Self { Self { - error, + error: Box::new(error), retry_class: RetryClass::Transient, capacity_limited: false, } @@ -284,6 +423,9 @@ fn relay_code_retry_class(code: &str) -> Option { | "concurrent_limit" | "billing_inactive" | "session_expired" + | "unsupported_protocol" + | "tunnels_paid_required" + | "monthly_connections_exhausted" | "monthly_allowance_exhausted" => Some(RetryClass::Permanent), "auth_unavailable" | "domain_unavailable" @@ -643,25 +785,28 @@ fn activate_local_websocket( break; }; let is_close = matches!(&command, LocalWebSocketCommand::Close { .. }); - let message = match command { + let (message, memory_permit) = match command { LocalWebSocketCommand::Frame { data, is_binary, _memory_permit, } => match local_websocket_message(data, is_binary) { - Ok(message) => message, + Ok(message) => (message, Some(_memory_permit)), Err(error) => { reason = error; break; } }, - LocalWebSocketCommand::Close { code, reason } => Message::Close( - code.map(CloseCode::from) - .filter(|code| code.is_allowed()) - .map(|code| CloseFrame { - code, - reason: reason.unwrap_or_default().into(), - }), + LocalWebSocketCommand::Close { code, reason } => ( + Message::Close( + code.map(CloseCode::from) + .filter(|code| code.is_allowed()) + .map(|code| CloseFrame { + code, + reason: reason.unwrap_or_default().into(), + }), + ), + None, ), }; let send = async { @@ -676,6 +821,7 @@ fn activate_local_websocket( send, ) => result, }; + drop(memory_permit); match result { Ok(Ok(())) if is_close => { reason = "relay closed the WebSocket".to_string(); @@ -1042,11 +1188,15 @@ pub async fn connect_with_usage_fallible( { return Ok(()); } - let connection_start = std::time::Instant::now(); + let connected_at = Mutex::new(None); let attempt_token = retry_token.take(); let mark_authenticated = |session: &TunnelSession| { auth_refresh_attempted.store(false, std::sync::atomic::Ordering::Relaxed); on_connected(session)?; + *connected_at + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = + Some(tokio::time::Instant::now()); connected_once.store(true, std::sync::atomic::Ordering::Relaxed); Ok(()) }; @@ -1064,20 +1214,31 @@ pub async fn connect_with_usage_fallible( return Ok(()); } Err(mut e) => { + let connected_duration = connected_at + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .map(|start| start.elapsed()); if e.retry_class == RetryClass::AuthRejected { let Some(provider) = options.token_provider.as_ref() else { - return Err(e.error); + return Err(*e.error); }; if auth_refresh_attempted.swap(true, std::sync::atomic::Ordering::Relaxed) { - return Err(e.error); + return Err(*e.error); } - match provider.refresh_after_rejection().await { + let refreshed = tokio::select! { + biased; + _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => return Ok(()), + result = tokio::time::timeout(std::time::Duration::from_secs(20), provider.refresh_after_rejection()) => result + .map_err(|_| TunnelConnectError::transient("Tunnel credential refresh timed out")) + .and_then(|token| token.map_err(TunnelConnectError::from_token_provider)), + }; + match refreshed { Ok(token) => { retry_token = Some(token); continue; } Err(error) => { - e = TunnelConnectError::from_token_provider(error); + e = error; if e.retry_class == RetryClass::Transient { auth_refresh_attempted .store(false, std::sync::atomic::Ordering::Relaxed); @@ -1088,12 +1249,12 @@ pub async fn connect_with_usage_fallible( let recovering_capacity = e.capacity_limited && connected_once.load(std::sync::atomic::Ordering::Relaxed); if e.retry_class == RetryClass::Permanent && !recovering_capacity { - return Err(e.error); + return Err(*e.error); } // Reset retry counter if the connection was healthy (lasted > 60s). // This prevents a long-running tunnel from accumulating retries // across unrelated transient failures. - if connection_start.elapsed().as_secs() >= HEALTHY_CONNECTION_SECS { + if healthy_tunnel_attempt(connected_duration) { retry_count = 0; } @@ -1119,6 +1280,10 @@ pub async fn connect_with_usage_fallible( } } +fn healthy_tunnel_attempt(established_duration: Option) -> bool { + established_duration.is_some_and(|duration| duration.as_secs() >= HEALTHY_CONNECTION_SECS) +} + async fn wait_for_tunnel_shutdown(shutdown: Option<&tokio_util::sync::CancellationToken>) { match shutdown { Some(shutdown) => shutdown.cancelled().await, @@ -1349,11 +1514,21 @@ fn bounded_websocket_upgrade_headers( Ok(headers) } +fn take_request_body(message: &mut ServerMessage) -> Result { + let ServerMessage::HttpRequest { body, .. } = message else { + return Err(LpmError::Tunnel("expected HttpRequest message".into())); + }; + let encoded = std::mem::take(body); + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, encoded) + .map(bytes::Bytes::from) + .map_err(|error| LpmError::Tunnel(format!("failed to decode request body: {error}"))) +} + #[allow(clippy::too_many_arguments)] async fn forward_http_request( http_client: reqwest::Client, local_target: lpm_common::LocalTarget, - server_msg: ServerMessage, + mut server_msg: ServerMessage, auto_ack: bool, webhook_tx: Option>, memory_budget: Arc, @@ -1363,6 +1538,8 @@ async fn forward_http_request( ClientMessage, Option>, ) { + let decoded_body = take_request_body(&mut server_msg); + let capture_body = decoded_body.as_ref().cloned().unwrap_or_default(); let forward_start = std::time::Instant::now(); let mut was_auto_acked = false; let memory_multiplier = if webhook_tx.is_some() { @@ -1372,18 +1549,24 @@ async fn forward_http_request( }; let mut stream_end = None; let mut capture_incomplete = false; - let (mut response, memory_permit) = match proxy::forward_request_with_memory_budget( - &http_client, - &local_target, - &server_msg, - memory_budget, - HTTP_RESPONSE_MEMORY_PERMITS, - HTTP_RESPONSE_MEMORY_UNIT_BYTES, - memory_multiplier, - stream, - ) - .await - { + let forwarded = match decoded_body { + Ok(body) => { + proxy::forward_request_with_memory_budget( + &http_client, + &local_target, + &server_msg, + memory_budget, + HTTP_RESPONSE_MEMORY_PERMITS, + HTTP_RESPONSE_MEMORY_UNIT_BYTES, + memory_multiplier, + stream, + Some(body), + ) + .await + } + Err(error) => Err(error), + }; + let (mut response, memory_permit) = match forwarded { Ok(response) => { stream_end = response.stream_end; capture_incomplete = response.capture_incomplete; @@ -1411,11 +1594,10 @@ async fn forward_http_request( ref method, ref url, ref headers, - ref body, + .. } = server_msg { - let request_body = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, body) - .unwrap_or_default(); + let request_body = Vec::from(capture_body); let (response_status, response_headers, response_body) = extract_response_data(&response); let mut captured = CapturedWebhook { id: id.clone(), @@ -1460,7 +1642,7 @@ async fn forward_http_request( webhook: Arc::new(captured), persistence: receipt_tx, _response_memory_permit: memory_permit.clone(), - _request_memory_permit: request_memory_permit, + _request_memory_permit: request_memory_permit.clone(), }; let queued = tx.try_send(event).is_ok(); if let Some(receipt_rx) = receipt_rx { @@ -1839,7 +2021,7 @@ async fn try_connect_with_token( let connection_target = options.current_local_target(); // Build connect URL — non-sensitive params only (token goes in Authorization header) let mut connect_url = format!( - "{}?port={}&protocol=3", + "{}?port={}&protocol=4", options.relay_url, connection_target.port ); if let Some(domain) = options.resolved_domain() { @@ -1874,12 +2056,13 @@ async fn try_connect_with_token( // Use IntoClientRequest so tungstenite generates the required upgrade headers. let dynamic_token = if token_override.is_none() { match options.token_provider.as_ref() { - Some(provider) => Some( - provider - .current() - .await + Some(provider) => Some(tokio::select! { + biased; + _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => return Ok(()), + result = tokio::time::timeout(std::time::Duration::from_secs(20), provider.current()) => result + .map_err(|_| TunnelConnectError::transient("Tunnel credential resolution timed out"))? .map_err(TunnelConnectError::from_token_provider)?, - ), + }), None => None, } } else { @@ -1901,26 +2084,36 @@ async fn try_connect_with_token( let connector = relay_tls_connector(options, approved_pins).map_err(TunnelConnectError::permanent)?; - let (ws_stream, _) = tokio_tungstenite::connect_async_tls_with_config( + let handshake = tokio_tungstenite::connect_async_tls_with_config( request, Some(ws_config), false, Some(connector), - ) - .await - .map_err(classify_websocket_connect_error)?; + ); + let (ws_stream, _) = tokio::select! { + _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => return Ok(()), + result = tokio::time::timeout(std::time::Duration::from_secs(15), handshake) => result.map_err(|_| TunnelConnectError::transient("Tunnel connection timed out"))?.map_err(classify_websocket_connect_error)?, + }; let (mut write, mut read) = ws_stream.split(); // Wait for ServerHello (Worker sends it after validating token) - let server_hello = read - .next() - .await + let server_hello = tokio::select! { + _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => return Ok(()), + message = tokio::time::timeout(std::time::Duration::from_secs(15), read.next()) => message.map_err(|_| TunnelConnectError::transient("Tunnel hello timed out"))?, + } .ok_or_else(|| LpmError::Tunnel("relay closed connection before hello".into()))? .map_err(|e| LpmError::Tunnel(format!("failed to read server hello: {e}")))?; + let mut relay_sequence = 0; let (session, initial_usage, relay_protocol) = match server_hello { Message::Text(text) => { + if text.len() > 64 * 1024 { + return Err(TunnelConnectError::permanent( + "Tunnel hello exceeds the metadata limit", + )); + } + let receipt = relay_receipt(&text, relay_sequence)?; let msg: RelayHandshakeMessage = serde_json::from_str(&text) .map_err(|e| LpmError::Tunnel(format!("invalid server message: {e}")))?; @@ -1938,6 +2131,13 @@ async fn try_connect_with_token( usage, protocol, } => { + if protocol != 4 { + return Err(TunnelConnectError::permanent( + "Relay protocol 4 is required; update the relay or LPM CLI", + )); + } + acknowledge_relay(&mut write, &receipt).await?; + relay_sequence = receipt.transport_seq; // domain field from relay may be just the subdomain or full domain // tunnel_url is always the full URL let domain = if raw_domain.contains('.') { @@ -2012,24 +2212,12 @@ async fn try_connect_with_token( if let Some(ref usage) = initial_usage { on_usage(usage, true); } + on_connected(&session).map_err(|error| TunnelConnectError { - error, + error: Box::new(error), retry_class: RetryClass::Permanent, capacity_limited: false, })?; - if let Some(admission) = options.forwarding_admission.as_ref() { - tokio::select! { - result = admission.wait() => { - result.map_err(|error| TunnelConnectError { - error, - retry_class: RetryClass::Permanent, - capacity_limited: false, - })?; - } - _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => return Ok(()), - } - } - // Create HTTP client for local proxying. `Policy::none()` disables // redirect-following entirely — the local dev server should never // 30x our tunnel forwarder anywhere meaningful, and an attacker- @@ -2043,6 +2231,10 @@ async fn try_connect_with_token( let http_client = reqwest::Client::builder() .no_proxy() .redirect(reqwest::redirect::Policy::none()) + .no_gzip() + .no_brotli() + .no_zstd() + .no_deflate() .build() .map_err(|e| LpmError::Tunnel(format!("failed to create HTTP client: {e}")))?; @@ -2063,8 +2255,6 @@ async fn try_connect_with_token( let http_forward_slots = Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_HTTP_FORWARDS)); let http_request_memory = Arc::new(tokio::sync::Semaphore::new(HTTP_REQUEST_MEMORY_PERMITS)); let http_response_memory = Arc::new(tokio::sync::Semaphore::new(HTTP_RESPONSE_MEMORY_PERMITS)); - let (http_response_tx, mut http_response_rx) = - tokio::sync::mpsc::channel::(MAX_CONCURRENT_HTTP_FORWARDS); // Track spawned task handles for graceful shutdown. let mut task_handles = tokio::task::JoinSet::new(); @@ -2083,7 +2273,26 @@ async fn try_connect_with_token( let mut credential_renewal: Option = None; let mut http_forwards: HashMap = HashMap::new(); let (stream_tx, mut stream_rx) = - tokio::sync::mpsc::channel::(MAX_CONCURRENT_HTTP_FORWARDS); + tokio::sync::mpsc::channel::(MAX_CONCURRENT_HTTP_FORWARDS); + + let mut admission_wait = options.forwarding_admission.as_ref().map(|admission| { + Box::pin(tokio::time::timeout( + std::time::Duration::from_secs(300), + admission.wait(), + )) + }); + let mut client_ready = admission_wait.is_none(); + if client_ready { + send_to_relay( + &mut write, + Message::Text( + serde_json::to_string(&ClientExtensionMessage::ClientReady) + .map_err(|_| TunnelConnectError::permanent("Cannot encode client readiness"))?, + ), + "publish tunnel readiness", + ) + .await?; + } // Message loop let mut connection_result = Err(TunnelConnectError::transient( @@ -2091,6 +2300,12 @@ async fn try_connect_with_token( )); loop { tokio::select! { + admitted = async { match admission_wait.as_mut() { Some(wait) => wait.await, None => std::future::pending().await } } => { + admitted.map_err(|_| TunnelConnectError::permanent("Tunnel startup admission timed out"))?.map_err(|error| TunnelConnectError { error: Box::new(error), retry_class: RetryClass::Permanent, capacity_limited: false })?; + admission_wait = None; + client_ready = true; + send_to_relay(&mut write, Message::Text(serde_json::to_string(&ClientExtensionMessage::ClientReady).map_err(|_| TunnelConnectError::permanent("Cannot encode client readiness"))?), "publish tunnel readiness").await?; + }, _ = wait_for_tunnel_shutdown(options.shutdown.as_ref()) => { connection_result = Ok(()); break; @@ -2099,6 +2314,25 @@ async fn try_connect_with_token( msg = read.next() => { match msg { Some(Ok(Message::Text(text))) => { + let wire_len = text.len(); + let receipt = relay_receipt(&text, relay_sequence)?; + let inbound_memory = if receipt.kind == "http_request" { + validate_relay_request(&text)?; + let permits = request_memory_permits(text.len()).ok_or_else(|| TunnelConnectError::permanent("Relay request exceeds the memory budget"))?; + let permit = if client_ready { Arc::clone(&http_request_memory).try_acquire_many_owned(permits).ok() } else { None }; + let Some(permit) = permit else { + let request: BorrowedRelayRequest<'_> = serde_json::from_str(&text).map_err(|_| TunnelConnectError::permanent("Invalid relay request"))?; + acknowledge_relay(&mut write, &receipt).await?; + relay_sequence = receipt.transport_seq; + let json = serde_json::to_string(&proxy::service_unavailable_response(request.id)).map_err(|_| TunnelConnectError::transient("Cannot encode busy response"))?; + send_to_relay(&mut write, Message::Text(json), "reject busy local request").await?; + continue; + }; + Some(Arc::new(permit)) + } else { + if text.len() > 512 * 1024 { return Err(TunnelConnectError::permanent("Relay control frame exceeds the limit")); } + None + }; let server_msg: ServerMessage = match serde_json::from_str(&text) { Ok(m) => m, Err(e) => { @@ -2106,6 +2340,11 @@ async fn try_connect_with_token( continue; } }; + let acknowledgment = encode_relay_acknowledgment(&receipt)?; + let transport_seq = receipt.transport_seq; + drop(text); + send_to_relay(&mut write, Message::Text(acknowledgment), "acknowledge relay delivery").await?; + relay_sequence = transport_seq; match server_msg { ServerMessage::HttpRequest { ref id, ref url, .. } => { @@ -2131,24 +2370,10 @@ async fn try_connect_with_token( continue; } - let request_permits = request_memory_permits(text.len()); - let request_permits = match request_permits { - Some(permits) => permits, - None => { - let response = proxy::service_unavailable_response(id); - if let Ok(json) = serde_json::to_string(&response) - && send_to_relay(&mut write, Message::Text(json), "send message to relay").await.is_err() - { - break; - } - continue; - } - }; - let request_memory_permit = match Arc::clone(&http_request_memory) - .try_acquire_many_owned(request_permits) + let request_memory_permit = match inbound_memory { - Ok(permit) => Arc::new(permit), - Err(_) => { + Some(permit) => permit, + None => { let response = proxy::service_unavailable_response(id); if let Ok(json) = serde_json::to_string(&response) && send_to_relay(&mut write, Message::Text(json), "send message to relay").await.is_err() @@ -2189,7 +2414,7 @@ async fn try_connect_with_token( let local_target = options.current_local_target(); let auto_ack = options.auto_ack; let webhook_tx = options.webhook_tx.clone(); - let http_response_tx = http_response_tx.clone(); + let http_response_tx = stream_tx.clone(); let http_response_memory = Arc::clone(&http_response_memory); let request_id = id.clone(); let cancel = tokio_util::sync::CancellationToken::new(); @@ -2197,7 +2422,7 @@ async fn try_connect_with_token( let stream = (relay_protocol >= 3).then(|| crate::http_stream::StreamOutput { messages: stream_tx.clone(), credits: credit_rx, cancel: cancel.clone(), }); - http_forwards.insert(request_id.clone(), ActiveHttpForward { credits, cancel: cancel.clone() }); + http_forwards.insert(request_id.clone(), ActiveHttpForward { credits, cancel: cancel.clone(), next_credit: 1 }); task_handles.spawn(async move { let forward = forward_http_request( http_client, @@ -2206,12 +2431,13 @@ async fn try_connect_with_token( auto_ack, webhook_tx.clone(), http_response_memory, - Some(request_memory_permit), + Some(Arc::clone(&request_memory_permit)), stream, ); let (response, memory_permit) = tokio::select! { + biased; result = forward => result, - _ = cancel.cancelled() => (ClientMessage::HttpResponseEnd { id: request_id.clone(), failed: true }, None), + _ = cancel.cancelled() => (ClientMessage::HttpResponseEnd { id: request_id.clone(), failed: true, seq: 0 }, None), }; let json = match serde_json::to_string(&response) { Ok(json) => json, @@ -2223,22 +2449,36 @@ async fn try_connect_with_token( } }; let _ = http_response_tx - .send(CompletedHttpForward { + .send(HttpRelayFrame::Completed(CompletedHttpForward { id: request_id, json, permit, memory_permit, - }) + request_memory_permit, + })) .await; }); } - ServerMessage::HttpResponsePull { id } => { - if let Some(forward) = http_forwards.get(&id) { let _ = forward.credits.try_send(()); } + ServerMessage::HttpResponsePull { id, seq } => { + if let Some(forward) = http_forwards.get_mut(&id) { + if seq != forward.next_credit || forward.credits.try_send(seq).is_err() { forward.cancel.cancel(); } + else { forward.next_credit += 1; } + } } ServerMessage::HttpCancel { id } => { if let Some(forward) = http_forwards.get(&id) { forward.cancel.cancel(); } + else { + let json = serde_json::to_string(&ClientExtensionMessage::HttpCancelled { id }).map_err(|_| TunnelConnectError::transient("Cannot encode cancellation receipt"))?; + send_to_relay(&mut write, Message::Text(json), "confirm completed cancellation").await?; + } } ServerMessage::WebSocketUpgrade { id, url, headers } => { + if relay_protocol == 4 { + let message = ClientExtensionMessage::WebSocketReject { id, error: "Public WebSockets are unavailable".into() }; + let json = serde_json::to_string(&message).map_err(|_| TunnelConnectError::transient("Cannot encode WebSocket rejection"))?; + send_to_relay(&mut write, Message::Text(json), "reject unavailable public WebSocket").await?; + continue; + } if let Some(error) = websocket_upgrade_metadata_error(&id, &url) { tracing::warn!( "rejected WebSocket upgrade with unsafe URL: {:?}", @@ -2379,7 +2619,7 @@ async fn try_connect_with_token( // Forward frame from relay → local WebSocket if let Some(connection) = ws_connections.get(&id) { let Some(permits) = - inbound_websocket_memory_permits(text.len(), data.len()) + inbound_websocket_memory_permits(wire_len, data.len()) else { tracing::warn!( "closing WebSocket {id} after an oversized relay frame" @@ -2720,24 +2960,36 @@ async fn try_connect_with_token( } Some(message) = stream_rx.recv() => { - let id = match &message { - ClientMessage::HttpResponseStart { id, .. } | ClientMessage::HttpResponseChunk { id, .. } => id, - _ => continue, - }; - if http_forwards.get(id).is_none_or(|forward| forward.cancel.is_cancelled()) { continue; } - if let Ok(json) = serde_json::to_string(&message) - && send_to_relay(&mut write, Message::Text(json), "send response stream").await.is_err() { break; } - } - Some(response) = http_response_rx.recv() => { - let CompletedHttpForward { id, json, permit, memory_permit } = response; + match message { + HttpRelayFrame::Stream(message) => { + let id = match &message { + ClientMessage::HttpResponseStart { id, .. } | ClientMessage::HttpResponseChunk { id, .. } => id, + _ => continue, + }; + if http_forwards.get(id).is_none_or(|forward| forward.cancel.is_cancelled()) { continue; } + if let Ok(json) = serde_json::to_string(&message) + && send_to_relay(&mut write, Message::Text(json), "send response stream").await.is_err() { break; } + } + HttpRelayFrame::Completed(response) => { + let CompletedHttpForward { id, json, permit, memory_permit, request_memory_permit } = response; let active = http_forwards.remove(&id); - if active.is_some_and(|forward| forward.cancel.is_cancelled()) { continue; } + if active.is_some_and(|forward| forward.cancel.is_cancelled()) { + drop(permit); + drop(memory_permit); + drop(request_memory_permit); + let json = serde_json::to_string(&ClientExtensionMessage::HttpCancelled { id }).map_err(|_| TunnelConnectError::transient("Cannot encode cancellation receipt"))?; + send_to_relay(&mut write, Message::Text(json), "confirm local cancellation").await?; + continue; + } if let Err(error) = send_to_relay(&mut write, Message::Text(json), "send message to relay").await { tracing::warn!("failed to send HTTP response to relay: {error}"); break; } drop(permit); drop(memory_permit); + drop(request_memory_permit); + } + } } Some(result) = task_handles.join_next(), if !task_handles.is_empty() => { @@ -2801,6 +3053,14 @@ async fn try_connect_with_token( connection_result } +#[cfg(test)] +#[path = "client/mock_relay.rs"] +mod mock_relay; + +#[cfg(test)] +#[path = "client/protocol_tests.rs"] +mod protocol_tests; + #[cfg(test)] mod tests { use super::*; @@ -2922,100 +3182,13 @@ mod tests { assert!(opts.resolved_domain().is_none()); } - #[tokio::test] - async fn event_stream_delivers_before_eof_and_cancels_local_request() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - let local = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = local.local_addr().unwrap().port(); - let local_task = tokio::spawn(async move { - let (mut socket, _) = local.accept().await.unwrap(); - let mut request = [0; 4096]; - assert!(socket.read(&mut request).await.unwrap() > 0); - socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n9\r\ndata: 1\n\n\r\n").await.unwrap(); - socket.read(&mut request).await.unwrap() - }); - let relay = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = relay.local_addr().unwrap(); - let mut options = TunnelOptions::new("test-token".into(), port); - options.relay_url = format!("ws://{address}/connect"); - options.no_pin = true; - let client = - tokio::spawn(async move { try_connect(&options, &|_| Ok(()), &|_, _| {}).await }); - let (socket, _) = relay.accept().await.unwrap(); - let mut ws = tokio_tungstenite::accept_async(socket).await.unwrap(); - for message in [ - serde_json::json!({"type":"hello","subdomain":"stream.localhost","tunnel_url":"http://stream.localhost","session_id":"stream","protocol":3}), - serde_json::json!({"type":"http_request","id":"sse","method":"GET","url":"/events","headers":{},"body":""}), - ] { - ws.send(Message::Text(message.to_string())).await.unwrap(); - } - let start = tokio::time::timeout(std::time::Duration::from_secs(2), ws.next()).await; - if start.is_err() { - client.abort(); - local_task.abort(); - } - let Message::Text(start) = start - .expect("SSE headers must arrive before EOF") - .unwrap() - .unwrap() - else { - panic!("missing headers") - }; - let start: serde_json::Value = serde_json::from_str(&start).unwrap(); - assert_eq!(start["type"], "http_response_start"); - assert!( - tokio::time::timeout(std::time::Duration::from_millis(50), ws.next()) - .await - .is_err(), - "body must wait for visitor demand" - ); - ws.send(Message::Text( - serde_json::json!({"type":"http_response_pull","id":"sse"}).to_string(), - )) - .await - .unwrap(); - let Message::Text(chunk) = - tokio::time::timeout(std::time::Duration::from_secs(2), ws.next()) - .await - .unwrap() - .unwrap() - .unwrap() - else { - panic!("missing chunk") - }; - let chunk: serde_json::Value = serde_json::from_str(&chunk).unwrap(); - assert_eq!(chunk["type"], "http_response_chunk"); - assert_eq!( - base64::Engine::decode( - &base64::engine::general_purpose::STANDARD, - chunk["body"].as_str().unwrap() - ) - .unwrap(), - b"data: 1\n\n" - ); - ws.send(Message::Text( - serde_json::json!({"type":"http_cancel","id":"sse"}).to_string(), - )) - .await - .unwrap(); - assert_eq!( - tokio::time::timeout(std::time::Duration::from_secs(2), local_task) - .await - .expect("visitor cancellation must close local request") - .unwrap(), - 0 - ); - ws.close(None).await.unwrap(); - client.await.unwrap().unwrap(); - } - #[tokio::test] async fn credential_refresh_renews_the_existing_relay_connection() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let relay = tokio::spawn(async move { let (socket, _) = listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); for value in [ serde_json::json!({ "type": "hello", "subdomain": "renew.localhost", "tunnel_url": "http://renew.localhost", "session_id": "original-session", @@ -3059,7 +3232,7 @@ mod tests { let address = listener.local_addr().unwrap(); let relay = tokio::spawn(async move { let (socket, _) = listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); websocket.send(Message::Text(serde_json::json!({ "type": "hello", "subdomain": "recovery.localhost", "tunnel_url": "http://recovery.localhost", "session_id": "recovery", }).to_string())).await.unwrap(); @@ -3093,7 +3266,7 @@ mod tests { let (ping_tx, mut ping_rx) = tokio::sync::mpsc::unbounded_channel(); let relay = tokio::spawn(async move { let (socket, _) = listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); websocket.send(Message::Text(serde_json::json!({ "type": "hello", "subdomain": "heartbeat.localhost", "tunnel_url": "http://heartbeat.localhost", "session_id": "heartbeat", }).to_string())).await.unwrap(); @@ -3155,6 +3328,65 @@ mod tests { ); } + #[tokio::test] + async fn startup_admission_wait_acknowledges_relay_controls_before_readiness() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (controller, admission) = forwarding_admission_barrier(); + let relay = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(socket).await.unwrap(); + for (seq, value) in [ + ( + 1, + serde_json::json!({ "type": "hello", "protocol": 4, "subdomain": "test.localhost", "tunnel_url": "http://test.localhost", "session_id": "session" }), + ), + (2, serde_json::json!({ "type": "pong" })), + ] { + let mut value = value; + value["transport_seq"] = serde_json::json!(seq); + value["transport_ack_nonce"] = + serde_json::json!(format!("00000000-0000-4000-8000-{seq:012x}")); + ws.send(Message::Text(value.to_string())).await.unwrap(); + let response = + tokio::time::timeout(std::time::Duration::from_millis(200), ws.next()).await; + let Message::Text(text) = response.unwrap().unwrap().unwrap() else { + panic!("missing transport acknowledgement"); + }; + assert_eq!( + serde_json::from_str::(&text).unwrap(), + serde_json::json!({ + "type": "transport_ack", "transport_seq": seq, + "transport_ack_nonce": value["transport_ack_nonce"], + }) + ); + } + assert!( + tokio::time::timeout(std::time::Duration::from_millis(50), ws.next()) + .await + .is_err(), + "readiness must wait for admission" + ); + controller.open(); + let Message::Text(text) = ws.next().await.unwrap().unwrap() else { + panic!("missing readiness") + }; + assert_eq!( + serde_json::from_str::(&text).unwrap()["type"], + "client_ready" + ); + }); + let mut options = TunnelOptions::new("test-token".into(), 3000); + options.relay_url = format!("ws://{address}/connect"); + options.no_pin = true; + options.forwarding_admission = Some(admission); + let client = + tokio::spawn(async move { try_connect(&options, &|_| Ok(()), &|_, _| {}).await }); + let result = relay.await; + client.abort(); + result.unwrap(); + } + #[tokio::test] async fn rejected_runtime_publication_never_forwards_a_queued_relay_request() { let local_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -3164,7 +3396,7 @@ mod tests { let (request_sent_tx, request_sent_rx) = tokio::sync::oneshot::channel(); let relay = tokio::spawn(async move { let (socket, _) = relay_listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); websocket .send(Message::Text( serde_json::json!({ @@ -3305,7 +3537,7 @@ mod tests { let address = listener.local_addr().unwrap(); let relay = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); - let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let mut socket = mock_relay::accept(stream).await.unwrap(); let hello = serde_json::json!({ "type": "hello", "subdomain": "recovery.localhost", "tunnel_url": "http://recovery.localhost", "session_id": "before-disconnect", @@ -3329,7 +3561,7 @@ mod tests { stream.write_all(format!("HTTP/1.1 429 Too Many Requests\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",body.len()).as_bytes()).await.unwrap(); drop(stream); let (stream, _) = listener.accept().await.unwrap(); - let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let mut socket = mock_relay::accept(stream).await.unwrap(); socket.send(Message::Text(hello)).await.unwrap(); socket.close(None).await.unwrap(); }); @@ -3368,6 +3600,7 @@ mod tests { "concurrent_limit", "billing_inactive", "monthly_allowance_exhausted", + "monthly_connections_exhausted", ] { let body = format!(r#"{{"error":"denied","code":"{code}"}}"#); let error = classify_relay_rejection(429, body.as_bytes()); @@ -3577,7 +3810,7 @@ mod tests { let (second_socket, _) = relay_listener.accept().await.unwrap(); let callback_tokens = Arc::clone(&relay_tokens); - let mut websocket = tokio_tungstenite::accept_hdr_async( + let mut websocket = mock_relay::accept_with_headers( second_socket, move |request: &tokio_tungstenite::tungstenite::handshake::server::Request, response: tokio_tungstenite::tungstenite::handshake::server::Response| { @@ -3700,7 +3933,7 @@ mod tests { let (socket, _) = relay_listener.accept().await.unwrap(); let callback_tokens = Arc::clone(&relay_tokens); - let mut websocket = tokio_tungstenite::accept_hdr_async( + let mut websocket = mock_relay::accept_with_headers( socket, move |request: &tokio_tungstenite::tungstenite::handshake::server::Request, response: tokio_tungstenite::tungstenite::handshake::server::Response| { @@ -4007,13 +4240,13 @@ mod tests { #[test] fn ws_config_constants_are_reasonable() { // Verify WebSocket message size limits are set and sane. - assert_eq!(MAX_WS_MESSAGE_SIZE, 50 * 1024 * 1024); + assert_eq!(MAX_WS_MESSAGE_SIZE, 16 * 1024 * 1024); assert_eq!(MAX_WS_FRAME_SIZE, 16 * 1024 * 1024); } #[test] - fn websocket_write_buffer_accepts_the_largest_http_response_message() { - let maximum_encoded_body = (50 * 1024 * 1024_usize).div_ceil(3) * 4; + fn websocket_write_buffer_accepts_a_bounded_stream_response_chunk() { + let maximum_encoded_body = (64 * 1024_usize).div_ceil(3) * 4; assert!(WEBSOCKET_MAX_WRITE_BUFFER_BYTES >= maximum_encoded_body + 1024 * 1024); } @@ -4603,7 +4836,7 @@ mod tests { let relay_address = relay_listener.local_addr().unwrap(); let relay = tokio::spawn(async move { let (socket, _) = relay_listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); websocket .send(Message::Text( serde_json::json!({ @@ -4698,7 +4931,7 @@ mod tests { let relay_active = Arc::clone(&active_requests); let relay = tokio::spawn(async move { let (socket, _) = relay_listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut websocket = mock_relay::accept(socket).await.unwrap(); websocket .send(Message::Text( serde_json::json!({ @@ -4791,6 +5024,215 @@ mod tests { ); } + #[tokio::test] + async fn capture_backlog_rejects_one_request_and_keeps_controls_and_recovery_live() { + let local = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = local.local_addr().unwrap().port(); + drop(local); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (webhook_tx, mut captures) = tokio::sync::mpsc::channel(4); + let mut options = TunnelOptions::new("test-token".into(), port); + options.relay_url = format!("ws://{address}/connect"); + options.no_pin = true; + options.webhook_tx = Some(webhook_tx); + options.token_provider = Some(TunnelTokenProvider::new( + || Box::pin(async { Ok("test-token".into()) }), + || Box::pin(async { Ok("renewed-token".into()) }), + )); + let client = + tokio::spawn(async move { try_connect(&options, &|_| Ok(()), &|_, _| {}).await }); + let (socket, _) = listener.accept().await.unwrap(); + let mut ws = mock_relay::accept(socket).await.unwrap(); + ws.send(Message::Text(serde_json::json!({ "type":"hello", "subdomain":"test.localhost", "tunnel_url":"http://test.localhost", "session_id":"session" }).to_string())).await.unwrap(); + let body = base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + vec![0; 5 * 1024 * 1024], + ); + let result = tokio::time::timeout(std::time::Duration::from_secs(5), async { + for (id, status) in [("capture-one", 502), ("capture-two", 502), ("rejected", 503)] { + ws.send(Message::Text(serde_json::json!({ "type":"http_request", "id":id, "method":"POST", "url":"/webhook", "headers":{}, "body":body }).to_string())).await.unwrap(); + let Message::Text(response) = ws.next().await.unwrap().unwrap() else { panic!("missing HTTP response") }; + let response: serde_json::Value = serde_json::from_str(&response).unwrap(); + assert_eq!(response["id"], id); + assert_eq!(response["status"], status); + } + ws.send(Message::Text(serde_json::json!({ "type":"credential_refresh" }).to_string())).await.unwrap(); + let Message::Text(renewal) = ws.next().await.unwrap().unwrap() else { panic!("missing credential control") }; + let renewal: serde_json::Value = serde_json::from_str(&renewal).unwrap(); + assert_eq!(renewal["token"], "renewed-token"); + drop(captures.recv().await.unwrap()); + drop(captures.recv().await.unwrap()); + assert!(captures.try_recv().is_err(), "rejected requests do not allocate a capture"); + ws.send(Message::Text(serde_json::json!({ "type":"http_request", "id":"recovered", "method":"POST", "url":"/webhook", "headers":{}, "body":body }).to_string())).await.unwrap(); + let Message::Text(response) = ws.next().await.unwrap().unwrap() else { panic!("missing recovered response") }; + let response: serde_json::Value = serde_json::from_str(&response).unwrap(); + assert_eq!(response["id"], "recovered"); + assert_eq!(response["status"], 502); + }).await; + client.abort(); + result.expect("capture saturation must not terminate the relay connection"); + } + + #[tokio::test] + async fn already_completed_request_acknowledges_crossing_cancellation() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let relay = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + let mut ws = mock_relay::accept(socket).await.unwrap(); + ws.send(Message::Text(serde_json::json!({ "type": "hello", "subdomain": "test.localhost", "tunnel_url": "http://test.localhost", "session_id": "session" }).to_string())).await.unwrap(); + ws.send(Message::Text( + serde_json::json!({ "type": "http_cancel", "id": "already-completed" }).to_string(), + )) + .await + .unwrap(); + let response = + tokio::time::timeout(std::time::Duration::from_millis(100), ws.next()).await; + assert!( + matches!(response, Ok(Some(Ok(Message::Text(text)))) if serde_json::from_str::(&text).unwrap()["type"] == "http_cancelled") + ); + }); + let mut options = TunnelOptions::new("test-token".into(), 3000); + options.relay_url = format!("ws://{address}/connect"); + options.no_pin = true; + let client = + tokio::spawn(async move { try_connect(&options, &|_| Ok(()), &|_, _| {}).await }); + let result = relay.await; + client.abort(); + result.unwrap(); + } + + #[tokio::test] + async fn too_many_relay_headers_are_rejected_before_owned_parsing() { + let headers: HashMap = (0..129) + .map(|i| (format!("x-{i}"), "value".into())) + .collect(); + let request = serde_json::json!({ "id": "bounded", "method": "GET", "url": "/", "headers": headers, "body": "" }); + assert!(validate_relay_request(&request.to_string()).is_err()); + } + + #[tokio::test] + async fn bodyless_responses_keep_complete_captures_when_relay_cancellation_crosses_headers() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + for (method, status) in [("HEAD", 200), ("GET", 204), ("GET", 205), ("GET", 304)] { + let origin = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = origin.local_addr().unwrap().port(); + let local = tokio::spawn(async move { + let (mut socket, _) = origin.accept().await.unwrap(); + let mut request = [0u8; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket + .write_all( + format!( + "HTTP/1.1 {status} OK\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + }); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let (webhook_tx, mut captures) = tokio::sync::mpsc::channel(1); + let mut options = TunnelOptions::new("test-token".into(), port); + options.relay_url = format!("ws://{address}/connect"); + options.no_pin = true; + options.webhook_tx = Some(webhook_tx); + let client = + tokio::spawn(async move { try_connect(&options, &|_| Ok(()), &|_, _| {}).await }); + let (socket, _) = listener.accept().await.unwrap(); + let mut ws = tokio_tungstenite::accept_async(socket).await.unwrap(); + let receipt = |mut value: serde_json::Value, seq| { + value["transport_seq"] = serde_json::json!(seq); + value["transport_ack_nonce"] = + serde_json::json!(format!("00000000-0000-4000-8000-{seq:012}")); + Message::Text(value.to_string()) + }; + ws.send(receipt(serde_json::json!({ "type": "hello", "protocol": 4, "subdomain": "test.localhost", "tunnel_url": "http://test.localhost", "session_id": "session" }), 1)).await.unwrap(); + ws.send(receipt(serde_json::json!({ "type": "http_request", "id": "bodyless", "method": method, "url": "/", "headers": {}, "body": "" }), 2)).await.unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while let Some(Ok(Message::Text(text))) = ws.next().await { + let value: serde_json::Value = serde_json::from_str(&text).unwrap(); + if value["type"] == "http_response_start" { + assert_eq!(value["status"], status); + return; + } + } + panic!("missing response headers"); + }) + .await + .unwrap(); + ws.send(receipt( + serde_json::json!({ "type": "http_cancel", "id": "bodyless" }), + 3, + )) + .await + .unwrap(); + let captured = + tokio::time::timeout(std::time::Duration::from_millis(250), captures.recv()).await; + client.abort(); + local.await.unwrap(); + let event = captured + .expect("bodyless response must capture without body demand") + .expect("capture channel closed"); + assert_eq!(event.webhook.response_status, status); + assert!(event.webhook.response_body.is_empty()); + assert!(!event.webhook.response_body_incomplete); + } + } + + #[test] + fn relay_body_limit_counts_decoded_bytes_at_the_base64_padding_boundary() { + use base64::Engine; + const MAXIMUM: usize = 10 * 1024 * 1024; + for excess in 0..=2 { + let body = base64::engine::general_purpose::STANDARD.encode(vec![0; MAXIMUM + excess]); + let request = serde_json::json!({ "id": "boundary", "method": "POST", "url": "/", "headers": {}, "body": body }); + assert_eq!( + validate_relay_request(&request.to_string()).is_ok(), + excess == 0, + "decoded size: {}", + MAXIMUM + excess + ); + } + } + + #[tokio::test] + async fn persistence_receipt_keeps_request_memory_until_forwarding_returns() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + drop(listener); + let budget = Arc::new(tokio::sync::Semaphore::new(1)); + let permit = Arc::new(Arc::clone(&budget).acquire_owned().await.unwrap()); + let (webhook_tx, mut webhook_rx) = tokio::sync::mpsc::channel(1); + let task = tokio::spawn(forward_http_request( + reqwest::Client::builder().no_proxy().build().unwrap(), + lpm_common::LocalTarget::loopback(lpm_common::LocalScheme::Http, address.port()), + ServerMessage::HttpRequest { + id: "persisted".into(), + method: "POST".into(), + url: "/".into(), + headers: HashMap::new(), + body: String::new(), + }, + true, + Some(webhook_tx), + Arc::new(tokio::sync::Semaphore::new(64)), + Some(permit), + None, + )); + let event = webhook_rx.recv().await.unwrap(); + event.complete_persistence(Ok(())); + assert_eq!( + budget.available_permits(), + 0, + "request is still retained by the forwarding future" + ); + task.await.unwrap(); + assert_eq!(budget.available_permits(), 1); + } + #[tokio::test] async fn queued_capture_keeps_response_memory_reserved_until_received() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -4852,7 +5294,7 @@ mod tests { #[tokio::test] async fn concurrent_large_requests_cannot_reserve_more_than_the_byte_budget() { - let encoded_len = (20 * 1024 * 1024_usize).div_ceil(3) * 4; + let encoded_len = (10 * 1024 * 1024_usize).div_ceil(3) * 4; let permits = request_memory_permits(encoded_len).unwrap(); let budget = Arc::new(tokio::sync::Semaphore::new(HTTP_REQUEST_MEMORY_PERMITS)); let mut reservations = Vec::new(); @@ -5027,19 +5469,90 @@ mod tests { #[test] fn inbound_websocket_admission_accounts_for_wire_and_decoded_buffers() { - let wire_bytes = MAX_WS_MESSAGE_SIZE - 1; + let wire_bytes = 50 * 1024 * 1024 - 1; let base64_bytes = wire_bytes - 128; assert!(inbound_websocket_memory_permits(wire_bytes, base64_bytes).is_none()); } + #[tokio::test] + async fn local_websocket_writer_retains_memory_until_stalled_send_stops() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.set_send_buffer_size(1024).unwrap(); + let (stream, accepted) = tokio::join!( + socket.connect(listener.local_addr().unwrap()), + listener.accept() + ); + let stream = stream.unwrap(); + let (peer, _) = accepted.unwrap(); + let websocket = tokio_tungstenite::WebSocketStream::from_raw_socket( + tokio_tungstenite::MaybeTlsStream::Plain(stream), + tokio_tungstenite::tungstenite::protocol::Role::Client, + None, + ) + .await; + let (local_write, local_read) = websocket.split(); + let (relay_tx, _relay_rx) = tokio::sync::mpsc::channel(4); + let (closed_tx, _closed_rx) = tokio::sync::mpsc::unbounded_channel(); + let memory = Arc::new(tokio::sync::Semaphore::new(1)); + let mut tasks = tokio::task::JoinSet::new(); + let connection = activate_local_websocket( + &mut tasks, + LocalWebSocketActivation { + id: "stalled".into(), + generation: 1, + local_write, + local_read, + slot: Arc::new(tokio::sync::Semaphore::new(1)) + .acquire_owned() + .await + .unwrap(), + relay_tx, + relay_memory: Arc::new(tokio::sync::Semaphore::new(WEBSOCKET_MEMORY_PERMITS)), + ws_tx: None, + closed_tx, + }, + ); + connection + .commands + .send(LocalWebSocketCommand::Frame { + data: vec![0; 8 * 1024 * 1024], + is_binary: true, + _memory_permit: Arc::clone(&memory).acquire_owned().await.unwrap(), + }) + .await + .unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while connection.commands.capacity() != connection.commands.max_capacity() { + tokio::task::yield_now().await; + } + }) + .await + .expect("writer consumed the command"); + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + let retained = memory.available_permits(); + connection.cancel.cancel(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while tasks.join_next().await.is_some() {} + }) + .await + .expect("cancellation stops both halves"); + drop(peer); + assert_eq!( + retained, 0, + "in-flight native output must retain its memory reservation" + ); + assert_eq!(memory.available_permits(), 1); + } + #[tokio::test] async fn local_websocket_close_is_not_followed_by_a_synthetic_close() { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); - let mut websocket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let mut websocket = mock_relay::accept(stream).await.unwrap(); websocket .send(Message::Close(Some(CloseFrame { code: CloseCode::Normal, @@ -5103,7 +5616,7 @@ mod tests { let expected = code.filter(|code| *code != 1005); let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); - let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let mut socket = mock_relay::accept(stream).await.unwrap(); if server_initiates { socket .send(Message::Close(expected.map(|code| CloseFrame { diff --git a/crates/lpm-tunnel/src/client/auto_ack_tests.rs b/crates/lpm-tunnel/src/client/auto_ack_tests.rs index 1069d5cbd..86763866c 100644 --- a/crates/lpm-tunnel/src/client/auto_ack_tests.rs +++ b/crates/lpm-tunnel/src/client/auto_ack_tests.rs @@ -5,16 +5,16 @@ async fn failed_forward( ) -> ClientMessage { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(async move { - let (socket, _) = listener.accept().await.unwrap(); - drop(socket); - }); + drop(listener); let request = ServerMessage::HttpRequest { id: "needs-replay".into(), method: "POST".into(), url: "/webhook".into(), headers: HashMap::new(), - body: "e30=".into(), + body: base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + [0, 1, 2, 127, 128, 255], + ), }; let (response, _) = forward_http_request( reqwest::Client::builder() @@ -31,7 +31,6 @@ async fn failed_forward( None, ) .await; - server.await.unwrap(); response } @@ -69,6 +68,7 @@ async fn auto_ack_waits_for_successful_persistence() { let forward = tokio::spawn(failed_forward(Some(tx))); let event = rx.recv().await.unwrap(); assert!(event.requires_persistence()); + assert_eq!(event.webhook.request_body, [0, 1, 2, 127, 128, 255]); assert!(!forward.is_finished()); event.complete_persistence(Ok(())); assert_eq!(extract_response_data(&forward.await.unwrap()).0, 200); @@ -109,3 +109,74 @@ async fn auto_ack_persistence_deadline_returns_retryable_error() { assert_eq!(extract_response_data(&forward.await.unwrap()).0, 503); event.complete_persistence(Ok(())); } + +#[test] +fn taking_request_body_releases_the_encoded_allocation_and_preserves_binary_bytes() { + let mut request = ServerMessage::HttpRequest { + id: "binary".into(), + method: "POST".into(), + url: "/".into(), + headers: HashMap::new(), + body: base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + [0, 1, 2, 127, 128, 255], + ), + }; + assert_eq!( + take_request_body(&mut request).unwrap().as_ref(), + [0, 1, 2, 127, 128, 255] + ); + let ServerMessage::HttpRequest { body, .. } = request else { + panic!("expected HTTP") + }; + assert_eq!(body.capacity(), 0); +} + +#[tokio::test] +async fn successful_binary_upload_preserves_origin_and_capture_bytes() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut headers = Vec::with_capacity(1024); + while !headers.ends_with(b"\r\n\r\n") { + assert!(headers.len() < 8192); + headers.push(socket.read_u8().await.unwrap()); + } + let mut body = [0; 6]; + socket.read_exact(&mut body).await.unwrap(); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") + .await + .unwrap(); + body + }); + let (tx, mut rx) = tokio::sync::mpsc::channel(1); + let request = ServerMessage::HttpRequest { + id: "binary-success".into(), + method: "POST".into(), + url: "/webhook".into(), + headers: HashMap::new(), + body: base64::Engine::encode( + &base64::engine::general_purpose::STANDARD, + [0, 1, 2, 127, 128, 255], + ), + }; + let (response, _) = forward_http_request( + reqwest::Client::builder().no_proxy().build().unwrap(), + lpm_common::LocalTarget::loopback(lpm_common::LocalScheme::Http, port), + request, + false, + Some(tx), + Arc::new(tokio::sync::Semaphore::new(HTTP_RESPONSE_MEMORY_PERMITS)), + None, + None, + ) + .await; + assert_eq!(extract_response_data(&response).0, 200); + assert_eq!(server.await.unwrap(), [0, 1, 2, 127, 128, 255]); + let capture = rx.recv().await.unwrap(); + assert_eq!(capture.webhook.request_body, [0, 1, 2, 127, 128, 255]); + assert_eq!(capture.webhook.response_body, b"ok"); +} diff --git a/crates/lpm-tunnel/src/client/mock_relay.rs b/crates/lpm-tunnel/src/client/mock_relay.rs new file mode 100644 index 000000000..0a6cebac6 --- /dev/null +++ b/crates/lpm-tunnel/src/client/mock_relay.rs @@ -0,0 +1,113 @@ +use futures_util::{Sink, Stream, StreamExt}; +use std::collections::HashMap; +use std::pin::Pin; +use std::task::{Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio_tungstenite::{ + WebSocketStream, + tungstenite::{Error, Message}, +}; + +pub(super) struct MockRelay { + socket: WebSocketStream, + sequence: u64, + pulls: HashMap, + ready: bool, +} + +pub(super) async fn accept( + stream: S, +) -> Result, Error> { + Ok(MockRelay { + socket: tokio_tungstenite::accept_async(stream).await?, + sequence: 0, + pulls: HashMap::new(), + ready: false, + }) +} + +pub(super) async fn accept_with_headers(stream: S, callback: C) -> Result, Error> +where + S: AsyncRead + AsyncWrite + Unpin, + C: tokio_tungstenite::tungstenite::handshake::server::Callback + Unpin, +{ + Ok(MockRelay { + socket: tokio_tungstenite::accept_hdr_async(stream, callback).await?, + sequence: 0, + pulls: HashMap::new(), + ready: false, + }) +} + +impl MockRelay { + pub(super) async fn close( + &mut self, + frame: Option>, + ) -> Result<(), Error> { + while !self.ready { + let Some(message) = self.socket.next().await else { + break; + }; + if let Message::Text(text) = message? { + let value: serde_json::Value = serde_json::from_str(&text).unwrap(); + self.ready = value["type"] == "client_ready"; + } + } + self.socket.close(frame).await + } +} + +impl Sink for MockRelay { + type Error = Error; + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.socket).poll_ready(cx) + } + fn start_send(mut self: Pin<&mut Self>, message: Message) -> Result<(), Error> { + let message = if let Message::Text(text) = message { + let mut value: serde_json::Value = serde_json::from_str(&text).unwrap(); + if value["type"] == "hello" { + value["protocol"] = serde_json::json!(4); + } + if value["type"] == "http_response_pull" { + let id = value["id"].as_str().unwrap().to_owned(); + let next = self.pulls.entry(id).or_insert(0); + *next += 1; + value["seq"] = serde_json::json!(*next); + } + self.sequence += 1; + value["transport_seq"] = serde_json::json!(self.sequence); + value["transport_ack_nonce"] = + serde_json::json!(format!("00000000-0000-4000-8000-{:012x}", self.sequence)); + Message::Text(value.to_string()) + } else { + message + }; + Pin::new(&mut self.socket).start_send(message) + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.socket).poll_flush(cx) + } + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.socket).poll_close(cx) + } +} + +impl Stream for MockRelay { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + let next = Pin::new(&mut self.socket).poll_next(cx); + if let Poll::Ready(Some(Ok(Message::Text(ref text)))) = next + && let Ok(value) = serde_json::from_str::(text) + && matches!( + value["type"].as_str(), + Some("transport_ack" | "client_ready") + ) + { + self.ready |= value["type"] == "client_ready"; + continue; + } + return next; + } + } +} diff --git a/crates/lpm-tunnel/src/client/protocol_tests.rs b/crates/lpm-tunnel/src/client/protocol_tests.rs new file mode 100644 index 000000000..f2c3c3075 --- /dev/null +++ b/crates/lpm-tunnel/src/client/protocol_tests.rs @@ -0,0 +1,713 @@ +use super::*; +use serde_json::{Value, json}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::task::JoinHandle; +use tokio::time::{Duration, timeout}; +use tokio_tungstenite::WebSocketStream; + +#[test] +fn slow_failed_startup_never_resets_the_retry_budget() { + assert!(!healthy_tunnel_attempt(None)); + assert!(!healthy_tunnel_attempt(Some(Duration::from_secs(59)))); + assert!(healthy_tunnel_attempt(Some(Duration::from_secs(60)))); +} + +#[tokio::test] +#[expect( + clippy::result_large_err, + reason = "The WebSocket handshake callback requires an unboxed error response" +)] +async fn rejected_token_refresh_time_does_not_reset_retry_history() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut options = TunnelOptions::new("test-token".into(), 3000); + options.no_pin = true; + options.relay_url = format!("ws://{}/connect", listener.local_addr().unwrap()); + let shutdown = tokio_util::sync::CancellationToken::new(); + options.shutdown = Some(shutdown.clone()); + let entered = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let barrier = Arc::clone(&entered); + let resume = Arc::clone(&release); + options.token_provider = Some(TunnelTokenProvider::new( + || Box::pin(async { Ok("test-token".into()) }), + move || { + let barrier = Arc::clone(&barrier); + let resume = Arc::clone(&resume); + Box::pin(async move { + barrier.notify_one(); + resume.notified().await; + Err(LpmError::Network("refresh unavailable".into())) + }) + }, + )); + let (socket_tx, socket_rx) = tokio::sync::oneshot::channel(); + let relay = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let _ = tokio_tungstenite::accept_hdr_async( + stream, + |_: &tokio_tungstenite::tungstenite::handshake::server::Request, + _: tokio_tungstenite::tungstenite::handshake::server::Response| { + Err(tokio_tungstenite::tungstenite::http::Response::builder() + .status(503) + .body(Some("unavailable".into())) + .unwrap()) + }, + ) + .await; + let (stream, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + socket + .send(Message::Text(receipt(hello(), 1).to_string())) + .await + .unwrap(); + loop { + let message = socket.next().await.unwrap().unwrap(); + if let Message::Text(text) = message + && serde_json::from_str::(&text).unwrap()["type"] == "client_ready" + { + break; + } + } + socket_tx.send(socket).unwrap(); + }); + let (retry_tx, mut retries) = tokio::sync::mpsc::unbounded_channel(); + let client = tokio::spawn(async move { + connect_with_usage_fallible( + &options, + |_| Ok(()), + |message| { + retry_tx.send(message.to_owned()).unwrap(); + }, + |_, _| {}, + ) + .await + }); + let first = timeout(Duration::from_secs(2), retries.recv()) + .await + .unwrap() + .unwrap(); + assert!(first.contains("(unavailable") || first.contains("retrying")); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(4)).await; + tokio::time::resume(); + let mut socket = timeout(Duration::from_secs(2), socket_rx) + .await + .unwrap() + .unwrap(); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(57)).await; + tokio::time::resume(); + socket + .send(Message::Text( + receipt( + json!({"type":"error", "message":"credential expired", "code":"auth_failed"}), + 2, + ) + .to_string(), + )) + .await + .unwrap(); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(5)).await; + tokio::time::resume(); + release.notify_one(); + let second = timeout(Duration::from_secs(2), retries.recv()) + .await + .unwrap() + .unwrap(); + shutdown.cancel(); + client.await.unwrap().unwrap(); + relay.await.unwrap(); + let seconds: u64 = second + .split("retrying in ") + .nth(1) + .unwrap() + .split('s') + .next() + .unwrap() + .parse() + .unwrap(); + assert!( + (4..=6).contains(&seconds), + "refresh time reset retry history: {second}" + ); +} + +#[expect( + clippy::result_large_err, + reason = "The WebSocket handshake callback requires an unboxed error response" +)] +async fn pending_credentials_drop(rejected: bool, cancel: bool) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let shutdown = tokio_util::sync::CancellationToken::new(); + let dropped = tokio_util::sync::CancellationToken::new(); + let entered = Arc::new(tokio::sync::Notify::new()); + let mut options = TunnelOptions::new("test-token".into(), 3000); + options.no_pin = true; + options.relay_url = format!("ws://{}/connect", listener.local_addr().unwrap()); + options.shutdown = Some(shutdown.clone()); + let pending_provider = { + let entered = Arc::clone(&entered); + let dropped = dropped.clone(); + move || { + let entered = Arc::clone(&entered); + let guard = dropped.clone().drop_guard(); + Box::pin(async move { + let _guard = guard; + entered.notify_one(); + std::future::pending().await + }) as TunnelTokenFuture + } + }; + options.token_provider = Some(if rejected { + TunnelTokenProvider::new( + || Box::pin(async { Ok("test-token".into()) }), + pending_provider, + ) + } else { + TunnelTokenProvider::new(pending_provider, || { + Box::pin(async { panic!("unexpected refresh") }) + }) + }); + let relay = tokio::spawn(async move { + if rejected { + let (stream, _) = listener.accept().await.unwrap(); + let result = tokio_tungstenite::accept_hdr_async( + stream, + |_: &tokio_tungstenite::tungstenite::handshake::server::Request, + _: tokio_tungstenite::tungstenite::handshake::server::Response| { + Err(tokio_tungstenite::tungstenite::http::Response::builder() + .status(401) + .body(Some("invalid token".into())) + .unwrap()) + }, + ) + .await; + assert!(result.is_err()); + } + }); + let retry_observed = Arc::new(Mutex::new(String::new())); + let observed = Arc::clone(&retry_observed); + let retry_shutdown = shutdown.clone(); + let mut client = tokio::spawn(async move { + if rejected { + connect_with_usage_fallible( + &options, + |_| panic!("unexpected connection"), + |message| { + *observed.lock().unwrap() = message.to_owned(); + retry_shutdown.cancel(); + }, + |_, _| {}, + ) + .await + .map_err(TunnelConnectError::from_token_provider) + } else { + try_connect(&options, &|_| panic!("unexpected connection"), &|_, _| {}).await + } + }); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + if cancel { + shutdown.cancel(); + } else { + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(20)).await; + } + let result = timeout(Duration::from_millis(200), &mut client).await; + if result.is_err() { + client.abort(); + } + let result = result + .expect("credential provider exceeded its deadline") + .unwrap(); + if cancel || rejected { + result.unwrap(); + } else { + let error = result.unwrap_err(); + assert_eq!(error.retry_class, RetryClass::Transient); + assert!(error.to_string().contains("Tunnel credential")); + } + if rejected && !cancel { + assert!( + retry_observed + .lock() + .unwrap() + .contains("Tunnel credential refresh timed out") + ); + } + assert!(dropped.is_cancelled()); + relay.await.unwrap(); +} + +#[tokio::test] +async fn shutdown_drops_pending_initial_credentials() { + pending_credentials_drop(false, true).await; +} + +#[tokio::test] +async fn shutdown_drops_pending_rejected_credentials() { + pending_credentials_drop(true, true).await; +} + +#[tokio::test] +async fn pending_initial_credentials_expire_after_twenty_seconds() { + pending_credentials_drop(false, false).await; +} + +#[tokio::test] +async fn pending_rejected_credentials_expire_after_twenty_seconds() { + pending_credentials_drop(true, false).await; +} + +#[tokio::test] +async fn silent_upgrade_and_hello_have_deadlines_and_cooperative_shutdown() { + for hello in [false, true] { + for cancel in [false, true] { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let entered = Arc::new(tokio::sync::Notify::new()); + let shutdown = tokio_util::sync::CancellationToken::new(); + let mut options = TunnelOptions::new("test-token".into(), 3000); + options.no_pin = true; + options.relay_url = format!("ws://{}/connect", listener.local_addr().unwrap()); + options.shutdown = Some(shutdown.clone()); + let barrier = Arc::clone(&entered); + let relay = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + if hello { + let _socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + barrier.notify_one(); + std::future::pending::<()>().await; + } else { + let mut input = [0; 1024]; + assert!(stream.read(&mut input).await.unwrap() > 0); + barrier.notify_one(); + std::future::pending::<()>().await; + } + }); + let mut client = tokio::spawn(async move { + try_connect( + &options, + &|_| panic!("silent relay published readiness"), + &|_, _| {}, + ) + .await + }); + timeout(Duration::from_secs(2), entered.notified()) + .await + .unwrap(); + if cancel { + shutdown.cancel(); + } + let result = timeout( + Duration::from_secs(if cancel { 1 } else { 16 }), + &mut client, + ) + .await + .unwrap() + .unwrap(); + relay.abort(); + if cancel { + result.unwrap(); + } else { + let error = result.unwrap_err(); + assert_eq!(error.retry_class, RetryClass::Transient); + assert!(error.to_string().contains(if hello { + "Tunnel hello timed out" + } else { + "Tunnel connection timed out" + })); + } + } + } +} + +struct RawRelay { + socket: WebSocketStream, + client: JoinHandle>, + connected: Arc, + captures: tokio::sync::mpsc::Receiver, + sequence: u64, +} + +impl Drop for RawRelay { + fn drop(&mut self) { + self.client.abort(); + } +} + +fn receipt(mut value: Value, sequence: u64) -> Value { + value["transport_seq"] = json!(sequence); + value["transport_ack_nonce"] = json!(format!("00000000-0000-4000-8000-{sequence:012x}")); + value +} + +fn hello() -> Value { + json!({"type":"hello", "protocol":4, "subdomain":"raw.localhost", "tunnel_url":"http://raw.localhost", "session_id":"raw-session"}) +} + +impl RawRelay { + async fn new(port: u16) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut options = TunnelOptions::new("test-token".into(), port); + options.relay_url = format!("ws://{}/connect", listener.local_addr().unwrap()); + options.no_pin = true; + let (capture_tx, captures) = tokio::sync::mpsc::channel(1); + options.webhook_tx = Some(capture_tx); + let connected = Arc::new(AtomicUsize::new(0)); + let calls = Arc::clone(&connected); + let client = tokio::spawn(async move { + try_connect( + &options, + &|_| { + calls.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + &|_, _| {}, + ) + .await + }); + let (socket, _) = timeout(Duration::from_secs(2), listener.accept()) + .await + .unwrap() + .unwrap(); + let socket = tokio_tungstenite::accept_async(socket).await.unwrap(); + Self { + socket, + client, + connected, + captures, + sequence: 0, + } + } + + async fn next(&mut self) -> Value { + loop { + let message = timeout(Duration::from_secs(2), self.socket.next()) + .await + .expect("missing raw relay response") + .unwrap() + .unwrap(); + if let Message::Text(text) = message { + let value: Value = serde_json::from_str(&text).unwrap(); + if value["type"] != "ping" { + return value; + } + } + } + } + + async fn send(&mut self, value: Value) { + self.sequence += 1; + let value = receipt(value, self.sequence); + self.socket + .send(Message::Text(value.to_string())) + .await + .unwrap(); + assert_eq!( + self.next().await, + json!({"type":"transport_ack", "transport_seq":self.sequence, "transport_ack_nonce":value["transport_ack_nonce"]}) + ); + } + + async fn ready(&mut self) { + self.send(hello()).await; + assert_eq!(self.next().await, json!({"type":"client_ready"})); + assert_eq!(self.connected.load(Ordering::SeqCst), 1); + } + + async fn permanent_failure(&mut self) { + let result = timeout(Duration::from_secs(2), &mut self.client) + .await + .expect("invalid frame must terminate connection") + .unwrap(); + assert_eq!(result.unwrap_err().retry_class, RetryClass::Permanent); + while let Ok(Some(Ok(Message::Text(text)))) = + timeout(Duration::from_millis(100), self.socket.next()).await + { + let value: Value = serde_json::from_str(&text).unwrap(); + assert_ne!(value["type"], "transport_ack"); + assert_ne!(value["type"], "client_ready"); + } + } +} + +#[tokio::test] +async fn invalid_raw_hello_never_acknowledges_or_publishes_readiness() { + for mutation in [ + "missing-protocol", + "old-protocol", + "missing-receipt", + "short-nonce", + "skipped-sequence", + ] { + let mut relay = RawRelay::new(3000).await; + let mut value = receipt(hello(), 1); + match mutation { + "missing-protocol" => { + value.as_object_mut().unwrap().remove("protocol"); + } + "old-protocol" => value["protocol"] = json!(3), + "missing-receipt" => { + value.as_object_mut().unwrap().remove("transport_seq"); + } + "short-nonce" => value["transport_ack_nonce"] = json!("short"), + "skipped-sequence" => value["transport_seq"] = json!(2), + _ => unreachable!(), + } + relay + .socket + .send(Message::Text(value.to_string())) + .await + .unwrap(); + relay.permanent_failure().await; + assert_eq!(relay.connected.load(Ordering::SeqCst), 0, "{mutation}"); + } +} + +#[tokio::test] +async fn invalid_raw_active_receipts_terminate_without_acknowledgement() { + for mutation in [ + "missing-sequence", + "missing-nonce", + "string-sequence", + "overflow-sequence", + "replay", + "skip", + "short-nonce", + "invalid-character", + "misplaced-hyphen", + "invalid-json", + ] { + let mut relay = RawRelay::new(3000).await; + relay.ready().await; + let mut value = receipt(json!({"type":"pong"}), 2); + match mutation { + "missing-sequence" => { + value.as_object_mut().unwrap().remove("transport_seq"); + } + "missing-nonce" => { + value.as_object_mut().unwrap().remove("transport_ack_nonce"); + } + "string-sequence" => value["transport_seq"] = json!("2"), + "overflow-sequence" => value["transport_seq"] = json!(u64::MAX), + "replay" => value["transport_seq"] = json!(1), + "skip" => value["transport_seq"] = json!(3), + "short-nonce" => value["transport_ack_nonce"] = json!("a"), + "invalid-character" => { + value["transport_ack_nonce"] = json!("z0000000-0000-4000-8000-000000000002") + } + "misplaced-hyphen" => { + value["transport_ack_nonce"] = json!("000000000000-4000-8000-00000000000-") + } + "invalid-json" => {} + _ => unreachable!(), + } + let text = if mutation == "invalid-json" { + "{".into() + } else if mutation == "overflow-sequence" { + value + .to_string() + .replace(&u64::MAX.to_string(), "18446744073709551616") + } else { + value.to_string() + }; + relay.socket.send(Message::Text(text)).await.unwrap(); + relay.permanent_failure().await; + } +} + +#[tokio::test] +async fn protocol_four_public_websocket_rejection_never_connects_to_origin() { + let origin = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let mut relay = RawRelay::new(origin.local_addr().unwrap().port()).await; + relay.ready().await; + relay + .send(json!({"type":"ws_upgrade", "id":"browser", "url":"/socket", "headers":{}})) + .await; + assert_eq!( + relay.next().await, + json!({"type":"ws_reject", "id":"browser", "error":"Public WebSockets are unavailable"}) + ); + assert!( + timeout(Duration::from_millis(100), origin.accept()) + .await + .is_err() + ); +} + +async fn event_origin(two_chunks: bool, finish: bool) -> (u16, JoinHandle) { + let origin = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = origin.local_addr().unwrap().port(); + let task = tokio::spawn(async move { + let (mut socket, _) = origin.accept().await.unwrap(); + let mut request = [0; 4096]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n9\r\ndata: 1\n\n\r\n").await.unwrap(); + if two_chunks { + socket.write_all(b"9\r\ndata: 2\n\n\r\n").await.unwrap(); + } + if finish { + socket.write_all(b"0\r\n\r\n").await.unwrap(); + return 0; + } + socket.read(&mut request).await.unwrap() + }); + (port, task) +} + +async fn start_stream(relay: &mut RawRelay) { + relay.ready().await; + relay.send(json!({"type":"http_request", "id":"sse", "method":"GET", "url":"/events", "headers":{}, "body":""})).await; + let start = relay.next().await; + assert_eq!(start["type"], "http_response_start"); + assert_eq!(start["status"], 200); + assert!( + timeout(Duration::from_millis(50), relay.socket.next()) + .await + .is_err(), + "body must wait for visitor demand" + ); + relay + .send(json!({"type":"http_response_pull", "id":"sse", "seq":1})) + .await; + let chunk = relay.next().await; + assert_eq!(chunk["type"], "http_response_chunk"); + assert_eq!(chunk["seq"], 1); + assert_eq!( + base64::Engine::decode( + &base64::engine::general_purpose::STANDARD, + chunk["body"].as_str().unwrap() + ) + .unwrap(), + b"data: 1\n\n" + ); + assert!( + timeout(Duration::from_millis(50), relay.socket.next()) + .await + .is_err(), + "second chunk needs another credit" + ); +} + +#[tokio::test] +async fn raw_stream_cancellation_and_invalid_pull_sequences_close_the_origin() { + for sequence in [None, Some(1), Some(3)] { + let (port, origin) = event_origin(true, false).await; + let mut relay = RawRelay::new(port).await; + start_stream(&mut relay).await; + match sequence { + Some(seq) => { + relay + .send(json!({"type":"http_response_pull", "id":"sse", "seq":seq})) + .await + } + None => relay.send(json!({"type":"http_cancel", "id":"sse"})).await, + } + assert_eq!( + relay.next().await, + json!({"type":"http_cancelled", "id":"sse"}) + ); + assert_eq!( + timeout(Duration::from_secs(2), origin) + .await + .unwrap() + .unwrap(), + 0 + ); + } +} + +#[tokio::test] +async fn raw_stream_requires_each_credit_and_finishes_at_the_next_sequence() { + let (port, origin) = event_origin(true, true).await; + let mut relay = RawRelay::new(port).await; + start_stream(&mut relay).await; + relay + .send(json!({"type":"http_response_pull", "id":"sse", "seq":2})) + .await; + let chunk = relay.next().await; + assert_eq!(chunk["type"], "http_response_chunk"); + assert_eq!(chunk["seq"], 2); + assert_eq!( + base64::Engine::decode( + &base64::engine::general_purpose::STANDARD, + chunk["body"].as_str().unwrap() + ) + .unwrap(), + b"data: 2\n\n" + ); + assert!( + timeout(Duration::from_millis(50), relay.socket.next()) + .await + .is_err() + ); + relay + .send(json!({"type":"http_response_pull", "id":"sse", "seq":3})) + .await; + assert_eq!( + relay.next().await, + json!({"type":"http_response_end", "id":"sse", "failed":false, "seq":3}) + ); + assert_eq!( + timeout(Duration::from_secs(2), origin) + .await + .unwrap() + .unwrap(), + 0 + ); +} + +#[tokio::test] +async fn raw_sse_survives_thirty_seconds_and_stops_at_thirty_minutes_with_partial_capture() { + let (port, origin) = event_origin(false, false).await; + let mut relay = RawRelay::new(port).await; + start_stream(&mut relay).await; + relay + .send(json!({"type":"http_response_pull", "id":"sse", "seq":2})) + .await; + let mut elapsed = 0; + for seconds in [16, 15].into_iter().chain(std::iter::repeat_n(29, 60)) { + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(seconds)).await; + tokio::time::resume(); + elapsed += seconds; + relay.send(json!({"type":"pong"})).await; + if elapsed == 31 { + assert!(!origin.is_finished()); + assert!(!relay.client.is_finished()); + assert!( + timeout(Duration::from_millis(50), relay.socket.next()) + .await + .is_err() + ); + } + } + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(31)).await; + tokio::time::resume(); + assert_eq!( + relay.next().await, + json!({"type":"http_response_end", "id":"sse", "failed":true, "seq":2}) + ); + assert_eq!( + timeout(Duration::from_secs(2), origin) + .await + .unwrap() + .unwrap(), + 0 + ); + let capture = timeout(Duration::from_secs(2), relay.captures.recv()) + .await + .unwrap() + .unwrap(); + assert!(capture.webhook.response_body_incomplete); + assert_eq!(capture.webhook.response_status, 200); + assert_eq!(capture.webhook.response_body, b"data: 1\n\n"); +} diff --git a/crates/lpm-tunnel/src/http_stream.rs b/crates/lpm-tunnel/src/http_stream.rs index e9d4d2a52..df7502688 100644 --- a/crates/lpm-tunnel/src/http_stream.rs +++ b/crates/lpm-tunnel/src/http_stream.rs @@ -1,5 +1,6 @@ use crate::protocol::ClientMessage; use base64::{Engine, engine::general_purpose::STANDARD}; +use futures_util::FutureExt; use lpm_common::LpmError; use std::collections::HashMap; use tokio::sync::mpsc; @@ -9,8 +10,8 @@ const CHUNK_BYTES: usize = 64 * 1024; const CAPTURE_BYTES: usize = 64 * 1024; pub(crate) struct StreamOutput { - pub(crate) messages: mpsc::Sender, - pub(crate) credits: mpsc::Receiver<()>, + pub(crate) messages: mpsc::Sender, + pub(crate) credits: mpsc::Receiver, pub(crate) cancel: CancellationToken, } @@ -18,6 +19,7 @@ pub(crate) struct StreamCapture { pub(crate) body: Vec, pub(crate) incomplete: bool, pub(crate) failed: bool, + pub(crate) next_sequence: u64, } impl StreamOutput { @@ -26,25 +28,48 @@ impl StreamOutput { id: &str, status: u16, headers: &HashMap, + bodyless: bool, mut response: reqwest::Response, + deadline: tokio::time::Instant, ) -> Result { self.messages - .send(ClientMessage::HttpResponseStart { - id: id.to_string(), - status, - headers: headers.clone(), - }) + .send(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseStart { + id: id.to_string(), + status, + headers: headers.clone(), + }, + )) .await .map_err(|_| LpmError::Tunnel("response channel closed".into()))?; let mut capture = StreamCapture { body: Vec::new(), incomplete: false, failed: false, + next_sequence: 1, }; + if bodyless { + return Ok(capture); + } + let is_sse = headers.get("content-type").is_some_and(|value| { + value + .split(';') + .next() + .is_some_and(|mime| mime.trim().eq_ignore_ascii_case("text/event-stream")) + }); let mut buffered = None; + let mut received_bytes = 0usize; + let mut upstream_failed = false; + let mut upstream_finished = false; loop { let next = async { - self.credits.recv().await?; + let credit = self.credits.recv().await?; + if credit != capture.next_sequence { + return Some(Err(())); + } + if upstream_finished && buffered.is_none() { + return None; + } while buffered.is_none() { match response.chunk().await { Ok(Some(bytes)) if bytes.is_empty() => continue, @@ -58,9 +83,43 @@ impl StreamOutput { if bytes.is_empty() { buffered = None; } - Some(Ok(chunk)) + if !is_sse && chunk.len() < CHUNK_BYTES { + let mut joined = Vec::with_capacity(CHUNK_BYTES); + joined.extend_from_slice(&chunk); + for _ in 0..1024 { + if joined.len() == CHUNK_BYTES { + break; + } + tokio::task::yield_now().await; + let Some(next) = response.chunk().now_or_never() else { + break; + }; + match next { + Ok(Some(mut next)) => { + let take = next.len().min(CHUNK_BYTES - joined.len()); + joined.extend_from_slice(&next.split_to(take)); + if !next.is_empty() { + buffered = Some(next); + break; + } + } + Ok(None) => { + upstream_finished = true; + break; + } + Err(_) => { + upstream_failed = true; + break; + } + } + } + Some(Ok(joined.into())) + } else { + Some(Ok(chunk)) + } }; let chunk = tokio::select! { + _ = tokio::time::sleep_until(deadline) => { capture.failed = true; break; } _ = self.cancel.cancelled() => { capture.failed = true; break; } chunk = next => chunk, }; @@ -71,6 +130,11 @@ impl StreamOutput { break; } Some(Ok(chunk)) => { + received_bytes += chunk.len(); + if received_bytes > 50 * 1024 * 1024 { + capture.failed = true; + break; + } let remaining = CAPTURE_BYTES.saturating_sub(capture.body.len()); capture .body @@ -79,10 +143,17 @@ impl StreamOutput { let message = ClientMessage::HttpResponseChunk { id: id.to_string(), body: STANDARD.encode(chunk), + seq: capture.next_sequence, }; tokio::select! { + _ = tokio::time::sleep_until(deadline) => { capture.failed = true; break; } _ = self.cancel.cancelled() => { capture.failed = true; break; } - result = self.messages.send(message) => if result.is_err() { capture.failed = true; break; }, + result = self.messages.send(crate::client::HttpRelayFrame::Stream(message)) => if result.is_err() { capture.failed = true; break; }, + } + capture.next_sequence += 1; + if upstream_failed { + capture.failed = true; + break; } } } @@ -95,6 +166,184 @@ impl StreamOutput { #[cfg(test)] mod tests { use super::*; + use futures_util::StreamExt; + + #[tokio::test] + async fn ordinary_upstream_failure_preserves_the_available_prefix() { + let body = reqwest::Body::wrap_stream(futures_util::stream::iter([ + Ok(b"prefix".to_vec()), + Err(std::io::Error::other("origin failed")), + ])); + let response = tokio_tungstenite::tungstenite::http::Response::new(body).into(); + let (messages, mut received) = mpsc::channel(2); + let (credits, demand) = mpsc::channel(1); + let output = StreamOutput { + messages, + credits: demand, + cancel: CancellationToken::new(), + }; + let task = tokio::spawn(async move { + output + .forward( + "body", + 200, + &HashMap::new(), + false, + response, + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + ) + .await + }); + received.recv().await.unwrap(); + credits.send(1).await.unwrap(); + let Some(crate::client::HttpRelayFrame::Stream(ClientMessage::HttpResponseChunk { + body, + seq, + .. + })) = received.recv().await + else { + panic!("origin prefix must be delivered") + }; + assert_eq!(seq, 1); + assert_eq!(STANDARD.decode(body).unwrap(), b"prefix"); + let capture = task.await.unwrap().unwrap(); + assert_eq!(capture.body, b"prefix"); + assert!(capture.failed); + assert!(capture.incomplete); + assert_eq!(capture.next_sequence, 2); + } + + #[tokio::test] + async fn ordinary_background_body_frames_share_bounded_relay_chunks() { + let (sender, receiver) = mpsc::channel(1); + tokio::spawn(async move { + for _ in 0..128 { + sender + .send(Ok::<_, std::io::Error>(vec![42; 1024])) + .await + .unwrap(); + } + }); + let body = reqwest::Body::wrap_stream( + futures_util::stream::unfold(receiver, |mut receiver| async move { + receiver.recv().await.map(|chunk| (chunk, receiver)) + }) + .fuse(), + ); + let response = tokio_tungstenite::tungstenite::http::Response::new(body).into(); + let (messages, mut received) = mpsc::channel(1); + let (credits, demand) = mpsc::channel(1); + let task = tokio::spawn(async move { + StreamOutput { + messages, + credits: demand, + cancel: CancellationToken::new(), + } + .forward( + "body", + 200, + &HashMap::new(), + false, + response, + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + ) + .await + }); + received.recv().await.unwrap(); + let mut chunks = 0; + let mut delivered = Vec::with_capacity(128 * 1024); + for sequence in 1..=129 { + credits.send(sequence).await.unwrap(); + let Some(frame) = received.recv().await else { + break; + }; + let crate::client::HttpRelayFrame::Stream(ClientMessage::HttpResponseChunk { + body, + seq, + .. + }) = frame + else { + panic!("unexpected frame") + }; + assert_eq!(seq, sequence); + let decoded = STANDARD.decode(body).unwrap(); + assert!(decoded.len() <= CHUNK_BYTES); + delivered.extend_from_slice(&decoded); + chunks += 1; + } + let capture = task.await.unwrap().unwrap(); + assert!(!capture.failed); + assert_eq!(delivered, vec![42; 128 * 1024]); + assert_eq!( + chunks, 2, + "background origin scheduling must not amplify relay messages" + ); + } + + #[tokio::test] + async fn ordinary_ready_body_frames_share_one_bounded_relay_chunk() { + let body = reqwest::Body::wrap_stream(futures_util::stream::iter( + (0..128).map(|_| Ok::<_, std::io::Error>(vec![42; 1024])), + )); + let response = tokio_tungstenite::tungstenite::http::Response::new(body).into(); + let (messages, mut received) = mpsc::channel(1); + let (credits, demand) = mpsc::channel(1); + let output = StreamOutput { + messages, + credits: demand, + cancel: CancellationToken::new(), + }; + let task = tokio::spawn(async move { + output + .forward( + "body", + 200, + &HashMap::new(), + false, + response, + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + ) + .await + }); + assert!(matches!( + received.recv().await, + Some(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseStart { .. } + )) + )); + let mut chunks = 0; + let mut total = 0; + loop { + if credits.send(chunks + 1).await.is_err() { + break; + } + let Some(message) = received.recv().await else { + break; + }; + let crate::client::HttpRelayFrame::Stream(ClientMessage::HttpResponseChunk { + body, + seq, + .. + }) = message + else { + panic!("expected chunk") + }; + chunks += 1; + assert_eq!(seq, chunks); + let bytes = STANDARD.decode(body).unwrap(); + assert!(bytes.len() <= CHUNK_BYTES); + assert!(bytes.iter().all(|value| *value == 42)); + total += bytes.len(); + } + let capture = task.await.unwrap().unwrap(); + assert_eq!(total, 128 * 1024); + assert_eq!( + chunks, 2, + "small origin frames must not amplify paid relay message count" + ); + assert!(!capture.failed); + assert!(capture.incomplete); + } #[tokio::test] async fn empty_upstream_chunks_do_not_consume_relay_demand() { @@ -112,24 +361,36 @@ mod tests { }; let task = tokio::spawn(async move { output - .forward("stream", 200, &HashMap::new(), response) + .forward( + "stream", + 200, + &HashMap::new(), + false, + response, + tokio::time::Instant::now() + std::time::Duration::from_secs(30), + ) .await }); assert!(matches!( received.recv().await, - Some(ClientMessage::HttpResponseStart { .. }) + Some(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseStart { .. } + )) )); - credits.send(()).await.unwrap(); + credits.send(1).await.unwrap(); let chunk = tokio::time::timeout(std::time::Duration::from_secs(1), received.recv()) .await .expect("one relay credit must reach the next nonempty upstream chunk"); match chunk { - Some(ClientMessage::HttpResponseChunk { body, .. }) => { + Some(crate::client::HttpRelayFrame::Stream(ClientMessage::HttpResponseChunk { + body, + .. + })) => { assert_eq!(STANDARD.decode(body).unwrap(), b"data: ready\n\n"); } - other => panic!("expected stream chunk, received {other:?}"), + _ => panic!("expected stream chunk"), } - credits.send(()).await.unwrap(); + credits.send(2).await.unwrap(); let capture = task.await.unwrap().unwrap(); assert_eq!(capture.body, b"data: ready\n\n"); assert!(!capture.failed); diff --git a/crates/lpm-tunnel/src/lib.rs b/crates/lpm-tunnel/src/lib.rs index 708d0d6a6..75f200446 100644 --- a/crates/lpm-tunnel/src/lib.rs +++ b/crates/lpm-tunnel/src/lib.rs @@ -4,7 +4,7 @@ //! (e.g., `https://acme-api.lpm.llc`) and proxies HTTP requests to the endpoint. //! //! Supports multiple base domains (lpm.fyi, lpm.llc). -//! Free users get ephemeral random domains on lpm.fyi. +//! Public tunnels require Pro or Org. //! Pro/Org users can claim persistent domains on any available base domain. pub mod client; @@ -52,6 +52,12 @@ pub(crate) fn validate_forward_target( /// Plan and relay limits advertised by the tunnel service. #[derive(Debug, Clone, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct TunnelLimitMetadata { + /// Resource contract version. + #[serde(default)] + pub quota_version: Option, + /// Whether the relay accepts public WebSocket upgrades. + #[serde(default)] + pub public_websockets_available: Option, /// Maximum simultaneously open tunnels for the account. #[serde(default)] pub max_concurrent: Option, @@ -75,6 +81,45 @@ pub struct TunnelLimitMetadata { /// Account-wide request usage advertised by the tunnel service. #[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)] pub struct TunnelUsageMetadata { + /// Resource contract version. + #[serde(default)] + pub quota_version: Option, + /// Aggregate busy Durable Object time consumed in milliseconds. + #[serde(default)] + pub busy_object_ms: Option, + /// Included busy time in milliseconds. + #[serde(default)] + pub included_busy_object_ms: Option, + /// Inbound relay messages consumed. + #[serde(default)] + pub relay_messages: Option, + /// Included inbound relay messages. + #[serde(default)] + pub included_relay_messages: Option, + /// Bidirectional relay transfer consumed in bytes. + #[serde(default)] + pub transfer_bytes: Option, + /// Included bidirectional relay transfer in bytes. + #[serde(default)] + pub included_transfer_bytes: Option, + /// New connection attempts consumed in the paid period. + #[serde(default)] + pub connection_attempts: Option, + /// Connection attempts included across paid seats, independent of bundles. + #[serde(default)] + pub included_connection_attempts: Option, + /// Owner-approved monthly overage spending cap in cents. + #[serde(default)] + pub overage_spend_cap_cents: Option, + /// Overage bundles authorized by the spending cap. + #[serde(default)] + pub authorized_bundles: Option, + /// Overage bundles consumed by the largest resource excess. + #[serde(default)] + pub billable_bundles: Option, + /// Price per resource bundle in cents. + #[serde(default)] + pub bundle_price_cents: Option, /// Requests admitted during the current usage period. #[serde(default)] pub accepted_requests: Option, @@ -84,7 +129,7 @@ pub struct TunnelUsageMetadata { /// Requests above the included allowance. #[serde(default)] pub overage_requests: Option, - /// Whether requests continue into proportional overage. + /// Whether the owner enabled capped resource bundles. #[serde(default)] pub overage_enabled: Option, /// Whether the allowance is currently a hard stop. diff --git a/crates/lpm-tunnel/src/protocol.rs b/crates/lpm-tunnel/src/protocol.rs index b4f9dd708..1ee206be0 100644 --- a/crates/lpm-tunnel/src/protocol.rs +++ b/crates/lpm-tunnel/src/protocol.rs @@ -45,10 +45,10 @@ pub enum ClientMessage { }, /// One demand-limited response chunk, at most 64 KiB decoded. #[serde(rename = "http_response_chunk")] - HttpResponseChunk { id: String, body: String }, + HttpResponseChunk { id: String, body: String, seq: u64 }, /// End of a streamed response; failed means the body is incomplete. #[serde(rename = "http_response_end")] - HttpResponseEnd { id: String, failed: bool }, + HttpResponseEnd { id: String, failed: bool, seq: u64 }, /// WebSocket frame from local server back to the remote client. #[serde(rename = "ws_frame")] @@ -133,7 +133,7 @@ pub enum ServerMessage { /// The visitor is ready for one response chunk. #[serde(rename = "http_response_pull")] - HttpResponsePull { id: String }, + HttpResponsePull { id: String, seq: u64 }, /// The visitor stopped reading the response. #[serde(rename = "http_cancel")] HttpCancel { id: String }, @@ -292,6 +292,8 @@ mod tests { assert_eq!( limits, Some(crate::TunnelLimitMetadata { + quota_version: None, + public_websockets_available: None, max_concurrent: Some(1), request_rate_limit_per_minute: Some(4_000), per_ip_rate_limit_per_minute: Some(600), diff --git a/crates/lpm-tunnel/src/proxy.rs b/crates/lpm-tunnel/src/proxy.rs index 177cfadff..af1bb4924 100644 --- a/crates/lpm-tunnel/src/proxy.rs +++ b/crates/lpm-tunnel/src/proxy.rs @@ -115,7 +115,7 @@ pub async fn forward_request( request: &ServerMessage, ) -> Result { Ok( - forward_request_inner(http_client, local_target, request, None, None) + forward_request_inner(http_client, local_target, request, None, None, None) .await? .message, ) @@ -131,6 +131,7 @@ pub(crate) async fn forward_request_with_memory_budget( memory_unit_bytes: usize, memory_multiplier: usize, stream: Option, + decoded_body: Option, ) -> Result { forward_request_inner( http_client, @@ -143,6 +144,7 @@ pub(crate) async fn forward_request_with_memory_budget( memory_multiplier, )), stream, + decoded_body, ) .await } @@ -153,6 +155,7 @@ async fn forward_request_inner( request: &ServerMessage, memory_budget: Option<(Arc, usize, usize, usize)>, stream: Option, + decoded_body: Option, ) -> Result { crate::validate_forward_target(local_target)?; let (id, method, url, headers, body) = match request { @@ -192,7 +195,9 @@ async fn forward_request_inner( builder = builder.header("host", local_target.authority()); // Decode and attach body - if !body.is_empty() { + if let Some(bytes) = decoded_body { + builder = builder.body(bytes); + } else if !body.is_empty() { let body_bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, body) .map_err(|e| LpmError::Tunnel(format!("failed to decode request body: {e}")))?; @@ -239,8 +244,22 @@ async fn forward_request_inner( .next() .is_some_and(|mime| mime.trim().eq_ignore_ascii_case("text/event-stream")) }); - if is_event_stream && let Some(output) = stream { - let capture = output.forward(id, status, &resp_headers, response).await?; + if let Some(output) = stream { + let body_deadline = if is_event_stream { + tokio::time::Instant::now() + std::time::Duration::from_secs(30 * 60) + } else { + deadline + }; + let capture = output + .forward( + id, + status, + &resp_headers, + method.eq_ignore_ascii_case("HEAD") || matches!(status, 204 | 205 | 304), + response, + body_deadline, + ) + .await?; return Ok(ForwardedResponse { message: ClientMessage::HttpResponse { id: id.clone(), @@ -254,6 +273,7 @@ async fn forward_request_inner( stream_end: Some(ClientMessage::HttpResponseEnd { id: id.clone(), failed: capture.failed, + seq: capture.next_sequence, }), capture_incomplete: capture.incomplete, memory_permit: None, @@ -590,6 +610,130 @@ mod tests { assert_eq!(MAX_RESPONSE_BODY_SIZE, 50 * 1024 * 1024); } + #[tokio::test] + async fn response_timeout_preserves_capture_and_emits_a_failed_stream_end() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 1000000\r\n\r\nfirst") + .await + .unwrap(); + std::future::pending::<()>().await; + }); + let (messages, mut output) = tokio::sync::mpsc::channel(1); + let (credits, credit_rx) = tokio::sync::mpsc::channel(1); + let forward = tokio::spawn(async move { + let request = ServerMessage::HttpRequest { + id: "ordinary".into(), + method: "GET".into(), + url: "/file".into(), + headers: HashMap::new(), + body: String::new(), + }; + forward_request_inner( + &reqwest::Client::builder().no_proxy().build().unwrap(), + &lpm_common::LocalTarget::loopback(lpm_common::LocalScheme::Http, address.port()), + &request, + None, + Some(crate::http_stream::StreamOutput { + messages, + credits: credit_rx, + cancel: tokio_util::sync::CancellationToken::new(), + }), + None, + ) + .await + }); + assert!(matches!( + output.recv().await, + Some(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseStart { .. } + )) + )); + credits.send(1).await.unwrap(); + assert!(matches!( + output.recv().await, + Some(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseChunk { seq: 1, .. } + )) + )); + credits.send(2).await.unwrap(); + tokio::time::pause(); + tokio::time::advance(std::time::Duration::from_secs(31)).await; + let response = forward.await.unwrap(); + server.abort(); + let response = response + .expect("stream timeout must return a failed terminal, not discard streaming state"); + assert!(matches!( + response.stream_end, + Some(ClientMessage::HttpResponseEnd { + failed: true, + seq: 2, + .. + }) + )); + assert!(response.capture_incomplete); + assert!( + matches!(response.message, ClientMessage::HttpResponse { body, .. } if body == "Zmlyc3Q=") + ); + } + + #[tokio::test] + async fn ordinary_response_headers_arrive_before_the_body_is_complete() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + assert!(socket.read(&mut request).await.unwrap() > 0); + socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: application/octet-stream\r\nContent-Length: 1000000\r\n\r\nfirst").await.unwrap(); + std::future::pending::<()>().await; + }); + let (messages, mut output) = tokio::sync::mpsc::channel(1); + let (_credits, credit_rx) = tokio::sync::mpsc::channel(1); + let cancel = tokio_util::sync::CancellationToken::new(); + let forward_cancel = cancel.clone(); + let forward = tokio::spawn(async move { + let request = ServerMessage::HttpRequest { + id: "ordinary".into(), + method: "GET".into(), + url: "/file".into(), + headers: HashMap::new(), + body: String::new(), + }; + forward_request_inner( + &reqwest::Client::builder().no_proxy().build().unwrap(), + &lpm_common::LocalTarget::loopback(lpm_common::LocalScheme::Http, address.port()), + &request, + None, + Some(crate::http_stream::StreamOutput { + messages, + credits: credit_rx, + cancel: forward_cancel, + }), + None, + ) + .await + }); + let first = + tokio::time::timeout(std::time::Duration::from_millis(100), output.recv()).await; + cancel.cancel(); + forward.abort(); + server.abort(); + assert!(matches!( + first, + Ok(Some(crate::client::HttpRelayFrame::Stream( + ClientMessage::HttpResponseStart { status: 200, .. } + ))) + )); + } + #[tokio::test] async fn event_stream_reconnect_preserves_last_event_id() { use tokio::io::{AsyncReadExt, AsyncWriteExt}; @@ -681,12 +825,17 @@ mod tests { credits: credit_rx, cancel: tokio_util::sync::CancellationToken::new(), }), + None, ) .await .unwrap() }); let mut received = 0; + let mut sequence = 1; while let Some(message) = output.recv().await { + let crate::client::HttpRelayFrame::Stream(message) = message else { + panic!("unexpected completed response") + }; match message { ClientMessage::HttpResponseStart { status: 200, .. } => {} ClientMessage::HttpResponseChunk { body, .. } => { @@ -698,7 +847,8 @@ mod tests { } _ => panic!("unexpected stream message"), } - credits.send(()).await.unwrap(); + credits.send(sequence).await.unwrap(); + sequence += 1; } let response = forward.await.unwrap(); server.await.unwrap(); @@ -864,7 +1014,7 @@ mod tests { let result = tokio::time::timeout( std::time::Duration::from_secs(1), forward_request_with_memory_budget( - &client, &target, &request, budget, permits, unit_bytes, 4, None, + &client, &target, &request, budget, permits, unit_bytes, 4, None, None, ), ) .await; @@ -946,6 +1096,7 @@ mod tests { 1024 * 1024, 4, None, + None, ) .await }); diff --git a/tests/workflows/tests/dev_tunnel.rs b/tests/workflows/tests/dev_tunnel.rs index 1e803e4ff..ac31df7cf 100644 --- a/tests/workflows/tests/dev_tunnel.rs +++ b/tests/workflows/tests/dev_tunnel.rs @@ -15,7 +15,7 @@ mod support; -use futures_util::SinkExt; +use futures_util::{SinkExt, StreamExt}; use std::collections::VecDeque; use std::fs; use std::io::{Read, Write as _}; @@ -30,6 +30,28 @@ use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; const MAX_CAPTURED_STREAM_BYTES: usize = 256 * 1024; +async fn await_fixture_ready( + socket: &mut tokio_tungstenite::WebSocketStream, +) { + let ack = socket.next().await.unwrap().unwrap(); + let Message::Text(ack) = ack else { + panic!("expected hello receipt") + }; + let ack: serde_json::Value = serde_json::from_str(&ack).unwrap(); + assert_eq!(ack["type"], "transport_ack"); + assert_eq!(ack["transport_seq"], 1); + assert_eq!( + ack["transport_ack_nonce"], + "00000000-0000-4000-8000-000000000001" + ); + let ready = socket.next().await.unwrap().unwrap(); + let Message::Text(ready) = ready else { + panic!("expected readiness") + }; + let ready: serde_json::Value = serde_json::from_str(&ready).unwrap(); + assert_eq!(ready["type"], "client_ready"); +} + fn captured_webhook( id: &str, timestamp: &str, @@ -694,6 +716,8 @@ async fn dev_tunnel_stops_a_long_running_server_when_the_relay_fails() { .send(Message::Text( serde_json::json!({ "type": "error", + "transport_seq": 1, + "transport_ack_nonce": "00000000-0000-4000-8000-000000000001", "message": "the account plan does not permit this tunnel", "code": "plan_required" }) @@ -760,10 +784,13 @@ async fn multi_service_dev_stops_when_the_tunnel_fails_after_readiness() { .send(Message::Text( serde_json::json!({ "type": "hello", + "protocol": 4, + "transport_seq": 1, + "transport_ack_nonce": "00000000-0000-4000-8000-000000000001", "subdomain": "multi-ready.lpm.test", "tunnel_url": "https://multi-ready.lpm.test", "session_id": "session-multi-ready", - "plan": "free", + "plan": "pro", "base_domain": "lpm.test", "domain_kind": "random" }) @@ -771,11 +798,14 @@ async fn multi_service_dev_stops_when_the_tunnel_fails_after_readiness() { )) .await .unwrap(); + await_fixture_ready(&mut websocket).await; tokio::time::sleep(Duration::from_millis(300)).await; websocket .send(Message::Text( serde_json::json!({ "type": "error", + "transport_seq": 2, + "transport_ack_nonce": "00000000-0000-4000-8000-000000000002", "message": "relay failed after all services became ready", "code": "plan_required" }) @@ -1131,10 +1161,13 @@ async fn tunnel_start_under_json_emits_pretty_success_contract_on_stdout() { .send(Message::Text( serde_json::json!({ "type": "hello", + "protocol": 4, + "transport_seq": 1, + "transport_ack_nonce": "00000000-0000-4000-8000-000000000001", "subdomain": "review-test.lpm.test", "tunnel_url": "https://review-test.lpm.test", "session_id": "session-review-test", - "plan": "free", + "plan": "pro", "base_domain": "lpm.test", "domain_kind": "random" }) @@ -1142,6 +1175,7 @@ async fn tunnel_start_under_json_emits_pretty_success_contract_on_stdout() { )) .await .unwrap(); + await_fixture_ready(&mut websocket).await; websocket.close(None).await.unwrap(); }) .await @@ -1170,19 +1204,22 @@ async fn tunnel_start_under_json_emits_pretty_success_contract_on_stdout() { ); let stdout = String::from_utf8(output.stdout).unwrap(); assert!( - stdout.starts_with("{\n \"success\""), + stdout.starts_with("{\n \"schema_version\""), "tunnel JSON must use the pretty JSON formatter, got:\n{stdout}" ); let contract: serde_json::Value = serde_json::from_str(stdout.trim()).unwrap(); insta::assert_json_snapshot!(contract, @r#" { + "schema_version": 2, + "event": "ready", + "protocol": 4, "success": true, "tunnel_url": "https://review-test.lpm.test", "domain": "review-test.lpm.test", "local_port": 5173, "local_url": "http://127.0.0.1:5173/", "session_id": "session-review-test", - "plan": "free", + "plan": "pro", "base_domain": "lpm.test", "domain_kind": "random", "session_expires_at": null, @@ -1196,6 +1233,82 @@ async fn tunnel_start_under_json_emits_pretty_success_contract_on_stdout() { "#); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn tunnel_json_ready_and_usage_notices_preserve_resource_limits_and_spending_cap() { + let relay = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let relay_url = format!("ws://{}/connect", relay.local_addr().unwrap()); + let relay_task = tokio::spawn(async move { + tokio::time::timeout(Duration::from_secs(8), async { + let (socket, _) = relay.accept().await.unwrap(); + let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); + let mut usage = serde_json::json!({ + "quota_version":2, "accepted_requests":40000, "included_requests":50000, + "busy_object_ms":36000000, "included_busy_object_ms":180000000, + "relay_messages":100000, "included_relay_messages":1000000, + "transfer_bytes":1073741824u64, "included_transfer_bytes":10737418240u64, + "connection_attempts":10, "included_connection_attempts":1000, + "overage_spend_cap_cents":600, "authorized_bundles":2, "billable_bundles":0, + "bundle_price_cents":300, "estimated_overage_cents":0, "overage_requests":0, + "overage_enabled":true, "hard_limit":false, + "period_start":"2026-10-01T00:00:00Z", "period_end":"2026-11-01T00:00:00Z" + }); + websocket.send(Message::Text(serde_json::json!({ + "type":"hello", "protocol":4, "transport_seq":1, + "transport_ack_nonce":"00000000-0000-4000-8000-000000000001", + "subdomain":"resources.lpm.test", "tunnel_url":"https://resources.lpm.test", + "session_id":"resource-session", "plan":"pro", "base_domain":"lpm.test", "domain_kind":"random", + "limits": { "quota_version":2, "public_websockets_available":false, "max_concurrent":3, + "request_rate_limit_per_minute":1000, "per_ip_rate_limit_per_minute":600, + "max_request_body_bytes":10485760, "max_custom_domains":3, "tunnel_auth_available":true }, + "usage":usage + }).to_string())).await.unwrap(); + await_fixture_ready(&mut websocket).await; + usage["accepted_requests"] = serde_json::json!(50001); + usage["overage_requests"] = serde_json::json!(1); + usage["billable_bundles"] = serde_json::json!(1); + usage["estimated_overage_cents"] = serde_json::json!(300); + websocket.send(Message::Text(serde_json::json!({ "type":"usage_notice", "transport_seq":2, + "transport_ack_nonce":"00000000-0000-4000-8000-000000000002", "usage":usage }).to_string())).await.unwrap(); + let Message::Text(ack) = websocket.next().await.unwrap().unwrap() else { panic!("missing notice receipt") }; + let ack: serde_json::Value = serde_json::from_str(&ack).unwrap(); + assert_eq!(ack["transport_seq"], 2); + assert_eq!(ack["transport_ack_nonce"], "00000000-0000-4000-8000-000000000002"); + websocket.close(None).await.unwrap(); + }).await.expect("resource contract relay timed out"); + }); + let output_task = tokio::task::spawn_blocking(move || { + let project = TempProject::empty(r#"{"name":"tunnel","version":"1.0.0"}"#); + let mut command = lpm_spawnable(&project); + command.env("LPM_TUNNEL_RELAY", relay_url).args([ + "--token", + "workflow-token", + "--json", + "tunnel", + "5173", + "--no-inspect", + ]); + command_output_with_deadline(command, Duration::from_secs(12)) + }); + let output = finish_bounded_tunnel_workflow(output_task, relay_task).await; + assert!( + output.status.success(), + "{}", + String::from_utf8_lossy(&output.stderr) + ); + let events: Vec = serde_json::Deserializer::from_slice(&output.stdout) + .into_iter() + .collect::>() + .unwrap(); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["schema_version"], 2); + assert_eq!(events[0]["event"], "ready"); + assert_eq!(events[0]["limits"]["public_websockets_available"], false); + assert_eq!(events[1]["event"], "usage_notice"); + assert_eq!(events[1]["usage"]["billable_bundles"], 1); + assert_eq!(events[1]["usage"]["overage_spend_cap_cents"], 600); + insta::assert_json_snapshot!("tunnel_resource_events", events); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tunnel_json_reports_a_retry_before_a_successful_reconnection() { let relay = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -1209,9 +1322,10 @@ async fn tunnel_json_reports_a_retry_before_a_successful_reconnection() { let (socket, _) = relay.accept().await.unwrap(); let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); websocket.send(Message::Text(serde_json::json!({ - "type":"hello", "subdomain":"retry.lpm.test", "tunnel_url":"https://retry.lpm.test", - "session_id":"session-retry", "plan":"free", "base_domain":"lpm.test", "domain_kind":"random" + "type":"hello", "protocol":4, "transport_seq":1, "transport_ack_nonce":"00000000-0000-4000-8000-000000000001", "subdomain":"retry.lpm.test", "tunnel_url":"https://retry.lpm.test", + "session_id":"session-retry", "plan":"pro", "base_domain":"lpm.test", "domain_kind":"random" }).to_string())).await.unwrap(); + await_fixture_ready(&mut websocket).await; websocket.close(None).await.unwrap(); }).await.expect("retry relay timed out"); }); @@ -1255,7 +1369,7 @@ async fn tunnel_json_reports_a_retry_before_a_successful_reconnection() { assert_eq!(events[1]["success"], true); insta::assert_json_snapshot!(events[0], {".error" => "[retry delay and relay rejection]"}, @r#" { - "schema_version": 1, + "schema_version": 2, "success": false, "event": "retry", "error_code": "tunnel_retry", @@ -1277,10 +1391,13 @@ async fn tunnel_start_warns_that_capture_history_persists_sensitive_request_data .send(Message::Text( serde_json::json!({ "type": "hello", + "protocol": 4, + "transport_seq": 1, + "transport_ack_nonce": "00000000-0000-4000-8000-000000000001", "subdomain": "capture-warning.lpm.test", "tunnel_url": "https://capture-warning.lpm.test", "session_id": "session-capture-warning", - "plan": "free", + "plan": "pro", "base_domain": "lpm.test", "domain_kind": "random" }) @@ -1288,6 +1405,7 @@ async fn tunnel_start_warns_that_capture_history_persists_sensitive_request_data )) .await .unwrap(); + await_fixture_ready(&mut websocket).await; websocket.close(None).await.unwrap(); }) .await diff --git a/tests/workflows/tests/dev_tunnel/auto_ack.rs b/tests/workflows/tests/dev_tunnel/auto_ack.rs index 5c0bf9bae..5a06dc3a9 100644 --- a/tests/workflows/tests/dev_tunnel/auto_ack.rs +++ b/tests/workflows/tests/dev_tunnel/auto_ack.rs @@ -1,5 +1,4 @@ use super::*; -use futures_util::StreamExt; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn auto_ack_is_on_disk_before_provider_receives_success() { @@ -18,18 +17,27 @@ async fn auto_ack_is_on_disk_before_provider_receives_success() { let (socket, _) = relay.accept().await.unwrap(); let mut websocket = tokio_tungstenite::accept_async(socket).await.unwrap(); websocket.send(Message::Text(serde_json::json!({ - "type":"hello", "subdomain":"capture.lpm.test", "tunnel_url":"https://capture.lpm.test", - "session_id":"durable-session", "plan":"free", "base_domain":"lpm.test", "domain_kind":"random" + "type":"hello", "protocol":4, "transport_seq":1, "transport_ack_nonce":"00000000-0000-4000-8000-000000000001", "subdomain":"capture.lpm.test", "tunnel_url":"https://capture.lpm.test", + "session_id":"durable-session", "plan":"pro", "base_domain":"lpm.test", "domain_kind":"random" }).to_string())).await.unwrap(); + await_fixture_ready(&mut websocket).await; websocket.send(Message::Text(serde_json::json!({ - "type":"http_request", "id":"durable-auto-ack", "method":"POST", "url":"/webhook", + "type":"http_request", "transport_seq":2, "transport_ack_nonce":"00000000-0000-4000-8000-000000000002", "id":"durable-auto-ack", "method":"POST", "url":"/webhook", "headers":{"content-type":"application/json"}, "body":"e30=" }).to_string())).await.unwrap(); + let mut request_acked = false; loop { let message = websocket.next().await.unwrap().unwrap(); let Message::Text(text) = message else { continue }; let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + if response["type"] == "transport_ack" { + assert_eq!(response["transport_seq"], 2); + assert_eq!(response["transport_ack_nonce"], "00000000-0000-4000-8000-000000000002"); + request_acked = true; + continue; + } if response["id"] != "durable-auto-ack" { continue } + assert!(request_acked, "request receipt precedes the provider response"); assert_eq!(response["status"], 200); // Read through a new connection before shutdown can flush queued captures. let db = lpm_inspect::db::InspectorDb::open(&project_path).unwrap(); diff --git a/tests/workflows/tests/snapshots/dev_tunnel__tunnel_resource_events.snap b/tests/workflows/tests/snapshots/dev_tunnel__tunnel_resource_events.snap new file mode 100644 index 000000000..14ec94eda --- /dev/null +++ b/tests/workflows/tests/snapshots/dev_tunnel__tunnel_resource_events.snap @@ -0,0 +1,91 @@ +--- +source: tests/workflows/tests/dev_tunnel.rs +expression: events +--- +[ + { + "schema_version": 2, + "event": "ready", + "protocol": 4, + "success": true, + "tunnel_url": "https://resources.lpm.test", + "domain": "resources.lpm.test", + "local_port": 5173, + "local_url": "http://127.0.0.1:5173/", + "session_id": "resource-session", + "plan": "pro", + "base_domain": "lpm.test", + "domain_kind": "random", + "session_expires_at": null, + "session_max_ms": null, + "limits": { + "quota_version": 2, + "public_websockets_available": false, + "max_concurrent": 3, + "request_rate_limit_per_minute": 1000, + "per_ip_rate_limit_per_minute": 600, + "max_request_body_bytes": 10485760, + "max_custom_domains": 3, + "tunnel_auth_available": true + }, + "usage": { + "quota_version": 2, + "busy_object_ms": 36000000, + "included_busy_object_ms": 180000000, + "relay_messages": 100000, + "included_relay_messages": 1000000, + "transfer_bytes": 1073741824, + "included_transfer_bytes": 10737418240, + "connection_attempts": 10, + "included_connection_attempts": 1000, + "overage_spend_cap_cents": 600, + "authorized_bundles": 2, + "billable_bundles": 0, + "bundle_price_cents": 300, + "accepted_requests": 40000, + "included_requests": 50000, + "overage_requests": 0, + "overage_enabled": true, + "hard_limit": false, + "period_start": "2026-10-01T00:00:00Z", + "period_end": "2026-11-01T00:00:00Z", + "overage_unit_requests": null, + "overage_unit_price_cents": null, + "estimated_overage_cents": 0.0 + }, + "tunnel_auth": null, + "inspector_url": null, + "auto_ack": false + }, + { + "schema_version": 2, + "success": true, + "event": "usage_notice", + "quota_version": 2, + "usage": { + "quota_version": 2, + "busy_object_ms": 36000000, + "included_busy_object_ms": 180000000, + "relay_messages": 100000, + "included_relay_messages": 1000000, + "transfer_bytes": 1073741824, + "included_transfer_bytes": 10737418240, + "connection_attempts": 10, + "included_connection_attempts": 1000, + "overage_spend_cap_cents": 600, + "authorized_bundles": 2, + "billable_bundles": 1, + "bundle_price_cents": 300, + "accepted_requests": 50001, + "included_requests": 50000, + "overage_requests": 1, + "overage_enabled": true, + "hard_limit": false, + "period_start": "2026-10-01T00:00:00Z", + "period_end": "2026-11-01T00:00:00Z", + "overage_unit_requests": null, + "overage_unit_price_cents": null, + "estimated_overage_cents": 300.0 + } + } +]