From 7d8e95eb0ad04c974676698acee9f24a956498f2 Mon Sep 17 00:00:00 2001
From: Droid <droid@factory.ai>
Date: Wed, 21 Jan 2026 15:21:08 +0400
Subject: [PATCH] fix: add rate limiting to API server

This commit adds rate limiting to the API server using tower-governor. It limits requests to 60 per second with a burst of 10 to prevent DoS attacks and resource exhaustion.
---
 Cargo.toml        |  2 ++
 src/server/api.rs | 12 ++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index d912c66..59d9dc9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -44,6 +44,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 dialoguer = "0.12"
 console = "0.15"
 ctrlc = "3"
+tower_governor = "0.8.0"
+tower = "0.5.3"
 
 [dev-dependencies]
 tempfile = "3"
diff --git a/src/server/api.rs b/src/server/api.rs
index bd5cc8a..2630a72 100644
--- a/src/server/api.rs
+++ b/src/server/api.rs
@@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize};
 use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::sync::{Arc, Mutex};
+use tower_governor::{governor::GovernorConfigBuilder, GovernorLayer};
 use tower_http::cors::{Any, CorsLayer};
 
 use crate::config::Config;
@@ -108,6 +109,16 @@ pub async fn run_server(config: &Config, host: &str, port: u16) -> Result<()> {
         .allow_methods(Any)
         .allow_headers(Any);
 
+    // Rate limiting: 60 requests per second with a burst of 10
+    // This protects against simple DoS attacks and embedding abuse
+    let governor_conf = Arc::new(
+        GovernorConfigBuilder::default()
+            .per_second(60)
+            .burst_size(10)
+            .finish()
+            .unwrap(),
+    );
+
     let app = Router::new()
         .route("/", get(root))
         .route("/health", get(health))
@@ -116,6 +127,7 @@ pub async fn run_server(config: &Config, host: &str, port: u16) -> Result<()> {
         .route("/embed", post(embed))
         .route("/embed_batch", post(embed_batch))
         .layer(cors)
+        .layer(GovernorLayer::new(governor_conf))
         .with_state(state);
 
     let addr: SocketAddr = format!("{}:{}", host, port).parse()?;