diff --git a/cmux-tui/crates/chatmux-relay/src/journal_forwarder.rs b/cmux-tui/crates/chatmux-relay/src/journal_forwarder.rs index 3b534b3b645f..9b4551770e80 100644 --- a/cmux-tui/crates/chatmux-relay/src/journal_forwarder.rs +++ b/cmux-tui/crates/chatmux-relay/src/journal_forwarder.rs @@ -987,6 +987,13 @@ fn enqueue_pending( } let total = pool.pending.iter().map(|entry| entry.records.len()).sum::(); if total < MAX_BATCH_RECORDS { + // During an in-flight POST the completion path re-arms the debounce + // for any leftover pending records (flush_cycle reads the total under + // the same lock that clears `flushing`), so waking here would only + // store a spurious permit — the exact wake the deferral pin forbids. + if pool.flushing { + return; + } drop(pool); shared.flush_wake.notify_one(); return; @@ -1909,7 +1916,7 @@ mod tests { #[cfg(unix)] #[tokio::test] - async fn pooled_arm_wakes_are_coalesced_while_the_flusher_is_busy() { + async fn pooled_arm_wakes_defer_while_the_flusher_is_busy_and_coalesce_after() { let (root, path) = cursor_test_path("pool-arm-coalesce").await; let (shared, wake) = test_shared(String::from("http://127.0.0.1:9"), path); let generation = Some(String::from("gen_a")); @@ -1927,9 +1934,32 @@ mod tests { ); } + // While a POST is in flight, below-threshold arms never wake: the + // completion path re-arms the debounce from the leftover total under + // the same lock that clears `flushing` (the #11034 deferral rule — + // this pin and the pooled_threshold deferral pin share it). + assert!( + tokio::time::timeout(Duration::from_millis(10), wake.notified()).await.is_err(), + "arms during an in-flight POST must defer to the completion re-arm" + ); + + // After the POST settles, arm wakes flow again — and coalesce. + { + let mut pool = shared.pool.lock().expect("lock pool"); + pool.flushing = false; + } + for sequence in 4..=6 { + enqueue_pending( + &shared, + "alpha", + "alpha", + &generation, + record("gen_a", &sequence.to_string()), + ); + } tokio::time::timeout(Duration::from_millis(100), wake.notified()) .await - .expect("arm request must wake the flusher"); + .expect("an idle-flusher arm request must wake the flusher"); assert!( tokio::time::timeout(Duration::from_millis(10), wake.notified()).await.is_err(), "repeated arm requests must use one pending wake" diff --git a/cmux-tui/crates/chatmux-relay/src/lib.rs b/cmux-tui/crates/chatmux-relay/src/lib.rs index a78d379e5948..e4999864eabb 100644 --- a/cmux-tui/crates/chatmux-relay/src/lib.rs +++ b/cmux-tui/crates/chatmux-relay/src/lib.rs @@ -22,6 +22,7 @@ pub mod pty_deps; pub mod relay_wire; pub mod session; pub mod trust; +pub mod tunnel_terminal; pub mod watch; pub mod wire; pub mod workspace; diff --git a/cmux-tui/crates/chatmux-relay/src/pty.rs b/cmux-tui/crates/chatmux-relay/src/pty.rs index a8960ee632f8..14ade2264861 100644 --- a/cmux-tui/crates/chatmux-relay/src/pty.rs +++ b/cmux-tui/crates/chatmux-relay/src/pty.rs @@ -55,6 +55,17 @@ const MAX_ENUM_SURFACES: usize = 8; const RAW_ATTACH_BACKLOG_CAP: usize = 1024 * 1024; const PTY_INPUT_B64_CAP: usize = 4 * 1024 * 1024; +/// Random lowercase-hex identity for transports and tunnel attachments. +pub fn random_hex(bytes: usize) -> String { + let mut buffer = vec![0_u8; bytes]; + let _ = getrandom::fill(&mut buffer); + let mut out = String::with_capacity(bytes * 2); + for byte in buffer { + out.push_str(&format!("{byte:02x}")); + } + out +} + pub fn session_name_ok(name: &str) -> bool { let invalid = name.is_empty() || matches!(name, "." | "..") @@ -259,6 +270,13 @@ pub struct FrameContext { pub trust: String, pub local_roots: Option>, pub owner_user_id: Option, + /// Identity of the transport this frame arrived on. The PtyManager is + /// shared between the relay WebSocket and the managed tunnel listener; + /// an attachment may only be written to, resized, flow-controlled, or + /// closed by the transport that opened it, and a dropped transport + /// detaches only its own attachments. `None` preserves the legacy + /// owns-everything behavior for callers that own the whole manager. + pub transport_id: Option, } #[derive(Clone)] @@ -309,6 +327,8 @@ struct Attachment { /// kill a viewer PTY) — never kills a shared session. control: Arc, actor_id: String, + /// Transport that opened this attachment (see FrameContext::transport_id). + transport_id: Option, } struct Inner { @@ -319,7 +339,8 @@ struct Inner { scrollback_limit: usize, output_cap: u64, attachments: Mutex>, - opening_ids: Mutex>, + /// ptyId -> transport that reserved it (None = legacy whole-manager owner). + opening_ids: Mutex>>, cancelled_openings: Mutex>, shell_sessions: Mutex>>, shell_starting: Mutex>>, @@ -370,7 +391,7 @@ impl PtyManager { scrollback_limit: SCROLLBACK_LIMIT, output_cap: OUTPUT_BUFFER_CAP, attachments: Mutex::new(HashMap::new()), - opening_ids: Mutex::new(std::collections::HashSet::new()), + opening_ids: Mutex::new(HashMap::new()), cancelled_openings: Mutex::new(std::collections::HashSet::new()), shell_sessions: Mutex::new(HashMap::new()), shell_starting: Mutex::new(HashMap::new()), @@ -396,7 +417,7 @@ impl PtyManager { scrollback_limit, output_cap, attachments: Mutex::new(HashMap::new()), - opening_ids: Mutex::new(std::collections::HashSet::new()), + opening_ids: Mutex::new(HashMap::new()), cancelled_openings: Mutex::new(std::collections::HashSet::new()), shell_sessions: Mutex::new(HashMap::new()), shell_starting: Mutex::new(HashMap::new()), @@ -418,6 +439,9 @@ impl PtyManager { "pty_open" => self.inner.clone().open(frame, context).await, "pty_input" => { let Some(pty_id) = frame.get("ptyId").and_then(Value::as_str) else { return }; + if !self.inner.transport_owns(pty_id, context.transport_id.as_deref()) { + return; + } let Some(data) = frame .get("dataB64") .and_then(Value::as_str) @@ -432,6 +456,9 @@ impl PtyManager { } "pty_resize" => { let Some(pty_id) = frame.get("ptyId").and_then(Value::as_str) else { return }; + if !self.inner.transport_owns(pty_id, context.transport_id.as_deref()) { + return; + } let (Some(cols), Some(rows)) = (clamp_dim(frame.get("cols")), clamp_dim(frame.get("rows"))) else { @@ -443,6 +470,9 @@ impl PtyManager { } "pty_flow" => { let Some(pty_id) = frame.get("ptyId").and_then(Value::as_str) else { return }; + if !self.inner.transport_owns(pty_id, context.transport_id.as_deref()) { + return; + } let pause = frame.get("pause").and_then(Value::as_bool).unwrap_or(false); if let Some(attachment) = self.inner.authorize(pty_id, context, "flow") { if pause { @@ -454,6 +484,9 @@ impl PtyManager { } "pty_close" => { let Some(pty_id) = frame.get("ptyId").and_then(Value::as_str) else { return }; + if !self.inner.transport_owns(pty_id, context.transport_id.as_deref()) { + return; + } self.inner.close_authorized(pty_id, context); } "surface_list" => self.inner.clone().list_surfaces(frame, context).await, @@ -461,10 +494,52 @@ impl PtyManager { } } + /// True while `pty_id` has a live attachment. The tunnel listener uses + /// this after a pty_error reply to tell a fatal refusal (attachment gone, + /// connection ends) from a non-fatal one (oversized input, stream lives). + pub fn has_attachment(&self, pty_id: &str) -> bool { + self.inner.attachments.lock().expect("attach lock").contains_key(pty_id) + } + + /// Live attachment count (viewers, not sessions). Diagnostics and tests. + pub fn attachment_count(&self) -> usize { + self.inner.attachments.lock().expect("attach lock").len() + } + /// The relay socket dropped: release every attachment (sessions live on). + /// Callers that own the whole manager only; a per-connection transport + /// must use `detach_transport` so it cannot detach attachments the + /// managed tunnel listener (or another socket) owns. pub fn detach_all(&self) { - let ids: Vec = - self.inner.attachments.lock().expect("attach lock").keys().cloned().collect(); + self.detach_matching(|_| true); + } + + /// One transport dropped: release only its attachments and cancel only + /// its in-flight opens. Sessions live on either way (docs/TERMINAL.md). + pub fn detach_transport(&self, transport_id: &str) { + self.detach_matching(|owner| owner == Some(transport_id)); + } + + fn detach_matching(&self, owns: impl Fn(Option<&str>) -> bool) { + // Openings first: close() records cancellation for a reserved id, so + // a late open cannot install an attachment after its transport died. + let mut ids: Vec = { + let opening = self.inner.opening_ids.lock().expect("opening lock"); + opening + .iter() + .filter(|(_, owner)| owns(owner.as_deref())) + .map(|(id, _)| id.clone()) + .collect() + }; + { + let attachments = self.inner.attachments.lock().expect("attach lock"); + ids.extend( + attachments + .iter() + .filter(|(_, attachment)| owns(attachment.transport_id.as_deref())) + .map(|(id, _)| id.clone()), + ); + } for id in ids { self.inner.close(&id); } @@ -507,7 +582,7 @@ impl Inner { let reservation_result = { let mut opening = self.opening_ids.lock().expect("opening lock"); let attached = self.attachments.lock().expect("attach lock").contains_key(&pty_id); - if attached || opening.contains(&pty_id) { + if attached || opening.contains_key(&pty_id) { Err(("bad_request", "ptyId is already attached".to_owned())) } else if self.attachments.lock().expect("attach lock").len() + opening.len() >= self.max_ptys @@ -517,7 +592,7 @@ impl Inner { format!("this relay caps concurrent terminals at {}", self.max_ptys), )) } else { - opening.insert(pty_id.clone()); + opening.insert(pty_id.clone(), context.transport_id.clone()); Ok(()) } }; @@ -689,6 +764,7 @@ impl Inner { closing: opened.closing, control: opened.control, actor_id: actor.to_owned(), + transport_id: context.transport_id.clone(), }, ); if let Some(previous) = previous { @@ -795,7 +871,7 @@ impl Inner { // Match `open`'s lock order. If opening still owns the reservation, // record cancellation and let it dispose the newly opened PTY. let opening = self.opening_ids.lock().expect("opening lock"); - if opening.contains(pty_id) { + if opening.contains_key(pty_id) { self.cancelled_openings .lock() .expect("cancelled openings lock") @@ -810,6 +886,20 @@ impl Inner { } } + /// Frame-level transport fence. Unknown ids retain the protocol's silent + /// no-op behavior; once an id is reserved or attached, a different + /// transport may not act on it. A `None` caller owns everything (legacy). + fn transport_owns(&self, pty_id: &str, transport_id: Option<&str>) -> bool { + let Some(transport_id) = transport_id else { return true }; + if let Some(attachment) = self.attachments.lock().expect("attach lock").get(pty_id) { + return attachment.transport_id.as_deref() == Some(transport_id); + } + if let Some(owner) = self.opening_ids.lock().expect("opening lock").get(pty_id) { + return owner.as_deref() == Some(transport_id); + } + true + } + fn authorize(&self, pty_id: &str, context: &FrameContext, action: &str) -> Option { let auth = self.auth.lock().expect("auth lock").clone()?; self.authorize_snapshot(pty_id, &auth, context, action) @@ -2151,9 +2241,38 @@ mod tests { trust: trust.to_owned(), local_roots: None, owner_user_id: owner, + transport_id: None, } } + fn context_with_transport( + &self, + trust: &str, + owner: Option, + transport_id: Option<&str>, + ) -> FrameContext { + let mut context = self.context(trust, owner); + context.transport_id = transport_id.map(str::to_owned); + context + } + + async fn open_with_transport(&self, pty_id: &str, session: &str, transport_id: &str) { + let frame = serde_json::json!({ + "version": 4, + "type": "pty_open", + "ptyId": pty_id, + "session": session, + "cols": 80, + "rows": 24, + "actorId": "user_owner", + "trust": "supervised", + "allowedRoots": Value::Null, + }); + let context = + self.context_with_transport("supervised", self.owner.clone(), Some(transport_id)); + self.manager.handle_frame(&frame, &context).await; + } + async fn open( &self, pty_id: &str, @@ -2650,6 +2769,42 @@ mod tests { assert_eq!(h.spawned().len(), 1); } + #[tokio::test] + async fn a_foreign_transport_cannot_write_resize_or_close_an_owned_pty() { + let h = harness(None, None); + h.open_with_transport("p1", "main", "transport-a").await; + let foreign = h.context_with_transport("supervised", h.owner.clone(), Some("transport-b")); + let input = serde_json::json!({ + "version": 4, + "type": "pty_input", + "ptyId": "p1", + "dataB64": b64("stolen"), + }); + h.manager.handle_frame(&input, &foreign).await; + assert!(h.spawned()[0].state.lock().unwrap().written.is_empty()); + let close = serde_json::json!({ "version": 4, "type": "pty_close", "ptyId": "p1" }); + h.manager.handle_frame(&close, &foreign).await; + assert!(h.manager.has_attachment("p1"), "a foreign close must be a silent no-op"); + let owner = h.context_with_transport("supervised", h.owner.clone(), Some("transport-a")); + h.manager.handle_frame(&input, &owner).await; + assert_eq!(h.spawned()[0].written_string(0), "stolen"); + // A caller with no transport identity owns the whole manager (legacy). + h.manager.handle_frame(&close, &h.context("supervised", h.owner.clone())).await; + assert!(!h.manager.has_attachment("p1")); + } + + #[tokio::test] + async fn detach_transport_releases_only_that_transports_attachments() { + let h = harness(None, None); + h.open_with_transport("p-relay", "relay-side", "transport-relay").await; + h.open_with_transport("p-tunnel", "tunnel-side", "transport-tunnel").await; + h.manager.detach_transport("transport-relay"); + assert!(!h.manager.has_attachment("p-relay"), "the relay transport's viewer must detach"); + assert!(h.manager.has_attachment("p-tunnel"), "the tunnel viewer must survive"); + h.manager.detach_all(); + assert!(!h.manager.has_attachment("p-tunnel")); + } + #[test] fn pty_env_scrubs_secrets_but_keeps_a_real_term() { let home = TestDirectory::new("env"); diff --git a/cmux-tui/crates/chatmux-relay/src/session.rs b/cmux-tui/crates/chatmux-relay/src/session.rs index 654a12044852..e5723b74ccaa 100644 --- a/cmux-tui/crates/chatmux-relay/src/session.rs +++ b/cmux-tui/crates/chatmux-relay/src/session.rs @@ -414,6 +414,26 @@ pub async fn stay_online( cancellation: CancellationToken, ) -> Result<(), RelayError> { let runtime = SessionRuntime::with_roots(config.allowed_roots.clone()); + // Tunnel-direct terminal data plane: serve terminals to spliced tunnel + // connections on loopback 127.0.0.1:9776. Managed sandboxes ONLY — this + // branch is the gate; paired human machines never start the listener. + // Best-effort: a failed bind degrades to the relay-socket terminal path. + #[cfg(unix)] + if state.managed { + match crate::tunnel_terminal::start_tunnel_terminal_listener( + Arc::clone(&runtime.pty), + cancellation.child_token(), + crate::tunnel_terminal::TUNNEL_TERMINAL_HOST, + crate::tunnel_terminal::TUNNEL_TERMINAL_PORT, + ) + .await + { + Ok(_) => eprintln!("Tunnel terminal listener is up on loopback."), + Err(error) => eprintln!( + "Tunnel terminal listener bind failed: {error}. Terminals stay on the relay socket path." + ), + } + } let mut attempt: u32 = 0; loop { if cancellation.is_cancelled() { @@ -487,7 +507,12 @@ fn unsupported_platform_pty_reply(frame_type: &str, raw: &Value) -> Option, auth: &AuthSnapshot) -> FrameContext { +fn make_context( + out: &OutboundSink, + pending: &Arc, + auth: &AuthSnapshot, + transport_id: &str, +) -> FrameContext { let sender = out.clone(); let pending_send = Arc::clone(pending); let pending_probe = Arc::clone(pending); @@ -518,6 +543,7 @@ fn make_context(out: &OutboundSink, pending: &Arc, auth: &AuthSnapsho trust: auth.trust.clone(), local_roots: auth.roots.clone(), owner_user_id: auth.owner.clone(), + transport_id: Some(transport_id.to_owned()), } } @@ -577,6 +603,12 @@ async fn relay_session( let connection_cancellation = cancellation.child_token(); let mut connection_tasks = JoinSet::new(); + // The PtyManager is shared with the managed tunnel listener. Every PTY + // this socket opens carries this connection's identity, so closing or + // reconnecting the socket cannot detach an independent tunnel attachment. + #[cfg(unix)] + let transport_id = format!("relay-{}", crate::pty::random_hex(16)); + // Ordered PTY frame dispatch on its own task so a slow open (daemon // spawn) never stalls heartbeats or other frames. #[cfg(unix)] @@ -591,6 +623,7 @@ async fn relay_session( let pending = Arc::clone(&pending); let auth = Arc::clone(&auth); let cancellation = connection_cancellation.clone(); + let transport = transport_id.clone(); connection_tasks.spawn(async move { loop { let frame = tokio::select! { @@ -602,7 +635,7 @@ async fn relay_session( } }; let snapshot = auth.lock().expect("auth lock").clone(); - let context = make_context(&out, &pending, &snapshot); + let context = make_context(&out, &pending, &snapshot, &transport); tokio::select! { biased; _ = cancellation.cancelled() => break, @@ -1034,7 +1067,8 @@ async fn relay_session( // bounded work queue is saturated. The manager close path is // synchronous and short, so this cannot create an unbounded wait. let snapshot = auth_direct.lock().expect("auth lock").clone(); - let context = make_context(&out_tx, &pending, &snapshot); + let context = + make_context(&out_tx, &pending, &snapshot, &transport_id); tokio::select! { biased; _ = cancellation.cancelled() => break Ok(connected), @@ -1155,9 +1189,11 @@ async fn relay_session( ); } - // Attachments die with the socket; sessions persist (docs/TERMINAL.md). + // This socket's attachments die with it; sessions persist, and the + // managed tunnel listener's attachments are another transport's — a + // reconnect must never detach them (docs/TERMINAL.md). #[cfg(unix)] - runtime.pty.detach_all(); + runtime.pty.detach_transport(&transport_id); result } diff --git a/cmux-tui/crates/chatmux-relay/src/tunnel_terminal.rs b/cmux-tui/crates/chatmux-relay/src/tunnel_terminal.rs new file mode 100644 index 000000000000..0df16cafbf67 --- /dev/null +++ b/cmux-tui/crates/chatmux-relay/src/tunnel_terminal.rs @@ -0,0 +1,1081 @@ +//! Tunnel-direct terminal listener (managed sandboxes only). +//! +//! The chatmux tunnel gateway's `/tcp` endpoint splices a browser WebSocket +//! onto a raw byte stream that dials 127.0.0.1: inside this sandbox. +//! This module is that port: a loopback TCP server that serves terminals +//! through the SAME PtyManager the relay-socket path uses (whole-session +//! attach, W86 single-terminal attach, scrollback replay, sizing, +//! backpressure caps — all shared, none duplicated). Rust port of the Node +//! reference `packages/relay/bin/tunnel-terminal.mjs` in chatmux; the wire +//! shapes are pinned there (docs/TERMINAL.md) and by the tests below. +//! +//! FRAMING — the splice is a raw byte pipe (the browser's WebSocket message +//! boundaries are lost in transit), so every message is length-prefixed: +//! +//! u32 big-endian payloadLength | u8 kind | payload[payloadLength] +//! +//! kind 0 = UTF-8 JSON control frame +//! kind 1 = raw PTY bytes (client->server stdin, server->client output) +//! +//! payloadLength is bounded by MAX_TUNNEL_FRAME_BYTES (1 MiB). An oversized +//! length, an unknown kind, or an undecodable control frame is a protocol +//! error: the server hard-closes the connection (a desynced length stream +//! can never be re-synchronized). +//! +//! CONTROL FRAMES mirror the browser terminal wire, minus the auth step: +//! +//! c->s first frame {t:"open", session?, surface?, cols, rows} +//! s->c {t:"opened", session, surface?, created, cols, rows} +//! or {t:"error", code, message?} +//! then kind-1 byte flow both ways; later control frames: +//! c->s {t:"resize", cols, rows}, {t:"detach"} +//! s->c {t:"exit", code}, {t:"error", code, message?} +//! +//! THREAT MODEL — there is deliberately NO auth frame. The tunnel gateway +//! already enforced a capability token minted by the Worker (policyd client +//! token bound to this sandbox + port) before splicing the connection, and +//! this listener binds loopback only, inside a sandbox where local code can +//! already attach terminals through the cmux CLI. Paired human machines +//! never run this listener: it starts from the managed branch only. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::time::Duration; + +use base64::Engine as _; +use base64::engine::general_purpose::STANDARD as BASE64; +use serde_json::{Value, json}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use crate::pty::{ + FrameContext, PTY_PROTOCOL_VERSION, PtyManager, random_hex, session_name_ok, surface_ref_ok, +}; + +/// Loopback port the gateway's spliced streams dial. The chatmux Worker +/// mirrors this value in apps/backend/src/tunnel/terminal-ticket-route.ts — +/// KEEP IN SYNC. +pub const TUNNEL_TERMINAL_PORT: u16 = 9776; +/// Loopback ONLY: the tunnel backhaul dials local ports; nothing else may. +pub const TUNNEL_TERMINAL_HOST: &str = "127.0.0.1"; +/// Hard per-frame payload bound (matches the relay JSON frame maximum). +pub const MAX_TUNNEL_FRAME_BYTES: usize = 1_048_576; +pub const FRAME_KIND_CONTROL: u8 = 0; +pub const FRAME_KIND_PTY: u8 = 1; +/// u32 length + u8 kind. +const HEADER_BYTES: usize = 5; +/// The opened (or refused) reply must arrive within this budget. +const OPEN_TIMEOUT: Duration = Duration::from_secs(10); +/// Writer flow control: pause the PTY source above the high-water mark of +/// bytes queued toward the socket, resume below the low-water mark. The +/// manager's own 1 MiB output cap stays the hard boundary above this. +const FLOW_PAUSE_BYTES: u64 = 262_144; +const FLOW_RESUME_BYTES: u64 = 32_768; + +/// Relay pty_error codes -> browser wire codes. Mirrors the Worker's +/// browserErrorCode map (apps/backend/src/terminal/relay-pty.ts). KEEP IN +/// SYNC; `wire_error_codes_match_the_worker_map` pins the shape here. +pub fn wire_error_code(code: &str) -> &'static str { + match code { + "bad_request" => "bad_request", + "trust_refused" => "trust_blocked", + "session_limit" => "session_limit", + "terminal_gone" => "terminal_gone", + "overflow" => "overflow", + "trust_revoked" => "trust_revoked", + "busy" => "busy", + _ => "failed", + } +} + +// --------------------------------------------------------------------------- +// Framing codec +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +pub struct TunnelFrame { + pub kind: u8, + pub payload: Vec, +} + +/// Encode one frame: u32be payload length, u8 kind, payload. +pub fn encode_tunnel_frame(kind: u8, payload: &[u8]) -> Vec { + debug_assert!(payload.len() <= MAX_TUNNEL_FRAME_BYTES); + let mut frame = Vec::with_capacity(HEADER_BYTES + payload.len()); + frame.extend_from_slice(&u32::try_from(payload.len()).unwrap_or(0).to_be_bytes()); + frame.push(kind); + frame.extend_from_slice(payload); + frame +} + +pub fn encode_control_frame(frame: &Value) -> Vec { + encode_tunnel_frame(FRAME_KIND_CONTROL, frame.to_string().as_bytes()) +} + +pub fn encode_pty_frame(bytes: &[u8]) -> Vec { + encode_tunnel_frame(FRAME_KIND_PTY, bytes) +} + +/// Incremental frame decoder. After one failure the decoder is poisoned: a +/// length-prefixed stream that desynced once can never be trusted again, so +/// the caller must close the connection. +pub struct TunnelFrameDecoder { + buffer: Vec, + failed: bool, + max_frame_bytes: usize, +} + +impl TunnelFrameDecoder { + pub fn new(max_frame_bytes: usize) -> TunnelFrameDecoder { + TunnelFrameDecoder { + buffer: Vec::new(), + failed: false, + max_frame_bytes: max_frame_bytes.clamp(1, MAX_TUNNEL_FRAME_BYTES), + } + } + + pub fn push(&mut self, chunk: &[u8]) -> Result, &'static str> { + if self.failed { + return Err("decoder_poisoned"); + } + self.buffer.extend_from_slice(chunk); + let mut frames = Vec::new(); + while self.buffer.len() >= HEADER_BYTES { + let length = u32::from_be_bytes([ + self.buffer[0], + self.buffer[1], + self.buffer[2], + self.buffer[3], + ]) as usize; + let kind = self.buffer[4]; + if length > self.max_frame_bytes { + self.failed = true; + return Err("frame_too_large"); + } + if kind != FRAME_KIND_CONTROL && kind != FRAME_KIND_PTY { + self.failed = true; + return Err("unknown_frame_kind"); + } + if self.buffer.len() < HEADER_BYTES + length { + break; + } + let payload = self.buffer[HEADER_BYTES..HEADER_BYTES + length].to_vec(); + self.buffer.drain(..HEADER_BYTES + length); + frames.push(TunnelFrame { kind, payload }); + } + Ok(frames) + } +} + +// --------------------------------------------------------------------------- +// Client control frame parsing (mirror of the browser wire minus `auth`) +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +pub enum ClientFrame { + Open { session: Option, surface: Option, cols: u16, rows: u16 }, + Resize { cols: u16, rows: u16 }, + Detach, +} + +fn valid_dims(raw: &Value) -> Option<(u16, u16)> { + let dim = |key: &str| { + raw.get(key) + .and_then(Value::as_u64) + .filter(|value| (1..=10_000).contains(value)) + .and_then(|value| u16::try_from(value).ok()) + }; + Some((dim("cols")?, dim("rows")?)) +} + +/// None = malformed (a protocol error; the caller closes). +pub fn parse_tunnel_client_frame(payload: &[u8]) -> Option { + let raw: Value = serde_json::from_slice(payload).ok()?; + if !raw.is_object() { + return None; + } + match raw.get("t").and_then(Value::as_str) { + Some("open") => { + let (cols, rows) = valid_dims(&raw)?; + let session = match raw.get("session") { + None => None, + Some(value) => { + let name = value.as_str().filter(|name| session_name_ok(name))?; + Some(name.to_owned()) + } + }; + let surface = match raw.get("surface") { + None => None, + Some(value) => { + // A surface ref without a session has nothing to resolve + // against. + session.as_ref()?; + let surface = value.as_str().filter(|surface| surface_ref_ok(surface))?; + Some(surface.to_owned()) + } + }; + Some(ClientFrame::Open { session, surface, cols, rows }) + } + Some("resize") => { + let (cols, rows) = valid_dims(&raw)?; + Some(ClientFrame::Resize { cols, rows }) + } + Some("detach") => Some(ClientFrame::Detach), + _ => None, + } +} + +/// Server-generated session names: same alphabet and prefix the Worker route +/// uses, so pickers and process tables read consistently. +pub fn generate_session_name() -> String { + const ALPHABET: &[u8] = b"abcdefghjkmnpqrstuvwxyz23456789"; + let mut bytes = [0_u8; 4]; + let _ = getrandom::fill(&mut bytes); + let suffix: String = + bytes.iter().map(|byte| ALPHABET[*byte as usize % ALPHABET.len()] as char).collect(); + format!("web-{suffix}") +} + +// --------------------------------------------------------------------------- +// One spliced connection = one terminal attachment +// --------------------------------------------------------------------------- + +enum WriterMessage { + Frame(Vec), + /// Flush what is queued, then close the write half. + End, +} + +/// State shared between the reader task, the writer task, and the manager's +/// synchronous reply sink. +struct Connection { + pty_id: String, + manager: Arc, + writer_tx: mpsc::UnboundedSender, + /// pty_flow requests from the writer's water marks (true = pause). + flow_tx: mpsc::UnboundedSender, + /// Bytes queued toward the socket and not yet written. + pending_out: AtomicU64, + paused: AtomicBool, + /// The open frame was forwarded to the manager (pty_close owed on exit). + open_sent: AtomicBool, + /// The manager answered pty_opened (clears the open deadline). + opened_seen: AtomicBool, + finished: AtomicBool, + done: CancellationToken, +} + +impl Connection { + fn send_control(&self, frame: &Value) { + self.enqueue(WriterMessage::Frame(encode_control_frame(frame))); + } + + fn enqueue(&self, message: WriterMessage) { + if self.finished.load(Ordering::SeqCst) { + return; + } + if let WriterMessage::Frame(frame) = &message { + self.pending_out.fetch_add(frame.len() as u64, Ordering::SeqCst); + } + let _ = self.writer_tx.send(message); + } + + /// Idempotent shutdown: flush queued frames, close the socket, and let + /// the reader task settle the owed pty_close (detach, never kill — the + /// session lives on for a later re-attach, the same rule a dropped + /// relay-socket viewer follows). + fn finish(&self) { + if self.finished.swap(true, Ordering::SeqCst) { + return; + } + let _ = self.writer_tx.send(WriterMessage::End); + self.done.cancel(); + } + + fn protocol_error(&self, code: &str, message: &str) { + if self.finished.load(Ordering::SeqCst) { + return; + } + self.send_control(&json!({ "t": "error", "code": code, "message": message })); + self.finish(); + } + + /// The manager's reply sink (FrameContext::send). Synchronous: enqueue + /// only, never block. + fn on_manager_frame(&self, frame: &Value) { + if self.finished.load(Ordering::SeqCst) { + return; + } + if frame.get("ptyId").and_then(Value::as_str) != Some(self.pty_id.as_str()) { + return; + } + match frame.get("type").and_then(Value::as_str) { + Some("pty_opened") => { + self.opened_seen.store(true, Ordering::SeqCst); + let mut opened = json!({ + "t": "opened", + "session": frame.get("session").cloned().unwrap_or(Value::Null), + "created": frame.get("created").and_then(Value::as_bool) == Some(true), + "cols": frame.get("cols").cloned().unwrap_or(Value::Null), + "rows": frame.get("rows").cloned().unwrap_or(Value::Null), + }); + if let Some(surface) = frame.get("surface").and_then(Value::as_str) { + opened["surface"] = Value::from(surface); + } + self.send_control(&opened); + } + Some("pty_output") => { + let Some(bytes) = frame + .get("dataB64") + .and_then(Value::as_str) + .and_then(|b64| BASE64.decode(b64).ok()) + .filter(|bytes| !bytes.is_empty()) + else { + return; + }; + self.enqueue(WriterMessage::Frame(encode_pty_frame(&bytes))); + // Socket-side congestion: pause the source through the + // manager's own flow verb; the writer resumes it below the + // low-water mark. + if self.pending_out.load(Ordering::SeqCst) > FLOW_PAUSE_BYTES + && !self.paused.swap(true, Ordering::SeqCst) + { + let _ = self.flow_tx.send(true); + } + } + Some("pty_exit") => { + let code = frame.get("code").and_then(Value::as_i64).unwrap_or(0); + self.send_control(&json!({ "t": "exit", "code": code })); + self.finish(); + } + Some("pty_error") => { + let code = + wire_error_code(frame.get("code").and_then(Value::as_str).unwrap_or("failed")); + let mut error = json!({ "t": "error", "code": code }); + if let Some(message) = frame.get("message").and_then(Value::as_str) { + error["message"] = Value::from(message); + } + self.send_control(&error); + // Non-fatal errors (an oversized input frame) keep the + // attachment; a refused open or a dropped attachment ends + // the connection. + if !self.manager.has_attachment(&self.pty_id) { + self.finish(); + } + } + _ => {} + } + } + + fn frame_context(self: &Arc) -> FrameContext { + let sink = Arc::clone(self); + let probe = Arc::clone(self); + FrameContext { + send: Arc::new(move |frame: Value| sink.on_manager_frame(&frame)), + buffered_amount: Arc::new(move || probe.pending_out.load(Ordering::SeqCst)), + trust: "supervised".to_owned(), + local_roots: None, + owner_user_id: None, + transport_id: Some(self.pty_id.clone()), + } + } +} + +async fn handle_client_frame( + connection: &Arc, + context: &FrameContext, + frame: TunnelFrame, +) { + if connection.finished.load(Ordering::SeqCst) { + return; + } + if frame.kind == FRAME_KIND_PTY { + if !connection.open_sent.load(Ordering::SeqCst) { + connection.protocol_error("bad_request", "bytes before open"); + return; + } + let input = json!({ + "version": PTY_PROTOCOL_VERSION, + "type": "pty_input", + "ptyId": connection.pty_id, + "dataB64": BASE64.encode(&frame.payload), + }); + connection.manager.handle_frame(&input, context).await; + return; + } + let Some(parsed) = parse_tunnel_client_frame(&frame.payload) else { + connection.protocol_error("bad_request", "invalid terminal request"); + return; + }; + match parsed { + ClientFrame::Open { session, surface, cols, rows } => { + if connection.open_sent.swap(true, Ordering::SeqCst) { + connection.protocol_error("bad_request", "duplicate open"); + return; + } + let mut open = json!({ + "version": PTY_PROTOCOL_VERSION, + "type": "pty_open", + "ptyId": connection.pty_id, + "session": session.unwrap_or_else(generate_session_name), + "cols": cols, + "rows": rows, + }); + if let Some(surface) = surface { + open["surface"] = Value::from(surface); + } + connection.manager.handle_frame(&open, context).await; + } + ClientFrame::Resize { cols, rows } => { + if !connection.open_sent.load(Ordering::SeqCst) { + return; + } + let resize = json!({ + "version": PTY_PROTOCOL_VERSION, + "type": "pty_resize", + "ptyId": connection.pty_id, + "cols": cols, + "rows": rows, + }); + connection.manager.handle_frame(&resize, context).await; + } + ClientFrame::Detach => connection.finish(), + } +} + +async fn serve_connection(stream: TcpStream, manager: Arc, parent: CancellationToken) { + let _ = stream.set_nodelay(true); + let (mut read_half, mut write_half) = stream.into_split(); + let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::(); + let (flow_tx, mut flow_rx) = mpsc::unbounded_channel::(); + let connection = Arc::new(Connection { + pty_id: format!("tunnel-{}", random_hex(8)), + manager: Arc::clone(&manager), + writer_tx, + flow_tx, + pending_out: AtomicU64::new(0), + paused: AtomicBool::new(false), + open_sent: AtomicBool::new(false), + opened_seen: AtomicBool::new(false), + finished: AtomicBool::new(false), + done: CancellationToken::new(), + }); + let context = connection.frame_context(); + + // Writer: the only task that touches the write half. Applies the flow + // water marks as the queue drains. + let mut writer = { + let connection = Arc::clone(&connection); + tokio::spawn(async move { + while let Some(message) = writer_rx.recv().await { + match message { + WriterMessage::Frame(frame) => { + let written = write_half.write_all(&frame).await; + // Every dequeued frame added exactly its length at + // enqueue, so this never underflows. + let length = frame.len() as u64; + let previous = connection.pending_out.fetch_sub(length, Ordering::SeqCst); + if written.is_err() { + connection.finish(); + break; + } + if previous.saturating_sub(length) < FLOW_RESUME_BYTES + && connection.paused.swap(false, Ordering::SeqCst) + { + let _ = connection.flow_tx.send(false); + } + } + WriterMessage::End => break, + } + } + let _ = write_half.shutdown().await; + }) + }; + + // Flow verbs need the async manager; drain them on their own task so a + // slow open never delays a pause. + let flow = { + let connection = Arc::clone(&connection); + let context = context.clone(); + tokio::spawn(async move { + while let Some(pause) = flow_rx.recv().await { + if connection.finished.load(Ordering::SeqCst) { + break; + } + let frame = json!({ + "version": PTY_PROTOCOL_VERSION, + "type": "pty_flow", + "ptyId": connection.pty_id, + "pause": pause, + }); + connection.manager.handle_frame(&frame, &context).await; + } + }) + }; + + // Reader: strictly in arrival order, so input received while an open + // settles lands after the attachment exists. Awaiting each frame is the + // ingest backpressure (the socket is simply not read meanwhile). + let mut buffer = vec![0_u8; 65_536]; + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let open_deadline = tokio::time::sleep(OPEN_TIMEOUT); + tokio::pin!(open_deadline); + loop { + tokio::select! { + biased; + _ = parent.cancelled() => { + connection.finish(); + break; + } + _ = connection.done.cancelled() => break, + _ = &mut open_deadline, if !connection.opened_seen.load(Ordering::SeqCst) => { + connection.protocol_error("bad_request", "no open frame"); + break; + } + read = read_half.read(&mut buffer) => { + let count = match read { + Ok(0) | Err(_) => { + // A torn splice is a detach, exactly like a dropped + // browser socket. + connection.finish(); + break; + } + Ok(count) => count, + }; + match decoder.push(&buffer[..count]) { + Ok(frames) => { + for frame in frames { + handle_client_frame(&connection, &context, frame).await; + } + } + Err(_) => { + connection.protocol_error("bad_request", "malformed frame"); + break; + } + } + } + } + } + connection.finish(); + // Detach, never kill: the owed close releases only this connection's + // attachment (transport-fenced), and the session lives on. + if connection.open_sent.load(Ordering::SeqCst) { + let close = json!({ + "version": PTY_PROTOCOL_VERSION, + "type": "pty_close", + "ptyId": connection.pty_id, + }); + manager.handle_frame(&close, &context).await; + } + flow.abort(); + // A peer that stopped reading can wedge the final flush forever; the + // attachment is already released above, so cap the flush and reap. + if tokio::time::timeout(Duration::from_secs(30), &mut writer).await.is_err() { + writer.abort(); + } + let _ = flow.await; +} + +/// Start the loopback listener. Managed mode only — the caller's managed +/// branch is the gate; paired human machines never reach this. Returns the +/// bound port; a bind failure is the caller's cue to degrade (the relay +/// socket path still serves terminals). +pub async fn start_tunnel_terminal_listener( + manager: Arc, + cancellation: CancellationToken, + host: &str, + port: u16, +) -> std::io::Result { + let listener = TcpListener::bind((host, port)).await?; + let bound = listener.local_addr()?.port(); + tokio::spawn(async move { + loop { + let accepted = tokio::select! { + biased; + _ = cancellation.cancelled() => break, + accepted = listener.accept() => accepted, + }; + match accepted { + Ok((stream, _)) => { + let manager = Arc::clone(&manager); + let child = cancellation.child_token(); + tokio::spawn(serve_connection(stream, manager, child)); + } + Err(_) => { + // Transient accept errors (EMFILE and friends) must not + // spin; the listener itself stays up. + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + } + }); + Ok(bound) +} + +// --------------------------------------------------------------------------- +// Tests — mirror packages/relay/test/tunnel-terminal.test.mjs in chatmux. +// A fake PtyDeps drives the real PtyManager over a real loopback socket. +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use crate::pty::{ + CmuxTui, DataSink, EnsureDaemon, ExitSink, PtyControl, PtyDeps, PtyHandle, PtyOutput, + SpawnSpec, + }; + use async_trait::async_trait; + use bytes::Bytes; + use std::collections::HashMap; + use std::path::{Path, PathBuf}; + use std::sync::Mutex as StdMutex; + use tokio::net::tcp::OwnedReadHalf; + + #[derive(Default)] + struct FakeState { + on_data: Option, + on_exit: Option, + written: Vec>, + resized: Vec<(u16, u16)>, + killed: bool, + } + + #[derive(Clone)] + struct FakePty { + state: Arc>, + } + + impl FakePty { + fn emit(&self, text: &str) { + let sink = self.state.lock().unwrap().on_data.clone(); + if let Some(sink) = sink { + sink(Bytes::copy_from_slice(text.as_bytes())); + } + } + fn exit(&self, code: i64) { + let sink = self.state.lock().unwrap().on_exit.clone(); + if let Some(sink) = sink { + sink(code); + } + } + } + + impl PtyControl for FakePty { + fn write(&self, data: &[u8]) { + self.state.lock().unwrap().written.push(data.to_vec()); + } + fn resize(&self, cols: u16, rows: u16) { + self.state.lock().unwrap().resized.push((cols, rows)); + } + fn pause(&self) {} + fn resume(&self) {} + fn kill(&self) { + self.state.lock().unwrap().killed = true; + } + } + + impl PtyOutput for FakePty { + fn subscribe(&self, on_data: DataSink, on_exit: ExitSink) { + let mut state = self.state.lock().unwrap(); + state.on_data = Some(on_data); + state.on_exit = Some(on_exit); + } + } + + struct FakeDeps { + spawned: Arc>>, + } + + #[async_trait] + impl PtyDeps for FakeDeps { + async fn spawn_pty(&self, _spec: SpawnSpec) -> PtyHandle { + let pty = FakePty { state: Arc::new(StdMutex::new(FakeState::default())) }; + self.spawned.lock().unwrap().push(pty.clone()); + PtyHandle { control: Arc::new(pty.clone()), output: Arc::new(pty), banner: None } + } + async fn resolve_cmux_tui(&self) -> Option { + None + } + async fn ensure_daemon( + &self, + _cmux_tui: &CmuxTui, + _session: &str, + _socket_dir: &Path, + _cwd: &Path, + _env: &HashMap, + ) -> Result { + Err("no daemon in tunnel tests".to_owned()) + } + async fn connect_control( + &self, + _socket_path: &Path, + ) -> Result, String> { + Err("no control in tunnel tests".to_owned()) + } + async fn read_dir(&self, _path: &Path) -> Result, ()> { + Err(()) + } + fn socket_dir(&self) -> PathBuf { + std::env::temp_dir() + } + fn shell(&self) -> String { + "/bin/fakesh".to_owned() + } + } + + struct Rig { + manager: Arc, + spawned: Arc>>, + port: u16, + cancel: CancellationToken, + } + + async fn rig_with_limits(max_ptys: usize) -> Rig { + let spawned = Arc::new(StdMutex::new(Vec::new())); + let deps = Arc::new(FakeDeps { spawned: Arc::clone(&spawned) }); + let env = HashMap::from([ + ("SHELL".to_owned(), "/bin/fakesh".to_owned()), + ("HOME".to_owned(), std::env::temp_dir().to_string_lossy().into_owned()), + ]); + let manager = Arc::new(PtyManager::with_limits( + deps, + std::env::temp_dir(), + env, + max_ptys, + 32, + 1_048_576, + )); + let cancel = CancellationToken::new(); + let port = start_tunnel_terminal_listener( + Arc::clone(&manager), + cancel.clone(), + TUNNEL_TERMINAL_HOST, + 0, + ) + .await + .expect("bind test listener"); + Rig { manager, spawned, port, cancel } + } + + async fn rig() -> Rig { + rig_with_limits(8).await + } + + async fn connect(rig: &Rig) -> TcpStream { + TcpStream::connect((TUNNEL_TERMINAL_HOST, rig.port)).await.expect("connect") + } + + /// Read whole frames off the socket with a deadline; panics on EOF. + async fn next_frame( + read: &mut OwnedReadHalf, + decoder: &mut TunnelFrameDecoder, + queue: &mut Vec, + ) -> TunnelFrame { + loop { + if !queue.is_empty() { + return queue.remove(0); + } + let mut buffer = vec![0_u8; 65_536]; + let count = tokio::time::timeout(Duration::from_secs(5), read.read(&mut buffer)) + .await + .expect("frame deadline") + .expect("read"); + assert!(count > 0, "peer closed while a frame was expected"); + queue.extend(decoder.push(&buffer[..count]).expect("decode")); + } + } + + fn control_json(frame: &TunnelFrame) -> Value { + assert_eq!(frame.kind, FRAME_KIND_CONTROL); + serde_json::from_slice(&frame.payload).expect("control json") + } + + /// Wait until the fake spawn landed (open settles asynchronously). + async fn spawned_pty(rig: &Rig) -> FakePty { + for _ in 0..100 { + if let Some(pty) = rig.spawned.lock().unwrap().first().cloned() { + return pty; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("no PTY spawned"); + } + + async fn read_eof(read: &mut OwnedReadHalf) { + let mut buffer = vec![0_u8; 4_096]; + loop { + let count = tokio::time::timeout(Duration::from_secs(5), read.read(&mut buffer)) + .await + .expect("eof deadline") + .expect("read"); + if count == 0 { + return; + } + } + } + + // -- pure codec/parse --------------------------------------------------- + + #[test] + fn codec_round_trips_frames_split_at_every_byte_boundary() { + let control = encode_control_frame(&json!({ "t": "open", "cols": 80, "rows": 24 })); + let pty = encode_pty_frame(b"echo hi\r"); + let stream = [control, pty].concat(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut frames = Vec::new(); + for byte in stream { + frames.extend(decoder.push(&[byte]).expect("clean stream")); + } + assert_eq!(frames.len(), 2); + assert_eq!(frames[0].kind, FRAME_KIND_CONTROL); + assert_eq!(frames[1].kind, FRAME_KIND_PTY); + assert_eq!(frames[1].payload, b"echo hi\r"); + } + + #[test] + fn oversized_length_poisons_the_decoder() { + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut header = ((MAX_TUNNEL_FRAME_BYTES + 1) as u32).to_be_bytes().to_vec(); + header.push(FRAME_KIND_PTY); + assert_eq!(decoder.push(&header), Err("frame_too_large")); + assert_eq!(decoder.push(b"anything"), Err("decoder_poisoned")); + } + + #[test] + fn unknown_kind_poisons_the_decoder() { + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut frame = 1_u32.to_be_bytes().to_vec(); + frame.push(7); + frame.push(b'x'); + assert_eq!(decoder.push(&frame), Err("unknown_frame_kind")); + } + + #[test] + fn parse_accepts_the_wire_shapes_and_rejects_malformed_requests() { + let open = parse_tunnel_client_frame(br#"{"t":"open","cols":80,"rows":24}"#); + assert_eq!( + open, + Some(ClientFrame::Open { session: None, surface: None, cols: 80, rows: 24 }) + ); + let full = parse_tunnel_client_frame( + br#"{"t":"open","session":"web-abc2","surface":"s:1.2","cols":1,"rows":10000}"#, + ); + assert_eq!( + full, + Some(ClientFrame::Open { + session: Some("web-abc2".to_owned()), + surface: Some("s:1.2".to_owned()), + cols: 1, + rows: 10_000, + }) + ); + assert_eq!( + parse_tunnel_client_frame(br#"{"t":"resize","cols":120,"rows":40}"#), + Some(ClientFrame::Resize { cols: 120, rows: 40 }) + ); + assert_eq!(parse_tunnel_client_frame(br#"{"t":"detach"}"#), Some(ClientFrame::Detach)); + for bad in [ + &br#"{"t":"open","cols":0,"rows":24}"#[..], + br#"{"t":"open","cols":10001,"rows":24}"#, + br#"{"t":"open","cols":80.5,"rows":24}"#, + br#"{"t":"open","cols":80}"#, + br#"{"t":"open","surface":"s:1.2","cols":80,"rows":24}"#, + br#"{"t":"open","session":"bad/name","cols":80,"rows":24}"#, + br#"{"t":"open","session":null,"cols":80,"rows":24}"#, + br#"{"t":"nope"}"#, + br#"[]"#, + br#"not json"#, + ] { + assert_eq!(parse_tunnel_client_frame(bad), None, "{}", String::from_utf8_lossy(bad)); + } + } + + #[test] + fn wire_error_codes_match_the_worker_map() { + assert_eq!(wire_error_code("trust_refused"), "trust_blocked"); + assert_eq!(wire_error_code("bad_request"), "bad_request"); + assert_eq!(wire_error_code("session_limit"), "session_limit"); + assert_eq!(wire_error_code("terminal_gone"), "terminal_gone"); + assert_eq!(wire_error_code("overflow"), "overflow"); + assert_eq!(wire_error_code("trust_revoked"), "trust_revoked"); + assert_eq!(wire_error_code("busy"), "busy"); + assert_eq!(wire_error_code("failed"), "failed"); + assert_eq!(wire_error_code("brand_new_code"), "failed"); + } + + #[test] + fn generated_session_names_use_the_web_prefix_and_alphabet() { + for _ in 0..32 { + let name = generate_session_name(); + let suffix = name.strip_prefix("web-").expect("web- prefix"); + assert_eq!(suffix.len(), 4); + assert!(suffix.chars().all(|c| "abcdefghjkmnpqrstuvwxyz23456789".contains(c))); + } + } + + // -- live listener ------------------------------------------------------ + + #[tokio::test] + async fn handshake_streams_both_ways_and_a_drop_detaches_without_killing() { + let rig = rig().await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + + write + .write_all(&encode_control_frame(&json!({ "t": "open", "cols": 80, "rows": 24 }))) + .await + .unwrap(); + let opened = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(opened["t"], "opened"); + assert_eq!(opened["created"], true); + assert_eq!(opened["cols"], 80); + assert_eq!(opened["rows"], 24); + let session = opened["session"].as_str().expect("session name").to_owned(); + assert!(session.starts_with("web-"), "server-minted name: {session}"); + + let pty = spawned_pty(&rig).await; + pty.emit("hello from the shell"); + let output = next_frame(&mut read, &mut decoder, &mut queue).await; + assert_eq!(output.kind, FRAME_KIND_PTY); + assert_eq!(output.payload, b"hello from the shell"); + + write.write_all(&encode_pty_frame(b"ls\r")).await.unwrap(); + write + .write_all(&encode_control_frame(&json!({ "t": "resize", "cols": 132, "rows": 43 }))) + .await + .unwrap(); + for _ in 0..100 { + if !pty.state.lock().unwrap().written.is_empty() + && !pty.state.lock().unwrap().resized.is_empty() + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(pty.state.lock().unwrap().written, vec![b"ls\r".to_vec()]); + assert_eq!(pty.state.lock().unwrap().resized, vec![(132, 43)]); + + // A torn splice detaches; the shell session must survive for a + // later re-attach (created:false proves it was found again). + drop(write); + drop(read); + for _ in 0..100 { + if rig.manager.attachment_count() == 0 { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(rig.manager.attachment_count(), 0, "drop must release the attachment"); + assert!(!pty.state.lock().unwrap().killed, "detach must not kill the session"); + + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + write + .write_all(&encode_control_frame( + &json!({ "t": "open", "session": session, "cols": 80, "rows": 24 }), + )) + .await + .unwrap(); + let reopened = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(reopened["t"], "opened"); + assert_eq!(reopened["created"], false); + rig.cancel.cancel(); + } + + #[tokio::test] + async fn bytes_before_open_are_a_protocol_error() { + let rig = rig().await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + write.write_all(&encode_pty_frame(b"sneaky")).await.unwrap(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + let error = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(error["t"], "error"); + assert_eq!(error["code"], "bad_request"); + read_eof(&mut read).await; + rig.cancel.cancel(); + } + + #[tokio::test] + async fn duplicate_open_is_a_protocol_error() { + let rig = rig().await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + let open = encode_control_frame(&json!({ "t": "open", "cols": 80, "rows": 24 })); + write.write_all(&open).await.unwrap(); + write.write_all(&open).await.unwrap(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + loop { + let frame = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + if frame["t"] == "error" { + assert_eq!(frame["code"], "bad_request"); + break; + } + assert_eq!(frame["t"], "opened"); + } + read_eof(&mut read).await; + rig.cancel.cancel(); + } + + #[tokio::test] + async fn a_malformed_control_frame_closes_the_connection() { + let rig = rig().await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + write.write_all(&encode_tunnel_frame(FRAME_KIND_CONTROL, b"{not json")).await.unwrap(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + let error = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(error["t"], "error"); + assert_eq!(error["code"], "bad_request"); + read_eof(&mut read).await; + rig.cancel.cancel(); + } + + #[tokio::test] + async fn a_pty_exit_reaches_the_client_and_ends_the_connection() { + let rig = rig().await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + write + .write_all(&encode_control_frame(&json!({ "t": "open", "cols": 80, "rows": 24 }))) + .await + .unwrap(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + let opened = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(opened["t"], "opened"); + spawned_pty(&rig).await.exit(3); + let exit = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(exit["t"], "exit"); + assert_eq!(exit["code"], 3); + read_eof(&mut read).await; + rig.cancel.cancel(); + } + + #[tokio::test] + async fn a_refused_open_maps_the_error_code_and_closes() { + let rig = rig_with_limits(0).await; + let stream = connect(&rig).await; + let (mut read, mut write) = stream.into_split(); + write + .write_all(&encode_control_frame(&json!({ "t": "open", "cols": 80, "rows": 24 }))) + .await + .unwrap(); + let mut decoder = TunnelFrameDecoder::new(MAX_TUNNEL_FRAME_BYTES); + let mut queue = Vec::new(); + let error = control_json(&next_frame(&mut read, &mut decoder, &mut queue).await); + assert_eq!(error["t"], "error"); + assert_eq!(error["code"], "session_limit"); + read_eof(&mut read).await; + rig.cancel.cancel(); + } +}