diff --git a/README.md b/README.md index 5d82d8d..3cf0a69 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,25 @@ history are untouched. `propose_admin_transfer` rejects proposing the current admin as the new admin (panics with `InvalidAdminProposal`). This surfaces no-op handovers as caller mistakes rather than silently storing a pending entry equal to the active admin. +### Per-agent rate limiting (fixed window) + +`record_usage` supports an optional per-agent rate limit anchored to +`env.ledger().timestamp()`. It is configured by two admin settings and is +**disabled by default** (both default to `0`): + +- `set_max_requests_per_window(max)` — max `requests` an agent may accumulate + per window (`get_max_requests_per_window`). +- `set_rate_window_seconds(seconds)` — the **fixed** window length + (`get_rate_window_seconds`). + +The limiter is active only when **both** are non-zero. Semantics are a +**fixed window** (not sliding): the window opens at an agent's first in-window +call and rolls forward as a whole once `now >= window_start + window_seconds`, +resetting the count. A call that would push the in-window count above the cap +is rejected with `RateLimitExceeded` (#15). State is per-agent +(`DataKey::RateWindow(agent)`), and an agent can never reset its own window +early — `window_start` only advances. Window arithmetic is saturating. + ### Schema version: fresh v2 init vs. legacy v1→v2 migration `init` stamps the current storage schema version (v2) directly, so a freshly diff --git a/contracts/escrow/src/lib.rs b/contracts/escrow/src/lib.rs index 82ccfaf..55097f5 100644 --- a/contracts/escrow/src/lib.rs +++ b/contracts/escrow/src/lib.rs @@ -81,6 +81,16 @@ pub enum DataKey { /// disabled without unregistering, preserving the metadata and the /// per-(agent, service) usage history. ServiceDisabled(Symbol), + /// Max `requests` an agent may accumulate within one rate-limit + /// window. `0` (the default) disables the limiter entirely. + MaxRequestsPerWindow, + /// Length of the fixed rate-limit window in seconds. `0` (the + /// default) disables the limiter entirely. + WindowSeconds, + /// Per-agent fixed-window rate-limit state: `(window_start, count)` + /// where `window_start` is the ledger timestamp the current window + /// opened and `count` is the requests accumulated in it so far. + RateWindow(Address), } /// Typed contract errors. Codes are append-only to keep client SDKs stable. @@ -124,6 +134,9 @@ pub enum EscrowError { /// proposed new admin — a no-op handover that is rejected to surface /// caller mistakes early. InvalidAdminProposal = 14, + /// `record_usage` would push the agent's per-window request count + /// above the configured `MaxRequestsPerWindow` cap. + RateLimitExceeded = 15, } #[contracttype] @@ -230,6 +243,45 @@ impl Escrow { { panic_with_error!(&env, EscrowError::AgentNotAllowed); } + + // Per-agent fixed-window rate limit. Enabled only when both the cap + // and the window length are non-zero. The window is anchored at the + // first in-window call's timestamp and rolls forward whole-window + // (fixed, not sliding) once `now >= window_start + window_seconds`. + let max_per_window: u32 = env + .storage() + .persistent() + .get(&DataKey::MaxRequestsPerWindow) + .unwrap_or(0); + let window_seconds: u64 = env + .storage() + .persistent() + .get(&DataKey::WindowSeconds) + .unwrap_or(0); + if max_per_window > 0 && window_seconds > 0 { + let now = env.ledger().timestamp(); + let (window_start, count): (u64, u32) = env + .storage() + .persistent() + .get(&DataKey::RateWindow(agent.clone())) + .unwrap_or((0, 0)); + // Roll the window forward if the current one has expired. The + // agent can never reset it early: window_start only advances. + let (window_start, count) = if now >= window_start.saturating_add(window_seconds) { + (now, 0u32) + } else { + (window_start, count) + }; + let new_count = count.saturating_add(requests); + if new_count > max_per_window { + panic_with_error!(&env, EscrowError::RateLimitExceeded); + } + env.storage().persistent().set( + &DataKey::RateWindow(agent.clone()), + &(window_start, new_count), + ); + } + let key = DataKey::Usage(agent.clone(), service_id.clone()); let prev: u32 = env.storage().persistent().get(&key).unwrap_or(0); let total = prev.saturating_add(requests); @@ -454,6 +506,55 @@ impl Escrow { .unwrap_or(u32::MAX) } + /// Read the configured per-window request cap, or `0` (limiter + /// disabled) when unset. + pub fn get_max_requests_per_window(env: Env) -> u32 { + env.storage() + .persistent() + .get(&DataKey::MaxRequestsPerWindow) + .unwrap_or(0) + } + + /// Admin sets the per-agent, per-window request cap. The limiter is + /// active only when both this cap and the window length + /// ([`Self::set_rate_window_seconds`]) are non-zero. Pass `0` to + /// disable. + pub fn set_max_requests_per_window(env: Env, max_requests: u32) { + let admin: Address = env + .storage() + .persistent() + .get(&DataKey::Admin) + .unwrap_or_else(|| panic_with_error!(&env, EscrowError::NotInitialized)); + admin.require_auth(); + env.storage() + .persistent() + .set(&DataKey::MaxRequestsPerWindow, &max_requests); + } + + /// Read the configured rate-limit window length in seconds, or `0` + /// (limiter disabled) when unset. + pub fn get_rate_window_seconds(env: Env) -> u64 { + env.storage() + .persistent() + .get(&DataKey::WindowSeconds) + .unwrap_or(0) + } + + /// Admin sets the fixed rate-limit window length in seconds. The + /// limiter is active only when both this and the per-window cap are + /// non-zero. Pass `0` to disable. + pub fn set_rate_window_seconds(env: Env, window_seconds: u64) { + let admin: Address = env + .storage() + .persistent() + .get(&DataKey::Admin) + .unwrap_or_else(|| panic_with_error!(&env, EscrowError::NotInitialized)); + admin.require_auth(); + env.storage() + .persistent() + .set(&DataKey::WindowSeconds, &window_seconds); + } + /// Admin sets the per-call upper bound on `requests` accepted by /// `record_usage`. Pass `u32::MAX` to effectively disable the cap. pub fn set_max_requests_per_call(env: Env, max_requests: u32) { diff --git a/contracts/escrow/src/test.rs b/contracts/escrow/src/test.rs index 2d924db..291fbbd 100644 --- a/contracts/escrow/src/test.rs +++ b/contracts/escrow/src/test.rs @@ -697,3 +697,91 @@ fn test_pause_pause_unpause_ends_unpaused() { assert!(!client.is_paused()); } + +// --------------------------------------------------------------------------- +// Per-agent, per-window rate limiting (#10) +// --------------------------------------------------------------------------- + +/// By default the limiter is disabled (cap 0, window 0): an agent can record +/// far more than any cap would allow. +#[test] +fn test_rate_limit_disabled_by_default() { + let env = Env::default(); + let (client, _admin) = setup_initialized(&env); + assert_eq!(client.get_max_requests_per_window(), 0); + assert_eq!(client.get_rate_window_seconds(), 0); + + let agent = Address::generate(&env); + let svc = Symbol::new(&env, "infer"); + for _ in 0..50 { + client.record_usage(&agent, &svc, &100u32); + } + assert_eq!(client.get_usage(&agent, &svc), 5_000); +} + +/// Config setters round-trip. +#[test] +fn test_rate_limit_config_round_trips() { + let env = Env::default(); + let (client, _admin) = setup_initialized(&env); + client.set_max_requests_per_window(&10u32); + client.set_rate_window_seconds(&60u64); + assert_eq!(client.get_max_requests_per_window(), 10); + assert_eq!(client.get_rate_window_seconds(), 60); +} + +/// Accumulating exactly up to the cap is allowed; one more request in the +/// same window is rejected with RateLimitExceeded (#15). +#[test] +#[should_panic(expected = "Error(Contract, #15)")] +fn test_rate_limit_rejects_over_cap_in_window() { + let env = Env::default(); + env.ledger().with_mut(|li| li.timestamp = 1_000); + let (client, _admin) = setup_initialized(&env); + client.set_max_requests_per_window(&10u32); + client.set_rate_window_seconds(&100u64); + + let agent = Address::generate(&env); + let svc = Symbol::new(&env, "infer"); + client.record_usage(&agent, &svc, &6u32); // count = 6 + client.record_usage(&agent, &svc, &4u32); // count = 10 (exactly at cap) + client.record_usage(&agent, &svc, &1u32); // count = 11 → reject #15 +} + +/// After the window expires the counter resets and the agent can record +/// again (fixed-window rollover). +#[test] +fn test_rate_limit_window_rollover_resets_count() { + let env = Env::default(); + env.ledger().with_mut(|li| li.timestamp = 1_000); + let (client, _admin) = setup_initialized(&env); + client.set_max_requests_per_window(&10u32); + client.set_rate_window_seconds(&100u64); + + let agent = Address::generate(&env); + let svc = Symbol::new(&env, "infer"); + client.record_usage(&agent, &svc, &10u32); // fills the window + + // Advance past the window; the count resets. + env.ledger().with_mut(|li| li.timestamp = 1_100); + let rec = client.record_usage(&agent, &svc, &10u32); + // Usage is cumulative (20), but the rate window accepted the new 10. + assert_eq!(rec.requests, 20); +} + +/// The limiter is per-agent: one agent hitting the cap does not block another. +#[test] +fn test_rate_limit_is_per_agent() { + let env = Env::default(); + env.ledger().with_mut(|li| li.timestamp = 1_000); + let (client, _admin) = setup_initialized(&env); + client.set_max_requests_per_window(&5u32); + client.set_rate_window_seconds(&100u64); + + let a = Address::generate(&env); + let b = Address::generate(&env); + let svc = Symbol::new(&env, "infer"); + client.record_usage(&a, &svc, &5u32); // a at cap + let rec_b = client.record_usage(&b, &svc, &5u32); // b independent + assert_eq!(rec_b.requests, 5); +}