From 67cf60a56e82ec8b37a5a63eb72a646ba12dc1a1 Mon Sep 17 00:00:00 2001 From: Glenn Gore Date: Sat, 12 Sep 2026 08:52:01 +0200 Subject: [PATCH] feat(store): expire, cap and rate-limit the anonymous registration path push/register needs no credentials, so the handle registry was a map an unauthenticated caller could grow without bound, and every mutation reserialised the whole map under the write lock. Four changes, in order of how much they matter. Expiry is the root-cause fix. A handle is inert until its VTA provisions a trigger, so HandleRecord gains created_at (#[serde(default)], so existing snapshots load) and a tokio sweeper drops handles still unprovisioned after GATEWAY_UNPROVISIONED_TTL_SECS, default 24h. Anonymous growth becomes bounded churn. 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 or Web Push endpoint, so one token cannot occupy the registry. The per-token count is an index maintained under the same lock as the map, not a scan, so the check stays O(1) under the flood it exists to stop; it keys on a truncated SHA-256 rather than the token, to avoid a second copy of a bearer credential in another map's keys. Rate limits in two layers, because DIDComm bypasses HTTP middleware and is the preferred transport. A tower_governor layer limits POST /trust-tasks per peer IP, which needs into_make_service_with_connect_info; and the transport-agnostic dispatch core limits register against a global budget and provision/wake against a budget keyed by the authenticated DID. The keyed buckets are reclaimed on a timer, since they are themselves keyed by caller-chosen input. Persistence is debounced: mutations set a dirty flag and a background flusher writes at most once per GATEWAY_SNAPSHOT_FLUSH_MS, keeping temp-file + fsync + rename. Drop flushes, so a clean shutdown is durable. Signed-off-by: Glenn Gore --- Cargo.lock | 249 ++++++++++++++-- Cargo.toml | 10 + README.md | 50 ++++ src/api.rs | 68 ++++- src/lib.rs | 1 + src/limits.rs | 286 +++++++++++++++++++ src/main.rs | 131 ++++++++- src/store.rs | 767 ++++++++++++++++++++++++++++++++++++++++++++------ tests/api.rs | 234 ++++++++++++++- 9 files changed, 1669 insertions(+), 127 deletions(-) create mode 100644 src/limits.rs 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"); +}