rate limiting

pxc1984 · pxc1984 · commit 29c584a43467 · 2026-02-03T17:51:11.000+03:00
diff --git a/backend/src/config.rs b/backend/src/config.rs
@@ -4,6 +4,7 @@ pub struct Config {
     pub db: DbConfig,
     pub jwt: JwtConfig,
     pub redis: RedisConfig,
+    pub rate_limit: RateLimitConfig,
     pub cors: CorsConfig,
     pub storage: StorageConfig,
 }
@@ -14,6 +15,7 @@ impl Config {
             db: DbConfig::init_from_env()?,
             jwt: JwtConfig::init_from_env()?,
             redis: RedisConfig::init_from_env()?,
+            rate_limit: RateLimitConfig::init_from_env()?,
             cors: CorsConfig::init_from_env()?,
             storage: StorageConfig::init_from_env()?,
         })
@@ -50,6 +52,15 @@ pub struct RedisConfig {
     pub url: String,
 }
 
+#[derive(envconfig::Envconfig)]
+pub struct RateLimitConfig {
+    #[envconfig(from = "RATE_LIMIT_CAPACITY", default = "60")]
+    pub capacity: u64,
+
+    #[envconfig(from = "RATE_LIMIT_REFILL_PER_SEC", default = "30")]
+    pub refill_per_sec: u64,
+}
+
 #[derive(envconfig::Envconfig)]
 pub struct CorsConfig {
     #[envconfig(from = "CORS_ALLOWED_ORIGIN")]
diff --git a/backend/src/http/mod.rs b/backend/src/http/mod.rs
@@ -3,6 +3,7 @@ use anyhow::Context;
 use axum::Router;
 use redis::aio::ConnectionManager;
 use sqlx::PgPool;
+use std::net::SocketAddr;
 use std::sync::Arc;
 
 // Utility modules.
@@ -12,6 +13,7 @@ mod error;
 mod extractor;
 
 mod types;
+mod rate_limit;
 
 // Api
 
@@ -34,6 +36,7 @@ struct ApiContext {
     config: Arc<Config>,
     db: PgPool,
     redis: ConnectionManager,
+    rate_limit_ttl_seconds: u64,
 }
 
 pub async fn serve(config: Config, db: PgPool, redis: ConnectionManager) -> anyhow::Result<()> {
@@ -52,12 +55,24 @@ pub async fn serve(config: Config, db: PgPool, redis: ConnectionManager) -> anyh
             .allow_headers(tower_http::cors::Any),
     };
 
+    let rate_limit_ttl_seconds = rate_limit::bucket_ttl_seconds(
+        config.rate_limit.capacity,
+        config.rate_limit.refill_per_sec,
+    );
+
+    let context = ApiContext {
+        config: Arc::new(config),
+        db,
+        redis,
+        rate_limit_ttl_seconds,
+    };
+
     let app = api_router()
-        .with_state(ApiContext {
-            config: Arc::new(config),
-            db,
-            redis: redis,
-        })
+        .with_state(context.clone())
+        .layer(axum::middleware::from_fn_with_state(
+            context.clone(),
+            rate_limit::rate_limit_middleware,
+        ))
         .layer(cors)
         .layer(axum::middleware::from_fn(metrics::metrics_middleware))
         .layer(TraceLayer::new_for_http());
@@ -78,7 +93,7 @@ pub async fn serve(config: Config, db: PgPool, redis: ConnectionManager) -> anyh
         .await
         .context("failed to bind HTTP server")?;
 
-    axum::serve(listener, app)
+    axum::serve(listener, app.into_make_service_with_connect_info::<SocketAddr>())
         .await
         .context("error running HTTP server")
 }
