diff --git a/README.md b/README.md index d1587583..e75625f6 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,16 @@ Useful environment variables: - `WAF_IDS_STATE_PATH`: optional JSON state path. When omitted, the service runs with seeded in-memory state. - `DNSBL_ORIGIN`: DNSBL zone origin, default `dnsbl.local` - `EVENT_LIMIT`: retained event count, default `1000`; must be greater than zero +- `RATE_LIMIT`: optional per-client gateway request budget; `0` disables local limiting +- `RATE_LIMIT_WINDOW`: fixed-window length in seconds for `RATE_LIMIT`, default `60` +- `RATE_LIMIT_MAX_CLIENTS`: maximum in-memory client buckets retained by the local limiter, default `4096`; may also be bootstrapped via `WAF_IDS_CREDENTIALS_PATH` as `rate_limit_max_clients` +- `TRUSTED_PROXY_IPS`: optional comma-separated proxy peer IPs allowed to supply `X-Forwarded-For` + +When the local limiter returns HTTP `429`, the response includes a `Retry-After` +header plus JSON `reason` codes that distinguish per-client quota exhaustion +from local limiter saturation (`local_rate_limiter_capacity_exceeded`). +Forwarded client-IP headers are ignored unless the connected peer IP is present +in `TRUSTED_PROXY_IPS`. Example with persistent local state: diff --git a/docs/papers/overload-management-service-design-primitive-ew10-2002.pdf b/docs/papers/overload-management-service-design-primitive-ew10-2002.pdf new file mode 100644 index 00000000..6f8fb890 Binary files /dev/null and b/docs/papers/overload-management-service-design-primitive-ew10-2002.pdf differ diff --git a/docs/runbooks/operations.md b/docs/runbooks/operations.md index 9b6b7015..56eb778c 100644 --- a/docs/runbooks/operations.md +++ b/docs/runbooks/operations.md @@ -78,6 +78,30 @@ The smoke test starts the service on a temporary port with a temporary JSON stat When `WAF_IDS_STATE_PATH` is enabled, the process writes a temporary sibling file and atomically replaces the configured state path. If a management write cannot be persisted, the in-memory mutation is rolled back and the API returns `500`. +## Local Admission Control + +`RATE_LIMIT` and `RATE_LIMIT_WINDOW` enable a per-client fixed-window limiter for +`/gateway` traffic. `RATE_LIMIT_MAX_CLIENTS` bounds the number of in-memory +client buckets the process will retain; stale buckets age out after one full +window. When the map is full, unseen clients receive `429 Too Many Requests` +with `Retry-After` and reason +`local_rate_limiter_capacity_exceeded` until older buckets expire. +By default the limiter keys on the connected peer IP. `X-Forwarded-For` is +considered only when the peer IP is listed in `TRUSTED_PROXY_IPS`, which +prevents direct clients from manufacturing new local limiter buckets with +spoofed forwarding headers. `X-Real-IP` is not trusted for limiter identity +because the header does not carry a verifiable proxy chain. + +This is a local emergency guard, not the distributed quota authority described +in issue `#83`. The contract follows RFC 6585's guidance that `429` responses +may include `Retry-After`, and it aligns with OWASP ASVS 5.0 availability +controls by failing with a bounded, operator-visible response instead of +allowing attacker-controlled client cardinality to grow process memory without +limit. Welsh et al. (2001) and Welsh and Culler (2002) motivate this choice: +overload control should be explicit in the service design, and excess demand +should be shed early enough to bound queue growth, latency, and memory rather +than letting attacker-driven request cardinality expand invisibly. + ## Safe Change Procedure 1. Start new routes in `monitor` mode. @@ -107,3 +131,11 @@ This baseline is suitable for local and controlled lab deployments. Internet-fac - Live Suricata EVE tailing / shipper (HTTP ingest of EVE alerts is available at `POST /api/ids/suricata/eve`) - Live MISP REST pull or live OpenCTI GraphQL pull (HTTP STIX/MISP/OpenCTI document ingest and TAXII 2.1 poll are available at `POST /api/threat-intel/stix`, `POST /api/threat-intel/misp`, `POST /api/threat-intel/opencti`, and `POST /api/threat-intel/taxii/poll`) - human approval workflow for AI SOC recommendations that change enforcement + +## References + +- Nottingham, M., & Fielding, R. (2012). *Additional HTTP status codes* (RFC 6585). https://www.rfc-editor.org/info/rfc6585 +- OWASP Foundation. (2025). *OWASP Application Security Verification Standard 5.0.0*. https://owasp.org/www-project-application-security-verification-standard/ +- Souppaya, M., Scarfone, K., & Dodson, D. (2022). *Secure Software Development Framework (SSDF) version 1.1* (NIST SP 800-218). https://doi.org/10.6028/NIST.SP.800-218 +- Welsh, M., Culler, D., & Brewer, E. (2001). *SEDA: An architecture for well-conditioned, scalable Internet services*. Proceedings of the Eighteenth ACM Symposium on Operating Systems Principles, 230-243. https://www.sosp.org/2001/papers/welsh.pdf +- Welsh, M., & Culler, D. (2002). *Overload management as a fundamental service design primitive*. Proceedings of the 10th ACM SIGOPS European Workshop, 63-69. Local PDF: `docs/papers/overload-management-service-design-primitive-ew10-2002.pdf` diff --git a/src/credentials.rs b/src/credentials.rs index bcc07e50..d7d1b6e0 100644 --- a/src/credentials.rs +++ b/src/credentials.rs @@ -1,5 +1,5 @@ -//! Secret-bearing and fetch-sensitive configuration via a process-local -//! credential registry. +//! Secret-bearing configuration plus narrowly scoped runtime bootstrap values +//! via a process-local credential registry. //! //! Org guidance: runtime code must not treat raw environment variables as the //! source of secrets. Environment (and optional credentials file) are bootstrap @@ -12,6 +12,8 @@ use std::{collections::HashMap, io::ErrorKind, path::Path}; /// Well-known credentials loaded into the registry at bootstrap. pub const CRED_ADMIN_TOKEN: &str = "admin_token"; pub const CRED_ADMIN_TOKENS: &str = "admin_tokens"; +pub const CRED_RATE_LIMIT_MAX_CLIENTS: &str = "rate_limit_max_clients"; +pub const CRED_TRUSTED_PROXY_IPS: &str = "trusted_proxy_ips"; /// Where secret-bearing credentials were loaded from (never includes values). #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] @@ -27,6 +29,7 @@ pub enum CredentialSource { } impl CredentialSource { + /// Return the stable telemetry label exposed in health/support surfaces. pub fn as_str(self) -> &'static str { match self { Self::File => "file", @@ -37,7 +40,7 @@ impl CredentialSource { } /// In-process map of secret credentials. Values are never logged or serialized -/// into health/support surfaces — only the source label is exposed. +/// into health/support surfaces; only the source label is exposed. #[derive(Debug, Clone, Default)] pub struct CredentialRegistry { values: HashMap, @@ -45,41 +48,77 @@ pub struct CredentialRegistry { } impl CredentialRegistry { + /// Construct an empty registry with no bootstrapped credentials. pub fn empty() -> Self { Self::default() } + /// Read a previously bootstrapped credential or config override by name. pub fn get_credential(&self, name: &str) -> Option<&str> { self.values.get(name).map(String::as_str) } + /// Report where secret-bearing bootstrap values came from. pub fn source(&self) -> CredentialSource { self.source } + /// True when at least one admin authentication path is configured. pub fn has_admin_auth(&self) -> bool { self.get_credential(CRED_ADMIN_TOKEN) - .is_some_and(|v| !v.is_empty()) + .is_some_and(|value| !value.is_empty()) || self .get_credential(CRED_ADMIN_TOKENS) - .is_some_and(|v| !v.trim().is_empty()) + .is_some_and(|value| !value.trim().is_empty()) } - /// Bootstrap secret-bearing credentials plus the optional KEV fetch override. + /// Bootstrap secret-bearing credentials. /// /// Precedence: JSON credentials file (when present) wins per-key; missing - /// keys are filled from the env bootstrap values. Operational non-secret - /// config (bind address, limits, DNSBL origin) stays on env. The KEV URL - /// defaults to the built-in CISA endpoint and is accepted here only as a - /// server-side override that must still satisfy the runtime allowlist. + /// keys are filled from the env bootstrap values. pub fn bootstrap_secrets( credentials_path: Option<&Path>, env_admin_token: Option, env_admin_tokens: Option, + ) -> Result { + let mut registry = Self::bootstrap_with_runtime_overrides( + credentials_path, + env_admin_token, + env_admin_tokens, + None, + None, + )?; + registry.values.remove(CRED_RATE_LIMIT_MAX_CLIENTS); + registry.values.remove(CRED_TRUSTED_PROXY_IPS); + Ok(registry) + } + + /// Bootstrap secret-bearing credentials plus env-transported runtime + /// policy overrides, keeping env access localized to registry bootstrap. + pub fn bootstrap_runtime_registry( + credentials_path: Option<&Path>, + env_admin_token: Option, + env_admin_tokens: Option, + ) -> Result { + Self::bootstrap_with_runtime_overrides( + credentials_path, + env_admin_token, + env_admin_tokens, + std::env::var("RATE_LIMIT_MAX_CLIENTS").ok(), + std::env::var("TRUSTED_PROXY_IPS").ok(), + ) + } + + /// Bootstrap secret-bearing credentials plus runtime policy overrides that + /// must flow through the registry before handlers consume them. + pub fn bootstrap_with_runtime_overrides( + credentials_path: Option<&Path>, + env_admin_token: Option, + env_admin_tokens: Option, + env_rate_limit_max_clients: Option, + env_trusted_proxy_ips: Option, ) -> Result { let mut values = HashMap::new(); - // CredentialSource is documented (and reported via HealthStatus/support - // bundle) as admin-secret provenance specifically. let mut admin_from_file = false; let mut admin_from_env = false; @@ -93,12 +132,19 @@ impl CredentialRegistry { path.display() ) })?; - for key in [CRED_ADMIN_TOKEN, CRED_ADMIN_TOKENS] { + for key in [ + CRED_ADMIN_TOKEN, + CRED_ADMIN_TOKENS, + CRED_RATE_LIMIT_MAX_CLIENTS, + CRED_TRUSTED_PROXY_IPS, + ] { if let Some(raw) = file_map.get(key) { let text = json_value_as_nonempty_string(raw); if let Some(text) = text { values.insert(key.to_string(), text); - admin_from_file = true; + if matches!(key, CRED_ADMIN_TOKEN | CRED_ADMIN_TOKENS) { + admin_from_file = true; + } } } } @@ -125,6 +171,17 @@ impl CredentialRegistry { values.insert(CRED_ADMIN_TOKENS.to_string(), tokens); admin_from_env = true; } + if !values.contains_key(CRED_RATE_LIMIT_MAX_CLIENTS) + && let Some(max_clients) = env_rate_limit_max_clients.filter(|value| !value.is_empty()) + { + values.insert(CRED_RATE_LIMIT_MAX_CLIENTS.to_string(), max_clients); + } + if !values.contains_key(CRED_TRUSTED_PROXY_IPS) + && let Some(trusted_proxy_ips) = env_trusted_proxy_ips.filter(|value| !value.is_empty()) + { + values.insert(CRED_TRUSTED_PROXY_IPS.to_string(), trusted_proxy_ips); + } + let source = if admin_from_file { CredentialSource::File } else if admin_from_env { @@ -137,6 +194,7 @@ impl CredentialRegistry { } } +/// Convert supported JSON credential values into non-empty strings. fn json_value_as_nonempty_string(value: &serde_json::Value) -> Option { match value { serde_json::Value::String(text) if !text.is_empty() => Some(text.clone()), @@ -159,10 +217,12 @@ mod tests { #[test] fn bootstrap_from_env_only() { - let registry = CredentialRegistry::bootstrap_secrets( + let registry = CredentialRegistry::bootstrap_with_runtime_overrides( None, Some("secret".to_string()), Some("tok:alice".to_string()), + Some("2048".to_string()), + Some("198.51.100.10,198.51.100.11".to_string()), ) .unwrap(); assert_eq!(registry.source(), CredentialSource::Env); @@ -171,6 +231,14 @@ mod tests { registry.get_credential(CRED_ADMIN_TOKENS), Some("tok:alice") ); + assert_eq!( + registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), + Some("2048") + ); + assert_eq!( + registry.get_credential(CRED_TRUSTED_PROXY_IPS), + Some("198.51.100.10,198.51.100.11") + ); assert!(registry.has_admin_auth()); } @@ -182,6 +250,23 @@ mod tests { assert!(!registry.has_admin_auth()); } + #[test] + fn bootstrap_secrets_excludes_runtime_policy_overrides() { + let registry = CredentialRegistry::bootstrap_secrets( + None, + Some("secret".to_string()), + Some("tok:alice".to_string()), + ) + .unwrap(); + assert_eq!(registry.get_credential(CRED_ADMIN_TOKEN), Some("secret")); + assert_eq!( + registry.get_credential(CRED_ADMIN_TOKENS), + Some("tok:alice") + ); + assert_eq!(registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), None); + assert_eq!(registry.get_credential(CRED_TRUSTED_PROXY_IPS), None); + } + #[test] fn file_overrides_env_per_key() { let dir = std::env::temp_dir().join(format!( @@ -202,10 +287,12 @@ mod tests { .unwrap(); drop(file); - let registry = CredentialRegistry::bootstrap_secrets( + let registry = CredentialRegistry::bootstrap_with_runtime_overrides( Some(&path), Some("from-env".to_string()), Some("envtok:env".to_string()), + Some("1024".to_string()), + Some("198.51.100.10".to_string()), ) .unwrap(); assert_eq!(registry.source(), CredentialSource::File); @@ -214,6 +301,14 @@ mod tests { registry.get_credential(CRED_ADMIN_TOKENS), Some("filetok:operator") ); + assert_eq!( + registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), + Some("1024") + ); + assert_eq!( + registry.get_credential(CRED_TRUSTED_PROXY_IPS), + Some("198.51.100.10") + ); let _ = std::fs::remove_dir_all(&dir); } @@ -232,10 +327,12 @@ mod tests { let path = dir.join("credentials.json"); std::fs::write(&path, r#"{"admin_token":"file-only"}"#).unwrap(); - let registry = CredentialRegistry::bootstrap_secrets( + let registry = CredentialRegistry::bootstrap_with_runtime_overrides( Some(&path), Some("ignored".to_string()), Some("envtok:bob".to_string()), + Some("3072".to_string()), + Some("198.51.100.10".to_string()), ) .unwrap(); assert_eq!(registry.source(), CredentialSource::File); @@ -244,6 +341,14 @@ mod tests { registry.get_credential(CRED_ADMIN_TOKENS), Some("envtok:bob") ); + assert_eq!( + registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), + Some("3072") + ); + assert_eq!( + registry.get_credential(CRED_TRUSTED_PROXY_IPS), + Some("198.51.100.10") + ); let _ = std::fs::remove_dir_all(&dir); } @@ -288,4 +393,68 @@ mod tests { assert!(err.contains("not valid JSON")); let _ = std::fs::remove_dir_all(&dir); } + + #[test] + fn file_overrides_rate_limit_max_clients_without_affecting_secret_source() { + let dir = std::env::temp_dir().join(format!( + "wardnet-creds-config-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("credentials.json"); + std::fs::write(&path, r#"{"rate_limit_max_clients":"512"}"#).unwrap(); + + let registry = CredentialRegistry::bootstrap_with_runtime_overrides( + Some(&path), + None, + None, + Some("2048".to_string()), + Some("198.51.100.10".to_string()), + ) + .unwrap(); + assert_eq!(registry.source(), CredentialSource::None); + assert_eq!( + registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), + Some("512") + ); + assert_eq!( + registry.get_credential(CRED_TRUSTED_PROXY_IPS), + Some("198.51.100.10") + ); + + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn bootstrap_secrets_ignores_policy_keys_from_credentials_file() { + let dir = std::env::temp_dir().join(format!( + "wardnet-creds-secret-only-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("credentials.json"); + std::fs::write( + &path, + r#"{"admin_token":"file-secret","rate_limit_max_clients":"512","trusted_proxy_ips":"198.51.100.10"}"#, + ) + .unwrap(); + + let registry = CredentialRegistry::bootstrap_secrets(Some(&path), None, None).unwrap(); + assert_eq!( + registry.get_credential(CRED_ADMIN_TOKEN), + Some("file-secret") + ); + assert_eq!(registry.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), None); + assert_eq!(registry.get_credential(CRED_TRUSTED_PROXY_IPS), None); + + let _ = std::fs::remove_dir_all(&dir); + } } diff --git a/src/lib.rs b/src/lib.rs index ab902cae..8d6fad33 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ use axum::{ - Json, Router, + Extension, Json, Router, body::Bytes, - extract::{DefaultBodyLimit, Path as PathParam, Query, State}, + extract::{ConnectInfo, DefaultBodyLimit, Path as PathParam, Query, State}, http::{HeaderMap, Method, StatusCode, Uri}, response::{Html, IntoResponse, Response}, routing::{any, get, post}, @@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize}; use std::{ collections::{HashMap, HashSet}, io::ErrorKind, - net::{IpAddr, Ipv4Addr}, + net::{IpAddr, Ipv4Addr, SocketAddr}, path::{Path, PathBuf}, sync::Arc, time::{SystemTime, UNIX_EPOCH}, @@ -44,7 +44,12 @@ mod opencti_import; mod stix_import; mod suricata_eve; mod taxii; -pub use credentials::{CRED_ADMIN_TOKEN, CRED_ADMIN_TOKENS, CredentialRegistry, CredentialSource}; +pub use credentials::{ + CRED_ADMIN_TOKEN, CRED_ADMIN_TOKENS, CRED_RATE_LIMIT_MAX_CLIENTS, CRED_TRUSTED_PROXY_IPS, + CredentialRegistry, CredentialSource, +}; + +const DEFAULT_RATE_LIMIT_MAX_CLIENTS: usize = 4_096; #[derive(Clone)] pub struct AppState { @@ -62,10 +67,11 @@ pub struct AppState { dnsbl_origin: String, event_limit: usize, // Ephemeral per-client-IP fixed-window counters (not persisted). - // ponytail: unbounded map — add TTL eviction if client-IP cardinality grows. - rate_limiter: Arc>>, + rate_limiter: Arc>>, rate_limit: u32, rate_limit_window: u64, + rate_limit_max_clients: usize, + trusted_proxies: HashSet, // Max accepted request body size in bytes; oversized requests get 413. max_body_bytes: usize, // Optional Clearfolio document-viewer integration. `None` unless configured. @@ -102,6 +108,43 @@ pub struct ClearfolioConfig { pub permissions: String, } +#[derive(Debug, Clone, Copy)] +struct RateLimitBucket { + window_start: u64, + count: u32, + last_seen: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct RateLimitDecision { + allowed: bool, + retry_after_seconds: u64, + reason: &'static str, +} + +impl RateLimitDecision { + const WINDOW_EXCEEDED: &'static str = "rate_limit_exceeded"; + const CAPACITY_EXCEEDED: &'static str = "local_rate_limiter_capacity_exceeded"; + + /// Build the success result for an admitted request. + fn allowed() -> Self { + Self { + allowed: true, + retry_after_seconds: 0, + reason: "", + } + } + + /// Build a denied decision with a stable reason and retry horizon. + fn denied(reason: &'static str, retry_after_seconds: u64) -> Self { + Self { + allowed: false, + retry_after_seconds: retry_after_seconds.max(1), + reason, + } + } +} + impl AppState { pub fn seeded(admin_token: Option) -> Self { Self::new(AppData::seeded(), AppConfig::memory(admin_token)) @@ -138,6 +181,8 @@ impl AppState { rate_limiter: Arc::new(Mutex::new(HashMap::new())), rate_limit: 0, rate_limit_window: 60, + rate_limit_max_clients: DEFAULT_RATE_LIMIT_MAX_CLIENTS, + trusted_proxies: HashSet::new(), max_body_bytes: 1_048_576, clearfolio: None, soc_llm: None, @@ -196,6 +241,19 @@ impl AppState { self } + /// Bound the number of local in-memory client buckets retained by the rate + /// limiter. When full, unseen clients receive 429 until stale buckets age out. + pub fn with_rate_limit_max_clients(mut self, max_clients: usize) -> Self { + self.rate_limit_max_clients = max_clients.max(1); + self + } + + /// Trust forwarded client-IP headers only from these connected proxy IPs. + pub fn with_trusted_proxies(mut self, trusted_proxies: HashSet) -> Self { + self.trusted_proxies = trusted_proxies; + self + } + /// Configure RBAC admin tokens (token -> principal). A non-empty map takes /// precedence over the single `admin_token`. Builder-style. pub fn with_admin_tokens(mut self, tokens: HashMap) -> Self { @@ -224,25 +282,50 @@ impl AppState { .map(|principal| principal.actor.clone()) } - /// Records one gateway request for `client_ip` and returns `true` if it is - /// within the configured rate limit. Unknown IPs share one bucket. - async fn allow_request(&self, client_ip: Option) -> bool { + /// Records one gateway request for `client_ip` and returns the local + /// admission decision. Unknown IPs share one bucket. + async fn allow_request(&self, client_ip: Option) -> RateLimitDecision { if self.rate_limit == 0 { - return true; + return RateLimitDecision::allowed(); } let key = client_ip.unwrap_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED)); let now = now_unix(); let mut map = self.rate_limiter.lock().await; - let (window_start, count) = map.get(&key).copied().unwrap_or((now, 0)); + prune_rate_limit_buckets(&mut map, now, self.rate_limit_window); + if !map.contains_key(&key) && map.len() >= self.rate_limit_max_clients { + return RateLimitDecision::denied( + RateLimitDecision::CAPACITY_EXCEEDED, + self.rate_limit_window, + ); + } + let bucket = map.get(&key).copied().unwrap_or(RateLimitBucket { + window_start: now, + count: 0, + last_seen: now, + }); let (allowed, new_start, new_count) = rate_limit_step( now, - window_start, - count, + bucket.window_start, + bucket.count, self.rate_limit, self.rate_limit_window, ); - map.insert(key, (new_start, new_count)); - allowed + map.insert( + key, + RateLimitBucket { + window_start: new_start, + count: new_count, + last_seen: now, + }, + ); + if allowed { + RateLimitDecision::allowed() + } else { + RateLimitDecision::denied( + RateLimitDecision::WINDOW_EXCEEDED, + retry_after_seconds(now, new_start, self.rate_limit_window), + ) + } } async fn mutate_and_persist( @@ -2321,6 +2404,7 @@ async fn dnsbl_zone(State(state): State) -> impl IntoResponse { async fn gateway( State(state): State, + peer_addr: Option>>, method: Method, uri: Uri, headers: HeaderMap, @@ -2343,33 +2427,54 @@ async fn gateway( (route.clone(), data.threats.clone(), data.dnsbl.clone()) }; - let client_ip = client_ip_from_headers(&headers); + let peer_ip = peer_addr.map(|Extension(ConnectInfo(peer_addr))| peer_addr.ip()); + let client_ip = client_ip_from_headers(&headers, peer_ip, &state.trusted_proxies); // Rate limiting runs before scoring/proxying so floods are shed cheaply. - if !state.allow_request(client_ip).await { + let rate_limit = state.allow_request(client_ip).await; + if !rate_limit.allowed { + let action = if rate_limit.reason == RateLimitDecision::CAPACITY_EXCEEDED { + "rate_limiter_saturated" + } else { + "rate_limited" + }; record_event( &state, client_ip, Some(route.id.clone()), - "rate_limited", - format!( - "rate limit exceeded ({} requests per {}s)", - state.rate_limit, state.rate_limit_window - ), + action, + match rate_limit.reason { + RateLimitDecision::CAPACITY_EXCEEDED => format!( + "local rate limiter saturated (max {} client buckets, {}s TTL)", + state.rate_limit_max_clients, state.rate_limit_window + ), + _ => format!( + "rate limit exceeded ({} requests per {}s)", + state.rate_limit, state.rate_limit_window + ), + }, 0, gateway_path, ) .await; - return ( + let mut response = ( StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ - "action": "rate_limited", + "action": action, + "reason": rate_limit.reason, "route_id": route.id, "limit": state.rate_limit, - "window_seconds": state.rate_limit_window + "window_seconds": state.rate_limit_window, + "max_clients": state.rate_limit_max_clients, + "retry_after_seconds": rate_limit.retry_after_seconds })), ) .into_response(); + let retry_after = + axum::http::HeaderValue::from_str(&rate_limit.retry_after_seconds.to_string()) + .expect("retry-after header is numeric"); + response.headers_mut().insert("retry-after", retry_after); + return response; } let body_text = String::from_utf8_lossy(&body); @@ -2440,11 +2545,61 @@ async fn gateway( } } -fn client_ip_from_headers(headers: &HeaderMap) -> Option { +/// Resolve the client IP for rate limiting, trusting forwarded headers only +/// when the connected peer is an explicitly trusted proxy. +fn client_ip_from_headers( + headers: &HeaderMap, + peer_ip: Option, + trusted_proxies: &HashSet, +) -> Option { + if let Some(peer_ip) = peer_ip { + if trusted_proxies.contains(&peer_ip) { + return trusted_forwarded_client_ip(headers, peer_ip, trusted_proxies) + .or(Some(peer_ip)); + } + return Some(peer_ip); + } + forwarded_client_ip(headers) +} + +/// Extract the first untrusted client IP from a trusted proxy chain. +fn trusted_forwarded_client_ip( + headers: &HeaderMap, + peer_ip: IpAddr, + trusted_proxies: &HashSet, +) -> Option { + let Some(value) = headers + .get("x-forwarded-for") + .and_then(|value| value.to_str().ok()) + else { + return headers + .get("x-real-ip") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()); + }; + + let mut current_hop = peer_ip; + for hop in value.rsplit(',').map(str::trim) { + let candidate = hop.parse::().ok()?; + if !trusted_proxies.contains(¤t_hop) { + return Some(current_hop); + } + current_hop = candidate; + } + + if trusted_proxies.contains(¤t_hop) { + None + } else { + Some(current_hop) + } +} + +/// Extract a forwarded client IP from standard proxy headers. +fn forwarded_client_ip(headers: &HeaderMap) -> Option { headers .get("x-forwarded-for") .and_then(|value| value.to_str().ok()) - .and_then(|value| value.split(',').next()) + .and_then(|value| value.rsplit(',').next()) .map(str::trim) .or_else(|| { headers @@ -2624,6 +2779,23 @@ fn audit_actor(state: &AppState, headers: &HeaderMap) -> String { .to_string() } +/// Drop client buckets that have been idle for at least one rate-limit window. +fn prune_rate_limit_buckets( + map: &mut HashMap, + now: u64, + window_secs: u64, +) { + map.retain(|_, bucket| now.saturating_sub(bucket.last_seen) < window_secs); +} + +/// Compute the `Retry-After` value for a fixed-window rate-limit rejection. +fn retry_after_seconds(now: u64, window_start: u64, window_secs: u64) -> u64 { + window_start + .saturating_add(window_secs) + .saturating_sub(now) + .max(1) +} + /// Parses an `ADMIN_TOKENS` string into a token -> [`AdminPrincipal`] map. /// /// Format is comma-separated items: @@ -3268,6 +3440,59 @@ pub fn parse_u64_env( } } +/// Parse a `usize` environment value (already read as an optional string), +/// returning `default` when absent and rejecting zero or malformed values. +pub fn parse_usize_env( + name: &str, + raw: Option<&str>, + default: usize, +) -> Result> { + match raw { + Some(raw) => { + let value = raw.parse::().map_err(|error| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("{name} must be a positive integer, got {raw:?}: {error}"), + ) + })?; + if value == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("{name} must be greater than 0"), + ) + .into()); + } + Ok(value) + } + None => Ok(default), + } +} + +/// Parse a comma-separated set of IP addresses from an optional env value. +pub fn parse_ip_set_env( + name: &str, + raw: Option<&str>, +) -> Result, Box> { + let mut values = HashSet::new(); + let Some(raw) = raw.map(str::trim).filter(|raw| !raw.is_empty()) else { + return Ok(values); + }; + for candidate in raw + .split(',') + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let ip = candidate.parse::().map_err(|error| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("{name} must be a comma-separated list of IP addresses, got {candidate:?}: {error}"), + ) + })?; + values.insert(ip); + } + Ok(values) +} + /// Read gateway configuration from the process environment, bind the listener, /// and serve until `shutdown` resolves. The binary entrypoint is a thin shim /// over this function so every branch is reachable from tests (the parse/error @@ -3282,11 +3507,15 @@ pub async fn run_from_env( let credentials_path = std::env::var("WAF_IDS_CREDENTIALS_PATH") .ok() .map(PathBuf::from); - let credentials = CredentialRegistry::bootstrap_secrets( + let credentials = CredentialRegistry::bootstrap_runtime_registry( credentials_path.as_deref(), std::env::var("ADMIN_TOKEN").ok(), std::env::var("ADMIN_TOKENS").ok(), )?; + let trusted_proxies = parse_ip_set_env( + "TRUSTED_PROXY_IPS", + credentials.get_credential(CRED_TRUSTED_PROXY_IPS), + )?; let config = AppConfig { admin_token: credentials .get_credential(CRED_ADMIN_TOKEN) @@ -3302,6 +3531,11 @@ pub async fn run_from_env( std::env::var("RATE_LIMIT_WINDOW").ok().as_deref(), 60, )?; + let rate_limit_max_clients = parse_usize_env( + "RATE_LIMIT_MAX_CLIENTS", + credentials.get_credential(CRED_RATE_LIMIT_MAX_CLIENTS), + DEFAULT_RATE_LIMIT_MAX_CLIENTS, + )?; let admin_tokens = parse_admin_tokens( credentials .get_credential(CRED_ADMIN_TOKENS) @@ -3322,12 +3556,17 @@ pub async fn run_from_env( .await .map_err(|message| std::io::Error::new(std::io::ErrorKind::InvalidData, message))? .with_rate_limit(rate_limit, rate_limit_window) + .with_rate_limit_max_clients(rate_limit_max_clients) + .with_trusted_proxies(trusted_proxies) .with_admin_tokens(admin_tokens) .with_credentials_source(credentials.source()) .with_max_body_size(max_body_bytes); - let served = axum::serve(listener, build_app(state)) - .with_graceful_shutdown(shutdown) - .await; + let served = axum::serve( + listener, + build_app(state).into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown) + .await; served?; Ok(()) } @@ -3343,7 +3582,7 @@ mod tests { use std::{ future::IntoFuture, io::{Read, Write}, - net::TcpListener as StdTcpListener, + net::{Ipv4Addr, SocketAddr, SocketAddrV4, TcpListener as StdTcpListener}, thread, time::{SystemTime, UNIX_EPOCH}, }; @@ -3367,6 +3606,8 @@ mod tests { "EVENT_LIMIT", "RATE_LIMIT", "RATE_LIMIT_WINDOW", + "RATE_LIMIT_MAX_CLIENTS", + "TRUSTED_PROXY_IPS", "MAX_BODY_BYTES", ] { unsafe { std::env::remove_var(name) }; @@ -3401,6 +3642,37 @@ mod tests { assert!(parse_u64_env("RATE_LIMIT_WINDOW", Some("abc"), 60).is_err()); } + #[test] + fn parse_usize_env_reads_optional_env() { + assert_eq!( + parse_usize_env("RATE_LIMIT_MAX_CLIENTS", None, 7).unwrap(), + 7 + ); + assert_eq!( + parse_usize_env("RATE_LIMIT_MAX_CLIENTS", Some("120"), 1).unwrap(), + 120 + ); + assert!(parse_usize_env("RATE_LIMIT_MAX_CLIENTS", Some("0"), 1).is_err()); + assert!(parse_usize_env("RATE_LIMIT_MAX_CLIENTS", Some("abc"), 1).is_err()); + } + + #[test] + fn parse_ip_set_env_reads_optional_env() { + assert!( + parse_ip_set_env("TRUSTED_PROXY_IPS", None) + .unwrap() + .is_empty() + ); + assert_eq!( + parse_ip_set_env("TRUSTED_PROXY_IPS", Some("203.0.113.9, 198.51.100.7")).unwrap(), + HashSet::from([ + "203.0.113.9".parse::().unwrap(), + "198.51.100.7".parse::().unwrap(), + ]) + ); + assert!(parse_ip_set_env("TRUSTED_PROXY_IPS", Some("not-an-ip")).is_err()); + } + #[test] fn parses_and_limits_phishing_database_feeds() { let domains = parse_phishing_domains( @@ -3460,6 +3732,22 @@ mod tests { clear_run_env(); } + #[tokio::test] + async fn run_from_env_rejects_malformed_rate_limit_max_clients() { + let _guard = ENV_GUARD.lock().await; + clear_run_env(); + unsafe { + std::env::set_var("BIND_ADDR", "127.0.0.1:0"); + std::env::set_var("RATE_LIMIT_MAX_CLIENTS", "not-a-number"); + } + assert!( + run_from_env(Box::pin(std::future::ready(()))) + .await + .is_err() + ); + clear_run_env(); + } + #[tokio::test] async fn run_from_env_rejects_malformed_max_body_bytes() { let _guard = ENV_GUARD.lock().await; @@ -3588,6 +3876,17 @@ mod tests { } } + async fn app_request(app: &Router, mut request: Request) -> Response { + if request + .extensions() + .get::>() + .is_none() + { + insert_peer(&mut request, IpAddr::V4(Ipv4Addr::LOCALHOST)); + } + app.clone().oneshot(request).await.unwrap() + } + fn kev_import_request() -> KevImportRequest { KevImportRequest { feed_id: "cisa-kev-seoul".to_string(), @@ -3596,8 +3895,16 @@ mod tests { } } - async fn app_request(app: &Router, request: Request) -> Response { - app.clone().oneshot(request).await.unwrap() + fn insert_peer(request: &mut Request, ip: IpAddr) { + request + .extensions_mut() + .insert(ConnectInfo(SocketAddr::V4(SocketAddrV4::new( + match ip { + IpAddr::V4(ip) => ip, + IpAddr::V6(_) => Ipv4Addr::LOCALHOST, + }, + 443, + )))); } fn empty_request(method: Method, uri: &str) -> Request { @@ -3636,12 +3943,26 @@ mod tests { } fn gateway_get_from_ip(uri: &str, ip: &str) -> Request { - Request::builder() + let ip = ip.parse::().unwrap(); + let mut request = Request::builder() .method(Method::GET) .uri(uri) - .header("x-forwarded-for", ip) + .header("x-forwarded-for", ip.to_string()) .body(Body::empty()) - .unwrap() + .unwrap(); + insert_peer(&mut request, ip); + request + } + + fn gateway_get_via_proxy(uri: &str, forwarded_ip: &str, proxy_ip: &str) -> Request { + let mut request = Request::builder() + .method(Method::GET) + .uri(uri) + .header("x-forwarded-for", forwarded_ip) + .body(Body::empty()) + .unwrap(); + insert_peer(&mut request, proxy_ip.parse::().unwrap()); + request } #[test] @@ -3657,6 +3978,33 @@ mod tests { assert_eq!(rate_limit_step(160, 100, 2, 2, 60), (true, 160, 1)); } + #[test] + fn prune_rate_limit_buckets_drops_expired_clients() { + let mut map = HashMap::from([ + ( + "203.0.113.10".parse().unwrap(), + RateLimitBucket { + window_start: 10, + count: 2, + last_seen: 10, + }, + ), + ( + "203.0.113.11".parse().unwrap(), + RateLimitBucket { + window_start: 11, + count: 1, + last_seen: 11, + }, + ), + ]); + + prune_rate_limit_buckets(&mut map, 70, 60); + + assert!(!map.contains_key(&"203.0.113.10".parse::().unwrap())); + assert!(map.contains_key(&"203.0.113.11".parse::().unwrap())); + } + #[tokio::test] async fn gateway_rate_limits_per_client_ip() { let app = build_app(AppState::seeded(None).with_rate_limit(2, 60)); @@ -3674,6 +4022,208 @@ mod tests { assert_eq!(other.status(), StatusCode::OK); } + #[tokio::test] + async fn gateway_rejects_new_clients_when_local_limiter_is_full() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(2, 60) + .with_rate_limit_max_clients(1), + ); + + let first = app_request(&app, gateway_get_from_ip("/gateway/demo", "203.0.113.9")).await; + assert_eq!(first.status(), StatusCode::OK); + + let saturated = + app_request(&app, gateway_get_from_ip("/gateway/demo", "198.51.100.7")).await; + assert_eq!(saturated.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!( + saturated + .headers() + .get("retry-after") + .and_then(|value| value.to_str().ok()), + Some("60") + ); + + let body: serde_json::Value = json_body(saturated).await; + assert_eq!(body["action"], "rate_limiter_saturated"); + assert_eq!(body["reason"], RateLimitDecision::CAPACITY_EXCEEDED); + assert_eq!(body["max_clients"], 1); + assert_eq!(body["retry_after_seconds"], 60); + } + + #[tokio::test] + async fn gateway_ignores_forwarded_ip_from_untrusted_peer_for_rate_limiting() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(1, 60) + .with_rate_limit_max_clients(2), + ); + let first = app_request( + &app, + gateway_get_via_proxy("/gateway/demo", "203.0.113.9", "198.51.100.10"), + ) + .await; + assert_eq!(first.status(), StatusCode::OK); + + let blocked = app_request( + &app, + gateway_get_via_proxy("/gateway/demo", "203.0.113.11", "198.51.100.10"), + ) + .await; + assert_eq!(blocked.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn gateway_uses_forwarded_ip_from_trusted_proxy_for_rate_limiting() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(1, 60) + .with_trusted_proxies(HashSet::from(["198.51.100.10".parse().unwrap()])), + ); + let first = app_request( + &app, + gateway_get_via_proxy("/gateway/demo", "203.0.113.9", "198.51.100.10"), + ) + .await; + assert_eq!(first.status(), StatusCode::OK); + + let second = app_request( + &app, + gateway_get_via_proxy("/gateway/demo", "203.0.113.11", "198.51.100.10"), + ) + .await; + assert_eq!(second.status(), StatusCode::OK); + } + + #[tokio::test] + async fn gateway_uses_real_ip_from_trusted_proxy_when_forwarded_for_is_absent() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(1, 60) + .with_trusted_proxies(HashSet::from(["198.51.100.10".parse().unwrap()])), + ); + let mut first = Request::builder() + .method(Method::GET) + .uri("/gateway/demo") + .header("x-real-ip", "203.0.113.9") + .body(Body::empty()) + .unwrap(); + insert_peer(&mut first, "198.51.100.10".parse::().unwrap()); + let first = app_request(&app, first).await; + assert_eq!(first.status(), StatusCode::OK); + + let mut second = Request::builder() + .method(Method::GET) + .uri("/gateway/demo") + .header("x-real-ip", "203.0.113.11") + .body(Body::empty()) + .unwrap(); + insert_peer(&mut second, "198.51.100.10".parse::().unwrap()); + let second = app_request(&app, second).await; + assert_eq!(second.status(), StatusCode::OK); + } + + #[tokio::test] + async fn gateway_uses_rightmost_forwarded_hop_from_trusted_proxy() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(1, 60) + .with_trusted_proxies(HashSet::from(["198.51.100.10".parse().unwrap()])), + ); + let first = app_request( + &app, + gateway_get_via_proxy( + "/gateway/demo", + "203.0.113.200, 203.0.113.9", + "198.51.100.10", + ), + ) + .await; + assert_eq!(first.status(), StatusCode::OK); + + let blocked = app_request( + &app, + gateway_get_via_proxy( + "/gateway/demo", + "198.51.100.200, 203.0.113.9", + "198.51.100.10", + ), + ) + .await; + assert_eq!(blocked.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn gateway_walks_trusted_proxy_chain_right_to_left() { + let app = build_app( + AppState::seeded(None) + .with_rate_limit(1, 60) + .with_trusted_proxies(HashSet::from([ + "198.51.100.10".parse().unwrap(), + "198.51.100.11".parse().unwrap(), + ])), + ); + + let first = app_request( + &app, + gateway_get_via_proxy( + "/gateway/demo", + "203.0.113.9, 198.51.100.10", + "198.51.100.11", + ), + ) + .await; + assert_eq!(first.status(), StatusCode::OK); + + let second = app_request( + &app, + gateway_get_via_proxy( + "/gateway/demo", + "203.0.113.20, 198.51.100.10", + "198.51.100.11", + ), + ) + .await; + assert_eq!(second.status(), StatusCode::OK); + + let blocked = app_request( + &app, + gateway_get_via_proxy( + "/gateway/demo", + "203.0.113.9, 198.51.100.10", + "198.51.100.11", + ), + ) + .await; + assert_eq!(blocked.status(), StatusCode::TOO_MANY_REQUESTS); + } + + #[tokio::test] + async fn gateway_accepts_requests_without_connect_info() { + let app = build_app(AppState::seeded(None).with_rate_limit(2, 60)); + + let first = app + .clone() + .oneshot(empty_request(Method::GET, "/gateway/demo")) + .await + .unwrap(); + assert_eq!(first.status(), StatusCode::OK); + + let second = app + .clone() + .oneshot(empty_request(Method::GET, "/gateway/demo")) + .await + .unwrap(); + assert_eq!(second.status(), StatusCode::OK); + + let third = app + .clone() + .oneshot(empty_request(Method::GET, "/gateway/demo")) + .await + .unwrap(); + assert_eq!(third.status(), StatusCode::TOO_MANY_REQUESTS); + } + async fn body_text(response: Response) -> String { let bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap(); String::from_utf8(bytes.to_vec()).unwrap() @@ -4475,12 +5025,13 @@ mod tests { .await; assert_eq!(saved_dnsbl.code, "127.0.0.9"); - let gateway_request = Request::builder() + let mut gateway_request = Request::builder() .method(Method::POST) .uri("/gateway/secure/login?q=DROP%20TABLE") .header("x-forwarded-for", "198.51.100.7, 10.0.0.1") .body(Body::from("payload")) .unwrap(); + insert_peer(&mut gateway_request, "198.51.100.7".parse().unwrap()); let response = app_request(&app, gateway_request).await; assert_eq!(response.status(), StatusCode::FORBIDDEN); assert!(body_text(response).await.contains("\"action\":\"blocked\"")); @@ -7300,7 +7851,20 @@ mod tests { let mut headers = HeaderMap::new(); headers.insert("x-forwarded-for", HeaderValue::from_static("not-an-ip")); - assert_eq!(client_ip_from_headers(&headers), None); + assert_eq!( + client_ip_from_headers(&headers, None, &HashSet::new()), + None + ); + + let trusted_peer = "198.51.100.10".parse().unwrap(); + assert_eq!( + client_ip_from_headers( + &HeaderMap::new(), + Some(trusted_peer), + &HashSet::from([trusted_peer]), + ), + Some(trusted_peer) + ); let valid_path = temp_state_path("valid-load"); fs::write(