diff --git a/AGENTS.md b/AGENTS.md index e21de91..15463cb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,7 +16,8 @@ - `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-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/*` - `docker/`, `services/`, `scripts/`: container, systemd, build/release @@ -41,10 +42,28 @@ 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. `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 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. +- 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 - Commits: short, imperative (e.g., "Fix localhost resolution panic", "add network perms", "change to StreamBuilder") 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/Cargo.lock b/Cargo.lock index 3e44e6f..e206df9 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,12 +1045,12 @@ version = "0.4.0" dependencies = [ "better_mimalloc_rs", "clap", - "dotenvy", "pb-mapper-auth", "pb-mapper-client", "pb-mapper-core", "pb-mapper-protocol", "pb-mapper-server", + "pb-mapper-testkit", "rand 0.10.0", "serde_json", "tokio", @@ -1152,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 4848d5f..e37aa76 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,13 +24,14 @@ 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"] } 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/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/` — 部署与工具 diff --git a/crates/pb-mapper-cli/Cargo.toml b/crates/pb-mapper-cli/Cargo.toml index f45b297..0ce404b 100644 --- a/crates/pb-mapper-cli/Cargo.toml +++ b/crates/pb-mapper-cli/Cargo.toml @@ -25,7 +25,8 @@ tracing.workspace = true uni-stream.workspace = true [dev-dependencies] -dotenvy.workspace = true +pb-mapper-testkit.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/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; +} diff --git a/crates/pb-mapper-cli/tests/test_delay.rs b/crates/pb-mapper-cli/tests/test_delay.rs index 78c747b..afed920 100644 --- a/crates/pb-mapper-cli/tests/test_delay.rs +++ b/crates/pb-mapper-cli/tests/test_delay.rs @@ -1,417 +1,44 @@ // See the note in `regression.rs`: the whole file is test code. #![allow(clippy::unwrap_used, clippy::expect_used)] -use std::env; -use std::sync::LazyLock; -use std::time::Duration; +//! 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. The scaffolding +//! lives in `pb-mapper-testkit`, so any other test file can stand up the same +//! flow — see `temporary_credential_e2e.rs`. -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_core::config::init_tracing; -use pb_mapper_protocol::{MessageReader, MessageWriter, NormalMessageReader, NormalMessageWriter}; -use pb_mapper_server::run_server_with_auth_config; -use rand::RngExt; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::net::UdpSocket; -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::udp::tune_udp_socket; +use pb_mapper_testkit::{Transport, TunnelHarness, run_echo_delay, run_udp_datagram_echo}; +use uni_stream::stream::TcpStreamProvider; -struct TimerTickGurad<'a> { - ins: Instant, - mut_duration: &'a mut Duration, -} - -impl<'a> TimerTickGurad<'a> { - fn new(mut_duration: &'a mut Duration) -> Self { - Self { - ins: Instant::now(), - mut_duration, - } - } -} - -impl<'a> Drop for TimerTickGurad<'a> { - fn drop(&mut self) { - let end = Instant::now(); - let duration = end - self.ins; - *self.mut_duration += duration; - println!("duration:{duration:?}"); - } -} - -use uni_stream::stream::StreamAccept; - -const UDP_TEST_PAYLOAD_MAX: usize = 1200; - -async fn echo_server( - server_addr: &str, -) -> Result<(), Box> { - let listener = P::bind(server_addr).await?; - println!("run echo server:{server_addr}"); - loop { - // Accept incoming connections - let (mut stream, addr) = listener.accept().await?; - println!("Connected from {addr}"); - - // Process each connection concurrently - tokio::spawn(async move { - // Read data from client - let mut buf = vec![0; 1024]; - loop { - let n = match stream.read(&mut buf).await { - Ok(n) => n, - Err(e) => { - println!("Error reading: {e}"); - return; - } - }; - - // 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}"); - } - }); +/// 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, None).await, } } -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?; - 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?; - if let Err(err) = socket.send_to(&buf[..len], peer).await { - println!("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}"); - } -} - -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 - } - 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 - } - } -} - -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, - ) - .await - } - ServerType::Tcp => { - run_client_side_cli::( - local_addr.to_string(), - remote_addr.to_string(), - key.into(), - false, - ) - .await - } - } -} - -/// 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)); - } - vec -} - -async fn run_udp_datagram_echo(addr: &str, rounds: usize, burst: usize) { - let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); - tune_udp_socket(&socket); - socket.connect(addr).await.unwrap(); - let mut buf = vec![0u8; 65_507]; - - // 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; - } - tokio::time::sleep(Duration::from_millis(50)).await; - } - assert!(ready, "udp echo path not ready"); - - 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(3); - 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; - } - let recv_seq = u32::from_be_bytes(buf[..4].try_into().unwrap()); - if recv_seq != 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 = TimerTickGurad::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:?}"); +#[tokio::test] +async fn tcp_tunnel_echoes_without_codec() { + assert_tunnel_echoes(Transport::Tcp, false).await; } -#[derive(Debug, Clone, Copy)] -enum ServerType { - Udp, - Tcp, +#[tokio::test] +async fn tcp_tunnel_echoes_with_codec() { + assert_tunnel_echoes(Transport::Tcp, true).await; } -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 - } -}); - -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, - } - - // abort all thread - echo_server_handle.abort(); - pb_mapper_server_handle.abort(); - pb_mapper_server_cli_handle.abort(); - pb_mapper_client_cli_handle.abort(); +async fn udp_tunnel_echoes_without_codec() { + assert_tunnel_echoes(Transport::Udp, false).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, - } - - // abort all thread - echo_server_handle.abort(); - pb_mapper_server_handle.abort(); - pb_mapper_server_cli_handle.abort(); - pb_mapper_client_cli_handle.abort(); +async fn udp_tunnel_echoes_with_codec() { + assert_tunnel_echoes(Transport::Udp, true).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..00bfb72 --- /dev/null +++ b/crates/pb-mapper-testkit/src/echo.rs @@ -0,0 +1,65 @@ +//! 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` 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}; +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]; + // 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(pending_tag.take()); + 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..e6b03c7 --- /dev/null +++ b/crates/pb-mapper-testkit/src/lib.rs @@ -0,0 +1,186 @@ +//! 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; + +/// 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 +/// 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. +/// +/// 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(); + 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 + } +} + +/// 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. +/// +/// 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 => { + 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..fb691b3 --- /dev/null +++ b/crates/pb-mapper-testkit/src/traffic.rs @@ -0,0 +1,230 @@ +//! 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, RAW_TCP_PAYLOAD_MAX, READY_TIMEOUT, UDP_TEST_PAYLOAD_MAX}; + +/// 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`]. 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); + 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(); + // 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 = if payload.is_empty() { + vec![7u8] + } else { + 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 + .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"); + } +} + +/// 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..2eb780f --- /dev/null +++ b/crates/pb-mapper-testkit/src/tunnel.rs @@ -0,0 +1,463 @@ +//! 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 identifying the echo server: per datagram on UDP, once per + /// connection on TCP. + 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 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 + } + + 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, + /// 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>, +} + +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 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; + 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 identifies itself with, if any. + pub fn echo_tag(&self) -> Option { + self.echo_tag + } + + /// 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); + 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 = 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 + } +} + +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() + } +} + +#[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; + } +} 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