diff --git a/backend/src/http/rate_limit.rs b/backend/src/http/rate_limit.rs
@@ -0,0 +1,133 @@
+use crate::http::ApiContext;
+use axum::body::Body;
+use axum::extract::State;
+use axum::extract::connect_info::ConnectInfo;
+use axum::http::header::HeaderName;
+use axum::http::{Request, StatusCode};
+use axum::middleware::Next;
+use axum::response::{IntoResponse, Response};
+use redis::Script;
+use std::net::{IpAddr, SocketAddr};
+use std::sync::OnceLock;
+use time::OffsetDateTime;
+
+const LUA_SCRIPT: &str = r#"
+local key = KEYS[1]
+local capacity = tonumber(ARGV[1])
+local refill_per_sec = tonumber(ARGV[2])
+local now_ms = tonumber(ARGV[3])
+local ttl = tonumber(ARGV[4])
+
+local data = redis.call('HMGET', key, 'tokens', 'ts')
+local tokens = tonumber(data[1])
+local ts = tonumber(data[2])
+
+if tokens == nil then
+  tokens = capacity
+  ts = now_ms
+end
+
+if refill_per_sec > 0 then
+  local delta = math.max(0, now_ms - ts)
+  local refill = (delta / 1000) * refill_per_sec
+  if refill > 0 then
+    tokens = math.min(capacity, tokens + refill)
+  end
+end
+
+local allowed = 0
+if tokens >= 1 then
+  tokens = tokens - 1
+  allowed = 1
+end
+
+redis.call('HMSET', key, 'tokens', tokens, 'ts', now_ms)
+if ttl and ttl > 0 then
+  redis.call('EXPIRE', key, ttl)
+end
+
+return allowed
+"#;
+
+static SCRIPT: OnceLock<Script> = OnceLock::new();
+static FORWARDED_FOR_HEADER: HeaderName = HeaderName::from_static("x-forwarded-for");
+static REAL_IP_HEADER: HeaderName = HeaderName::from_static("x-real-ip");
+
+pub async fn rate_limit_middleware(
+    State(ctx): State<ApiContext>,
+    req: Request<Body>,
+    next: Next,
+) -> Response {
+    let ip = match extract_ip(&req) {
+        Some(ip) => ip,
+        None => return next.run(req).await,
+    };
+
+    let capacity = ctx.config.rate_limit.capacity;
+    let refill_per_sec = ctx.config.rate_limit.refill_per_sec;
+    if capacity == 0 || refill_per_sec == 0 {
+        return next.run(req).await;
+    }
+
+    let ttl_seconds = ctx.rate_limit_ttl_seconds;
+    let now_ms = (OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as i64;
+    let key = format!("rate_limit:{ip}");
+
+    let mut redis = ctx.redis.clone();
+    let script = SCRIPT.get_or_init(|| Script::new(LUA_SCRIPT));
+    let allowed: i32 = match script
+        .key(key)
+        .arg(capacity)
+        .arg(refill_per_sec)
+        .arg(now_ms)
+        .arg(ttl_seconds)
+        .invoke_async(&mut redis)
+        .await
+    {
+        Ok(allowed) => allowed,
+        Err(err) => {
+            log::error!("rate limit redis error: {err}");
+            return next.run(req).await;
+        }
+    };
+
+    if allowed == 1 {
+        next.run(req).await
+    } else {
+        (StatusCode::TOO_MANY_REQUESTS, "rate limit exceeded").into_response()
+    }
+}
+
+fn extract_ip(req: &Request<Body>) -> Option<IpAddr> {
+    if let Some(value) = req.headers().get(&FORWARDED_FOR_HEADER) {
+        if let Ok(value) = value.to_str() {
+            if let Some(first) = value.split(',').next() {
+                if let Ok(ip) = first.trim().parse::<IpAddr>() {
+                    return Some(ip);
+                }
+            }
+        }
+    }
+
+    if let Some(value) = req.headers().get(&REAL_IP_HEADER) {
+        if let Ok(value) = value.to_str() {
+            if let Ok(ip) = value.trim().parse::<IpAddr>() {
+                return Some(ip);
+            }
+        }
+    }
+
+    req.extensions()
+        .get::<ConnectInfo<SocketAddr>>()
+        .map(|info| info.0.ip())
+}
+
+pub fn bucket_ttl_seconds(capacity: u64, refill_per_sec: u64) -> u64 {
+    if refill_per_sec == 0 {
+        return 60;
+    }
+
+    let refill_time = capacity.saturating_add(refill_per_sec - 1) / refill_per_sec;
+    let ttl = refill_time.saturating_mul(2);
+    ttl.max(60)
+}
diff --git a/stress-test/src/api_usecase.rs b/stress-test/src/api_usecase.rs
@@ -1,11 +1,12 @@
 use rand::prelude::*;
+use reqwest::StatusCode;
 use reqwest::header::AUTHORIZATION;
 use serde::de::DeserializeOwned;
 
 const API_URL: &str = "http://localhost:8080";
 
 // this function should imitate use case of api: full user cycle. Registration,
-pub async fn api_usecase() -> Result<(), reqwest::Error> {
+pub async fn api_usecase() -> Result<(), ApiError> {
     let data = MockData::generate();
     let client = reqwest::Client::new();
 
@@ -239,12 +240,43 @@ struct Comment {
     id: i64,
 }
 
-async fn send_json<T: DeserializeOwned>(
-    response: reqwest::Response,
-) -> Result<T, reqwest::Error> {
-    response.error_for_status()?.json().await
+async fn send_json<T: DeserializeOwned>(response: reqwest::Response) -> Result<T, ApiError> {
+    let status = response.status();
+    if !status.is_success() {
+        return Err(ApiError::Http(status));
+    }
+
+    response.json().await.map_err(ApiError::Transport)
+}
+
+async fn send_empty(response: reqwest::Response) -> Result<(), ApiError> {
+    let status = response.status();
+    if !status.is_success() {
+        return Err(ApiError::Http(status));
+    }
+
+    Ok(())
 }
 
-async fn send_empty(response: reqwest::Response) -> Result<(), reqwest::Error> {
-    response.error_for_status().map(|_| ())
+#[derive(Debug)]
+pub enum ApiError {
+    Http(StatusCode),
+
+    #[allow(dead_code)]
+    Transport(reqwest::Error),
+}
+
+impl ApiError {
+    pub fn status(&self) -> Option<StatusCode> {
+        match self {
+            Self::Http(status) => Some(*status),
+            Self::Transport(_) => None,
+        }
+    }
+}
+
+impl From<reqwest::Error> for ApiError {
+    fn from(err: reqwest::Error) -> Self {
+        Self::Transport(err)
+    }
 }
diff --git a/stress-test/src/main.rs b/stress-test/src/main.rs
@@ -1,21 +1,22 @@
-use api_usecase::api_usecase;
+use api_usecase::{api_usecase, ApiError};
 use std::sync::atomic;
 use std::time::Instant;
 use tokio::task::JoinSet;
 
 mod api_usecase;
 
-const THREAD_COUNT: usize = 1;
-const WOKER_COUNT: usize = 8; // per-thread
+const THREAD_COUNT: usize = 8;
+const WOKER_COUNT: usize = 1; // per-thread
 
 static TOTAL_API_REQUESTS: atomic::AtomicU64 = atomic::AtomicU64::new(0);
 static API_FAILURES: atomic::AtomicU64 = atomic::AtomicU64::new(0);
 static CYCLES_SUCCESS: atomic::AtomicU64 = atomic::AtomicU64::new(0);
+static RATE_LIMITED_REQUESTS: atomic::AtomicU64 = atomic::AtomicU64::new(0);
 
 async fn worker() {
-    let mut set: JoinSet<Result<(), reqwest::Error>> = JoinSet::new();
+    let mut set: JoinSet<Result<(), ApiError>> = JoinSet::new();
 
-    let spawn_one = |set: &mut JoinSet<Result<(), reqwest::Error>>| {
+    let spawn_one = |set: &mut JoinSet<Result<(), ApiError>>| {
         TOTAL_API_REQUESTS.fetch_add(1, atomic::Ordering::Relaxed);
         set.spawn(async { api_usecase().await });
     };
@@ -26,10 +27,17 @@ async fn worker() {
 
     loop {
         match set.join_next().await {
-            Some(Ok(_)) => {
+            Some(Ok(Ok(()))) => {
                 CYCLES_SUCCESS.fetch_add(1, atomic::Ordering::Relaxed);
                 spawn_one(&mut set);
             }
+            Some(Ok(Err(err))) => {
+                if err.status() == Some(reqwest::StatusCode::TOO_MANY_REQUESTS) {
+                    RATE_LIMITED_REQUESTS.fetch_add(1, atomic::Ordering::Relaxed);
+                }
+                API_FAILURES.fetch_add(1, atomic::Ordering::Relaxed);
+                spawn_one(&mut set);
+            }
             Some(Err(_)) => {
                 API_FAILURES.fetch_add(1, atomic::Ordering::Relaxed);
                 spawn_one(&mut set);
@@ -44,6 +52,7 @@ fn print_stats(started_at: Instant) {
     let total = TOTAL_API_REQUESTS.load(atomic::Ordering::Relaxed);
     let failures = API_FAILURES.load(atomic::Ordering::Relaxed);
     let cycles = CYCLES_SUCCESS.load(atomic::Ordering::Relaxed);
+    let rate_limited = RATE_LIMITED_REQUESTS.load(atomic::Ordering::Relaxed);
     let success = total.saturating_sub(failures);
     let rps = total as f64 / elapsed;
     let failure_rate = if total == 0 {
@@ -57,6 +66,7 @@ fn print_stats(started_at: Instant) {
     println!("requests_total={total}");
     println!("requests_ok={success}");
     println!("requests_failed={failures}");
+    println!("requests_rate_limited={rate_limited}");
     println!("request_rate_rps={rps:.2}");
     println!("failure_rate_pct={failure_rate:.2}");
     println!("cycles_success={cycles}");
