Something went wrong. Try again.
A lexicon-driven AppView for ATProto. forked from gamesgamesgamesgames.games/happyview
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356use dashmap::DashMap;use std::sync::Arc;use std::time::{Instant, SystemTime, UNIX_EPOCH};
/// Hardcoded seed values for the per-instance default token costs. These are/// only used by the startup seeding step in `main.rs` to populate fresh/// `instance_settings` rows; at runtime the values are read from the DB into/// `RateLimitDefaults`.pub const SEED_DEFAULT_QUERY_COST: u32 = 1;pub const SEED_DEFAULT_PROCEDURE_COST: u32 = 1;pub const SEED_DEFAULT_PROXY_COST: u32 = 1;
/// `instance_settings` keys for the seeded defaults.pub const SETTING_DEFAULT_QUERY_COST: &str = "rate_limit.default_query_cost";pub const SETTING_DEFAULT_PROCEDURE_COST: &str = "rate_limit.default_procedure_cost";pub const SETTING_DEFAULT_PROXY_COST: &str = "rate_limit.default_proxy_cost";
/// Default token costs per XRPC request type, loaded from `instance_settings`/// at startup. Owned by the `RateLimiter`.#[derive(Clone, Copy)]pub struct RateLimitDefaults { pub query_cost: u32, pub procedure_cost: u32, pub proxy_cost: u32,}
pub struct RateLimitConfig { pub capacity: u32, pub refill_rate: f64, pub default_query_cost: u32, pub default_procedure_cost: u32, pub default_proxy_cost: u32,}
pub enum CheckResult { Allowed { remaining: u32, limit: u32, reset: u64, }, Limited { retry_after: u64, limit: u32, reset: u64, }, Disabled,}
struct TokenBucket { tokens: f64, capacity: u32, refill_rate: f64, last_refill: Instant, last_access: Instant,}
/// Metadata about a registered API client, used for request validation.pub struct ClientIdentity { /// SHA-256 hash of the client secret pub secret_hash: String, /// The client's registered URI (for Origin header validation) pub client_uri: String,}
pub struct RateLimiter { defaults: RateLimitDefaults, buckets: DashMap<String, TokenBucket>, /// Per-client config, keyed by client_key (e.g. "hvc_..."). Presence in /// this map is the *only* thing that enables rate limiting for a key — /// unregistered keys are always allowed. client_configs: DashMap<String, RateLimitConfig>, /// Registered client identities, keyed by client_key client_identities: DashMap<String, ClientIdentity>,}
fn now_unix() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs()}
impl RateLimiter { pub fn new(defaults: RateLimitDefaults) -> Arc<Self> { Arc::new(Self { defaults, buckets: DashMap::new(), client_configs: DashMap::new(), client_identities: DashMap::new(), }) }
pub fn defaults(&self) -> RateLimitDefaults { self.defaults }
pub fn check(&self, key: &str, cost: u32) -> CheckResult { let (capacity, refill_rate) = match self.client_configs.get(key) { Some(cfg) => (cfg.capacity, cfg.refill_rate), None => return CheckResult::Disabled, }; let cost_f64 = cost as f64;
let now = Instant::now();
let mut bucket = self .buckets .entry(key.to_string()) .or_insert_with(|| TokenBucket { tokens: capacity as f64, capacity, refill_rate, last_refill: now, last_access: now, });
bucket.capacity = capacity; bucket.refill_rate = refill_rate;
let elapsed = now.duration_since(bucket.last_refill).as_secs_f64(); bucket.tokens = (bucket.tokens + elapsed * refill_rate).min(capacity as f64); bucket.last_refill = now; bucket.last_access = now;
let reset_secs = if bucket.tokens < capacity as f64 { ((capacity as f64 - bucket.tokens) / refill_rate).ceil() as u64 } else { 0 }; let reset = now_unix() + reset_secs;
if bucket.tokens >= cost_f64 { bucket.tokens -= cost_f64; CheckResult::Allowed { remaining: bucket.tokens.floor() as u32, limit: capacity, reset, } } else { let retry_after = ((cost_f64 - bucket.tokens) / refill_rate).ceil() as u64; CheckResult::Limited { retry_after, limit: capacity, reset: now_unix() + ((capacity as f64) / refill_rate).ceil() as u64, } } }
/// Get the default cost for a request type. Looks up the per-client /// override if one is registered, otherwise falls back to the seeded /// instance defaults. pub fn default_cost_for_type(&self, client_key: &str, request_type: &str) -> u32 { if let Some(cfg) = self.client_configs.get(client_key) { return match request_type { "query" => cfg.default_query_cost, "procedure" => cfg.default_procedure_cost, "proxy" => cfg.default_proxy_cost, _ => 1, }; } match request_type { "query" => self.defaults.query_cost, "procedure" => self.defaults.procedure_cost, "proxy" => self.defaults.proxy_cost, _ => 1, } }
pub fn register_client_config(&self, client_key: String, config: RateLimitConfig) { self.client_configs.insert(client_key, config); }
pub fn get_client_config(&self, client_key: &str) -> Option<RateLimitConfig> { self.client_configs .get(client_key) .map(|cfg| RateLimitConfig { capacity: cfg.capacity, refill_rate: cfg.refill_rate, default_query_cost: cfg.default_query_cost, default_procedure_cost: cfg.default_procedure_cost, default_proxy_cost: cfg.default_proxy_cost, }) }
pub fn remove_client_config(&self, client_key: &str) { self.client_configs.remove(client_key); }
pub fn register_client_identity(&self, client_key: String, identity: ClientIdentity) { self.client_identities.insert(client_key, identity); }
pub fn remove_client_identity(&self, client_key: &str) { self.client_identities.remove(client_key); }
pub fn validate_client_secret(&self, client_key: &str, secret: &str) -> bool { use sha2::{Digest, Sha256}; if let Some(identity) = self.client_identities.get(client_key) { let hash = hex::encode(Sha256::digest(secret.as_bytes())); hash == identity.secret_hash } else { false } }
pub fn validate_client_origin(&self, client_key: &str, origin: &str) -> bool { if let Some(identity) = self.client_identities.get(client_key) { let registered = identity.client_uri.trim_end_matches('/'); let provided = origin.trim_end_matches('/'); registered == provided } else { false } }
pub fn is_valid_client_key(&self, client_key: &str) -> bool { self.client_identities.contains_key(client_key) }
pub async fn spawn_cleanup(self: Arc<Self>) { let interval = tokio::time::Duration::from_secs(60); let stale_threshold = std::time::Duration::from_secs(300); loop { tokio::time::sleep(interval).await; let now = Instant::now(); self.buckets .retain(|_, bucket| now.duration_since(bucket.last_access) < stale_threshold); } }}
#[cfg(test)]mod tests { use super::*;
fn defaults() -> RateLimitDefaults { RateLimitDefaults { query_cost: 1, procedure_cost: 1, proxy_cost: 1, } }
fn cfg(capacity: u32, refill_rate: f64) -> RateLimitConfig { RateLimitConfig { capacity, refill_rate, default_query_cost: 1, default_procedure_cost: 1, default_proxy_cost: 1, } }
#[test] fn unregistered_key_is_not_rate_limited() { let rl = RateLimiter::new(defaults()); for _ in 0..1000 { assert!(matches!(rl.check("anything", 1), CheckResult::Disabled)); } }
#[test] fn registered_client_is_rate_limited() { let rl = RateLimiter::new(defaults()); rl.register_client_config("hvc_a".to_string(), cfg(3, 0.001));
for _ in 0..3 { assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); } assert!(matches!(rl.check("hvc_a", 1), CheckResult::Limited { .. })); }
#[test] fn cost_deducts_multiple_tokens() { let rl = RateLimiter::new(defaults()); rl.register_client_config("hvc_a".to_string(), cfg(10, 0.001));
assert!(matches!( rl.check("hvc_a", 5), CheckResult::Allowed { remaining: 5, .. } )); assert!(matches!( rl.check("hvc_a", 5), CheckResult::Allowed { remaining: 0, .. } )); assert!(matches!(rl.check("hvc_a", 5), CheckResult::Limited { .. })); }
#[test] fn different_clients_get_separate_buckets() { let rl = RateLimiter::new(defaults()); rl.register_client_config("hvc_a".to_string(), cfg(2, 0.001)); rl.register_client_config("hvc_b".to_string(), cfg(2, 0.001));
assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); assert!(matches!(rl.check("hvc_a", 1), CheckResult::Allowed { .. })); assert!(matches!(rl.check("hvc_a", 1), CheckResult::Limited { .. }));
assert!(matches!(rl.check("hvc_b", 1), CheckResult::Allowed { .. })); assert!(matches!(rl.check("hvc_b", 1), CheckResult::Allowed { .. })); assert!(matches!(rl.check("hvc_b", 1), CheckResult::Limited { .. })); }
#[test] fn remove_client_config_disables_limiting() { let rl = RateLimiter::new(defaults()); rl.register_client_config("hvc_temp".to_string(), cfg(1, 0.001)); assert!(matches!( rl.check("hvc_temp", 1), CheckResult::Allowed { .. } )); assert!(matches!( rl.check("hvc_temp", 1), CheckResult::Limited { .. } ));
rl.remove_client_config("hvc_temp"); assert!(matches!(rl.check("hvc_temp", 1), CheckResult::Disabled)); }
#[test] fn default_cost_for_type_uses_seeded_defaults_when_no_client_override() { let rl = RateLimiter::new(RateLimitDefaults { query_cost: 2, procedure_cost: 5, proxy_cost: 3, }); assert_eq!(rl.default_cost_for_type("nope", "query"), 2); assert_eq!(rl.default_cost_for_type("nope", "procedure"), 5); assert_eq!(rl.default_cost_for_type("nope", "proxy"), 3); }
#[test] fn default_cost_for_type_uses_per_client_override() { let rl = RateLimiter::new(RateLimitDefaults { query_cost: 2, procedure_cost: 5, proxy_cost: 3, }); rl.register_client_config( "hvc_a".to_string(), RateLimitConfig { capacity: 100, refill_rate: 1.0, default_query_cost: 7, default_procedure_cost: 8, default_proxy_cost: 9, }, ); assert_eq!(rl.default_cost_for_type("hvc_a", "query"), 7); assert_eq!(rl.default_cost_for_type("hvc_a", "procedure"), 8); assert_eq!(rl.default_cost_for_type("hvc_a", "proxy"), 9); }}