diff --git a/.github/workflows/backend-ci.yml b/.github/workflows/backend-ci.yml index da0ffde..091139f 100644 --- a/.github/workflows/backend-ci.yml +++ b/.github/workflows/backend-ci.yml @@ -92,6 +92,8 @@ jobs: PGPASSWORD=utility_secret psql -h localhost -U utility -d utility_test -f src/time_series/compress.sql PGPASSWORD=utility_secret psql -h localhost -U utility -d utility_test -f src/soroban/sync.sql PGPASSWORD=utility_secret psql -h localhost -U utility -d utility_test -f src/settlement/schema.sql + PGPASSWORD=utility_secret psql -h localhost -U utility -d utility_test -f db/rate_limits.sql + PGPASSWORD=utility_secret psql -h localhost -U utility -d utility_test -f db/audit_events.sql env: PGPASSWORD: utility_secret diff --git a/db/rate_limits.sql b/db/rate_limits.sql new file mode 100644 index 0000000..ee6c2e6 --- /dev/null +++ b/db/rate_limits.sql @@ -0,0 +1,13 @@ +CREATE TABLE IF NOT EXISTS rate_limit_configs ( + id BIGSERIAL PRIMARY KEY, + scope_type VARCHAR(16) NOT NULL CHECK (scope_type IN ('global', 'service', 'user')), + scope_key VARCHAR(255) NOT NULL, + max_tokens BIGINT NOT NULL CHECK (max_tokens > 0), + refill_rate BIGINT NOT NULL CHECK (refill_rate >= 0), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + UNIQUE (scope_type, scope_key) +); + +CREATE INDEX IF NOT EXISTS idx_rate_limit_configs_scope + ON rate_limit_configs (scope_type, scope_key); diff --git a/src/api/handlers.rs b/src/api/handlers.rs index b80f70e..52a0d59 100644 --- a/src/api/handlers.rs +++ b/src/api/handlers.rs @@ -773,3 +773,148 @@ pub async fn retry_dlq( ) -> Result, StatusCode> { retry_dlq_entry(pool, id).await } + +// --------------------------------------------------------------------------- +// Rate limit configuration handlers +// --------------------------------------------------------------------------- + +#[derive(Deserialize)] +pub struct RateLimitConfigListQuery { + pub scope_type: Option, +} + +#[derive(Deserialize)] +pub struct RateLimitAuditQuery { + pub limit: Option, +} + +fn map_rate_limit_error(error: crate::api::rate_limit_config::RateLimitConfigError) -> StatusCode { + match error { + crate::api::rate_limit_config::RateLimitConfigError::InvalidScopeKey + | crate::api::rate_limit_config::RateLimitConfigError::InvalidMaxTokens + | crate::api::rate_limit_config::RateLimitConfigError::InvalidRefillRate + | crate::api::rate_limit_config::RateLimitConfigError::InvalidGlobalScopeKey => { + StatusCode::BAD_REQUEST + } + crate::api::rate_limit_config::RateLimitConfigError::NotFound => StatusCode::NOT_FOUND, + crate::api::rate_limit_config::RateLimitConfigError::Conflict => StatusCode::CONFLICT, + crate::api::rate_limit_config::RateLimitConfigError::Database(_) => { + StatusCode::INTERNAL_SERVER_ERROR + } + } +} + +pub async fn list_rate_limit_configs( + State(pool): State>, + Query(query): Query, +) -> Result>, StatusCode> { + let scope_type = query + .scope_type + .as_deref() + .and_then(crate::api::rate_limit_config::RateLimitScopeType::parse); + if query.scope_type.is_some() && scope_type.is_none() { + return Err(StatusCode::BAD_REQUEST); + } + + crate::api::rate_limit_config::list_configs(&pool, scope_type) + .await + .map(Json) + .map_err(|e| { + tracing::error!(error = %e, "failed to list rate limit configs"); + StatusCode::INTERNAL_SERVER_ERROR + }) +} + +pub async fn get_rate_limit_config( + State(pool): State>, + Path(id): Path, +) -> Result, StatusCode> { + crate::api::rate_limit_config::get_config(&pool, id) + .await + .map_err(|e| { + tracing::error!(error = %e, "failed to get rate limit config"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .map(Json) + .ok_or(StatusCode::NOT_FOUND) +} + +pub async fn create_rate_limit_config( + State(state): State, + Json(body): Json, +) -> Result<(StatusCode, Json), StatusCode> { + crate::api::rate_limit_config::create_config( + &state.pool, + &state.rate_limiter, + &state.tenant_rate_limiter, + &state.service_rate_limiter, + body, + ) + .await + .map(|config| (StatusCode::CREATED, Json(config))) + .map_err(|e| { + tracing::error!(error = %e, "failed to create rate limit config"); + map_rate_limit_error(e) + }) +} + +pub async fn update_rate_limit_config( + State(state): State, + Path(id): Path, + Json(body): Json, +) -> Result, StatusCode> { + crate::api::rate_limit_config::update_config( + &state.pool, + &state.rate_limiter, + &state.tenant_rate_limiter, + &state.service_rate_limiter, + id, + body, + ) + .await + .map(Json) + .map_err(|e| { + tracing::error!(error = %e, "failed to update rate limit config"); + map_rate_limit_error(e) + }) +} + +#[derive(Deserialize)] +pub struct DeleteRateLimitConfigQuery { + #[serde(default = "crate::api::rate_limit_config::default_actor")] + pub actor: String, +} + +pub async fn delete_rate_limit_config( + State(state): State, + Path(id): Path, + Query(query): Query, +) -> Result { + crate::api::rate_limit_config::delete_config( + &state.pool, + &state.rate_limiter, + &state.tenant_rate_limiter, + &state.service_rate_limiter, + id, + &query.actor, + ) + .await + .map(|_| StatusCode::NO_CONTENT) + .map_err(|e| { + tracing::error!(error = %e, "failed to delete rate limit config"); + map_rate_limit_error(e) + }) +} + +pub async fn list_rate_limit_config_audit( + State(pool): State>, + Query(query): Query, +) -> Result>, StatusCode> { + crate::api::rate_limit_config::list_audit_entries(&pool, query.limit.unwrap_or(50)) + .await + .map(Json) + .map_err(|e| { + tracing::error!(error = %e, "failed to list rate limit audit entries"); + StatusCode::INTERNAL_SERVER_ERROR + }) +} diff --git a/src/api/middleware.rs b/src/api/middleware.rs index c4bd517..9b90da2 100644 --- a/src/api/middleware.rs +++ b/src/api/middleware.rs @@ -6,7 +6,7 @@ use axum::{ response::Response, }; use dashmap::DashMap; -use parking_lot::Mutex; +use parking_lot::{Mutex, RwLock}; use std::collections::VecDeque; use std::net::SocketAddr; use std::sync::atomic::{AtomicU64, Ordering}; @@ -128,7 +128,8 @@ pub struct FraudContext { } pub struct DynamicRateLimiter { - pub(crate) global_bucket: TokenBucket, + global_bucket: RwLock>, + default_source_limit: RwLock<(u64, u64)>, pub(crate) per_source_buckets: DashMap>, pub(crate) sliding_windows: DashMap>>, pub(crate) fraud_contexts: DashMap>>, @@ -136,10 +137,22 @@ pub struct DynamicRateLimiter { pub(crate) last_accessed: DashMap, } +const DEFAULT_GLOBAL_MAX_TOKENS: u64 = 10_000; +const DEFAULT_GLOBAL_REFILL_RATE: u64 = 10_000; +const DEFAULT_SOURCE_MAX_TOKENS: u64 = 100; +const DEFAULT_SOURCE_REFILL_RATE: u64 = 100; + impl DynamicRateLimiter { pub fn new() -> Arc { let limiter = Arc::new(Self { - global_bucket: TokenBucket::new(10000, 10000), + global_bucket: RwLock::new(Arc::new(TokenBucket::new( + DEFAULT_GLOBAL_MAX_TOKENS, + DEFAULT_GLOBAL_REFILL_RATE, + ))), + default_source_limit: RwLock::new(( + DEFAULT_SOURCE_MAX_TOKENS, + DEFAULT_SOURCE_REFILL_RATE, + )), per_source_buckets: DashMap::new(), sliding_windows: DashMap::new(), fraud_contexts: DashMap::new(), @@ -159,6 +172,20 @@ impl DynamicRateLimiter { limiter } + /// Hot-reload the global token bucket without restarting the service. + pub fn set_global_limit(&self, max_tokens: u64, refill_rate: u64) { + *self.global_bucket.write() = Arc::new(TokenBucket::new(max_tokens, refill_rate)); + } + + /// Restore the built-in global defaults after a config row is deleted. + pub fn reset_global_limit(&self) { + self.set_global_limit(DEFAULT_GLOBAL_MAX_TOKENS, DEFAULT_GLOBAL_REFILL_RATE); + } + + fn default_source_limit(&self) -> (u64, u64) { + *self.default_source_limit.read() + } + fn prune_inactive_sources(&self) { let now = Instant::now(); let timeout = Duration::from_secs(300); @@ -211,13 +238,19 @@ impl DynamicRateLimiter { }; // 2. Global rate limit - if !self.global_bucket.try_consume(1) { + let global_bucket = self.global_bucket.read().clone(); + if !global_bucket.try_consume(1) { self.increment_rejection("global"); return false; } // 3. Per-source rate limit - let limit = if is_flagged { 10 } else { 100 }; + let (default_max, default_refill) = self.default_source_limit(); + let limit = if is_flagged { + 10 + } else { + default_max.max(default_refill) + }; let bucket = { let b = self @@ -225,7 +258,7 @@ impl DynamicRateLimiter { .get(source_id) .map(|e| e.value().clone()); if let Some(b) = b { - if b.refill_rate == limit || source_id.starts_with("test-large-") { + if (b.max_tokens == limit && b.refill_rate == limit) || source_id.starts_with("test-large-") { b } else { let new_b = Arc::new(TokenBucket::new(limit, limit)); @@ -338,12 +371,20 @@ impl DynamicRateLimiter { } } +pub fn is_rate_limit_admin_path(path: &str) -> bool { + path.starts_with("/api/v1/rate-limits/configs") +} + pub async fn rate_limit_layer( State(limiter): State>, connect_info: Option>, req: Request, next: Next, ) -> Response { + if is_rate_limit_admin_path(req.uri().path()) { + return next.run(req).await; + } + let source_id = connect_info .map(|ConnectInfo(addr)| addr.ip().to_string()) .unwrap_or_else(|| "unknown".to_string()); @@ -418,18 +459,19 @@ mod tests { let source = "test-large-source"; // Use a very high global limit and per-source limit to ensure we can reach 1000 - limiter.global_bucket.tokens.store(100000, Ordering::SeqCst); + limiter.set_global_limit(100_000, 100_000); limiter .per_source_buckets .insert(source.to_string(), Arc::new(TokenBucket::new(2000, 2000))); for i in 0..1000 { let ok = limiter.try_consume(source); + let global_tokens = limiter.global_bucket.read().tokens.load(Ordering::SeqCst); assert!( ok, "Failed at request {} - Global: {}, Per: {}", i, - limiter.global_bucket.tokens.load(Ordering::SeqCst), + global_tokens, limiter .per_source_buckets .get(source) @@ -750,6 +792,10 @@ pub async fn tenant_rate_limit_layer( req: Request, next: Next, ) -> Response { + if is_rate_limit_admin_path(req.uri().path()) { + return next.run(req).await; + } + let tenant_id = req .headers() .get("x-tenant-id") @@ -766,6 +812,67 @@ pub async fn tenant_rate_limit_layer( next.run(req).await } +/// Per-service tier rate limiter (same token-bucket engine as tenant limiting). +#[derive(Clone)] +pub struct ServiceRateLimiter(Arc); + +impl ServiceRateLimiter { + pub fn new(default_max_tokens: u64, default_refill_rate: u64) -> Arc { + Arc::new(Self(TenantRateLimiter::new( + default_max_tokens, + default_refill_rate, + ))) + } + + pub fn inner(&self) -> &Arc { + &self.0 + } + + pub fn set_service_limit(&self, service_id: &str, limit: TenantLimit) { + self.0.set_tenant_limit(service_id, limit); + } + + pub fn remove_service_override(&self, service_id: &str) { + self.0.remove_tenant_override(service_id); + } + + pub fn try_consume(&self, service_id: &str, tokens: u64) -> bool { + self.0.try_consume(service_id, tokens) + } +} + +/// Axum middleware that enforces per-service tier rate limits. +/// +/// The service tier is identified via the `X-Service-ID` request header. When +/// the header is absent the request bypasses service-tier limiting. +pub async fn service_rate_limit_layer( + State(limiter): State>, + req: Request, + next: Next, +) -> Response { + if is_rate_limit_admin_path(req.uri().path()) { + return next.run(req).await; + } + + let Some(service_id) = req + .headers() + .get("x-service-id") + .and_then(|v| v.to_str().ok()) + .filter(|value| !value.is_empty()) + else { + return next.run(req).await; + }; + + if !limiter.try_consume(service_id, 1) { + warn!(service = %service_id, "service rate limit exceeded"); + return Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .body(Body::from("service rate limit exceeded")) + .unwrap(); + } + next.run(req).await +} + pub async fn slo_monitoring_layer(req: Request, next: Next) -> Response { let route = req.uri().path().to_string(); let started = Instant::now(); diff --git a/src/api/mod.rs b/src/api/mod.rs index a61d9d6..fdbdd7c 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,4 +1,4 @@ -use crate::api::middleware::{DynamicRateLimiter, TenantRateLimiter}; +use crate::api::middleware::{DynamicRateLimiter, ServiceRateLimiter, TenantRateLimiter}; use crate::gateway::hlc::HybridLogicalClock; use crate::gateway::lock::AdvisoryLock; use crate::soroban::rpc::CircuitBreaker; @@ -12,6 +12,7 @@ pub mod alloc_tracker; pub mod handlers; pub mod metrics; pub mod middleware; +pub mod rate_limit_config; pub mod router; pub mod slo_state; @@ -23,6 +24,7 @@ pub struct AppState { pub breaker: Arc>, pub rate_limiter: Arc, pub tenant_rate_limiter: Arc, + pub service_rate_limiter: Arc, pub hlc: Arc, } @@ -62,6 +64,12 @@ impl FromRef for Arc { } } +impl FromRef for Arc { + fn from_ref(state: &AppState) -> Self { + state.service_rate_limiter.clone() + } +} + impl FromRef for Arc { fn from_ref(state: &AppState) -> Self { state.hlc.clone() diff --git a/src/api/rate_limit_config.rs b/src/api/rate_limit_config.rs new file mode 100644 index 0000000..4621c0e --- /dev/null +++ b/src/api/rate_limit_config.rs @@ -0,0 +1,473 @@ +//! Rate limit configuration persistence and hot-reload. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::Arc; + +use crate::api::middleware::{DynamicRateLimiter, ServiceRateLimiter, TenantLimit, TenantRateLimiter}; +use crate::audit::store::audit_rate_limit_change_in_tx; + +pub const GLOBAL_SCOPE_KEY: &str = "_global"; +pub const AUDIT_ACTION_CREATE: &str = "rate_limit.create"; +pub const AUDIT_ACTION_UPDATE: &str = "rate_limit.update"; +pub const AUDIT_ACTION_DELETE: &str = "rate_limit.delete"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum RateLimitScopeType { + Global, + Service, + User, +} + +impl RateLimitScopeType { + pub fn as_str(self) -> &'static str { + match self { + Self::Global => "global", + Self::Service => "service", + Self::User => "user", + } + } + + pub fn parse(value: &str) -> Option { + match value { + "global" => Some(Self::Global), + "service" => Some(Self::Service), + "user" => Some(Self::User), + _ => None, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct RateLimitConfig { + pub id: i64, + pub scope_type: RateLimitScopeType, + pub scope_key: String, + pub max_tokens: i64, + pub refill_rate: i64, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct CreateRateLimitConfigRequest { + pub scope_type: RateLimitScopeType, + pub scope_key: Option, + pub max_tokens: i64, + pub refill_rate: i64, + #[serde(default = "default_actor")] + pub actor: String, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct UpdateRateLimitConfigRequest { + pub max_tokens: i64, + pub refill_rate: i64, + #[serde(default = "default_actor")] + pub actor: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RateLimitAuditEntry { + pub sequence: i64, + pub occurred_at: DateTime, + pub actor: String, + pub action: String, + pub resource: String, + pub payload_hash: String, +} + +pub fn default_actor() -> String { + "operator".to_string() +} + +#[derive(Debug, thiserror::Error)] +pub enum RateLimitConfigError { + #[error("invalid scope_key")] + InvalidScopeKey, + #[error("max_tokens must be between 1 and 1000000")] + InvalidMaxTokens, + #[error("refill_rate must be between 0 and 1000000")] + InvalidRefillRate, + #[error("global scope_key must be '{GLOBAL_SCOPE_KEY}'")] + InvalidGlobalScopeKey, + #[error("configuration not found")] + NotFound, + #[error("configuration already exists for this scope")] + Conflict, + #[error(transparent)] + Database(#[from] sqlx::Error), +} + +pub fn validate_create_request(req: &CreateRateLimitConfigRequest) -> Result { + validate_limits(req.max_tokens, req.refill_rate)?; + let scope_key = normalize_scope_key(req.scope_type, req.scope_key.as_deref())?; + Ok(scope_key) +} + +pub fn validate_update_request(req: &UpdateRateLimitConfigRequest) -> Result<(), RateLimitConfigError> { + validate_limits(req.max_tokens, req.refill_rate) +} + +fn validate_limits(max_tokens: i64, refill_rate: i64) -> Result<(), RateLimitConfigError> { + if !(1..=1_000_000).contains(&max_tokens) { + return Err(RateLimitConfigError::InvalidMaxTokens); + } + if !(0..=1_000_000).contains(&refill_rate) { + return Err(RateLimitConfigError::InvalidRefillRate); + } + Ok(()) +} + +fn normalize_scope_key( + scope_type: RateLimitScopeType, + scope_key: Option<&str>, +) -> Result { + match scope_type { + RateLimitScopeType::Global => { + match scope_key { + None | Some(GLOBAL_SCOPE_KEY) => Ok(GLOBAL_SCOPE_KEY.to_string()), + Some(_) => Err(RateLimitConfigError::InvalidGlobalScopeKey), + } + } + RateLimitScopeType::Service | RateLimitScopeType::User => { + let key = scope_key.unwrap_or("").trim(); + if key.is_empty() || key.len() > 255 || !key.chars().all(is_scope_char) { + return Err(RateLimitConfigError::InvalidScopeKey); + } + Ok(key.to_string()) + } + } +} + +fn is_scope_char(ch: char) -> bool { + ch.is_ascii_alphanumeric() || matches!(ch, '-' | '_' | '.') +} + +pub async fn ensure_schema(pool: &PgPool) -> Result<(), sqlx::Error> { + for statement in include_str!("../../db/rate_limits.sql").split(';') { + let statement = statement.trim(); + if !statement.is_empty() { + sqlx::query(statement).execute(pool).await?; + } + } + Ok(()) +} + +pub async fn list_configs( + pool: &PgPool, + scope_type: Option, +) -> Result, sqlx::Error> { + let rows = if let Some(scope) = scope_type { + sqlx::query_as::<_, RateLimitConfigRow>( + "SELECT id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at + FROM rate_limit_configs + WHERE scope_type = $1 + ORDER BY scope_type, scope_key", + ) + .bind(scope.as_str()) + .fetch_all(pool) + .await? + } else { + sqlx::query_as::<_, RateLimitConfigRow>( + "SELECT id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at + FROM rate_limit_configs + ORDER BY scope_type, scope_key", + ) + .fetch_all(pool) + .await? + }; + + Ok(rows.into_iter().map(RateLimitConfigRow::into_config).collect()) +} + +pub async fn get_config(pool: &PgPool, id: i64) -> Result, sqlx::Error> { + let row = sqlx::query_as::<_, RateLimitConfigRow>( + "SELECT id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at + FROM rate_limit_configs + WHERE id = $1", + ) + .bind(id) + .fetch_optional(pool) + .await?; + + Ok(row.map(RateLimitConfigRow::into_config)) +} + +pub async fn create_config( + pool: &PgPool, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, + req: CreateRateLimitConfigRequest, +) -> Result { + let scope_key = validate_create_request(&req)?; + let now = Utc::now(); + + let mut tx = pool.begin().await.map_err(RateLimitConfigError::Database)?; + + let row = sqlx::query_as::<_, RateLimitConfigRow>( + "INSERT INTO rate_limit_configs (scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at", + ) + .bind(req.scope_type.as_str()) + .bind(&scope_key) + .bind(req.max_tokens) + .bind(req.refill_rate) + .bind(now) + .bind(now) + .fetch_one(&mut *tx) + .await + .map_err(|e| { + if e.to_string().contains("duplicate key") { + RateLimitConfigError::Conflict + } else { + RateLimitConfigError::Database(e) + } + })?; + + let config = row.into_config(); + audit_rate_limit_change_in_tx( + &mut tx, + &req.actor, + AUDIT_ACTION_CREATE, + &resource_id(&config), + &config, + ) + .await + .map_err(RateLimitConfigError::Database)?; + + tx.commit().await.map_err(RateLimitConfigError::Database)?; + apply_config(&config, dynamic, tenant, service); + + Ok(config) +} + +pub async fn update_config( + pool: &PgPool, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, + id: i64, + req: UpdateRateLimitConfigRequest, +) -> Result { + validate_update_request(&req)?; + let now = Utc::now(); + + let mut tx = pool.begin().await.map_err(RateLimitConfigError::Database)?; + + let row = sqlx::query_as::<_, RateLimitConfigRow>( + "UPDATE rate_limit_configs + SET max_tokens = $1, refill_rate = $2, updated_at = $3 + WHERE id = $4 + RETURNING id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at", + ) + .bind(req.max_tokens) + .bind(req.refill_rate) + .bind(now) + .bind(id) + .fetch_optional(&mut *tx) + .await + .map_err(RateLimitConfigError::Database)?; + + let Some(row) = row else { + return Err(RateLimitConfigError::NotFound); + }; + + let config = row.into_config(); + audit_rate_limit_change_in_tx( + &mut tx, + &req.actor, + AUDIT_ACTION_UPDATE, + &resource_id(&config), + &config, + ) + .await + .map_err(RateLimitConfigError::Database)?; + + tx.commit().await.map_err(RateLimitConfigError::Database)?; + apply_config(&config, dynamic, tenant, service); + + Ok(config) +} + +pub async fn delete_config( + pool: &PgPool, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, + id: i64, + actor: &str, +) -> Result { + let mut tx = pool.begin().await.map_err(RateLimitConfigError::Database)?; + + let row = sqlx::query_as::<_, RateLimitConfigRow>( + "DELETE FROM rate_limit_configs + WHERE id = $1 + RETURNING id, scope_type, scope_key, max_tokens, refill_rate, created_at, updated_at", + ) + .bind(id) + .fetch_optional(&mut *tx) + .await + .map_err(RateLimitConfigError::Database)?; + + let Some(row) = row else { + return Err(RateLimitConfigError::NotFound); + }; + + let config = row.into_config(); + audit_rate_limit_change_in_tx( + &mut tx, + actor, + AUDIT_ACTION_DELETE, + &resource_id(&config), + &config, + ) + .await + .map_err(RateLimitConfigError::Database)?; + + tx.commit().await.map_err(RateLimitConfigError::Database)?; + remove_config(&config, dynamic, tenant, service); + + Ok(config) +} + +pub async fn list_audit_entries( + pool: &PgPool, + limit: i64, +) -> Result, sqlx::Error> { + let rows = sqlx::query_as::<_, (i64, DateTime, String, String, String, String)>( + "SELECT sequence, occurred_at, actor, action, resource, payload_hash + FROM audit_events + WHERE action LIKE 'rate_limit.%' + ORDER BY sequence DESC + LIMIT $1", + ) + .bind(limit.clamp(1, 500)) + .fetch_all(pool) + .await?; + + Ok(rows + .into_iter() + .map( + |(sequence, occurred_at, actor, action, resource, payload_hash)| RateLimitAuditEntry { + sequence, + occurred_at, + actor, + action, + resource, + payload_hash, + }, + ) + .collect()) +} + +pub async fn hydrate_from_db( + pool: &PgPool, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, +) -> Result<(), sqlx::Error> { + let configs = list_configs(pool, None).await?; + for config in configs { + apply_config(&config, dynamic, tenant, service); + } + Ok(()) +} + +pub fn apply_config( + config: &RateLimitConfig, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, +) { + let limit = TenantLimit::new(config.max_tokens as u64, config.refill_rate as u64); + match config.scope_type { + RateLimitScopeType::Global => { + dynamic.set_global_limit(config.max_tokens as u64, config.refill_rate as u64); + } + RateLimitScopeType::User => { + tenant.set_tenant_limit(&config.scope_key, limit); + } + RateLimitScopeType::Service => { + service.set_service_limit(&config.scope_key, limit); + } + } +} + +pub fn remove_config( + config: &RateLimitConfig, + dynamic: &Arc, + tenant: &Arc, + service: &Arc, +) { + match config.scope_type { + RateLimitScopeType::Global => dynamic.reset_global_limit(), + RateLimitScopeType::User => tenant.remove_tenant_override(&config.scope_key), + RateLimitScopeType::Service => service.remove_service_override(&config.scope_key), + } +} + +fn resource_id(config: &RateLimitConfig) -> String { + format!("rate_limit/{}/{}", config.scope_type.as_str(), config.scope_key) +} + +#[derive(sqlx::FromRow)] +struct RateLimitConfigRow { + id: i64, + scope_type: String, + scope_key: String, + max_tokens: i64, + refill_rate: i64, + created_at: DateTime, + updated_at: DateTime, +} + +impl RateLimitConfigRow { + fn into_config(self) -> RateLimitConfig { + RateLimitConfig { + id: self.id, + scope_type: RateLimitScopeType::parse(&self.scope_type).unwrap_or(RateLimitScopeType::User), + scope_key: self.scope_key, + max_tokens: self.max_tokens, + refill_rate: self.refill_rate, + created_at: self.created_at, + updated_at: self.updated_at, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn validates_global_scope_key() { + let req = CreateRateLimitConfigRequest { + scope_type: RateLimitScopeType::Global, + scope_key: None, + max_tokens: 100, + refill_rate: 10, + actor: "ops".into(), + }; + assert_eq!(validate_create_request(&req).unwrap(), GLOBAL_SCOPE_KEY); + } + + #[test] + fn rejects_invalid_user_scope_key() { + let req = CreateRateLimitConfigRequest { + scope_type: RateLimitScopeType::User, + scope_key: Some("".into()), + max_tokens: 100, + refill_rate: 10, + actor: "ops".into(), + }; + assert!(matches!( + validate_create_request(&req), + Err(RateLimitConfigError::InvalidScopeKey) + )); + } +} diff --git a/src/api/router.rs b/src/api/router.rs index 8a4119f..fa95ab7 100644 --- a/src/api/router.rs +++ b/src/api/router.rs @@ -57,6 +57,20 @@ pub async fn build_router(state: AppState) -> anyhow::Result { "/api/v1/tenant-rate-limiter/status", get(handlers::tenant_rate_limiter_status), ) + .route( + "/api/v1/rate-limits/configs", + get(handlers::list_rate_limit_configs).post(handlers::create_rate_limit_config), + ) + .route( + "/api/v1/rate-limits/configs/audit", + get(handlers::list_rate_limit_config_audit), + ) + .route( + "/api/v1/rate-limits/configs/:id", + get(handlers::get_rate_limit_config) + .put(handlers::update_rate_limit_config) + .delete(handlers::delete_rate_limit_config), + ) .route( "/api/v1/webhooks/endpoints", get(handlers::list_webhook_endpoints).post(handlers::create_webhook_endpoint), @@ -81,6 +95,10 @@ pub async fn build_router(state: AppState) -> anyhow::Result { state.clone(), crate::api::middleware::tenant_rate_limit_layer, )) + .layer(axum_mw::from_fn_with_state( + state.clone(), + crate::api::middleware::service_rate_limit_layer, + )) .layer(axum_mw::from_fn_with_state( state.clone(), crate::api::middleware::rate_limit_layer, diff --git a/src/audit.rs b/src/audit.rs index 3a90f42..f597394 100644 --- a/src/audit.rs +++ b/src/audit.rs @@ -4,6 +4,8 @@ use sha2::{Digest, Sha256}; use crate::api::metrics; +pub mod store; + pub const GENESIS_PREVIOUS_HASH: &str = "0000000000000000000000000000000000000000000000000000000000000000"; diff --git a/src/audit/store.rs b/src/audit/store.rs new file mode 100644 index 0000000..21d08e9 --- /dev/null +++ b/src/audit/store.rs @@ -0,0 +1,103 @@ +//! Persist tamper-evident audit events to PostgreSQL. + +use sqlx::PgPool; + +use super::{payload_hash, AuditEvent, NewAuditEvent, GENESIS_PREVIOUS_HASH}; + +const AUDIT_APPEND_LOCK_KEY: i64 = 4_000_000_002; + +/// Append a new audit event to the hash chain inside a short transaction. +pub async fn append_audit_event( + pool: &PgPool, + event: NewAuditEvent, +) -> Result { + let mut tx = pool.begin().await?; + let audit = append_audit_event_in_tx(&mut tx, event).await?; + tx.commit().await?; + Ok(audit) +} + +/// Append within an existing transaction so config mutations can commit atomically. +pub async fn append_audit_event_in_tx( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + event: NewAuditEvent, +) -> Result { + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(AUDIT_APPEND_LOCK_KEY) + .execute(&mut **tx) + .await?; + + let (previous_hash, next_sequence) = match sqlx::query_as::<_, (i64, String)>( + "SELECT sequence, hash FROM audit_events ORDER BY sequence DESC LIMIT 1", + ) + .fetch_optional(&mut **tx) + .await? + { + Some((seq, hash)) => (hash, seq + 1), + None => (GENESIS_PREVIOUS_HASH.to_string(), 1), + }; + + let audit = AuditEvent::append(next_sequence as u64, previous_hash, event); + + sqlx::query( + "INSERT INTO audit_events + (sequence, occurred_at, actor, service, action, resource, payload_hash, previous_hash, hash) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)", + ) + .bind(audit.sequence as i64) + .bind(audit.occurred_at) + .bind(&audit.actor) + .bind(&audit.service) + .bind(&audit.action) + .bind(&audit.resource) + .bind(&audit.payload_hash) + .bind(&audit.previous_hash) + .bind(&audit.hash) + .execute(&mut **tx) + .await?; + + Ok(audit) +} + +/// Convenience helper for rate-limit configuration mutations. +pub async fn audit_rate_limit_change( + pool: &PgPool, + actor: &str, + action: &str, + resource: &str, + payload: &T, +) -> Result { + append_audit_event( + pool, + NewAuditEvent { + occurred_at: chrono::Utc::now(), + actor: actor.to_string(), + service: "api".to_string(), + action: action.to_string(), + resource: resource.to_string(), + payload_hash: payload_hash(payload).map_err(|e| sqlx::Error::Decode(Box::new(e)))?, + }, + ) + .await +} + +pub async fn audit_rate_limit_change_in_tx( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + actor: &str, + action: &str, + resource: &str, + payload: &T, +) -> Result { + append_audit_event_in_tx( + tx, + NewAuditEvent { + occurred_at: chrono::Utc::now(), + actor: actor.to_string(), + service: "api".to_string(), + action: action.to_string(), + resource: resource.to_string(), + payload_hash: payload_hash(payload).map_err(|e| sqlx::Error::Decode(Box::new(e)))?, + }, + ) + .await +} diff --git a/src/main.rs b/src/main.rs index e451aa4..8ea8067 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,7 +4,7 @@ use std::time::Duration; use tokio::sync::Mutex; use tracing_subscriber::EnvFilter; -use utility_backend::api::middleware::{DynamicRateLimiter, TenantRateLimiter}; +use utility_backend::api::middleware::{DynamicRateLimiter, ServiceRateLimiter, TenantRateLimiter}; use utility_backend::api::AppState; use utility_backend::gateway::hlc::HybridLogicalClock; use utility_backend::gateway::lock::AdvisoryLock; @@ -83,6 +83,30 @@ async fn main() -> anyhow::Result<()> { let breaker = Arc::new(Mutex::new(CircuitBreaker::new(5))); let rate_limiter = DynamicRateLimiter::new(); let tenant_rate_limiter = TenantRateLimiter::new(1000, 1000); + let service_rate_limiter = ServiceRateLimiter::new(1000, 1000); + + if let Err(e) = utility_backend::api::rate_limit_config::ensure_schema(&db_pool).await { + tracing::warn!("rate limit config schema init failed: {}", e); + } + for statement in include_str!("../db/audit_events.sql").split(';') { + let statement = statement.trim(); + if !statement.is_empty() { + if let Err(e) = sqlx::query(statement).execute(&db_pool).await { + tracing::warn!("audit events schema init failed: {}", e); + } + } + } + if let Err(e) = utility_backend::api::rate_limit_config::hydrate_from_db( + &db_pool, + &rate_limiter, + &tenant_rate_limiter, + &service_rate_limiter, + ) + .await + { + tracing::warn!("failed to hydrate rate limit configs: {}", e); + } + let hlc = Arc::new(HybridLogicalClock::new()); let pd_client = utility_backend::incident::PagerDutyClient::from_env(); @@ -116,6 +140,7 @@ async fn main() -> anyhow::Result<()> { breaker, rate_limiter, tenant_rate_limiter, + service_rate_limiter, hlc, }; diff --git a/tests/dlq_tests.rs b/tests/dlq_tests.rs index 750cb40..4e8e4ef 100644 --- a/tests/dlq_tests.rs +++ b/tests/dlq_tests.rs @@ -226,6 +226,7 @@ async fn test_dlq_admin_api_endpoints() { breaker, rate_limiter, tenant_rate_limiter: utility_backend::api::middleware::TenantRateLimiter::new(100, 10), + service_rate_limiter: utility_backend::api::middleware::ServiceRateLimiter::new(100, 10), hlc: Arc::new(utility_backend::gateway::hlc::HybridLogicalClock::new()), }; diff --git a/tests/rate_limit_config_integration.rs b/tests/rate_limit_config_integration.rs new file mode 100644 index 0000000..5cd12eb --- /dev/null +++ b/tests/rate_limit_config_integration.rs @@ -0,0 +1,371 @@ +use axum::http::StatusCode; +use axum_test::TestServer; +use std::sync::Arc; +use tokio::sync::Mutex; +use utility_backend::api::middleware::{DynamicRateLimiter, ServiceRateLimiter, TenantRateLimiter}; +use utility_backend::api::rate_limit_config::{ + ensure_schema, RateLimitConfig, RateLimitScopeType, GLOBAL_SCOPE_KEY, +}; +use utility_backend::api::router::build_router; +use utility_backend::api::AppState; +use utility_backend::audit::verify_hash_chain; +use utility_backend::gateway::lock::AdvisoryLock; +use utility_backend::soroban::rpc::CircuitBreaker; + +async fn setup_test_db() -> Option<(sqlx::PgPool, sqlx::pool::PoolConnection)> { + let db_url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://utility:utility_secret@localhost:5432/utility_test".into()); + + match sqlx::PgPool::connect(&db_url).await { + Ok(pool) => { + let mut lock_conn = match pool.acquire().await { + Ok(conn) => conn, + Err(e) => { + eprintln!("Could not acquire DB lock connection: {e}"); + return None; + } + }; + if let Err(e) = sqlx::query("SELECT pg_advisory_lock(4000000003)") + .execute(&mut *lock_conn) + .await + { + eprintln!("Could not acquire DB advisory lock: {e}"); + return None; + } + + for statement in include_str!("../db/audit_events.sql").split(';') { + let statement = statement.trim(); + if !statement.is_empty() { + let _ = sqlx::query(statement).execute(&pool).await; + } + } + let _ = ensure_schema(&pool).await; + let _ = sqlx::query("DELETE FROM rate_limit_configs") + .execute(&pool) + .await; + let _ = sqlx::query("DELETE FROM audit_events") + .execute(&pool) + .await; + + Some((pool, lock_conn)) + } + Err(_) => { + eprintln!("Skipping integration test: DATABASE_URL not available"); + None + } + } +} + +async fn build_test_server( + pool: sqlx::PgPool, + tenant_limiter: Arc, + service_limiter: Arc, +) -> TestServer { + let state = AppState { + sequencer: Arc::new(utility_backend::soroban::sequencer::NonceSequencer::new()), + pool: pool.clone(), + advisory_lock: Arc::new(AdvisoryLock::postgres(pool)), + breaker: Arc::new(Mutex::new(CircuitBreaker::new(5))), + rate_limiter: DynamicRateLimiter::new(), + tenant_rate_limiter: tenant_limiter, + service_rate_limiter: service_limiter, + hlc: Arc::new(utility_backend::gateway::hlc::HybridLogicalClock::new()), + }; + + TestServer::new(build_router(state).await.unwrap()).unwrap() +} + +#[tokio::test] +async fn test_rate_limit_config_crud_and_audit() { + let Some((pool, _guard)) = setup_test_db().await else { + return; + }; + + let server = build_test_server( + pool.clone(), + TenantRateLimiter::new(1000, 1000), + ServiceRateLimiter::new(1000, 1000), + ) + .await; + + let create_body = serde_json::json!({ + "scope_type": "user", + "scope_key": "tenant-alpha", + "max_tokens": 3, + "refill_rate": 0, + "actor": "ops@example.com" + }); + + let response = server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await; + response.assert_status(StatusCode::CREATED); + let created: RateLimitConfig = response.json(); + assert_eq!(created.scope_type, RateLimitScopeType::User); + assert_eq!(created.scope_key, "tenant-alpha"); + assert_eq!(created.max_tokens, 3); + + let response = server.get("/api/v1/rate-limits/configs").await; + response.assert_status_ok(); + let list: Vec = response.json(); + assert_eq!(list.len(), 1); + + let response = server + .get(&format!("/api/v1/rate-limits/configs/{}", created.id)) + .await; + response.assert_status_ok(); + + let update_body = serde_json::json!({ + "max_tokens": 5, + "refill_rate": 0, + "actor": "ops@example.com" + }); + let response = server + .put(&format!("/api/v1/rate-limits/configs/{}", created.id)) + .json(&update_body) + .await; + response.assert_status_ok(); + let updated: RateLimitConfig = response.json(); + assert_eq!(updated.max_tokens, 5); + + let response = server + .delete(&format!( + "/api/v1/rate-limits/configs/{}?actor=ops@example.com", + created.id + )) + .await; + response.assert_status(StatusCode::NO_CONTENT); + + let response = server + .get(&format!("/api/v1/rate-limits/configs/{}", created.id)) + .await; + response.assert_status(StatusCode::NOT_FOUND); + + let response = server.get("/api/v1/rate-limits/configs/audit?limit=10").await; + response.assert_status_ok(); + let audit_entries: Vec = + response.json(); + assert_eq!(audit_entries.len(), 3); + assert!(audit_entries.iter().any(|entry| entry.action == "rate_limit.create")); + assert!(audit_entries.iter().any(|entry| entry.action == "rate_limit.update")); + assert!(audit_entries.iter().any(|entry| entry.action == "rate_limit.delete")); + + let rows = sqlx::query_as::< + _, + ( + i64, + chrono::DateTime, + String, + String, + String, + String, + String, + String, + String, + ), + >( + "SELECT sequence, occurred_at, actor, service, action, resource, payload_hash, previous_hash, hash + FROM audit_events + ORDER BY sequence", + ) + .fetch_all(&pool) + .await + .unwrap(); + + let events: Vec<_> = rows + .into_iter() + .map( + |( + sequence, + occurred_at, + actor, + service, + action, + resource, + payload_hash, + previous_hash, + hash, + )| { + utility_backend::audit::AuditEvent { + sequence: sequence as u64, + occurred_at, + actor, + service, + action, + resource, + payload_hash, + previous_hash, + hash, + } + }, + ) + .collect(); + let report = verify_hash_chain(&events); + assert!(report.verified, "{:?}", report.reason); +} + +#[tokio::test] +async fn test_user_rate_limit_hot_reload_via_api() { + let Some((pool, _guard)) = setup_test_db().await else { + return; + }; + + let server = build_test_server( + pool, + TenantRateLimiter::new(100, 0), + ServiceRateLimiter::new(100, 0), + ) + .await; + + let create_body = serde_json::json!({ + "scope_type": "user", + "scope_key": "grid-east", + "max_tokens": 2, + "refill_rate": 0, + "actor": "ops" + }); + server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await + .assert_status(StatusCode::CREATED); + + for _ in 0..2 { + let response = server + .get("/health") + .add_header("x-tenant-id", "grid-east") + .await; + response.assert_status_ok(); + } + + let response = server + .get("/health") + .add_header("x-tenant-id", "grid-east") + .await; + response.assert_status(StatusCode::TOO_MANY_REQUESTS); + + let response = server + .get("/health") + .add_header("x-tenant-id", "grid-west") + .await; + response.assert_status_ok(); +} + +#[tokio::test] +async fn test_global_rate_limit_hot_reload_via_api() { + let Some((pool, _guard)) = setup_test_db().await else { + return; + }; + + let state = AppState { + sequencer: Arc::new(utility_backend::soroban::sequencer::NonceSequencer::new()), + pool: pool.clone(), + advisory_lock: Arc::new(AdvisoryLock::postgres(pool.clone())), + breaker: Arc::new(Mutex::new(CircuitBreaker::new(5))), + rate_limiter: DynamicRateLimiter::new(), + tenant_rate_limiter: TenantRateLimiter::new(10_000, 10_000), + service_rate_limiter: ServiceRateLimiter::new(10_000, 10_000), + hlc: Arc::new(utility_backend::gateway::hlc::HybridLogicalClock::new()), + }; + let server = TestServer::new(build_router(state).await.unwrap()).unwrap(); + + let create_body = serde_json::json!({ + "scope_type": "global", + "max_tokens": 1, + "refill_rate": 0, + "actor": "ops" + }); + server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await + .assert_status(StatusCode::CREATED); + + let response = server.get("/health").await; + response.assert_status_ok(); + + let response = server.get("/health").await; + response.assert_status(StatusCode::TOO_MANY_REQUESTS); +} + +#[tokio::test] +async fn test_service_rate_limit_hot_reload_via_api() { + let Some((pool, _guard)) = setup_test_db().await else { + return; + }; + + let server = build_test_server( + pool, + TenantRateLimiter::new(100, 100), + ServiceRateLimiter::new(100, 0), + ) + .await; + + let create_body = serde_json::json!({ + "scope_type": "service", + "scope_key": "readings", + "max_tokens": 2, + "refill_rate": 0, + "actor": "ops" + }); + server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await + .assert_status(StatusCode::CREATED); + + for _ in 0..2 { + let response = server + .get("/health") + .add_header("x-service-id", "readings") + .await; + response.assert_status_ok(); + } + + let response = server + .get("/health") + .add_header("x-service-id", "readings") + .await; + response.assert_status(StatusCode::TOO_MANY_REQUESTS); +} + +#[tokio::test] +async fn test_global_scope_key_validation() { + let Some((pool, _guard)) = setup_test_db().await else { + return; + }; + + let server = build_test_server( + pool, + TenantRateLimiter::new(100, 100), + ServiceRateLimiter::new(100, 100), + ) + .await; + + let create_body = serde_json::json!({ + "scope_type": "global", + "scope_key": "wrong-key", + "max_tokens": 10, + "refill_rate": 1, + "actor": "ops" + }); + server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await + .assert_status(StatusCode::BAD_REQUEST); + + let create_body = serde_json::json!({ + "scope_type": "global", + "scope_key": GLOBAL_SCOPE_KEY, + "max_tokens": 10, + "refill_rate": 1, + "actor": "ops" + }); + server + .post("/api/v1/rate-limits/configs") + .json(&create_body) + .await + .assert_status(StatusCode::CREATED); +}