From 7789348906c6ae7d2ed6719043c4a2ff4b99b9d0 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Tue, 11 Aug 2026 18:19:02 +0200 Subject: [PATCH 1/2] feat(relay): embed the relay-tunnel receiver behind a `tunnel` feature MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring the SDK-side relay stack into shared-core so SDKs can act as relay receivers without reimplementing the protocol, and without the relay runtime owning the SDK's threads. Everything lives under `src/relay/` gated by `#[cfg(feature = "tunnel")]` (off by default), so the crate stays strictly sans-IO and WASM-compilable when the feature is off. Contents (a fork of the relay repo's crates + a new engine): * protocol/ — was relay-protocol: env/tunnel, path grammar, /whoami, h2 tuning. Sans-IO. Internal `crate::` refs rewritten to `super::`. * bridge.rs — was relay-bridge: the bidirectional h2 bridge (failpoints hooks stripped; shared-core has no `failpoints` feature). * receiver/ — was relay-receiver: dial-out, role-flip, /whoami, redial, multi-homing (R4) + instance-id affinity (R5), InvokeHandler. * loopback.rs — new: the engine. Owns a tokio runtime, runs the receiver with a LoopbackHandler that h2c-dials 127.0.0.1: and bridges each forwarded stream, replaying the tail-only `:path` (preserves request-identity) + headers verbatim. JSON `Config` (serde), `Engine::start` -> `Handle` (status/stop), `Status`. The `tunnel` feature pulls in optional tokio/h2/http/tokio-rustls/ rustls-pki-types/async-trait/serde_json. `http` is now an explicit feature (`dep:http` elsewhere would otherwise suppress its implicit feature and break `#[cfg(feature = "http")]` in headers.rs + any external `features=["http"]`). Verified: default `cargo tree` has no tokio (sans-IO preserved); default tests green; `--features tunnel` builds, `tests/relay_loopback.rs` round-trips a forwarded request end-to-end; clippy `-D warnings`, rustfmt, and rustdoc all clean on both feature sets. No cross-repo coupling: the relay repo keeps its own copies untouched (forked at move time). Design + rationale in docs/relay-receiver-in-shared-core.md. Co-Authored-By: Claude Opus 4.8 (1M context) --- Cargo.toml | 28 + docs/relay-receiver-in-shared-core.md | 176 ++ src/lib.rs | 2 + src/relay/bridge.rs | 2208 +++++++++++++++++++++++++ src/relay/loopback.rs | 471 ++++++ src/relay/mod.rs | 32 + src/relay/protocol/env.rs | 120 ++ src/relay/protocol/error.rs | 91 + src/relay/protocol/h2_tuning.rs | 194 +++ src/relay/protocol/mod.rs | 33 + src/relay/protocol/path.rs | 601 +++++++ src/relay/protocol/tunnel.rs | 165 ++ src/relay/protocol/whoami.rs | 291 ++++ src/relay/receiver/handler.rs | 41 + src/relay/receiver/mod.rs | 44 + src/relay/receiver/runtime.rs | 579 +++++++ tests/relay_loopback.rs | 182 ++ 17 files changed, 5258 insertions(+) create mode 100644 docs/relay-receiver-in-shared-core.md create mode 100644 src/relay/bridge.rs create mode 100644 src/relay/loopback.rs create mode 100644 src/relay/mod.rs create mode 100644 src/relay/protocol/env.rs create mode 100644 src/relay/protocol/error.rs create mode 100644 src/relay/protocol/h2_tuning.rs create mode 100644 src/relay/protocol/mod.rs create mode 100644 src/relay/protocol/path.rs create mode 100644 src/relay/protocol/tunnel.rs create mode 100644 src/relay/protocol/whoami.rs create mode 100644 src/relay/receiver/handler.rs create mode 100644 src/relay/receiver/mod.rs create mode 100644 src/relay/receiver/runtime.rs create mode 100644 tests/relay_loopback.rs diff --git a/Cargo.toml b/Cargo.toml index 6181bf5..d57b6f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,26 @@ sha2_random_seed = ["dep:sha2"] rust_crypto = ["jsonwebtoken?/rust_crypto"] aws_lc_rs = ["jsonwebtoken?/aws_lc_rs"] tracing_pretty = ["dep:tracing-subscriber"] +# `http` was an implicit feature of the optional `http` dependency. Referencing +# that dependency with `dep:` (below, via `tunnel`) suppresses the implicit +# feature, so it is declared explicitly here to keep `#[cfg(feature = "http")]` +# (src/headers.rs) and any external `features = ["http"]` working unchanged. +http = ["dep:http"] + +# Embedded relay-tunnel receiver (see docs/relay-receiver-in-shared-core.md). +# OFF by default so the crate stays strictly sans-IO and WASM-compilable: none +# of these deps enter the default dependency graph. ON pulls in the relay stack +# (src/relay/) + a tokio runtime so the crate can act as a relay receiver, +# bridging each forwarded stream to a local HTTP/2 server over a loopback socket. +tunnel = [ + "dep:tokio", + "dep:h2", + "http", + "dep:tokio-rustls", + "dep:rustls-pki-types", + "dep:async-trait", + "dep:serde_json", +] [dependencies] thiserror = "2.0" @@ -35,6 +55,14 @@ jsonwebtoken = { version = "10.3", default-features = false, optional = true } http = { version = "1.3", optional = true } +# --- `tunnel` feature deps (all optional; absent from the default/WASM graph) --- +tokio = { version = "1", optional = true, features = ["rt-multi-thread", "net", "io-util", "sync", "time", "macros"] } +h2 = { version = "0.4", optional = true } +tokio-rustls = { version = "0.26", optional = true, default-features = false, features = ["ring", "tls12"] } +rustls-pki-types = { version = "1", optional = true } +async-trait = { version = "0.1", optional = true } +serde_json = { version = "1", optional = true } + [dev-dependencies] googletest = "0.11" test-log = { version = "0.2", default-features = false, features = ["trace", "color"] } diff --git a/docs/relay-receiver-in-shared-core.md b/docs/relay-receiver-in-shared-core.md new file mode 100644 index 0000000..40c59e4 --- /dev/null +++ b/docs/relay-receiver-in-shared-core.md @@ -0,0 +1,176 @@ +# Relay receiver in shared-core — plan + +**Status:** **landed** (Phase 0–1). **Date:** 2026-08-11. +**Scope:** this repo (`sdk-shared-core`) absorbs the whole SDK-side relay Rust +stack — **`relay-protocol` + `relay-bridge` + `relay-receiver`** — plus a +**loopback engine**, as internal modules gated behind a `tunnel` cargo feature +so the crate still compiles to WASM (sans-IO) when the feature is off. + +> **Implemented.** The stack lives under `src/relay/` (`protocol/`, `bridge.rs`, +> `receiver/`, `loopback.rs`) behind `#[cfg(feature = "tunnel")]`; the feature +> pulls in the optional `tokio`/`h2`/`http`/`tokio-rustls`/`rustls-pki-types`/ +> `async-trait`/`serde_json` deps. Default build is unchanged & sans-IO +> (verified: no tokio in the default `cargo tree`); `--features tunnel` builds +> the engine, and `tests/relay_loopback.rs` proves a forwarded request round- +> trips relay → engine → loopback → local h2c server → back with the tail-only +> `:path` preserved. The native binding surface (`relay::{Config, Engine, +> Handle, Status}`) is what sdk-java/sdk-typescript wrap. Remaining: Phase 2 +> (SDK bindings — sdk-java landed on its own branch) and Phase 3 (CI/release). + +Sibling doc (consuming side, Java): `sdk-java/development/relay-receiver-plan.md`. + +## What this repo is today + +`restate-sdk-shared-core` is a **single, strictly sans-IO crate**: the +invocation VM (`CoreVM` / the `VM` trait) with **zero networking**. That's why +it ships as **WASM** (TS SDK + edge: Workers/Deno) and as a **no-runtime FFM +cdylib** (Java: `librestate_sdk_core.so`, driven by `FfmStateMachine`). +Bindings are vendored per-SDK (TS wasm-bindings; Java `sdk-core/src/main/rust`), +not in this repo. + +## The approach: bring the relay stack in, behind a `tunnel` feature + +Copy the three relay crates + the loopback engine into this crate as modules +gated by `#[cfg(feature = "tunnel")]`, with **no `relay-*` external deps** — only +third-party networking deps, made **optional** and pulled in by the feature: + +```toml +[features] +default = [] # sans-IO; WASM/edge use this +tunnel = ["dep:tokio", "dep:h2", "dep:tokio-rustls", "dep:rustls-pki-types", + "dep:async-trait"] + +[dependencies] +# already present, non-optional: bytes, http (optional today), serde, thiserror, tracing +tokio = { version = "1", optional = true, features = ["rt-multi-thread","net","io-util","sync","time","macros"] } +h2 = { version = "0.4", optional = true } +tokio-rustls = { version = "0.26", optional = true, default-features = false, features = ["ring","tls12"] } +rustls-pki-types = { version = "1", optional = true } +async-trait = { version = "0.1", optional = true } +``` + +Module layout (all `#[cfg(feature = "tunnel")]`): + +``` +src/relay/ +├── mod.rs +├── protocol/ # was relay-protocol: Env/Tunnel, path grammar, whoami/start-tunnel, h2 tuning (sans-IO) +├── bridge.rs # was relay-bridge: bidirectional h2 bridge +├── receiver/ # was relay-receiver: dial-out, role-flip, /whoami, liveness, redial, R4/R5, InvokeHandler +├── loopback.rs # LoopbackHandler + Engine{start,stop,status} + Config +└── ffi.rs # optional: control-plane facade (see Bindings) +``` + +- **Feature OFF (default):** no tokio in the dependency graph, no `src/relay` + compiled → the crate is exactly as sans-IO as today → `wasm-pack` / edge / + Lambda unaffected. This is the load-bearing invariant; the feature gate + the + optional deps are all that enforce it. +- **Feature ON:** the relay stack + loopback engine compile in and the crate + exposes the tunnel API. Native builds (Java cdylib, Node napi) enable it. + +Payoff: the whole SDK-side relay stack is single-sourced here, the VM stays +WASM-able, and there is **no cross-repo coupling** — see the note below. + +## No cross-repo coupling (the deliberate trade) + +The three crates are **copied in**, not shared with `restatedev/relay`. The +relay repo keeps its own `relay-protocol` / `relay-bridge` / `relay-receiver` +untouched — the broker still builds, its `receiver_e2e` still uses the relay +repo's own copy, nobody publishes anything, no dependency inverts. The one +honest cost is that the code is **forked at move time**; keeping the two copies +in sync (if ever needed) is a manual/later concern — the same kind of workflow +the existing `translate-from-shared-core` skill already handles. Independence +was chosen over DRY, on purpose. + +## Artifact × consumer matrix + +| Consumer | Build | Contents | +|---|---|---| +| TS edge / Workers / Deno / Lambda | WASM, **default features** | VM only — sans-IO, unchanged | +| Java (all) | FFM cdylib, `--features tunnel` | one `librestate_sdk_core.so` exporting `vm_*` + `relay_tunnel_*`; tunnel **inert** unless started | +| TS Node **server** | napi addon, `--features tunnel` | tunnel engine (bridges to the loopback port where the WASM-VM node-http2 server listens) | + +Java folds the tunnel into its existing cdylib (it already links the VM +natively). TS keeps VM-via-WASM (default) + tunnel-via-napi (feature on) as two +independent artifacts — the tunnel addon doesn't need the VM. + +## The one-lib decision for Java (measured) + +Enabling `tunnel` grows the Java cdylib, it doesn't add a second one. Measured +tunnel cdylib (relay workspace `lto=fat`, stripped, `ring`, **TLS linked**): + +| | size | +|---|---| +| VM lib today | 1.60 MB | +| + tunnel, stripped | ~1.5–1.8 MB added | +| + tunnel, gzipped (jar-stored) | ~0.7–0.85 MB added | + +An order of magnitude under the ~15–20 MB that would argue for a split. Lambda +ships ~0.7 MB gzipped of dead bytes (tokio spawns only on `relay_tunnel_start`). + +## The engine (`src/relay/loopback.rs`) + +- `LoopbackHandler: InvokeHandler` — per forwarded `Invocation`: h2c-dial + `127.0.0.1:`, replay method + **tail-only `:path`** (preserves + request-identity, signed over the SDK-relative path) + headers verbatim, + bridge both bodies via the in-crate `bridge`. +- `Engine::start(Config) -> Handle` builds a tokio runtime, spawns + `receiver::Receiver::run(shutdown)`, holds the runtime + `Arc`; `stop` + fires the shutdown future + joins; `status` reads the shared state. +- `Config { relay: Vec, env, tunnel, api_key, local_port, tls, + connections (R4), instance_id (R5), h2_tuning }`. + +## Binding surfaces + +Bindings stay vendored per-SDK and just enable `shared-core/tunnel`: + +- **Java** — `sdk-core/src/main/rust` enables the feature and adds the + `relay_tunnel_*` C ABI (`start`/`stop`/`status`, JSON config, `free_string`) + next to the VM's `vm_*` → one cdylib. jextract's header covers both. +- **TS** — a new napi wrapper enables the feature and exports the same three + control functions. No `ThreadsafeFunction`, no per-request marshalling — data + rides the loopback socket. + +Config crosses as a **UTF-8 JSON string**; `status()` reads an `Arc`; +**zero upcalls**. (Optionally a `src/relay/ffi.rs` facade could live here behind +the feature to single-source the C ABI; not required for v1.) + +## Phases + +- **Phase 0 — spike.** Prove the mechanism (can be done before the move, in the + relay repo): loopback engine + a `relay_tunnel_*` C ABI round-tripping a + relay-forwarded request → loopback → an HTTP/2 server → back, vs `cargo run + --bin relay examples/relay.toml` (whoami mode). Retires (a) encapsulated-runtime + start/stop over FFI, (b) prior-knowledge h2c bidi into the SDK server. +- **Phase 1 — bring the stack in.** Copy `relay-protocol` / `relay-bridge` / + `relay-receiver` into `src/relay/` + add `loopback.rs`, all behind + `#[cfg(feature = "tunnel")]`; make the networking deps optional. **Gate:** + default build + WASM compile unchanged (sans-IO, no tokio in the default + graph); `--features tunnel` builds the engine; existing VM tests green; + add tunnel tests under the feature. +- **Phase 2 — SDK bindings.** sdk-java enables the feature + adds the + `relay_tunnel_*` C ABI (one cdylib); sdk-typescript adds the napi package; its + WASM/VM path is untouched. +- **Phase 3 — CI/release.** Build the Java cdylib + Node napi `--features + tunnel` across the platform matrix; WASM keeps building default. crates.io + source publish unchanged (name + semver + default API stable). + +## Risks + +1. **WASM contamination** (#1). A networking dep leaking into a non-optional + path, or a `src/relay` module not gated, breaks the WASM/edge build. Enforce + with a CI check: the **default-feature** dependency graph must contain no + tokio. +2. **Published-crate stability.** `restate-sdk-shared-core` v7.x has real + crates.io consumers; the **default** build + public API must stay identical. + A new optional feature is additive and safe. +3. **napi ⇄ FFM parity.** One engine, two bindings — keep the control surface + (start/stop/status + JSON config) identical so behaviour can't drift. + +## Scope cuts for v1 + +- **`/whoami` mode** (target the relay; `/_/start-tunnel` / real Cloud is a + later receiver extension). +- **Loopback TCP**, not UDS. +- **No drain** capability advertised. +- Java **FFM only** (23+); napi for Node servers. diff --git a/src/lib.rs b/src/lib.rs index b32cc61..901a40e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,8 @@ pub mod error; pub mod fmt; mod headers; +#[cfg(feature = "tunnel")] +pub mod relay; #[cfg(feature = "request_identity")] mod request_identity; mod retries; diff --git a/src/relay/bridge.rs b/src/relay/bridge.rs new file mode 100644 index 0000000..a346640 --- /dev/null +++ b/src/relay/bridge.rs @@ -0,0 +1,2208 @@ +//! Bidirectional HTTP/2 stream bridge. +//! +//! Given: +//! * A sender-side h2 stream — the cluster received this as an HTTP/2 +//! *server* on :9080. We have a `RecvStream` for the request body and a +//! `SendResponse` for writing the response. +//! * A receiver-side h2 stream — the cluster opened this as an HTTP/2 +//! *client* on the reversed connection. We have a `SendStream` for +//! writing the request body and a `ResponseFuture`/`RecvStream` for +//! reading the response. +//! +//! The bridge pumps bytes in both directions until either: +//! * Both streams terminate cleanly (success path), or +//! * Either side errors (failure path — propagate per design §6). +//! +//! ## Coupling the two halves ([`run_duplex`]) +//! +//! Bridging upstream (sender→receiver) and downstream (receiver→sender) are +//! concurrent, but they are NOT independent at termination. A plain +//! `tokio::join!` waits for both regardless of outcome, so a failure on one +//! side leaves the other lingering until its own I/O happens to error — or +//! forever, if it is parked on an idle read (`reader.data().await` has no +//! timeout). [`run_duplex`] couples them: +//! +//! * **One half errors** → the other is dropped immediately. Dropping a +//! pump future drops its h2 send/recv halves, which resets the stream, +//! so the failure propagates in both directions at once. +//! * **One half completes cleanly (EOS)** → the other is *not* cancelled. +//! HTTP/2 half-close is legal and the surviving direction may keep +//! streaming (e.g. a server-streaming response after the request body +//! finished, or a long upload after the response headers). We keep +//! pumping it. +//! +//! The remaining gap is a *clean* half-close where the surviving direction's +//! peer then goes silent without closing (a silently-dead or adversarial +//! sender holding a half-open stream). For receiver connections this is +//! already covered by the PING liveness loop (`src/liveness.rs`), which +//! tears the connection down and errors its streams. Sender connections +//! have no such probe yet — see `BridgeConfig::half_close_drain_timeout` +//! for an opt-in bound; sender-side PING liveness is the proper fix and a +//! known follow-up. A *default* drain/idle timeout is +//! deliberately avoided: it would cut off legitimate long or sparse +//! one-directional streaming. +//! +//! ## Flow control +//! +//! For each direction, the sending side has an h2 send-window. We must +//! reserve capacity before sending DATA. The receiving side has a window +//! that we acknowledge as we consume bytes. Both `h2::SendStream::reserve_capacity` / +//! `available_capacity` and `RecvStream::flow_control().release_capacity()` +//! are the primitives. +//! +//! Per design §8: respect both windows; don't ack upstream faster than +//! downstream drains; bound memory held inside the cluster. + +use bytes::{Bytes, BytesMut}; +use std::future::Future; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; +use std::time::Duration; +use tracing::{debug, warn}; + +/// Cap on bytes coalesced into a single `send_data` call. Picked at +/// `max_frame_size` (64 KiB) so the resulting DATA frame doesn't +/// straddle h2's wire-frame boundary — h2 would just re-fragment a +/// larger payload internally, which would erase the saving. Sources +/// of "many small chunks" worth coalescing (drip producers, +/// fragmented application-level frames) typically have total +/// in-flight bytes well under this cap anyway. +const COALESCE_CAP_BYTES: usize = 64 * 1024; + +/// A pump flushes its locally-accumulated held-bytes to the shared +/// [`BufferMeter`] only when the value moves by at least this much. Healthy +/// flow (held ≈ 0 on the fast path) does zero shared writes; a backpressured +/// stream flushes a handful of times as its backlog climbs/drains. +/// +/// It is also the **per-pump quantization bound**: [`BufferMeterGuard::observe`] +/// keeps `|held − published| < BUFFER_FLUSH_THRESHOLD` at all times, so each pump +/// under-reports its held bytes by strictly less than this. `pub` so node +/// admission can reserve headroom for that residual (`PUMPS_PER_BRIDGE × +/// BUFFER_FLUSH_THRESHOLD` per bridged stream) as a true bound rather than a +/// heuristic. +pub const BUFFER_FLUSH_THRESHOLD: i64 = 64 * 1024; + +/// Pumps per bridged stream (upstream + downstream), each with its own +/// [`BufferMeterGuard`]. The worst-case unreported buffered bytes for one +/// bridge is therefore `PUMPS_PER_BRIDGE × BUFFER_FLUSH_THRESHOLD`. +pub const PUMPS_PER_BRIDGE: usize = 2; + +/// Cheap, cloneable handle to a node-global "buffered bytes" accumulator that +/// the bridge feeds so memory-aware admission can see transient stall-buffering +/// (`Σ used_capacity` across streams). See `relay::node_admission`. +/// +/// **Contention:** the per-chunk hot path touches NO shared memory. Each pump +/// accumulates its stream's held bytes locally in a [`BufferMeterGuard`] and +/// writes this shared atomic only on coarse [`BUFFER_FLUSH_THRESHOLD`] +/// crossings — and only on the backpressured slow path (a healthy fast-path +/// stream releases each chunk the same iteration, so it never registers held +/// bytes and never flushes). So the shared atomic sees writes only from +/// genuinely-backpressured streams, at a low rate. `None` disables it entirely +/// (a non-broker embedder, and the broker when admission is off). +#[derive(Clone, Default)] +pub struct BufferMeter { + counter: Option>, +} + +impl BufferMeter { + /// A meter backed by a shared node-global counter. + pub fn new(counter: Arc) -> Self { + Self { + counter: Some(counter), + } + } + + /// A no-op meter: `add` does nothing, `is_enabled` is false. + pub fn disabled() -> Self { + Self { counter: None } + } + + /// Whether this meter is backed by a real counter. + pub fn is_enabled(&self) -> bool { + self.counter.is_some() + } + + /// Apply a signed delta to the shared counter. `Relaxed` is sufficient — + /// the value is a soft admission signal read on a separate slow path, not a + /// synchronization point. Public so feeders other than the bridge pump (and + /// tests) can contribute; the pump goes through [`BufferMeterGuard`]. + #[inline] + pub fn add(&self, delta: i64) { + if let Some(c) = &self.counter { + if delta != 0 { + c.fetch_add(delta, Ordering::Relaxed); + } + } + } +} + +impl std::fmt::Debug for BufferMeter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "BufferMeter({})", + if self.counter.is_some() { "on" } else { "off" } + ) + } +} + +/// Per-pump local accumulator that feeds a [`BufferMeter`] with this stream's +/// held (received-but-not-released) bytes, flushing to the shared counter only +/// on coarse threshold crossings. `Drop` zeroes this pump's contribution on +/// ANY exit path (mirrors the owner's `StreamFinishGuard`), so a reset/cancel +/// can never leak the counter. +struct BufferMeterGuard<'a> { + meter: &'a BufferMeter, + /// The value we last pushed to the shared counter (our current contribution). + published: i64, +} + +impl<'a> BufferMeterGuard<'a> { + fn new(meter: &'a BufferMeter) -> Self { + Self { + meter, + published: 0, + } + } + + /// Reconcile against this stream's TRUE current `used_capacity`. Pushes the + /// delta to the shared counter only when it has moved by at least + /// [`BUFFER_FLUSH_THRESHOLD`] since the last push — so the shared write is + /// rare. A no-op when the meter is disabled. + #[inline] + fn observe(&mut self, used_capacity: usize) { + if !self.meter.is_enabled() { + return; + } + let held = used_capacity as i64; + if (held - self.published).abs() >= BUFFER_FLUSH_THRESHOLD { + self.meter.add(held - self.published); + self.published = held; + } + } +} + +impl Drop for BufferMeterGuard<'_> { + fn drop(&mut self) { + // Remove exactly what we contributed (the adds telescope to `published`). + if self.published != 0 { + self.meter.add(-self.published); + } + } +} + +/// Errors surfaced from the bridge. The handler converts these to sender- +/// visible HTTP status codes per design §6. +#[derive(Debug, thiserror::Error)] +pub enum BridgeError { + /// h2 error reading or writing on the sender-side stream. + #[error("sender-side h2 error: {0}")] + Sender(h2::Error), + /// h2 error reading or writing on the receiver-side stream. + #[error("receiver-side h2 error: {0}")] + Receiver(h2::Error), + /// Sender capacity-reservation timed out (upstream stalled). + #[error("flow control wait timed out")] + FlowControlTimeout, + /// This half was cancelled by [`run_duplex`] because the other + /// direction ended first — either it errored (immediate teardown) or it + /// closed cleanly and this half exceeded the half-close drain timeout. + /// The *cause* (if any) is carried by the other half's result. + #[error("bridge half cancelled (other direction ended first)")] + Cancelled, +} + +impl BridgeError { + /// True if this is a **sender-side** error that means the sender peer + /// closed its own stream on its own terms — a remote reset/GOAWAY (any + /// reason), an explicit `CANCEL`, or the `InactiveStreamId` race where our + /// write to the sender lost to the sender's concurrent close. Per design + /// §6 these are normal client closes, not relay/receiver faults, so the + /// handler counts them as `SenderClosed` (not `BridgeError`) and the + /// bridge logs them at `debug`. + /// + /// Returns `false` for `Receiver`/`FlowControlTimeout`/`Cancelled` and for + /// relay-*caused* `Sender` errors (library resets, `PayloadTooBig`, etc.) + /// — those are genuine faults. See `sender_went_away` for the predicate. + pub fn is_benign_sender_close(&self) -> bool { + matches!(self, BridgeError::Sender(e) if sender_went_away(e)) + } + + /// A coarse, stable bucket for WHY this bridge half failed — the label for + /// the `relay_bridge_reset_total{reason}` breakdown. Maps the h2 RST_STREAM + /// code (or a transport/IO error) to a fixed string so operators can tell, + /// e.g., a receiver `protocol_error` / `cancel` apart from a `broken_pipe` + /// connection drop or a flow-control stall. The set of returned values is + /// finite and mirrored by `metrics_defs::BRIDGE_RESET_REASONS`. + pub fn reset_reason(&self) -> &'static str { + match self { + BridgeError::FlowControlTimeout => "flow_control_timeout", + BridgeError::Cancelled => "cancelled", + BridgeError::Sender(e) | BridgeError::Receiver(e) => h2_reset_reason(e), + } + } +} + +/// Bucket an `h2::Error` by its RST_STREAM reason (or transport/IO nature). +/// Kept in sync with `metrics_defs::BRIDGE_RESET_REASONS`. +fn h2_reset_reason(e: &h2::Error) -> &'static str { + match e.reason() { + Some(r) if r == h2::Reason::NO_ERROR => "no_error", + Some(r) if r == h2::Reason::PROTOCOL_ERROR => "protocol_error", + Some(r) if r == h2::Reason::INTERNAL_ERROR => "internal_error", + Some(r) if r == h2::Reason::FLOW_CONTROL_ERROR => "flow_control_error", + Some(r) if r == h2::Reason::STREAM_CLOSED => "stream_closed", + Some(r) if r == h2::Reason::REFUSED_STREAM => "refused_stream", + Some(r) if r == h2::Reason::CANCEL => "cancel", + Some(r) if r == h2::Reason::ENHANCE_YOUR_CALM => "enhance_your_calm", + // Any other RST_STREAM code (FRAME_SIZE_ERROR, COMPRESSION_ERROR, …). + Some(_) => "other_h2", + // No RST_STREAM reason: a transport/IO error (e.g. broken pipe / reset + // connection) or a library-level error the SDK never turned into a code. + None if e.is_io() => "broken_pipe", + None => "other", + } +} + +/// Response headers/trailers whose (lowercased) name starts with this reserved +/// prefix are **relay-internal control signals** (see the relay crate's +/// `cluster::InternalResult` — `x-relay-internal-version` / `-result`). When +/// [`BridgeConfig::strip_reserved_relay_headers`] is set, the bridge removes +/// every such field from a forwarded *application* response so a receiver can +/// never smuggle one to the origin and forge a retry (double-execute). The relay +/// crate references THIS constant when emitting the signals, so the two sides +/// cannot drift. +pub const RESERVED_RELAY_RESPONSE_HEADER_PREFIX: &str = "x-relay-internal-"; + +/// Remove every field whose name starts with [`RESERVED_RELAY_RESPONSE_HEADER_PREFIX`]. +/// Used on both the response head and the trailers. Allocation-free in the +/// common case (no reserved fields → the `to_remove` collect yields an empty, +/// non-allocating `Vec`). +fn strip_reserved_relay_headers(headers: &mut http::HeaderMap) { + let to_remove: Vec = headers + .keys() + .filter(|n| { + n.as_str() + .starts_with(RESERVED_RELAY_RESPONSE_HEADER_PREFIX) + }) + .cloned() + .collect(); + for name in to_remove { + headers.remove(&name); + } +} + +/// Configuration knobs for the bridge. +#[derive(Debug, Clone)] +pub struct BridgeConfig { + /// Maximum time to wait for flow-control capacity on a single chunk. + /// Prevents pathological wedge if a peer never opens its window. + pub flow_control_timeout: Duration, + /// Optional bound on how long the *surviving* bridge half may keep + /// running after the other half has closed **cleanly** (a legal HTTP/2 + /// half-close). `None` (the default) leaves it unbounded — correct for + /// full-duplex workloads where one direction can legitimately stream + /// long after the other finishes (server-streaming, sparse event + /// streams). Set to `Some(_)` to bound a half-open stream's lifetime + /// when a peer goes silent without closing — at the cost of cutting off + /// legitimately-long one-directional streaming past the bound. See + /// [`run_duplex`]. This does NOT apply when a half *errors* — that + /// always tears the other down immediately. The proper fix for + /// silently-dead senders is sender-side PING liveness (TODO F2), not + /// this knob. + pub half_close_drain_timeout: Option, + /// Optional node-global buffered-bytes meter for memory-aware admission. + /// When enabled, the pump feeds this stream's held (received-but-not- + /// released) bytes into it so the admission gate can see transient + /// stall-buffering on top of its per-stream footprint accounting. Fed + /// contention-light (local accumulation, coarse flush, slow-path only — + /// see [`BufferMeter`]). Default [`BufferMeter::disabled`] (no accounting; + /// a non-broker embedder and the admission-off broker leave it here). + pub buffer_meter: BufferMeter, + /// When `true`, strip every response header **and trailer** whose name + /// starts with [`RESERVED_RELAY_RESPONSE_HEADER_PREFIX`] before forwarding + /// the response. The cross-node forward server sets this so a receiver + /// *application* response can never carry a forged relay-internal control + /// signal (`x-relay-internal-*`) to the origin, where it could drive a + /// spurious retry (double-execute). Relay-*generated* internal replies + /// bypass the bridge, so they are never stripped. Default `false` — a + /// non-broker embedder forwards headers verbatim. + pub strip_reserved_relay_headers: bool, +} + +impl Default for BridgeConfig { + fn default() -> Self { + Self { + flow_control_timeout: Duration::from_secs(60), + half_close_drain_timeout: None, + buffer_meter: BufferMeter::disabled(), + strip_reserved_relay_headers: false, + } + } +} + +/// Run the two bridge halves as a coupled duplex pair, replacing a bare +/// `tokio::join!` so termination propagates correctly (see the module-level +/// "Coupling the two halves" docs). +/// +/// Returns `(upstream_result, downstream_result)` — the same tuple shape +/// `tokio::join!(up, down)` produced, so callers' per-direction metric +/// recording is unchanged. A half cancelled because the other ended first +/// is reported as [`BridgeError::Cancelled`]; the real cause (if any) is in +/// the other half's result. +/// +/// Semantics: +/// * first half to finish **errors** → the other is dropped at once +/// (resetting its stream) and reported `Cancelled`; +/// * first half to finish is **Ok** (clean half-close) → the survivor is +/// awaited, bounded by [`BridgeConfig::half_close_drain_timeout`] if +/// set (else unbounded). +pub async fn run_duplex( + up: U, + down: D, + config: &BridgeConfig, +) -> (Result<(), BridgeError>, Result<(), BridgeError>) +where + U: Future>, + D: Future>, +{ + tokio::pin!(up); + tokio::pin!(down); + + tokio::select! { + up_res = &mut up => { + if up_res.is_err() { + // Upstream errored — don't let downstream linger on an idle + // read. Dropping `down` at return resets its stream. + (up_res, Err(BridgeError::Cancelled)) + } else { + let down_res = drain_survivor(&mut down, config).await; + (up_res, down_res) + } + } + down_res = &mut down => { + if down_res.is_err() { + (Err(BridgeError::Cancelled), down_res) + } else { + let up_res = drain_survivor(&mut up, config).await; + (up_res, down_res) + } + } + } +} + +/// Await the surviving bridge half after the other closed cleanly, bounded +/// by the optional half-close drain timeout. On timeout the survivor is +/// dropped by the caller (resetting its stream) and reported `Cancelled`. +async fn drain_survivor(fut: F, config: &BridgeConfig) -> Result<(), BridgeError> +where + F: Future>, +{ + match config.half_close_drain_timeout { + Some(t) => tokio::time::timeout(t, fut) + .await + .unwrap_or(Err(BridgeError::Cancelled)), + None => fut.await, + } +} + +/// A type-erased, `Send` drop-guard the caller wants released **only after** the +/// upstream pump has taken over buffered-bytes accounting for this stream — the +/// "peek→bridge accounting handoff" (see the [`bridge_upstream_with_handoff`] +/// docs). +/// +/// The bridge crate does not know what the guard *is* (in the broker it is a +/// node-admission `PeekReservation`); it only needs that dropping it runs the +/// concrete guard's `Drop`. The pump publishes whatever the pre-bridge window +/// buffered into the [`BufferMeter`] and *then* drops this guard, so the held +/// bytes are covered by the meter the instant the transient reservation goes +/// away — never a window where they are unaccounted. If the pump future is +/// dropped before it runs (e.g. the other bridge half errored first), the guard +/// is simply dropped too, releasing the reservation — exactly once either way. +pub type AccountingHandoff = Box; + +/// Pump request body bytes from the sender's `RecvStream` to the receiver's +/// `SendStream`. Returns when either side ends or errors. +#[tracing::instrument(level = "info", name = "bridge.upstream", skip_all)] +pub async fn bridge_upstream( + sender_body: h2::RecvStream, + receiver_send: h2::SendStream, + config: &BridgeConfig, +) -> Result<(), BridgeError> { + pump( + sender_body, + receiver_send, + config, + PumpDirection::Upstream, + None, + ) + .await +} + +/// Like [`bridge_upstream`] but carries an [`AccountingHandoff`] guard that is +/// released only once this pump has taken over buffered-bytes accounting. +/// +/// This closes the **peek→bridge handoff gap**. Before the bridge starts, the +/// sender's request body can buffer in the h2 `RecvStream` (up to one stream +/// flow-control window) while the broker peeks the receiver's response head; the +/// broker covers that transient with a separate node-admission *peek +/// reservation*. If that reservation were released the instant the peek +/// resolved — before this pump observed the still-buffered bytes — there would +/// be a scheduling window where those bytes are accounted by neither the peek +/// reservation nor the meter. Instead the broker hands the reservation guard +/// here: the pump publishes the currently-held bytes into the [`BufferMeter`] +/// and *then* drops the guard, so `used` never dips below the true held amount +/// during the transition (it is momentarily double-counted, which is safe). +/// +/// The handoff is a single, once-per-stream meter reconcile at pump start — not +/// a per-chunk write — so the healthy fast path (peek resolved within an RTT, +/// ~0 buffered) publishes nothing and stays free of shared writes. +#[tracing::instrument(level = "info", name = "bridge.upstream", skip_all)] +pub async fn bridge_upstream_with_handoff( + sender_body: h2::RecvStream, + receiver_send: h2::SendStream, + config: &BridgeConfig, + handoff: Option, +) -> Result<(), BridgeError> { + pump( + sender_body, + receiver_send, + config, + PumpDirection::Upstream, + handoff, + ) + .await +} + +/// Pump response body bytes from the receiver's `RecvStream` to the sender's +/// `SendStream`. Sends the response head first. +/// +/// Awaits the response future internally — callers that need to inspect +/// the response status before committing to the bridge (e.g. for +/// admission-driven retry) should use [`bridge_response_body`] instead, +/// which takes an already-resolved response. +#[tracing::instrument(level = "info", name = "bridge.downstream", skip_all)] +pub async fn bridge_downstream( + receiver_response: h2::client::ResponseFuture, + sender_respond: &mut h2::server::SendResponse, + config: &BridgeConfig, +) -> Result<(), BridgeError> { + let response = receiver_response.await.map_err(BridgeError::Receiver)?; + bridge_response_body(response, sender_respond, config).await +} + +/// Like [`bridge_downstream`] but takes the response *already resolved*. +/// Used by the sender handler's retry loop, which awaits the response +/// future itself so it can inspect status (and trigger a retry on 503 +/// before any sender body has been pumped upstream). +/// +/// The receiver-side `cluster::forward` also uses this on the branch where its +/// bounded response-head peek RESOLVES: it inspects the head to detect a +/// capable owner's negotiated client-drain refusal (`503 + +/// x-restate-tunnel-draining`, Gate A) and translate it into a relay-internal +/// `client-draining` result instead of bridging the sentinel back. It falls +/// back to [`bridge_downstream`] only on the deferred / RPC-timeout branch, +/// where the response head has not arrived within the peek window (so no drain +/// decision is possible and the stream simply proceeds). +#[tracing::instrument(level = "info", name = "bridge.response_body", skip_all)] +pub async fn bridge_response_body( + response: http::Response, + sender_respond: &mut h2::server::SendResponse, + config: &BridgeConfig, +) -> Result<(), BridgeError> { + let (parts, receiver_body) = response.into_parts(); + // Reconstruct an http::Response<()> from the parts to hand to the + // sender side. + let mut response_head = http::Response::new(()); + *response_head.status_mut() = parts.status; + *response_head.headers_mut() = parts.headers; + *response_head.version_mut() = parts.version; + + // Scrub the relay-internal namespace from an *application* response before + // it leaves this node (cross-node forward sets this). A receiver could + // otherwise stamp `x-relay-internal-result` on its own response and forge a + // retry on the origin — the double-execute vector. Relay-generated internal + // replies bypass the bridge and are unaffected. + if config.strip_reserved_relay_headers { + strip_reserved_relay_headers(response_head.headers_mut()); + } + + // Send the response head. `end_of_stream = false`: body follows. + let sender_send = sender_respond + .send_response(response_head, false) + .map_err(BridgeError::Sender)?; + + pump( + receiver_body, + sender_send, + config, + PumpDirection::Downstream, + None, + ) + .await +} + +/// Which direction a [`pump`] call is moving bytes in. +/// +/// Both directions share the same `read → reserve → send → release` shape; +/// the direction only determines (a) which side a read-error vs a +/// write-error blames in [`BridgeError`] and (b) the log prefix. +#[derive(Copy, Clone)] +enum PumpDirection { + /// Sender's request body → receiver's send stream. + Upstream, + /// Receiver's response body → sender's send stream. + Downstream, +} + +impl PumpDirection { + /// Error variant to use when the *read* side fails. + fn read_err(self, e: h2::Error) -> BridgeError { + match self { + Self::Upstream => BridgeError::Sender(e), + Self::Downstream => BridgeError::Receiver(e), + } + } + + /// Error variant to use when the *write* side fails. + fn write_err(self, e: h2::Error) -> BridgeError { + match self { + Self::Upstream => BridgeError::Receiver(e), + Self::Downstream => BridgeError::Sender(e), + } + } + + /// Log prefix. + fn label(self) -> &'static str { + match self { + Self::Upstream => "bridge upstream", + Self::Downstream => "bridge downstream", + } + } + + /// RST_STREAM reason to send on the *writer* when the *reader* fails, + /// per design §6. The two directions mean different things: + /// * Upstream read failure = the sender vanished mid-request. Per §6 + /// we `CANCEL` the receiver stream ("we no longer need this"); the + /// reversed connection stays up. + /// * Downstream read failure = the receiver died mid-response. Per §6 + /// we send `INTERNAL_ERROR` to the sender ("forwarding broke"). + fn read_error_reset_reason(self) -> h2::Reason { + match self { + Self::Upstream => h2::Reason::CANCEL, + Self::Downstream => h2::Reason::INTERNAL_ERROR, + } + } +} + +/// Shared pump loop: read DATA from `reader`, forward to `writer`, +/// release the reader's window. Returns when the reader ends or either +/// side errors. +/// +/// ## Zero-copy on the happy path +/// +/// `bytes::Bytes` is refcounted. `RecvStream::data()` returns a `Bytes` +/// that refs into h2's recv buffer; passing the same `Bytes` to +/// `SendStream::send_data` queues the refcount in h2's send buffer. +/// The bridge never copies the payload itself — the only copies are the +/// two irreducible kernel↔user transitions (recv into h2's buffer, send +/// from h2's buffer). True zero-copy across HTTP/2 would need a +/// buffer-owning h2 impl (doesn't exist in mainline Rust — see +/// `docs/research-2026-05-off-the-shelf.md`). +/// +/// ## Flow control +/// +/// Two strategies stacked on each other: +/// * **Fast path** (the typical case with tuned h2 windows): if +/// `writer.capacity()` already covers the incoming chunk, just call +/// `send_data` — no `reserve_capacity` / `poll_capacity` await dance. +/// That's the common case at small-to-medium chunk sizes with the +/// recommended `initial_window_size` (1 MiB+). +/// * **Slow path** (window has filled): explicitly `reserve_capacity` + +/// `poll_capacity`-with-timeout. Partial grants are handled correctly +/// via `Bytes::split_to(granted)` (refcounted slice — no allocation, +/// no copy, no tail loss); the leftover tail is carried into the next +/// iteration via `pending`. +/// +/// Either way, `release_capacity(n)` on the reader's flow control +/// happens *after* the bytes have been handed to the writer, so a +/// stalled writer applies backpressure all the way back to the reader. +/// True if this h2 error means the **sender peer closed its own stream** — +/// not a relay/receiver fault. Three shapes, all "the sender went away": +/// * a remote-initiated reset/GOAWAY (`is_remote()`), any reason — the peer +/// explicitly tore the stream/connection down; +/// * an explicit `CANCEL` reason — benign regardless of initiator ("no +/// longer needed"); this is how a completed Restate bidi invocation ends; +/// * the `InactiveStreamId` user-error race — our write to the sender lost +/// to the sender's own concurrent close. This is the suspension/reconnect +/// case (finding 5.7): Restate suspends mid-invocation, drops its stream, +/// and our `send-end`/`send_data` then hits an already-inactive stream. +/// +/// Relay-*caused* sender errors are deliberately EXCLUDED so they keep +/// warning + counting as bridge failures: library-initiated resets +/// (`is_library()` — a protocol error the relay's own h2 stack emitted) and +/// every other `UserError` (PayloadTooBig, ReleaseCapacityTooBig, +/// MalformedHeaders, PollResetAfterSendResponse, …) are real defects, not the +/// sender leaving. +/// +/// `InactiveStreamId` is matched by its `Display` string because h2's +/// `UserError` enum is not public and the variant carries `reason() == None`, +/// `is_remote() == false`, `is_library() == false` — there is nothing else to +/// match on. `h2` is a `=`-pinned dependency so the string is stable; a dep +/// bump must re-confirm it. The guard test `non_cancel_reason_is_not_benign` +/// pins the inverse (a non-CANCEL `Reason` error must NOT be treated benign). +fn sender_went_away(e: &h2::Error) -> bool { + e.is_remote() + || e.reason() == Some(h2::Reason::CANCEL) + || e.to_string() == "user error: inactive stream" +} + +/// True if `e` is a **graceful remote end-of-stream** on the *response* reader: +/// a remote `RST_STREAM(NO_ERROR)`. +/// +/// Some HTTP/2 servers finish a **unary** response by sending the full body and +/// then `RST_STREAM(NO_ERROR)` — *not* a bare END_STREAM. hyper does exactly +/// this (and the Restate Cloud tunnel client is built on hyper): when it has a +/// complete response but the peer's request half is still open, it answers, +/// then resets the stream with NO_ERROR to reclaim it. RFC 7540 §8.1 defines +/// NO_ERROR as "no error" and h2 surfaces queued DATA *before* the reset, so +/// everything already received is the complete message. The relay must deliver +/// what it has and close the sender stream cleanly, NOT propagate an +/// INTERNAL_ERROR (which truncates the body — the sender would see headers but +/// an empty/aborted body). Discovered running the unmodified +/// `restate-cloud-tunnel-client` against the relay; see +/// `docs/cloud-tunnel-compat.md`. +/// +/// Scoped to the **downstream** (receiver→sender) direction. The upstream +/// (sender→receiver) reader already classifies a remote reset as a benign +/// sender close via [`sender_went_away`]; only the downstream path currently +/// treats a remote reset as a hard receiver fault, and only a NO_ERROR reason +/// is unambiguously graceful — INTERNAL_ERROR / CANCEL / etc. stay faults. +fn is_graceful_response_eos(dir: PumpDirection, e: &h2::Error) -> bool { + matches!(dir, PumpDirection::Downstream) + && e.is_remote() + && e.reason() == Some(h2::Reason::NO_ERROR) +} + +/// True if a *write* failure on the **upstream** (request → receiver) direction +/// means the receiver has gracefully closed its side of the stream — so the +/// request-forwarding direction is cleanly *done*, not faulted. +/// +/// The Restate Cloud tunnel client proxies **unary** requests: it answers the +/// forwarded request, then resets the stream (NO_ERROR). The relay is often +/// still finishing the request half (e.g. sending END_STREAM for a bodyless +/// GET), and that write loses the race to the receiver's close — surfacing as +/// either a remote `RST_STREAM(NO_ERROR)`/`CANCEL` or the `InactiveStreamId` +/// user-error (our write hit an already-closed stream). All three mean "the +/// receiver took what it needed and closed", not "forwarding broke". Treating +/// them as a clean upstream completion lets [`run_duplex`] *drain* the response +/// half (deliver the body) instead of tearing it down — without this, the +/// upstream error wins the `run_duplex` race and the response is lost. +/// +/// Scoped to upstream. A downstream (response → sender) write failure is a +/// sender close, classified separately by [`sender_went_away`] / +/// [`BridgeError::is_benign_sender_close`]. +fn write_hit_graceful_peer_close(dir: PumpDirection, e: &h2::Error) -> bool { + matches!(dir, PumpDirection::Upstream) + && ((e.is_remote() && e.reason() == Some(h2::Reason::NO_ERROR)) + || e.reason() == Some(h2::Reason::CANCEL) + || e.to_string() == "user error: inactive stream") +} + +/// Map a write error to a pump result: `Ok(())` when the peer gracefully closed +/// (see [`write_hit_graceful_peer_close`]) so the caller returns clean; a +/// genuine `Err` otherwise (logged at the level its cause warrants). +fn finish_on_write_err( + dir: PumpDirection, + e: h2::Error, + label: &str, + ctx: &str, +) -> Result<(), BridgeError> { + if write_hit_graceful_peer_close(dir, &e) { + debug!("{label}: {ctx}, but the receiver had already closed the request stream — request forwarding complete"); + return Ok(()); + } + let be = dir.write_err(e); + log_bridge_error(label, ctx, &be); + Err(be) +} + +/// Close the writer with a bare END_STREAM frame after the reader signalled a +/// graceful reset (no trailers can follow a reset, so we don't probe for them). +fn finish_writer_after_graceful_reset( + writer: &mut h2::SendStream, + dir: PumpDirection, +) -> Result<(), BridgeError> { + debug!( + "{}: receiver ended the response with RST_STREAM(NO_ERROR) — \ + treating as clean end-of-stream", + dir.label() + ); + if let Err(e) = writer.send_data(Bytes::new(), true) { + let be = dir.write_err(e); + log_bridge_error(dir.label(), "send-end failed", &be); + return Err(be); + } + Ok(()) +} + +/// Release `n` bytes on the reader's flow-control window after forwarding them. +/// Returns `Ok(true)` when the reader has *gracefully* reset (remote NO_ERROR): +/// the bytes are already queued to the writer, so the caller should close the +/// writer with a bare END_STREAM frame and finish. `Ok(false)` is a normal +/// release; `Err` is a genuine read fault. `try_peek_data` uses a no-op waker +/// and so may not have observed a just-arrived reset during coalescing — it can +/// instead surface here, on the post-send `release_capacity`, which is why this +/// site needs the same graceful-EOS treatment as the DATA read. +fn release_or_graceful_eos( + reader: &mut h2::RecvStream, + n: usize, + dir: PumpDirection, + label: &str, +) -> Result { + match reader.flow_control().release_capacity(n) { + Ok(()) => Ok(false), + Err(e) if is_graceful_response_eos(dir, &e) => Ok(true), + Err(e) => { + let be = dir.read_err(e); + log_bridge_error(label, "release_capacity failed", &be); + Err(be) + } + } +} + +/// Log a bridge error at the level its cause warrants. A benign sender close +/// (see [`BridgeError::is_benign_sender_close`]) drops to `debug!`: the +/// primary workload ends every invocation in a sender CANCEL, and each +/// suspension step ends in the `InactiveStreamId` race, so warning on these +/// would flood the log and pollute WARN-rate alerts one-line-per-request. +/// Everything else — receiver faults, flow-control stalls, and relay-caused +/// sender errors — stays a genuine `warn!`. +fn log_bridge_error(label: &str, ctx: &str, err: &BridgeError) { + if err.is_benign_sender_close() { + debug!("{label}: {ctx}: {err} (sender closed its stream — normal)"); + } else { + warn!("{label}: {ctx}: {err}"); + } +} + +async fn pump( + mut reader: h2::RecvStream, + mut writer: h2::SendStream, + config: &BridgeConfig, + dir: PumpDirection, + handoff: Option, +) -> Result<(), BridgeError> { + let label = dir.label(); + // Bytes carried over from a partial-grant slow-path iteration. + let mut pending: Option = None; + // Feeds this stream's held (received-but-not-released) bytes into the + // node-global buffered-bytes meter for memory-aware admission. Only the + // backpressured slow path observes (a healthy fast-path chunk is released + // the same iteration and never held), so healthy flow costs nothing and + // the shared counter sees writes only under real buffering. Drop zeroes + // this stream's contribution on any exit. + let mut bmeter = BufferMeterGuard::new(&config.buffer_meter); + + // Peek→bridge accounting handoff (upstream only; see + // [`bridge_upstream_with_handoff`] and [`AccountingHandoff`]). Before this + // pump ran, the caller held a transient reservation (the node-admission peek + // reservation) covering whatever the sender body buffered in `reader`'s recv + // window during the response-head peek. Reconcile that held amount into the + // meter FIRST, then drop the guard — so the bytes are covered by the meter + // the instant the reservation is released, never unaccounted. `observe` + // respects the flush threshold, so a healthy ~0-buffered handoff writes + // nothing to the shared counter (fast path stays free); a genuinely-buffered + // handoff publishes the held bytes (already bounded by the residual reserve + // for the sub-threshold remainder). Runs once, at pump start. + if let Some(guard) = handoff { + bmeter.observe(reader.flow_control().used_capacity()); + drop(guard); + } + loop { + // Acquire the next chunk. If we have a leftover tail from a + // partial grant, use that; otherwise pull from the reader. + let mut chunk = if let Some(p) = pending.take() { + p + } else { + match reader.data().await { + Some(Ok(b)) => b, + Some(Err(e)) => { + // A graceful remote NO_ERROR reset on the response reader is + // a clean end-of-stream (hyper's unary-response idiom): any + // DATA already forwarded is the complete body, so close the + // sender stream normally instead of resetting it. + if is_graceful_response_eos(dir, &e) { + return finish_writer_after_graceful_reset(&mut writer, dir); + } + let be = dir.read_err(e); + log_bridge_error(label, "read body error", &be); + writer.send_reset(dir.read_error_reset_reason()); + return Err(be); + } + None => { + // Reader's DATA stream is done. Forward the reader's + // trailers (if any) or an empty end-of-stream frame. + // Design §4.2(e)/§4.3 require trailers — the HEADERS + // frame that can follow DATA — to round-trip verbatim. + // `data()` returning None only signals the end of DATA; + // trailers are retrieved separately via `trailers()`. + debug!("{label}: read body finished"); + return finish_writer(&mut reader, &mut writer, dir, config).await; + } + } + }; + + // Coalesce any DATA frames that are *already buffered* on the + // reader's recv queue into a single payload. This costs one + // allocation + memcpy per accumulated chunk, in exchange for + // collapsing many `send_data` calls into one. The win shows + // up under "drip" producers (sender sends many small frames + // back-to-back within a stream). For the fixed-large-body + // common case, `try_peek_data` returns Pending on the second + // call and we don't allocate at all. + // + // `eos_after_coalesce` carries the reader-side EOF discovered + // during peek so we can emit a single `send_data(.., true)` + // closing call instead of going around the outer loop once + // more just to send a zero-length end-of-stream frame. + let mut coalesce_end = CoalesceEnd::Open; + if pending.is_none() && chunk.len() < COALESCE_CAP_BYTES { + (chunk, coalesce_end) = coalesce_pending(&mut reader, chunk, dir)?; + } + + // Graceful NO_ERROR reset surfaced during coalescing: the reader is + // dead. Deliver everything buffered as a single closing frame and + // finish — do NOT fall through to release_capacity / finish_writer, + // both of which would error on the reset stream. The buffer is bounded + // by COALESCE_CAP_BYTES (64 KiB), so a single `send_data` is safe (h2 + // buffers past the peer window; the slow-path backpressure dance is + // only needed for unbounded streaming bodies). + if coalesce_end == CoalesceEnd::GracefulResetEos { + debug!( + "{label}: receiver ended the response with RST_STREAM(NO_ERROR) \ + mid-body — delivering {} buffered byte(s) as clean EOS", + chunk.len() + ); + if let Err(e) = writer.send_data(chunk, true) { + let be = dir.write_err(e); + log_bridge_error(label, "send-end failed", &be); + return Err(be); + } + return Ok(()); + } + let eos_after_coalesce = coalesce_end == CoalesceEnd::CleanEos; + + let len = chunk.len(); + + if len == 0 { + // Zero-length chunk: forward nothing, but still ack the + // reader's window. (`poll_capacity` on a 0-byte reservation + // can never resolve, so the slow path would deadlock; skip + // the whole capacity dance for empty chunks.) + if release_or_graceful_eos(&mut reader, 0, dir, label)? { + return finish_writer_after_graceful_reset(&mut writer, dir); + } + continue; + } + + // Fast path: writer already has enough granted capacity from the + // peer's initial window (or a prior WINDOW_UPDATE). Send directly + // — no `reserve_capacity` / `poll_capacity` await round-trip. + // The common case with tuned `initial_window_size` (1 MiB+) is + // that capacity is plenty for any single chunk. + if writer.capacity() >= len { + // Never merge end-of-stream into this DATA frame, even when the + // coalesce peek saw EOF: trailers may follow, and once an + // end-of-stream frame is sent the stream is closed and trailers + // can no longer be forwarded. `finish_writer` checks for + // trailers and emits them (or an empty EOS frame). The cost is + // one extra zero-length DATA frame on the trailers-absent + // coalesced path — negligible next to losing trailers. + if let Err(e) = writer.send_data(chunk, false) { + return finish_on_write_err(dir, e, label, "send_data failed"); + } + if release_or_graceful_eos(&mut reader, len, dir, label)? { + return finish_writer_after_graceful_reset(&mut writer, dir); + } + // Fast path = the writer kept up, so nothing is held. Only touch the + // meter to flush DOWN if this stream had previously backpressured + // (published != 0) and has now recovered — otherwise skip the + // `used_capacity` read entirely, keeping the healthy path free. + if bmeter.published != 0 { + bmeter.observe(reader.flow_control().used_capacity()); + } + if eos_after_coalesce { + return finish_writer(&mut reader, &mut writer, dir, config).await; + } + continue; + } + + // Slow path: ask h2 for capacity, wait (with timeout) for it to + // be granted, then forward what was granted. + writer.reserve_capacity(len); + // Slow path = the writer is backpressured and we're about to block, + // holding `chunk` (and possibly more h2-buffered). This is exactly the + // buffering that pressures node memory, so record it before the block. + bmeter.observe(reader.flow_control().used_capacity()); + let granted = + match tokio::time::timeout(config.flow_control_timeout, poll_capacity(&mut writer)) + .await + { + Ok(Ok(n)) => n, + Ok(Err(e)) => { + return finish_on_write_err(dir, e, label, "write capacity error"); + } + Err(_elapsed) => { + warn!("{label}: write capacity timed out"); + writer.send_reset(h2::Reason::CANCEL); + return Err(BridgeError::FlowControlTimeout); + } + }; + + // `granted` may be less than `len`. Split the chunk via + // `Bytes::split_to` — refcounted slice, zero copy. The tail + // rides the next iteration via `pending`. + let to_send_len = granted.min(len); + let (to_send, leftover) = if to_send_len >= len { + (chunk, None) + } else { + let mut head = chunk; + let tail = head.split_off(to_send_len); + (head, Some(tail)) + }; + + if let Err(e) = writer.send_data(to_send, false) { + return finish_on_write_err(dir, e, label, "send_data failed"); + } + + // Release exactly what we just forwarded. + if release_or_graceful_eos(&mut reader, to_send_len, dir, label)? { + return finish_writer_after_graceful_reset(&mut writer, dir); + } + // Update the held-bytes meter after draining part of the backlog. + bmeter.observe(reader.flow_control().used_capacity()); + + pending = leftover; + // Slow-path partial grant: don't try to attach the coalesced + // EOS here. Carry on; the next outer iteration will pull the + // EOF naturally via `reader.data().await -> None`. + let _ = eos_after_coalesce; + } +} + +/// Terminate the writer side once the reader's DATA stream has ended +/// (`data()` returned `None`, or the coalesce peek observed EOF). +/// +/// HTTP/2 carries trailers as a HEADERS frame *after* the DATA frames; +/// `data()` returning `None` means only that DATA is done, not that the +/// stream is over. Design §4.2(e)/§4.3 require trailers to be forwarded +/// verbatim, so we pull them via `RecvStream::trailers()` and, if present, +/// emit them with `SendStream::send_trailers` (which carries END_STREAM). +/// With no trailers we send an empty end-of-stream DATA frame, exactly as +/// before. A receiver that puts its status in trailers (e.g. gRPC's +/// `grpc-status`) depends on this. +async fn finish_writer( + reader: &mut h2::RecvStream, + writer: &mut h2::SendStream, + dir: PumpDirection, + config: &BridgeConfig, +) -> Result<(), BridgeError> { + let label = dir.label(); + match reader.trailers().await { + Ok(Some(mut trailers)) => { + // Trailers are a second forgery/leak channel the head-strip misses: + // scrub the relay-internal namespace here too. (The origin never + // reads trailers as a control signal, so this is defense-in-depth, + // not a retry vector — but it keeps the namespace off the wire.) + if config.strip_reserved_relay_headers { + strip_reserved_relay_headers(&mut trailers); + } + if trailers.is_empty() { + // Every trailer field was in the reserved namespace (or there + // were none left): send a bare END_STREAM instead of an empty + // HEADERS frame, so we don't emit a meaningless trailers frame. + if let Err(e) = writer.send_data(Bytes::new(), true) { + return finish_on_write_err(dir, e, label, "send-end failed"); + } + return Ok(()); + } + debug!("{label}: forwarding {} trailer field(s)", trailers.len()); + if let Err(e) = writer.send_trailers(trailers) { + return finish_on_write_err(dir, e, label, "send_trailers failed"); + } + Ok(()) + } + Ok(None) => { + if let Err(e) = writer.send_data(Bytes::new(), true) { + return finish_on_write_err(dir, e, label, "send-end failed"); + } + Ok(()) + } + Err(e) => { + // A graceful NO_ERROR reset surfaces here when the coalesce peek + // already saw EOS via that reset: no trailers follow a reset, so + // close the sender stream cleanly (the body was forwarded above). + if is_graceful_response_eos(dir, &e) { + return finish_writer_after_graceful_reset(writer, dir); + } + // The reader errored while delivering trailers. Reset the + // writer and surface it as a read-side error, mirroring the + // main loop's DATA read-error handling (direction-dependent + // reason, design §6). + let be = dir.read_err(e); + log_bridge_error(label, "read trailers error", &be); + writer.send_reset(dir.read_error_reset_reason()); + Err(be) + } + } +} + +/// How a [`coalesce_pending`] peek left the reader. +#[derive(Copy, Clone, PartialEq)] +enum CoalesceEnd { + /// More DATA may still arrive; keep pumping. The reader stays usable. + Open, + /// Clean end-of-stream (`data()` would return `None`) observed during the + /// peek. The reader is still usable — `finish_writer` reads trailers next. + CleanEos, + /// The reader was gracefully reset (remote `RST_STREAM(NO_ERROR)`) during + /// the peek — hyper's unary-response close. Everything buffered so far is + /// the complete body, and the reader is **dead**: the caller must deliver + /// the buffer with END_STREAM and touch the reader no further (a + /// `release_capacity` or `trailers()` on a reset stream errors). + GracefulResetEos, +} + +/// Drain any DATA chunks the reader has *already buffered* and +/// concatenate them with `first` into a single `Bytes`, up to +/// [`COALESCE_CAP_BYTES`]. Returns the coalesced bytes plus how the reader +/// ended (see [`CoalesceEnd`]). +/// +/// Errors observed during peek propagate as a read error, mirroring +/// the main loop's error handling — except a graceful NO_ERROR reset, which +/// is reported as [`CoalesceEnd::GracefulResetEos`] so the already-buffered +/// body is delivered rather than discarded. +/// +/// Uses a no-op waker to peek `RecvStream::poll_data` without +/// suspending: only chunks that are *already* on h2's recv queue are +/// pulled in. The next `data().await` call re-registers a real waker +/// for future arrivals, so no deadlock. +fn coalesce_pending( + reader: &mut h2::RecvStream, + first: Bytes, + dir: PumpDirection, +) -> Result<(Bytes, CoalesceEnd), BridgeError> { + // Peek once before allocating. The common (no-coalesce) case + // returns Pending here and we hand `first` back as-is — zero + // allocation, zero memcpy. + let next = match try_peek_data(reader) { + PeekOutcome::Pending => return Ok((first, CoalesceEnd::Open)), + PeekOutcome::Eof => return Ok((first, CoalesceEnd::CleanEos)), + PeekOutcome::Chunk(b) => b, + PeekOutcome::Error(e) => { + // Graceful NO_ERROR reset while peeking: `first` is the complete + // body (h2 hands us all queued DATA before the reset). Deliver it + // rather than discarding the chunk we already hold. + if is_graceful_response_eos(dir, &e) { + return Ok((first, CoalesceEnd::GracefulResetEos)); + } + let be = dir.read_err(e); + log_bridge_error(dir.label(), "read body error during peek", &be); + return Err(be); + } + }; + + // At least two chunks present — build an accumulator. Capacity + // hint is generous: this branch only runs when there are >=2 + // chunks already buffered, so we'll likely drain more. + let mut buf = BytesMut::with_capacity(COALESCE_CAP_BYTES.min(first.len() + next.len() + 4096)); + buf.extend_from_slice(&first); + buf.extend_from_slice(&next); + + while buf.len() < COALESCE_CAP_BYTES { + match try_peek_data(reader) { + PeekOutcome::Pending => return Ok((buf.freeze(), CoalesceEnd::Open)), + PeekOutcome::Eof => return Ok((buf.freeze(), CoalesceEnd::CleanEos)), + PeekOutcome::Chunk(more) => { + buf.extend_from_slice(&more); + } + PeekOutcome::Error(e) => { + // Graceful NO_ERROR reset mid-coalesce: everything buffered so + // far is the complete body. Deliver it. + if is_graceful_response_eos(dir, &e) { + return Ok((buf.freeze(), CoalesceEnd::GracefulResetEos)); + } + let be = dir.read_err(e); + log_bridge_error(dir.label(), "read body error during peek", &be); + return Err(be); + } + } + } + Ok((buf.freeze(), CoalesceEnd::Open)) +} + +/// Outcome of a non-suspending `poll_data` peek. +enum PeekOutcome { + /// A chunk was already buffered; consumed by the peek. + Chunk(Bytes), + /// Reader has signalled end-of-stream. + Eof, + /// No chunk currently buffered; would have to await. + Pending, + /// Peek surfaced a read error. + Error(h2::Error), +} + +/// Try to pull a chunk from the reader without yielding. Uses a +/// `Waker::noop()` so h2's internal "wake on new data" doesn't fire +/// anything — fine because the outer loop's next `data().await` will +/// re-register a real waker. +fn try_peek_data(reader: &mut h2::RecvStream) -> PeekOutcome { + let waker = Waker::noop(); + let mut cx = Context::from_waker(waker); + match reader.poll_data(&mut cx) { + Poll::Ready(Some(Ok(b))) => PeekOutcome::Chunk(b), + Poll::Ready(Some(Err(e))) => PeekOutcome::Error(e), + Poll::Ready(None) => PeekOutcome::Eof, + Poll::Pending => PeekOutcome::Pending, + } +} + +/// Wait until the SendStream's `poll_capacity` reports any non-zero value. +/// Returns the granted capacity in bytes, or an h2 error if the stream +/// errors before any is granted. +async fn poll_capacity(stream: &mut h2::SendStream) -> Result { + std::future::poll_fn(|cx| stream.poll_capacity(cx)) + .await + .unwrap_or(Ok(0)) +} + +#[cfg(test)] +mod tests { + //! Bridge tests use a full sender↔cluster↔receiver triangle. They are + //! larger and slower than unit tests; their value is in catching real + //! flow-control and ordering bugs that wouldn't show up in isolation. + + use super::*; + use bytes::Bytes; + use std::time::Duration; + use tokio::net::{TcpListener, TcpStream}; + + // --- is_benign_sender_close (finding 5.6/5.7 classification) ---------- + // + // Only the `reason()`-constructible and variant cases are unit-testable: + // h2 exposes no public constructor for remote resets or `User` errors, so + // the `is_remote()` and `InactiveStreamId` branches are exercised by the + // live Restate e2e (suspension test). The guard test below pins the + // load-bearing inverse: a *non-CANCEL* `Reason` error must NOT be benign, + // so a relay-caused fault keeps warning + counting as bridge_error. + + #[test] + fn strip_reserved_removes_only_reserved_prefix() { + // The sanitizer used by bridge_response_body (head) and finish_writer + // (trailers) must drop the whole x-relay-internal-* namespace while + // leaving every ordinary header/trailer field intact. + let mut h = http::HeaderMap::new(); + h.insert("content-type", "text/plain".parse().unwrap()); + h.insert("x-relay-internal-result", "at-capacity".parse().unwrap()); + h.insert("x-relay-internal-version", "1".parse().unwrap()); + h.insert("x-custom", "keep".parse().unwrap()); + h.insert("grpc-status", "0".parse().unwrap()); // a trailer field + strip_reserved_relay_headers(&mut h); + assert!(!h.contains_key("x-relay-internal-result")); + assert!(!h.contains_key("x-relay-internal-version")); + assert_eq!(h.get("content-type").unwrap(), "text/plain"); + assert_eq!(h.get("x-custom").unwrap(), "keep"); + assert_eq!(h.get("grpc-status").unwrap(), "0"); + } + + #[test] + fn strip_reserved_is_noop_without_reserved_fields() { + let mut h = http::HeaderMap::new(); + h.insert("content-type", "application/json".parse().unwrap()); + strip_reserved_relay_headers(&mut h); + assert_eq!(h.len(), 1); + assert_eq!(h.get("content-type").unwrap(), "application/json"); + } + + #[test] + fn cancel_reason_is_benign_sender_close() { + let e = BridgeError::Sender(h2::Error::from(h2::Reason::CANCEL)); + assert!(e.is_benign_sender_close()); + } + + #[test] + fn non_cancel_reason_is_not_benign() { + // A non-CANCEL, non-remote sender error is a genuine fault. + for reason in [ + h2::Reason::INTERNAL_ERROR, + h2::Reason::PROTOCOL_ERROR, + h2::Reason::FRAME_SIZE_ERROR, + ] { + let e = BridgeError::Sender(h2::Error::from(reason)); + assert!(!e.is_benign_sender_close(), "{reason:?} must not be benign"); + } + } + + #[test] + fn non_sender_variants_are_not_benign() { + // Even a CANCEL on the *receiver* side is a receiver fault, and + // FlowControlTimeout / Cancelled are never sender closes. + assert!( + !BridgeError::Receiver(h2::Error::from(h2::Reason::CANCEL)).is_benign_sender_close() + ); + assert!(!BridgeError::FlowControlTimeout.is_benign_sender_close()); + assert!(!BridgeError::Cancelled.is_benign_sender_close()); + } + + /// Set up: + /// * a "fake sender" — an h2 client that hits the cluster on a TCP listener + /// * a "fake receiver" — an h2 server that's dialed by the cluster on another TCP listener + /// * the bridge wired between them + /// + /// Returns the body the receiver captured and the response body the + /// sender captured. Both ends pump h2 properly (server stays in + /// accept-loop, client drains response fully) so flow control and + /// flushing work end-to-end. + async fn run_bridge_test(request_body: &[u8], response_body: &[u8]) -> (Vec, Vec) { + run_bridge_test_chunked(request_body, response_body, 1, 1).await + } + + /// Same as [`run_bridge_test`] but the sender splits its request + /// body into `req_chunks` `send_data` calls and the receiver + /// splits its response into `resp_chunks` calls. Used to exercise + /// the bridge's DATA-frame coalescing path under "drip" producers. + async fn run_bridge_test_chunked( + request_body: &[u8], + response_body: &[u8], + req_chunks: usize, + resp_chunks: usize, + ) -> (Vec, Vec) { + run_bridge_test_full( + request_body, + response_body, + req_chunks, + resp_chunks, + ReceiverEnding::Clean, + ) + .await + } + + /// How the fake receiver terminates the response stream. + #[derive(Copy, Clone)] + enum ReceiverEnding { + /// Last DATA frame carries END_STREAM — the normal HTTP/2 close. + Clean, + /// Body DATA frames carry `end_of_stream=false`, then the receiver + /// sends `RST_STREAM(NO_ERROR)`. Models hyper's unary-response idiom + /// (used by the Restate Cloud tunnel client): a complete response + /// followed by a graceful reset to reclaim the still-open request half. + ResetNoError, + } + + async fn run_bridge_test_full( + request_body: &[u8], + response_body: &[u8], + req_chunks: usize, + resp_chunks: usize, + ending: ReceiverEnding, + ) -> (Vec, Vec) { + let receiver_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let receiver_addr = receiver_listener.local_addr().unwrap(); + let response_body_vec = response_body.to_vec(); + + // Fake receiver: handle one request, send the response, then keep + // accept-polling to flush + drive the connection until the cluster + // closes its side. + let (recv_body_tx, recv_body_rx) = tokio::sync::oneshot::channel::>(); + let recv_task = tokio::spawn(async move { + let (socket, _) = receiver_listener.accept().await.unwrap(); + let mut server = h2::server::handshake(socket).await.unwrap(); + + // Spawn the request handler so the main loop can keep polling. + let mut recv_body_tx = Some(recv_body_tx); + while let Some(req) = server.accept().await { + let (request, mut respond) = req.unwrap(); + let response_body_vec = response_body_vec.clone(); + let tx = recv_body_tx.take(); + tokio::spawn(async move { + let (_parts, mut body) = request.into_parts(); + let mut got_body = Vec::new(); + while let Some(chunk) = body.data().await { + let chunk = chunk.unwrap(); + body.flow_control().release_capacity(chunk.len()).unwrap(); + got_body.extend_from_slice(&chunk); + } + if let Some(tx) = tx { + let _ = tx.send(got_body); + } + + let response = http::Response::builder() + .status(http::StatusCode::OK) + .body(()) + .unwrap(); + let mut send = respond.send_response(response, false).unwrap(); + match ending { + ReceiverEnding::Clean => { + // Send the response body either in one frame or in + // `resp_chunks` smaller ones — the latter exercises + // the bridge's coalescing path on the downstream + // direction. With `resp_chunks == 1` this is a + // single `send_data` call. + send_in_chunks(&mut send, &response_body_vec, resp_chunks); + } + ReceiverEnding::ResetNoError => { + // Body without END_STREAM, then a graceful reset — + // hyper's unary-response close. The bridge must + // still deliver the full body to the sender. + // + // Yield first so the response HEADERS frame is + // flushed and the relay's response future resolves + // `Ok(200)` *before* the body + reset arrive — + // matching the real networked ordering (the live + // run sees a clean 200, then the reset hits the body + // pump). Without the gap the in-process h2 races and + // the reset can pre-empt the headers future, which is + // a harness artifact, not the path under test. + tokio::time::sleep(Duration::from_millis(25)).await; + if !response_body_vec.is_empty() { + send.send_data(Bytes::from(response_body_vec.clone()), false) + .unwrap(); + // Let the DATA frame flush to the relay before + // the reset — hyper resets only *after* the + // complete body is on the wire. Sending the + // reset immediately would drop the unflushed + // DATA frame (an artifact of the synchronous + // in-process send, not the behaviour under test). + tokio::time::sleep(Duration::from_millis(25)).await; + } + send.send_reset(h2::Reason::NO_ERROR); + } + } + }); + } + }); + + // Cluster side: dial the fake receiver as h2 client. + let cluster_recv_sock = TcpStream::connect(receiver_addr).await.unwrap(); + let (send_request, connection) = h2::client::handshake(cluster_recv_sock).await.unwrap(); + let recv_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut send_request = send_request.ready().await.unwrap(); + + let req = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/test") + .body(()) + .unwrap(); + let (receiver_response, receiver_send) = send_request.send_request(req, false).unwrap(); + + // Fake sender + cluster's sender-side server. + let sender_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let sender_addr = sender_listener.local_addr().unwrap(); + let request_body_vec = request_body.to_vec(); + + let sender_task = tokio::spawn(async move { + let socket = TcpStream::connect(sender_addr).await.unwrap(); + let (h2, connection) = h2::client::handshake(socket).await.unwrap(); + let sender_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut h2 = h2.ready().await.unwrap(); + + let request = http::Request::builder() + .method(http::Method::POST) + .uri("/c1/invoke/test") + .body(()) + .unwrap(); + let (resp, mut body) = h2.send_request(request, false).unwrap(); + send_in_chunks(&mut body, &request_body_vec, req_chunks); + + let response = resp.await.unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + let (_parts, mut response_body) = response.into_parts(); + let mut got_response_body = Vec::new(); + while let Some(chunk) = response_body.data().await { + let chunk = chunk.unwrap(); + response_body + .flow_control() + .release_capacity(chunk.len()) + .unwrap(); + got_response_body.extend_from_slice(&chunk); + } + sender_conn_task.abort(); + got_response_body + }); + + // Cluster's sender-facing h2 server: accept the sender's connection. + let (cluster_sender_sock, _) = sender_listener.accept().await.unwrap(); + let mut sender_server = h2::server::handshake(cluster_sender_sock).await.unwrap(); + let (sender_request, mut sender_respond) = sender_server.accept().await.unwrap().unwrap(); + let (_sender_parts, sender_body) = sender_request.into_parts(); + + // Run the bridge. The sender-server must continue to be polled + // to drive the h2 connection on the sender side; spawn a tiny task + // that keeps draining accept() (it'll return None once the sender + // closes). + let sender_server_drain = + tokio::spawn(async move { while sender_server.accept().await.is_some() {} }); + + let config = BridgeConfig::default(); + let up = bridge_upstream(sender_body, receiver_send, &config); + let down = bridge_downstream(receiver_response, &mut sender_respond, &config); + let (up_res, down_res) = tokio::join!(up, down); + down_res.unwrap(); + match ending { + ReceiverEnding::Clean => up_res.unwrap(), + // Under a graceful full-stream reset, the request pump may race the + // receiver's reset and see a benign `Receiver(Reset NO_ERROR)` write + // error. In production `forward.rs` cancels the request pump the + // moment the (unary) response completes, so this never surfaces; the + // load-bearing contract here is only that DOWNSTREAM delivered the + // body cleanly (asserted above + by the captured response body). + ReceiverEnding::ResetNoError => { + let _ = up_res; + } + } + + let received_response_body = tokio::time::timeout(Duration::from_secs(5), sender_task) + .await + .unwrap() + .unwrap(); + let received_request_body = tokio::time::timeout(Duration::from_secs(5), recv_body_rx) + .await + .unwrap() + .unwrap(); + + // Clean up background tasks. + sender_server_drain.abort(); + recv_conn_task.abort(); + recv_task.abort(); + + (received_request_body, received_response_body) + } + + /// Send `body` over `stream` in `chunks` near-equal `send_data` + /// calls; only the last carries `end_of_stream=true`. With + /// `chunks == 1` this is the original one-shot send. + fn send_in_chunks(stream: &mut h2::SendStream, body: &[u8], chunks: usize) { + if body.is_empty() { + stream.send_data(Bytes::new(), true).unwrap(); + return; + } + let chunks = chunks.max(1); + let chunk_size = body.len().div_ceil(chunks); + let mut offset = 0; + while offset < body.len() { + let end = (offset + chunk_size).min(body.len()); + let eos = end == body.len(); + stream + .send_data(Bytes::copy_from_slice(&body[offset..end]), eos) + .unwrap(); + offset = end; + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn round_trip_small_body() { + let (req_seen_by_receiver, resp_seen_by_sender) = + run_bridge_test(b"hello world", b"goodbye world").await; + assert_eq!(req_seen_by_receiver, b"hello world"); + assert_eq!(resp_seen_by_sender, b"goodbye world"); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn round_trip_empty_bodies() { + let (req_seen_by_receiver, resp_seen_by_sender) = run_bridge_test(b"", b"").await; + assert_eq!(req_seen_by_receiver, b""); + assert_eq!(resp_seen_by_sender, b""); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn round_trip_medium_body() { + let req: Vec = (0..10_000).map(|i| (i % 256) as u8).collect(); + let resp: Vec = (0..20_000).map(|i| ((i * 7) % 256) as u8).collect(); + let (req_seen, resp_seen) = run_bridge_test(&req, &resp).await; + assert_eq!(req_seen, req); + assert_eq!(resp_seen, resp); + } + + /// A receiver that finishes a unary response with `RST_STREAM(NO_ERROR)` + /// instead of END_STREAM (hyper's idiom — what the unmodified Restate Cloud + /// tunnel client does) must still have its full body delivered to the + /// sender. Regression for the body-truncation found running the real + /// client against the relay: the graceful reset surfaced during the + /// coalesce peek and discarded the already-received body. The harness's + /// `down_res.unwrap()` asserts the downstream bridge returns `Ok`. + #[test] + fn upstream_graceful_peer_close_predicate() { + use PumpDirection::{Downstream, Upstream}; + // Upstream: a receiver CANCEL while we're still writing the request is a + // graceful close (it took what it needed) — request forwarding is done. + assert!(write_hit_graceful_peer_close( + Upstream, + &h2::Error::from(h2::Reason::CANCEL) + )); + // Upstream: a genuine receiver fault is NOT a graceful close. + assert!(!write_hit_graceful_peer_close( + Upstream, + &h2::Error::from(h2::Reason::INTERNAL_ERROR) + )); + assert!(!write_hit_graceful_peer_close( + Upstream, + &h2::Error::from(h2::Reason::PROTOCOL_ERROR) + )); + // Downstream (response → sender) is out of scope here — a sender close + // is classified separately. Even a CANCEL must return false. + assert!(!write_hit_graceful_peer_close( + Downstream, + &h2::Error::from(h2::Reason::CANCEL) + )); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn downstream_no_error_reset_delivers_full_body() { + let (_req_seen, resp_seen) = run_bridge_test_full( + b"req", + b"hello-from-target\n", + 1, + 1, + ReceiverEnding::ResetNoError, + ) + .await; + assert_eq!(resp_seen, b"hello-from-target\n"); + } + + /// Same graceful-reset close, but with an empty response body: the reset + /// surfaces on the very first `data().await` (no buffered DATA), exercising + /// the main-pump-loop arm rather than the coalesce peek. The sender must + /// see a clean, empty `200`. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn downstream_no_error_reset_empty_body_is_clean_eos() { + let (_req_seen, resp_seen) = + run_bridge_test_full(b"req", b"", 1, 1, ReceiverEnding::ResetNoError).await; + assert_eq!(resp_seen, b""); + } + + // ── Peek→bridge accounting handoff (remediation Gate C, Item 4) ─────── + + /// A test [`AccountingHandoff`] guard: on Drop it increments a shared + /// counter (release count) and snapshots the meter's value at that instant, + /// so a test can prove the pump published held bytes into the meter BEFORE + /// releasing the guard (publish-then-release ordering) and released it + /// exactly once. + struct ReleaseSpy { + released: Arc, + meter_at_release: Arc, + meter: Arc, + } + impl Drop for ReleaseSpy { + fn drop(&mut self) { + self.meter_at_release + .store(self.meter.load(Ordering::SeqCst), Ordering::SeqCst); + self.released.fetch_add(1, Ordering::SeqCst); + } + } + + /// The upstream pump, given an [`AccountingHandoff`], releases it exactly + /// once while forwarding a normal request body, and the shared buffer meter + /// returns to zero at the end (the handoff never leaks the meter). Proves the + /// peek→bridge handoff runs and releases on the healthy path. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn handoff_guard_released_exactly_once_by_upstream_pump() { + let meter_counter = Arc::new(AtomicI64::new(0)); + let released = Arc::new(AtomicI64::new(0)); + let meter_at_release = Arc::new(AtomicI64::new(-1)); + + // Wire a real sender↔relay↔receiver triangle, running the UPSTREAM pump + // with a handoff guard and the DOWNSTREAM pump normally. + let receiver_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let receiver_addr = receiver_listener.local_addr().unwrap(); + let (recv_body_tx, recv_body_rx) = tokio::sync::oneshot::channel::>(); + let recv_task = tokio::spawn(async move { + let (socket, _) = receiver_listener.accept().await.unwrap(); + let mut server = h2::server::handshake(socket).await.unwrap(); + let mut recv_body_tx = Some(recv_body_tx); + while let Some(req) = server.accept().await { + let (request, mut respond) = req.unwrap(); + let tx = recv_body_tx.take(); + tokio::spawn(async move { + let (_p, mut body) = request.into_parts(); + let mut got = Vec::new(); + while let Some(chunk) = body.data().await { + let chunk = chunk.unwrap(); + body.flow_control().release_capacity(chunk.len()).unwrap(); + got.extend_from_slice(&chunk); + } + if let Some(tx) = tx { + let _ = tx.send(got); + } + let response = http::Response::builder() + .status(http::StatusCode::OK) + .body(()) + .unwrap(); + let mut send = respond.send_response(response, false).unwrap(); + send.send_data(Bytes::new(), true).unwrap(); + }); + } + }); + + let cluster_recv_sock = TcpStream::connect(receiver_addr).await.unwrap(); + let (send_request, connection) = h2::client::handshake(cluster_recv_sock).await.unwrap(); + let recv_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut send_request = send_request.ready().await.unwrap(); + let req = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/test") + .body(()) + .unwrap(); + let (receiver_response, receiver_send) = send_request.send_request(req, false).unwrap(); + + let sender_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let sender_addr = sender_listener.local_addr().unwrap(); + let body_bytes: Vec = (0..40_000).map(|i| (i % 251) as u8).collect(); + let body_for_sender = body_bytes.clone(); + let sender_task = tokio::spawn(async move { + let socket = TcpStream::connect(sender_addr).await.unwrap(); + let (h2, connection) = h2::client::handshake(socket).await.unwrap(); + let sender_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut h2 = h2.ready().await.unwrap(); + let request = http::Request::builder() + .method(http::Method::POST) + .uri("/c1/invoke/test") + .body(()) + .unwrap(); + let (resp, mut body) = h2.send_request(request, false).unwrap(); + body.send_data(Bytes::from(body_for_sender), true).unwrap(); + let response = resp.await.unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + let (_p, mut rb) = response.into_parts(); + while let Some(chunk) = rb.data().await { + let chunk = chunk.unwrap(); + rb.flow_control().release_capacity(chunk.len()).unwrap(); + } + sender_conn_task.abort(); + }); + + let (cluster_sender_sock, _) = sender_listener.accept().await.unwrap(); + let mut sender_server = h2::server::handshake(cluster_sender_sock).await.unwrap(); + let (sender_request, mut sender_respond) = sender_server.accept().await.unwrap().unwrap(); + let (_p, sender_body) = sender_request.into_parts(); + let sender_server_drain = + tokio::spawn(async move { while sender_server.accept().await.is_some() {} }); + + let config = BridgeConfig { + buffer_meter: BufferMeter::new(meter_counter.clone()), + ..BridgeConfig::default() + }; + let spy = ReleaseSpy { + released: released.clone(), + meter_at_release: meter_at_release.clone(), + meter: meter_counter.clone(), + }; + let up = bridge_upstream_with_handoff( + sender_body, + receiver_send, + &config, + Some(Box::new(spy) as AccountingHandoff), + ); + let down = bridge_downstream(receiver_response, &mut sender_respond, &config); + let (up_res, down_res) = tokio::join!(up, down); + up_res.unwrap(); + down_res.unwrap(); + + let got_req = tokio::time::timeout(Duration::from_secs(5), recv_body_rx) + .await + .unwrap() + .unwrap(); + tokio::time::timeout(Duration::from_secs(5), sender_task) + .await + .unwrap() + .unwrap(); + + // The guard was released exactly once by the pump… + assert_eq!( + released.load(Ordering::SeqCst), + 1, + "the handoff guard must be released exactly once" + ); + // …the release read a non-negative meter (publish-then-release: the meter + // is reconciled at the handoff before the guard drops — never negative)… + assert!( + meter_at_release.load(Ordering::SeqCst) >= 0, + "meter must be reconciled (>=0) at the instant the guard is released" + ); + // …the full body round-tripped… + assert_eq!( + got_req, body_bytes, + "the request body must round-trip intact" + ); + // …and the shared meter telescopes back to zero (no leak from handoff). + assert_eq!( + meter_counter.load(Ordering::SeqCst), + 0, + "the buffer meter must return to zero once the bridge completes" + ); + + sender_server_drain.abort(); + recv_conn_task.abort(); + recv_task.abort(); + } + + /// If the upstream-pump future is dropped **before its first poll** (the + /// cancellation-before-bridge case — e.g. `run_duplex` dropping it because + /// the other half errored first), the handoff guard is still released + /// exactly once. Proves the RAII fallback: no leak when the pump never runs. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn handoff_guard_released_when_pump_future_dropped_unpolled() { + // Build a real upstream stream pair, then drop the pump future without + // ever awaiting it. + let receiver_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let receiver_addr = receiver_listener.local_addr().unwrap(); + let recv_task = tokio::spawn(async move { + let (socket, _) = receiver_listener.accept().await.unwrap(); + let mut server = h2::server::handshake(socket).await.unwrap(); + while server.accept().await.is_some() {} + }); + let cluster_recv_sock = TcpStream::connect(receiver_addr).await.unwrap(); + let (send_request, connection) = h2::client::handshake(cluster_recv_sock).await.unwrap(); + let recv_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut send_request = send_request.ready().await.unwrap(); + let req = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/test") + .body(()) + .unwrap(); + let (_receiver_response, receiver_send) = send_request.send_request(req, false).unwrap(); + + let sender_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let sender_addr = sender_listener.local_addr().unwrap(); + let sender_task = tokio::spawn(async move { + let socket = TcpStream::connect(sender_addr).await.unwrap(); + let (h2, connection) = h2::client::handshake(socket).await.unwrap(); + let drive = tokio::spawn(async move { + let _ = connection.await; + }); + let mut h2 = h2.ready().await.unwrap(); + let request = http::Request::builder() + .method(http::Method::POST) + .uri("/c1/invoke/test") + .body(()) + .unwrap(); + let (_resp, mut body) = h2.send_request(request, false).unwrap(); + let _ = body.send_data(Bytes::new(), true); + // Keep the connection alive briefly. + tokio::time::sleep(Duration::from_millis(50)).await; + drive.abort(); + }); + let (cluster_sender_sock, _) = sender_listener.accept().await.unwrap(); + let mut sender_server = h2::server::handshake(cluster_sender_sock).await.unwrap(); + let (sender_request, _sender_respond) = sender_server.accept().await.unwrap().unwrap(); + let (_p, sender_body) = sender_request.into_parts(); + + let released = Arc::new(AtomicI64::new(0)); + let spy = ReleaseSpy { + released: released.clone(), + meter_at_release: Arc::new(AtomicI64::new(0)), + meter: Arc::new(AtomicI64::new(0)), + }; + let config = BridgeConfig::default(); + let up = bridge_upstream_with_handoff( + sender_body, + receiver_send, + &config, + Some(Box::new(spy) as AccountingHandoff), + ); + // Drop the future WITHOUT polling it — the guard it captured must drop. + drop(up); + assert_eq!( + released.load(Ordering::SeqCst), + 1, + "dropping the un-polled pump future must release the handoff guard exactly once" + ); + + recv_conn_task.abort(); + recv_task.abort(); + let _ = sender_task.await; + } + + /// Coalescing correctness: sender drips the request body in many + /// small DATA frames; bridge must deliver the full payload + /// byte-identical to the receiver. Each chunk is sent in its own + /// `send_data` call which lands as a separate DATA frame on the + /// wire; the bridge's coalescing peek then concatenates whatever + /// is already buffered into a single forwarded `send_data` call. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn coalesces_dripped_request_body_lossless() { + // 32 chunks × 128 bytes = 4 KiB body. Distinct byte values + // per chunk so a re-ordering bug would show up in the diff. + let req: Vec = (0..4096).map(|i| (i % 251) as u8).collect(); + let (req_seen, resp_seen) = run_bridge_test_chunked( + &req, b"ok", /*req_chunks=*/ 32, /*resp_chunks=*/ 1, + ) + .await; + assert_eq!( + req_seen, req, + "dripped request body must round-trip byte-exact" + ); + assert_eq!(resp_seen, b"ok"); + } + + /// Same but on the response side: receiver drips, bridge coalesces + /// downstream. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn coalesces_dripped_response_body_lossless() { + let resp: Vec = (0..8192).map(|i| ((i * 11) % 256) as u8).collect(); + let (req_seen, resp_seen) = run_bridge_test_chunked( + b"ping", &resp, /*req_chunks=*/ 1, /*resp_chunks=*/ 64, + ) + .await; + assert_eq!(req_seen, b"ping"); + assert_eq!( + resp_seen, resp, + "dripped response body must round-trip byte-exact" + ); + } + + /// Bidirectional drip — both ends fragment their bodies. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn coalesces_both_directions_lossless() { + let req: Vec = (0..16_384).map(|i| (i % 256) as u8).collect(); + let resp: Vec = (0..16_384).map(|i| ((i * 13) % 256) as u8).collect(); + let (req_seen, resp_seen) = run_bridge_test_chunked(&req, &resp, 16, 32).await; + assert_eq!(req_seen, req); + assert_eq!(resp_seen, resp); + } + + /// Boundary: a body larger than the COALESCE_CAP_BYTES (64 KiB). + /// Coalescing must stop at the cap and the remainder ride the + /// next outer-loop iteration; no bytes are lost or duplicated. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn coalesces_at_cap_lossless() { + // 128 KiB total, dripped in 256 chunks of 512 bytes each. + // After coalescing the first ~128 chunks (~64 KiB) we hit the + // cap, send, then loop and pick up the rest. + let req: Vec = (0..131_072).map(|i| ((i * 17) % 256) as u8).collect(); + let (req_seen, resp_seen) = run_bridge_test_chunked(&req, b"k", 256, 1).await; + assert_eq!(req_seen, req); + assert_eq!(resp_seen, b"k"); + } + + /// Trailers must round-trip verbatim in BOTH directions (design + /// §4.2(e)/§4.3). The sender attaches a request trailer after its + /// body; the receiver attaches a response trailer after its body. + /// Each side must observe the other's trailer. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn forwards_trailers_both_directions() { + use http::HeaderMap; + + let receiver_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let receiver_addr = receiver_listener.local_addr().unwrap(); + + // Receiver: drain the request body, capture its trailers, then + // respond with a body + a response trailer. + let (req_tr_tx, req_tr_rx) = tokio::sync::oneshot::channel::>(); + let recv_task = tokio::spawn(async move { + let (socket, _) = receiver_listener.accept().await.unwrap(); + let mut server = h2::server::handshake(socket).await.unwrap(); + let mut req_tr_tx = Some(req_tr_tx); + while let Some(req) = server.accept().await { + let (request, mut respond) = req.unwrap(); + let tx = req_tr_tx.take(); + tokio::spawn(async move { + let (_p, mut body) = request.into_parts(); + while let Some(c) = body.data().await { + let c = c.unwrap(); + body.flow_control().release_capacity(c.len()).unwrap(); + } + let trailers = body.trailers().await.unwrap(); + if let Some(tx) = tx { + let _ = tx.send(trailers); + } + let response = http::Response::builder() + .status(http::StatusCode::OK) + .body(()) + .unwrap(); + let mut send = respond.send_response(response, false).unwrap(); + send.send_data(Bytes::from_static(b"resp-body"), false) + .unwrap(); + let mut tr = HeaderMap::new(); + tr.insert("x-resp-trailer", "resp-val".parse().unwrap()); + send.send_trailers(tr).unwrap(); + }); + } + }); + + let cluster_recv_sock = TcpStream::connect(receiver_addr).await.unwrap(); + let (send_request, connection) = h2::client::handshake(cluster_recv_sock).await.unwrap(); + let recv_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut send_request = send_request.ready().await.unwrap(); + let req = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/h/9080/test") + .body(()) + .unwrap(); + let (receiver_response, receiver_send) = send_request.send_request(req, false).unwrap(); + + let sender_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let sender_addr = sender_listener.local_addr().unwrap(); + + // Sender: send a body + a request trailer; read the response body + // and capture the response trailer. + let (resp_tr_tx, resp_tr_rx) = tokio::sync::oneshot::channel::>(); + let sender_task = tokio::spawn(async move { + let socket = TcpStream::connect(sender_addr).await.unwrap(); + let (h2, connection) = h2::client::handshake(socket).await.unwrap(); + let sender_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut h2 = h2.ready().await.unwrap(); + let request = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/c/h/9080/test") + .body(()) + .unwrap(); + let (resp, mut body) = h2.send_request(request, false).unwrap(); + body.send_data(Bytes::from_static(b"req-body"), false) + .unwrap(); + let mut tr = HeaderMap::new(); + tr.insert("x-req-trailer", "req-val".parse().unwrap()); + body.send_trailers(tr).unwrap(); + + let response = resp.await.unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + let (_p, mut rb) = response.into_parts(); + while let Some(c) = rb.data().await { + let c = c.unwrap(); + rb.flow_control().release_capacity(c.len()).unwrap(); + } + let resp_trailers = rb.trailers().await.unwrap(); + let _ = resp_tr_tx.send(resp_trailers); + sender_conn_task.abort(); + }); + + let (cluster_sender_sock, _) = sender_listener.accept().await.unwrap(); + let mut sender_server = h2::server::handshake(cluster_sender_sock).await.unwrap(); + let (sender_request, mut sender_respond) = sender_server.accept().await.unwrap().unwrap(); + let (_p, sender_body) = sender_request.into_parts(); + let sender_server_drain = + tokio::spawn(async move { while sender_server.accept().await.is_some() {} }); + + let config = BridgeConfig::default(); + let up = bridge_upstream(sender_body, receiver_send, &config); + let down = bridge_downstream(receiver_response, &mut sender_respond, &config); + let (up_res, down_res) = tokio::join!(up, down); + up_res.unwrap(); + down_res.unwrap(); + + let req_trailers = tokio::time::timeout(Duration::from_secs(5), req_tr_rx) + .await + .unwrap() + .unwrap() + .expect("receiver should have seen request trailers"); + let resp_trailers = tokio::time::timeout(Duration::from_secs(5), resp_tr_rx) + .await + .unwrap() + .unwrap() + .expect("sender should have seen response trailers"); + + sender_server_drain.abort(); + recv_conn_task.abort(); + recv_task.abort(); + let _ = sender_task.await; + + assert_eq!(req_trailers.get("x-req-trailer").unwrap(), "req-val"); + assert_eq!(resp_trailers.get("x-resp-trailer").unwrap(), "resp-val"); + } + + /// A bodyless response carrying only trailers (HEADERS → trailing + /// HEADERS, no DATA) must still deliver the trailers — this exercises + /// `finish_writer` when `data()` returns None immediately. gRPC's + /// trailers-only responses have this shape. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn forwards_trailers_on_bodyless_response() { + use http::HeaderMap; + + let receiver_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let receiver_addr = receiver_listener.local_addr().unwrap(); + let recv_task = tokio::spawn(async move { + let (socket, _) = receiver_listener.accept().await.unwrap(); + let mut server = h2::server::handshake(socket).await.unwrap(); + while let Some(req) = server.accept().await { + let (request, mut respond) = req.unwrap(); + tokio::spawn(async move { + let (_p, mut body) = request.into_parts(); + while let Some(c) = body.data().await { + let c = c.unwrap(); + body.flow_control().release_capacity(c.len()).unwrap(); + } + let _ = body.trailers().await; + // Response head, then trailers — no DATA frame at all. + let response = http::Response::builder() + .status(http::StatusCode::OK) + .body(()) + .unwrap(); + let mut send = respond.send_response(response, false).unwrap(); + let mut tr = HeaderMap::new(); + tr.insert("grpc-status", "0".parse().unwrap()); + send.send_trailers(tr).unwrap(); + }); + } + }); + + let cluster_recv_sock = TcpStream::connect(receiver_addr).await.unwrap(); + let (send_request, connection) = h2::client::handshake(cluster_recv_sock).await.unwrap(); + let recv_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut send_request = send_request.ready().await.unwrap(); + let req = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/h/9080/test") + .body(()) + .unwrap(); + // eos=false: the bridge owns closing this stream (it pumps the + // sender's empty body through and emits the end-of-stream frame). + let (receiver_response, receiver_send) = send_request.send_request(req, false).unwrap(); + + let sender_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let sender_addr = sender_listener.local_addr().unwrap(); + let (resp_tr_tx, resp_tr_rx) = tokio::sync::oneshot::channel::>(); + let sender_task = tokio::spawn(async move { + let socket = TcpStream::connect(sender_addr).await.unwrap(); + let (h2, connection) = h2::client::handshake(socket).await.unwrap(); + let sender_conn_task = tokio::spawn(async move { + let _ = connection.await; + }); + let mut h2 = h2.ready().await.unwrap(); + let request = http::Request::builder() + .method(http::Method::POST) + .uri("/invoke/c/h/9080/test") + .body(()) + .unwrap(); + let (resp, mut send_body) = h2.send_request(request, false).unwrap(); + send_body.send_data(Bytes::new(), true).unwrap(); + let response = resp.await.unwrap(); + let (_p, mut rb) = response.into_parts(); + while let Some(c) = rb.data().await { + let c = c.unwrap(); + rb.flow_control().release_capacity(c.len()).unwrap(); + } + let resp_trailers = rb.trailers().await.unwrap(); + let _ = resp_tr_tx.send(resp_trailers); + sender_conn_task.abort(); + }); + + let (cluster_sender_sock, _) = sender_listener.accept().await.unwrap(); + let mut sender_server = h2::server::handshake(cluster_sender_sock).await.unwrap(); + let (sender_request, mut sender_respond) = sender_server.accept().await.unwrap().unwrap(); + let (_p, sender_body) = sender_request.into_parts(); + let sender_server_drain = + tokio::spawn(async move { while sender_server.accept().await.is_some() {} }); + + let config = BridgeConfig::default(); + let up = bridge_upstream(sender_body, receiver_send, &config); + let down = bridge_downstream(receiver_response, &mut sender_respond, &config); + let (up_res, down_res) = tokio::join!(up, down); + up_res.unwrap(); + down_res.unwrap(); + + let resp_trailers = tokio::time::timeout(Duration::from_secs(5), resp_tr_rx) + .await + .unwrap() + .unwrap() + .expect("sender should have seen response trailers on a bodyless response"); + + sender_server_drain.abort(); + recv_conn_task.abort(); + recv_task.abort(); + let _ = sender_task.await; + + assert_eq!(resp_trailers.get("grpc-status").unwrap(), "0"); + } + + // --- run_duplex coupling (F2) --- + + /// When one half errors, the other is cancelled immediately — even with + /// no drain timeout and even if the other would otherwise never finish. + /// This is the property a bare `tokio::join!` lacked (it would hang on + /// the pending half forever). + #[tokio::test] + async fn errored_half_cancels_the_other_immediately() { + let cfg = BridgeConfig::default(); // half_close_drain_timeout = None + let up = async { Err::<(), BridgeError>(BridgeError::FlowControlTimeout) }; + let down = std::future::pending::>(); + + let res = tokio::time::timeout(Duration::from_secs(2), run_duplex(up, down, &cfg)).await; + let (up_res, down_res) = res.expect("run_duplex must return promptly, not hang on `down`"); + assert!(matches!(up_res, Err(BridgeError::FlowControlTimeout))); + assert!(matches!(down_res, Err(BridgeError::Cancelled))); + } + + /// Symmetric: a downstream error cancels a never-finishing upstream. + #[tokio::test] + async fn errored_downstream_cancels_upstream_immediately() { + let cfg = BridgeConfig::default(); + let up = std::future::pending::>(); + let down = async { Err::<(), BridgeError>(BridgeError::FlowControlTimeout) }; + + let res = tokio::time::timeout(Duration::from_secs(2), run_duplex(up, down, &cfg)).await; + let (up_res, down_res) = res.expect("run_duplex must return promptly"); + assert!(matches!(up_res, Err(BridgeError::Cancelled))); + assert!(matches!(down_res, Err(BridgeError::FlowControlTimeout))); + } + + /// A clean half-close with the drain timeout set bounds an otherwise + /// idle survivor — the half-open stream can't pin resources forever. + #[tokio::test] + async fn clean_close_bounds_idle_survivor_when_drain_timeout_set() { + let cfg = BridgeConfig { + half_close_drain_timeout: Some(Duration::from_millis(50)), + ..BridgeConfig::default() + }; + let up = async { Ok::<(), BridgeError>(()) }; + let down = std::future::pending::>(); + + let res = tokio::time::timeout(Duration::from_secs(2), run_duplex(up, down, &cfg)).await; + let (up_res, down_res) = res.expect("run_duplex must return within the drain timeout"); + assert!(up_res.is_ok()); + assert!(matches!(down_res, Err(BridgeError::Cancelled))); + } + + /// No-regression: with the default (no drain timeout), a clean half-close + /// does NOT cancel the survivor — it keeps running to completion. This is + /// the legitimate server-streaming / long-upload case. + #[tokio::test] + async fn clean_close_lets_survivor_finish_when_unbounded() { + let cfg = BridgeConfig::default(); // None → unbounded survivor + let up = async { Ok::<(), BridgeError>(()) }; + let down = async { + tokio::time::sleep(Duration::from_millis(50)).await; + Ok::<(), BridgeError>(()) + }; + let (up_res, down_res) = run_duplex(up, down, &cfg).await; + assert!(up_res.is_ok()); + assert!(down_res.is_ok()); + } + + /// Both halves completing cleanly returns `(Ok, Ok)`. + #[tokio::test] + async fn both_clean_returns_ok_ok() { + let cfg = BridgeConfig::default(); + let up = async { Ok::<(), BridgeError>(()) }; + let down = async { Ok::<(), BridgeError>(()) }; + let (up_res, down_res) = run_duplex(up, down, &cfg).await; + assert!(up_res.is_ok() && down_res.is_ok()); + } +} diff --git a/src/relay/loopback.rs b/src/relay/loopback.rs new file mode 100644 index 0000000..6581c6f --- /dev/null +++ b/src/relay/loopback.rs @@ -0,0 +1,471 @@ +//! The loopback engine: run the embedded relay [`super::receiver`] on its own +//! tokio runtime and bridge each forwarded stream to a local HTTP/2 (h2c) +//! server over a loopback socket. +//! +//! This is the piece that lets an SDK act as a relay *receiver* without +//! reimplementing the protocol and without letting the relay runtime own the +//! SDK's threads. The FFI/native binding only ever drives the control plane — +//! [`Engine::start`], [`Handle::status_json`], [`Handle::stop`]; every +//! forwarded request rides the loopback socket, never the FFI boundary. +//! +//! # Shape +//! +//! ```text +//! relay :8080 host process +//! ┌────────┐ h2 (role-flip) ┌──────────────────────────────────────────┐ +//! │ relay │◄────────────────►│ Receiver + LoopbackHandler (this crate) │ +//! └────────┘ │ owns its OWN tokio runtime │ +//! │ │ h2c dial per stream │ +//! │ ▼ │ +//! │ HTTP/2 server on 127.0.0.1: │ +//! └──────────────────────────────────────────┘ +//! ``` +//! +//! The `LoopbackHandler` replays the sender's method + **tail-only `:path`** +//! (so the request-identity signature, which covers only the SDK-relative +//! path, stays valid) + headers verbatim, then bridges both bodies with the +//! in-crate [`super::bridge`]. + +use super::bridge::{self, BridgeConfig}; +use super::protocol::{Env, EnvError, H2Tuning, Tunnel, TunnelError}; +use super::receiver::{Invocation, InvokeHandler, Receiver, ReceiverConfig, TlsClient}; +use bytes::Bytes; +use serde::Deserialize; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::net::TcpStream; +use tokio::sync::oneshot; +use tracing::{debug, warn}; + +/// Tunnel configuration, as it crosses a native binding (JSON) or a Rust +/// caller. Deliberately version-tolerant (extra fields ignored, most fields +/// optional) so a binding can add knobs without breaking the wire shape. +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + /// `host:port` of the relay's receiver port (`:8080`) or the LB in front + /// of it. A single name that resolves to many A-records fans out across + /// nodes (multi-homing); see [`connections`](Config::connections). + pub relay_addr: String, + /// The env (routing namespace) this receiver registers under. + pub env: String, + /// The tunnel this receiver registers under (scoped by `env`). + pub tunnel: String, + /// The API key presented in the `/whoami` body. + pub api_key: String, + /// The local HTTP/2 (h2c) server port to bridge forwarded requests to — + /// the SDK's own server, listening on `127.0.0.1:`. + pub local_port: u16, + /// Optional client TLS to the relay (PEM material). Usually absent — + /// receiver TLS is normally LB-terminated so the relay sees plain TCP. + #[serde(default)] + pub tls: Option, + /// R4 multi-homing: how many concurrent connection slots to fan out across + /// the relay nodes `relay_addr` resolves to. `None` = one per resolved node. + #[serde(default)] + pub connections: Option, + /// R5 receiver-instance id, held across reconnects so the broker treats + /// this receiver's connections as one logical receiver for affinity. + /// `None` = auto-generate an ephemeral, process-unique id. + #[serde(default)] + pub instance_id: Option, + /// Delay between a connection ending and the next redial, in milliseconds. + /// `None` = 1000ms. + #[serde(default)] + pub reconnect_backoff_millis: Option, + /// tokio worker-thread count for the embedded runtime. Bridging is + /// I/O-bound, so a small pool is plenty and avoids spawning one thread per + /// core inside the host process. `None` = 2. + #[serde(default)] + pub worker_threads: Option, +} + +/// Client-TLS material for the receiver→relay hop (PEM). See +/// [`TlsClient::from_pem`]. +#[derive(Debug, Clone, Deserialize)] +pub struct TlsConfig { + /// CA cert(s) (PEM) verifying the relay's server cert. + pub ca_pem: String, + /// DNS name to verify on the relay's cert (SNI). + pub server_name: String, + /// Optional client-cert PEM for mTLS (rare on `:8080`). + #[serde(default)] + pub client_cert_pem: Option, + /// Optional client-key PEM for mTLS. Required iff `client_cert_pem` is set. + #[serde(default)] + pub client_key_pem: Option, +} + +impl TlsConfig { + fn build(&self) -> Result { + let identity = match (&self.client_cert_pem, &self.client_key_pem) { + (Some(cert), Some(key)) => Some((cert.as_bytes(), key.as_bytes())), + (None, None) => None, + _ => { + return Err(ConfigError::Tls( + "client_cert_pem and client_key_pem must be set together".into(), + )) + } + }; + TlsClient::from_pem(self.ca_pem.as_bytes(), identity, &self.server_name) + .map_err(|e| ConfigError::Tls(e.to_string())) + } +} + +/// Error building a running [`Handle`] from a [`Config`]. +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("invalid env: {0}")] + Env(#[from] EnvError), + #[error("invalid tunnel: {0}")] + Tunnel(#[from] TunnelError), + #[error("tls: {0}")] + Tls(String), + #[error("failed to build tokio runtime: {0}")] + Runtime(String), +} + +impl Config { + /// Turn the wire config into a [`ReceiverConfig`], validating the env, + /// tunnel, and TLS material. + fn into_receiver_config(self) -> Result { + let env = Env::new(self.env)?; + let tunnel = Tunnel::new(self.tunnel)?; + let tls = match &self.tls { + Some(t) => Some(t.build()?), + None => None, + }; + let mut rc = ReceiverConfig::new(self.relay_addr, env, tunnel, self.api_key); + rc.tls = tls; + rc.connections = self.connections; + if let Some(id) = self.instance_id { + // An explicit id overrides the auto-generated one; an empty string + // is treated as "opt out of instance affinity". + rc.instance_id = if id.is_empty() { None } else { Some(id) }; + } + if let Some(ms) = self.reconnect_backoff_millis { + rc.reconnect_backoff = Duration::from_millis(ms); + } + Ok(rc) + } +} + +/// Live status of a running tunnel, readable at any time via +/// [`Handle::status`] / [`Handle::status_json`]. +#[derive(Debug)] +pub struct Status { + running: AtomicBool, + last_error: Mutex>, +} + +impl Status { + fn new() -> Self { + Self { + running: AtomicBool::new(false), + last_error: Mutex::new(None), + } + } + + fn set_running(&self, v: bool) { + self.running.store(v, Ordering::SeqCst); + } + + /// Whether the receiver task is currently running. + pub fn is_running(&self) -> bool { + self.running.load(Ordering::SeqCst) + } + + /// The last error observed by the engine, if any. + pub fn last_error(&self) -> Option { + self.last_error + .lock() + .expect("status mutex poisoned") + .clone() + } + + /// Serialise to the JSON shape the native bindings return: + /// `{"running": bool, "last_error": string|null}`. + pub fn to_json(&self) -> String { + serde_json::json!({ + "running": self.is_running(), + "last_error": self.last_error(), + }) + .to_string() + } +} + +/// The tunnel engine. Zero-sized; the running state lives in the [`Handle`] +/// that [`start`](Engine::start) returns. +pub struct Engine; + +impl Engine { + /// Build a tokio runtime, spawn the embedded [`Receiver`] driving a + /// `LoopbackHandler`, and return a live [`Handle`]. Returns immediately; + /// the receiver dials the relay in the background. + pub fn start(config: Config) -> Result { + let local_port = config.local_port; + let worker_threads = config.worker_threads.unwrap_or(2).max(1); + let receiver_config = config.into_receiver_config()?; + let h2_tuning = receiver_config.h2_tuning.clone(); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .enable_all() + .thread_name("relay-tunnel") + .build() + .map_err(|e| ConfigError::Runtime(e.to_string()))?; + + let status = Arc::new(Status::new()); + let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>(); + + let handler = Arc::new(LoopbackHandler { + local_port, + h2_tuning, + }); + let status_for_task = status.clone(); + runtime.spawn(async move { + Receiver::new(receiver_config, handler) + .run(async move { + let _ = shutdown_rx.await; + }) + .await; + status_for_task.set_running(false); + }); + status.set_running(true); + + Ok(Handle { + runtime: Some(runtime), + shutdown: Some(shutdown_tx), + status, + }) + } +} + +/// A running tunnel. Dropping it (or calling [`stop`](Handle::stop)) signals a +/// graceful shutdown and joins the runtime. +pub struct Handle { + runtime: Option, + shutdown: Option>, + status: Arc, +} + +impl Handle { + /// Read the shared [`Status`]. + pub fn status(&self) -> &Arc { + &self.status + } + + /// Read the status as the JSON the native bindings expose. + pub fn status_json(&self) -> String { + self.status.to_json() + } + + /// Signal a graceful shutdown and join the runtime (bounded). Idempotent. + /// + /// Must be called from a non-async context (it blocks while the runtime + /// drains) — the native bindings call it from a host control thread. It is + /// never called from within the tunnel's own runtime. + pub fn stop(&mut self) { + if let Some(tx) = self.shutdown.take() { + let _ = tx.send(()); + } + if let Some(rt) = self.runtime.take() { + rt.shutdown_timeout(Duration::from_secs(5)); + } + self.status.set_running(false); + } +} + +impl Drop for Handle { + fn drop(&mut self) { + self.stop(); + } +} + +/// The [`InvokeHandler`] that bridges each forwarded relay stream to the local +/// HTTP/2 server. +struct LoopbackHandler { + local_port: u16, + h2_tuning: H2Tuning, +} + +#[async_trait::async_trait] +impl InvokeHandler for LoopbackHandler { + async fn handle(&self, invocation: Invocation) { + self.serve(invocation).await; + } +} + +impl LoopbackHandler { + async fn serve(&self, invocation: Invocation) { + let Invocation { + target, + method, + headers, + body, + mut respond, + } = invocation; + + // Dial the local h2c server. A failure here is ours to surface: reply + // 502 so the sender sees a clean status rather than a bare stream reset. + let tcp = match TcpStream::connect(("127.0.0.1", self.local_port)).await { + Ok(s) => s, + Err(e) => { + warn!( + port = self.local_port, + ?e, + "loopback: connect to local server failed" + ); + reply_status(&mut respond, http::StatusCode::BAD_GATEWAY); + return; + } + }; + let _ = tcp.set_nodelay(true); + + let mut builder = h2::client::Builder::new(); + self.h2_tuning.apply_client(&mut builder); + let (h2, connection) = match builder.handshake::<_, Bytes>(tcp).await { + Ok(pair) => pair, + Err(e) => { + warn!(?e, "loopback: h2 handshake to local server failed"); + reply_status(&mut respond, http::StatusCode::BAD_GATEWAY); + return; + } + }; + // Drive the local connection in the background. + tokio::spawn(async move { + if let Err(e) = connection.await { + debug!(?e, "loopback: local h2 connection ended"); + } + }); + + // Replay method + tail-only `:path` (preserves the request-identity + // signature) + headers verbatim onto an absolute URI so h2 derives the + // `:scheme` / `:authority` / `:path` pseudo-headers. + let uri = format!("http://127.0.0.1:{}{}", self.local_port, target.tail); + let mut req_builder = http::Request::builder().method(method).uri(uri); + if let Some(hs) = req_builder.headers_mut() { + *hs = headers; + } + let request = match req_builder.body(()) { + Ok(r) => r, + Err(e) => { + warn!(?e, "loopback: building forwarded request failed"); + reply_status(&mut respond, http::StatusCode::BAD_GATEWAY); + return; + } + }; + + let mut h2 = match h2.ready().await { + Ok(h) => h, + Err(e) => { + warn!(?e, "loopback: local h2 client not ready"); + reply_status(&mut respond, http::StatusCode::BAD_GATEWAY); + return; + } + }; + // `end_of_stream = false`: the request body (if any) follows via the + // send stream; the upstream pump closes it (END_STREAM / trailers). + let (response_fut, send_stream) = match h2.send_request(request, false) { + Ok(pair) => pair, + Err(e) => { + warn!(?e, "loopback: send_request to local server failed"); + reply_status(&mut respond, http::StatusCode::BAD_GATEWAY); + return; + } + }; + + // Bridge both directions until they terminate. The relay-facing stream + // plays the bridge's "sender side" (we are the h2 server on the + // reversed connection); the local dial plays the "receiver side". + let cfg = BridgeConfig::default(); + let up = bridge::bridge_upstream(body, send_stream, &cfg); + let down = bridge::bridge_downstream(response_fut, &mut respond, &cfg); + let (up_res, down_res) = bridge::run_duplex(up, down, &cfg).await; + if let Err(e) = up_res { + debug!(?e, "loopback: upstream bridge ended with error"); + } + if let Err(e) = down_res { + debug!(?e, "loopback: downstream bridge ended with error"); + } + } +} + +/// Send an empty-body response with the given status back to the relay. +fn reply_status(respond: &mut h2::server::SendResponse, status: http::StatusCode) { + let response = http::Response::builder() + .status(status) + .body(()) + .expect("status response builds"); + if let Ok(mut send) = respond.send_response(response, false) { + let _ = send.send_data(Bytes::new(), true); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn config_parses_minimal_json() { + let json = r#"{ + "relay_addr": "127.0.0.1:8080", + "env": "e1", + "tunnel": "c1", + "api_key": "k", + "local_port": 9080 + }"#; + let cfg: Config = serde_json::from_str(json).expect("parses"); + assert_eq!(cfg.local_port, 9080); + assert!(cfg.connections.is_none()); + let rc = cfg.into_receiver_config().expect("builds"); + assert_eq!(rc.env.as_str(), "e1"); + assert_eq!(rc.tunnel.as_str(), "c1"); + // An instance_id is auto-generated when not supplied. + assert!(rc.instance_id.is_some()); + } + + #[test] + fn config_rejects_bad_env() { + let json = r#"{"relay_addr":"h:8080","env":"has-dash","tunnel":"c1","api_key":"k","local_port":1}"#; + let cfg: Config = serde_json::from_str(json).unwrap(); + assert!(matches!( + cfg.into_receiver_config(), + Err(ConfigError::Env(_)) + )); + } + + #[test] + fn config_honours_explicit_and_empty_instance_id() { + let with_id: Config = serde_json::from_str( + r#"{"relay_addr":"h:8080","env":"e1","tunnel":"c1","api_key":"k","local_port":1,"instance_id":"fixed"}"#, + ) + .unwrap(); + assert_eq!( + with_id + .into_receiver_config() + .unwrap() + .instance_id + .as_deref(), + Some("fixed") + ); + + let empty_id: Config = serde_json::from_str( + r#"{"relay_addr":"h:8080","env":"e1","tunnel":"c1","api_key":"k","local_port":1,"instance_id":""}"#, + ) + .unwrap(); + assert_eq!(empty_id.into_receiver_config().unwrap().instance_id, None); + } + + #[test] + fn status_json_shape() { + let s = Status::new(); + let v: serde_json::Value = serde_json::from_str(&s.to_json()).unwrap(); + assert_eq!(v["running"], false); + assert!(v["last_error"].is_null()); + + s.set_running(true); + let v: serde_json::Value = serde_json::from_str(&s.to_json()).unwrap(); + assert_eq!(v["running"], true); + } +} diff --git a/src/relay/mod.rs b/src/relay/mod.rs new file mode 100644 index 0000000..ac59e51 --- /dev/null +++ b/src/relay/mod.rs @@ -0,0 +1,32 @@ +//! Embedded relay-tunnel receiver — **compiled only under the `tunnel` +//! feature** (see `docs/relay-receiver-in-shared-core.md`). +//! +//! With the feature off (the default), none of this compiles and the crate +//! stays strictly sans-IO / WASM-compilable — no tokio, no networking in the +//! dependency graph. With it on, an SDK can act as a relay *receiver*: register +//! with the relay, accept the HTTP/2 role-flip, and bridge each forwarded +//! request to the SDK's own local HTTP/2 server over a loopback socket. +//! +//! The stack is a fork of the relay repo's `relay-protocol` / `relay-bridge` / +//! `relay-receiver` crates plus the loopback engine that ties them together: +//! +//! * [`protocol`] — the runtime-agnostic wire contract (env/tunnel, path +//! grammar, `/whoami`, h2 tuning). Sans-IO. +//! * [`bridge`] — the bidirectional HTTP/2 stream bridge. +//! * [`receiver`] — the relay-facing receiver runtime (dial-out, role-flip, +//! `/whoami`, liveness, redial, multi-homing) + the [`receiver::InvokeHandler`] +//! seam. +//! * [`loopback`] — the engine: own a tokio runtime, run the receiver, bridge +//! each forwarded stream to `127.0.0.1:`. This is the API a +//! native binding drives. + +// `bridge` is a faithful fork of the relay repo's `relay-bridge`; its public +// docs reference the private `BufferMeterGuard` (harmless internal cross-refs). +// Allow it here rather than diverging the forked file. +#[allow(rustdoc::private_intra_doc_links)] +pub mod bridge; +pub mod loopback; +pub mod protocol; +pub mod receiver; + +pub use loopback::{Config, ConfigError, Engine, Handle, Status, TlsConfig}; diff --git a/src/relay/protocol/env.rs b/src/relay/protocol/env.rs new file mode 100644 index 0000000..7bed5b4 --- /dev/null +++ b/src/relay/protocol/env.rs @@ -0,0 +1,120 @@ +//! The `Env` newtype: an opaque, alphanumeric environment id. +//! +//! An `Env` scopes tunnels: a tunnel is only meaningful *within* an env, so +//! the rendezvous routing identity is the `(env, tunnel)` pair, not the +//! tunnel alone. Two receivers under the same tunnel string but different +//! envs are distinct tunnels. Like [`super::tunnel::Tunnel`], an env is +//! validated at construction and otherwise treated as an opaque label. + +use super::error::EnvError; +use std::fmt; +use std::str::FromStr; + +/// Maximum accepted env length in bytes. An env is an externally-controlled +/// identifier (sender path segment + receiver `/whoami` field), so it is +/// length-bounded like every other ingress input. Restate Cloud env ids are +/// `env_` → ~28 chars after the `env_` strip, so 64 is generous. +pub const MAX_ENV_LEN: usize = 64; + +/// A validated environment id. Always non-empty, alphanumeric +/// (`[A-Za-z0-9]+`) — the same charset as [`super::tunnel::Tunnel`] — and at +/// most [`MAX_ENV_LEN`] bytes. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Env(String); + +impl Env { + /// Construct an `Env`, validating the input. + pub fn new(s: impl Into) -> Result { + let s = s.into(); + if s.is_empty() { + return Err(EnvError::Empty); + } + if s.len() > MAX_ENV_LEN { + return Err(EnvError::TooLong(MAX_ENV_LEN)); + } + if !s.chars().all(|c| c.is_ascii_alphanumeric()) { + return Err(EnvError::NonAlphanumeric); + } + Ok(Self(s)) + } + + /// The env as a string slice. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Display for Env { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl FromStr for Env { + type Err = EnvError; + + fn from_str(s: &str) -> Result { + Self::new(s) + } +} + +impl AsRef for Env { + fn as_ref(&self) -> &str { + &self.0 + } +} + +// Serialises as a plain string so the `/whoami` wire body carries +// `"env": "..."`. No `Deserialize`: inbound envs go through `Env::new`. +impl serde::Serialize for Env { + fn serialize(&self, s: S) -> Result { + s.serialize_str(&self.0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn accepts_alphanumeric() { + assert!(Env::new("201kqw76hvvdkhsvhdajh77zxek").is_ok()); + assert!(Env::new("E1").is_ok()); + assert!(Env::new("0").is_ok()); + } + + #[test] + fn rejects_empty() { + assert_eq!(Env::new(""), Err(EnvError::Empty)); + } + + #[test] + fn rejects_too_long() { + // At the cap is fine; one over is rejected before the charset check. + assert!(Env::new("a".repeat(MAX_ENV_LEN)).is_ok()); + assert_eq!( + Env::new("a".repeat(MAX_ENV_LEN + 1)), + Err(EnvError::TooLong(MAX_ENV_LEN)) + ); + } + + #[test] + fn rejects_special_characters() { + for bad in ["e-1", "e 1", "e_1", "e/1", "e.1"] { + assert_eq!( + Env::new(bad), + Err(EnvError::NonAlphanumeric), + "reject {bad:?}" + ); + } + } + + #[test] + fn round_trips_via_display_and_from_str() { + let e = Env::new("env123").unwrap(); + assert_eq!(e.to_string(), "env123"); + assert_eq!(e.as_str(), "env123"); + let p: Env = "env123".parse().unwrap(); + assert_eq!(p, e); + } +} diff --git a/src/relay/protocol/error.rs b/src/relay/protocol/error.rs new file mode 100644 index 0000000..c02567e --- /dev/null +++ b/src/relay/protocol/error.rs @@ -0,0 +1,91 @@ +//! Protocol error types shared across the workspace. +//! +//! These were originally in the broker's `error.rs`; the broker re-exports +//! them from here now so there's one definition. Broker-specific errors +//! (e.g. `AuthError`) stay in the broker. + +use thiserror::Error; + +/// Errors that can occur while validating a tunnel string. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum TunnelError { + #[error("tunnel is empty")] + Empty, + #[error("tunnel contains characters outside [A-Za-z0-9._-]")] + InvalidChar, + #[error("tunnel exceeds the maximum length ({0} bytes)")] + TooLong(usize), +} + +/// Errors that can occur while validating an env string. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum EnvError { + #[error("env is empty")] + Empty, + #[error("env contains non-alphanumeric characters")] + NonAlphanumeric, + #[error("env exceeds the maximum length ({0} bytes)")] + TooLong(usize), +} + +/// Errors that can occur during path parsing. +/// +/// Most variants arise from the sender-form parser ([`super::path::parse`], +/// grammar `//////`); `InvalidHost` / +/// `InvalidPort` are specific to the forwarded-form parser +/// ([`super::path::parse_forwarded`], grammar `////`), +/// which the receiver uses and which validates the target the relay treated +/// as opaque. +#[derive(Debug, Error, PartialEq, Eq)] +pub enum PathError { + #[error("path is empty or does not start with '/'")] + Malformed, + #[error("env segment is missing or empty")] + MissingEnv, + #[error("env segment contains non-alphanumeric characters")] + InvalidEnv, + #[error("tunnel segment is missing or empty")] + MissingTunnel, + #[error("tunnel segment contains non-alphanumeric characters")] + InvalidTunnel, + #[error("proto segment is missing or empty")] + MissingProto, + /// The `` segment was present and non-empty but is not the literal + /// `http` (the only implemented transport). Payload-free by design: the + /// rejected value is attacker-controlled, so it is neither retained nor + /// logged — carrying it as a `String` would let a caller stuff an arbitrary + /// (up to `MAX_PATH_LEN`) value into an allocation and any error log. The + /// segment is bounded by the path-length cap regardless; there is no + /// diagnostic value in echoing it back. + #[error("unsupported target protocol (only \"http\" is implemented)")] + UnsupportedProto, + #[error("path is missing the target host and/or port segments")] + MissingTarget, + #[error("target host segment is empty")] + InvalidHost, + #[error("target host segment exceeds the maximum length ({0} bytes)")] + HostTooLong(usize), + #[error("path exceeds the maximum length ({0} bytes)")] + PathTooLong(usize), + #[error("target port segment is not a valid port: {0:?}")] + InvalidPort(String), +} + +/// Errors from parsing a `/whoami` response body. +#[derive(Debug, Error)] +pub enum WhoamiError { + #[error("response body is not valid JSON: {0}")] + InvalidJson(#[from] serde_json::Error), + #[error("response is missing required field: {0}")] + MissingField(&'static str), + #[error("tunnel field is invalid: {0}")] + InvalidTunnel(#[from] TunnelError), + #[error("env field is invalid: {0}")] + InvalidEnv(#[from] EnvError), + #[error("api_key field is empty")] + EmptyApiKey, + #[error("api_key field exceeds the maximum length ({0} bytes)")] + ApiKeyTooLong(usize), + #[error("instance_id field exceeds the maximum length ({0} bytes)")] + InstanceIdTooLong(usize), +} diff --git a/src/relay/protocol/h2_tuning.rs b/src/relay/protocol/h2_tuning.rs new file mode 100644 index 0000000..6625f05 --- /dev/null +++ b/src/relay/protocol/h2_tuning.rs @@ -0,0 +1,194 @@ +//! Shared HTTP/2 protocol-level tuning knobs. +//! +//! Every h2 endpoint in the system accepts the same set of frame-size and +//! flow-control window settings via this struct: the broker's sender-facing +//! server (`:9080`) and receiver-facing client (`:8080` after role-flip), +//! and — once they exist — the receiver's relay-facing server and the +//! sender's client. +//! +//! For large-payload forwarded requests, raising +//! [`H2Tuning::initial_window_size`] reduces `WINDOW_UPDATE` round-trips +//! per stream; raising [`H2Tuning::max_frame_size`] lets each DATA frame +//! carry more bytes. The connection window is the binding constraint at +//! high stream-fan-in (see the field doc). + +use serde::Deserialize; + +/// Default `SETTINGS_MAX_HEADER_LIST_SIZE` (64 KiB) applied to **every** h2 +/// endpoint when [`H2Tuning::max_header_list_size`] is unset. Load-bearing for +/// pre-auth DoS resistance: the h2 crate's own default is **16 MiB**, so a +/// listener that never sets this would accept a 16 MiB header block from an +/// unauthenticated peer. Because [`H2Tuning::apply_server`] / +/// [`apply_client`](H2Tuning::apply_client) fall back to this constant (not to +/// h2's default) when the field is `None`, the cap can never silently vanish — +/// not on a partial `[h2.*]` config section, not on the raw-defaults path. +/// +/// NB: h2 enforces this *softly* up to 4×: a header block between `max` and +/// `4 × max` is a per-stream reject (RST), only `> 4 × max` aborts the whole +/// connection. So the real per-block ceiling is `4 × 64 KiB = 256 KiB`. +pub const DEFAULT_MAX_HEADER_LIST_SIZE: u32 = 65536; + +/// HTTP/2 tuning. Frame-size / window fields are optional; `None` leaves the h2 +/// crate's default in place. `max_header_list_size` is special — it is *always* +/// applied (falling back to [`DEFAULT_MAX_HEADER_LIST_SIZE`] when `None`) so the +/// pre-auth header cap can never silently vanish. Use [`H2Tuning::default`] for +/// the recommended relay tuning (see field docs) or [`H2Tuning::h2_raw_defaults`] +/// to fall back to the h2 crate's wire windows (still with the header cap). +#[derive(Debug, Clone, Deserialize)] +pub struct H2Tuning { + /// Largest frame size we accept (and prefer to send), in bytes. + /// HTTP/2 caps this at `2^24 - 1` (16 MiB - 1). h2 crate default is + /// 16 KiB; relay default is 64 KiB — for large-payload forwarded + /// requests, larger frames mean less framing overhead per byte. + pub max_frame_size: Option, + /// Initial flow-control window per stream, in bytes. h2 crate + /// default is 64 KiB; relay default is 1 MiB. Raising this reduces + /// `WINDOW_UPDATE` chatter on long-lived streams and lets a single + /// stream push more bytes in flight before being window-blocked. + pub initial_window_size: Option, + /// Initial flow-control window for the whole connection, in bytes. + /// h2 crate default is 64 KiB; relay default is 16 MiB. This is the + /// **binding constraint at high stream-fan-in** — every stream on a + /// connection shares it, so with hundreds-to-thousands of streams per + /// connection the per-stream window (above) is unreachable and this + /// window governs aggregate throughput. A high-fan-in bench (2000 + /// streams / 2 connections, 64 KiB frames) measured +28% request-body + /// throughput going 4 MiB → 16 MiB and +79% at 64 MiB; 16 MiB is the + /// balanced default (the window is a *ceiling* on un-acked buffered + /// bytes per connection, not steady-state use, and connection count is + /// O(senders+receivers), not O(streams)). Deployments with extreme + /// fan-in and memory headroom can raise this to 64 MiB+ via `[h2]`. + pub initial_connection_window_size: Option, + /// `SETTINGS_MAX_HEADER_LIST_SIZE` — the largest request/response header + /// block (sum of all header name+value+overhead bytes) this endpoint + /// accepts, in bytes. Unlike the fields above, `None` here does **not** mean + /// "h2 default" — [`apply_server`](Self::apply_server) / + /// [`apply_client`](Self::apply_client) fall back to + /// [`DEFAULT_MAX_HEADER_LIST_SIZE`] (64 KiB) so the pre-auth header cap is + /// always in force (h2's own default is 16 MiB — unsafe for a public + /// listener). Set explicitly to override; validated to `[1 KiB, 16 MiB]`. + pub max_header_list_size: Option, +} + +impl Default for H2Tuning { + /// Relay's recommended tuning, applied by the broker's `AcceptorConfig` + /// and `ServerConfig` defaults (and the receiver/sender): + /// `max_frame_size = 64 KiB`, + /// `initial_window_size = 1 MiB`, + /// `initial_connection_window_size = 16 MiB`. + fn default() -> Self { + Self { + max_frame_size: Some(65536), + initial_window_size: Some(1024 * 1024), + initial_connection_window_size: Some(16 * 1024 * 1024), + max_header_list_size: Some(DEFAULT_MAX_HEADER_LIST_SIZE), + } + } +} + +impl H2Tuning { + /// All-fields-`None`: leave every setting at the h2 crate's wire + /// default. Useful for tests that pin behaviour to the crate's + /// defaults regardless of what we choose as recommended values. + /// Frame-size / window fields all-`None` (h2 crate wire defaults), but the + /// header cap still applies via [`DEFAULT_MAX_HEADER_LIST_SIZE`] — used by + /// the cluster-forward raw path so it keeps h2's small windows yet still + /// bounds the pre-auth header block. Pin-crate-default tests should note the + /// header cap is intentionally present here. + pub fn h2_raw_defaults() -> Self { + Self { + max_frame_size: None, + initial_window_size: None, + initial_connection_window_size: None, + max_header_list_size: None, + } + } + + /// The effective header-list cap: the explicit value or + /// [`DEFAULT_MAX_HEADER_LIST_SIZE`]. Never h2's 16 MiB default. + fn effective_max_header_list_size(&self) -> u32 { + self.max_header_list_size + .unwrap_or(DEFAULT_MAX_HEADER_LIST_SIZE) + } + + /// Validate the tuning knobs. `max_header_list_size`, when set, must be in + /// `[1 KiB, 16 MiB]` (below 1 KiB would reject legitimate requests; above + /// 16 MiB exceeds any sane header block and defeats the DoS bound). + pub fn validate(&self) -> Result<(), &'static str> { + if let Some(v) = self.max_header_list_size { + if !(1024..=16 * 1024 * 1024).contains(&v) { + return Err("max_header_list_size must be between 1024 and 16777216 bytes"); + } + } + Ok(()) + } + + /// Apply our settings to an `h2::server::Builder`. Frame-size / window fields + /// left `None` stay at the builder's default; the header-list cap is + /// **always** applied (explicit value or [`DEFAULT_MAX_HEADER_LIST_SIZE`]). + pub fn apply_server(&self, builder: &mut h2::server::Builder) { + if let Some(v) = self.max_frame_size { + builder.max_frame_size(v); + } + if let Some(v) = self.initial_window_size { + builder.initial_window_size(v); + } + if let Some(v) = self.initial_connection_window_size { + builder.initial_connection_window_size(v); + } + builder.max_header_list_size(self.effective_max_header_list_size()); + } + + /// Apply our settings to an `h2::client::Builder`. Frame-size / window fields + /// left `None` stay at the builder's default; the header-list cap is + /// **always** applied (explicit value or [`DEFAULT_MAX_HEADER_LIST_SIZE`]). + pub fn apply_client(&self, builder: &mut h2::client::Builder) { + if let Some(v) = self.max_frame_size { + builder.max_frame_size(v); + } + if let Some(v) = self.initial_window_size { + builder.initial_window_size(v); + } + if let Some(v) = self.initial_connection_window_size { + builder.initial_connection_window_size(v); + } + builder.max_header_list_size(self.effective_max_header_list_size()); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_carries_the_header_cap() { + assert_eq!( + H2Tuning::default().max_header_list_size, + Some(DEFAULT_MAX_HEADER_LIST_SIZE) + ); + } + + #[test] + fn raw_defaults_still_apply_the_header_cap() { + // The field is None, but the *effective* cap is the const default, not + // h2's 16 MiB — the whole point of the always-apply design. + let raw = H2Tuning::h2_raw_defaults(); + assert_eq!(raw.max_header_list_size, None); + assert_eq!( + raw.effective_max_header_list_size(), + DEFAULT_MAX_HEADER_LIST_SIZE + ); + } + + #[test] + fn validate_rejects_out_of_range_header_cap() { + let with_cap = |v: Option| H2Tuning { + max_header_list_size: v, + ..Default::default() + }; + assert!(with_cap(Some(512)).validate().is_err()); + assert!(with_cap(Some(32 * 1024 * 1024)).validate().is_err()); + assert!(with_cap(Some(65536)).validate().is_ok()); + assert!(with_cap(None).validate().is_ok()); + } +} diff --git a/src/relay/protocol/mod.rs b/src/relay/protocol/mod.rs new file mode 100644 index 0000000..9353a13 --- /dev/null +++ b/src/relay/protocol/mod.rs @@ -0,0 +1,33 @@ +//! `relay-protocol` — the runtime-agnostic wire contract of the M:N +//! HTTP/2 reverse-tunnel rendezvous system. +//! +//! Everything here is pure data + parsing: no tokio, no sockets. It's the +//! single source of truth for the pieces the broker, the receiver, and the +//! sender must all agree on: +//! +//! * [`Tunnel`] — the rendezvous token (`[A-Za-z0-9._-]+`). +//! * [`path`] — the URL grammar, both directions: +//! - [`path::parse`] for the sender form `/invoke////` +//! (used by the broker), +//! - [`path::parse_forwarded`] for the form the broker forwards to the +//! receiver, `/invoke///` (used by the receiver). +//! * [`whoami`] — the `/whoami` JSON body the receiver returns and the +//! broker parses. +//! * [`H2Tuning`] — frame-size / flow-control window knobs applied by +//! every h2 endpoint in the system. +//! +//! See `design.md` §4 for the protocol these types encode. + +pub mod env; +pub mod error; +pub mod h2_tuning; +pub mod path; +pub mod tunnel; +pub mod whoami; + +pub use env::Env; +pub use error::{EnvError, PathError, TunnelError, WhoamiError}; +pub use h2_tuning::H2Tuning; +pub use path::{parse, parse_forwarded, ForwardedTarget, ParsedPath}; +pub use tunnel::Tunnel; +pub use whoami::WhoamiBody; diff --git a/src/relay/protocol/path.rs b/src/relay/protocol/path.rs new file mode 100644 index 0000000..1807653 --- /dev/null +++ b/src/relay/protocol/path.rs @@ -0,0 +1,601 @@ +//! The URL path grammar, both directions. +//! +//! * Sender → broker (`:9080`): +//! `//////`, parsed by [`parse`]. +//! The broker routes on the **`(env, tunnel)` pair** — a tunnel is only +//! meaningful within an env — and forwards the rest. +//! * Broker → receiver: `////` (the env + tunnel +//! removed, the proto + target kept), parsed by [`parse_forwarded`]. The +//! receiver uses `host`/`port` to dial its local target and forwards +//! `tail` to it. +//! +//! `` must be the literal `http` — the only transport implemented. Both +//! parsers reject any other value ([`SUPPORTED_PROTO`]), so an unimplemented +//! transport (e.g. `https`) can never be accepted and silently dialed as +//! cleartext; the grammar still keeps the segment so a real transport can be +//! added by widening the gate. `` and `` +//! are **opaque** to the broker (it neither validates nor restricts them — +//! see `design.md` §4.3 / §10.1). The receiver is the first validator: +//! [`parse_forwarded`] is the receiver's tool, so it *does* reject an empty +//! host and a non-numeric/out-of-range port. +//! +//! No path normalisation is performed in either direction: `..`, `.`, `//`, +//! and percent-encoding pass through verbatim (design §4.4 rule 7). + +use super::env::Env; +use super::error::{EnvError, PathError, TunnelError}; +use super::tunnel::Tunnel; + +/// The only `` the relay implements. Both parsers reject any other value +/// (including `https` and mixed case) so an unimplemented transport can never be +/// accepted and silently dialed as cleartext. Adding a transport means widening +/// this gate deliberately, not just carrying the string through. +pub const SUPPORTED_PROTO: &str = "http"; + +/// Maximum accepted `:path` length in bytes (both parsers). The h2 +/// `max_header_list_size` cap is the primary bound on the request header block; +/// this is defense-in-depth so a single pseudo-header can't drive an outsized +/// allocation before that cap applies, and it bounds the `String`s the parser +/// builds. 8 KiB comfortably covers any legitimate `///http///`. +pub const MAX_PATH_LEN: usize = 8192; + +/// Maximum accepted target host length in bytes (forwarded form). 255 is the +/// DNS name limit; an IP literal is far shorter. +pub const MAX_HOST_LEN: usize = 255; + +/// Result of parsing a sender request path. +#[derive(Debug, PartialEq, Eq)] +pub struct ParsedPath { + /// The alphanumeric env from segment 0 — the routing namespace. + pub env: Env, + /// The tunnel from segment 1 (`[A-Za-z0-9._-]+`). Meaningful only within `env`. + pub tunnel: Tunnel, + /// The path to forward to the receiver, including any query string. + /// Always begins with `///` (env + tunnel removed). + pub forwarded_path: String, +} + +/// Parse a sender `:path` of the form +/// `//////`. +/// +/// Returns the env + tunnel (the `(env, tunnel)` routing key) and the path to +/// forward to the receiver. The forwarded path is `///` +/// followed by any suffix segments joined by `/`, plus any query string from +/// the original path. The `` and `` segments are dropped — the +/// receiver never sees them (design §4.3); the proto + target are kept. +pub fn parse(path: &str) -> Result { + if path.len() > MAX_PATH_LEN { + return Err(PathError::PathTooLong(MAX_PATH_LEN)); + } + let (raw_path, query) = split_query(path); + + if !raw_path.starts_with('/') { + return Err(PathError::Malformed); + } + + let segments: Vec<&str> = raw_path[1..].split('/').collect(); + + // Segment 0 is the env — the routing namespace. Validated alphanumeric. + let Some(env_seg) = segments.first() else { + return Err(PathError::MissingEnv); + }; + if env_seg.is_empty() { + return Err(PathError::MissingEnv); + } + let env = Env::new(*env_seg).map_err(|e| match e { + EnvError::Empty => PathError::MissingEnv, + // A too-long or out-of-charset env segment is an invalid env either way. + EnvError::NonAlphanumeric | EnvError::TooLong(_) => PathError::InvalidEnv, + })?; + + // Segment 1 is the tunnel — meaningful only within the env. Charset + // [A-Za-z0-9._-] (validated by `Tunnel`). + let Some(tunnel_seg) = segments.get(1) else { + return Err(PathError::MissingTunnel); + }; + if tunnel_seg.is_empty() { + return Err(PathError::MissingTunnel); + } + let tunnel = Tunnel::new(*tunnel_seg).map_err(|e| match e { + TunnelError::Empty => PathError::MissingTunnel, + // A too-long or out-of-charset tunnel segment is an invalid tunnel. + TunnelError::InvalidChar | TunnelError::TooLong(_) => PathError::InvalidTunnel, + })?; + + // Segment 2 is the proto. It MUST be the literal `http` — the only + // transport implemented. Reject anything else at this shared boundary + // (e.g. `https`, mixed case) rather than carry it forward, so it can never + // be silently dialed as cleartext downstream (remediation Phase 8.2). + let Some(proto_seg) = segments.get(2) else { + return Err(PathError::MissingProto); + }; + if proto_seg.is_empty() { + return Err(PathError::MissingProto); + } + if *proto_seg != SUPPORTED_PROTO { + // Payload-free: do not retain or echo the attacker-controlled proto + // value (see `PathError::UnsupportedProto`). + return Err(PathError::UnsupportedProto); + } + + // Segments 3 (host) and 4 (port) MUST be present structurally; content + // is opaque here (the receiver validates). + if segments.len() < 5 { + return Err(PathError::MissingTarget); + } + + // Forwarded path: drop env (0) + tunnel (1), keep + // `////` (segments 2..) + query. + let mut forwarded = String::with_capacity(raw_path.len()); + forwarded.push('/'); + forwarded.push_str(&segments[2..].join("/")); + if let Some(q) = query { + forwarded.push_str(q); + } + + Ok(ParsedPath { + env, + tunnel, + forwarded_path: forwarded, + }) +} + +/// The target a receiver should dial, parsed from the forwarded path +/// `////`. +#[derive(Debug, PartialEq, Eq)] +pub struct ForwardedTarget { + /// Transport the sender named. Always `http` — [`parse_forwarded`] rejects + /// anything else ([`SUPPORTED_PROTO`]). Kept as a field so a future + /// transport can branch on it once the gate is widened. + pub proto: String, + /// Target host (opaque string: IP or DNS name). Non-empty. + pub host: String, + /// Target port. Validated `1..=65535`. + pub port: u16, + /// The path to forward to the target, beginning with `/` and including + /// any query string. `/http/h/9080` (no suffix) yields `tail = "/"`. + pub tail: String, +} + +/// Parse the forwarded path `////` the broker sends +/// to a receiver. Unlike [`parse`], this validates the target the broker +/// treated as opaque: empty host → [`PathError::InvalidHost`], non-numeric +/// or out-of-range port → [`PathError::InvalidPort`]. +pub fn parse_forwarded(path: &str) -> Result { + if path.len() > MAX_PATH_LEN { + return Err(PathError::PathTooLong(MAX_PATH_LEN)); + } + let (raw_path, query) = split_query(path); + + if !raw_path.starts_with('/') { + return Err(PathError::Malformed); + } + + let segments: Vec<&str> = raw_path[1..].split('/').collect(); + + // Segment 0 = proto (must be literal `http`), 1 = host, 2 = port, 3.. = tail. + let proto = match segments.first() { + Some(p) if !p.is_empty() => *p, + _ => return Err(PathError::MissingProto), + }; + if proto != SUPPORTED_PROTO { + // Defense in depth: the broker only builds `http` forwarded paths, but a + // receiver must never dial a non-`http` proto as cleartext either. + // Payload-free (see `PathError::UnsupportedProto`). + return Err(PathError::UnsupportedProto); + } + let host = match segments.get(1) { + Some(h) => *h, + None => return Err(PathError::MissingTarget), + }; + let port_seg = match segments.get(2) { + Some(p) => *p, + None => return Err(PathError::MissingTarget), + }; + if host.is_empty() { + return Err(PathError::InvalidHost); + } + if host.len() > MAX_HOST_LEN { + return Err(PathError::HostTooLong(MAX_HOST_LEN)); + } + let port: u16 = port_seg + .parse() + .map_err(|_| PathError::InvalidPort(port_seg.to_string()))?; + + // Tail = "/" + segments[3..] joined, + query. No suffix → "/". + let mut tail = String::with_capacity(raw_path.len()); + tail.push('/'); + if segments.len() > 3 { + tail.push_str(&segments[3..].join("/")); + } + if let Some(q) = query { + tail.push_str(q); + } + + Ok(ForwardedTarget { + proto: proto.to_string(), + host: host.to_string(), + port, + tail, + }) +} + +/// Split a path into `(path_without_query, Some("?..."))`. +fn split_query(path: &str) -> (&str, Option<&str>) { + match path.find('?') { + Some(idx) => (&path[..idx], Some(&path[idx..])), + None => (path, None), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ── Sender form: parse ──────────────────────────────────────────── + + fn ok(path: &str, env: &str, tunnel: &str, forwarded: &str) { + let parsed = parse(path).unwrap_or_else(|e| panic!("expected ok for {path:?}, got {e:?}")); + assert_eq!(parsed.env.as_str(), env, "env for {path:?}"); + assert_eq!(parsed.tunnel.as_str(), tunnel, "tunnel for {path:?}"); + assert_eq!(parsed.forwarded_path, forwarded, "forwarded for {path:?}"); + } + + fn err(path: &str, expected: PathError) { + assert_eq!(parse(path), Err(expected), "for path {path:?}"); + } + + #[test] + fn basic() { + ok( + "/e1/abc123/http/127.0.0.1/9080/foo", + "e1", + "abc123", + "/http/127.0.0.1/9080/foo", + ); + } + + #[test] + fn example_from_spec() { + // The real-world example: env / tunnel / http / host / port / (empty tail). + ok( + "/201kqw76hvvdkhsvhdajh77zxek/mytunnel/http/replit-agent.restate-apps/9080/", + "201kqw76hvvdkhsvhdajh77zxek", + "mytunnel", + "/http/replit-agent.restate-apps/9080/", + ); + } + + #[test] + fn nested_path_with_query() { + ok( + "/e1/abc123/http/10.0.0.5/8080/users/42?fields=name", + "e1", + "abc123", + "/http/10.0.0.5/8080/users/42?fields=name", + ); + } + + #[test] + fn trailing_slash_preserved() { + ok("/e1/abc/http/h/9080/", "e1", "abc", "/http/h/9080/"); + } + + #[test] + fn empty_suffix() { + ok("/e1/abc/http/h/9080", "e1", "abc", "/http/h/9080"); + } + + #[test] + fn deep_path() { + ok( + "/env/key/http/host/1/a/b/c/d/e", + "env", + "key", + "/http/host/1/a/b/c/d/e", + ); + } + + #[test] + fn dns_host_with_dots_and_dashes() { + ok( + "/e1/c/http/my-svc.ns.svc.cluster.local/8080/x", + "e1", + "c", + "/http/my-svc.ns.svc.cluster.local/8080/x", + ); + } + + #[test] + fn host_and_port_content_not_validated_by_sender_parse() { + ok( + "/e1/c/http/not_a_host/not_a_port/x", + "e1", + "c", + "/http/not_a_host/not_a_port/x", + ); + } + + #[test] + fn no_path_normalization() { + ok( + "/e1/c/http/h/9080/../etc/passwd", + "e1", + "c", + "/http/h/9080/../etc/passwd", + ); + ok("/e1/c/http/h/9080//foo", "e1", "c", "/http/h/9080//foo"); + } + + #[test] + fn percent_encoding_forwarded_verbatim() { + ok( + "/e1/c/http/h/9080/hello%20world", + "e1", + "c", + "/http/h/9080/hello%20world", + ); + } + + #[test] + fn query_on_empty_suffix() { + ok("/e1/abc/http/h/9080?q=1", "e1", "abc", "/http/h/9080?q=1"); + } + + #[test] + fn missing_env() { + err("/", PathError::MissingEnv); + err("", PathError::Malformed); + } + + #[test] + fn invalid_env_rejected() { + err("/e-1/c/http/h/9080/x", PathError::InvalidEnv); + } + + #[test] + fn missing_tunnel() { + err("/e1", PathError::MissingTunnel); + err("/e1//http/h/9080/x", PathError::MissingTunnel); + } + + #[test] + fn dashed_dotted_underscored_tunnel_accepted() { + // The relaxed tunnel charset `[A-Za-z0-9._-]+` (Cloud tunnel-name + // compatibility) accepts these; none collide with the `/` delimiter. + ok( + "/e1/abc-123/http/h/9080/foo", + "e1", + "abc-123", + "/http/h/9080/foo", + ); + ok( + "/e1/a.b_c-d/http/h/9080/x", + "e1", + "a.b_c-d", + "/http/h/9080/x", + ); + } + + #[test] + fn tunnel_with_out_of_charset_char_rejected() { + // `~` is URL-safe (unreserved) but outside `[A-Za-z0-9._-]`. + err("/e1/abc~123/http/h/9080/foo", PathError::InvalidTunnel); + } + + #[test] + fn missing_proto() { + err("/e1/abc", PathError::MissingProto); + err("/e1/abc//h/9080/x", PathError::MissingProto); + } + + #[test] + fn missing_host_or_port() { + err("/e1/abc/http", PathError::MissingTarget); + err("/e1/abc/http/h", PathError::MissingTarget); + } + + #[test] + fn missing_leading_slash() { + err("e1/abc/http/h/9080/foo", PathError::Malformed); + } + + #[test] + fn rejects_over_long_path() { + // A path longer than MAX_PATH_LEN is rejected up front (defense-in-depth + // beyond the h2 header-list-size cap). + let long_tail = "a".repeat(MAX_PATH_LEN); + let path = format!("/e1/abc/http/h/9080/{long_tail}"); + assert!(path.len() > MAX_PATH_LEN); + err(&path, PathError::PathTooLong(MAX_PATH_LEN)); + } + + #[test] + fn parse_rejects_non_http_proto() { + // Only literal `http` is accepted; anything else (incl. mixed case) is + // rejected here so it can't be carried forward + dialed as cleartext. + // The error is payload-free — it never echoes the rejected proto value. + err("/e1/abc/https/h/9080/x", PathError::UnsupportedProto); + err("/e1/abc/HTTP/h/9080/x", PathError::UnsupportedProto); + err("/e1/abc/Http/h/9080/x", PathError::UnsupportedProto); + err("/e1/abc/tcp/h/9080/x", PathError::UnsupportedProto); + err("/e1/abc/ftp/h/9080/x", PathError::UnsupportedProto); + // Empty proto is still MissingProto, not UnsupportedProto. + err("/e1/abc//h/9080/x", PathError::MissingProto); + } + + #[test] + fn unsupported_proto_error_never_echoes_the_rejected_value() { + // Payload-free by design: an attacker-controlled (up to MAX_PATH_LEN) + // proto value must not be retained or rendered. Feed a large, distinctive + // proto and assert the formatted error does not contain any of its bytes. + let attacker = "SECRETPROTO".repeat(64); + let path = format!("/e1/abc/{attacker}/h/9080/x"); + let e = parse(&path).unwrap_err(); + assert_eq!(e, PathError::UnsupportedProto); + let rendered = e.to_string(); + assert!( + !rendered.contains("SECRETPROTO"), + "UnsupportedProto must not echo the rejected value, got: {rendered}" + ); + } + + // ── Forwarded form: parse_forwarded ─────────────────────────────── + + fn fwd_ok(path: &str, proto: &str, host: &str, port: u16, tail: &str) { + let t = + parse_forwarded(path).unwrap_or_else(|e| panic!("expected ok for {path:?}, got {e:?}")); + assert_eq!(t.proto, proto, "proto for {path:?}"); + assert_eq!(t.host, host, "host for {path:?}"); + assert_eq!(t.port, port, "port for {path:?}"); + assert_eq!(t.tail, tail, "tail for {path:?}"); + } + + fn fwd_err(path: &str, expected: PathError) { + assert_eq!(parse_forwarded(path), Err(expected), "for path {path:?}"); + } + + #[test] + fn forwarded_rejects_non_http_proto() { + // Defense in depth on the receiver side: never dial a non-http proto. + // Payload-free error (no rejected value echoed). + fwd_err("/https/h/9080/x", PathError::UnsupportedProto); + fwd_err("/HTTP/h/9080/x", PathError::UnsupportedProto); + fwd_err("/Http/h/9080/x", PathError::UnsupportedProto); + fwd_err("//h/9080/x", PathError::MissingProto); + } + + #[test] + fn forwarded_basic() { + fwd_ok( + "/http/127.0.0.1/9080/orders", + "http", + "127.0.0.1", + 9080, + "/orders", + ); + } + + #[test] + fn forwarded_with_query() { + fwd_ok( + "/http/127.0.0.1/9080/orders?x=1", + "http", + "127.0.0.1", + 9080, + "/orders?x=1", + ); + } + + #[test] + fn forwarded_no_suffix_is_root() { + fwd_ok("/http/h/9080", "http", "h", 9080, "/"); + fwd_ok("/http/h/9080?q=1", "http", "h", 9080, "/?q=1"); + } + + #[test] + fn forwarded_trailing_slash() { + fwd_ok("/http/h/9080/", "http", "h", 9080, "/"); + } + + #[test] + fn forwarded_deep_tail_and_dns_host() { + fwd_ok( + "/http/svc.ns.svc.cluster.local/8080/a/b/c", + "http", + "svc.ns.svc.cluster.local", + 8080, + "/a/b/c", + ); + } + + #[test] + fn forwarded_rejects_empty_host() { + fwd_err("/http//9080/x", PathError::InvalidHost); + } + + #[test] + fn forwarded_rejects_over_long_host() { + let host = "h".repeat(MAX_HOST_LEN + 1); + let path = format!("/http/{host}/9080/x"); + fwd_err(&path, PathError::HostTooLong(MAX_HOST_LEN)); + // At the cap is fine. + let host = "h".repeat(MAX_HOST_LEN); + let path = format!("/http/{host}/9080/x"); + assert!(parse_forwarded(&path).is_ok()); + } + + #[test] + fn forwarded_rejects_over_long_path() { + let long_tail = "a".repeat(MAX_PATH_LEN); + let path = format!("/http/h/9080/{long_tail}"); + assert!(path.len() > MAX_PATH_LEN); + fwd_err(&path, PathError::PathTooLong(MAX_PATH_LEN)); + } + + #[test] + fn forwarded_rejects_non_numeric_port() { + fwd_err( + "/http/h/notaport/x", + PathError::InvalidPort("notaport".to_string()), + ); + } + + #[test] + fn forwarded_rejects_out_of_range_port() { + fwd_err( + "/http/h/70000/x", + PathError::InvalidPort("70000".to_string()), + ); + } + + #[test] + fn forwarded_rejects_missing_port() { + fwd_err("/http/h", PathError::MissingTarget); + } + + #[test] + fn forwarded_rejects_malformed() { + fwd_err("http/h/9080", PathError::Malformed); + } + + // --- Restate-as-sender contract (the e2e example rests on this) --- + // + // Restate registers a base URL + // `http://relay:9080///http//` and appends + // `/discover` and `/invoke//`. The relay strips its + // `//` routing prefix and forwards the (proto + target + + // possibly `invoke`-containing) tail verbatim. + + #[test] + fn restate_discovery_path() { + ok( + "/e1/demo/http/127.0.0.1/9080/discover", + "e1", + "demo", + "/http/127.0.0.1/9080/discover", + ); + } + + #[test] + fn restate_invocation_path() { + ok( + "/e1/demo/http/127.0.0.1/9080/invoke/greeter/greet", + "e1", + "demo", + "/http/127.0.0.1/9080/invoke/greeter/greet", + ); + } + + #[test] + fn restate_forwarded_tail_round_trips_at_receiver() { + let t = parse_forwarded("/http/127.0.0.1/9080/invoke/greeter/greet").unwrap(); + assert_eq!(t.proto, "http"); + assert_eq!(t.host, "127.0.0.1"); + assert_eq!(t.port, 9080); + assert_eq!(t.tail, "/invoke/greeter/greet"); + + let d = parse_forwarded("/http/127.0.0.1/9080/discover").unwrap(); + assert_eq!(d.tail, "/discover"); + } +} diff --git a/src/relay/protocol/tunnel.rs b/src/relay/protocol/tunnel.rs new file mode 100644 index 0000000..643895d --- /dev/null +++ b/src/relay/protocol/tunnel.rs @@ -0,0 +1,165 @@ +//! The `Tunnel` newtype: an opaque rendezvous token, `[A-Za-z0-9._-]+`. +//! +//! Tunnels are validated at construction and are otherwise treated as +//! opaque routing labels. The system never interprets tunnel contents. +//! +//! The charset is alphanumerics plus `.`, `_`, and `-`. The dot/dash/ +//! underscore allowance exists for **Restate Cloud tunnel compatibility** +//! (`docs/cloud-tunnel-compat.md`): a Cloud tunnel client may present a +//! self-chosen `tunnel-name` matching Cloud's `^[A-Za-z0-9._-]+$`, and the +//! relay routes on it as an opaque label. None of those characters collide +//! with the path delimiter (`/`), so the sender-path grammar is unaffected. +//! `Env`, by contrast, stays strictly alphanumeric (Cloud env ids are +//! ``, alnum after the `env_` strip). + +use super::error::TunnelError; +use std::fmt; +use std::str::FromStr; + +/// Maximum accepted tunnel length in bytes. A tunnel is externally controlled +/// (sender path segment + receiver `/whoami`/`start-tunnel` field), so it is +/// length-bounded. Cloud tunnel names are short; 128 is generous headroom. +pub const MAX_TUNNEL_LEN: usize = 128; + +/// A validated tunnel value. Always non-empty, `[A-Za-z0-9._-]+`, and at most +/// [`MAX_TUNNEL_LEN`] bytes. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Tunnel(String); + +impl Tunnel { + /// Whether `c` is allowed in a tunnel: ASCII alphanumeric or one of + /// `.`, `_`, `-`. Kept as a single predicate so the charset has exactly + /// one definition. + fn is_valid_char(c: char) -> bool { + c.is_ascii_alphanumeric() || matches!(c, '.' | '_' | '-') + } + + /// Construct a `Tunnel`, validating the input. + pub fn new(s: impl Into) -> Result { + let s = s.into(); + if s.is_empty() { + return Err(TunnelError::Empty); + } + if s.len() > MAX_TUNNEL_LEN { + return Err(TunnelError::TooLong(MAX_TUNNEL_LEN)); + } + if !s.chars().all(Self::is_valid_char) { + return Err(TunnelError::InvalidChar); + } + Ok(Self(s)) + } + + /// The tunnel as a string slice. + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl fmt::Display for Tunnel { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl FromStr for Tunnel { + type Err = TunnelError; + + fn from_str(s: &str) -> Result { + Self::new(s) + } +} + +impl AsRef for Tunnel { + fn as_ref(&self) -> &str { + &self.0 + } +} + +// Serialises as a plain string so the `/whoami` wire body is +// `{"api_key": "...", "tunnel": "..."}`. Used by `whoami::WhoamiBody`. +// (No `Deserialize`: inbound tunnels go through `Tunnel::new` validation.) +impl serde::Serialize for Tunnel { + fn serialize(&self, s: S) -> Result { + s.serialize_str(&self.0) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn accepts_alphanumeric() { + assert!(Tunnel::new("abc123").is_ok()); + assert!(Tunnel::new("ABC").is_ok()); + assert!(Tunnel::new("0").is_ok()); + assert!(Tunnel::new("a").is_ok()); + } + + #[test] + fn accepts_dot_dash_underscore() { + // The relaxed charset (Cloud tunnel-name compatibility, + // `^[A-Za-z0-9._-]+$`). A self-chosen Cloud tunnel name like these + // must round-trip as an opaque routing label. + for ok in ["abc-123", "abc_123", "abc.123", "a.b-c_d", "-", ".", "_"] { + assert!(Tunnel::new(ok).is_ok(), "should accept {ok:?}"); + } + } + + #[test] + fn rejects_empty() { + assert_eq!(Tunnel::new(""), Err(TunnelError::Empty)); + } + + #[test] + fn rejects_too_long() { + assert!(Tunnel::new("a".repeat(MAX_TUNNEL_LEN)).is_ok()); + assert_eq!( + Tunnel::new("a".repeat(MAX_TUNNEL_LEN + 1)), + Err(TunnelError::TooLong(MAX_TUNNEL_LEN)) + ); + } + + #[test] + fn rejects_out_of_charset_characters() { + // Everything outside `[A-Za-z0-9._-]` is rejected, including the + // path delimiter `/`, whitespace, and other punctuation. + for bad in ["abc 123", "abc/123", "abc!", "abc~", "a:b", "a@b", "a+b"] { + assert_eq!( + Tunnel::new(bad), + Err(TunnelError::InvalidChar), + "should reject {bad:?}" + ); + } + } + + #[test] + fn rejects_non_ascii() { + // Unicode letters are not ASCII alphanumeric. + assert_eq!(Tunnel::new("café"), Err(TunnelError::InvalidChar)); + assert_eq!(Tunnel::new("日本"), Err(TunnelError::InvalidChar)); + } + + #[test] + fn round_trips_via_display() { + let c = Tunnel::new("hello123").unwrap(); + assert_eq!(c.to_string(), "hello123"); + assert_eq!(c.as_str(), "hello123"); + } + + #[test] + fn equality_and_hash() { + use std::collections::HashSet; + let mut set = HashSet::new(); + set.insert(Tunnel::new("a").unwrap()); + assert!(set.contains(&Tunnel::new("a").unwrap())); + assert!(!set.contains(&Tunnel::new("b").unwrap())); + } + + #[test] + fn from_str_works() { + let c: Tunnel = "abc".parse().unwrap(); + assert_eq!(c.as_str(), "abc"); + assert!("".parse::().is_err()); + } +} diff --git a/src/relay/protocol/whoami.rs b/src/relay/protocol/whoami.rs new file mode 100644 index 0000000..5d7e7d7 --- /dev/null +++ b/src/relay/protocol/whoami.rs @@ -0,0 +1,291 @@ +//! The `/whoami` response body — produced by the receiver, parsed by the +//! broker. +//! +//! Per design §4.1, the receiver's response to the cluster's `GET /whoami` +//! has a JSON body of the shape: +//! +//! ```json +//! { "env": "...", "api_key": "...", "tunnel": "..." } +//! ``` +//! +//! `env`, `api_key`, and `tunnel` are required and non-empty. The env is +//! alphanumeric and the tunnel is `[A-Za-z0-9._-]+` (validated by the +//! [`Env`] / [`Tunnel`] types); a tunnel is meaningful only within its env, +//! so the receiver reports both. Additional +//! fields are ignored (forward compatibility — `/whoami` carries `instance_id`, +//! and may grow `version`, etc., per design §3 / open questions). + +use super::env::Env; +use super::error::WhoamiError; +use super::tunnel::Tunnel; +use serde::{Deserialize, Serialize}; + +/// Maximum accepted `/whoami` `api_key` length in bytes. This is the relay's own +/// API key (e.g. Restate Cloud `key_.`), which is short — a few tens +/// of bytes — so 512 is generous. The `start_tunnel` bearer token (which can be +/// a multi-KB JWT) is a *separate* credential with its own, larger cap; do not +/// reuse this one for it. +pub const MAX_API_KEY_LEN: usize = 512; + +/// Maximum accepted `/whoami` `instance_id` length in bytes. An ephemeral, +/// receiver-generated identifier; 128 is ample. +pub const MAX_INSTANCE_ID_LEN: usize = 128; + +/// Maximum accepted `start_tunnel` `authorization: Bearer ` length in +/// bytes. This is a *different* credential from the `/whoami` [`MAX_API_KEY_LEN`] +/// key — a Restate Cloud tunnel client may present a multi-KB JWT — so it has +/// its own, larger cap. Still far below the h2 `max_header_list_size` (64 KiB) +/// that bounds the whole header block, so a max-length token is never rejected +/// by the enclosing header cap first. +pub const MAX_BEARER_TOKEN_LEN: usize = 8192; + +/// Parsed contents of a valid `/whoami` response body. +/// +/// `Serialize` is derived so the receiver can build the body it returns from +/// the same type the broker parses — one struct, both directions. (`Tunnel` +/// serialises as a plain string; see its impl in `tunnel.rs`.) +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct WhoamiBody { + /// The env this receiver registers under — the routing namespace that + /// scopes the tunnel. Required: a tunnel is only meaningful within an env. + pub env: Env, + pub api_key: String, + pub tunnel: Tunnel, + /// Optional ephemeral receiver-instance id, generated once per receiver + /// *process* and presented on every connection + reconnection + node + /// (so the broker recognises the receiver's many connections — when it + /// multi-homes for HA — as one logical receiver). Scoped under the + /// tunnel. Drives request affinity: all of a receiver's connections + /// share one affinity bucket `(tunnel, instance_id)`, so a key sticks to + /// the *receiver* (served on whichever node the sender hits) rather than + /// to one specific connection. Absent ⇒ the broker falls back to a + /// per-connection affinity bucket (legacy behaviour). Not persisted: a + /// restarted receiver is a new instance and affinity re-pins. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub instance_id: Option, +} + +/// Internal raw deserialisation type. We keep it private so that all +/// callers go through [`parse`], which enforces non-emptiness and the +/// `Env` / `Tunnel` charsets. +#[derive(Debug, Deserialize)] +struct Raw { + #[serde(default)] + env: Option, + #[serde(default)] + api_key: Option, + #[serde(default)] + tunnel: Option, + #[serde(default)] + instance_id: Option, +} + +/// Parse a `/whoami` response body. +pub fn parse(body: &[u8]) -> Result { + let raw: Raw = serde_json::from_slice(body)?; + + let env_str = raw.env.ok_or(WhoamiError::MissingField("env"))?; + let env = Env::new(env_str)?; + + let api_key = raw.api_key.ok_or(WhoamiError::MissingField("api_key"))?; + if api_key.is_empty() { + return Err(WhoamiError::EmptyApiKey); + } + if api_key.len() > MAX_API_KEY_LEN { + return Err(WhoamiError::ApiKeyTooLong(MAX_API_KEY_LEN)); + } + + let tunnel_str = raw.tunnel.ok_or(WhoamiError::MissingField("tunnel"))?; + let tunnel = Tunnel::new(tunnel_str)?; + + // An empty instance_id is treated as absent (no identity supplied). + let instance_id = raw.instance_id.filter(|s| !s.is_empty()); + if let Some(id) = &instance_id { + if id.len() > MAX_INSTANCE_ID_LEN { + return Err(WhoamiError::InstanceIdTooLong(MAX_INSTANCE_ID_LEN)); + } + } + + Ok(WhoamiBody { + env, + api_key, + tunnel, + instance_id, + }) +} + +impl WhoamiBody { + /// Serialise to the JSON wire body. Used by the receiver to answer + /// `GET /whoami`. Infallible for our fixed shape. + pub fn to_json(&self) -> Vec { + serde_json::to_vec(self).expect("WhoamiBody serialises") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_valid_body() { + let body = br#"{"env":"e1","api_key":"k1","tunnel":"abc123"}"#; + let parsed = parse(body).unwrap(); + assert_eq!(parsed.env.as_str(), "e1"); + assert_eq!(parsed.api_key, "k1"); + assert_eq!(parsed.tunnel.as_str(), "abc123"); + } + + #[test] + fn round_trips_through_json() { + let b = WhoamiBody { + env: Env::new("e1").unwrap(), + api_key: "k1".to_string(), + tunnel: Tunnel::new("abc123").unwrap(), + instance_id: None, + }; + let bytes = b.to_json(); + let back = parse(&bytes).unwrap(); + assert_eq!(back, b); + } + + #[test] + fn round_trips_with_instance_id() { + let b = WhoamiBody { + env: Env::new("e1").unwrap(), + api_key: "k1".to_string(), + tunnel: Tunnel::new("abc123").unwrap(), + instance_id: Some("inst-7f3a".to_string()), + }; + let bytes = b.to_json(); + // The id is on the wire and parses back. + assert!(std::str::from_utf8(&bytes).unwrap().contains("instance_id")); + let back = parse(&bytes).unwrap(); + assert_eq!(back, b); + assert_eq!(back.instance_id.as_deref(), Some("inst-7f3a")); + } + + #[test] + fn empty_instance_id_parses_as_absent() { + let body = br#"{"env":"e1","api_key":"k1","tunnel":"abc","instance_id":""}"#; + assert_eq!(parse(body).unwrap().instance_id, None); + } + + #[test] + fn parses_with_whitespace_and_extra_fields() { + let body = br#" + { + "env": "e1", + "api_key": "k1", + "tunnel": "abc", + "instance_id": "i-future-field", + "version": "1.2.3" + } + "#; + // Extra fields are accepted (forward compat). + let parsed = parse(body).unwrap(); + assert_eq!(parsed.env.as_str(), "e1"); + assert_eq!(parsed.api_key, "k1"); + assert_eq!(parsed.tunnel.as_str(), "abc"); + } + + #[test] + fn rejects_missing_env() { + let body = br#"{"api_key":"k1","tunnel":"abc"}"#; + assert!(matches!(parse(body), Err(WhoamiError::MissingField("env")))); + } + + #[test] + fn rejects_invalid_env() { + let body = br#"{"env":"has-dash","api_key":"k1","tunnel":"abc"}"#; + assert!(matches!(parse(body), Err(WhoamiError::InvalidEnv(_)))); + } + + #[test] + fn rejects_missing_api_key() { + let body = br#"{"env":"e1","tunnel":"abc"}"#; + assert!(matches!( + parse(body), + Err(WhoamiError::MissingField("api_key")) + )); + } + + #[test] + fn rejects_missing_tunnel() { + let body = br#"{"env":"e1","api_key":"k1"}"#; + assert!(matches!( + parse(body), + Err(WhoamiError::MissingField("tunnel")) + )); + } + + #[test] + fn rejects_empty_api_key() { + let body = br#"{"env":"e1","api_key":"","tunnel":"abc"}"#; + assert!(matches!(parse(body), Err(WhoamiError::EmptyApiKey))); + } + + #[test] + fn rejects_over_long_api_key() { + let key = "k".repeat(MAX_API_KEY_LEN + 1); + let body = format!(r#"{{"env":"e1","api_key":"{key}","tunnel":"abc"}}"#); + assert!(matches!( + parse(body.as_bytes()), + Err(WhoamiError::ApiKeyTooLong(_)) + )); + // Exactly at the cap is accepted. + let key = "k".repeat(MAX_API_KEY_LEN); + let body = format!(r#"{{"env":"e1","api_key":"{key}","tunnel":"abc"}}"#); + assert!(parse(body.as_bytes()).is_ok()); + } + + #[test] + fn rejects_over_long_instance_id() { + let id = "i".repeat(MAX_INSTANCE_ID_LEN + 1); + let body = format!(r#"{{"env":"e1","api_key":"k1","tunnel":"abc","instance_id":"{id}"}}"#); + assert!(matches!( + parse(body.as_bytes()), + Err(WhoamiError::InstanceIdTooLong(_)) + )); + } + + #[test] + fn rejects_invalid_tunnel_via_tunnel_type() { + // `~` is outside the relaxed tunnel charset `[A-Za-z0-9._-]`. + let body = br#"{"env":"e1","api_key":"k1","tunnel":"has~tilde"}"#; + assert!(matches!(parse(body), Err(WhoamiError::InvalidTunnel(_)))); + } + + #[test] + fn rejects_empty_tunnel_via_tunnel_type() { + let body = br#"{"env":"e1","api_key":"k1","tunnel":""}"#; + assert!(matches!(parse(body), Err(WhoamiError::InvalidTunnel(_)))); + } + + #[test] + fn rejects_invalid_json() { + let body = b"not json"; + assert!(matches!(parse(body), Err(WhoamiError::InvalidJson(_)))); + } + + #[test] + fn rejects_wrong_type_for_field() { + // api_key is a number, not a string. + let body = br#"{"env":"e1","api_key":42,"tunnel":"abc"}"#; + assert!(matches!(parse(body), Err(WhoamiError::InvalidJson(_)))); + } + + #[test] + fn rejects_empty_body() { + assert!(matches!(parse(b""), Err(WhoamiError::InvalidJson(_)))); + } + + #[test] + fn rejects_null_fields() { + let body = br#"{"env":"e1","api_key":null,"tunnel":"abc"}"#; + // serde sees null as None for Option, which we treat as missing. + assert!(matches!( + parse(body), + Err(WhoamiError::MissingField("api_key")) + )); + } +} diff --git a/src/relay/receiver/handler.rs b/src/relay/receiver/handler.rs new file mode 100644 index 0000000..a2121d0 --- /dev/null +++ b/src/relay/receiver/handler.rs @@ -0,0 +1,41 @@ +//! The `InvokeHandler` seam — the one thing an embedder customises: whether to +//! dial a target `host:port`, serve the invocation in-process, or anything else. +//! The runtime owns all the relay-facing machinery and calls a handler per +//! forwarded stream. + +use crate::relay::protocol::ForwardedTarget; +use bytes::Bytes; + +/// One forwarded invocation handed to an [`InvokeHandler`]. +/// +/// The runtime has already validated the path into `target` (host/port/tail) +/// and stripped the relay framing. The handler owns both halves of the h2 +/// stream — `body` (the request bytes coming from the relay) and `respond` +/// (where it writes the response) — so it can stream in both directions +/// without buffering, matching the bidi model. +/// +/// **Contract:** the handler MUST send a response on `respond` (a success +/// status or an error status). If it drops `respond` without sending, the +/// relay observes the stream reset and the sender gets a 500. +pub struct Invocation { + /// The target the sender named: where to dial and what path to use. + pub target: ForwardedTarget, + /// The sender's method, forwarded verbatim. + pub method: http::Method, + /// The sender's headers, forwarded verbatim (relay framing already + /// stripped). The handler is responsible for any hop-by-hop hygiene. + pub headers: http::HeaderMap, + /// Request body bytes streaming in from the relay. + pub body: h2::RecvStream, + /// Response sink back to the relay (and thence the sender). + pub respond: h2::server::SendResponse, +} + +/// Handles forwarded invocations. One handler instance is shared +/// (`Arc`-wrapped by the runtime) across all streams on all connections, so +/// implementations are `Send + Sync` and hold their own shared state. +#[async_trait::async_trait] +pub trait InvokeHandler: Send + Sync + 'static { + /// Serve one invocation. See [`Invocation`] for the response contract. + async fn handle(&self, invocation: Invocation); +} diff --git a/src/relay/receiver/mod.rs b/src/relay/receiver/mod.rs new file mode 100644 index 0000000..9dfb2ef --- /dev/null +++ b/src/relay/receiver/mod.rs @@ -0,0 +1,44 @@ +//! The relay-facing receiver runtime. +//! +//! A receiver dials *out* to the relay's `:8080`, lets the relay drive the +//! HTTP/2 role-flip (so the receiver is the h2 *server* on its own +//! outbound socket), answers `GET /whoami` with its `{api_key, tunnel}`, +//! and then serves the forwarded streams the relay opens +//! (`/invoke///`). See the relay `design.md` §4.1. +//! +//! Everything relay-facing — dial-in, role-flip, /whoami, stream +//! multiplexing, redial-on-drop — lives here. What to *do* with a +//! forwarded invocation is the one thing left open, via the +//! [`InvokeHandler`] trait: +//! +//! ```ignore +//! use restate_sdk_shared_core::relay::receiver::{Receiver, ReceiverConfig, InvokeHandler, Invocation}; +//! use restate_sdk_shared_core::relay::protocol::{Env, Tunnel}; +//! use std::sync::Arc; +//! +//! struct MyHandler; +//! #[async_trait::async_trait] +//! impl InvokeHandler for MyHandler { +//! async fn handle(&self, inv: Invocation) { +//! // dial inv.target.host:inv.target.port, forward inv.tail, … +//! } +//! } +//! +//! let cfg = ReceiverConfig::new("relay.internal:8080", Env::new("e1")?, Tunnel::new("c1")?, "api-key"); +//! Receiver::new(cfg, Arc::new(MyHandler)).run(shutdown_future).await; +//! ``` +//! +//! An embedder chooses what the handler does: dial the named `host:port` and +//! bridge, serve the invocation in-process, or anything else. The embedded +//! [`crate::relay::loopback`] engine supplies a handler that bridges each +//! forwarded stream to a local HTTP/2 server over a loopback socket. + +mod handler; +mod runtime; + +pub use handler::{Invocation, InvokeHandler}; +pub use runtime::{ConnError, Receiver, ReceiverConfig, TlsClient, TlsConfigError}; + +// Re-export the protocol types that appear in this module's public API so +// callers don't need to reach into `crate::relay::protocol` just to name them. +pub use super::protocol::{Env, ForwardedTarget, H2Tuning, Tunnel}; diff --git a/src/relay/receiver/runtime.rs b/src/relay/receiver/runtime.rs new file mode 100644 index 0000000..9a78940 --- /dev/null +++ b/src/relay/receiver/runtime.rs @@ -0,0 +1,579 @@ +//! The relay-facing receiver runtime: dial the relay, accept the h2 +//! role-flip (we are the *server*), answer `/whoami`, and dispatch each +//! forwarded stream to the [`InvokeHandler`]. Redials on connection loss. + +use super::handler::{Invocation, InvokeHandler}; +use crate::relay::protocol::{Env, H2Tuning, Tunnel, WhoamiBody}; +use bytes::Bytes; +use std::future::Future; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::net::TcpStream; +use tokio::sync::watch; +use tokio_rustls::TlsConnector; +use tracing::{debug, info, warn}; + +/// Client TLS for the receiver→relay (or receiver→LB) hop. Usually unset — +/// receiver TLS is typically terminated at the LB (design §10.2) so the +/// relay sees plain TCP — but supported for direct-to-relay deployments. +#[derive(Clone)] +pub struct TlsClient { + pub connector: TlsConnector, + pub server_name: rustls_pki_types::ServerName<'static>, +} + +/// Error building a [`TlsClient`] from PEM material. +#[derive(Debug, thiserror::Error)] +#[error("TLS config: {0}")] +pub struct TlsConfigError(String); + +impl TlsClient { + /// Build receiver→relay client TLS from PEM bytes — the ergonomic + /// constructor so a caller doesn't have to hand-assemble a + /// `rustls::ClientConfig` + `TlsConnector`. + /// + /// * `ca_pem` — CA cert(s) verifying the relay's **server** cert. + /// * `client_identity` — `Some((cert_pem, key_pem))` presents a client cert + /// (rare on `:8080` — receiver auth is the `/whoami` API key, design + /// §10.2); `None` = server-only TLS, the usual case. + /// * `server_name` — DNS name to verify on the relay's cert. + /// + /// Receiver→relay TLS is typically terminated at the LB (so the relay sees + /// plain TCP and no `TlsClient` is needed); this is for direct-to-relay + /// deployments. PEM parsing uses `rustls-pki-types`' `PemObject`; the ring + /// crypto provider is installed once (idempotent). + pub fn from_pem( + ca_pem: &[u8], + client_identity: Option<(&[u8], &[u8])>, + server_name: &str, + ) -> Result { + use rustls_pki_types::pem::PemObject; + use rustls_pki_types::{CertificateDer, PrivateKeyDer, ServerName}; + use std::sync::Arc; + use tokio_rustls::rustls::{ClientConfig, RootCertStore}; + + let _ = tokio_rustls::rustls::crypto::ring::default_provider().install_default(); + + let mut roots = RootCertStore::empty(); + let mut added = 0usize; + for ca in CertificateDer::pem_slice_iter(ca_pem) { + let ca = ca.map_err(|e| TlsConfigError(format!("CA PEM: {e}")))?; + roots + .add(ca) + .map_err(|e| TlsConfigError(format!("add CA: {e}")))?; + added += 1; + } + if added == 0 { + return Err(TlsConfigError("no CERTIFICATE block in ca_pem".into())); + } + + let builder = ClientConfig::builder().with_root_certificates(roots); + let mut config = match client_identity { + Some((cert_pem, key_pem)) => { + let certs = CertificateDer::pem_slice_iter(cert_pem) + .collect::, _>>() + .map_err(|e| TlsConfigError(format!("client cert PEM: {e}")))?; + let key = PrivateKeyDer::from_pem_slice(key_pem) + .map_err(|e| TlsConfigError(format!("client key PEM: {e}")))?; + builder + .with_client_auth_cert(certs, key) + .map_err(|e| TlsConfigError(format!("client auth cert: {e}")))? + } + None => builder.with_no_client_auth(), + }; + config.alpn_protocols = vec![b"h2".to_vec()]; + + let server_name = ServerName::try_from(server_name.to_string()) + .map_err(|e| TlsConfigError(format!("server name {server_name:?}: {e}")))?; + Ok(Self { + connector: TlsConnector::from(Arc::new(config)), + server_name, + }) + } +} + +/// Receiver configuration. +pub struct ReceiverConfig { + /// `host:port` of the relay's receiver port (`:8080`) or the LB in front + /// of it. + pub relay_addr: String, + /// The env (routing namespace) this receiver registers under. Reported in + /// `/whoami`. A tunnel is only meaningful within its env, so the broker + /// keys membership on the `(env, tunnel)` pair. + pub env: Env, + /// The tunnel this receiver registers under (scoped by `env`). + pub tunnel: Tunnel, + /// The API key presented in the `/whoami` body. + pub api_key: String, + /// Optional client TLS. `None` = plain TCP. + pub tls: Option, + /// h2 tuning for our server side of the reversed connection. + pub h2_tuning: H2Tuning, + /// Delay between a connection ending and the next redial. + pub reconnect_backoff: Duration, + /// Ephemeral receiver-instance id sent in `/whoami` (see + /// [`crate::relay::protocol::WhoamiBody::instance_id`]). Generated once per + /// process by [`ReceiverConfig::new`] and held across reconnections so + /// the broker treats this receiver's connections as one logical receiver + /// for affinity. `None` opts out (legacy per-connection affinity). + pub instance_id: Option, + /// How many concurrent connections to fan out across the relay nodes that + /// `relay_addr` resolves to (multi-homing — R4 in the reconnection + /// contract). `None` (the default) opens **one connection per resolved + /// node** (K=N) — every relay node gets a local owner for this receiver, + /// so affinity traffic is served locally wherever the sender lands. `Some(k)` + /// pins exactly `k` connection slots; when `k` exceeds the number of + /// distinct resolved nodes the surplus slots double up on nodes + /// (best-effort spread). Always at least 1. All slots share the one + /// [`instance_id`](ReceiverConfig::instance_id), so the broker collapses + /// them into a single affinity bucket — matching the native Node + /// `connect()` client. + pub connections: Option, +} + +/// Generate an ephemeral, process-unique receiver-instance id. Uses the +/// OS-seeded entropy behind `RandomState` (the same source `HashMap` uses for +/// its DoS-resistance seed) so we get ~128 bits of randomness with **no new +/// dependency** — deliberately avoiding a `uuid`/`rand` pull-in for one id. +/// Ephemeral by design: a fresh id per process means a restarted receiver is +/// a new instance and request affinity re-pins, which is the intended +/// behaviour (a restarted process is genuinely new). +pub fn generate_instance_id() -> String { + use std::hash::{BuildHasher, Hasher}; + let chunk = || { + let mut h = std::collections::hash_map::RandomState::new().build_hasher(); + h.write_u8(0xA5); + h.finish() + }; + format!("{:016x}{:016x}", chunk(), chunk()) +} + +impl ReceiverConfig { + /// Construct with sensible defaults (plain TCP, recommended h2 tuning, + /// 1s reconnect backoff, a freshly-generated ephemeral `instance_id`). + pub fn new( + relay_addr: impl Into, + env: Env, + tunnel: Tunnel, + api_key: impl Into, + ) -> Self { + Self { + relay_addr: relay_addr.into(), + env, + tunnel, + api_key: api_key.into(), + tls: None, + h2_tuning: H2Tuning::default(), + reconnect_backoff: Duration::from_secs(1), + instance_id: Some(generate_instance_id()), + connections: None, + } + } +} + +/// Errors that end a single relay connection (the run loop redials after). +#[derive(Debug, thiserror::Error)] +pub enum ConnError { + #[error("connect failed: {0}")] + Connect(#[from] std::io::Error), + #[error("h2 server handshake / connection error: {0}")] + H2(#[from] h2::Error), +} + +/// A configured receiver: the runtime plus the handler it dispatches to. +pub struct Receiver { + config: ReceiverConfig, + handler: Arc, +} + +impl Receiver { + pub fn new(config: ReceiverConfig, handler: Arc) -> Self { + Self { config, handler } + } + + /// Run until `shutdown` resolves. Resolves `relay_addr` to its relay nodes + /// and fans out **K connection slots** (multi-homing — R4): each slot + /// repeatedly dials a node, serves the reversed connection until it ends, + /// then redials after the backoff, re-resolving DNS each time so the slot + /// follows node churn. `shutdown` firing cancels every in-flight + /// connection and ends all slots. + /// + /// K is [`ReceiverConfig::connections`] when set, else the number of nodes + /// `relay_addr` resolves to at startup (one connection per node). All slots + /// share the one `/whoami` body — including the `instance_id` — so the + /// broker treats them as a single logical receiver for affinity. + pub async fn run(self, shutdown: impl Future + Send) { + let Receiver { config, handler } = self; + // Precompute the /whoami body once — it never changes for a given + // receiver identity, and every slot presents the same one (shared + // instance_id → one affinity bucket across all K connections). + let whoami = Arc::new( + WhoamiBody { + env: config.env.clone(), + api_key: config.api_key.clone(), + tunnel: config.tunnel.clone(), + instance_id: config.instance_id.clone(), + } + .to_json(), + ); + let config = Arc::new(config); + + // Decide the slot count from an initial resolution. `connections` + // overrides; otherwise one slot per resolved node (K=N). If resolution + // fails right now we still start a single slot, which keeps retrying. + let initial = resolve(&config.relay_addr).await; + let slot_count = config.connections.unwrap_or(initial.len()).max(1); + + info!( + relay = %config.relay_addr, + tunnel = %config.tunnel, + resolved = initial.len(), + slots = slot_count, + "receiver: starting" + ); + + // Shared record of which resolved addresses currently have a live + // connection, so each slot can prefer a node no sibling already holds. + let live = Arc::new(LiveTargets::default()); + + // One watch channel broadcasts shutdown to every slot; the driver + // (this task) owns the `shutdown` future and flips it. + let (stop_tx, stop_rx) = watch::channel(false); + let mut slots = Vec::with_capacity(slot_count); + for idx in 0..slot_count { + let config = config.clone(); + let handler = handler.clone(); + let whoami = whoami.clone(); + let live = live.clone(); + let stop_rx = stop_rx.clone(); + slots.push(tokio::spawn(async move { + slot_loop(idx, config, handler, whoami, live, stop_rx).await; + })); + } + drop(stop_rx); + + shutdown.await; + let _ = stop_tx.send(true); + for slot in slots { + let _ = slot.await; + } + info!("receiver: shutdown"); + } +} + +/// One connection slot: pick a node, dial, serve until the connection ends, +/// then re-resolve and redial after the backoff — until `stop` flips. +async fn slot_loop( + idx: usize, + config: Arc, + handler: Arc, + whoami: Arc>, + live: Arc, + mut stop: watch::Receiver, +) { + loop { + if *stop.borrow_and_update() { + break; + } + // Re-resolve every iteration so a slot follows DNS / node churn. + let targets = resolve(&config.relay_addr).await; + match live.pick(&targets, idx) { + Some(addr) => { + // Hold the address for the lifetime of this connection so + // sibling slots steer toward other nodes. + let _guard = live.hold(addr); + let outcome = tokio::select! { + _ = stop.changed() => break, + r = dial_and_serve(addr, &config, &handler, &whoami) => r, + }; + match outcome { + Ok(()) => debug!(slot = idx, %addr, "receiver: connection closed, redialing"), + Err(e) => { + warn!(slot = idx, %addr, ?e, "receiver: connection ended, redialing after backoff") + } + } + } + None => { + debug!( + slot = idx, + relay = %config.relay_addr, + "receiver: no relay address resolved, retrying after backoff" + ); + } + } + tokio::select! { + _ = stop.changed() => break, + _ = tokio::time::sleep(config.reconnect_backoff) => {} + } + } +} + +/// Resolve `relay_addr` to its socket addresses (all A/AAAA records, via the +/// platform resolver — no new dependency). Returns empty on failure; the +/// caller treats that as "nothing to dial yet" and backs off. +async fn resolve(relay_addr: &str) -> Vec { + match tokio::net::lookup_host(relay_addr).await { + Ok(addrs) => addrs.collect(), + Err(e) => { + debug!(%relay_addr, ?e, "receiver: DNS resolution failed"); + Vec::new() + } + } +} + +/// One connection lifecycle: dial the chosen node, role-flip into an h2 +/// server, serve streams until the connection ends. +async fn dial_and_serve( + addr: SocketAddr, + config: &ReceiverConfig, + handler: &Arc, + whoami: &Arc>, +) -> Result<(), ConnError> { + let socket = TcpStream::connect(addr).await?; + let _ = socket.set_nodelay(true); + debug!(%addr, "receiver: connected, awaiting role-flip"); + + match &config.tls { + Some(tls) => { + let tls_socket = tls + .connector + .connect(tls.server_name.clone(), socket) + .await?; + serve_connection(tls_socket, config, handler, whoami).await + } + None => serve_connection(socket, config, handler, whoami).await, + } +} + +/// Shared record of which resolved relay addresses currently have a live +/// connection. Each slot consults this to **prefer a node no sibling already +/// holds**, turning "K connections" into "K *distinct* nodes" (best-effort — +/// when K exceeds the resolved-node count, surplus slots deterministically +/// double up by slot index). Mirrors the Node `connect()` client's +/// `pickTarget`. The held/pick window is racy by design: two slots starting at +/// once may briefly pick the same node; they re-spread on the next redial. +#[derive(Default)] +struct LiveTargets { + held: Mutex>, +} + +impl LiveTargets { + /// Choose a target for `slot` from `targets`: the first not currently held + /// by a live sibling, else a deterministic fallback by slot index (all + /// nodes already held, i.e. K > distinct nodes). `None` iff `targets` is + /// empty. + fn pick(&self, targets: &[SocketAddr], slot: usize) -> Option { + if targets.is_empty() { + return None; + } + let held = self.held.lock().expect("LiveTargets mutex poisoned"); + if let Some(addr) = targets.iter().find(|a| !held.contains(a)) { + return Some(*addr); + } + Some(targets[slot % targets.len()]) + } + + /// Record `addr` as held; the returned guard releases it on drop (when the + /// connection ends), so the node becomes preferable to siblings again. + fn hold(self: &Arc, addr: SocketAddr) -> HoldGuard { + self.held + .lock() + .expect("LiveTargets mutex poisoned") + .push(addr); + HoldGuard { + book: self.clone(), + addr, + } + } +} + +/// RAII release of a held address (see [`LiveTargets::hold`]). +struct HoldGuard { + book: Arc, + addr: SocketAddr, +} + +impl Drop for HoldGuard { + fn drop(&mut self) { + let mut held = self.book.held.lock().expect("LiveTargets mutex poisoned"); + if let Some(pos) = held.iter().position(|a| *a == self.addr) { + held.swap_remove(pos); + } + } +} + +/// Drive the h2 server side of the reversed connection: the relay is the h2 +/// client (role-flip), we accept its streams. Spawns a task per stream so +/// concurrent invocations run in parallel and `accept()` keeps the +/// connection (and its PING acks) flowing. +async fn serve_connection( + socket: IO, + config: &ReceiverConfig, + handler: &Arc, + whoami: &Arc>, +) -> Result<(), ConnError> +where + IO: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let mut builder = h2::server::Builder::new(); + config.h2_tuning.apply_server(&mut builder); + let mut connection = builder.handshake::<_, Bytes>(socket).await?; + debug!("receiver: role-flip complete, serving"); + + while let Some(accepted) = connection.accept().await { + let (request, respond) = match accepted { + Ok(pair) => pair, + Err(e) => { + warn!(?e, "receiver: stream accept error"); + return Err(ConnError::H2(e)); + } + }; + let handler = handler.clone(); + let whoami = whoami.clone(); + tokio::spawn(serve_stream(request, respond, handler, whoami)); + } + Ok(()) +} + +/// Serve a single accepted stream: answer `/whoami`, otherwise parse the +/// forwarded target and dispatch to the handler. +async fn serve_stream( + request: http::Request, + mut respond: h2::server::SendResponse, + handler: Arc, + whoami: Arc>, +) { + let (parts, body) = request.into_parts(); + + // /whoami: the relay opens this first to learn our identity. + if parts.method == http::Method::GET && parts.uri.path() == "/whoami" { + if let Err(e) = reply_whoami(&mut respond, &whoami) { + warn!(?e, "receiver: failed to answer /whoami"); + } + return; + } + + // Everything else is a forwarded invocation. Parse the target from the + // full path+query; a malformed target is the relay/sender's error → 400 + // (we never reach the handler). + let path_and_query = parts + .uri + .path_and_query() + .map(|pq| pq.as_str()) + .unwrap_or_else(|| parts.uri.path()); + match crate::relay::protocol::parse_forwarded(path_and_query) { + Ok(target) => { + handler + .handle(Invocation { + target, + method: parts.method, + headers: parts.headers, + body, + respond, + }) + .await; + } + Err(e) => { + warn!(?e, path = %path_and_query, "receiver: unparseable forwarded path → 400"); + let _ = reply_status(&mut respond, http::StatusCode::BAD_REQUEST); + } + } +} + +/// Send the `/whoami` JSON body. +fn reply_whoami( + respond: &mut h2::server::SendResponse, + whoami: &Arc>, +) -> Result<(), h2::Error> { + let response = http::Response::builder() + .status(http::StatusCode::OK) + .header(http::header::CONTENT_TYPE, "application/json") + .body(()) + .expect("static /whoami response builds"); + let mut send = respond.send_response(response, false)?; + send.send_data(Bytes::from(whoami.as_ref().clone()), true)?; + Ok(()) +} + +/// Send an empty-body response with the given status. +fn reply_status( + respond: &mut h2::server::SendResponse, + status: http::StatusCode, +) -> Result<(), h2::Error> { + let response = http::Response::builder() + .status(status) + .body(()) + .expect("status response builds"); + let mut send = respond.send_response(response, false)?; + let _ = send.send_data(Bytes::new(), true); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::{HoldGuard, LiveTargets, TlsClient}; + use std::net::SocketAddr; + use std::sync::Arc; + + fn addr(s: &str) -> SocketAddr { + s.parse().unwrap() + } + + #[test] + fn pick_prefers_an_unheld_node_then_spreads_by_slot() { + let nodes = [addr("10.0.0.1:8080"), addr("10.0.0.2:8080")]; + let live = Arc::new(LiveTargets::default()); + + // Slot 0 takes node 0 (first unheld). + let a0 = live.pick(&nodes, 0).unwrap(); + assert_eq!(a0, nodes[0]); + let g0 = live.hold(a0); + + // Slot 1: node 0 is held, so it must land on node 1. + let a1 = live.pick(&nodes, 1).unwrap(); + assert_eq!(a1, nodes[1]); + let g1 = live.hold(a1); + + // Slot 2 (K > nodes): everything held → deterministic by index. + assert_eq!(live.pick(&nodes, 2).unwrap(), nodes[2 % nodes.len()]); + assert_eq!(live.pick(&nodes, 3).unwrap(), nodes[3 % nodes.len()]); + + // Releasing a slot frees its node for the next picker. + drop(g0); + assert_eq!(live.pick(&nodes, 5).unwrap(), nodes[0]); + + drop(g1); + } + + #[test] + fn pick_returns_none_when_nothing_resolves() { + let live = Arc::new(LiveTargets::default()); + assert!(live.pick(&[], 0).is_none()); + } + + #[test] + fn hold_guard_is_droppable() { + // Compile-time check that HoldGuard's Drop releases cleanly even when + // the address was already swap-removed (idempotent position lookup). + let live = Arc::new(LiveTargets::default()); + let g: HoldGuard = live.hold(addr("127.0.0.1:1")); + live.held.lock().unwrap().clear(); + drop(g); // position() returns None → no panic + assert!(live.held.lock().unwrap().is_empty()); + } + + #[test] + fn from_pem_rejects_invalid_ca() { + // No CERTIFICATE block in the CA input → error (not a panic). + assert!(TlsClient::from_pem(b"not a pem", None, "relay.example").is_err()); + assert!(TlsClient::from_pem(b"", None, "relay.example").is_err()); + assert!(TlsClient::from_pem( + b"-----BEGIN FOO-----\nAAAA\n-----END FOO-----\n", + None, + "relay.example" + ) + .is_err()); + } +} diff --git a/tests/relay_loopback.rs b/tests/relay_loopback.rs new file mode 100644 index 0000000..e58169c --- /dev/null +++ b/tests/relay_loopback.rs @@ -0,0 +1,182 @@ +//! End-to-end test of the embedded relay-tunnel loopback engine. +//! +//! Compiled only under the `tunnel` feature (the whole file is gated, so a +//! default `cargo test` sees an empty test binary). Run with: +//! `cargo test --features tunnel --test relay_loopback`. +//! +//! The test stands up two fakes and the real engine between them: +//! * a **fake relay** — an h2 *client* (role-flip) that the receiver dials, +//! asks `/whoami`, then opens a forwarded `/http/host/port/tail` stream on; +//! * a **fake local server** — an h2 *server* on `127.0.0.1:` playing +//! the SDK's own HTTP/2 server, which echoes the tail. +//! +//! It asserts the forwarded request round-trips relay → engine → loopback → +//! local server → back, with the tail-only `:path` preserved and the body +//! delivered. +#![cfg(feature = "tunnel")] + +use bytes::Bytes; +use restate_sdk_shared_core::relay::{Config, Engine}; +use std::sync::Arc; +use tokio::net::TcpListener; + +/// A minimal fake local HTTP/2 (h2c) server: accept one connection, serve +/// streams by echoing `served ` and draining the request body. Returns +/// the (host, port, body) it observed on the first stream via `tx`. +async fn spawn_fake_local_server( + listener: TcpListener, + tx: tokio::sync::mpsc::UnboundedSender<(String, Vec)>, +) { + tokio::spawn(async move { + let (socket, _) = listener.accept().await.expect("local server accept"); + let mut conn = h2::server::handshake(socket) + .await + .expect("local h2 handshake"); + while let Some(accepted) = conn.accept().await { + let (request, mut respond) = match accepted { + Ok(pair) => pair, + Err(_) => break, + }; + let tx = tx.clone(); + tokio::spawn(async move { + // Full path + query (the `:path` pseudo-header verbatim), so + // the test can confirm the query survived the tail-only replay. + let path = request + .uri() + .path_and_query() + .map(|pq| pq.as_str().to_string()) + .unwrap_or_else(|| request.uri().path().to_string()); + let mut body = request.into_body(); + let mut got = Vec::new(); + while let Some(chunk) = body.data().await { + let chunk = chunk.expect("local body chunk"); + let _ = body.flow_control().release_capacity(chunk.len()); + got.extend_from_slice(&chunk); + } + let _ = tx.send((path.clone(), got)); + + let response = http::Response::builder() + .status(http::StatusCode::OK) + .body(()) + .unwrap(); + let mut send = respond.send_response(response, false).unwrap(); + let _ = send.send_data(Bytes::from(format!("served {path}")), true); + }); + } + }); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn forwarded_request_round_trips_through_loopback() { + // ── Fake local server (the SDK's own h2c server) ────────────────────── + let local_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let local_port = local_listener.local_addr().unwrap().port(); + let (obs_tx, mut obs_rx) = tokio::sync::mpsc::unbounded_channel(); + spawn_fake_local_server(local_listener, obs_tx).await; + + // ── Fake relay (an h2 client, per the role-flip) ────────────────────── + let relay_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let relay_addr = relay_listener.local_addr().unwrap(); + + // ── Start the engine (its own runtime) pointed at both ──────────────── + let config: Config = serde_json::from_value(serde_json::json!({ + "relay_addr": relay_addr.to_string(), + "env": "e1", + "tunnel": "c1", + "api_key": "api-key", + "local_port": local_port, + "reconnect_backoff_millis": 50, + })) + .unwrap(); + let handle = Arc::new(std::sync::Mutex::new(Some( + Engine::start(config).expect("engine starts"), + ))); + assert!(handle + .lock() + .unwrap() + .as_ref() + .unwrap() + .status() + .is_running()); + + // The receiver dials us; accept and drive it as the relay (h2 client). + let (socket, _) = + tokio::time::timeout(std::time::Duration::from_secs(5), relay_listener.accept()) + .await + .expect("receiver should dial the fake relay") + .unwrap(); + let (send_request, connection) = h2::client::handshake(socket).await.unwrap(); + tokio::spawn(async move { + let _ = connection.await; + }); + + // /whoami first (the relay learns our identity). + let mut sr = send_request.ready().await.unwrap(); + let (resp, _) = sr + .send_request( + http::Request::builder() + .method(http::Method::GET) + .uri("/whoami") + .body(()) + .unwrap(), + true, + ) + .unwrap(); + let response = resp.await.unwrap(); + assert_eq!(response.status(), http::StatusCode::OK); + let mut wbody = response.into_body(); + let mut buf = Vec::new(); + while let Some(c) = wbody.data().await { + let c = c.unwrap(); + wbody.flow_control().release_capacity(c.len()).unwrap(); + buf.extend_from_slice(&c); + } + let json: serde_json::Value = serde_json::from_slice(&buf).unwrap(); + assert_eq!(json["tunnel"], "c1"); + assert_eq!(json["env"], "e1"); + assert_eq!(json["api_key"], "api-key"); + + // Forwarded invocation with a body — should be bridged to the local server. + let mut sr = sr.ready().await.unwrap(); + let (resp2, mut req_body) = sr + .send_request( + http::Request::builder() + .method(http::Method::POST) + .uri("/http/127.0.0.1/9999/orders?x=1") + .body(()) + .unwrap(), + false, + ) + .unwrap(); + req_body + .send_data(Bytes::from_static(b"hello"), true) + .unwrap(); + + let response2 = resp2.await.unwrap(); + assert_eq!(response2.status(), http::StatusCode::OK); + let mut rbody = response2.into_body(); + let mut buf2 = Vec::new(); + while let Some(c) = rbody.data().await { + let c = c.unwrap(); + rbody.flow_control().release_capacity(c.len()).unwrap(); + buf2.extend_from_slice(&c); + } + // The local server echoed the tail-only path — proving `:path` was the tail, + // not the full `/e1/c1/http/...` sender path. + assert_eq!(buf2, b"served /orders?x=1"); + + // The local server observed the forwarded tail + body. + let (path, body) = tokio::time::timeout(std::time::Duration::from_secs(5), obs_rx.recv()) + .await + .expect("local server should have been invoked") + .expect("channel open"); + assert_eq!(path, "/orders?x=1"); + assert_eq!(body, b"hello"); + + // Stop the engine from a blocking context (it drops a tokio runtime). + tokio::task::spawn_blocking(move || { + drop(handle.lock().unwrap().take()); + }) + .await + .unwrap(); +} From 74b116983caa675a5bbae313288fe47fbbad8f26 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Thu, 29 Oct 2026 11:45:00 +0100 Subject: [PATCH 2/2] feat(relay): carry a receiver version in /whoami `Config.version` (engine JSON) -> `ReceiverConfig.version` -> the `/whoami` body's optional `version` ([A-Za-z0-9._-], at most 64 bytes; empty = absent). A relay routes a versioned receiver as the tunnel `.` and, with auto-registration, registers each version as its own deployment. Co-Authored-By: Claude Opus 5.5 --- src/relay/loopback.rs | 21 +++++++++++++++++ src/relay/protocol/error.rs | 2 ++ src/relay/protocol/whoami.rs | 44 +++++++++++++++++++++++++++++++++++ src/relay/receiver/runtime.rs | 5 ++++ 4 files changed, 72 insertions(+) diff --git a/src/relay/loopback.rs b/src/relay/loopback.rs index 6581c6f..3a855b3 100644 --- a/src/relay/loopback.rs +++ b/src/relay/loopback.rs @@ -69,6 +69,12 @@ pub struct Config { /// `None` = auto-generate an ephemeral, process-unique id. #[serde(default)] pub instance_id: Option, + /// The version of the code being served (e.g. `1.2.3` or a build id), + /// sent in `/whoami`: the relay routes each version apart and, with + /// auto-registration, registers each as its own deployment. `None` (or + /// empty) = unversioned. + #[serde(default)] + pub version: Option, /// Delay between a connection ending and the next redial, in milliseconds. /// `None` = 1000ms. #[serde(default)] @@ -138,6 +144,7 @@ impl Config { let mut rc = ReceiverConfig::new(self.relay_addr, env, tunnel, self.api_key); rc.tls = tls; rc.connections = self.connections; + rc.version = self.version.filter(|v| !v.is_empty()); if let Some(id) = self.instance_id { // An explicit id overrides the auto-generated one; an empty string // is treated as "opt out of instance affinity". @@ -457,6 +464,20 @@ mod tests { assert_eq!(empty_id.into_receiver_config().unwrap().instance_id, None); } + #[test] + fn config_passes_the_version_through() { + let parse = |version: &str| -> Option { + let config: Config = serde_json::from_str(&format!( + r#"{{"relay_addr":"h:8080","env":"e1","tunnel":"c1","api_key":"k","local_port":1{version}}}"# + )) + .unwrap(); + config.into_receiver_config().unwrap().version + }; + assert_eq!(parse(r#","version":"v1""#).as_deref(), Some("v1")); + assert_eq!(parse(r#","version":"""#), None); + assert_eq!(parse(""), None); + } + #[test] fn status_json_shape() { let s = Status::new(); diff --git a/src/relay/protocol/error.rs b/src/relay/protocol/error.rs index c02567e..fb519d3 100644 --- a/src/relay/protocol/error.rs +++ b/src/relay/protocol/error.rs @@ -88,4 +88,6 @@ pub enum WhoamiError { ApiKeyTooLong(usize), #[error("instance_id field exceeds the maximum length ({0} bytes)")] InstanceIdTooLong(usize), + #[error("version field must be 1 to {0} bytes of [A-Za-z0-9._-]")] + InvalidVersion(usize), } diff --git a/src/relay/protocol/whoami.rs b/src/relay/protocol/whoami.rs index 5d7e7d7..8d8eb92 100644 --- a/src/relay/protocol/whoami.rs +++ b/src/relay/protocol/whoami.rs @@ -63,6 +63,12 @@ pub struct WhoamiBody { /// restarted receiver is a new instance and affinity re-pins. #[serde(default, skip_serializing_if = "Option::is_none")] pub instance_id: Option, + /// The version of the code this receiver serves (the relay's tunnel + /// charset, at most [`MAX_VERSION_LEN`] bytes). A relay routes a versioned + /// receiver as the tunnel `.` and, with auto-registration, + /// registers it in its schema registry as its own deployment. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub version: Option, } /// Internal raw deserialisation type. We keep it private so that all @@ -78,6 +84,8 @@ struct Raw { tunnel: Option, #[serde(default)] instance_id: Option, + #[serde(default)] + version: Option, } /// Parse a `/whoami` response body. @@ -106,14 +114,36 @@ pub fn parse(body: &[u8]) -> Result { } } + // Empty = absent, like `instance_id`. + let version = raw.version.filter(|s| !s.is_empty()); + if let Some(v) = &version { + if !valid_version(v) { + return Err(WhoamiError::InvalidVersion(MAX_VERSION_LEN)); + } + } + Ok(WhoamiBody { env, api_key, tunnel, instance_id, + version, }) } +/// Maximum accepted `/whoami` `version` length in bytes. +pub const MAX_VERSION_LEN: usize = 64; + +/// A `/whoami` `version`: 1 to [`MAX_VERSION_LEN`] bytes of the [`Tunnel`] +/// charset, so `.` is a tunnel too. +pub fn valid_version(version: &str) -> bool { + !version.is_empty() + && version.len() <= MAX_VERSION_LEN + && version + .bytes() + .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-')) +} + impl WhoamiBody { /// Serialise to the JSON wire body. Used by the receiver to answer /// `GET /whoami`. Infallible for our fixed shape. @@ -142,6 +172,7 @@ mod tests { api_key: "k1".to_string(), tunnel: Tunnel::new("abc123").unwrap(), instance_id: None, + version: None, }; let bytes = b.to_json(); let back = parse(&bytes).unwrap(); @@ -155,6 +186,7 @@ mod tests { api_key: "k1".to_string(), tunnel: Tunnel::new("abc123").unwrap(), instance_id: Some("inst-7f3a".to_string()), + version: None, }; let bytes = b.to_json(); // The id is on the wire and parses back. @@ -288,4 +320,16 @@ mod tests { Err(WhoamiError::MissingField("api_key")) )); } + + #[test] + fn version_round_trips_and_is_validated() { + let body = br#"{"env":"e1","api_key":"k1","tunnel":"checkout","version":"1.2.3-rc_1"}"#; + let parsed = parse(body).unwrap(); + assert_eq!(parsed.version.as_deref(), Some("1.2.3-rc_1")); + assert_eq!(parse(&parsed.to_json()).unwrap(), parsed); + let empty = br#"{"env":"e1","api_key":"k1","tunnel":"t","version":""}"#; + assert_eq!(parse(empty).unwrap().version, None); + let bad = br#"{"env":"e1","api_key":"k1","tunnel":"t","version":"v/1"}"#; + assert!(matches!(parse(bad), Err(WhoamiError::InvalidVersion(_)))); + } } diff --git a/src/relay/receiver/runtime.rs b/src/relay/receiver/runtime.rs index 9a78940..90066f7 100644 --- a/src/relay/receiver/runtime.rs +++ b/src/relay/receiver/runtime.rs @@ -131,6 +131,9 @@ pub struct ReceiverConfig { /// them into a single affinity bucket — matching the native Node /// `connect()` client. pub connections: Option, + /// The code version this receiver serves, sent in `/whoami` (see + /// [`crate::relay::protocol::WhoamiBody::version`]). `None` = unversioned. + pub version: Option, } /// Generate an ephemeral, process-unique receiver-instance id. Uses the @@ -169,6 +172,7 @@ impl ReceiverConfig { reconnect_backoff: Duration::from_secs(1), instance_id: Some(generate_instance_id()), connections: None, + version: None, } } } @@ -215,6 +219,7 @@ impl Receiver { api_key: config.api_key.clone(), tunnel: config.tunnel.clone(), instance_id: config.instance_id.clone(), + version: config.version.clone(), } .to_json(), );