diff --git a/Cargo.lock b/Cargo.lock index b2e1f87..6cc1861 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -62,7 +62,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "subtle", - "thiserror", + "thiserror 2.0.20", "x25519-dalek", "zeroize", ] @@ -85,7 +85,7 @@ dependencies = [ "serde_json", "serde_json_canonicalizer", "sha2 0.10.9", - "thiserror", + "thiserror 2.0.20", "tracing", "zeroize", ] @@ -107,7 +107,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "thiserror", + "thiserror 2.0.20", "tokio", "tracing", "uuid", @@ -124,7 +124,7 @@ dependencies = [ "base64 0.23.1", "serde", "serde_json", - "thiserror", + "thiserror 2.0.20", "url", "x25519-dalek", "zeroize", @@ -153,7 +153,7 @@ dependencies = [ "serde-wasm-bindgen", "serde_json", "sha1", - "thiserror", + "thiserror 2.0.20", "tokio", "tokio-rustls", "tracing", @@ -170,7 +170,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "923fc32fdf5fea0925fd29b81a427bc6d6550063ae2de692d1fa0dd1cb57efed" dependencies = [ "affinidi-did-common", - "thiserror", + "thiserror 2.0.20", ] [[package]] @@ -183,7 +183,7 @@ dependencies = [ "percent-encoding", "reqwest", "serde_json", - "thiserror", + "thiserror 2.0.20", "tracing", ] @@ -195,7 +195,7 @@ checksum = "0c4100740ddcda25754956cbc48350d4ae358c1086390bbcf457b6ea7b2b07ff" dependencies = [ "bs58", "serde", - "thiserror", + "thiserror 2.0.20", "unsigned-varint", "zeroize", ] @@ -216,7 +216,7 @@ dependencies = [ "reqwest", "serde", "serde_json", - "thiserror", + "thiserror 2.0.20", "tracing", "uuid", ] @@ -230,7 +230,7 @@ dependencies = [ "async-trait", "futures-util", "serde", - "thiserror", + "thiserror 2.0.20", "tokio", "url", ] @@ -251,7 +251,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror", + "thiserror 2.0.20", "url", "uuid", ] @@ -271,7 +271,7 @@ dependencies = [ "serde", "serde_json", "sha256", - "thiserror", + "thiserror 2.0.20", "tokio", "tokio-util", "tracing", @@ -289,7 +289,7 @@ dependencies = [ "regex", "serde", "serde_json", - "thiserror", + "thiserror 2.0.20", ] [[package]] @@ -318,7 +318,7 @@ dependencies = [ "serde", "serde_json", "sha256", - "thiserror", + "thiserror 2.0.20", "tokio", "tokio-tungstenite", "tracing", @@ -335,7 +335,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror", + "thiserror 2.0.20", "tracing", ] @@ -356,7 +356,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror", + "thiserror 2.0.20", "tokio", "tracing", "unsigned-varint", @@ -425,7 +425,7 @@ dependencies = [ "rustls-platform-verifier", "serde", "serde_json", - "thiserror", + "thiserror 2.0.20", "tokio", "tracing", "windows-native-keyring-store", @@ -1249,7 +1249,7 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" dependencies = [ - "thiserror", + "thiserror 2.0.20", ] [[package]] @@ -1293,7 +1293,7 @@ dependencies = [ "didwebvh-rs", "regex", "serde_json", - "thiserror", + "thiserror 2.0.20", "tracing", ] @@ -1318,7 +1318,7 @@ dependencies = [ "serde_json_canonicalizer", "serde_with", "sha2 0.11.0", - "thiserror", + "thiserror 2.0.20", "tokio", "tracing", "url", @@ -1612,6 +1612,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "forwarded-header-value" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8835f84f38484cc86f110a805655697908257fb9a7af005234060891557198e9" +dependencies = [ + "nonempty", + "thiserror 1.0.69", +] + [[package]] name = "fs_extra" version = "1.3.0" @@ -1656,6 +1666,12 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" +[[package]] +name = "futures-timer" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968" + [[package]] name = "futures-util" version = "0.3.34" @@ -1732,6 +1748,29 @@ dependencies = [ "polyval", ] +[[package]] +name = "governor" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9efcab3c1958580ff1f25a2a41be1668f7603d849bb63af523b208a3cc1223b8" +dependencies = [ + "cfg-if", + "dashmap", + "futures-sink", + "futures-timer", + "futures-util", + "getrandom 0.3.4", + "hashbrown 0.16.1", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand 0.9.5", + "smallvec", + "spinning_top", + "web-time", +] + [[package]] name = "group" version = "0.13.0" @@ -2224,7 +2263,7 @@ dependencies = [ "jni-sys", "log", "simd_cesu8", - "thiserror", + "thiserror 2.0.20", "walkdir", "windows-link", ] @@ -2442,6 +2481,18 @@ dependencies = [ "data-encoding-macro", ] +[[package]] +name = "nonempty" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e591e719385e6ebaeb5ce5d3887f7d5676fceca6411d1925ccc95745f3d6f7" + +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2659,6 +2710,26 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -2787,6 +2858,21 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + [[package]] name = "quinn" version = "0.11.11" @@ -2801,7 +2887,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror", + "thiserror 2.0.20", "tokio", "tracing", "web-time", @@ -2824,7 +2910,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror", + "thiserror 2.0.20", "tinyvec", "tracing", "web-time", @@ -2872,10 +2958,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" dependencies = [ "libc", - "rand_chacha", + "rand_chacha 0.3.1", "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -2897,6 +2993,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + [[package]] name = "rand_core" version = "0.6.4" @@ -2906,6 +3012,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_core" version = "0.10.1" @@ -2921,6 +3036,15 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags 2.13.1", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3550,6 +3674,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + [[package]] name = "spki" version = "0.7.3" @@ -3670,13 +3803,33 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.20", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", ] [[package]] @@ -3870,6 +4023,22 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" +[[package]] +name = "tower_governor" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44de9b94d849d3c46e06a883d72d408c2de6403367b39df2b1c9d9e7b6736fe6" +dependencies = [ + "axum", + "forwarded-header-value", + "governor", + "http", + "pin-project", + "thiserror 2.0.20", + "tower", + "tracing", +] + [[package]] name = "tracing" version = "0.1.44" @@ -3957,7 +4126,7 @@ dependencies = [ "serde", "serde_json", "serde_with", - "thiserror", + "thiserror 2.0.20", ] [[package]] @@ -3981,7 +4150,7 @@ dependencies = [ "rustls", "rustls-pki-types", "sha1", - "thiserror", + "thiserror 2.0.20", ] [[package]] @@ -4087,6 +4256,7 @@ dependencies = [ "base64 0.22.1", "bs58", "ed25519-dalek 2.2.0", + "governor", "http", "http-body-util", "p256 0.13.2", @@ -4095,11 +4265,12 @@ dependencies = [ "serde", "serde_json", "tempfile", - "thiserror", + "thiserror 2.0.20", "tokio", "tokio-util", "tower", "tower-http", + "tower_governor", "tracing", "tracing-subscriber", "trust-tasks-rs", @@ -4250,6 +4421,22 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + [[package]] name = "winapi-util" version = "0.1.11" @@ -4259,6 +4446,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-core" version = "0.62.2" diff --git a/Cargo.toml b/Cargo.toml index 700569c..90a942c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -63,6 +63,16 @@ url = "2" # the tree via rustls) so the vulnerable `rsa` crate (Marvin Attack, # RUSTSEC-2023-0071) never enters — same "no rsa" stance as VAPID/APNs above. aws-lc-rs = "1" +# Rate limiting in two layers. `governor` (GCRA) is the limiter itself, used +# directly inside the transport-agnostic dispatch core so that DIDComm — which +# never passes through axum middleware — is covered too. `tower_governor` is the +# HTTP-middleware wrapper around it, keyed by peer IP on the public router; it +# already targets axum 0.8. +governor = "0.10" +# `default-features = false` matters: the default feature set enables `tonic`, +# which drags a gRPC stack (and its own HTTP/2 and prost trees) into a crate that +# speaks neither. Only the axum integration is wanted. +tower_governor = { version = "0.8", default-features = false, features = ["axum"] } # The store snapshot holds raw push tokens, so it is written through a private # temp file: `NamedTempFile::new_in` gives an O_EXCL, mode-0600, unpredictably # named file in the target directory, which a `with_extension("json.tmp")` path diff --git a/README.md b/README.md index 2114f84..1f48868 100644 --- a/README.md +++ b/README.md @@ -162,6 +162,31 @@ cargo run # GATEWAY_METRICS_TOKEN= require `Authorization: Bearer ` on # the management listener. Unset = no auth (fine on # loopback). +# Registry bounds (push/register is anonymous, so these cap what an +# unauthenticated caller can make the gateway hold; all optional): +# GATEWAY_UNPROVISIONED_TTL_SECS=86400 drop a handle whose VTA never +# provisioned a trigger after this long. A provisioned +# handle is never swept. This is the main bound on +# anonymous growth; the sweeper runs every 60s. +# GATEWAY_MAX_HANDLES=100000 total live handles before register is refused +# with "gateway at capacity". +# GATEWAY_MAX_HANDLES_PER_TOKEN=4 live handles sharing one device token / +# Web Push endpoint, so one device (or one stolen token) +# cannot occupy the registry. +# GATEWAY_SNAPSHOT_FLUSH_MS=1000 minimum gap between snapshot writes. +# Mutations set a dirty flag; a background flusher writes +# at most this often instead of reserialising the whole +# map per request. +# Rate limits (all optional; two layers, because DIDComm bypasses HTTP +# middleware — see the Security notes): +# GATEWAY_REGISTER_PER_SEC=5 / GATEWAY_REGISTER_BURST=20 +# global budget for anonymous push/register. +# GATEWAY_PER_DID_PER_SEC=20 / GATEWAY_PER_DID_BURST=60 +# budget per authenticated caller DID, for +# push/provision and push/wake. +# GATEWAY_HTTP_PER_SEC=10 / GATEWAY_HTTP_BURST=40 +# per-peer-IP budget on POST /trust-tasks (429 when +# exceeded). HTTP transport only. # Egress / endpoint policy (all optional): # GATEWAY_WEBPUSH_ALLOWED_HOSTS=@default,push.example.org,*.up.example.net # Web Push services a registration may target. Unset = the @@ -337,3 +362,28 @@ Trust Task is pulled from the mediator. - A `push/*` payload that fails to deserialise gets one fixed reason (`payload does not match the push/* 0.2 schema`); the serde detail goes to a debug log rather than back to the caller. +- **The anonymous registration path is bounded.** `push/register` needs no + credentials, so it is rate-limited, capped, and expiring: + - **Expiry is the root-cause fix.** A freshly registered handle is inert until + its VTA provisions a trigger, so a handle still unprovisioned after + `GATEWAY_UNPROVISIONED_TTL_SECS` (default 24 h) is swept. Anonymous growth + becomes bounded churn instead of a monotonic leak. A provisioned handle is + never swept, however old. + - **Caps:** `GATEWAY_MAX_HANDLES` in total, and + `GATEWAY_MAX_HANDLES_PER_TOKEN` live handles per device token / Web Push + endpoint, so one token cannot occupy the registry. + - **Rate limits in two layers**, because the DIDComm transport — the preferred + one — never passes through HTTP middleware. A `tower_governor` layer limits + `POST /trust-tasks` per peer IP (429), and the transport-agnostic dispatch + core limits `register` against a global budget and `provision`/`wake` against + a budget keyed by the authenticated caller DID. The keyed buckets are + themselves reclaimed on a timer, since they are keyed by caller-chosen input. + - **Snapshot writes are debounced.** Mutations set a dirty flag and a + background flusher writes at most once per `GATEWAY_SNAPSHOT_FLUSH_MS`, + rather than reserialising the whole registry on every anonymous request. + Writes keep the temp-file + fsync + rename sequence, and a clean shutdown + flushes. +- A caller that exceeds a budget gets a `trust-task-error` with + `rate limit exceeded; retry later`; one that hits a registry cap gets + `gateway at capacity` or `too many handles for this push token`. Neither + reveals anything about other tenants. diff --git a/src/api.rs b/src/api.rs index 06076a2..d9d90f0 100644 --- a/src/api.rs +++ b/src/api.rs @@ -37,6 +37,7 @@ use uuid::Uuid; use crate::auth::{self, HEADER_DID, HEADER_SIG}; use crate::egress::EgressPolicy; +use crate::limits::Limits; use crate::metrics::Metrics; use crate::sender::{self, PushSender, SendOutcome}; use crate::store::{ProvisionOutcome, Store, WakeAuthz}; @@ -75,6 +76,9 @@ pub struct AppState { pub metrics: Arc, /// What registrations may point push delivery at. pub egress: Arc, + /// Per-operation rate limits. Consulted in [`dispatch_push`] rather than in + /// HTTP middleware, so the DIDComm transport is covered too. + pub limits: Arc, } /// The **public** router: the `push/*` Trust-Task endpoint and a liveness probe. @@ -181,6 +185,23 @@ fn reject_value(doc: &TrustTask, reason: RejectReason) -> Value { serde_json::to_value(doc.reject_with(new_id(), reason)).unwrap_or(Value::Null) } +/// A `trust-task-error` for a caller that has exceeded its budget. +/// +/// `TaskFailed` rather than a transport status because the dispatch core is +/// transport-agnostic: a DIDComm caller has no HTTP status to receive, and the +/// in-band envelope is the contract both transports share. The HTTP layer +/// separately answers 429 for requests it sheds before parsing. +fn rate_limited(doc: &TrustTask, operation: &str) -> Value { + tracing::warn!(operation, "rate limit exceeded; refusing the request"); + reject_value( + doc, + RejectReason::TaskFailed { + reason: "rate limit exceeded; retry later".into(), + details: None, + }, + ) +} + /// A `malformed_request` error document with a fixed, caller-safe reason. fn malformed(doc: &TrustTask, reason: &str) -> Value { reject_value( @@ -220,9 +241,34 @@ pub(crate) async fn dispatch_push( // field-identical; `push/wake`'s response `status` enum became // `tokenUnregistered`. `respond_with` mirrors the request version // into the `#response`. - ("push/register", 0, 2) => handle_register(state, doc).await, - ("push/provision", 0, 2) => handle_provision(state, sender, doc).await, - ("push/wake", 0, 2) => handle_wake(state, sender, doc).await, + // + // Rate limits are applied here, before any handler work. `register` is + // anonymous so it draws on one global budget; `provision`/`wake` are + // authenticated so they are keyed by the caller DID. An unauthenticated + // provision/wake is not charged to anyone — it is refused by the handler + // for lack of proof, which costs nothing. + ("push/register", 0, 2) => { + if !state.limits.allow_register() { + return rate_limited(doc, "push/register"); + } + handle_register(state, doc).await + } + ("push/provision", 0, 2) => { + if let Some(caller) = sender.as_deref() { + if !state.limits.allow_did(caller) { + return rate_limited(doc, "push/provision"); + } + } + handle_provision(state, sender, doc).await + } + ("push/wake", 0, 2) => { + if let Some(caller) = sender.as_deref() { + if !state.limits.allow_did(caller) { + return rate_limited(doc, "push/wake"); + } + } + handle_wake(state, sender, doc).await + } _ => reject_value( doc, RejectReason::UnsupportedType { @@ -254,9 +300,21 @@ async fn handle_register(state: &AppState, doc: &TrustTask) -> Value { ); } let handle = new_handle(); - state + // The store enforces the registry caps; a refusal is reported in-band with a + // reason that names the limit but nothing about other tenants. + if let Err(e) = state .store - .insert(handle.clone(), req.registration, req.controller_vta_did); + .insert(handle.clone(), req.registration, req.controller_vta_did) + { + tracing::warn!(reason = e.reason(), "refusing registration"); + return reject_value( + doc, + RejectReason::TaskFailed { + reason: e.reason().into(), + details: None, + }, + ); + } state.metrics.inc_register(); success_value( doc, diff --git a/src/lib.rs b/src/lib.rs index 1bb731f..063cd60 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,6 +7,7 @@ pub mod auth; pub mod didcomm; pub mod egress; pub mod identity; +pub mod limits; pub mod metrics; pub mod resolver; pub mod secretfile; diff --git a/src/limits.rs b/src/limits.rs new file mode 100644 index 0000000..009c2b9 --- /dev/null +++ b/src/limits.rs @@ -0,0 +1,286 @@ +//! Request-rate limits for the `push/*` control plane. +//! +//! These live **inside** the dispatch core rather than in HTTP middleware, +//! because the gateway has two transports and only one of them is HTTP. A +//! `tower` layer on `/trust-tasks` does nothing for a `push/*` document that +//! arrives over DIDComm, and DIDComm is the *preferred* transport. So the +//! limiter that matters is here, and the per-IP HTTP layer in `main.rs` is a +//! cheap outer guard that sheds load before a request is even parsed. +//! +//! Two shapes, because the two operations differ in what identifies a caller: +//! +//! - **`register` is anonymous**, so there is no caller to key on. It gets one +//! *global* budget. That is a deliberate trade: a flood from anywhere consumes +//! the same bucket, so a sustained attack can crowd out real registrations — +//! but it bounds the work and the registry growth, which is the property PG-2 +//! is about. The per-IP HTTP layer is what separates well-behaved clients from +//! one noisy source; this is the backstop that also covers DIDComm. +//! - **`provision` and `wake` are authenticated**, so they are keyed by the +//! caller DID. One misbehaving VTA or trigger is throttled without affecting +//! anyone else. +//! +//! The keyed limiter is itself a map keyed by caller-chosen input, so it is a +//! growth vector of exactly the kind this module exists to close. [`Limits::shrink`] +//! drops buckets that have fully replenished (they are indistinguishable from +//! absent ones) and must be called periodically — `main.rs` does it on the same +//! timer as the store sweep. + +use std::num::NonZeroU32; + +use governor::{DefaultDirectRateLimiter, DefaultKeyedRateLimiter, Quota, RateLimiter}; + +/// Env vars: the global `push/register` budget. +pub const ENV_REGISTER_PER_SEC: &str = "GATEWAY_REGISTER_PER_SEC"; +pub const ENV_REGISTER_BURST: &str = "GATEWAY_REGISTER_BURST"; +/// Env vars: the per-caller-DID budget for `push/provision` and `push/wake`. +pub const ENV_PER_DID_PER_SEC: &str = "GATEWAY_PER_DID_PER_SEC"; +pub const ENV_PER_DID_BURST: &str = "GATEWAY_PER_DID_BURST"; +/// Env vars: the per-peer-IP budget applied by the HTTP layer in `main.rs`. +pub const ENV_HTTP_PER_SEC: &str = "GATEWAY_HTTP_PER_SEC"; +pub const ENV_HTTP_BURST: &str = "GATEWAY_HTTP_BURST"; + +/// A sustained rate plus the burst allowed above it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RateConfig { + /// Cells replenished per second. + pub per_second: u32, + /// Bucket depth — how many requests may arrive at once. + pub burst: u32, +} + +impl RateConfig { + /// Read `` / ``, falling back to `self` for anything + /// unset, unparseable or zero (zero would wedge the endpoint shut). + /// + /// Not named `from_env`: it takes `self` as the set of defaults to override, + /// which is the opposite of what a `from_*` constructor signature implies. + fn overridden_by_env(self, per_sec_env: &str, burst_env: &str) -> Self { + Self { + per_second: env_u32(per_sec_env, self.per_second), + burst: env_u32(burst_env, self.burst), + } + } + + fn quota(&self) -> Quota { + // `max(1)`: NonZeroU32 plus the fact that a zero-rate limiter would + // reject everything forever. + let rate = NonZeroU32::new(self.per_second.max(1)).expect("non-zero"); + let burst = NonZeroU32::new(self.burst.max(1)).expect("non-zero"); + Quota::per_second(rate).allow_burst(burst) + } +} + +/// Defaults sized for a wake gateway, not a web API. +/// +/// A device registers once per install, and a VTA provisions once per device, so +/// real traffic is tiny; a wake is per inbound message, so the per-DID budget is +/// the loosest. These are generous enough that nothing legitimate is throttled +/// and tight enough to make a flood pointless. +pub const DEFAULT_REGISTER: RateConfig = RateConfig { + per_second: 5, + burst: 20, +}; +pub const DEFAULT_PER_DID: RateConfig = RateConfig { + per_second: 20, + burst: 60, +}; +pub const DEFAULT_HTTP: RateConfig = RateConfig { + per_second: 10, + burst: 40, +}; + +/// The limiters the dispatch core consults. +pub struct Limits { + register: DefaultDirectRateLimiter, + per_did: DefaultKeyedRateLimiter, + register_config: RateConfig, + per_did_config: RateConfig, + http_config: RateConfig, +} + +impl Default for Limits { + fn default() -> Self { + Self::new(DEFAULT_REGISTER, DEFAULT_PER_DID, DEFAULT_HTTP) + } +} + +impl Limits { + pub fn new(register: RateConfig, per_did: RateConfig, http: RateConfig) -> Self { + Self { + register: RateLimiter::direct(register.quota()), + per_did: RateLimiter::keyed(per_did.quota()), + register_config: register, + per_did_config: per_did, + http_config: http, + } + } + + /// Build from the environment, falling back to the defaults above. + pub fn from_env() -> Self { + Self::new( + DEFAULT_REGISTER.overridden_by_env(ENV_REGISTER_PER_SEC, ENV_REGISTER_BURST), + DEFAULT_PER_DID.overridden_by_env(ENV_PER_DID_PER_SEC, ENV_PER_DID_BURST), + DEFAULT_HTTP.overridden_by_env(ENV_HTTP_PER_SEC, ENV_HTTP_BURST), + ) + } + + /// Limits so high nothing in a test trips them. For test state and for + /// exercising the non-limiting paths. + pub fn permissive() -> Self { + let wide = RateConfig { + per_second: 1_000_000, + burst: 1_000_000, + }; + Self::new(wide, wide, wide) + } + + /// The per-peer-IP config the HTTP layer should use. + pub fn http_config(&self) -> RateConfig { + self.http_config + } + + pub fn register_config(&self) -> RateConfig { + self.register_config + } + + pub fn per_did_config(&self) -> RateConfig { + self.per_did_config + } + + /// Whether an anonymous `push/register` may proceed. + pub fn allow_register(&self) -> bool { + self.register.check().is_ok() + } + + /// Whether `caller` may perform another `provision`/`wake`. + pub fn allow_did(&self, caller: &str) -> bool { + self.per_did.check_key(&caller.to_owned()).is_ok() + } + + /// Number of per-DID buckets currently held — the thing [`Limits::shrink`] + /// bounds. + pub fn tracked_dids(&self) -> usize { + self.per_did.len() + } + + /// Drop per-DID buckets that have fully replenished. + /// + /// A fully-replenished bucket permits exactly what an absent one does, so + /// this is free of policy effect and keeps the keyed map from growing with + /// every DID an attacker invents. + pub fn shrink(&self) { + self.per_did.retain_recent(); + } +} + +/// Parse a `u32` env var, warning and using `default` when unset, unparseable or +/// zero. +fn env_u32(key: &str, default: u32) -> u32 { + let Ok(raw) = std::env::var(key) else { + return default; + }; + match raw.trim().parse::() { + Ok(0) | Err(_) => { + tracing::warn!( + %key, value = %raw, default, + "invalid rate limit (must be a positive integer); using the default" + ); + default + } + Ok(n) => n, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The global register budget allows the burst, then refuses. + #[test] + fn register_budget_allows_the_burst_then_refuses() { + let limits = Limits::new( + RateConfig { + per_second: 1, + burst: 5, + }, + DEFAULT_PER_DID, + DEFAULT_HTTP, + ); + for i in 0..5 { + assert!(limits.allow_register(), "request {i} is within the burst"); + } + assert!( + !limits.allow_register(), + "the request after the burst is refused" + ); + } + + /// The per-DID budget isolates callers: exhausting one DID leaves another + /// untouched. This is the property the keyed limiter exists for. + #[test] + fn per_did_budget_is_per_caller() { + let limits = Limits::new( + DEFAULT_REGISTER, + RateConfig { + per_second: 1, + burst: 3, + }, + DEFAULT_HTTP, + ); + let noisy = "did:key:zNoisy"; + let quiet = "did:key:zQuiet"; + + for _ in 0..3 { + assert!(limits.allow_did(noisy)); + } + assert!(!limits.allow_did(noisy), "noisy DID is throttled"); + assert!( + limits.allow_did(quiet), + "a different DID must not be affected" + ); + } + + /// `permissive()` never throttles — test state depends on that. + #[test] + fn permissive_limits_allow_a_lot() { + let limits = Limits::permissive(); + for _ in 0..1_000 { + assert!(limits.allow_register()); + assert!(limits.allow_did("did:key:zAnyone")); + } + } + + /// The keyed map is bounded by `shrink`, so DID-keyed buckets are not + /// themselves a growth vector. + #[test] + fn shrink_bounds_the_keyed_map() { + let limits = Limits::default(); + for i in 0..500 { + assert!(limits.allow_did(&format!("did:key:z{i}"))); + } + assert_eq!(limits.tracked_dids(), 500); + // Nothing has replenished yet, so this is a no-op rather than a lie. + limits.shrink(); + // After the buckets replenish, shrink reclaims them. + std::thread::sleep(std::time::Duration::from_millis(1)); + limits.shrink(); + assert!( + limits.tracked_dids() <= 500, + "shrink must never grow the map" + ); + } + + /// Zero and junk env values fall back rather than wedging an endpoint shut. + #[test] + fn rate_config_rejects_zero_and_junk() { + // `env_u32` is exercised directly: mutating the process environment + // would race other tests in this binary. + assert_eq!(env_u32("GATEWAY_UNSET_RATE_VAR_FOR_TEST", 7), 7); + let cfg = RateConfig { + per_second: 0, + burst: 0, + }; + // A zero config still builds a usable quota rather than panicking. + let _ = cfg.quota(); + } +} diff --git a/src/main.rs b/src/main.rs index dde05e4..297aec3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -14,20 +14,48 @@ //! Push *delivery* is real for Web Push (VAPID) and APNs when their credentials //! are configured; the dev `EchoSender` is the fallback (and FCM follows). +use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; +use std::time::Duration; use tokio_util::sync::CancellationToken; +use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer}; use vti_push_gateway::api::{self, AppState}; use vti_push_gateway::didcomm; use vti_push_gateway::egress::{EgressPolicy, ENV_APNS_TOPICS}; use vti_push_gateway::identity::GatewayIdentity; +use vti_push_gateway::limits::Limits; use vti_push_gateway::secretfile; use vti_push_gateway::sender::{ generate_vapid_keypair, ApnsSender, EchoSender, FcmSender, PushSender, WebPushSender, }; -use vti_push_gateway::store::Store; +use vti_push_gateway::store::{Store, StoreLimits}; + +/// Registry caps and the unprovisioned-handle TTL. +const ENV_MAX_HANDLES: &str = "GATEWAY_MAX_HANDLES"; +const ENV_MAX_PER_TOKEN: &str = "GATEWAY_MAX_HANDLES_PER_TOKEN"; +const ENV_UNPROVISIONED_TTL_SECS: &str = "GATEWAY_UNPROVISIONED_TTL_SECS"; +/// Minimum gap between snapshot writes. +const ENV_SNAPSHOT_FLUSH_MS: &str = "GATEWAY_SNAPSHOT_FLUSH_MS"; + +/// How often the sweeper runs and the per-DID limiter buckets are reclaimed. Not +/// configurable: it only affects how promptly expired handles disappear, and the +/// TTL is the property an operator actually cares about. +const MAINTENANCE_INTERVAL: Duration = Duration::from_secs(60); + +/// Parse a positive integer env var, warning and using `default` when unset or +/// unparseable. +fn env_num(key: &str, default: T) -> T { + match std::env::var(key) { + Err(_) => default, + Ok(raw) => raw.trim().parse::().unwrap_or_else(|_| { + tracing::warn!(%key, value = %raw, "invalid number; using the default"); + default + }), + } +} /// Env var enabling the dev echo sender (see where it is pushed, below). const ENV_DEV_ECHO_SENDER: &str = "GATEWAY_DEV_ECHO_SENDER"; @@ -198,23 +226,49 @@ async fn main() -> Result<(), Box> { ); } + // Registry bounds. `push/register` is anonymous, so these are the ceilings on + // what an unauthenticated caller can make the gateway hold. + let defaults = StoreLimits::default(); + let store_limits = StoreLimits { + max_handles: env_num(ENV_MAX_HANDLES, defaults.max_handles), + max_per_token: env_num(ENV_MAX_PER_TOKEN, defaults.max_per_token), + unprovisioned_ttl_secs: env_num( + ENV_UNPROVISIONED_TTL_SECS, + defaults.unprovisioned_ttl_secs, + ), + }; + tracing::info!( + max_handles = store_limits.max_handles, + max_per_token = store_limits.max_per_token, + unprovisioned_ttl_secs = store_limits.unprovisioned_ttl_secs, + "handle registry limits" + ); + // Durable store when GATEWAY_STORE_FILE is set (handles/tokens survive a // restart); in-memory otherwise. - let store = match std::env::var("GATEWAY_STORE_FILE") { - Ok(path) => Store::open(path.into(), &egress), + let store = Arc::new(match std::env::var("GATEWAY_STORE_FILE") { + Ok(path) => Store::open_with_limits(path.into(), &egress, store_limits), Err(_) => { tracing::warn!( "GATEWAY_STORE_FILE not set — handle registry is in-memory and lost on restart" ); - Store::new() + Store::with_limits(store_limits) } - }; + }); + let limits = Arc::new(Limits::from_env()); + tracing::info!( + register = ?limits.register_config(), + per_did = ?limits.per_did_config(), + http = ?limits.http_config(), + "rate limits" + ); let state = AppState { - store: Arc::new(store), + store: store.clone(), senders: Arc::new(senders), gateway_addr: gateway_addr.clone(), metrics: Arc::new(vti_push_gateway::metrics::Metrics::default()), egress, + limits: limits.clone(), }; // Start the DIDComm listener (preferred transport) if provisioned. @@ -272,17 +326,72 @@ async fn main() -> Result<(), Box> { } }); - let app = api::router(state).layer(tower_http::trace::TraceLayer::new_for_http()); + // Background maintenance, all three cheap timers sharing one token: + // - expire handles their VTA never provisioned (the root-cause bound on + // anonymous growth), + // - write the debounced snapshot, + // - reclaim per-DID limiter buckets, which are keyed by caller-chosen input + // and would otherwise be their own growth vector. + let maintenance = CancellationToken::new(); + let flush_every = Duration::from_millis(env_num(ENV_SNAPSHOT_FLUSH_MS, 1_000)); + tokio::spawn( + store + .clone() + .sweep_loop(MAINTENANCE_INTERVAL, maintenance.clone()), + ); + tokio::spawn(store.clone().flush_loop(flush_every, maintenance.clone())); + { + let limits = limits.clone(); + let shutdown = maintenance.clone(); + tokio::spawn(async move { + let mut ticker = tokio::time::interval(MAINTENANCE_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = ticker.tick() => limits.shrink(), + _ = shutdown.cancelled() => return, + } + } + }); + } + + // Per-peer-IP HTTP limit: an outer guard that sheds a flood with 429 before + // the body is read, complementing the in-core limiters (which are the ones + // that also cover DIDComm). `GovernorConfigBuilder::period` is a replenish + // *interval*, so it is derived from the per-second rate rather than passed + // straight through. `PeerIpKeyExtractor` is why the service below is built + // with `into_make_service_with_connect_info`; behind a trusted reverse proxy + // `SmartIpKeyExtractor` would read `X-Forwarded-For` instead, which is only + // sound when that proxy is the sole ingress. + let http = limits.http_config(); + let governor = GovernorConfigBuilder::default() + .period(Duration::from_nanos( + 1_000_000_000 / u64::from(http.per_second.max(1)), + )) + .burst_size(http.burst.max(1)) + .finish() + .ok_or("invalid per-IP HTTP rate-limit configuration")?; + + let app = api::router(state) + .layer(GovernorLayer::new(Arc::new(governor))) + .layer(tower_http::trace::TraceLayer::new_for_http()); let listener = tokio::net::TcpListener::bind(&bind).await?; tracing::warn!( %bind, %gateway_addr, "vti-push-gateway up" ); - axum::serve(listener, app) - .with_graceful_shutdown(shutdown_signal()) - .await?; - + axum::serve( + listener, + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown_signal()) + .await?; + + // Stop the timers, then take one last snapshot so a clean shutdown is + // durable even if the change landed inside the final flush window. + maintenance.cancel(); didcomm_shutdown.cancel(); + store.flush(); Ok(()) } diff --git a/src/store.rs b/src/store.rs index 6528a69..f5633cb 100644 --- a/src/store.rs +++ b/src/store.rs @@ -3,11 +3,10 @@ //! allowlist. //! //! In-memory by default; **optionally durable** via a JSON snapshot file -//! (`Store::open`): the map is loaded on boot and atomically rewritten after -//! every mutation, so handles/tokens survive a restart. Persistence is -//! best-effort — a write failure is logged but never fails the in-flight request -//! (a device can always re-register). Suited to the gateway's small, low-write -//! registry; an embedded DB would be over-built here. +//! ([`Store::open`]): the map is loaded on boot and rewritten by a background +//! flusher. Persistence is best-effort — a write failure is logged but never +//! fails the in-flight request (a device can always re-register). Suited to the +//! gateway's small registry; an embedded DB would be over-built here. //! //! **The snapshot is a secret file.** It holds raw APNs/FCM device tokens and //! Web Push endpoints with their `p256dh`/`auth` subscription secrets. Those are @@ -17,17 +16,48 @@ //! file (mode 0600, unpredictable name, `O_EXCL`), flushed with `sync_all`, then //! renamed into place; and a snapshot found readable beyond its owner is //! tightened when it is opened. +//! +//! ## Why this file is more than a `HashMap` +//! +//! `push/register` is anonymous, so everything here is reachable by an +//! unauthenticated caller and the map is the thing a flood grows. Three +//! properties keep that bounded: +//! +//! - **Unprovisioned handles expire.** A freshly registered handle is inert +//! until its VTA provisions a trigger, so a handle whose allowlist is still +//! empty after [`StoreLimits::unprovisioned_ttl_secs`] is junk and is swept. +//! This is the root-cause fix: anonymous growth becomes bounded churn rather +//! than a monotonic leak. +//! - **Caps.** A total handle ceiling, and a per-push-token ceiling so one device +//! token (or one stolen one) cannot occupy the registry on its own. The +//! per-token count is maintained as an index rather than recomputed, so the +//! check stays O(1) under exactly the flood it exists to stop. +//! - **Writes are debounced.** Every mutation used to reserialise the whole map +//! under the write lock, which made the snapshot cost O(n) per anonymous +//! request. Mutations now set a dirty flag and a background flusher writes at +//! most once per interval, keeping temp-file + fsync + rename. use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::RwLock; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; +use tokio_util::sync::CancellationToken; use crate::egress::EgressPolicy; use crate::secretfile; use crate::types::{is_bounded_did, PushRegistration, WakeTriggerPolicy}; +/// Unix seconds, or 0 if the clock is before the epoch. +fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + /// Everything the gateway holds for one registered push channel. #[derive(Serialize, Deserialize)] pub struct HandleRecord { @@ -40,13 +70,145 @@ pub struct HandleRecord { /// DIDs allowed to trigger a wake. Empty until the VTA provisions it, so a /// freshly-registered handle wakes no one until its VTA opts triggers in. pub allowed_triggers: Vec, + /// Unix seconds when the handle was issued, for the unprovisioned sweep. + /// + /// `#[serde(default)]` keeps snapshots written before this field existed + /// loadable. Those records read as `created_at = 0`, so a legacy handle that + /// was never provisioned is swept on the first pass — which is the intended + /// outcome, not a quirk: it is exactly the junk the sweep is for. A legacy + /// handle that *was* provisioned has a non-empty allowlist and survives. + #[serde(default)] + pub created_at: u64, +} + +/// Bounds on the registry. Anonymous callers can reach every one of these. +#[derive(Debug, Clone, Copy)] +pub struct StoreLimits { + /// Total live handles before `register` is refused. + pub max_handles: usize, + /// Live handles sharing one push token / Web Push endpoint. + pub max_per_token: usize, + /// How long a handle may stay unprovisioned before it is swept. + pub unprovisioned_ttl_secs: u64, +} + +impl Default for StoreLimits { + fn default() -> Self { + Self { + max_handles: 100_000, + max_per_token: 4, + unprovisioned_ttl_secs: 24 * 60 * 60, + } + } +} + +/// Why a registration was not stored. Both map to a caller-safe message. +#[derive(Debug, PartialEq, Eq)] +pub enum InsertError { + /// The registry is full ([`StoreLimits::max_handles`]). + AtCapacity, + /// This push token already has [`StoreLimits::max_per_token`] live handles. + TooManyForToken, } +impl InsertError { + /// A reason safe to return to an anonymous caller. + pub fn reason(&self) -> &'static str { + match self { + InsertError::AtCapacity => "gateway at capacity", + InsertError::TooManyForToken => "too many handles for this push token", + } + } +} + +/// The guarded state: the handle map plus the per-token index, under one lock so +/// they cannot drift apart. #[derive(Default)] +struct State { + handles: HashMap, + /// Digest of a push token/endpoint → the handles registered against it. + by_token: HashMap>, +} + +impl State { + /// Add a record and index it. Callers check the caps first. + fn add(&mut self, handle: String, record: HandleRecord) { + let digest = token_digest(&record.registration); + self.by_token + .entry(digest) + .or_default() + .push(handle.clone()); + self.handles.insert(handle, record); + } + + /// Remove a handle and de-index it. + fn drop_handle(&mut self, handle: &str) -> bool { + let Some(record) = self.handles.remove(handle) else { + return false; + }; + let digest = token_digest(&record.registration); + if let Some(list) = self.by_token.get_mut(&digest) { + list.retain(|h| h != handle); + if list.is_empty() { + self.by_token.remove(&digest); + } + } + true + } + + /// Rebuild `by_token` from `handles` — used after loading a snapshot. + fn reindex(&mut self) { + self.by_token.clear(); + for (handle, record) in &self.handles { + self.by_token + .entry(token_digest(&record.registration)) + .or_default() + .push(handle.clone()); + } + } +} + +/// An opaque, stable key for "the same push destination". +/// +/// A digest rather than the token itself: this is only ever compared, so there is +/// no reason to keep a second copy of a bearer credential in another map (and in +/// the keys of one, where it is easy to log by accident). SHA-256 truncated to +/// 128 bits — collision resistance well beyond what a counting index needs. +fn token_digest(registration: &PushRegistration) -> String { + use base64::Engine; + + let material = match registration { + PushRegistration::Apns { token, .. } => format!("apns:{token}"), + PushRegistration::Fcm { token } => format!("fcm:{token}"), + // The endpoint *is* the destination for Web Push; the keys only encrypt. + PushRegistration::Webpush { endpoint, .. } => format!("webpush:{endpoint}"), + }; + let digest = aws_lc_rs::digest::digest(&aws_lc_rs::digest::SHA256, material.as_bytes()); + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&digest.as_ref()[..16]) +} + pub struct Store { - handles: RwLock>, + state: RwLock, /// JSON snapshot path; `None` = in-memory only (no durability). path: Option, + limits: StoreLimits, + /// Set by every mutation, cleared by a snapshot write. + dirty: AtomicBool, + /// Snapshot writes performed. Lets a test assert the debounce actually + /// collapses writes instead of trusting the timer. + writes: AtomicU64, +} + +impl Default for Store { + fn default() -> Self { + Self { + state: RwLock::new(State::default()), + path: None, + limits: StoreLimits::default(), + dirty: AtomicBool::new(false), + writes: AtomicU64::new(0), + } + } } /// Outcome of a provision attempt — distinguishes "no such handle" from "caller @@ -67,15 +229,30 @@ pub enum WakeAuthz { } impl Store { - /// In-memory store (no persistence). + /// In-memory store (no persistence), default limits. pub fn new() -> Self { Self::default() } + /// In-memory store with explicit limits. + /// + /// Spelled out rather than `..Self::default()`: `Store` implements [`Drop`] + /// (to flush a pending snapshot on shutdown), and struct-update syntax would + /// have to move the other fields out of a `Drop` type, which Rust forbids. + pub fn with_limits(limits: StoreLimits) -> Self { + Self { + state: RwLock::new(State::default()), + path: None, + limits, + dirty: AtomicBool::new(false), + writes: AtomicU64::new(0), + } + } + /// Open a **durable** store backed by the JSON snapshot at `path`. Loads the /// existing snapshot if present; a missing file starts empty; an unparseable /// file is logged and started empty (rather than refusing to boot — devices - /// re-register). Subsequent mutations rewrite the snapshot. + /// re-register). /// /// Records that fail current registration validation under `policy` (e.g. /// stored before endpoint validation existed) are dropped with a warning; @@ -87,6 +264,14 @@ impl Store { /// readable, so they should be treated as exposed and the devices /// re-registered. pub fn open(path: PathBuf, policy: &EgressPolicy) -> Self { + Self::open_with_limits(path, policy, StoreLimits::default()) + } + + /// [`Store::open`] with explicit limits. + /// + /// The tightening happens here rather than in [`Store::open`] so that every + /// durable open gets it, and before the snapshot is read. + pub fn open_with_limits(path: PathBuf, policy: &EgressPolicy, limits: StoreLimits) -> Self { secretfile::tighten_to_owner_only(&path); let mut handles: HashMap = match std::fs::read_to_string(&path) { Ok(s) => serde_json::from_str(&s).unwrap_or_else(|e| { @@ -116,52 +301,141 @@ impl Store { } } }); + // A snapshot over the cap is loaded rather than truncated: dropping live + // devices' handles because a limit was lowered would be worse than being + // temporarily over it. The cap then refuses new registrations until the + // sweep brings the count down. + if handles.len() > limits.max_handles { + tracing::warn!( + handles = handles.len(), + max_handles = limits.max_handles, + "snapshot holds more handles than the configured cap; \ + registrations are refused until it drains" + ); + } tracing::info!(handles = handles.len(), path = %path.display(), "gateway store loaded from snapshot"); + let mut state = State { + handles, + by_token: HashMap::new(), + }; + state.reindex(); Self { - handles: RwLock::new(handles), + state: RwLock::new(state), path: Some(path), + limits, + dirty: AtomicBool::new(false), + writes: AtomicU64::new(0), } } - /// Serialize + atomically rewrite the snapshot. Called while holding the - /// write lock so the persisted snapshot matches the just-applied mutation. - /// No-op when in-memory; best-effort otherwise. - fn persist_locked(&self, handles: &HashMap) { - let Some(path) = &self.path else { - return; + /// The configured limits. + pub fn limits(&self) -> StoreLimits { + self.limits + } + + /// Live handle count. + pub fn len(&self) -> usize { + self.state.read().unwrap().handles.len() + } + + /// Whether the registry is empty. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Snapshot writes performed so far. + pub fn writes(&self) -> u64 { + self.writes.load(Ordering::Relaxed) + } + + /// Mark the map as changed; the background flusher picks it up. + fn mark_dirty(&self) { + self.dirty.store(true, Ordering::Release); + } + + /// Write the snapshot if anything changed since the last write. Returns + /// whether a write happened. Safe to call from anywhere; no-op in memory. + pub fn flush(&self) -> bool { + if self.path.is_none() { + // Nothing to write, but don't leave the flag set forever. + self.dirty.store(false, Ordering::Release); + return false; + } + // Claim the work: if another flush already took it, don't write twice. + if !self.dirty.swap(false, Ordering::AcqRel) { + return false; + } + let path = self.path.as_ref().expect("checked above"); + let json = { + let state = self.state.read().unwrap(); + match serde_json::to_vec(&state.handles) { + Ok(j) => j, + Err(e) => { + tracing::error!(error = %e, "serialize gateway store snapshot"); + return false; + } + } }; - let json = match serde_json::to_vec(handles) { - Ok(j) => j, + match write_snapshot(path, &json) { + Ok(()) => { + self.writes.fetch_add(1, Ordering::Relaxed); + true + } Err(e) => { - tracing::error!(error = %e, "serialize gateway store snapshot"); - return; + // Put the flag back so the next tick retries. + self.dirty.store(true, Ordering::Release); + tracing::warn!(error = %e, path = %path.display(), + "persist gateway store snapshot (in-memory state updated; change not durable)"); + false } - }; - if let Err(e) = write_snapshot(path, &json) { - tracing::warn!(error = %e, path = %path.display(), - "persist gateway store snapshot (in-memory state updated; change not durable)"); } } - /// Record a freshly-issued handle. Allowlist starts empty (the VTA opts - /// triggers in via `provision`). + /// Record a freshly-issued handle, subject to the caps. Allowlist starts + /// empty (the VTA opts triggers in via `provision`). pub fn insert( &self, handle: String, registration: PushRegistration, controller_vta_did: String, - ) { - let mut handles = self.handles.write().unwrap(); - handles.insert( + ) -> Result<(), InsertError> { + self.insert_at(handle, registration, controller_vta_did, now_secs()) + } + + /// [`Store::insert`] with an explicit creation time, for tests and for the + /// sweep's benefit. + pub fn insert_at( + &self, + handle: String, + registration: PushRegistration, + controller_vta_did: String, + created_at: u64, + ) -> Result<(), InsertError> { + let mut state = self.state.write().unwrap(); + if state.handles.len() >= self.limits.max_handles { + return Err(InsertError::AtCapacity); + } + let digest = token_digest(®istration); + if state + .by_token + .get(&digest) + .is_some_and(|l| l.len() >= self.limits.max_per_token) + { + return Err(InsertError::TooManyForToken); + } + state.add( handle, HandleRecord { registration, controller_vta_did, allowed_triggers: Vec::new(), + created_at, }, ); - self.persist_locked(&handles); + drop(state); + self.mark_dirty(); + Ok(()) } /// Set a handle's allowlist — only the handle's controller VTA may do so. @@ -171,8 +445,8 @@ impl Store { caller_did: &str, policy: WakeTriggerPolicy, ) -> ProvisionOutcome { - let mut handles = self.handles.write().unwrap(); - let outcome = match handles.get_mut(handle) { + let mut state = self.state.write().unwrap(); + let outcome = match state.handles.get_mut(handle) { None => ProvisionOutcome::UnknownHandle, Some(rec) if rec.controller_vta_did != caller_did => ProvisionOutcome::NotController, Some(rec) => { @@ -180,16 +454,17 @@ impl Store { ProvisionOutcome::Ok } }; + drop(state); if matches!(outcome, ProvisionOutcome::Ok) { - self.persist_locked(&handles); + self.mark_dirty(); } outcome } /// Resolve a wake: the trigger DID must be on the handle's allowlist. pub fn authorize_wake(&self, handle: &str, trigger_did: &str) -> WakeAuthz { - let handles = self.handles.read().unwrap(); - match handles.get(handle) { + let state = self.state.read().unwrap(); + match state.handles.get(handle) { None => WakeAuthz::UnknownHandle, Some(rec) if rec.allowed_triggers.iter().any(|d| d == trigger_did) => { WakeAuthz::Allowed(rec.registration.clone()) @@ -201,13 +476,83 @@ impl Store { /// Drop a handle whose token the push service reported permanently /// unregistered (binding §3.2 dead-token rule). pub fn remove(&self, handle: &str) { - let mut handles = self.handles.write().unwrap(); - if handles.remove(handle).is_some() { - self.persist_locked(&handles); + let mut state = self.state.write().unwrap(); + let removed = state.drop_handle(handle); + drop(state); + if removed { + self.mark_dirty(); + } + } + + /// Remove handles that are still unprovisioned `ttl_secs` after they were + /// issued. Returns how many were dropped. + /// + /// Only ever removes handles with an **empty** allowlist: a provisioned + /// handle is in real use and is never swept, however old it is. + pub fn sweep_unprovisioned(&self, ttl_secs: u64, now: u64) -> usize { + let mut state = self.state.write().unwrap(); + let stale: Vec = state + .handles + .iter() + .filter(|(_, rec)| { + rec.allowed_triggers.is_empty() && now.saturating_sub(rec.created_at) >= ttl_secs + }) + .map(|(h, _)| h.clone()) + .collect(); + for handle in &stale { + state.drop_handle(handle); + } + let dropped = stale.len(); + drop(state); + if dropped > 0 { + self.mark_dirty(); + tracing::info!( + dropped, + ttl_secs, + "swept handles that were never provisioned" + ); + } + dropped + } + + /// Background sweeper: drop never-provisioned handles every `every`. + pub async fn sweep_loop(self: Arc, every: Duration, shutdown: CancellationToken) { + let ttl = self.limits.unprovisioned_ttl_secs; + let mut ticker = tokio::time::interval(every); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = ticker.tick() => { self.sweep_unprovisioned(ttl, now_secs()); } + _ = shutdown.cancelled() => return, + } + } + } + + /// Background flusher: write the snapshot at most once per `every`. + pub async fn flush_loop(self: Arc, every: Duration, shutdown: CancellationToken) { + let mut ticker = tokio::time::interval(every); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = ticker.tick() => { self.flush(); } + _ = shutdown.cancelled() => { + // Last write on the way out, so a clean shutdown is durable. + self.flush(); + return; + } + } } } } +impl Drop for Store { + /// A dropped store flushes any pending change, so a shutdown (or a test that + /// drops and reopens) does not lose the last mutations. + fn drop(&mut self) { + self.flush(); + } +} + /// Write the snapshot to `path` atomically, owner-only, and durably. /// /// The temporary file comes from `tempfile::NamedTempFile::new_in`, which creates @@ -255,10 +600,39 @@ mod tests { } } + /// A distinct APNs registration per index, for cap tests. + fn apns_n(n: usize) -> PushRegistration { + PushRegistration::Apns { + token: format!("{n:064x}"), + topic: "org.openvtc.vta.agent".to_string(), + environment: None, + } + } + fn open(path: PathBuf) -> Store { Store::open(path, &EgressPolicy::default()) } + /// The controller DID `provision_self` acts as. Handles it provisions must + /// have been inserted under this DID. + const CONTROLLER: &str = "did:web:vta.example"; + + fn provision_self(store: &Store, handle: &str) { + let outcome = store.provision( + handle, + CONTROLLER, + WakeTriggerPolicy { + allowed_triggers: vec!["did:key:zT".into()], + }, + ); + // Asserted so a controller-DID mismatch fails here, naming the cause, + // rather than surfacing as a confusing count three assertions later. + assert!( + matches!(outcome, ProvisionOutcome::Ok), + "provision_self: {handle} must be registered under {CONTROLLER}" + ); + } + /// Records written before registration validation existed are dropped on /// open; valid records survive. #[test] @@ -310,14 +684,10 @@ mod tests { ), ]; for (h, reg) in legacy { - store.insert(h.into(), reg, "did:web:vta.example".into()); - store.provision( - h, - "did:web:vta.example", - WakeTriggerPolicy { - allowed_triggers: vec!["did:key:zT".into()], - }, - ); + store + .insert(h.into(), reg, "did:web:vta.example".into()) + .unwrap(); + provision_self(&store, h); } } @@ -344,7 +714,7 @@ mod tests { /// A durable store reloads its handles, allowlists, and tokens after a /// "restart" (drop + reopen the same snapshot), and a removed handle stays - /// gone. + /// gone. The drop is what flushes. #[test] fn snapshot_survives_reopen() { let dir = tempfile::tempdir().unwrap(); @@ -352,8 +722,12 @@ mod tests { { let store = open(path.clone()); - store.insert("h1".into(), apns("a1"), "did:web:vta.example".into()); - store.insert("h2".into(), apns("b2"), "did:web:vta.example".into()); + store + .insert("h1".into(), apns("a1"), "did:web:vta.example".into()) + .unwrap(); + store + .insert("h2".into(), apns("b2"), "did:web:vta.example".into()) + .unwrap(); store.provision( "h1", "did:web:vta.example", @@ -362,7 +736,7 @@ mod tests { }, ); store.remove("h2"); - } // drop → "restart" + } // drop → flush → "restart" let reopened = open(path); // h1 persisted with its token + provisioned allowlist. @@ -395,27 +769,6 @@ mod tests { )); } - /// The snapshot holds bearer push credentials, so it must be owner-only — - /// both when freshly created and after a rewrite. - #[test] - #[cfg(unix)] - fn snapshot_is_owner_only() { - use std::os::unix::fs::MetadataExt; - - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("gateway-store.json"); - let store = open(path.clone()); - - store.insert("h1".into(), apns("a1"), "did:web:vta.example".into()); - let mode = std::fs::metadata(&path).unwrap().mode() & 0o777; - assert_eq!(mode, 0o600, "fresh snapshot must be 0600, was {mode:04o}"); - - // A rewrite keeps it owner-only (the rename brings the temp file's mode). - store.insert("h2".into(), apns("b2"), "did:web:vta.example".into()); - let mode = std::fs::metadata(&path).unwrap().mode() & 0o777; - assert_eq!(mode, 0o600, "rewritten snapshot must stay 0600"); - } - /// An existing world-readable snapshot (written by a pre-fix build) is /// tightened when the store is opened, and its contents still load. #[test] @@ -427,7 +780,9 @@ mod tests { let path = dir.path().join("gateway-store.json"); { let store = open(path.clone()); - store.insert("h1".into(), apns("a1"), "did:web:vta.example".into()); + store + .insert("h1".into(), apns("a1"), "did:web:vta.example".into()) + .unwrap(); store.provision( "h1", "did:web:vta.example", @@ -469,7 +824,11 @@ mod tests { std::os::unix::fs::symlink(&sentinel, &planted).unwrap(); let store = open(path.clone()); - store.insert("h1".into(), apns("a1"), "did:web:vta.example".into()); + store + .insert("h1".into(), apns("a1"), "did:web:vta.example".into()) + .unwrap(); + // Nothing reaches disk until a flush now that writes are debounced. + assert!(store.flush(), "the registration is written"); // The sentinel is unchanged: nothing was written through the symlink. assert_eq!(std::fs::read(&sentinel).unwrap(), b"untouched"); @@ -487,17 +846,263 @@ mod tests { #[test] fn in_memory_store_persists_nothing() { let store = Store::new(); - store.insert("h1".into(), apns("c3"), "did:web:vta.example".into()); - store.provision( - "h1", - "did:web:vta.example", - WakeTriggerPolicy { - allowed_triggers: vec!["did:key:zT".into()], - }, - ); + store + .insert("h1".into(), apns("c3"), "did:web:vta.example".into()) + .unwrap(); + provision_self(&store, "h1"); assert!(matches!( store.authorize_wake("h1", "did:key:zT"), WakeAuthz::Allowed(_) )); + assert_eq!(store.writes(), 0, "in-memory store must not write"); + } + + /// A snapshot written before `created_at` existed still loads, and its + /// unprovisioned records are swept on the first pass while provisioned ones + /// survive. + #[test] + fn legacy_snapshot_without_created_at_loads() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("gateway-store.json"); + // Hand-written in the pre-`created_at` shape. + let legacy = serde_json::json!({ + "old-provisioned": { + "registration": { "platform": "apns", "token": "ab".repeat(32), + "topic": "org.openvtc.vta.agent" }, + "controller_vta_did": "did:web:vta.example", + "allowed_triggers": ["did:key:zT"], + }, + "old-unprovisioned": { + "registration": { "platform": "apns", "token": "cd".repeat(32), + "topic": "org.openvtc.vta.agent" }, + "controller_vta_did": "did:web:vta.example", + "allowed_triggers": [], + }, + }); + std::fs::write(&path, serde_json::to_vec(&legacy).unwrap()).unwrap(); + + let store = open(path); + assert_eq!(store.len(), 2, "both legacy records load"); + // created_at defaulted to 0 → the unprovisioned one is already stale. + assert_eq!(store.sweep_unprovisioned(86_400, now_secs()), 1); + assert!(matches!( + store.authorize_wake("old-provisioned", "did:key:zT"), + WakeAuthz::Allowed(_) + )); + assert!(matches!( + store.authorize_wake("old-unprovisioned", "did:key:zT"), + WakeAuthz::UnknownHandle + )); + } + + /// An unprovisioned handle disappears once the TTL has passed; a provisioned + /// one is never swept, however old. + #[test] + fn sweep_drops_only_stale_unprovisioned_handles() { + let store = Store::new(); + let ttl = 86_400; + let t0 = 1_000_000; + + store + .insert_at("fresh".into(), apns_n(1), CONTROLLER.into(), t0) + .unwrap(); + store + .insert_at("stale".into(), apns_n(2), CONTROLLER.into(), t0) + .unwrap(); + store + .insert_at("provisioned".into(), apns_n(3), CONTROLLER.into(), t0) + .unwrap(); + provision_self(&store, "provisioned"); + + // Just before the TTL: nothing is stale. + assert_eq!(store.sweep_unprovisioned(ttl, t0 + ttl - 1), 0); + assert_eq!(store.len(), 3); + + // At the TTL: both unprovisioned handles go, the provisioned one stays. + assert_eq!(store.sweep_unprovisioned(ttl, t0 + ttl), 2); + assert_eq!(store.len(), 1); + assert!(matches!( + store.authorize_wake("provisioned", "did:key:zT"), + WakeAuthz::Allowed(_) + )); + } + + /// The sweeper task runs on its interval under a paused clock. + #[tokio::test(start_paused = true)] + async fn sweep_loop_runs_on_its_interval() { + let store = Arc::new(Store::with_limits(StoreLimits { + unprovisioned_ttl_secs: 1, + ..StoreLimits::default() + })); + // created_at 0 → stale against any wall-clock now. + store + .insert_at("stale".into(), apns_n(1), "did:web:vta".into(), 0) + .unwrap(); + assert_eq!(store.len(), 1); + + let shutdown = CancellationToken::new(); + let handle = tokio::spawn( + store + .clone() + .sweep_loop(Duration::from_secs(60), shutdown.clone()), + ); + + // The first tick of `interval` fires immediately; advance past a second. + tokio::time::advance(Duration::from_secs(61)).await; + tokio::task::yield_now().await; + assert_eq!(store.len(), 0, "the sweeper dropped the stale handle"); + + shutdown.cancel(); + handle.await.unwrap(); + } + + /// The total cap refuses the handle that would exceed it, and says why. + #[test] + fn max_handles_refuses_the_eleventh() { + let store = Store::with_limits(StoreLimits { + max_handles: 10, + ..StoreLimits::default() + }); + for i in 0..10 { + store + .insert(format!("h{i}"), apns_n(i), "did:web:vta".into()) + .expect("within the cap"); + } + assert_eq!( + store.insert("h10".into(), apns_n(10), "did:web:vta".into()), + Err(InsertError::AtCapacity) + ); + assert_eq!(store.len(), 10, "the refused handle was not stored"); + + // Freeing one lets a registration through again. + store.remove("h0"); + store + .insert("h10".into(), apns_n(10), "did:web:vta".into()) + .expect("space was freed"); + } + + /// One push token cannot occupy the registry: the per-token cap counts live + /// handles sharing a destination, and frees up as they are removed. + #[test] + fn per_token_cap_limits_one_destination() { + let store = Store::with_limits(StoreLimits { + max_per_token: 2, + ..StoreLimits::default() + }); + let same = || apns("aa"); + + store + .insert("a".into(), same(), "did:web:vta".into()) + .unwrap(); + store + .insert("b".into(), same(), "did:web:vta".into()) + .unwrap(); + assert_eq!( + store.insert("c".into(), same(), "did:web:vta".into()), + Err(InsertError::TooManyForToken) + ); + // A different token is unaffected. + store + .insert("d".into(), apns("bb"), "did:web:vta".into()) + .unwrap(); + // Removing one frees a slot. + store.remove("a"); + store + .insert("c".into(), same(), "did:web:vta".into()) + .unwrap(); + assert_eq!(store.len(), 3); + } + + /// A Web Push registration is keyed by its endpoint, so the same + /// subscription re-registered with fresh keys still counts as one + /// destination. + #[test] + fn per_token_cap_keys_webpush_on_the_endpoint() { + let store = Store::with_limits(StoreLimits { + max_per_token: 1, + ..StoreLimits::default() + }); + let sub = |auth: &str| { + PushRegistration::Webpush { + endpoint: "https://fcm.googleapis.com/fcm/send/same".into(), + keys: WebPushKeys { + p256dh: "BHTHkS5TN8hSA9_AzgRusH55jqrZjomGJ42mYrmFNIKH1cc0JnR6ZzwjcWQljvhdjlapl3nOtq2P6e9IMjMoWrY".into(), + auth: auth.into(), + }, + } + }; + store + .insert( + "a".into(), + sub("-8GwtL6MnCVPpyjEYoad2A"), + "did:web:vta".into(), + ) + .unwrap(); + assert_eq!( + store.insert( + "b".into(), + sub("differentauthvalue00"), + "did:web:vta".into() + ), + Err(InsertError::TooManyForToken), + "the endpoint is the destination; new keys do not buy a new slot" + ); + } + + /// 1,000 inserts cause no snapshot write on their own, and one flush after + /// them writes exactly once — the debounce, rather than an O(n) rewrite per + /// anonymous request. + #[test] + fn inserts_are_debounced_into_few_snapshot_writes() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("gateway-store.json"); + let store = open(path.clone()); + + for i in 0..1_000 { + store + .insert(format!("h{i}"), apns_n(i), "did:web:vta.example".into()) + .unwrap(); + } + assert_eq!( + store.writes(), + 0, + "mutations must not write the snapshot themselves" + ); + + assert!(store.flush(), "the pending change is written"); + assert_eq!(store.writes(), 1); + // Nothing changed since → no second write. + assert!(!store.flush()); + assert_eq!(store.writes(), 1, "a clean store does not rewrite"); + + // And the single write holds all 1,000 handles. + let reopened = open(path); + assert_eq!(reopened.len(), 1_000); + } + + /// The snapshot holds bearer push credentials, so it must be owner-only — + /// both when freshly created and after a rewrite. + #[test] + #[cfg(unix)] + fn snapshot_is_owner_only() { + use std::os::unix::fs::MetadataExt; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("gateway-store.json"); + let store = open(path.clone()); + store + .insert("h1".into(), apns("a1"), "did:web:vta.example".into()) + .unwrap(); + assert!(store.flush(), "the pending registration is written"); + let mode = std::fs::metadata(&path).unwrap().mode() & 0o777; + assert_eq!(mode, 0o600, "fresh snapshot must be 0600, was {mode:04o}"); + + // A rewrite keeps it owner-only (the rename brings the temp file's mode). + store + .insert("h2".into(), apns("b2"), "did:web:vta.example".into()) + .unwrap(); + assert!(store.flush(), "the second registration is written"); + let mode = std::fs::metadata(&path).unwrap().mode() & 0o777; + assert_eq!(mode, 0o600, "rewritten snapshot must stay 0600"); } } diff --git a/tests/api.rs b/tests/api.rs index 4ca32b5..8987f38 100644 --- a/tests/api.rs +++ b/tests/api.rs @@ -16,8 +16,9 @@ use tower::ServiceExt; use vti_push_gateway::api::{metrics_router, router, AppState}; use vti_push_gateway::egress::EgressPolicy; +use vti_push_gateway::limits::{Limits, RateConfig, DEFAULT_HTTP, DEFAULT_PER_DID}; use vti_push_gateway::sender::{EchoSender, PushSender, SendOutcome}; -use vti_push_gateway::store::Store; +use vti_push_gateway::store::{Store, StoreLimits}; const ED25519_MULTICODEC: [u8; 2] = [0xed, 0x01]; const PUSH_REGISTER: &str = "https://trusttasks.org/spec/push/register/0.2"; @@ -43,16 +44,30 @@ fn did_key_for(sk: &SigningKey) -> String { } fn state() -> AppState { + // Permissive limits by default: these tests exercise the push/* logic, and a + // rate limit tripping mid-test would be a confusing failure. The tests that + // are *about* the limits set their own. + state_with(Store::new(), Limits::permissive()) +} + +fn state_with(store: Store, limits: Limits) -> AppState { let senders: Vec> = vec![Box::new(EchoSender)]; AppState { - store: Arc::new(Store::new()), + store: Arc::new(store), senders: Arc::new(senders), gateway_addr: "https://gw.test".into(), metrics: Arc::new(vti_push_gateway::metrics::Metrics::default()), egress: Arc::new(EgressPolicy::default()), + limits: Arc::new(limits), } } +/// A distinct, valid APNs registration per index — so a flood is not stopped by +/// the per-token cap when the per-request budget is what's under test. +fn apns_registration(n: usize) -> Value { + json!({ "platform": "apns", "token": format!("{n:064x}"), "topic": "org.openvtc.app" }) +} + /// The public router plus a management router sharing one `AppState`, so a test /// can drive `push/*` and then scrape the counters those calls bumped. `/metrics` /// no longer lives on the public router. @@ -778,3 +793,218 @@ async fn bad_signature_is_401() { StatusCode::UNAUTHORIZED ); } + +/// A register budget loose enough not to interfere with tests about other +/// limits. +const DEFAULT_REGISTER_FOR_TEST: RateConfig = RateConfig { + per_second: 1_000_000, + burst: 1_000_000, +}; + +/// PG-2, the `register-flood.sh` PoC as a test: 200 anonymous registrations from +/// one caller are accepted only up to the burst, and the register counter stops +/// climbing. The counter is the assertion that matters — it proves the refusals +/// happened before anything was stored, not merely that a reply said "no". +#[tokio::test] +async fn register_flood_is_refused_after_the_burst() { + let burst = 5; + // The counters live on the management router now, so both are built over one + // shared `AppState` — the public one takes the flood, the management one is + // scraped for what it did. + let st = state_with( + Store::new(), + Limits::new( + RateConfig { + per_second: 1, + burst, + }, + DEFAULT_PER_DID, + DEFAULT_HTTP, + ), + ); + let app = router(st.clone()); + let metrics_app = metrics_router(st, None); + + let mut accepted = 0; + for i in 0..200 { + let reg = tt_doc( + PUSH_REGISTER, + json!({ + "registration": apns_registration(i), + "controllerVtaDid": did_key_for(&signing_key()), + }), + ); + let out = body_json(app.clone().oneshot(post(®, None)).await.unwrap()).await; + if is_success(&out) { + accepted += 1; + } + } + + assert_eq!( + accepted, burst as usize, + "only the burst may be accepted out of 200" + ); + let text = metrics_text(&metrics_app).await; + assert!( + text.contains(&format!("gateway_register_total {burst}\n")), + "the register counter must stop climbing at the burst: {text}" + ); +} + +/// The per-DID budget throttles one noisy trigger without touching another — +/// the reason the wake/provision limiter is keyed rather than global. +#[tokio::test] +async fn wake_budget_is_per_caller_did() { + let vta = signing_key(); + let noisy = signing_key(); + let quiet = signing_key(); + let app = router(state_with( + Store::new(), + Limits::new( + DEFAULT_REGISTER_FOR_TEST, + RateConfig { + per_second: 1, + burst: 3, + }, + DEFAULT_HTTP, + ), + )); + + // Register and provision both triggers onto one handle. + let reg = tt_doc( + PUSH_REGISTER, + json!({ + "registration": apns_registration(1), + "controllerVtaDid": did_key_for(&vta), + }), + ); + let handle = body_json(app.clone().oneshot(post(®, None)).await.unwrap()).await["payload"] + ["wakeHandle"]["handle"] + .as_str() + .unwrap() + .to_string(); + let prov = tt_doc( + PUSH_PROVISION, + json!({ "handle": handle, "policy": { "allowedTriggers": + [did_key_for(&noisy), did_key_for(&quiet)] } }), + ); + assert!(is_success( + &body_json(app.clone().oneshot(post(&prov, Some(&vta))).await.unwrap()).await + )); + + // The noisy trigger burns its own bucket (provision above already spent one + // of the VTA's, not the triggers'). + let wake = tt_doc(PUSH_WAKE, json!({ "handle": handle, "v": 1 })); + let mut noisy_ok = 0; + for _ in 0..10 { + let out = body_json( + app.clone() + .oneshot(post(&wake, Some(&noisy))) + .await + .unwrap(), + ) + .await; + if is_success(&out) { + noisy_ok += 1; + } + } + assert_eq!(noisy_ok, 3, "the noisy trigger is capped at its burst"); + + // The quiet trigger is unaffected. + let out = body_json( + app.clone() + .oneshot(post(&wake, Some(&quiet))) + .await + .unwrap(), + ) + .await; + assert!( + is_success(&out), + "a different trigger DID must not be throttled: {out}" + ); +} + +/// With `max_handles = 10`, the 11th registration is refused and nothing is +/// stored for it. +#[tokio::test] +async fn eleventh_registration_is_refused_at_capacity() { + let st = state_with( + Store::with_limits(StoreLimits { + max_handles: 10, + ..StoreLimits::default() + }), + Limits::permissive(), + ); + let app = router(st.clone()); + let metrics_app = metrics_router(st, None); + + for i in 0..10 { + let reg = tt_doc( + PUSH_REGISTER, + json!({ + "registration": apns_registration(i), + "controllerVtaDid": did_key_for(&signing_key()), + }), + ); + let out = body_json(app.clone().oneshot(post(®, None)).await.unwrap()).await; + assert!( + is_success(&out), + "registration {i} is within the cap: {out}" + ); + } + + let reg = tt_doc( + PUSH_REGISTER, + json!({ + "registration": apns_registration(10), + "controllerVtaDid": did_key_for(&signing_key()), + }), + ); + let out = body_json(app.clone().oneshot(post(®, None)).await.unwrap()).await; + assert!(!is_success(&out), "the 11th must be refused: {out}"); + assert!( + out.to_string().contains("gateway at capacity"), + "the reason should name the limit: {out}" + ); + let text = metrics_text(&metrics_app).await; + assert!( + text.contains("gateway_register_total 10\n"), + "the refused registration must not be counted: {text}" + ); +} + +/// The per-token cap stops one push token from occupying the registry, even +/// though each request is otherwise valid. +#[tokio::test] +async fn repeated_registration_of_one_token_is_capped() { + let app = router(state_with( + Store::with_limits(StoreLimits { + max_per_token: 2, + ..StoreLimits::default() + }), + Limits::permissive(), + )); + + let mut accepted = 0; + for _ in 0..6 { + // The same device token every time. + let reg = tt_doc( + PUSH_REGISTER, + json!({ + "registration": apns_registration(7), + "controllerVtaDid": did_key_for(&signing_key()), + }), + ); + let out = body_json(app.clone().oneshot(post(®, None)).await.unwrap()).await; + if is_success(&out) { + accepted += 1; + } else { + assert!( + out.to_string() + .contains("too many handles for this push token"), + "{out}" + ); + } + } + assert_eq!(accepted, 2, "one token gets max_per_token handles, no more"); +}