From a4ffb0cfa4af349a501ddae506381619c9150b74 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 02:39:49 +0800 Subject: [PATCH 1/7] Cover the whole transport/codec matrix in the e2e tunnel test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `test_delay.rs` ran one case: UDP with `--codec`, on hardcoded ports read from a checked-in `tests/.env`, sequenced by three 200ms sleeps. The TCP case existed but was `#[ignore]`d. So three of the four real forwarding paths were never exercised, and the one that was could only run when no relay already held port 7666. Now every case builds its own relay, echo server, and tunnel: - Four cases — TCP and UDP, each with and without `--codec`. The codec flag is not cosmetic: the relay generates an AES key per stream and swaps `Normal*` for `Codec*` readers and writers (`pb-mapper-server/src/client.rs:226`), so these are four distinct paths through the forwarder. - The transport is a parameter in code, not `SERVER_TEST_TYPE` in a dotfile, which is what made the matrix impossible to express before. - Every port is `:0`. The relay takes a pre-bound listener via `run_server_on_listener`, so there is no window between reserving a port and binding it. `connect` binds inside the client, so its address is reserved by binding and dropping — with the matching protocol, since TCP and UDP have separate port spaces. - No sleeps. Registration waits on the relay's `Keys` status; forwarding waits on a probe payload making the full round trip. Both are true readiness rather than a guess at how long startup takes. The cases run concurrently, which needs every process-global read out of the path: `AuthRuntime::start` takes the key directly instead of reading the process credential, both tunnel ends use their pinned-credential entry points, and each relay gets its own auth state directory for the `auth.lock` flock. One `LazyLock` writes the process credential once, for the framing checksum that `NormalMessageWriter` falls back to. `tests/.env` and the `dotenvy` dev-dependency are gone with the last test that needed them. Test count goes 136 pass + 1 ignored to 139 pass + 0 ignored: the ignored TCP case is replaced by real coverage, and the single tunnel case became four. Co-Authored-By: Claude Opus 5 --- AGENTS.md | 15 +- Cargo.lock | 7 - Cargo.toml | 1 - crates/pb-mapper-cli/Cargo.toml | 1 - crates/pb-mapper-cli/tests/.env | 5 - crates/pb-mapper-cli/tests/test_delay.rs | 713 ++++++++++++++--------- 6 files changed, 440 insertions(+), 302 deletions(-) delete mode 100644 crates/pb-mapper-cli/tests/.env diff --git a/AGENTS.md b/AGENTS.md index e21de91..92f2e2c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,7 +16,7 @@ - `crates/`: the Rust workspace; the root `Cargo.toml` is a virtual manifest - `crates/pb-mapper-cli/src/bin/pb-mapper.rs`: unified CLI entry point - `crates/pb-mapper-{core,auth,protocol,server,client,cli}` - - `crates/pb-mapper-cli/tests/`: integration tests; loads env from `tests/.env` + - `crates/pb-mapper-cli/tests/`: integration tests; no env setup required - `crates/pb-mapper-cli/examples/`: runnable examples - `ui/`: Flutter UI; Rust bridge under `ui/native/*` - `docker/`, `services/`, `scripts/`: container, systemd, build/release @@ -41,10 +41,15 @@ Notes: CI builds release artifacts on tags `vX.Y.Z` (see `.github/workflows/rele - Naming: modules/functions `snake_case`, types/traits `PascalCase`, consts `SCREAMING_SNAKE_CASE` ## Testing Guidelines -- Framework: `tokio` async + integration tests under `tests/` -- Env vars (see `tests/.env`): `PB_MAPPER_TEST_SERVER`, `LOCAL_TEST_SERVER`, `ECHO_TEST_SERVER`, `SERVER_TEST_KEY`, `SERVER_TEST_TYPE` (`TCP`/`UDP`) -- Run ignored tests: `cargo test -- --ignored` -- Prefer new integration tests in `tests/` with reproducible env defaults +- Framework: `tokio` async + integration tests in `crates/pb-mapper-cli/tests/` +- No test needs environment setup. `test_delay.rs` runs the whole tunnel + (`server` + `register` + `connect`) over both transports with and without + `--codec`; each case reserves its own loopback ports and its own auth state + directory, so cases run concurrently and never collide with a live relay. +- Sequence components with a readiness probe, not a sleep: poll the relay's + `Keys` status for a registration, and round-trip a payload through the tunnel + for forwarding. `TunnelHarness` in `test_delay.rs` does both. +- Prefer new integration tests that need no external setup ## Commit & Pull Request Guidelines - Commits: short, imperative (e.g., "Fix localhost resolution panic", "add network perms", "change to StreamBuilder") diff --git a/Cargo.lock b/Cargo.lock index 3e44e6f..3d848a0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -331,12 +331,6 @@ dependencies = [ "syn 2.0.114", ] -[[package]] -name = "dotenvy" -version = "0.15.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" - [[package]] name = "either" version = "1.18.0" @@ -1051,7 +1045,6 @@ version = "0.4.0" dependencies = [ "better_mimalloc_rs", "clap", - "dotenvy", "pb-mapper-auth", "pb-mapper-client", "pb-mapper-core", diff --git a/Cargo.toml b/Cargo.toml index 4848d5f..47e43fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,7 +30,6 @@ better_mimalloc_rs = { version = "0.1.2", features = ["config"] } bytes = "1.11" clap = { version = "4.5", features = ["derive"] } dirs = "6.0.0" -dotenvy = "0.15.7" hashbrown = { version = "0.17.1" } hickory-resolver = { version = "0.26.1" } kanal = { git = "https://github.com/acking-you/kanal.git", branch = "dev/pb-mapper" } diff --git a/crates/pb-mapper-cli/Cargo.toml b/crates/pb-mapper-cli/Cargo.toml index f45b297..07fb2ee 100644 --- a/crates/pb-mapper-cli/Cargo.toml +++ b/crates/pb-mapper-cli/Cargo.toml @@ -25,7 +25,6 @@ tracing.workspace = true uni-stream.workspace = true [dev-dependencies] -dotenvy.workspace = true rand.workspace = true [features] diff --git a/crates/pb-mapper-cli/tests/.env b/crates/pb-mapper-cli/tests/.env deleted file mode 100644 index ce2796a..0000000 --- a/crates/pb-mapper-cli/tests/.env +++ /dev/null @@ -1,5 +0,0 @@ -PB_MAPPER_TEST_SERVER="127.0.0.1:7666" -LOCAL_TEST_SERVER="127.0.0.1:33333" -ECHO_TEST_SERVER="127.0.0.1:44444" -SERVER_TEST_KEY="echo" -SERVER_TEST_TYPE="UDP" diff --git a/crates/pb-mapper-cli/tests/test_delay.rs b/crates/pb-mapper-cli/tests/test_delay.rs index 78c747b..39e9806 100644 --- a/crates/pb-mapper-cli/tests/test_delay.rs +++ b/crates/pb-mapper-cli/tests/test_delay.rs @@ -1,32 +1,83 @@ // See the note in `regression.rs`: the whole file is test code. #![allow(clippy::unwrap_used, clippy::expect_used)] -use std::env; +//! End-to-end tunnel tests over the full `server` + `register` + `connect` path. +//! +//! Every case builds its own relay, echo server, and tunnel on reserved loopback +//! ports, so the four transport/codec combinations run concurrently and none of +//! them collides with a relay already running on the machine. + +use std::net::SocketAddr; +use std::path::PathBuf; use std::sync::LazyLock; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use pb_mapper_auth::{AuthConfig, LegacyProtocolPolicy}; -use pb_mapper_client::client::run_client_side_cli; -use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli}; +use pb_mapper_auth::{AuthConfig, AuthRuntime, LegacyProtocolPolicy}; +use pb_mapper_client::client::run_client_side_cli_with_pinned_credential; +use pb_mapper_client::client::status::get_status_with_credential; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; +use pb_mapper_core::checksum::{Credential, set_process_msg_header_key}; use pb_mapper_core::config::init_tracing; +use pb_mapper_protocol::command::{PbConnStatusReq, PbConnStatusResp}; use pb_mapper_protocol::{MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter}; -use pb_mapper_server::run_server_with_auth_config; +use pb_mapper_server::run_server_on_listener; use rand::RngExt; use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::UdpSocket; +use tokio::net::{TcpListener, TcpStream, UdpSocket}; +use tokio::task::JoinHandle; use tokio::time::{Instant, timeout}; use tokio_util::sync::CancellationToken; use uni_stream::addr::ToSocketAddrs; -use uni_stream::stream::{ListenerProvider, TcpListenerProvider, UdpListenerProvider}; -use uni_stream::stream::{StreamProvider, StreamSplit, TcpStreamProvider, UdpStreamProvider}; +use uni_stream::stream::{ + StreamProvider, StreamSplit, TcpListenerProvider, TcpStreamProvider, UdpListenerProvider, + UdpStreamProvider, +}; use uni_stream::udp::tune_udp_socket; -struct TimerTickGurad<'a> { +const TEST_ADMIN_KEY: &str = "0123456789abcdefghijklmnopqrstuv"; +const UDP_TEST_PAYLOAD_MAX: usize = 1200; +const PROBE: &[u8] = b"pb-mapper-probe"; +/// Every readiness probe polls until this deadline before failing the test. +const READY_TIMEOUT: Duration = Duration::from_secs(10); + +static TEST_ENV: LazyLock<()> = LazyLock::new(|| { + init_tracing(); + // `run_echo_delay` frames its payload with the process credential's checksum, + // so one must be configured before any case runs. Every case uses the same + // administrator key, which keeps this a one-time write. + set_process_msg_header_key(Some(TEST_ADMIN_KEY)).unwrap(); +}); + +fn admin_credential() -> Credential { + Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Transport { + Tcp, + Udp, +} + +impl Transport { + fn name(self) -> &'static str { + match self { + Self::Tcp => "tcp", + Self::Udp => "udp", + } + } + + fn is_datagram(self) -> bool { + self == Self::Udp + } +} + +struct TimerTickGuard<'a> { ins: Instant, mut_duration: &'a mut Duration, } -impl<'a> TimerTickGurad<'a> { +impl<'a> TimerTickGuard<'a> { fn new(mut_duration: &'a mut Duration) -> Self { Self { ins: Instant::now(), @@ -35,200 +86,398 @@ impl<'a> TimerTickGurad<'a> { } } -impl<'a> Drop for TimerTickGurad<'a> { +impl Drop for TimerTickGuard<'_> { fn drop(&mut self) { - let end = Instant::now(); - let duration = end - self.ins; + let duration = Instant::now() - self.ins; *self.mut_duration += duration; println!("duration:{duration:?}"); } } -use uni_stream::stream::StreamAccept; +/// Reserve a loopback port by binding it and immediately dropping the socket. +/// +/// TCP and UDP have separate port spaces, so the reservation has to use the same +/// protocol the caller will bind. The relay keeps its own listener (see +/// [`TunnelHarness::start`]); this is only for `connect`, whose bind happens +/// inside the client and cannot be handed a pre-bound socket. +async fn reserve_addr(transport: Transport) -> SocketAddr { + match transport { + Transport::Tcp => { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + addr + } + Transport::Udp => { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = socket.local_addr().unwrap(); + drop(socket); + addr + } + } +} -const UDP_TEST_PAYLOAD_MAX: usize = 1200; +fn auth_config(label: &str) -> AuthConfig { + static SEQUENCE: AtomicUsize = AtomicUsize::new(0); + let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed); + AuthConfig { + state_dir: std::env::temp_dir().join(format!( + "pb-mapper-delay-{}-{label}-{sequence}", + std::process::id() + )), + max_temporary_keys: 64, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + } +} -async fn echo_server( - server_addr: &str, -) -> Result<(), Box> { - let listener = P::bind(server_addr).await?; - println!("run echo server:{server_addr}"); +async fn tcp_echo_server(listener: TcpListener) { loop { - // Accept incoming connections - let (mut stream, addr) = listener.accept().await?; - println!("Connected from {addr}"); - - // Process each connection concurrently + let Ok((mut stream, peer)) = listener.accept().await else { + return; + }; tokio::spawn(async move { - // Read data from client - let mut buf = vec![0; 1024]; + let mut buf = vec![0u8; 4096]; loop { - let n = match stream.read(&mut buf).await { - Ok(n) => n, - Err(e) => { - println!("Error reading: {e}"); - return; + match stream.read(&mut buf).await { + Ok(0) | Err(_) => return, + Ok(n) => { + if stream.write_all(&buf[..n]).await.is_err() { + return; + } + tracing::debug!("echoed {n} bytes to {peer}"); } - }; - - // If no data received, assume disconnect - if n == 0 { - return; - } - - // Echo data back to client - if let Err(e) = stream.write_all(&buf[..n]).await { - println!("Error writing: {e}"); - return; } - - println!("Echoed {n} bytes to {addr}"); } }); } } -async fn run_echo_server( - server_type: ServerType, - addr: &str, -) -> Result<(), Box> { - match server_type { - ServerType::Udp => run_udp_echo_server(addr).await, - ServerType::Tcp => echo_server::(addr).await, - } -} - -async fn run_udp_echo_server(addr: &str) -> Result<(), Box> { - let socket = UdpSocket::bind(addr).await?; +async fn udp_echo_server(socket: UdpSocket) { tune_udp_socket(&socket); let mut buf = vec![0u8; 65_507]; - println!("run udp echo server:{addr}"); loop { - let (len, peer) = socket.recv_from(&mut buf).await?; + let Ok((len, peer)) = socket.recv_from(&mut buf).await else { + return; + }; if let Err(err) = socket.send_to(&buf[..len], peer).await { - println!("udp echo send error: {err}"); + tracing::warn!("udp echo send error: {err}"); } } } -async fn run_pb_mapper_server(addr: &str) { - let port = addr - .rsplit_once(':') - .and_then(|(_, port)| port.parse::().ok()) - .unwrap_or_default(); - let auth_config = AuthConfig { - state_dir: std::env::temp_dir() - .join(format!("pb-mapper-delay-{}-{port}", std::process::id())), - max_temporary_keys: 64, - max_temporary_key_ttl: Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - }; - if let Err(e) = - run_server_with_auth_config(addr, CancellationToken::new(), None, false, auth_config).await - { - eprintln!("pb-mapper server failed to start: {e}"); +/// One complete tunnel: relay, echo server, `register`, and `connect`. +/// +/// `start` returns only once the tunnel carries traffic, so the tests contain no +/// timing assumptions. `Drop` tears the four tasks down and removes the relay's +/// authentication state directory. +struct TunnelHarness { + /// Where a test client sends its traffic — the address `connect` listens on. + tunnel_addr: SocketAddr, + transport: Transport, + shutdown: CancellationToken, + tasks: Vec>, + state_dir: PathBuf, +} + +impl TunnelHarness { + async fn start(transport: Transport, need_codec: bool) -> Self { + *TEST_ENV; + let label = format!( + "{}-{}", + transport.name(), + if need_codec { "codec" } else { "plain" } + ); + let config = auth_config(&label); + let _ = std::fs::remove_dir_all(&config.state_dir); + let state_dir = config.state_dir.clone(); + let credential = admin_credential(); + + // The relay gets its listener pre-bound, so nothing can take the port + // between reserving it and the relay's own bind. + let relay_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let relay_addr = relay_listener.local_addr().unwrap(); + let auth = AuthRuntime::start( + *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), + config, + ) + .await + .unwrap(); + + let shutdown = CancellationToken::new(); + let relay_shutdown = shutdown.clone(); + let mut tasks = Vec::new(); + tasks.push(tokio::spawn(async move { + if let Err(error) = + run_server_on_listener(relay_listener, relay_shutdown, None, false, auth).await + { + tracing::error!("relay stopped: {error}"); + } + })); + + // The echo server also owns its socket from the start. + let echo_addr = match transport { + Transport::Tcp => { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tasks.push(tokio::spawn(tcp_echo_server(listener))); + addr + } + Transport::Udp => { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = socket.local_addr().unwrap(); + tasks.push(tokio::spawn(udp_echo_server(socket))); + addr + } + }; + + let service_key = format!("echo-{label}"); + let options = ServerTunnelOptions { + need_codec, + is_datagram: transport.is_datagram(), + keep_alive: false, + namespace: None, + force_namespace: false, + }; + let register_key = service_key.clone(); + tasks.push(tokio::spawn(async move { + match transport { + Transport::Tcp => { + run_server_side_cli_with_pinned_credential::( + echo_addr, + relay_addr, + register_key.into(), + options, + None, + credential, + ) + .await + } + Transport::Udp => { + run_server_side_cli_with_pinned_credential::( + echo_addr, + relay_addr, + register_key.into(), + options, + None, + credential, + ) + .await + } + } + })); + + // `register` must have published the key before `connect` probes for it; + // otherwise `connect` burns its backoff waiting. + wait_for_registration(relay_addr, &service_key, credential).await; + + let tunnel_addr = reserve_addr(transport).await; + let connect_key = service_key.clone(); + tasks.push(tokio::spawn(async move { + match transport { + Transport::Tcp => { + run_client_side_cli_with_pinned_credential::( + tunnel_addr, + relay_addr, + connect_key.into(), + false, + None, + credential, + ) + .await + } + Transport::Udp => { + run_client_side_cli_with_pinned_credential::( + tunnel_addr, + relay_addr, + connect_key.into(), + false, + None, + credential, + ) + .await + } + } + })); + + let harness = Self { + tunnel_addr, + transport, + shutdown, + tasks, + state_dir, + }; + harness.wait_until_forwarding().await; + harness } } -async fn run_pb_mapper_server_cli( - server_type: ServerType, - local_addr: &str, - remote_addr: &str, - key: &str, - need_codec: bool, -) { - match server_type { - ServerType::Udp => { - run_server_side_cli::( - local_addr, - remote_addr, - key.into(), - ServerTunnelOptions { - need_codec, - is_datagram: true, - keep_alive: false, - namespace: None, - force_namespace: false, - }, - ) - .await +impl TunnelHarness { + /// Send one probe payload through the tunnel and wait for it to come back. + /// + /// This replaces the fixed sleeps the earlier version of this file used: it is + /// true end-to-end readiness, covering the relay, `register`'s control + /// connection, `connect`'s local listener, and the echo server at once. + async fn wait_until_forwarding(&self) { + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match self.probe_once().await { + Ok(()) => return, + Err(error) => last_error = error, + } + tokio::time::sleep(Duration::from_millis(50)).await; } - ServerType::Tcp => { - run_server_side_cli::( - local_addr, - remote_addr, - key.into(), - ServerTunnelOptions { - need_codec, - is_datagram: false, - keep_alive: false, - namespace: None, - force_namespace: false, - }, - ) - .await + panic!( + "{} tunnel at {} never forwarded traffic: {last_error}", + self.transport.name(), + self.tunnel_addr + ); + } + + async fn probe_once(&self) -> Result<(), String> { + match self.transport { + Transport::Tcp => self.probe_tcp().await, + Transport::Udp => self.probe_udp().await, } } -} -async fn run_pb_mapper_client_cli( - server_type: ServerType, - local_addr: &str, - remote_addr: &str, - key: &str, -) { - match server_type { - ServerType::Udp => { - run_client_side_cli::( - local_addr.to_string(), - remote_addr.to_string(), - key.into(), - false, - ) + async fn probe_tcp(&self) -> Result<(), String> { + let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(self.tunnel_addr)) .await - } - ServerType::Tcp => { - run_client_side_cli::( - local_addr.to_string(), - remote_addr.to_string(), - key.into(), - false, - ) + .map_err(|_| "tcp connect timed out".to_string())? + .map_err(|error| format!("tcp connect failed: {error}"))?; + let (mut reader, mut writer) = stream.split(); + let mut reader = NormalMessageReader::new(&mut reader); + let mut writer = NormalMessageWriter::new(&mut writer); + writer + .write_msg(PROBE) .await + .map_err(|error| format!("tcp probe write failed: {error}"))?; + let echoed = timeout(Duration::from_secs(1), reader.read_msg()) + .await + .map_err(|_| "tcp probe read timed out".to_string())? + .map_err(|error| format!("tcp probe read failed: {error}"))?; + if echoed == PROBE { + Ok(()) + } else { + Err(format!("tcp probe echoed {} bytes", echoed.len())) } } -} -/// get random message -fn gen_random_msg(max_len: usize) -> Vec { - let len = rand::rng().random_range(0_usize..max_len); - let mut vec = Vec::new(); - for _ in 0..len { - vec.push(rand::rng().random_range(0..212)); + async fn probe_udp(&self) -> Result<(), String> { + let socket = connected_udp_socket(self.tunnel_addr).await; + probe_udp_socket(&socket).await } - vec } -async fn run_udp_datagram_echo(addr: &str, rounds: usize, burst: usize) { +async fn connected_udp_socket(addr: SocketAddr) -> UdpSocket { let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); tune_udp_socket(&socket); socket.connect(addr).await.unwrap(); + socket +} + +/// Round-trip one probe datagram on `socket`. +/// +/// The tunnel keys UDP streams by source address, so every new socket starts a +/// new stream and its first datagram can be dropped while the relay sets that +/// stream up. Callers that go on to assert payload equality warm their own +/// socket with this first. +async fn probe_udp_socket(socket: &UdpSocket) -> Result<(), String> { + socket + .send(PROBE) + .await + .map_err(|error| format!("udp probe send failed: {error}"))?; let mut buf = vec![0u8; 65_507]; + let len = timeout(Duration::from_millis(500), socket.recv(&mut buf)) + .await + .map_err(|_| "udp probe timed out".to_string())? + .map_err(|error| format!("udp probe recv failed: {error}"))?; + if &buf[..len] == PROBE { + Ok(()) + } else { + Err(format!("udp probe echoed {len} bytes")) + } +} - // Warm up UDP path to ensure listener and forwarding pipeline are ready. - let probe = b"pb-mapper-probe"; - let mut ready = false; - for _ in 0..10 { - socket.send(probe).await.unwrap(); - if let Ok(Ok(len)) = timeout(Duration::from_millis(300), socket.recv(&mut buf)).await - && &buf[..len] == probe - { - ready = true; - break; +async fn wait_until_udp_socket_forwards(socket: &UdpSocket) { + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match probe_udp_socket(socket).await { + Ok(()) => return, + Err(error) => last_error = error, } tokio::time::sleep(Duration::from_millis(50)).await; } - assert!(ready, "udp echo path not ready"); + panic!("udp socket never forwarded through the tunnel: {last_error}"); +} + +impl Drop for TunnelHarness { + fn drop(&mut self) { + self.shutdown.cancel(); + for task in &self.tasks { + task.abort(); + } + let _ = std::fs::remove_dir_all(&self.state_dir); + } +} + +/// Poll the relay's `Keys` status until `register` has published `service_key`. +async fn wait_for_registration(relay_addr: SocketAddr, service_key: &str, credential: Credential) { + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match registered_keys(relay_addr, credential).await { + Ok(keys) => { + if keys.iter().any(|key| key == service_key) { + return; + } + last_error = format!("relay reports keys {keys:?}"); + } + Err(error) => last_error = error, + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("`{service_key}` was never registered: {last_error}"); +} + +async fn registered_keys( + relay_addr: SocketAddr, + credential: Credential, +) -> Result, String> { + let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(relay_addr)) + .await + .map_err(|_| "status connect timed out".to_string())? + .map_err(|error| format!("status connect failed: {error}"))?; + let response = timeout( + Duration::from_secs(1), + get_status_with_credential(&mut stream, PbConnStatusReq::Keys, None, &credential), + ) + .await + .map_err(|_| "status request timed out".to_string())? + .map_err(|error| format!("status request failed: {error}"))?; + match response { + PbConnStatusResp::Keys(keys) => Ok(keys), + other => Err(format!("unexpected status response: {other:?}")), + } +} + +/// Random payload; the length is random too, so framing is exercised at many sizes. +fn gen_random_msg(max_len: usize) -> Vec { + let len = rand::rng().random_range(0_usize..max_len); + let mut vec = Vec::with_capacity(len); + for _ in 0..len { + vec.push(rand::rng().random_range(0..212)); + } + vec +} + +async fn run_udp_datagram_echo(addr: SocketAddr, rounds: usize, burst: usize) { + let socket = connected_udp_socket(addr).await; + wait_until_udp_socket_forwards(&socket).await; + let mut buf = vec![0u8; 65_507]; for round in 0..rounds { for seq in 0..burst { @@ -239,7 +488,7 @@ async fn run_udp_datagram_echo(addr: &str, rounds: usize, burst: usize) { msg.extend(gen_random_msg(UDP_TEST_PAYLOAD_MAX)); } socket.send(&msg).await.unwrap(); - let deadline = Instant::now() + Duration::from_secs(3); + let deadline = Instant::now() + Duration::from_secs(5); loop { let wait = deadline.saturating_duration_since(Instant::now()); let len = match timeout(wait, socket.recv(&mut buf)).await { @@ -250,8 +499,9 @@ async fn run_udp_datagram_echo(addr: &str, rounds: usize, burst: usize) { if len < 4 { continue; } - let recv_seq = u32::from_be_bytes(buf[..4].try_into().unwrap()); - if recv_seq != seq as u32 { + // A datagram from an earlier sequence number may still be in + // flight; skip it rather than failing the comparison. + if u32::from_be_bytes(buf[..4].try_into().unwrap()) != seq as u32 { continue; } assert_eq!(msg, &buf[..len]); @@ -271,7 +521,7 @@ async fn run_echo_delay(addr: A, tim let expected = gen_random_msg(2000); for _ in 0..10 { let msg = { - let _guard = TimerTickGurad::new(&mut duration); + let _guard = TimerTickGuard::new(&mut duration); writer.write_msg(&expected).await.unwrap(); reader.read_msg().await.unwrap() }; @@ -282,136 +532,33 @@ async fn run_echo_delay(addr: A, tim println!("{times} rounds of 10 random data echo delay tests each took a total of {duration:?}"); } -#[derive(Debug, Clone, Copy)] -enum ServerType { - Udp, - Tcp, -} - -static PB_MAPPER_SERVER: LazyLock = - LazyLock::new(|| env::var("PB_MAPPER_TEST_SERVER").unwrap()); - -static LOCAL_SERVER: LazyLock = LazyLock::new(|| env::var("LOCAL_TEST_SERVER").unwrap()); - -static ECHO_SERVER: LazyLock = LazyLock::new(|| env::var("ECHO_TEST_SERVER").unwrap()); - -static SERVER_KEY: LazyLock = LazyLock::new(|| env::var("SERVER_TEST_KEY").unwrap()); - -static SERVER_TYPE: LazyLock = LazyLock::new(|| { - if env::var("SERVER_TEST_TYPE").unwrap() == "UDP" { - ServerType::Udp - } else { - ServerType::Tcp +/// Push traffic through the tunnel and assert every payload comes back byte for byte. +/// +/// These cases verify the logic. For latency numbers, run a separate binary. +async fn assert_tunnel_echoes(transport: Transport, need_codec: bool) { + let harness = TunnelHarness::start(transport, need_codec).await; + match transport { + Transport::Tcp => run_echo_delay::(harness.tunnel_addr, 10).await, + Transport::Udp => run_udp_datagram_echo(harness.tunnel_addr, 10, 8).await, } -}); - -static INIT_TRACING: LazyLock<()> = LazyLock::new(|| { - println!("{:?}", env::current_dir().unwrap()); - dotenvy::from_filename(env::current_dir().unwrap().join("tests").join(".env")).unwrap(); - init_tracing(); -}); +} -/// This is only for testing the correctness of the logic, for performance testing of latency, -/// please run a separate binary. -#[ignore = "run codec test enough"] #[tokio::test] -async fn test_pb_mapper_server_no_codec() { - *INIT_TRACING; - // run echo server - let remote_echo = ECHO_SERVER.clone(); - let server_type = *SERVER_TYPE; - let pb_mapper_server = PB_MAPPER_SERVER.clone(); - let server_key = SERVER_KEY.clone(); - let echo_server = ECHO_SERVER.clone(); - let local_server = LOCAL_SERVER.clone(); - - let echo_server_handle = - tokio::spawn(async move { run_echo_server(server_type, &remote_echo).await.unwrap() }); - // run the pb-mapper server role - let pb_server = pb_mapper_server.clone(); - let pb_mapper_server_handle = tokio::spawn(async move { - run_pb_mapper_server(&pb_server).await; - }); - // slepp some time to wait for pb server - tokio::time::sleep(Duration::from_millis(200)).await; - // run subcribe server cli - let key = server_key.clone(); - let subcribe_remote = pb_mapper_server.clone(); - let pb_mapper_server_cli_handle = tokio::spawn(async move { - run_pb_mapper_server_cli(server_type, &echo_server, &subcribe_remote, &key, false).await; - }); - // slepp some time to wait for pb server cli - tokio::time::sleep(Duration::from_millis(200)).await; - // run register client cli - let key = server_key.clone(); - let local_echo = local_server.clone(); - let register_remote = pb_mapper_server.clone(); - let pb_mapper_client_cli_handle = tokio::spawn(async move { - run_pb_mapper_client_cli(server_type, &local_echo, ®ister_remote, &key).await; - }); - // slepp some time to wait for pb client cli - tokio::time::sleep(Duration::from_millis(200)).await; - // run echo test - match server_type { - ServerType::Udp => run_udp_datagram_echo(local_server.as_str(), 10, 8).await, - ServerType::Tcp => run_echo_delay::(local_server.as_str(), 10).await, - } +async fn tcp_tunnel_echoes_without_codec() { + assert_tunnel_echoes(Transport::Tcp, false).await; +} - // abort all thread - echo_server_handle.abort(); - pb_mapper_server_handle.abort(); - pb_mapper_server_cli_handle.abort(); - pb_mapper_client_cli_handle.abort(); +#[tokio::test] +async fn tcp_tunnel_echoes_with_codec() { + assert_tunnel_echoes(Transport::Tcp, true).await; } -/// This is only for testing the correctness of the logic, for performance testing of latency, -/// please run a separate binary. #[tokio::test] -async fn test_pb_mapper_server_codec() { - *INIT_TRACING; - // run echo server - let remote_echo = ECHO_SERVER.clone(); - let server_type = *SERVER_TYPE; - let pb_mapper_server = PB_MAPPER_SERVER.clone(); - let server_key = SERVER_KEY.clone(); - let echo_server = ECHO_SERVER.clone(); - let local_server = LOCAL_SERVER.clone(); - - let echo_server_handle = - tokio::spawn(async move { run_echo_server(server_type, &remote_echo).await.unwrap() }); - // run the pb-mapper server role - let pb_server = pb_mapper_server.clone(); - let pb_mapper_server_handle = tokio::spawn(async move { - run_pb_mapper_server(&pb_server).await; - }); - // slepp some time to wait for pb server - tokio::time::sleep(Duration::from_millis(200)).await; - // run subcribe server cli - let key = server_key.clone(); - let subcribe_remote = pb_mapper_server.clone(); - let pb_mapper_server_cli_handle = tokio::spawn(async move { - run_pb_mapper_server_cli(server_type, &echo_server, &subcribe_remote, &key, true).await; - }); - // slepp some time to wait for pb server cli - tokio::time::sleep(Duration::from_millis(200)).await; - // run register client cli - let key = server_key.clone(); - let local_echo = local_server.clone(); - let register_remote = pb_mapper_server.clone(); - let pb_mapper_client_cli_handle = tokio::spawn(async move { - run_pb_mapper_client_cli(server_type, &local_echo, ®ister_remote, &key).await; - }); - // slepp some time to wait for pb client cli - tokio::time::sleep(Duration::from_millis(200)).await; - // run echo test - match server_type { - ServerType::Udp => run_udp_datagram_echo(local_server.as_str(), 10, 8).await, - ServerType::Tcp => run_echo_delay::(local_server.as_str(), 10).await, - } +async fn udp_tunnel_echoes_without_codec() { + assert_tunnel_echoes(Transport::Udp, false).await; +} - // abort all thread - echo_server_handle.abort(); - pb_mapper_server_handle.abort(); - pb_mapper_server_cli_handle.abort(); - pb_mapper_client_cli_handle.abort(); +#[tokio::test] +async fn udp_tunnel_echoes_with_codec() { + assert_tunnel_echoes(Transport::Udp, true).await; } From aa5775c0a68adc2f3f08c5d26c142c7b4b8788b2 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:19:47 +0800 Subject: [PATCH 2/7] Extract the e2e tunnel harness into pb-mapper-testkit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The harness lived privately inside `test_delay.rs`, so a second test file wanting a full `server` + `register` + `connect` flow had two options: crowd into that file, or copy it. This makes it a crate any test can depend on. Two levels, because the cases need different things. `TunnelHarness::start` is the one-liner for "give me a working tunnel". `Relay` and `TunnelSpec` split beneath it for cases that must act on the relay first — a relay retains its `AuthRuntime` clone, which both lets a test issue, renew and revoke credentials without the admin wire protocol, and keeps the actor's command channel open, since dropping every clone cancels all leases. A real crate rather than `tests/common/mod.rs`: that module compiles separately into every test binary, and whatever a given binary does not use is reported as dead code — fatal under `-D warnings`. Traffic drivers come in framed and raw pairs. `NormalMessageReader` writes a checksum + length header, while the local side of a tunnel is byte transparent, so an echo server that prepends a tag byte shifts every frame header. Framed drivers need a transparent echo; tagged tunnels need the raw ones. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 17 + Cargo.toml | 2 + crates/pb-mapper-cli/Cargo.toml | 2 + crates/pb-mapper-cli/tests/test_delay.rs | 534 +---------------------- crates/pb-mapper-testkit/Cargo.toml | 30 ++ crates/pb-mapper-testkit/src/echo.rs | 54 +++ crates/pb-mapper-testkit/src/lib.rs | 163 +++++++ crates/pb-mapper-testkit/src/relay.rs | 182 ++++++++ crates/pb-mapper-testkit/src/traffic.rs | 223 ++++++++++ crates/pb-mapper-testkit/src/tunnel.rs | 378 ++++++++++++++++ 10 files changed, 1058 insertions(+), 527 deletions(-) create mode 100644 crates/pb-mapper-testkit/Cargo.toml create mode 100644 crates/pb-mapper-testkit/src/echo.rs create mode 100644 crates/pb-mapper-testkit/src/lib.rs create mode 100644 crates/pb-mapper-testkit/src/relay.rs create mode 100644 crates/pb-mapper-testkit/src/traffic.rs create mode 100644 crates/pb-mapper-testkit/src/tunnel.rs diff --git a/Cargo.lock b/Cargo.lock index 3d848a0..e206df9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1050,6 +1050,7 @@ dependencies = [ "pb-mapper-core", "pb-mapper-protocol", "pb-mapper-server", + "pb-mapper-testkit", "rand 0.10.0", "serde_json", "tokio", @@ -1145,6 +1146,22 @@ dependencies = [ "uni-stream", ] +[[package]] +name = "pb-mapper-testkit" +version = "0.4.0" +dependencies = [ + "pb-mapper-auth", + "pb-mapper-client", + "pb-mapper-core", + "pb-mapper-protocol", + "pb-mapper-server", + "rand 0.10.0", + "tokio", + "tokio-util", + "tracing", + "uni-stream", +] + [[package]] name = "percent-encoding" version = "2.3.2" diff --git a/Cargo.toml b/Cargo.toml index 47e43fb..e37aa76 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,6 +24,8 @@ pb-mapper-client = { path = "crates/pb-mapper-client" } pb-mapper-core = { path = "crates/pb-mapper-core" } pb-mapper-protocol = { path = "crates/pb-mapper-protocol" } pb-mapper-server = { path = "crates/pb-mapper-server" } +# Test support only; never a dependency of a shipped crate. +pb-mapper-testkit = { path = "crates/pb-mapper-testkit" } base64 = "0.23.1" better_mimalloc_rs = { version = "0.1.2", features = ["config"] } diff --git a/crates/pb-mapper-cli/Cargo.toml b/crates/pb-mapper-cli/Cargo.toml index 07fb2ee..0ce404b 100644 --- a/crates/pb-mapper-cli/Cargo.toml +++ b/crates/pb-mapper-cli/Cargo.toml @@ -25,6 +25,8 @@ tracing.workspace = true uni-stream.workspace = true [dev-dependencies] +pb-mapper-testkit.workspace = true + rand.workspace = true [features] diff --git a/crates/pb-mapper-cli/tests/test_delay.rs b/crates/pb-mapper-cli/tests/test_delay.rs index 39e9806..afed920 100644 --- a/crates/pb-mapper-cli/tests/test_delay.rs +++ b/crates/pb-mapper-cli/tests/test_delay.rs @@ -5,532 +5,12 @@ //! //! Every case builds its own relay, echo server, and tunnel on reserved loopback //! ports, so the four transport/codec combinations run concurrently and none of -//! them collides with a relay already running on the machine. +//! them collides with a relay already running on the machine. The scaffolding +//! lives in `pb-mapper-testkit`, so any other test file can stand up the same +//! flow — see `temporary_credential_e2e.rs`. -use std::net::SocketAddr; -use std::path::PathBuf; -use std::sync::LazyLock; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::time::Duration; - -use pb_mapper_auth::{AuthConfig, AuthRuntime, LegacyProtocolPolicy}; -use pb_mapper_client::client::run_client_side_cli_with_pinned_credential; -use pb_mapper_client::client::status::get_status_with_credential; -use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; -use pb_mapper_core::checksum::{Credential, set_process_msg_header_key}; -use pb_mapper_core::config::init_tracing; -use pb_mapper_protocol::command::{PbConnStatusReq, PbConnStatusResp}; -use pb_mapper_protocol::{MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter}; -use pb_mapper_server::run_server_on_listener; -use rand::RngExt; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::{TcpListener, TcpStream, UdpSocket}; -use tokio::task::JoinHandle; -use tokio::time::{Instant, timeout}; -use tokio_util::sync::CancellationToken; -use uni_stream::addr::ToSocketAddrs; -use uni_stream::stream::{ - StreamProvider, StreamSplit, TcpListenerProvider, TcpStreamProvider, UdpListenerProvider, - UdpStreamProvider, -}; -use uni_stream::udp::tune_udp_socket; - -const TEST_ADMIN_KEY: &str = "0123456789abcdefghijklmnopqrstuv"; -const UDP_TEST_PAYLOAD_MAX: usize = 1200; -const PROBE: &[u8] = b"pb-mapper-probe"; -/// Every readiness probe polls until this deadline before failing the test. -const READY_TIMEOUT: Duration = Duration::from_secs(10); - -static TEST_ENV: LazyLock<()> = LazyLock::new(|| { - init_tracing(); - // `run_echo_delay` frames its payload with the process credential's checksum, - // so one must be configured before any case runs. Every case uses the same - // administrator key, which keeps this a one-time write. - set_process_msg_header_key(Some(TEST_ADMIN_KEY)).unwrap(); -}); - -fn admin_credential() -> Credential { - Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum Transport { - Tcp, - Udp, -} - -impl Transport { - fn name(self) -> &'static str { - match self { - Self::Tcp => "tcp", - Self::Udp => "udp", - } - } - - fn is_datagram(self) -> bool { - self == Self::Udp - } -} - -struct TimerTickGuard<'a> { - ins: Instant, - mut_duration: &'a mut Duration, -} - -impl<'a> TimerTickGuard<'a> { - fn new(mut_duration: &'a mut Duration) -> Self { - Self { - ins: Instant::now(), - mut_duration, - } - } -} - -impl Drop for TimerTickGuard<'_> { - fn drop(&mut self) { - let duration = Instant::now() - self.ins; - *self.mut_duration += duration; - println!("duration:{duration:?}"); - } -} - -/// Reserve a loopback port by binding it and immediately dropping the socket. -/// -/// TCP and UDP have separate port spaces, so the reservation has to use the same -/// protocol the caller will bind. The relay keeps its own listener (see -/// [`TunnelHarness::start`]); this is only for `connect`, whose bind happens -/// inside the client and cannot be handed a pre-bound socket. -async fn reserve_addr(transport: Transport) -> SocketAddr { - match transport { - Transport::Tcp => { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - drop(listener); - addr - } - Transport::Udp => { - let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - let addr = socket.local_addr().unwrap(); - drop(socket); - addr - } - } -} - -fn auth_config(label: &str) -> AuthConfig { - static SEQUENCE: AtomicUsize = AtomicUsize::new(0); - let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed); - AuthConfig { - state_dir: std::env::temp_dir().join(format!( - "pb-mapper-delay-{}-{label}-{sequence}", - std::process::id() - )), - max_temporary_keys: 64, - max_temporary_key_ttl: Duration::from_secs(3600), - legacy_protocol: LegacyProtocolPolicy::Allow, - } -} - -async fn tcp_echo_server(listener: TcpListener) { - loop { - let Ok((mut stream, peer)) = listener.accept().await else { - return; - }; - tokio::spawn(async move { - let mut buf = vec![0u8; 4096]; - loop { - match stream.read(&mut buf).await { - Ok(0) | Err(_) => return, - Ok(n) => { - if stream.write_all(&buf[..n]).await.is_err() { - return; - } - tracing::debug!("echoed {n} bytes to {peer}"); - } - } - } - }); - } -} - -async fn udp_echo_server(socket: UdpSocket) { - tune_udp_socket(&socket); - let mut buf = vec![0u8; 65_507]; - loop { - let Ok((len, peer)) = socket.recv_from(&mut buf).await else { - return; - }; - if let Err(err) = socket.send_to(&buf[..len], peer).await { - tracing::warn!("udp echo send error: {err}"); - } - } -} - -/// One complete tunnel: relay, echo server, `register`, and `connect`. -/// -/// `start` returns only once the tunnel carries traffic, so the tests contain no -/// timing assumptions. `Drop` tears the four tasks down and removes the relay's -/// authentication state directory. -struct TunnelHarness { - /// Where a test client sends its traffic — the address `connect` listens on. - tunnel_addr: SocketAddr, - transport: Transport, - shutdown: CancellationToken, - tasks: Vec>, - state_dir: PathBuf, -} - -impl TunnelHarness { - async fn start(transport: Transport, need_codec: bool) -> Self { - *TEST_ENV; - let label = format!( - "{}-{}", - transport.name(), - if need_codec { "codec" } else { "plain" } - ); - let config = auth_config(&label); - let _ = std::fs::remove_dir_all(&config.state_dir); - let state_dir = config.state_dir.clone(); - let credential = admin_credential(); - - // The relay gets its listener pre-bound, so nothing can take the port - // between reserving it and the relay's own bind. - let relay_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let relay_addr = relay_listener.local_addr().unwrap(); - let auth = AuthRuntime::start( - *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap(), - config, - ) - .await - .unwrap(); - - let shutdown = CancellationToken::new(); - let relay_shutdown = shutdown.clone(); - let mut tasks = Vec::new(); - tasks.push(tokio::spawn(async move { - if let Err(error) = - run_server_on_listener(relay_listener, relay_shutdown, None, false, auth).await - { - tracing::error!("relay stopped: {error}"); - } - })); - - // The echo server also owns its socket from the start. - let echo_addr = match transport { - Transport::Tcp => { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tasks.push(tokio::spawn(tcp_echo_server(listener))); - addr - } - Transport::Udp => { - let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - let addr = socket.local_addr().unwrap(); - tasks.push(tokio::spawn(udp_echo_server(socket))); - addr - } - }; - - let service_key = format!("echo-{label}"); - let options = ServerTunnelOptions { - need_codec, - is_datagram: transport.is_datagram(), - keep_alive: false, - namespace: None, - force_namespace: false, - }; - let register_key = service_key.clone(); - tasks.push(tokio::spawn(async move { - match transport { - Transport::Tcp => { - run_server_side_cli_with_pinned_credential::( - echo_addr, - relay_addr, - register_key.into(), - options, - None, - credential, - ) - .await - } - Transport::Udp => { - run_server_side_cli_with_pinned_credential::( - echo_addr, - relay_addr, - register_key.into(), - options, - None, - credential, - ) - .await - } - } - })); - - // `register` must have published the key before `connect` probes for it; - // otherwise `connect` burns its backoff waiting. - wait_for_registration(relay_addr, &service_key, credential).await; - - let tunnel_addr = reserve_addr(transport).await; - let connect_key = service_key.clone(); - tasks.push(tokio::spawn(async move { - match transport { - Transport::Tcp => { - run_client_side_cli_with_pinned_credential::( - tunnel_addr, - relay_addr, - connect_key.into(), - false, - None, - credential, - ) - .await - } - Transport::Udp => { - run_client_side_cli_with_pinned_credential::( - tunnel_addr, - relay_addr, - connect_key.into(), - false, - None, - credential, - ) - .await - } - } - })); - - let harness = Self { - tunnel_addr, - transport, - shutdown, - tasks, - state_dir, - }; - harness.wait_until_forwarding().await; - harness - } -} - -impl TunnelHarness { - /// Send one probe payload through the tunnel and wait for it to come back. - /// - /// This replaces the fixed sleeps the earlier version of this file used: it is - /// true end-to-end readiness, covering the relay, `register`'s control - /// connection, `connect`'s local listener, and the echo server at once. - async fn wait_until_forwarding(&self) { - let deadline = Instant::now() + READY_TIMEOUT; - let mut last_error = String::from("no attempt completed"); - while Instant::now() < deadline { - match self.probe_once().await { - Ok(()) => return, - Err(error) => last_error = error, - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - panic!( - "{} tunnel at {} never forwarded traffic: {last_error}", - self.transport.name(), - self.tunnel_addr - ); - } - - async fn probe_once(&self) -> Result<(), String> { - match self.transport { - Transport::Tcp => self.probe_tcp().await, - Transport::Udp => self.probe_udp().await, - } - } - - async fn probe_tcp(&self) -> Result<(), String> { - let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(self.tunnel_addr)) - .await - .map_err(|_| "tcp connect timed out".to_string())? - .map_err(|error| format!("tcp connect failed: {error}"))?; - let (mut reader, mut writer) = stream.split(); - let mut reader = NormalMessageReader::new(&mut reader); - let mut writer = NormalMessageWriter::new(&mut writer); - writer - .write_msg(PROBE) - .await - .map_err(|error| format!("tcp probe write failed: {error}"))?; - let echoed = timeout(Duration::from_secs(1), reader.read_msg()) - .await - .map_err(|_| "tcp probe read timed out".to_string())? - .map_err(|error| format!("tcp probe read failed: {error}"))?; - if echoed == PROBE { - Ok(()) - } else { - Err(format!("tcp probe echoed {} bytes", echoed.len())) - } - } - - async fn probe_udp(&self) -> Result<(), String> { - let socket = connected_udp_socket(self.tunnel_addr).await; - probe_udp_socket(&socket).await - } -} - -async fn connected_udp_socket(addr: SocketAddr) -> UdpSocket { - let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - tune_udp_socket(&socket); - socket.connect(addr).await.unwrap(); - socket -} - -/// Round-trip one probe datagram on `socket`. -/// -/// The tunnel keys UDP streams by source address, so every new socket starts a -/// new stream and its first datagram can be dropped while the relay sets that -/// stream up. Callers that go on to assert payload equality warm their own -/// socket with this first. -async fn probe_udp_socket(socket: &UdpSocket) -> Result<(), String> { - socket - .send(PROBE) - .await - .map_err(|error| format!("udp probe send failed: {error}"))?; - let mut buf = vec![0u8; 65_507]; - let len = timeout(Duration::from_millis(500), socket.recv(&mut buf)) - .await - .map_err(|_| "udp probe timed out".to_string())? - .map_err(|error| format!("udp probe recv failed: {error}"))?; - if &buf[..len] == PROBE { - Ok(()) - } else { - Err(format!("udp probe echoed {len} bytes")) - } -} - -async fn wait_until_udp_socket_forwards(socket: &UdpSocket) { - let deadline = Instant::now() + READY_TIMEOUT; - let mut last_error = String::from("no attempt completed"); - while Instant::now() < deadline { - match probe_udp_socket(socket).await { - Ok(()) => return, - Err(error) => last_error = error, - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - panic!("udp socket never forwarded through the tunnel: {last_error}"); -} - -impl Drop for TunnelHarness { - fn drop(&mut self) { - self.shutdown.cancel(); - for task in &self.tasks { - task.abort(); - } - let _ = std::fs::remove_dir_all(&self.state_dir); - } -} - -/// Poll the relay's `Keys` status until `register` has published `service_key`. -async fn wait_for_registration(relay_addr: SocketAddr, service_key: &str, credential: Credential) { - let deadline = Instant::now() + READY_TIMEOUT; - let mut last_error = String::from("no attempt completed"); - while Instant::now() < deadline { - match registered_keys(relay_addr, credential).await { - Ok(keys) => { - if keys.iter().any(|key| key == service_key) { - return; - } - last_error = format!("relay reports keys {keys:?}"); - } - Err(error) => last_error = error, - } - tokio::time::sleep(Duration::from_millis(20)).await; - } - panic!("`{service_key}` was never registered: {last_error}"); -} - -async fn registered_keys( - relay_addr: SocketAddr, - credential: Credential, -) -> Result, String> { - let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(relay_addr)) - .await - .map_err(|_| "status connect timed out".to_string())? - .map_err(|error| format!("status connect failed: {error}"))?; - let response = timeout( - Duration::from_secs(1), - get_status_with_credential(&mut stream, PbConnStatusReq::Keys, None, &credential), - ) - .await - .map_err(|_| "status request timed out".to_string())? - .map_err(|error| format!("status request failed: {error}"))?; - match response { - PbConnStatusResp::Keys(keys) => Ok(keys), - other => Err(format!("unexpected status response: {other:?}")), - } -} - -/// Random payload; the length is random too, so framing is exercised at many sizes. -fn gen_random_msg(max_len: usize) -> Vec { - let len = rand::rng().random_range(0_usize..max_len); - let mut vec = Vec::with_capacity(len); - for _ in 0..len { - vec.push(rand::rng().random_range(0..212)); - } - vec -} - -async fn run_udp_datagram_echo(addr: SocketAddr, rounds: usize, burst: usize) { - let socket = connected_udp_socket(addr).await; - wait_until_udp_socket_forwards(&socket).await; - let mut buf = vec![0u8; 65_507]; - - for round in 0..rounds { - for seq in 0..burst { - let mut msg = (seq as u32).to_be_bytes().to_vec(); - if round == 0 && seq == 0 { - msg.extend(vec![0u8; UDP_TEST_PAYLOAD_MAX]); - } else { - msg.extend(gen_random_msg(UDP_TEST_PAYLOAD_MAX)); - } - socket.send(&msg).await.unwrap(); - let deadline = Instant::now() + Duration::from_secs(5); - loop { - let wait = deadline.saturating_duration_since(Instant::now()); - let len = match timeout(wait, socket.recv(&mut buf)).await { - Ok(Ok(len)) => len, - Ok(Err(err)) => panic!("udp recv error: {err}"), - Err(_) => panic!("udp echo timeout; missing seq: {seq}"), - }; - if len < 4 { - continue; - } - // A datagram from an earlier sequence number may still be in - // flight; skip it rather than failing the comparison. - if u32::from_be_bytes(buf[..4].try_into().unwrap()) != seq as u32 { - continue; - } - assert_eq!(msg, &buf[..len]); - break; - } - } - } -} - -async fn run_echo_delay(addr: A, times: usize) { - let mut stream = P::from_addr(addr).await.unwrap(); - let (mut reader, mut writer) = stream.split(); - let mut reader = NormalMessageReader::new(&mut reader); - let mut writer = NormalMessageWriter::new(&mut writer); - let mut duration = Duration::default(); - for _ in 0..times { - let expected = gen_random_msg(2000); - for _ in 0..10 { - let msg = { - let _guard = TimerTickGuard::new(&mut duration); - writer.write_msg(&expected).await.unwrap(); - reader.read_msg().await.unwrap() - }; - - assert_eq!(expected, msg); - } - } - println!("{times} rounds of 10 random data echo delay tests each took a total of {duration:?}"); -} +use pb_mapper_testkit::{Transport, TunnelHarness, run_echo_delay, run_udp_datagram_echo}; +use uni_stream::stream::TcpStreamProvider; /// Push traffic through the tunnel and assert every payload comes back byte for byte. /// @@ -538,8 +18,8 @@ async fn run_echo_delay(addr: A, tim async fn assert_tunnel_echoes(transport: Transport, need_codec: bool) { let harness = TunnelHarness::start(transport, need_codec).await; match transport { - Transport::Tcp => run_echo_delay::(harness.tunnel_addr, 10).await, - Transport::Udp => run_udp_datagram_echo(harness.tunnel_addr, 10, 8).await, + Transport::Tcp => run_echo_delay::(harness.tunnel_addr(), 10).await, + Transport::Udp => run_udp_datagram_echo(harness.tunnel_addr(), 10, 8, None).await, } } diff --git a/crates/pb-mapper-testkit/Cargo.toml b/crates/pb-mapper-testkit/Cargo.toml new file mode 100644 index 0000000..0bba95e --- /dev/null +++ b/crates/pb-mapper-testkit/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "pb-mapper-testkit" +version.workspace = true +edition.workspace = true +authors.workspace = true +publish = false + +# Test-only support crate. It stands up a complete relay + `register` + `connect` +# flow so that any integration test file can build one, instead of a single test +# file owning the harness privately. +# +# A real crate rather than `tests/common/mod.rs`: that module is compiled into +# every test binary separately, and whatever a given binary does not use is +# reported as dead code — fatal under `-D warnings`. + +[dependencies] +pb-mapper-auth.workspace = true +pb-mapper-client.workspace = true +pb-mapper-core.workspace = true +pb-mapper-protocol.workspace = true +pb-mapper-server.workspace = true + +rand.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +uni-stream.workspace = true + +[lints] +workspace = true diff --git a/crates/pb-mapper-testkit/src/echo.rs b/crates/pb-mapper-testkit/src/echo.rs new file mode 100644 index 0000000..10abc8d --- /dev/null +++ b/crates/pb-mapper-testkit/src/echo.rs @@ -0,0 +1,54 @@ +//! Echo servers a tunnel forwards to. +//! +//! Both take an already-bound socket, so the caller can learn the address without +//! a window in which another test could take the port. +//! +//! `tag` prepends one byte to every reply. Two tunnels that share a service name +//! in different namespaces get different tags, so a test can tell which echo +//! server actually received the traffic — without it, a namespace leak would +//! still satisfy payload equality, since both servers echo identically. + +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, UdpSocket}; +use uni_stream::udp::tune_udp_socket; + +pub async fn tcp_echo_server(listener: TcpListener, tag: Option) { + loop { + let Ok((mut stream, peer)) = listener.accept().await else { + return; + }; + tokio::spawn(async move { + let mut buf = vec![0u8; 4096]; + loop { + match stream.read(&mut buf).await { + Ok(0) | Err(_) => return, + Ok(n) => { + let mut reply = Vec::with_capacity(n + 1); + reply.extend(tag); + reply.extend_from_slice(&buf[..n]); + if stream.write_all(&reply).await.is_err() { + return; + } + tracing::debug!("echoed {n} bytes to {peer}"); + } + } + } + }); + } +} + +pub async fn udp_echo_server(socket: UdpSocket, tag: Option) { + tune_udp_socket(&socket); + let mut buf = vec![0u8; 65_507]; + loop { + let Ok((len, peer)) = socket.recv_from(&mut buf).await else { + return; + }; + let mut reply = Vec::with_capacity(len + 1); + reply.extend(tag); + reply.extend_from_slice(&buf[..len]); + if let Err(err) = socket.send_to(&reply, peer).await { + tracing::warn!("udp echo send error: {err}"); + } + } +} diff --git a/crates/pb-mapper-testkit/src/lib.rs b/crates/pb-mapper-testkit/src/lib.rs new file mode 100644 index 0000000..86151f3 --- /dev/null +++ b/crates/pb-mapper-testkit/src/lib.rs @@ -0,0 +1,163 @@ +//! Shared end-to-end scaffolding for the integration tests. +//! +//! The unit of reuse is a [`Relay`] — one authenticated `pb-mapper server` on a +//! reserved loopback port with its own authentication state directory — and a +//! [`Tunnel`] on top of it: an echo server, a `register`, and a `connect`. Both +//! are ordinary values with a `Drop` that tears their tasks down, so a test file +//! stands up a complete flow in one line and needs no shared fixture, no test +//! ordering, and no cleanup code. +//! +//! ```ignore +//! let harness = TunnelHarness::start(Transport::Tcp, false).await; +//! run_echo_delay::(harness.tunnel_addr(), 10).await; +//! ``` +//! +//! For anything the one-liner cannot express — a temporary credential, two +//! tunnels sharing a relay, revoking a key under a live tunnel — drive the two +//! halves separately: +//! +//! ```ignore +//! let relay = Relay::start("temp-key").await; +//! let (key_id, credential) = relay.issue_credential(Duration::from_secs(60), "tcp").await; +//! let tunnel = relay +//! .start_tunnel(TunnelSpec::new(Transport::Tcp).credential(credential)) +//! .await; +//! ``` +//! +//! Everything here is a readiness probe rather than a sleep: [`Relay`] polls the +//! `Keys` status to see a registration, and [`Tunnel::start`] returns only once a +//! payload has made the full round trip. That is why the cases carry no timing +//! assumptions and can all run concurrently. +//! +//! This is a real crate rather than a `tests/common/mod.rs` module because that +//! module is compiled separately into every test binary, and whatever a given +//! binary happens not to use is reported as dead code — fatal under +//! `-D warnings`. + +// The entire crate is test support, so the workspace's `unwrap`/`expect` denial +// is lifted here the same way it is inside a `tests/` target. `clippy.toml`'s +// test exemptions do not reach a library crate, which is not `#[cfg(test)]`. +#![allow(clippy::unwrap_used, clippy::expect_used)] + +use std::sync::LazyLock; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use pb_mapper_auth::{AuthConfig, LegacyProtocolPolicy}; +use pb_mapper_core::checksum::{Credential, set_process_msg_header_key}; +use pb_mapper_core::config::init_tracing; + +mod echo; +mod relay; +mod traffic; +mod tunnel; + +pub use relay::Relay; +pub use traffic::{ + TimerTickGuard, connected_udp_socket, gen_random_msg, probe_udp_socket, raw_tcp_probe, + run_echo_delay, run_raw_tcp_echo, run_udp_datagram_echo, tagged_echo, + wait_until_udp_socket_forwards, +}; +pub use tunnel::{Tunnel, TunnelHarness, TunnelSpec}; + +/// The administrator key every test relay starts with. +/// +/// A fixed key rather than a random one: it is also the process credential (see +/// [`init_test_env`]), and one value keeps that a single write per process. +pub const TEST_ADMIN_KEY: &str = "0123456789abcdefghijklmnopqrstuv"; + +/// How long any readiness probe polls before failing the test. +pub const READY_TIMEOUT: Duration = Duration::from_secs(10); + +/// The payload every readiness probe round-trips. +pub const PROBE: &[u8] = b"pb-mapper-probe"; + +/// Upper bound on generated UDP payloads, comfortably inside one datagram. +pub const UDP_TEST_PAYLOAD_MAX: usize = 1200; + +/// Tracing plus the process credential, set up once per test process. +/// +/// The framing helpers in this crate ([`run_echo_delay`] and the TCP probe) write +/// their own checksummed frames, and `checksum_for` fails closed without a +/// process credential — so one has to exist before any case runs. It is +/// unrelated to which credential a tunnel authenticates with: `pb-mapper`'s local +/// side is byte-transparent, so these frames are only ever read back by the same +/// process that wrote them. +pub fn init_test_env() { + static TEST_ENV: LazyLock<()> = LazyLock::new(|| { + init_tracing(); + set_process_msg_header_key(Some(TEST_ADMIN_KEY)).unwrap(); + }); + *TEST_ENV +} + +/// The administrator credential matching [`TEST_ADMIN_KEY`]. +pub fn admin_credential() -> Credential { + Credential::Admin(*TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap()) +} + +/// The 32 raw bytes of [`TEST_ADMIN_KEY`], as [`pb_mapper_auth::AuthRuntime`] wants them. +pub fn admin_key_bytes() -> [u8; 32] { + *TEST_ADMIN_KEY.as_bytes().first_chunk::<32>().unwrap() +} + +/// A private authentication state directory, unique per relay. +/// +/// Each relay takes an exclusive `flock` on `auth.lock` inside its state +/// directory, so two relays must never share one. +pub fn auth_config(label: &str) -> AuthConfig { + static SEQUENCE: AtomicUsize = AtomicUsize::new(0); + let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed); + AuthConfig { + state_dir: std::env::temp_dir().join(format!( + "pb-mapper-testkit-{}-{label}-{sequence}", + std::process::id() + )), + max_temporary_keys: 64, + max_temporary_key_ttl: Duration::from_secs(3600), + legacy_protocol: LegacyProtocolPolicy::Allow, + } +} + +/// Which transport a tunnel carries, end to end. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Transport { + Tcp, + Udp, +} + +impl Transport { + pub fn name(self) -> &'static str { + match self { + Self::Tcp => "tcp", + Self::Udp => "udp", + } + } + + pub fn is_datagram(self) -> bool { + self == Self::Udp + } +} + +/// Reserve a loopback port by binding it and immediately dropping the socket. +/// +/// TCP and UDP have separate port spaces, so the reservation has to use the same +/// protocol the caller will bind. A relay and an echo server keep the socket they +/// bound; this is for `connect`, whose bind happens inside the client and cannot +/// be handed a pre-bound socket. +pub async fn reserve_addr(transport: Transport) -> std::net::SocketAddr { + match transport { + Transport::Tcp => { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + drop(listener); + addr + } + Transport::Udp => { + let socket = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = socket.local_addr().unwrap(); + drop(socket); + addr + } + } +} diff --git a/crates/pb-mapper-testkit/src/relay.rs b/crates/pb-mapper-testkit/src/relay.rs new file mode 100644 index 0000000..357b37a --- /dev/null +++ b/crates/pb-mapper-testkit/src/relay.rs @@ -0,0 +1,182 @@ +//! A `pb-mapper server` on a reserved loopback port, with its own auth state. + +use std::net::SocketAddr; +use std::path::PathBuf; +use std::time::Duration; + +use pb_mapper_auth::{ + ADMIN_KEY_ID, AuthContext, AuthRuntime, IssuedTemporaryKey, KeyId, TemporaryKeyMetadata, +}; +use pb_mapper_client::client::status::get_status_with_credential; +use pb_mapper_core::checksum::{Credential, parse_credential}; +use pb_mapper_protocol::command::{PbConnStatusReq, PbConnStatusResp}; +use pb_mapper_server::run_server_on_listener; +use tokio::net::{TcpListener, TcpStream}; +use tokio::task::JoinHandle; +use tokio::time::{Instant, timeout}; +use tokio_util::sync::CancellationToken; + +use crate::tunnel::{Tunnel, TunnelSpec}; +use crate::{READY_TIMEOUT, admin_key_bytes, auth_config, init_test_env}; + +/// One relay: a listener, an [`AuthRuntime`], and the server loop over both. +/// +/// The runtime is kept, not moved into the server, for two reasons. Tests need it +/// to issue, renew, and revoke credentials without going through the admin wire +/// protocol; and dropping every clone closes the actor's command channel, which +/// makes it cancel every outstanding lease — so a live relay has to hold one. +pub struct Relay { + addr: SocketAddr, + auth: AuthRuntime, + admin: AuthContext, + shutdown: CancellationToken, + task: JoinHandle<()>, + state_dir: PathBuf, +} + +impl Relay { + /// Start a relay. `label` only shows up in the state directory's name. + pub async fn start(label: &str) -> Self { + init_test_env(); + let config = auth_config(label); + let _ = std::fs::remove_dir_all(&config.state_dir); + let state_dir = config.state_dir.clone(); + + // Pre-bound, so nothing can take the port between reserving it and the + // relay's own bind. + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let auth = AuthRuntime::start(admin_key_bytes(), config).await.unwrap(); + let admin = auth + .authenticate_presented(ADMIN_KEY_ID, &admin_key_bytes()) + .unwrap(); + + let shutdown = CancellationToken::new(); + let server_auth = auth.clone(); + let server_shutdown = shutdown.clone(); + let task = tokio::spawn(async move { + if let Err(error) = + run_server_on_listener(listener, server_shutdown, None, false, server_auth).await + { + tracing::error!("relay stopped: {error}"); + } + }); + + Self { + addr, + auth, + admin, + shutdown, + task, + state_dir, + } + } + + pub fn addr(&self) -> SocketAddr { + self.addr + } + + /// The administrator context, for auth calls this type does not wrap. + pub fn admin_context(&self) -> &AuthContext { + &self.admin + } + + pub fn auth(&self) -> &AuthRuntime { + &self.auth + } + + /// Issue a temporary credential and return it ready to authenticate with. + /// + /// The key ID doubles as the credential's namespace, which is what makes + /// `namespace: None` on the register and connect paths resolve to the + /// credential's own namespace. + pub async fn issue_credential(&self, ttl: Duration, label: &str) -> (KeyId, Credential) { + let issued = self.issue(ttl, label).await; + let credential = parse_credential(&issued.credential).unwrap(); + (issued.metadata.key_id, credential) + } + + /// Issue a temporary credential, keeping the lifecycle metadata. + pub async fn issue(&self, ttl: Duration, label: &str) -> IssuedTemporaryKey { + self.auth + .issue(&self.admin, ttl, Some(label.to_string())) + .await + .unwrap() + } + + /// Extend a credential's lifetime. The credential text does not change. + pub async fn renew(&self, key_id: KeyId, ttl: Duration) -> IssuedTemporaryKey { + self.auth.renew(&self.admin, key_id, ttl).await.unwrap() + } + + /// Revoke a credential, cancelling its lease and every connection under it. + pub async fn revoke(&self, key_id: KeyId) -> TemporaryKeyMetadata { + self.auth.revoke(&self.admin, key_id).await.unwrap() + } + + /// Build a tunnel — echo server, `register`, and `connect` — against this relay. + /// + /// Returns once traffic has made the full round trip. + pub async fn start_tunnel(&self, spec: TunnelSpec) -> Tunnel { + Tunnel::start(self, spec).await + } + + /// Poll the `Keys` status until `register` has published `service_key`. + pub async fn wait_for_registration( + &self, + service_key: &str, + credential: Credential, + namespace: Option, + ) { + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match self.registered_keys(credential, namespace).await { + Ok(keys) => { + if keys.iter().any(|key| key == service_key) { + return; + } + last_error = format!("relay reports keys {keys:?}"); + } + Err(error) => last_error = error, + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + panic!("`{service_key}` was never registered: {last_error}"); + } + + /// The service names visible to `credential` in `namespace`. + /// + /// `Keys` is namespace-scoped and answers with bare service names, so a + /// temporary credential sees exactly its own namespace with no extra work. + pub async fn registered_keys( + &self, + credential: Credential, + namespace: Option, + ) -> Result, String> { + let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(self.addr)) + .await + .map_err(|_| "status connect timed out".to_string())? + .map_err(|error| format!("status connect failed: {error}"))?; + let response = timeout( + Duration::from_secs(1), + get_status_with_credential(&mut stream, PbConnStatusReq::Keys, namespace, &credential), + ) + .await + .map_err(|_| "status request timed out".to_string())? + .map_err(|error| format!("status request failed: {error}"))?; + match response { + PbConnStatusResp::Keys(keys) => Ok(keys), + other => Err(format!("unexpected status response: {other:?}")), + } + } +} + +impl Drop for Relay { + fn drop(&mut self) { + self.shutdown.cancel(); + self.task.abort(); + let _ = std::fs::remove_dir_all(&self.state_dir); + } +} diff --git a/crates/pb-mapper-testkit/src/traffic.rs b/crates/pb-mapper-testkit/src/traffic.rs new file mode 100644 index 0000000..81293e1 --- /dev/null +++ b/crates/pb-mapper-testkit/src/traffic.rs @@ -0,0 +1,223 @@ +//! Payload generators and the load drivers the cases assert with. + +use std::net::SocketAddr; +use std::time::Duration; + +use pb_mapper_protocol::{MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter}; +use rand::RngExt; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpStream, UdpSocket}; +use tokio::time::{Instant, timeout}; +use uni_stream::addr::ToSocketAddrs; +use uni_stream::stream::{StreamProvider, StreamSplit}; +use uni_stream::udp::tune_udp_socket; + +use crate::{PROBE, READY_TIMEOUT, UDP_TEST_PAYLOAD_MAX}; + +/// What an echo server tagged with `tag` replies to `payload`. +/// +/// The tag is a leading byte identifying which echo server answered; see +/// [`crate::TunnelSpec::echo_tag`]. Untagged servers echo the payload verbatim. +pub fn tagged_echo(tag: Option, payload: &[u8]) -> Vec { + let mut expected = Vec::with_capacity(payload.len() + 1); + expected.extend(tag); + expected.extend_from_slice(payload); + expected +} + +/// Accumulates the time spent inside its scope and prints that slice. +pub struct TimerTickGuard<'a> { + ins: Instant, + mut_duration: &'a mut Duration, +} + +impl<'a> TimerTickGuard<'a> { + pub fn new(mut_duration: &'a mut Duration) -> Self { + Self { + ins: Instant::now(), + mut_duration, + } + } +} + +impl Drop for TimerTickGuard<'_> { + fn drop(&mut self) { + let duration = Instant::now() - self.ins; + *self.mut_duration += duration; + println!("duration:{duration:?}"); + } +} + +pub async fn connected_udp_socket(addr: SocketAddr) -> UdpSocket { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + tune_udp_socket(&socket); + socket.connect(addr).await.unwrap(); + socket +} + +/// Round-trip one probe datagram on `socket` and check it against `expected`. +/// +/// The tunnel keys UDP streams by source address, so every new socket starts a +/// new stream and its first datagram can be dropped while the relay sets that +/// stream up. Callers that go on to assert payload equality warm their own socket +/// with this first. +pub async fn probe_udp_socket(socket: &UdpSocket, expected: &[u8]) -> Result<(), String> { + socket + .send(PROBE) + .await + .map_err(|error| format!("udp probe send failed: {error}"))?; + let mut buf = vec![0u8; 65_507]; + let len = timeout(Duration::from_millis(500), socket.recv(&mut buf)) + .await + .map_err(|_| "udp probe timed out".to_string())? + .map_err(|error| format!("udp probe recv failed: {error}"))?; + if &buf[..len] == expected { + Ok(()) + } else { + Err(format!("udp probe echoed {len} bytes: {:?}", &buf[..len])) + } +} + +pub async fn wait_until_udp_socket_forwards(socket: &UdpSocket, tag: Option) { + let expected = tagged_echo(tag, PROBE); + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match probe_udp_socket(socket, &expected).await { + Ok(()) => return, + Err(error) => last_error = error, + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("udp socket never forwarded through the tunnel: {last_error}"); +} + +/// Random payload; the length is random too, so framing is exercised at many sizes. +pub fn gen_random_msg(max_len: usize) -> Vec { + let len = rand::rng().random_range(0_usize..max_len); + let mut vec = Vec::with_capacity(len); + for _ in 0..len { + vec.push(rand::rng().random_range(0..212)); + } + vec +} + +/// Send `rounds` × `burst` datagrams through the tunnel and assert each comes back. +pub async fn run_udp_datagram_echo(addr: SocketAddr, rounds: usize, burst: usize, tag: Option) { + let socket = connected_udp_socket(addr).await; + wait_until_udp_socket_forwards(&socket, tag).await; + let mut buf = vec![0u8; 65_507]; + + for round in 0..rounds { + for seq in 0..burst { + let mut msg = (seq as u32).to_be_bytes().to_vec(); + if round == 0 && seq == 0 { + msg.extend(vec![0u8; UDP_TEST_PAYLOAD_MAX]); + } else { + msg.extend(gen_random_msg(UDP_TEST_PAYLOAD_MAX)); + } + socket.send(&msg).await.unwrap(); + let expected = tagged_echo(tag, &msg); + let offset = usize::from(tag.is_some()); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + let wait = deadline.saturating_duration_since(Instant::now()); + let len = match timeout(wait, socket.recv(&mut buf)).await { + Ok(Ok(len)) => len, + Ok(Err(err)) => panic!("udp recv error: {err}"), + Err(_) => panic!("udp echo timeout; missing seq: {seq}"), + }; + if len < offset + 4 { + continue; + } + // A datagram from an earlier sequence number may still be in + // flight; skip it rather than failing the comparison. + let echoed_seq = u32::from_be_bytes(buf[offset..offset + 4].try_into().unwrap()); + if echoed_seq != seq as u32 { + continue; + } + assert_eq!(expected, &buf[..len]); + break; + } + } + } +} + +/// Round-trip `times` × 10 random payloads over one stream and assert each echo. +/// +/// This frames its payloads, so the echo server must be byte-transparent — an +/// echo tag would shift every frame header by a byte. Use [`run_raw_tcp_echo`] +/// for a tagged tunnel. +pub async fn run_echo_delay(addr: A, times: usize) { + let mut stream = P::from_addr(addr).await.unwrap(); + let (mut reader, mut writer) = stream.split(); + let mut reader = NormalMessageReader::new(&mut reader); + let mut writer = NormalMessageWriter::new(&mut writer); + let mut duration = Duration::default(); + for _ in 0..times { + let expected = gen_random_msg(2000); + for _ in 0..10 { + let msg = { + let _guard = TimerTickGuard::new(&mut duration); + writer.write_msg(&expected).await.unwrap(); + reader.read_msg().await.unwrap() + }; + + assert_eq!(expected, msg); + } + } + println!("{times} rounds of 10 random data echo delay tests each took a total of {duration:?}"); +} + +/// Round-trip `rounds` random payloads over one raw TCP stream, tag included. +/// +/// Unframed, so an echo tag stays a tag instead of shifting a frame header. The +/// tunnel's local side is byte-transparent either way, so the only thing framing +/// buys a test is the length delimiter — which this replaces by reading exactly +/// as many bytes as it expects. +pub async fn run_raw_tcp_echo(addr: SocketAddr, rounds: usize, tag: Option) { + let mut stream = TcpStream::connect(addr).await.unwrap(); + for _ in 0..rounds { + // At least one byte, so a payload is never indistinguishable from a bare tag. + let payload = gen_random_msg(2000); + let payload = if payload.is_empty() { + vec![7u8] + } else { + payload + }; + let expected = tagged_echo(tag, &payload); + stream.write_all(&payload).await.unwrap(); + let mut echoed = vec![0u8; expected.len()]; + timeout(Duration::from_secs(5), stream.read_exact(&mut echoed)) + .await + .expect("raw tcp echo timed out") + .expect("raw tcp echo failed"); + assert_eq!(expected, echoed); + } +} + +/// One raw payload through the tunnel, compared against `expected`. +pub async fn raw_tcp_probe( + addr: SocketAddr, + payload: &[u8], + expected: &[u8], +) -> Result<(), String> { + let mut stream = timeout(Duration::from_secs(1), TcpStream::connect(addr)) + .await + .map_err(|_| "tcp connect timed out".to_string())? + .map_err(|error| format!("tcp connect failed: {error}"))?; + stream + .write_all(payload) + .await + .map_err(|error| format!("tcp probe write failed: {error}"))?; + let mut echoed = vec![0u8; expected.len()]; + timeout(Duration::from_secs(1), stream.read_exact(&mut echoed)) + .await + .map_err(|_| "tcp probe read timed out".to_string())? + .map_err(|error| format!("tcp probe read failed: {error}"))?; + if echoed == expected { + Ok(()) + } else { + Err(format!("tcp probe echoed {echoed:?}")) + } +} diff --git a/crates/pb-mapper-testkit/src/tunnel.rs b/crates/pb-mapper-testkit/src/tunnel.rs new file mode 100644 index 0000000..6540569 --- /dev/null +++ b/crates/pb-mapper-testkit/src/tunnel.rs @@ -0,0 +1,378 @@ +//! One tunnel over a [`Relay`]: echo server, `register`, and `connect`. + +use std::net::SocketAddr; +use std::time::Duration; + +use pb_mapper_client::client::run_client_side_cli_with_callback_scoped; +use pb_mapper_client::server::{ServerTunnelOptions, run_server_side_cli_with_pinned_credential}; +use pb_mapper_core::checksum::Credential; +use tokio::net::{TcpListener, UdpSocket}; +use tokio::task::JoinHandle; +use tokio::time::Instant; +use uni_stream::stream::{ + TcpListenerProvider, TcpStreamProvider, UdpListenerProvider, UdpStreamProvider, +}; + +use crate::echo::{tcp_echo_server, udp_echo_server}; +use crate::relay::Relay; +use crate::traffic::{connected_udp_socket, probe_udp_socket, raw_tcp_probe}; +use crate::{PROBE, READY_TIMEOUT, Transport, admin_credential, reserve_addr}; + +/// What kind of tunnel to build. Everything but the transport has a default that +/// matches the plain administrator case. +#[derive(Debug, Clone)] +pub struct TunnelSpec { + transport: Transport, + need_codec: bool, + keep_alive: bool, + credential: Option, + /// Overrides `credential` on the `connect` side only, for the case where one + /// party registers a service and a different one subscribes to it. + connect_credential: Option, + service_key: Option, + /// The namespace `register` and `connect` ask for. `None` means "the + /// credential's own", which for a temporary credential is its key ID. + namespace: Option, + /// Required for an administrator registering into a temporary namespace. + force_namespace: bool, + /// A byte the echo server prepends to every reply, identifying it. + echo_tag: Option, +} + +impl TunnelSpec { + pub fn new(transport: Transport) -> Self { + Self { + transport, + need_codec: false, + keep_alive: false, + credential: None, + connect_credential: None, + service_key: None, + namespace: None, + force_namespace: false, + echo_tag: None, + } + } + + /// Encrypt forwarded traffic, as `register --codec` does. + pub fn codec(mut self, need_codec: bool) -> Self { + self.need_codec = need_codec; + self + } + + pub fn keep_alive(mut self, keep_alive: bool) -> Self { + self.keep_alive = keep_alive; + self + } + + /// Authenticate as `credential` instead of the administrator. + pub fn credential(mut self, credential: Credential) -> Self { + self.credential = Some(credential); + self + } + + /// Subscribe as a different credential than the one that registered. + /// + /// A temporary credential may name its own namespace explicitly, so this + /// composes with [`TunnelSpec::namespace`] without a second namespace knob. + pub fn connect_credential(mut self, credential: Credential) -> Self { + self.connect_credential = Some(credential); + self + } + + /// Use an explicit service name. Two tunnels may share one only if they are + /// in different namespaces. + pub fn service_key(mut self, key: impl Into) -> Self { + self.service_key = Some(key.into()); + self + } + + /// Register and connect inside an explicit namespace. + pub fn namespace(mut self, namespace: u64) -> Self { + self.namespace = Some(namespace); + self + } + + /// Acknowledge an administrator registration into a temporary namespace. + pub fn force_namespace(mut self, force: bool) -> Self { + self.force_namespace = force; + self + } + + /// Have this tunnel's echo server stamp every reply with `tag`. + /// + /// Two tunnels sharing a service name across namespaces need distinct tags: + /// without them, traffic reaching the wrong echo server would still come back + /// byte-identical and the test would pass on a leak. + pub fn echo_tag(mut self, tag: u8) -> Self { + self.echo_tag = Some(tag); + self + } + + fn label(&self) -> String { + format!( + "{}-{}", + self.transport.name(), + if self.need_codec { "codec" } else { "plain" } + ) + } +} + +/// A running tunnel. `Drop` aborts its tasks; the echo server, `register`, and +/// `connect` all stop with it. +pub struct Tunnel { + tunnel_addr: SocketAddr, + transport: Transport, + service_key: String, + echo_tag: Option, + tasks: Vec>, +} + +impl Tunnel { + /// Build the tunnel and return once it carries traffic. + pub async fn start(relay: &Relay, spec: TunnelSpec) -> Self { + let relay_addr = relay.addr(); + let transport = spec.transport; + let credential = spec.credential.unwrap_or_else(admin_credential); + let connect_credential = spec.connect_credential.unwrap_or(credential); + let label = spec.label(); + let service_key = spec.service_key.clone().unwrap_or(format!("echo-{label}")); + let echo_tag = spec.echo_tag; + let mut tasks = Vec::new(); + + // The echo server owns its socket from the start, so its address is known + // without a reservation window. + let echo_addr = match transport { + Transport::Tcp => { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tasks.push(tokio::spawn(tcp_echo_server(listener, echo_tag))); + addr + } + Transport::Udp => { + let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let addr = socket.local_addr().unwrap(); + tasks.push(tokio::spawn(udp_echo_server(socket, echo_tag))); + addr + } + }; + + let options = ServerTunnelOptions { + need_codec: spec.need_codec, + is_datagram: transport.is_datagram(), + keep_alive: spec.keep_alive, + namespace: spec.namespace, + force_namespace: spec.force_namespace, + }; + let register_key = service_key.clone(); + tasks.push(tokio::spawn(async move { + match transport { + Transport::Tcp => { + run_server_side_cli_with_pinned_credential::( + echo_addr, + relay_addr, + register_key.into(), + options, + None, + credential, + ) + .await + } + Transport::Udp => { + run_server_side_cli_with_pinned_credential::( + echo_addr, + relay_addr, + register_key.into(), + options, + None, + credential, + ) + .await + } + } + })); + + // `register` must have published the key before `connect` probes for it; + // otherwise `connect` burns its backoff waiting. + relay + .wait_for_registration(&service_key, credential, spec.namespace) + .await; + + let tunnel_addr = reserve_addr(transport).await; + let connect_key = service_key.clone(); + let keep_alive = spec.keep_alive; + let namespace = spec.namespace; + tasks.push(tokio::spawn(async move { + match transport { + Transport::Tcp => { + run_client_side_cli_with_callback_scoped::( + tunnel_addr, + relay_addr, + connect_key.into(), + keep_alive, + namespace, + None, + Some(connect_credential), + ) + .await + } + Transport::Udp => { + run_client_side_cli_with_callback_scoped::( + tunnel_addr, + relay_addr, + connect_key.into(), + keep_alive, + namespace, + None, + Some(connect_credential), + ) + .await + } + } + })); + + let tunnel = Self { + tunnel_addr, + transport, + service_key, + echo_tag, + tasks, + }; + tunnel.wait_until_forwarding().await; + tunnel + } + + /// Where a test client sends its traffic — the address `connect` listens on. + pub fn addr(&self) -> SocketAddr { + self.tunnel_addr + } + + pub fn transport(&self) -> Transport { + self.transport + } + + pub fn service_key(&self) -> &str { + &self.service_key + } + + /// The tag this tunnel's echo server prepends, if any. + pub fn echo_tag(&self) -> Option { + self.echo_tag + } + + /// What a reply to `payload` should look like, tag included. + fn expected_echo(&self, payload: &[u8]) -> Vec { + let mut expected = Vec::with_capacity(payload.len() + 1); + expected.extend(self.echo_tag); + expected.extend_from_slice(payload); + expected + } + + /// True if one probe payload round-trips right now. + pub async fn forwards_now(&self) -> bool { + self.probe_once().await.is_ok() + } + + /// Poll until a payload round-trips, or fail the test. + /// + /// This is true end-to-end readiness — the relay, `register`'s control + /// connection, `connect`'s local listener, and the echo server at once — which + /// is why no case here needs a sleep. + pub async fn wait_until_forwarding(&self) { + let deadline = Instant::now() + READY_TIMEOUT; + let mut last_error = String::from("no attempt completed"); + while Instant::now() < deadline { + match self.probe_once().await { + Ok(()) => return, + Err(error) => last_error = error, + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!( + "{} tunnel at {} never forwarded traffic: {last_error}", + self.transport.name(), + self.tunnel_addr + ); + } + + /// Poll until a probe stops round-tripping, or fail the test. + /// + /// A credential's expiry or revocation cancels its lease, which drops the + /// forwarding tasks; this is how a test observes that from the outside. + pub async fn wait_until_not_forwarding(&self, within: Duration) { + let deadline = Instant::now() + within; + while Instant::now() < deadline { + if self.probe_once().await.is_err() { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!( + "{} tunnel at {} still forwards traffic after {within:?}", + self.transport.name(), + self.tunnel_addr + ); + } + + async fn probe_once(&self) -> Result<(), String> { + match self.transport { + Transport::Tcp => self.probe_tcp().await, + Transport::Udp => self.probe_udp().await, + } + } + + /// Unframed on purpose: the tunnel's local side is byte-transparent, and a + /// framed probe would misread an echo tag as part of the length header. + async fn probe_tcp(&self) -> Result<(), String> { + raw_tcp_probe(self.tunnel_addr, PROBE, &self.expected_echo(PROBE)).await + } + + async fn probe_udp(&self) -> Result<(), String> { + let socket = connected_udp_socket(self.tunnel_addr).await; + probe_udp_socket(&socket, &self.expected_echo(PROBE)).await + } +} + +impl Drop for Tunnel { + fn drop(&mut self) { + for task in &self.tasks { + task.abort(); + } + } +} + +/// A relay and one tunnel on it, for the common case where a test wants a +/// complete flow and nothing else. +/// +/// The field order matters: `Tunnel` is dropped before `Relay`, so `register` and +/// `connect` stop before the relay they talk to. +pub struct TunnelHarness { + tunnel: Tunnel, + relay: Relay, +} + +impl TunnelHarness { + /// The one-liner: an administrator-credentialed tunnel over its own relay. + pub async fn start(transport: Transport, need_codec: bool) -> Self { + Self::with_spec(TunnelSpec::new(transport).codec(need_codec)).await + } + + pub async fn with_spec(spec: TunnelSpec) -> Self { + let relay = Relay::start(&spec.label()).await; + let tunnel = relay.start_tunnel(spec).await; + Self { tunnel, relay } + } + + pub fn relay(&self) -> &Relay { + &self.relay + } + + pub fn tunnel(&self) -> &Tunnel { + &self.tunnel + } + + /// Where a test client sends its traffic. + pub fn tunnel_addr(&self) -> SocketAddr { + self.tunnel.addr() + } +} From 03b5697515c4c36b5a6cd1e0d582cf8e4eae0c2c Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:19:58 +0800 Subject: [PATCH 3/7] Cover temporary credentials end to end MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `regression.rs` drives the credential lifecycle at the hand-rolled frame level, which left the real entry points untested with anything but the administrator key. These run `register` and `connect` themselves, so a temporary credential goes through connection pooling, the heartbeat window, control-connection reconnect, the stream-establishment handshake, and `connect`'s startup status probe. Nine cases: the transport/codec matrix on a temporary credential, two credentials sharing one service name in different namespaces, renew keeping a live tunnel forwarding past its original TTL, expiry and revocation each closing one, and the administrator placing a service inside a tenant's namespace. The namespace case tags each echo server, so a leak cannot satisfy payload equality — without it both servers echo identically and the assertion proves nothing. Co-Authored-By: Claude Opus 5 --- .../tests/temporary_credential_e2e.rs | 213 ++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 crates/pb-mapper-cli/tests/temporary_credential_e2e.rs diff --git a/crates/pb-mapper-cli/tests/temporary_credential_e2e.rs b/crates/pb-mapper-cli/tests/temporary_credential_e2e.rs new file mode 100644 index 0000000..46b2717 --- /dev/null +++ b/crates/pb-mapper-cli/tests/temporary_credential_e2e.rs @@ -0,0 +1,213 @@ +// See the note in `regression.rs`: the whole file is test code. +#![allow(clippy::unwrap_used, clippy::expect_used)] + +//! End-to-end tunnels authenticated by a temporary credential. +//! +//! `regression.rs` already covers the credential lifecycle at the hand-rolled +//! frame level. These cases run the real `register` and `connect` entry points +//! instead, so they also exercise connection pooling, the heartbeat window, +//! control-connection reconnect, the stream-establishment handshake, and +//! `connect`'s startup status probe — none of which a hand-written frame touches. +//! +//! Every case builds its own relay through `pb-mapper-testkit`, so they run +//! concurrently and never collide with a relay already on the machine. + +use std::time::Duration; + +use pb_mapper_auth::MIN_TEMP_KEY_TTL; +use pb_mapper_testkit::{ + Relay, Transport, TunnelSpec, admin_credential, run_echo_delay, run_raw_tcp_echo, + run_udp_datagram_echo, +}; +use uni_stream::stream::TcpStreamProvider; + +/// A TTL long enough that nothing expires mid-test. +const LONG_TTL: Duration = Duration::from_secs(600); + +/// Drive traffic through a tunnel a temporary credential registered and subscribed. +async fn assert_temporary_credential_echoes(transport: Transport, need_codec: bool) { + let label = format!( + "temp-{}-{}", + transport.name(), + if need_codec { "codec" } else { "plain" } + ); + let relay = Relay::start(&label).await; + let (_key_id, credential) = relay.issue_credential(LONG_TTL, &label).await; + + // `namespace: None` resolves to the credential's own namespace, which is its + // key ID — so the temporary tunnel needs no explicit namespace anywhere. + let tunnel = relay + .start_tunnel( + TunnelSpec::new(transport) + .codec(need_codec) + .credential(credential), + ) + .await; + + match transport { + Transport::Tcp => run_echo_delay::(tunnel.addr(), 10).await, + Transport::Udp => run_udp_datagram_echo(tunnel.addr(), 10, 8, None).await, + } +} + +#[tokio::test] +async fn temporary_credential_tcp_tunnel_echoes_without_codec() { + assert_temporary_credential_echoes(Transport::Tcp, false).await; +} + +#[tokio::test] +async fn temporary_credential_tcp_tunnel_echoes_with_codec() { + assert_temporary_credential_echoes(Transport::Tcp, true).await; +} + +#[tokio::test] +async fn temporary_credential_udp_tunnel_echoes_without_codec() { + assert_temporary_credential_echoes(Transport::Udp, false).await; +} + +#[tokio::test] +async fn temporary_credential_udp_tunnel_echoes_with_codec() { + assert_temporary_credential_echoes(Transport::Udp, true).await; +} + +/// Two credentials may register the same service name; neither sees the other's. +/// +/// The echo servers are tagged, because without a tag a leak between namespaces +/// would still echo the payload byte for byte and the assertion would pass. +#[tokio::test] +async fn two_temporary_credentials_share_a_service_name() { + let relay = Relay::start("temp-shared-name").await; + let (first_id, first) = relay.issue_credential(LONG_TTL, "first").await; + let (second_id, second) = relay.issue_credential(LONG_TTL, "second").await; + assert_ne!(first_id, second_id); + + const SHARED: &str = "shared-service"; + let first_tunnel = relay + .start_tunnel( + TunnelSpec::new(Transport::Tcp) + .credential(first) + .service_key(SHARED) + .echo_tag(b'1'), + ) + .await; + let second_tunnel = relay + .start_tunnel( + TunnelSpec::new(Transport::Tcp) + .credential(second) + .service_key(SHARED) + .echo_tag(b'2'), + ) + .await; + + // Each credential's `Keys` view holds the shared name exactly once. + for credential in [first, second] { + let keys = relay.registered_keys(credential, None).await.unwrap(); + assert_eq!( + keys.iter().filter(|key| *key == SHARED).count(), + 1, + "a credential should see the shared name once, saw {keys:?}" + ); + } + + // Both tunnels forward concurrently, each to its own echo server. + run_raw_tcp_echo(first_tunnel.addr(), 5, Some(b'1')).await; + run_raw_tcp_echo(second_tunnel.addr(), 5, Some(b'2')).await; +} + +/// Renewing keeps a live tunnel forwarding, and does not change the credential. +#[tokio::test] +async fn renew_keeps_a_live_tunnel_forwarding() { + let relay = Relay::start("temp-renew").await; + let issued = relay.issue(MIN_TEMP_KEY_TTL, "renew").await; + let key_id = issued.metadata.key_id; + let credential = pb_mapper_core::checksum::parse_credential(&issued.credential).unwrap(); + + let tunnel = relay + .start_tunnel(TunnelSpec::new(Transport::Tcp).credential(credential)) + .await; + + let renewed = relay.renew(key_id, LONG_TTL).await; + assert_eq!(renewed.metadata.key_id, key_id); + // Renewal moves the expiry without reissuing; a renewed credential that + // changed text would silently break every process already holding it. + assert_eq!(renewed.credential, issued.credential); + assert!(renewed.metadata.expires_at > issued.metadata.expires_at); + + // Past the original TTL, the renewed tunnel still carries traffic. + tokio::time::sleep(MIN_TEMP_KEY_TTL + Duration::from_secs(2)).await; + tunnel.wait_until_forwarding().await; + run_echo_delay::(tunnel.addr(), 3).await; +} + +/// Expiry — not revocation — tears a live tunnel down. +/// +/// The lifecycle actor drives expiry from a one-second tick, and the shortest +/// accepted TTL is [`MIN_TEMP_KEY_TTL`], so this case costs that plus a tick. +#[tokio::test] +async fn expiry_closes_a_live_tunnel() { + let relay = Relay::start("temp-expiry").await; + let (_key_id, credential) = relay.issue_credential(MIN_TEMP_KEY_TTL, "expiry").await; + + let tunnel = relay + .start_tunnel(TunnelSpec::new(Transport::Tcp).credential(credential)) + .await; + + tunnel + .wait_until_not_forwarding(MIN_TEMP_KEY_TTL + Duration::from_secs(10)) + .await; +} + +/// Revoking a credential closes the tunnel it authenticated. +#[tokio::test] +async fn revoke_closes_a_live_tunnel() { + let relay = Relay::start("temp-revoke").await; + let (key_id, credential) = relay.issue_credential(LONG_TTL, "revoke").await; + + let tunnel = relay + .start_tunnel(TunnelSpec::new(Transport::Tcp).credential(credential)) + .await; + + relay.revoke(key_id).await; + tunnel + .wait_until_not_forwarding(Duration::from_secs(10)) + .await; +} + +/// The administrator can register into a temporary namespace, with `--force`. +/// +/// The subscriber here is the temporary credential itself, so the case also shows +/// the service landing where the tenant can reach it. +#[tokio::test] +async fn admin_registers_into_a_temporary_namespace() { + let relay = Relay::start("temp-admin-namespace").await; + let (key_id, credential) = relay.issue_credential(LONG_TTL, "admin-force").await; + + let tunnel = relay + .start_tunnel( + TunnelSpec::new(Transport::Tcp) + .credential(admin_credential()) + .connect_credential(credential) + .namespace(key_id.as_u64()) + .force_namespace(true) + .service_key("admin-placed"), + ) + .await; + + // The tenant sees it as an ordinary service in its own namespace. + let keys = relay.registered_keys(credential, None).await.unwrap(); + assert!( + keys.iter().any(|key| key == "admin-placed"), + "the tenant should see the service the administrator placed, saw {keys:?}" + ); + // And the administrator's own namespace 0 stays empty. + let admin_keys = relay + .registered_keys(admin_credential(), None) + .await + .unwrap(); + assert!( + !admin_keys.iter().any(|key| key == "admin-placed"), + "namespace 0 should not hold the service, saw {admin_keys:?}" + ); + + run_echo_delay::(tunnel.addr(), 3).await; +} From 9fd1ab54e5c590b3f2cc5207452c20b1047f6625 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:20:10 +0800 Subject: [PATCH 4/7] Serialise the FFI state tests on the process credential MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Building a `PbMapperState` applies its stored `MSG_HEADER_KEY` to the process, and a state rooted at a fresh temporary directory has no stored key — so `temp_state` *clears* the process credential. The four tests share one binary, so a sibling constructing its state could wipe the credential that `a_failed_registration_releases_its_claim` had just set, and `register_service` reads it in phase 1, before the address parsing the case is actually asserting on. It failed with `InvalidArgument` instead of `InvalidAddress`. `temp_state` now takes `PROCESS_CREDENTIAL_TEST_LOCK` and hands the guard back, the way the tests under `crates/` already do. It was latent before — adding test load to the workspace run is what surfaced it. Co-Authored-By: Claude Opus 5 --- ui/native/pb_mapper_ffi/src/state.rs | 25 +++++++++++++++++++------ 1 file changed, 19 insertions(+), 6 deletions(-) diff --git a/ui/native/pb_mapper_ffi/src/state.rs b/ui/native/pb_mapper_ffi/src/state.rs index 6c437b2..dab51e1 100644 --- a/ui/native/pb_mapper_ffi/src/state.rs +++ b/ui/native/pb_mapper_ffi/src/state.rs @@ -620,12 +620,25 @@ pub async fn connect_service( // FFI boundary. In a test a panic *is* the failure report. #[allow(clippy::expect_used, clippy::unwrap_used)] mod tests { + use pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK; + use super::*; use crate::error::ErrorCode; + type ProcessCredentialGuard = tokio::sync::MutexGuard<'static, ()>; + /// A state rooted in a temporary directory, so a test never reads or writes /// the real user config. - fn temp_state(name: &str) -> (Arc>, PathBuf) { + /// + /// Building a state applies its stored `MSG_HEADER_KEY` to the process, and + /// a fresh temporary directory has no stored key — so this *clears* the + /// process credential. That is process-global, so the returned guard is what + /// keeps one test's state from wiping the credential another test just set. + /// Hold it for the whole test, not just the construction. + async fn temp_state( + name: &str, + ) -> (ProcessCredentialGuard, Arc>, PathBuf) { + let guard = PROCESS_CREDENTIAL_TEST_LOCK.lock().await; let root = std::env::temp_dir().join(format!( "pb-mapper-ffi-test-{name}-{}", SystemTime::now() @@ -634,12 +647,12 @@ mod tests { .unwrap_or_default() )); let state = PbMapperState::new(Some(root.to_string_lossy().into_owned())); - (Arc::new(Mutex::new(state)), root) + (guard, Arc::new(Mutex::new(state)), root) } #[tokio::test] async fn ui_server_uses_its_writable_config_directory_and_reports_readiness() { - let (state, root) = temp_state("server-auth-path"); + let (_process_credential_guard, state, root) = temp_state("server-auth-path").await; let auth_dir = { let mut state = state.lock().await; let auth_dir = state.config_dir.join("auth"); @@ -675,7 +688,7 @@ mod tests { /// nothing able to abort it. #[tokio::test] async fn a_second_registration_of_the_same_key_is_refused() { - let (state, root) = temp_state("claim"); + let (_process_credential_guard, state, root) = temp_state("claim").await; let held = { let guard = state.lock().await; @@ -717,7 +730,7 @@ mod tests { /// path can forget it — this pins that. #[tokio::test] async fn a_failed_registration_releases_its_claim() { - let (state, root) = temp_state("release"); + let (_process_credential_guard, state, root) = temp_state("release").await; struct RestoreProcessKey; impl Drop for RestoreProcessKey { fn drop(&mut self) { @@ -771,7 +784,7 @@ mod tests { /// must not block each other. #[tokio::test] async fn registering_and_connecting_claim_separately() { - let (state, root) = temp_state("separate"); + let (_process_credential_guard, state, root) = temp_state("separate").await; let guard = state.lock().await; let _registering = guard From 08b927186281795e5df3a13a0304246b561df273 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:20:10 +0800 Subject: [PATCH 5/7] Point the docs at the testkit as the way to build an e2e flow Co-Authored-By: Claude Opus 5 --- AGENTS.md | 20 ++++++++++++++++---- CLAUDE.md | 22 ++++++++++++++++++++-- README.md | 2 +- README.zh-CN.md | 2 +- 4 files changed, 38 insertions(+), 8 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 92f2e2c..1646f18 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,6 +16,7 @@ - `crates/`: the Rust workspace; the root `Cargo.toml` is a virtual manifest - `crates/pb-mapper-cli/src/bin/pb-mapper.rs`: unified CLI entry point - `crates/pb-mapper-{core,auth,protocol,server,client,cli}` + - `crates/pb-mapper-testkit/`: test support only; nothing shipped depends on it - `crates/pb-mapper-cli/tests/`: integration tests; no env setup required - `crates/pb-mapper-cli/examples/`: runnable examples - `ui/`: Flutter UI; Rust bridge under `ui/native/*` @@ -42,13 +43,24 @@ Notes: CI builds release artifacts on tags `vX.Y.Z` (see `.github/workflows/rele ## Testing Guidelines - Framework: `tokio` async + integration tests in `crates/pb-mapper-cli/tests/` -- No test needs environment setup. `test_delay.rs` runs the whole tunnel - (`server` + `register` + `connect`) over both transports with and without - `--codec`; each case reserves its own loopback ports and its own auth state +- No test needs environment setup. `pb-mapper-testkit` stands up a complete + tunnel (`server` + `register` + `connect`), so any test file can build one + rather than a single file owning the harness: `TunnelHarness::start(transport, + need_codec)` for the common case, or `Relay` + `TunnelSpec` when a case needs + to issue, renew, or revoke a credential first. `test_delay.rs` covers the + transport/codec matrix and `temporary_credential_e2e.rs` the credential + lifecycle; each case reserves its own loopback ports and its own auth state directory, so cases run concurrently and never collide with a live relay. - Sequence components with a readiness probe, not a sleep: poll the relay's `Keys` status for a registration, and round-trip a payload through the tunnel - for forwarding. `TunnelHarness` in `test_delay.rs` does both. + for forwarding. `Tunnel::start` does both before it returns. +- A framed driver (`run_echo_delay`) needs a byte-transparent echo server; a + tagged tunnel (`TunnelSpec::echo_tag`, which is how a namespace-isolation + assertion avoids passing on a leak) needs the raw drivers instead. +- Anything that sets the process credential takes + `pb_mapper_core::test_support::PROCESS_CREDENTIAL_TEST_LOCK` first — it is + process-global, and that includes indirect writers such as building a + `PbMapperState`. - Prefer new integration tests that need no external setup ## Commit & Pull Request Guidelines diff --git a/CLAUDE.md b/CLAUDE.md index 8cb35e4..a162c3f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -52,6 +52,7 @@ pb-mapper/ │ ├── pb-mapper-protocol/ # Message framing, v2 secure sessions, forwarding │ ├── pb-mapper-server/ # Central relay server, plus the task manager │ ├── pb-mapper-client/ # Both tunnel ends: `register` and `connect` +│ ├── pb-mapper-testkit/ # Test support: a complete e2e tunnel, for any test file │ └── pb-mapper-cli/ # The `pb-mapper` binary, integration tests, examples ├── ui/ # Flutter UI, talking to Rust over dart:ffi │ ├── lib/ # Flutter application code @@ -132,9 +133,23 @@ loader, two CMakeLists, four xcconfigs, and the release-ui hash checks expect. - `server/`: `register` — publishes a local service (`mod.rs`, `stream.rs`, `error.rs`) - `client/`: `connect` — subscribes and listens locally, plus `status.rs` +- **`pb-mapper-testkit/`**: Test support only; nothing shipped depends on it + - `relay.rs`: `Relay` — a live server that retains its `AuthRuntime`, so a case + can issue, renew, and revoke credentials without the admin wire protocol + - `tunnel.rs`: `TunnelSpec` / `Tunnel` / `TunnelHarness` — echo server plus + `register` plus `connect`, each on reserved loopback ports + - `echo.rs`, `traffic.rs`: Echo servers and the framed and raw traffic drivers + - A crate rather than `tests/common/mod.rs`: that module is compiled separately + into every test binary, and whatever a binary does not use is reported as + dead code — fatal under `-D warnings` + - **`pb-mapper-cli/`**: The binary, integration tests, and examples - `src/bin/pb-mapper.rs`: Argument parsing and the role commands - `src/bin/pb-mapper/admin.rs`: The `admin` subcommand + - `tests/test_delay.rs`: The transport/codec matrix over the whole tunnel + - `tests/temporary_credential_e2e.rs`: The credential lifecycle over the whole + tunnel — namespace isolation, renew, expiry, revoke + - `tests/regression.rs`: Protocol-level cases against hand-rolled frames #### Flutter UI (`ui/`) - **`lib/src/views/`**: One file per zone the shell can show @@ -346,7 +361,9 @@ part of the landing page and the setup wizard. - **Toolchain**: rust-toolchain.toml for reproducible builds - **Testing**: Unit tests live beside the code; integration tests are in `crates/pb-mapper-cli/tests/`, which is the crate that depends on every layer - they exercise + they exercise. The e2e scaffolding is `pb-mapper-testkit`, so a new test file + stands up its own full `server` + `register` + `connect` flow instead of + everything accumulating in one file. ### UI Development Guidelines - **Framework**: Flutter 3.44.9, Material 3. CI pins the same version. @@ -406,7 +423,8 @@ docker-compose -f docker/docker-compose.yml up - **Unit Tests**: `cargo test` for Rust components - **Widget Tests**: `flutter test` in ui/ directory -- **Integration Tests**: `tests/` directory contains end-to-end tests +- **Integration Tests**: `crates/pb-mapper-cli/tests/` — end-to-end tests built + on `pb-mapper-testkit` - **Examples**: `examples/` directory provides working usage examples ### Service Deployment diff --git a/README.md b/README.md index e94280c..8f4f3be 100644 --- a/README.md +++ b/README.md @@ -122,7 +122,7 @@ Open `http://localhost:3000` in the coffee-shop browser — traffic flows throug ## Repository layout -- `crates/` — the Rust workspace (six crates; the root manifest is virtual) +- `crates/` — the Rust workspace (six shipped crates plus a test-support one; the root manifest is virtual) - `ui/` — Flutter UI + native bridge - `docs/` — documentation and assets - `docker/`, `services/`, `scripts/` — deployment and tooling diff --git a/README.zh-CN.md b/README.zh-CN.md index 562351e..fbafd10 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -122,7 +122,7 @@ pb-mapper connect tcp --server :7666 --key web --addr 127.0.0.1:3000 ## 仓库结构 -- `crates/` — Rust workspace（六个 crate，根清单为虚拟清单） +- `crates/` — Rust workspace（六个发布 crate 加一个测试支撑 crate，根清单为虚拟清单） - `ui/` — Flutter UI + 原生桥接 - `docs/` — 文档与素材 - `docker/`、`services/`、`scripts/` — 部署与工具 From afd1237995c42f9bf0169bd976ae2f0e2b49bd0c Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:52:53 +0800 Subject: [PATCH 6/7] Tag a TCP echo once per connection, not once per read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The tagged echo server prepended its identifying byte to every reply, which on a stream means every `read` — and TCP preserves no write boundaries, so a payload arriving as two reads came back with a tag injected into its middle. Measured: correct through 4000-byte payloads, corrupt from 5000 up, four tags on 16 KiB. `run_raw_tcp_echo` generated at most 2000 bytes, so the bug sat under passing tests, waiting for whoever raised the size. A stream now carries the tag once, as the first byte of the connection. That is well defined however the traffic is chunked, and still identifies which echo server answered — all the tag is for. UDP keeps datagram boundaries, so it still tags every reply. The driver's payloads go to 20 KiB, well past the 8 KiB forwarding buffer, so this can no longer hold only because the sizes stayed small; a `const` assertion pins that. Two cases cover it, both verified to fail against the old server: one payload larger than a single read, and the driver and server agreeing over enough rounds to cross the threshold. Also reuse one socket across UDP probes. The relay keys UDP streams by source address, so a fresh socket per attempt meant `wait_until_not_forwarding` watched new streams fail to start rather than the established one going away. Co-Authored-By: Claude Opus 5 --- crates/pb-mapper-testkit/src/echo.rs | 21 ++++-- crates/pb-mapper-testkit/src/lib.rs | 33 +++++++-- crates/pb-mapper-testkit/src/traffic.rs | 25 ++++--- crates/pb-mapper-testkit/src/tunnel.rs | 97 +++++++++++++++++++++++-- 4 files changed, 151 insertions(+), 25 deletions(-) diff --git a/crates/pb-mapper-testkit/src/echo.rs b/crates/pb-mapper-testkit/src/echo.rs index 10abc8d..00bfb72 100644 --- a/crates/pb-mapper-testkit/src/echo.rs +++ b/crates/pb-mapper-testkit/src/echo.rs @@ -3,10 +3,19 @@ //! Both take an already-bound socket, so the caller can learn the address without //! a window in which another test could take the port. //! -//! `tag` prepends one byte to every reply. Two tunnels that share a service name -//! in different namespaces get different tags, so a test can tell which echo -//! server actually received the traffic — without it, a namespace leak would -//! still satisfy payload equality, since both servers echo identically. +//! `tag` identifies which echo server answered. Two tunnels that share a service +//! name in different namespaces get different tags, so a test can tell which one +//! actually received the traffic — without it, a namespace leak would still +//! satisfy payload equality, since both servers echo identically. +//! +//! Where the tag goes differs by transport, because the transports differ in what +//! they preserve. UDP keeps datagram boundaries, so every reply carries it. TCP +//! keeps none: one write can arrive as several reads and several writes as one, so +//! a tag per read would inject bytes into the middle of a payload. The TCP server +//! therefore sends it **once, as the first byte of the connection**, and the reply +//! stream is `tag` followed by the echoed bytes verbatim. That is well defined +//! however the traffic happens to be chunked, and it is still enough to identify +//! the server — which is all the tag is for. use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, UdpSocket}; @@ -19,12 +28,14 @@ pub async fn tcp_echo_server(listener: TcpListener, tag: Option) { }; tokio::spawn(async move { let mut buf = vec![0u8; 4096]; + // Consumed by the first reply on this connection; see the module note. + let mut pending_tag = tag; loop { match stream.read(&mut buf).await { Ok(0) | Err(_) => return, Ok(n) => { let mut reply = Vec::with_capacity(n + 1); - reply.extend(tag); + reply.extend(pending_tag.take()); reply.extend_from_slice(&buf[..n]); if stream.write_all(&reply).await.is_err() { return; diff --git a/crates/pb-mapper-testkit/src/lib.rs b/crates/pb-mapper-testkit/src/lib.rs index 86151f3..e6b03c7 100644 --- a/crates/pb-mapper-testkit/src/lib.rs +++ b/crates/pb-mapper-testkit/src/lib.rs @@ -75,6 +75,13 @@ pub const PROBE: &[u8] = b"pb-mapper-probe"; /// Upper bound on generated UDP payloads, comfortably inside one datagram. pub const UDP_TEST_PAYLOAD_MAX: usize = 1200; +/// Upper bound on generated raw-TCP payloads. +/// +/// Deliberately past both one segment and the 8 KiB initial forwarding buffer, so +/// a driver that assumed a write arrives as one read is caught here rather than by +/// whoever later raises a payload size. +pub const RAW_TCP_PAYLOAD_MAX: usize = 20_000; + /// Tracing plus the process credential, set up once per test process. /// /// The framing helpers in this crate ([`run_echo_delay`] and the TCP probe) write @@ -83,6 +90,16 @@ pub const UDP_TEST_PAYLOAD_MAX: usize = 1200; /// unrelated to which credential a tunnel authenticates with: `pb-mapper`'s local /// side is byte-transparent, so these frames are only ever read back by the same /// process that wrote them. +/// +/// This establishes the baseline for the whole test binary rather than taking +/// `PROCESS_CREDENTIAL_TEST_LOCK` around a write, which is why the `LazyLock` is +/// the whole synchronisation: it runs once, before any case can observe the +/// credential, and never writes again. **The consequence is a rule for test +/// files: a target that uses this crate must not also set the process credential +/// itself.** There is no lock to coordinate with — a case that wrote its own +/// would corrupt the framing under every concurrent case in the binary. A case +/// that needs a different credential should pass it to +/// [`TunnelSpec::credential`], which is per tunnel and touches nothing global. pub fn init_test_env() { static TEST_ENV: LazyLock<()> = LazyLock::new(|| { init_tracing(); @@ -139,12 +156,18 @@ impl Transport { } } -/// Reserve a loopback port by binding it and immediately dropping the socket. +/// Pick a free loopback port by binding it and immediately dropping the socket. +/// +/// TCP and UDP have separate port spaces, so this has to use the same protocol the +/// caller will bind. A relay and an echo server keep the socket they bound; this is +/// for `connect`, whose bind happens inside the client and cannot be handed a +/// pre-bound socket. /// -/// TCP and UDP have separate port spaces, so the reservation has to use the same -/// protocol the caller will bind. A relay and an echo server keep the socket they -/// bound; this is for `connect`, whose bind happens inside the client and cannot -/// be handed a pre-bound socket. +/// That leaves a window between the drop and the real bind, so this is not a +/// reservation and cannot be made into one. What keeps it from colliding in +/// practice: the ephemeral range holds tens of thousands of ports, and a kernel +/// does not hand back a just-released one while others are free. A test that can +/// own its socket outright should do that instead of calling this. pub async fn reserve_addr(transport: Transport) -> std::net::SocketAddr { match transport { Transport::Tcp => { diff --git a/crates/pb-mapper-testkit/src/traffic.rs b/crates/pb-mapper-testkit/src/traffic.rs index 81293e1..fb691b3 100644 --- a/crates/pb-mapper-testkit/src/traffic.rs +++ b/crates/pb-mapper-testkit/src/traffic.rs @@ -12,12 +12,14 @@ use uni_stream::addr::ToSocketAddrs; use uni_stream::stream::{StreamProvider, StreamSplit}; use uni_stream::udp::tune_udp_socket; -use crate::{PROBE, READY_TIMEOUT, UDP_TEST_PAYLOAD_MAX}; +use crate::{PROBE, RAW_TCP_PAYLOAD_MAX, READY_TIMEOUT, UDP_TEST_PAYLOAD_MAX}; -/// What an echo server tagged with `tag` replies to `payload`. +/// What an echo server replies to `payload`, given the tag it still owes. /// /// The tag is a leading byte identifying which echo server answered; see -/// [`crate::TunnelSpec::echo_tag`]. Untagged servers echo the payload verbatim. +/// [`crate::TunnelSpec::echo_tag`]. Pass `None` once a TCP connection has already +/// received it — on a stream it arrives once, not per reply — and for every +/// untagged server, which echoes the payload verbatim. pub fn tagged_echo(tag: Option, payload: &[u8]) -> Vec { let mut expected = Vec::with_capacity(payload.len() + 1); expected.extend(tag); @@ -177,22 +179,27 @@ pub async fn run_echo_delay(addr: A, /// as many bytes as it expects. pub async fn run_raw_tcp_echo(addr: SocketAddr, rounds: usize, tag: Option) { let mut stream = TcpStream::connect(addr).await.unwrap(); - for _ in 0..rounds { + // The tag arrives once, ahead of the first echoed byte on this connection — + // see the note in `echo.rs` for why it cannot be per reply on a stream. + let mut pending_tag = tag; + for round in 0..rounds { + // Deliberately past one segment and past the forwarding buffer, so a + // driver that assumed write boundaries survive would fail here. + let payload = gen_random_msg(RAW_TCP_PAYLOAD_MAX); // At least one byte, so a payload is never indistinguishable from a bare tag. - let payload = gen_random_msg(2000); let payload = if payload.is_empty() { vec![7u8] } else { payload }; - let expected = tagged_echo(tag, &payload); + let expected = tagged_echo(pending_tag.take(), &payload); stream.write_all(&payload).await.unwrap(); let mut echoed = vec![0u8; expected.len()]; timeout(Duration::from_secs(5), stream.read_exact(&mut echoed)) .await - .expect("raw tcp echo timed out") - .expect("raw tcp echo failed"); - assert_eq!(expected, echoed); + .unwrap_or_else(|_| panic!("raw tcp echo round {round} timed out")) + .unwrap_or_else(|error| panic!("raw tcp echo round {round} failed: {error}")); + assert_eq!(expected, echoed, "raw tcp echo round {round} mismatched"); } } diff --git a/crates/pb-mapper-testkit/src/tunnel.rs b/crates/pb-mapper-testkit/src/tunnel.rs index 6540569..2eb780f 100644 --- a/crates/pb-mapper-testkit/src/tunnel.rs +++ b/crates/pb-mapper-testkit/src/tunnel.rs @@ -35,7 +35,8 @@ pub struct TunnelSpec { namespace: Option, /// Required for an administrator registering into a temporary namespace. force_namespace: bool, - /// A byte the echo server prepends to every reply, identifying it. + /// A byte identifying the echo server: per datagram on UDP, once per + /// connection on TCP. echo_tag: Option, } @@ -99,11 +100,15 @@ impl TunnelSpec { self } - /// Have this tunnel's echo server stamp every reply with `tag`. + /// Have this tunnel's echo server identify itself with `tag`. /// /// Two tunnels sharing a service name across namespaces need distinct tags: /// without them, traffic reaching the wrong echo server would still come back /// byte-identical and the test would pass on a leak. + /// + /// A UDP server tags every datagram; a TCP server tags the connection once, + /// ahead of its first echoed byte. See `echo.rs` for why a stream cannot carry + /// it per reply. pub fn echo_tag(mut self, tag: u8) -> Self { self.echo_tag = Some(tag); self @@ -125,6 +130,15 @@ pub struct Tunnel { transport: Transport, service_key: String, echo_tag: Option, + /// The socket every UDP probe reuses. + /// + /// The relay keys UDP streams by source address, so a fresh socket starts a + /// fresh stream whose first datagram can be dropped during setup. Probing on + /// one socket means a retry exercises the stream the previous attempt + /// established, rather than opening another one — and it is what makes + /// [`Self::wait_until_not_forwarding`] observe the established stream going + /// away instead of a new stream failing to start. + udp_probe_socket: Option, tasks: Vec>, } @@ -231,11 +245,17 @@ impl Tunnel { } })); + let udp_probe_socket = match transport { + Transport::Tcp => None, + Transport::Udp => Some(connected_udp_socket(tunnel_addr).await), + }; + let tunnel = Self { tunnel_addr, transport, service_key, echo_tag, + udp_probe_socket, tasks, }; tunnel.wait_until_forwarding().await; @@ -255,12 +275,16 @@ impl Tunnel { &self.service_key } - /// The tag this tunnel's echo server prepends, if any. + /// The tag this tunnel's echo server identifies itself with, if any. pub fn echo_tag(&self) -> Option { self.echo_tag } - /// What a reply to `payload` should look like, tag included. + /// What a reply to a single probe should look like, tag included. + /// + /// Correct for both transports because a probe is one datagram, or the first + /// and only exchange on a fresh connection — which is exactly where a TCP + /// tag sits. fn expected_echo(&self, payload: &[u8]) -> Vec { let mut expected = Vec::with_capacity(payload.len() + 1); expected.extend(self.echo_tag); @@ -328,8 +352,11 @@ impl Tunnel { } async fn probe_udp(&self) -> Result<(), String> { - let socket = connected_udp_socket(self.tunnel_addr).await; - probe_udp_socket(&socket, &self.expected_echo(PROBE)).await + let socket = self + .udp_probe_socket + .as_ref() + .ok_or_else(|| "no udp probe socket on a tcp tunnel".to_string())?; + probe_udp_socket(socket, &self.expected_echo(PROBE)).await } } @@ -376,3 +403,61 @@ impl TunnelHarness { self.tunnel.addr() } } + +#[cfg(test)] +mod tests { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpStream; + + use super::*; + use crate::{RAW_TCP_PAYLOAD_MAX, run_raw_tcp_echo}; + + /// A tagged TCP tunnel carries a payload larger than one read intact. + /// + /// This is the case the harness got wrong: the echo server tagged every read, + /// so a payload that arrived as two reads came back with a tag injected into + /// its middle. It only showed up past roughly 4 KiB, which the driver's + /// payloads did not reach — so the bug sat under passing tests, waiting for + /// whoever raised the size. + #[tokio::test] + async fn a_tagged_tcp_tunnel_survives_a_payload_larger_than_one_read() { + let relay = Relay::start("tag-fragmentation").await; + let tunnel = relay + .start_tunnel(TunnelSpec::new(Transport::Tcp).echo_tag(b'1')) + .await; + + let mut stream = TcpStream::connect(tunnel.addr()).await.unwrap(); + // Past one segment and past the 8 KiB initial forwarding buffer, so the + // echo server is guaranteed more than one read. + let payload = vec![0xABu8; 64 * 1024]; + stream.write_all(&payload).await.unwrap(); + + let mut echoed = vec![0u8; payload.len() + 1]; + tokio::time::timeout(Duration::from_secs(5), stream.read_exact(&mut echoed)) + .await + .expect("tagged echo timed out") + .expect("tagged echo failed"); + + assert_eq!(echoed[0], b'1', "the reply should open with the tag"); + assert_eq!( + &echoed[1..], + &payload[..], + "the payload came back altered, so a tag was injected mid-stream" + ); + } + + /// The driver's payloads must be able to exceed one read, or the case below + /// proves nothing about fragmentation. + const _: () = assert!(RAW_TCP_PAYLOAD_MAX > 8 * 1024); + + /// The driver agrees with the server about where the tag goes, over enough + /// rounds that its random sizes cross the fragmentation threshold. + #[tokio::test] + async fn the_raw_tcp_driver_agrees_with_a_tagged_server() { + let relay = Relay::start("tag-driver").await; + let tunnel = relay + .start_tunnel(TunnelSpec::new(Transport::Tcp).echo_tag(b'7')) + .await; + run_raw_tcp_echo(tunnel.addr(), 12, Some(b'7')).await; + } +} From 607c424e8e4f77c33fabef71b7404ed0c743ac15 Mon Sep 17 00:00:00 2001 From: LB7666 Date: Sat, 22 Aug 2026 04:53:05 +0800 Subject: [PATCH 7/7] Say what the testkit actually guarantees MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two claims were stronger than the code. `reserve_addr` has to drop the socket before `connect` binds it — the bind happens inside the client and cannot be handed a pre-bound socket — so it picks a port rather than reserving one, and "never collide" was not something it could promise. What does hold is now written down, along with the advice to own a socket outright where a test can. `init_test_env` establishes the process credential for the whole binary and never writes again, which is why a `LazyLock` is the whole synchronisation and not an oversight. The rule that follows is what was missing: a target using this crate must not also set the credential itself, because there is no lock to coordinate with. A case wanting a different one passes it to `TunnelSpec::credential`, which touches nothing global. Co-Authored-By: Claude Opus 5 --- AGENTS.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/AGENTS.md b/AGENTS.md index 1646f18..15463cb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,8 +49,10 @@ Notes: CI builds release artifacts on tags `vX.Y.Z` (see `.github/workflows/rele need_codec)` for the common case, or `Relay` + `TunnelSpec` when a case needs to issue, renew, or revoke a credential first. `test_delay.rs` covers the transport/codec matrix and `temporary_credential_e2e.rs` the credential - lifecycle; each case reserves its own loopback ports and its own auth state + lifecycle; each case picks its own loopback ports and owns its auth state directory, so cases run concurrently and never collide with a live relay. + Prefer binding a socket and keeping it over `reserve_addr`, which has to drop + the socket before the real bind and so cannot rule out a race. - Sequence components with a readiness probe, not a sleep: poll the relay's `Keys` status for a registration, and round-trip a payload through the tunnel for forwarding. `Tunnel::start` does both before it returns.