143 lines
4.1 KiB
Rust
143 lines
4.1 KiB
Rust
//! HTTP rate limiting for git operations.
|
|
//!
|
|
//! Uses a token-bucket approach with per-IP and per-repo-write limits.
|
|
//! Cleanup runs every 5 minutes to prevent unbounded memory growth.
|
|
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::sync::RwLock;
|
|
use tokio::time::interval;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct RateLimitConfig {
|
|
/// Requests allowed per window for read operations.
|
|
pub read_requests_per_window: u32,
|
|
/// Requests allowed per window for write operations.
|
|
pub write_requests_per_window: u32,
|
|
/// Window duration in seconds.
|
|
pub window_secs: u64,
|
|
}
|
|
|
|
impl Default for RateLimitConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
read_requests_per_window: 120,
|
|
write_requests_per_window: 30,
|
|
window_secs: 60,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct RateLimitBucket {
|
|
read_count: u32,
|
|
write_count: u32,
|
|
reset_time: Instant,
|
|
}
|
|
|
|
pub struct RateLimiter {
|
|
buckets: Arc<RwLock<HashMap<String, RateLimitBucket>>>,
|
|
config: RateLimitConfig,
|
|
}
|
|
|
|
impl RateLimiter {
|
|
pub fn new(config: RateLimitConfig) -> Self {
|
|
Self {
|
|
buckets: Arc::new(RwLock::new(HashMap::new())),
|
|
config,
|
|
}
|
|
}
|
|
|
|
pub async fn is_ip_read_allowed(&self, ip: &str) -> bool {
|
|
let key = format!("ip:read:{}", ip);
|
|
self.is_allowed(&key, self.config.read_requests_per_window)
|
|
.await
|
|
}
|
|
|
|
pub async fn is_ip_write_allowed(&self, ip: &str) -> bool {
|
|
let key = format!("ip:write:{}", ip);
|
|
self.is_allowed(&key, self.config.write_requests_per_window)
|
|
.await
|
|
}
|
|
|
|
pub async fn is_repo_write_allowed(&self, ip: &str, repo_path: &str) -> bool {
|
|
let key = format!("repo:write:{}:{}", ip, repo_path);
|
|
self.is_allowed(&key, self.config.write_requests_per_window)
|
|
.await
|
|
}
|
|
|
|
async fn is_allowed(&self, key: &str, limit: u32) -> bool {
|
|
let now = Instant::now();
|
|
let mut buckets = self.buckets.write().await;
|
|
|
|
let bucket = buckets
|
|
.entry(key.to_string())
|
|
.or_insert_with(|| RateLimitBucket {
|
|
read_count: 0,
|
|
write_count: 0,
|
|
reset_time: now + Duration::from_secs(self.config.window_secs),
|
|
});
|
|
|
|
if now >= bucket.reset_time {
|
|
bucket.read_count = 0;
|
|
bucket.write_count = 0;
|
|
bucket.reset_time = now + Duration::from_secs(self.config.window_secs);
|
|
}
|
|
|
|
// Use read_count for both read/write since we don't distinguish in bucket
|
|
if bucket.read_count >= limit {
|
|
return false;
|
|
}
|
|
|
|
bucket.read_count += 1;
|
|
true
|
|
}
|
|
|
|
pub async fn retry_after(&self, ip: &str) -> u64 {
|
|
let key_read = format!("ip:read:{}", ip);
|
|
let now = Instant::now();
|
|
let buckets = self.buckets.read().await;
|
|
|
|
if let Some(bucket) = buckets.get(&key_read) {
|
|
if now < bucket.reset_time {
|
|
return bucket.reset_time.saturating_duration_since(now).as_secs() as u64;
|
|
}
|
|
}
|
|
0
|
|
}
|
|
|
|
/// Start a background cleanup task that removes expired entries.
|
|
/// Should be spawned once at startup.
|
|
pub fn start_cleanup(self: Arc<Self>) -> tokio::task::JoinHandle<()> {
|
|
tokio::spawn(async move {
|
|
let mut ticker = interval(Duration::from_secs(300)); // every 5 minutes
|
|
loop {
|
|
ticker.tick().await;
|
|
let now = Instant::now();
|
|
let mut buckets = self.buckets.write().await;
|
|
buckets.retain(|_, bucket| now < bucket.reset_time);
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[tokio::test]
|
|
async fn test_rate_limit_allows_requests_up_to_limit() {
|
|
let limiter = Arc::new(RateLimiter::new(RateLimitConfig {
|
|
read_requests_per_window: 3,
|
|
write_requests_per_window: 1,
|
|
window_secs: 60,
|
|
}));
|
|
|
|
for _ in 0..3 {
|
|
assert!(limiter.is_ip_read_allowed("1.2.3.4").await);
|
|
}
|
|
assert!(!limiter.is_ip_read_allowed("1.2.3.4").await);
|
|
}
|
|
}
|