From 03dfe354bd08f73037b436f3384d4e6570bfb222 Mon Sep 17 00:00:00 2001 From: Sheng Kun Chang Date: Tue, 5 May 2026 08:39:34 +0800 Subject: [PATCH 1/2] feat(dev/tunnel): WSS client + frame protocol (Phase 1, no integration yet) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds the WebSocket side of `mirrorstack dev`: - Frame envelope (type, id, corr_id, stream_id, payload) — wire-shape matches api-platform's internal/dispatch/ws/frames.go - Connect: opens wss_url?token=ttok, sends register, awaits register_ack - Background task: pings every 30s, replies to server pings, sends close frame on shutdown signal Deferred frame types (rpc.*, sql.*) are enumerated in the FrameType enum so server frames of those shapes deserialize without forcing a reconnect once those planes go live. This module is `#[allow(dead_code)]` for now — it compiles + tests but isn't wired into the dev lifecycle yet. End-to-end integration requires two pieces still in flight elsewhere: - api-platform AWS Sender impl (apigatewaymanagementapi.PostToConnection) - IaC for the WS API GW + Lambda routing Once those land, a small follow-up PR plugs this module into dev::run and adds a --tunnel flag that gates the connect. Deps added (used only inside this module): - tokio (rt + macros + time + sync + net + io-util) - tokio-tungstenite (connect + rustls-tls-webpki-roots) - futures-util (sink + std) Test plan: - [x] cargo test (4 new round-trip + uniqueness tests) - [x] cargo clippy --all-targets -- -D warnings clean - [x] cargo fmt --all -- --check clean Co-Authored-By: Claude Opus 4.7 (1M context) --- Cargo.lock | 135 ++++++++++++++++- Cargo.toml | 3 + src/commands/dev/mod.rs | 5 + src/commands/dev/tunnel.rs | 303 +++++++++++++++++++++++++++++++++++++ 4 files changed, 438 insertions(+), 8 deletions(-) create mode 100644 src/commands/dev/tunnel.rs diff --git a/Cargo.lock b/Cargo.lock index 78a64c7..36d9cb4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -119,6 +119,12 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.11.1" @@ -245,6 +251,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "dialoguer" version = "0.11.0" @@ -590,7 +602,7 @@ dependencies = [ "tokio", "tokio-rustls", "tower-service", - "webpki-roots", + "webpki-roots 1.0.7", ] [[package]] @@ -891,16 +903,19 @@ dependencies = [ "dialoguer", "dirs", "dotenvy", + "futures-util", "indicatif", "mockito", "open", - "rand", + "rand 0.9.4", "reqwest", "serde", "serde_json", "sha2", "tempfile", "thiserror 2.0.18", + "tokio", + "tokio-tungstenite", "url", ] @@ -921,7 +936,7 @@ dependencies = [ "hyper-util", "log", "pin-project-lite", - "rand", + "rand 0.9.4", "regex", "serde_json", "serde_urlencoded", @@ -1104,7 +1119,7 @@ dependencies = [ "bytes", "getrandom 0.3.4", "lru-slab", - "rand", + "rand 0.9.4", "ring", "rustc-hash", "rustls", @@ -1151,14 +1166,35 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -1168,7 +1204,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -1266,7 +1311,7 @@ dependencies = [ "wasm-bindgen", "wasm-bindgen-futures", "web-sys", - "webpki-roots", + "webpki-roots 1.0.7", ] [[package]] @@ -1416,6 +1461,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sha2" version = "0.10.9" @@ -1606,9 +1662,21 @@ dependencies = [ "parking_lot", "pin-project-lite", "socket2", + "tokio-macros", "windows-sys 0.61.2", ] +[[package]] +name = "tokio-macros" +version = "2.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -1619,6 +1687,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" +dependencies = [ + "futures-util", + "log", + "rustls", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tungstenite", + "webpki-roots 0.26.11", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -1702,6 +1786,26 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e5b8366ee7a95b16d32197d0b2604b43a0be89dc5fac9f8e96ccafbaedda8a" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.8.6", + "rustls", + "rustls-pki-types", + "sha1", + "thiserror 1.0.69", + "utf-8", +] + [[package]] name = "typenum" version = "1.20.0" @@ -1744,6 +1848,12 @@ dependencies = [ "serde", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" @@ -1904,6 +2014,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.7", +] + [[package]] name = "webpki-roots" version = "1.0.7" diff --git a/Cargo.toml b/Cargo.toml index dfbec57..c0ce623 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,9 @@ dialoguer = { version = "0.11", default-features = false } console = "0.15" indicatif = "0.17" ctrlc = "3" +tokio = { version = "1", features = ["rt", "macros", "time", "sync", "net", "io-util"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-webpki-roots"] } +futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } [dev-dependencies] mockito = "1.6" diff --git a/src/commands/dev/mod.rs b/src/commands/dev/mod.rs index c7d2328..34f97d0 100644 --- a/src/commands/dev/mod.rs +++ b/src/commands/dev/mod.rs @@ -22,6 +22,11 @@ use super::{ok_mark, warn_prefix}; mod compose; mod process; +// Used in a follow-up PR that integrates WSS into the dev lifecycle. +// Compiled + tested now so the wire format and protocol stay correct +// while the AWS Sender impl + IaC for the WS API GW catch up. +#[allow(dead_code)] +mod tunnel; #[derive(Args)] pub struct DevArgs { diff --git a/src/commands/dev/tunnel.rs b/src/commands/dev/tunnel.rs new file mode 100644 index 0000000..ff2d204 --- /dev/null +++ b/src/commands/dev/tunnel.rs @@ -0,0 +1,303 @@ +//! WSS tunnel client for `mirrorstack dev`. +//! +//! Two layers: +//! - `frames` — typed envelopes (mirrors api-platform's `internal/dispatch/ws/frames.go`) +//! - `connect` — open the WSS, send `register`, await `register_ack`, hold the +//! connection and ping every 30s until the host process tells us to stop +//! +//! Phase 1 ships the connect + register half. RPC frames + SQL frames are +//! deferred to Phase 2/3 per the design doc; this module's enum +//! intentionally enumerates them so future PRs only add handler arms. + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::{Context, Result, anyhow}; +use futures_util::{SinkExt, StreamExt}; +use serde::{Deserialize, Serialize}; +use tokio::sync::Notify; +use tokio_tungstenite::tungstenite::Message; +use uuid_lite::uuid_v4_string; + +mod uuid_lite { + //! Tiny v4 UUID generator. We don't pull `uuid` for one call site — the + //! envelope just wants a unique ID to correlate frames, not anything + //! cryptographic. Bytes come from `rand::random` (already a CLI dep). + pub(super) fn uuid_v4_string() -> String { + let mut b: [u8; 16] = rand::random(); + // RFC 4122 variant + version-4 bits. + b[6] = (b[6] & 0x0f) | 0x40; + b[8] = (b[8] & 0x3f) | 0x80; + format!( + "{:02x}{:02x}{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}-{:02x}{:02x}{:02x}{:02x}{:02x}{:02x}", + b[0], + b[1], + b[2], + b[3], + b[4], + b[5], + b[6], + b[7], + b[8], + b[9], + b[10], + b[11], + b[12], + b[13], + b[14], + b[15] + ) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(super) enum FrameType { + Register, + RegisterAck, + Ping, + Pong, + Close, + // RPC + SQL planes are reserved here so server frames of those types + // don't fail to deserialize and force a reconnect. + #[serde(rename = "rpc.req")] + RpcReq, + #[serde(rename = "rpc.resp")] + RpcResp, + #[serde(rename = "rpc.err")] + RpcErr, + #[serde(rename = "sql.req")] + SqlReq, + #[serde(rename = "sql.rows")] + SqlRows, + #[serde(rename = "sql.end")] + SqlEnd, + #[serde(rename = "sql.err")] + SqlErr, +} + +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct Frame { + #[serde(rename = "type")] + pub frame_type: FrameType, + pub id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub corr_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stream_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub payload: Option, +} + +impl Frame { + pub(super) fn new(frame_type: FrameType, payload: Option) -> Self { + Self { + frame_type, + id: format!("frm_{}", uuid_v4_string()), + corr_id: None, + stream_id: None, + payload, + } + } + + pub(super) fn with_corr(mut self, corr_id: String) -> Self { + self.corr_id = Some(corr_id); + self + } +} + +#[derive(Debug, Serialize)] +pub(super) struct RegisterPayload<'a> { + pub module_id: &'a str, + pub local_url: &'a str, + pub version: &'a str, +} + +#[derive(Debug, Deserialize)] +pub(super) struct RegisterAck { + pub session_id: String, + pub service_token: String, + pub expires_at: String, +} + +/// Spawned-task handle. Drop or call [`shutdown`] to close the WSS. +pub(super) struct TunnelHandle { + pub session_id: String, + pub service_token: String, + shutdown: Arc, +} + +impl TunnelHandle { + /// Ask the tunnel to close. The background task will send a `close` + /// frame and drop the connection. Idempotent. + pub(super) fn shutdown(&self) { + self.shutdown.notify_one(); + } +} + +/// Connect to `wss_url?token=ttok`, send `register`, await `register_ack`, +/// then hand the connection off to a background task that pings every 30s +/// and replies to inbound pings. Returns when the handshake completes; the +/// connection itself lives until [`TunnelHandle::shutdown`] is called. +pub(super) async fn open( + wss_url: &str, + ttok: &str, + register: RegisterPayload<'_>, +) -> Result { + let url = if wss_url.contains('?') { + format!("{wss_url}&token={ttok}") + } else { + format!("{wss_url}?token={ttok}") + }; + let (ws_stream, _resp) = tokio_tungstenite::connect_async(&url) + .await + .with_context(|| format!("dev: connect WSS {wss_url}"))?; + let (mut sink, mut stream) = ws_stream.split(); + + // Send register. + let payload = serde_json::to_value(®ister).context("dev: serialize register payload")?; + let register_frame = Frame::new(FrameType::Register, Some(payload)); + let register_id = register_frame.id.clone(); + sink.send(Message::Text(serde_json::to_string(®ister_frame)?)) + .await + .context("dev: send register frame")?; + + // Await register_ack with a generous timeout. + let ack = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let next = stream + .next() + .await + .ok_or_else(|| anyhow!("stream closed before register_ack"))?; + let msg = next.context("dev: read register_ack")?; + let text = match msg { + Message::Text(t) => t.to_string(), + Message::Binary(b) => String::from_utf8_lossy(&b).into_owned(), + Message::Ping(p) => { + sink.send(Message::Pong(p)).await.ok(); + continue; + } + Message::Close(_) => return Err(anyhow!("server closed before register_ack")), + _ => continue, + }; + let frame: Frame = serde_json::from_str(&text) + .with_context(|| format!("dev: parse server frame: {text}"))?; + if frame.frame_type == FrameType::RegisterAck + && frame.corr_id.as_deref() == Some(®ister_id) + { + let ack_payload = frame + .payload + .ok_or_else(|| anyhow!("register_ack missing payload"))?; + let ack: RegisterAck = + serde_json::from_value(ack_payload).context("dev: deserialize register_ack")?; + return Ok(ack); + } + // Ignore unrelated frames during handshake — server may send + // ping or status frames before our ack arrives. + } + }) + .await + .context("dev: register_ack timeout")? + .context("dev: register_ack")?; + + let shutdown = Arc::new(Notify::new()); + let shutdown_clone = shutdown.clone(); + + // Background task: ping every 30s, respond to server pings, close on shutdown. + tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + // First tick fires immediately; we don't want a ping at t=0. + interval.tick().await; + loop { + tokio::select! { + _ = shutdown_clone.notified() => { + let close = Frame::new(FrameType::Close, None); + if let Ok(body) = serde_json::to_string(&close) { + let _ = sink.send(Message::Text(body)).await; + } + let _ = sink.close().await; + return; + } + _ = interval.tick() => { + let ping = Frame::new(FrameType::Ping, None); + if let Ok(body) = serde_json::to_string(&ping) { + if sink.send(Message::Text(body)).await.is_err() { + return; + } + } + } + msg = stream.next() => { + match msg { + Some(Ok(Message::Text(_))) | Some(Ok(Message::Binary(_))) => { + // Server frames (pong, future RPC responses) ignored + // for now — Phase 1 doesn't act on them client-side. + } + Some(Ok(Message::Ping(p))) => { + let _ = sink.send(Message::Pong(p)).await; + } + Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Close(_))) | Some(Err(_)) | None => return, + Some(Ok(Message::Frame(_))) => {} + } + } + } + } + }); + + Ok(TunnelHandle { + session_id: ack.session_id, + service_token: ack.service_token, + shutdown, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn frame_serialization_round_trip() { + let f = Frame::new(FrameType::Register, Some(serde_json::json!({"hi": 1}))); + let s = serde_json::to_string(&f).unwrap(); + assert!(s.contains("\"type\":\"register\"")); + assert!(s.contains("\"id\":\"frm_")); + let back: Frame = serde_json::from_str(&s).unwrap(); + assert_eq!(back.frame_type, FrameType::Register); + } + + #[test] + fn frame_type_renames_match_server_protocol() { + let cases = [ + (FrameType::Register, "register"), + (FrameType::RegisterAck, "register_ack"), + (FrameType::Ping, "ping"), + (FrameType::Pong, "pong"), + (FrameType::RpcReq, "rpc.req"), + (FrameType::SqlRows, "sql.rows"), + ]; + for (variant, wire) in cases { + let s = serde_json::to_string(&variant).unwrap(); + assert_eq!(s, format!("\"{wire}\"")); + } + } + + #[test] + fn frame_id_is_unique_across_calls() { + let a = Frame::new(FrameType::Ping, None); + let b = Frame::new(FrameType::Ping, None); + assert_ne!(a.id, b.id); + assert!(a.id.starts_with("frm_")); + } + + #[test] + fn register_payload_serializes_with_snake_case_fields() { + let p = RegisterPayload { + module_id: "m_abc", + local_url: "http://localhost:8080", + version: "0.1.0", + }; + let s = serde_json::to_string(&p).unwrap(); + assert!(s.contains("\"module_id\":\"m_abc\"")); + assert!(s.contains("\"local_url\":\"http://localhost:8080\"")); + } +} From 479c5049f529d8ac27382ed75fa1677687ae558c Mon Sep 17 00:00:00 2001 From: Sheng Kun Chang Date: Wed, 6 May 2026 02:20:39 +0800 Subject: [PATCH 2/2] =?UTF-8?q?refactor(dev/tunnel):=20/simplify=20pass=20?= =?UTF-8?q?=E2=80=94=20URL=20builder,=20flatter=20handshake,=20observabili?= =?UTF-8?q?ty?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five concrete fixes from the multi-agent review: 1. URL building uses url::Url + query_pairs_mut, matching auth::authorize_url (auth/mod.rs:52-61). Drops the brittle "?" vs "&" substring check. 2. Handshake flattened into await_register_ack(sink, stream, register_id) — removes three-level nesting (timeout > async > loop > match). 3. Background task surfaces transport failures via warn_prefix() instead of silently returning. Stream close, send error, server close all log actionable lines now. (Liveness-signal-on-handle queued as #29 — that's a structural change worth its own PR.) 4. WebSocketConfig caps inbound frames at 1 MiB (was: tungstenite default 64 MiB). Per L1.5 a frame is at most 128 KB; 1 MiB gives streaming headroom without exposing a memory-bomb surface. 5. TunnelHandle drops the unused service_token field. RegisterAck still carries it (forward-compat with the wire shape); re-thread to the handle when the RPC plane lands and actually consumes it. Plus: MissedTickBehavior::Delay on the heartbeat interval so a paused runtime (laptop sleep) doesn't burst-ping on resume. Three new tests cover the URL builder. 54 tests pass, clippy + fmt clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/commands/dev/tunnel.rs | 251 ++++++++++++++++++++++++------------- 1 file changed, 166 insertions(+), 85 deletions(-) diff --git a/src/commands/dev/tunnel.rs b/src/commands/dev/tunnel.rs index ff2d204..903bc3f 100644 --- a/src/commands/dev/tunnel.rs +++ b/src/commands/dev/tunnel.rs @@ -13,12 +13,30 @@ use std::sync::Arc; use std::time::Duration; use anyhow::{Context, Result, anyhow}; +use futures_util::stream::{SplitSink, SplitStream}; use futures_util::{SinkExt, StreamExt}; use serde::{Deserialize, Serialize}; +use tokio::net::TcpStream; use tokio::sync::Notify; +use tokio::time::MissedTickBehavior; use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use url::Url; use uuid_lite::uuid_v4_string; +use super::super::warn_prefix; + +type WsStream = WebSocketStream>; +type WsSink = SplitSink; +type WsReader = SplitStream; + +/// Cap inbound frames at 1 MiB. tokio-tungstenite's default is 64 MiB — +/// way too generous for our wire format (L1.5 says 128 KB per frame; 1 MiB +/// gives an order of magnitude headroom for streaming row batches without +/// exposing a memory-bomb surface to a misbehaving server. +const MAX_INBOUND_FRAME_BYTES: usize = 1 << 20; + mod uuid_lite { //! Tiny v4 UUID generator. We don't pull `uuid` for one call site — the //! envelope just wants a unique ID to correlate frames, not anything @@ -121,9 +139,12 @@ pub(super) struct RegisterAck { } /// Spawned-task handle. Drop or call [`shutdown`] to close the WSS. +/// +/// `service_token` is intentionally NOT held here — it lives on the +/// returned [`RegisterAck`] until a real consumer (the future RPC plane) +/// needs it on the handle itself. Re-thread when wired. pub(super) struct TunnelHandle { pub session_id: String, - pub service_token: String, shutdown: Arc, } @@ -144,111 +165,151 @@ pub(super) async fn open( ttok: &str, register: RegisterPayload<'_>, ) -> Result { - let url = if wss_url.contains('?') { - format!("{wss_url}&token={ttok}") - } else { - format!("{wss_url}?token={ttok}") + let url = with_token_param(wss_url, ttok)?; + let config = WebSocketConfig { + max_message_size: Some(MAX_INBOUND_FRAME_BYTES), + max_frame_size: Some(MAX_INBOUND_FRAME_BYTES), + ..Default::default() }; - let (ws_stream, _resp) = tokio_tungstenite::connect_async(&url) - .await - .with_context(|| format!("dev: connect WSS {wss_url}"))?; + let (ws_stream, _resp) = + tokio_tungstenite::connect_async_with_config(&url, Some(config), false) + .await + .with_context(|| format!("dev: connect WSS {wss_url}"))?; let (mut sink, mut stream) = ws_stream.split(); - // Send register. - let payload = serde_json::to_value(®ister).context("dev: serialize register payload")?; - let register_frame = Frame::new(FrameType::Register, Some(payload)); + let register_frame = Frame::new( + FrameType::Register, + Some(serde_json::to_value(®ister).context("dev: serialize register payload")?), + ); let register_id = register_frame.id.clone(); sink.send(Message::Text(serde_json::to_string(®ister_frame)?)) .await .context("dev: send register frame")?; - // Await register_ack with a generous timeout. - let ack = tokio::time::timeout(Duration::from_secs(10), async { - loop { - let next = stream - .next() - .await - .ok_or_else(|| anyhow!("stream closed before register_ack"))?; - let msg = next.context("dev: read register_ack")?; - let text = match msg { - Message::Text(t) => t.to_string(), - Message::Binary(b) => String::from_utf8_lossy(&b).into_owned(), - Message::Ping(p) => { - sink.send(Message::Pong(p)).await.ok(); - continue; - } - Message::Close(_) => return Err(anyhow!("server closed before register_ack")), - _ => continue, - }; - let frame: Frame = serde_json::from_str(&text) - .with_context(|| format!("dev: parse server frame: {text}"))?; - if frame.frame_type == FrameType::RegisterAck - && frame.corr_id.as_deref() == Some(®ister_id) - { - let ack_payload = frame - .payload - .ok_or_else(|| anyhow!("register_ack missing payload"))?; - let ack: RegisterAck = - serde_json::from_value(ack_payload).context("dev: deserialize register_ack")?; - return Ok(ack); - } - // Ignore unrelated frames during handshake — server may send - // ping or status frames before our ack arrives. - } - }) + let ack = tokio::time::timeout( + Duration::from_secs(10), + await_register_ack(&mut sink, &mut stream, ®ister_id), + ) .await .context("dev: register_ack timeout")? .context("dev: register_ack")?; let shutdown = Arc::new(Notify::new()); - let shutdown_clone = shutdown.clone(); + tokio::spawn(run_tunnel_loop(sink, stream, shutdown.clone())); - // Background task: ping every 30s, respond to server pings, close on shutdown. - tokio::spawn(async move { - let mut interval = tokio::time::interval(Duration::from_secs(30)); - // First tick fires immediately; we don't want a ping at t=0. - interval.tick().await; - loop { - tokio::select! { - _ = shutdown_clone.notified() => { - let close = Frame::new(FrameType::Close, None); - if let Ok(body) = serde_json::to_string(&close) { - let _ = sink.send(Message::Text(body)).await; - } - let _ = sink.close().await; - return; + Ok(TunnelHandle { + session_id: ack.session_id, + shutdown, + }) +} + +/// Append `?token=` (or `&token=` if a query already exists) to +/// `wss_url` using the `url` crate's parser — same pattern as +/// `auth::authorize_url` for the OAuth consent URL. Avoids the fragile +/// `?` vs `&` substring check the previous version did. +fn with_token_param(wss_url: &str, ttok: &str) -> Result { + let mut u = Url::parse(wss_url).with_context(|| format!("dev: parse wss_url {wss_url}"))?; + u.query_pairs_mut().append_pair("token", ttok); + Ok(u.into()) +} + +/// Read frames off `stream` until a `register_ack` correlated with +/// `register_id` arrives. Server-side ping/binary/text-without-ack are +/// tolerated — the server may send heartbeat or status frames before +/// the ack lands. Wrapped in `tokio::time::timeout` by the caller. +async fn await_register_ack( + sink: &mut WsSink, + stream: &mut WsReader, + register_id: &str, +) -> Result { + loop { + let next = stream + .next() + .await + .ok_or_else(|| anyhow!("stream closed before register_ack"))?; + let msg = next.context("dev: read register_ack")?; + let text = match msg { + Message::Text(t) => t.to_string(), + Message::Binary(b) => String::from_utf8_lossy(&b).into_owned(), + Message::Ping(p) => { + let _ = sink.send(Message::Pong(p)).await; + continue; + } + Message::Close(_) => return Err(anyhow!("server closed before register_ack")), + _ => continue, + }; + let frame: Frame = serde_json::from_str(&text) + .with_context(|| format!("dev: parse server frame: {text}"))?; + if frame.frame_type == FrameType::RegisterAck + && frame.corr_id.as_deref() == Some(register_id) + { + let payload = frame + .payload + .ok_or_else(|| anyhow!("register_ack missing payload"))?; + return serde_json::from_value(payload).context("dev: deserialize register_ack"); + } + // Other frames during handshake are ignored — see fn comment. + } +} + +/// Long-running tunnel loop: ping every 30s, respond to server pings, +/// shut down on signal. Surfaces transport failures via `warn_prefix()` +/// so a silently-dying tunnel doesn't leave the user wondering why +/// inbound calls stop arriving. (See issue #29 for upgrading this to a +/// liveness signal callers can poll.) +async fn run_tunnel_loop(mut sink: WsSink, mut stream: WsReader, shutdown: Arc) { + let mut interval = tokio::time::interval(Duration::from_secs(30)); + // The first tick fires immediately — consume it so the first ping is at + // t=30s, not t=0. Delay the missed-tick behavior so a paused runtime + // (e.g. laptop sleep) doesn't burst-ping on resume. + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + interval.tick().await; + loop { + tokio::select! { + _ = shutdown.notified() => { + let close = Frame::new(FrameType::Close, None); + if let Ok(body) = serde_json::to_string(&close) { + let _ = sink.send(Message::Text(body)).await; } - _ = interval.tick() => { - let ping = Frame::new(FrameType::Ping, None); - if let Ok(body) = serde_json::to_string(&ping) { - if sink.send(Message::Text(body)).await.is_err() { - return; - } + let _ = sink.close().await; + return; + } + _ = interval.tick() => { + let ping = Frame::new(FrameType::Ping, None); + if let Ok(body) = serde_json::to_string(&ping) { + if let Err(e) = sink.send(Message::Text(body)).await { + eprintln!("{} tunnel: ping send failed ({e}); closing tunnel", warn_prefix()); + return; } } - msg = stream.next() => { - match msg { - Some(Ok(Message::Text(_))) | Some(Ok(Message::Binary(_))) => { - // Server frames (pong, future RPC responses) ignored - // for now — Phase 1 doesn't act on them client-side. - } - Some(Ok(Message::Ping(p))) => { - let _ = sink.send(Message::Pong(p)).await; - } - Some(Ok(Message::Pong(_))) => {} - Some(Ok(Message::Close(_))) | Some(Err(_)) | None => return, - Some(Ok(Message::Frame(_))) => {} + } + msg = stream.next() => { + match msg { + Some(Ok(Message::Text(_))) | Some(Ok(Message::Binary(_))) => { + // Server frames (pong, future RPC responses) ignored + // for now — Phase 1 doesn't act on them client-side. + } + Some(Ok(Message::Ping(p))) => { + let _ = sink.send(Message::Pong(p)).await; + } + Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Close(reason))) => { + eprintln!("{} tunnel closed by server: {reason:?}", warn_prefix()); + return; + } + Some(Err(e)) => { + eprintln!("{} tunnel: read failed ({e}); closing tunnel", warn_prefix()); + return; } + None => { + eprintln!("{} tunnel: stream ended", warn_prefix()); + return; + } + Some(Ok(Message::Frame(_))) => {} } } } - }); - - Ok(TunnelHandle { - session_id: ack.session_id, - service_token: ack.service_token, - shutdown, - }) + } } #[cfg(test)] @@ -300,4 +361,24 @@ mod tests { assert!(s.contains("\"module_id\":\"m_abc\"")); assert!(s.contains("\"local_url\":\"http://localhost:8080\"")); } + + #[test] + fn with_token_param_appends_when_no_query() { + let got = with_token_param("wss://api.example/ws", "ttok_abc").unwrap(); + assert_eq!(got, "wss://api.example/ws?token=ttok_abc"); + } + + #[test] + fn with_token_param_appends_when_query_present() { + let got = with_token_param("wss://api.example/ws?stage=prod", "ttok_abc").unwrap(); + // url crate may reorder pairs; just assert both made it. + assert!(got.contains("stage=prod")); + assert!(got.contains("token=ttok_abc")); + assert!(got.starts_with("wss://api.example/ws?")); + } + + #[test] + fn with_token_param_rejects_malformed_url() { + assert!(with_token_param("not a url", "ttok").is_err()); + } }