diff --git a/.env.example b/.env.example index 2d8ccd6..8bc338e 100644 --- a/.env.example +++ b/.env.example @@ -72,3 +72,8 @@ OMEM_EMBED_PROVIDER=noop # OMEM_LLM_PROVIDER=bedrock # OMEM_LLM_MODEL=anthropic.claude-3-haiku-20240307-v1:0 # AWS_REGION=us-east-1 + +# ─── Sharing rate limit ────────────────────────────────────────────────────── +# Per-user sharing operations allowed per minute (0 = disabled). Bursts up to +# this value are allowed, then refill at OMEM_SHARE_RATE_PER_MIN/60 per second. +OMEM_SHARE_RATE_PER_MIN=0 diff --git a/docs/SHARING.md b/docs/SHARING.md index b9c76a9..4908932 100644 --- a/docs/SHARING.md +++ b/docs/SHARING.md @@ -828,9 +828,9 @@ LanceDB doesn't support cross-database vector queries. Each space has its own ve `POST /v1/memories/batch-share` accepts at most 500 memory IDs per call. Requests exceeding this limit return 400 Bad Request. -### No rate limiting on sharing +### Rate limiting on sharing (opt-in) -There is no per-user or per-space rate limit on sharing operations. A user with write access can share thousands of memories in rapid succession. Rate limiting is a separate hardening task. +Per-user rate limiting on sharing operations is available via `OMEM_SHARE_RATE_PER_MIN` (operations per minute per user; **default 0 = disabled**). When enabled, each call to a sharing endpoint (`share`, `pull`, `reshare`, `batch-share`, `share-all`, `share-to-user`, `share-all-to-user`, `org/publish`) consumes one token from the caller's bucket; an empty bucket returns `429 Too Many Requests`. Tokens refill continuously at `OMEM_SHARE_RATE_PER_MIN / 60` per second, so bursts up to the per-minute value are allowed. State is process-local; a multi-replica deployment would need a shared store. ### Organization spaces are read-only for non-admins diff --git a/omem-server/src/api/handlers/sharing.rs b/omem-server/src/api/handlers/sharing.rs index 11fea67..8d64ba1 100644 --- a/omem-server/src/api/handlers/sharing.rs +++ b/omem-server/src/api/handlers/sharing.rs @@ -295,6 +295,7 @@ pub async fn share_memory( Path(id): Path, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; if body.target_space.is_empty() { return Err(OmemError::Validation( "target_space is required".to_string(), @@ -353,6 +354,7 @@ pub async fn pull_memory( Path(id): Path, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; if body.source_space.is_empty() { return Err(OmemError::Validation( "source_space is required".to_string(), @@ -479,6 +481,7 @@ pub async fn batch_share( Extension(auth): Extension, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; if body.memory_ids.is_empty() { return Err(OmemError::Validation( "memory_ids cannot be empty".to_string(), @@ -698,6 +701,7 @@ pub async fn reshare_memory( Path(id): Path, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; let spaces = state .space_store .list_spaces_for_user(&auth.tenant_id) @@ -882,6 +886,7 @@ pub async fn share_all( Extension(auth): Extension, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; let target_space_id = normalize_space_id(&body.target_space); if target_space_id.is_empty() { return Err(OmemError::Validation( @@ -982,6 +987,7 @@ pub async fn share_to_user( Path(id): Path, Json(body): Json, ) -> Result { + state.share_rate_limiter.check(&auth.tenant_id)?; if body.target_user.is_empty() { return Err(OmemError::Validation("target_user is required".to_string())); } @@ -1051,6 +1057,7 @@ pub async fn share_all_to_user( Extension(auth): Extension, Json(body): Json, ) -> Result, OmemError> { + state.share_rate_limiter.check(&auth.tenant_id)?; if body.target_user.is_empty() { return Err(OmemError::Validation("target_user is required".to_string())); } @@ -1214,6 +1221,7 @@ pub async fn org_publish( Path(org_id): Path, Json(body): Json, ) -> Result, OmemError> { + state.share_rate_limiter.check(&auth.tenant_id)?; let org_id = normalize_space_id(&org_id); let mut space = state .space_store diff --git a/omem-server/src/api/handlers/stats.rs b/omem-server/src/api/handlers/stats.rs index 89f57ce..18bd296 100644 --- a/omem-server/src/api/handlers/stats.rs +++ b/omem-server/src/api/handlers/stats.rs @@ -830,6 +830,7 @@ mod tests { config: OmemConfig::default(), import_semaphore: Arc::new(tokio::sync::Semaphore::new(3)), reconcile_semaphore: Arc::new(tokio::sync::Semaphore::new(1)), + share_rate_limiter: Arc::new(crate::api::rate_limit::RateLimiter::new(0)), }); (state, store_dir, space_dir, tenant_dir) diff --git a/omem-server/src/api/mod.rs b/omem-server/src/api/mod.rs index 6e0622a..722609c 100644 --- a/omem-server/src/api/mod.rs +++ b/omem-server/src/api/mod.rs @@ -1,6 +1,7 @@ pub mod error; pub mod handlers; pub mod middleware; +pub mod rate_limit; pub mod router; pub mod server; @@ -86,6 +87,7 @@ mod tests { config: OmemConfig::default(), import_semaphore: Arc::new(tokio::sync::Semaphore::new(3)), reconcile_semaphore: Arc::new(tokio::sync::Semaphore::new(1)), + share_rate_limiter: Arc::new(crate::api::rate_limit::RateLimiter::new(0)), }); (build_router(state), dir) diff --git a/omem-server/src/api/rate_limit.rs b/omem-server/src/api/rate_limit.rs new file mode 100644 index 0000000..c38fe98 --- /dev/null +++ b/omem-server/src/api/rate_limit.rs @@ -0,0 +1,111 @@ +//! In-memory per-user token-bucket rate limiter for sharing operations. +//! +//! Sharing has no built-in rate limit, so one user with write access can fire +//! thousands of share operations in rapid succession. This adds an opt-in +//! per-user limit: each sharing call consumes one token; a user gets `per_min` +//! tokens that refill continuously at `per_min / 60` per second, so bursts up +//! to `per_min` are allowed and then a steady rate. `per_min == 0` disables it. +//! +//! State is process-local (`Mutex`). A multi-replica deployment would +//! want a shared store (Redis, etc.); for single-instance omem this suffices. + +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::Instant; + +use crate::domain::error::OmemError; + +struct Bucket { + tokens: f64, + last: Instant, +} + +pub struct RateLimiter { + per_min: u32, + buckets: Mutex>, +} + +impl RateLimiter { + pub fn new(per_min: u32) -> Self { + Self { + per_min, + buckets: Mutex::new(HashMap::new()), + } + } + + /// Charge one token to `key`. Returns `Err(RateLimited)` if the bucket is + /// empty. Always `Ok` when disabled (`per_min == 0`). + pub fn check(&self, key: &str) -> Result<(), OmemError> { + self.check_at(key, Instant::now()) + } + + /// Time-injectable core, for deterministic tests. + fn check_at(&self, key: &str, now: Instant) -> Result<(), OmemError> { + if self.per_min == 0 { + return Ok(()); + } + let cap = self.per_min as f64; + let refill_per_sec = cap / 60.0; + let mut buckets = self.buckets.lock().expect("rate limiter mutex poisoned"); + let bucket = buckets.entry(key.to_string()).or_insert(Bucket { + tokens: cap, + last: now, + }); + let elapsed = now.saturating_duration_since(bucket.last).as_secs_f64(); + bucket.tokens = (bucket.tokens + elapsed * refill_per_sec).min(cap); + bucket.last = now; + if bucket.tokens >= 1.0 { + bucket.tokens -= 1.0; + Ok(()) + } else { + Err(OmemError::RateLimited) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn disabled_never_limits() { + let rl = RateLimiter::new(0); + for _ in 0..10_000 { + assert!(rl.check("u1").is_ok()); + } + } + + #[test] + fn allows_burst_then_blocks() { + let rl = RateLimiter::new(60); + let t0 = Instant::now(); + for _ in 0..60 { + assert!(rl.check_at("u1", t0).is_ok()); + } + assert!(matches!(rl.check_at("u1", t0), Err(OmemError::RateLimited))); + } + + #[test] + fn refills_over_time() { + let rl = RateLimiter::new(60); // 1 token/sec + let t0 = Instant::now(); + for _ in 0..60 { + let _ = rl.check_at("u1", t0); + } + assert!(rl.check_at("u1", t0).is_err()); + let t2 = t0 + Duration::from_secs(2); + assert!(rl.check_at("u1", t2).is_ok()); + assert!(rl.check_at("u1", t2).is_ok()); + assert!(rl.check_at("u1", t2).is_err()); + } + + #[test] + fn keys_are_independent() { + let rl = RateLimiter::new(1); + let t0 = Instant::now(); + assert!(rl.check_at("u1", t0).is_ok()); + assert!(rl.check_at("u1", t0).is_err()); + assert!(rl.check_at("u2", t0).is_ok()); + } +} diff --git a/omem-server/src/api/server.rs b/omem-server/src/api/server.rs index 5c5bcbc..78517cf 100644 --- a/omem-server/src/api/server.rs +++ b/omem-server/src/api/server.rs @@ -16,6 +16,7 @@ pub struct AppState { pub config: OmemConfig, pub import_semaphore: Arc, pub reconcile_semaphore: Arc, + pub share_rate_limiter: Arc, } /// Map tenant_id to their personal Space ID. diff --git a/omem-server/src/config.rs b/omem-server/src/config.rs index fbcde0f..b4d3913 100644 --- a/omem-server/src/config.rs +++ b/omem-server/src/config.rs @@ -16,6 +16,8 @@ pub struct OmemConfig { pub embed_model: String, pub embed_dim: usize, pub embed_timeout_secs: u64, + /// Per-user sharing rate limit (operations per minute). 0 = disabled. + pub share_rate_per_min: u32, } impl Default for OmemConfig { @@ -35,6 +37,7 @@ impl Default for OmemConfig { embed_model: String::new(), embed_dim: 1024, embed_timeout_secs: 10, + share_rate_per_min: 0, } } } @@ -66,6 +69,10 @@ impl OmemConfig { .ok() .and_then(|v| v.parse().ok()) .unwrap_or(defaults.embed_timeout_secs), + share_rate_per_min: env::var("OMEM_SHARE_RATE_PER_MIN") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(defaults.share_rate_per_min), } } diff --git a/omem-server/src/main.rs b/omem-server/src/main.rs index d46debc..a6c28d8 100644 --- a/omem-server/src/main.rs +++ b/omem-server/src/main.rs @@ -89,6 +89,9 @@ async fn main() { config: config.clone(), import_semaphore: Arc::new(tokio::sync::Semaphore::new(3)), reconcile_semaphore: Arc::new(tokio::sync::Semaphore::new(1)), + share_rate_limiter: Arc::new(omem_server::api::rate_limit::RateLimiter::new( + config.share_rate_per_min, + )), }); let app = build_router(state);