feat: add phase 6 auth and rate limiting

.env.example

@@ -30,6 +30,14 @@ CHUNK_OVERLAP=64
 # Retrieval
 TOP_K=5
 
+# Auth & Rate Limiting
+# If true, requests to ingest/query endpoints must include a valid API key.
+REQUIRE_AUTH=false
+# Comma-separated keys accepted via X-API-Key or Authorization: Bearer <key>.
+API_KEYS=dev-key-1,dev-key-2
+# Per-identity limit across protected endpoints, in requests/minute.
+RATE_LIMIT_PER_MINUTE=120
+
 # Cross-Encoder Reranker (optional, improves result relevance)
 RERANKER_API_URL=https://api-inference.huggingface.co/models/cross-encoder/ms-marco-MiniLM-L-6-v2
 # Falls back to HF_API_KEY if not set

README.md

@@ -192,6 +192,9 @@ All configuration is via environment variables (see `.env.example`):
 | `CHUNK_SIZE` | `512` | Max chars per chunk |
 | `CHUNK_OVERLAP` | `64` | Overlap between chunks |
 | `TOP_K` | `5` | Default top-k retrieval |
+| `REQUIRE_AUTH` | `false` | Require API key on ingest/query endpoints (`true`/`false`) |
+| `API_KEYS` |  | Comma-separated valid API keys (used by `X-API-Key` or `Authorization: Bearer`) |
+| `RATE_LIMIT_PER_MINUTE` | `120` | Per-identity request limit (ingest/query endpoints) |
 | `RERANKER_API_URL` | HuggingFace cross-encoder | Cross-encoder reranker endpoint (leave empty to disable) |
 | `RERANKER_API_KEY` |  | API key for reranker (falls back to `HF_API_KEY`) |
 
@@ -235,6 +238,11 @@ cargo test --all-features && cargo clippy -- -D warnings && cargo fmt --check
 | Header | Required | Default | Description |
 |--------|----------|---------|-------------|
 | `X-Tenant-ID` | No | `default` | Isolates data per tenant into separate Qdrant collections |
+| `X-API-Key` | Depends (`REQUIRE_AUTH`) |  | API key auth header for protected endpoints |
+| `Authorization: Bearer <key>` | Depends (`REQUIRE_AUTH`) |  | Alternate auth header for protected endpoints |
+
+When `REQUIRE_AUTH=true`, missing/invalid keys return `401 Unauthorized`.
+If request volume exceeds `RATE_LIMIT_PER_MINUTE`, protected endpoints return `429 Too Many Requests`.
 
 ### `POST /ingest`
 
@@ -434,7 +442,7 @@ blazerag/
 - [x] Phase 3: Reranking (cross-encoder)
 - [x] Phase 4: Batch ingestion (PDF, HTML, Markdown)
 - [x] Phase 5: Multi-tenant collections
-- [ ] Phase 6: Auth & rate limiting
+- [x] Phase 6: Auth & rate limiting
 - [ ] Phase 7: Web UI dashboard
 - [ ] Phase 8: Managed cloud offering
 

src/lib.rs

@@ -6,6 +6,7 @@ pub mod ingestor;
 pub mod llm;
 pub mod reranker;
 pub mod retriever;
+pub mod security;
 pub mod server;
 
 #[derive(Clone)]
@@ -15,4 +16,5 @@ pub struct AppState {
     pub llm_client: Arc<llm::LlmClient>,
     pub reranker: Arc<reranker::Reranker>,
     pub chunker_config: chunker::ChunkerConfig,
+    pub security: Arc<security::SecurityState>,
 }

src/main.rs

@@ -1,6 +1,6 @@
 use anyhow::Context;
 use axum::Router;
-use blazerag::{chunker, embedder, llm, reranker, retriever, server, AppState};
+use blazerag::{chunker, embedder, llm, reranker, retriever, security, server, AppState};
 use std::sync::Arc;
 use tracing_subscriber::EnvFilter;
 
@@ -60,6 +60,13 @@ async fn main() -> anyhow::Result<()> {
         .await
         .context("Failed to initialize reranker")?;
 
+    let security = security::SecurityState::from_env();
+    tracing::info!(
+        "Security config loaded: require_auth={}, rate_limit_per_minute={}",
+        security.require_auth,
+        security.rate_limit_per_minute
+    );
+
     let state = AppState {
         embedder: Arc::new(embedder),
         retriever: Arc::new(retriever),
@@ -69,6 +76,7 @@ async fn main() -> anyhow::Result<()> {
             chunk_size,
             chunk_overlap,
         },
+        security: Arc::new(security),
     };
 
     let app = Router::new().merge(server::routes()).with_state(state);

src/security/mod.rs

@@ -0,0 +1,157 @@
+use axum::http::HeaderMap;
+use std::collections::{HashMap, HashSet, VecDeque};
+use std::time::{Duration, Instant};
+use tokio::sync::Mutex;
+
+#[derive(Debug)]
+pub struct SecurityState {
+    pub require_auth: bool,
+    api_keys: HashSet<String>,
+    pub rate_limit_per_minute: usize,
+    limiter: Mutex<HashMap<String, VecDeque<Instant>>>,
+}
+
+#[derive(Debug)]
+pub enum SecurityError {
+    Unauthorized(String),
+    RateLimited(String),
+}
+
+impl SecurityState {
+    pub fn from_env() -> Self {
+        let require_auth = parse_bool_env("REQUIRE_AUTH", false);
+        let api_keys = std::env::var("API_KEYS")
+            .unwrap_or_default()
+            .split(',')
+            .map(|key| key.trim().to_string())
+            .filter(|key| !key.is_empty())
+            .collect::<HashSet<_>>();
+        let rate_limit_per_minute = std::env::var("RATE_LIMIT_PER_MINUTE")
+            .unwrap_or_else(|_| "120".to_string())
+            .parse::<usize>()
+            .ok()
+            .filter(|v| *v > 0)
+            .unwrap_or(120);
+
+        Self {
+            require_auth,
+            api_keys,
+            rate_limit_per_minute,
+            limiter: Mutex::new(HashMap::new()),
+        }
+    }
+
+    pub async fn authorize_and_check_rate_limit(
+        &self,
+        headers: &HeaderMap,
+        tenant_id: &str,
+    ) -> Result<(), SecurityError> {
+        let provided_key = extract_api_key(headers);
+        let identity = if let Some(key) = provided_key.as_ref() {
+            if !self.api_keys.is_empty() && !self.api_keys.contains(key) {
+                return Err(SecurityError::Unauthorized("Invalid API key".to_string()));
+            }
+            format!("api_key:{key}")
+        } else {
+            if self.require_auth {
+                return Err(SecurityError::Unauthorized("Missing API key".to_string()));
+            }
+            format!("tenant:{tenant_id}")
+        };
+
+        self.check_rate_limit(&identity).await
+    }
+
+    async fn check_rate_limit(&self, identity: &str) -> Result<(), SecurityError> {
+        let now = Instant::now();
+        let cutoff = now - Duration::from_secs(60);
+        let mut guard = self.limiter.lock().await;
+        let entries = guard.entry(identity.to_string()).or_default();
+
+        while let Some(ts) = entries.front() {
+            if *ts < cutoff {
+                entries.pop_front();
+            } else {
+                break;
+            }
+        }
+
+        if entries.len() >= self.rate_limit_per_minute {
+            return Err(SecurityError::RateLimited(format!(
+                "Rate limit exceeded: max {} requests per minute",
+                self.rate_limit_per_minute
+            )));
+        }
+
+        entries.push_back(now);
+        Ok(())
+    }
+}
+
+fn parse_bool_env(key: &str, default: bool) -> bool {
+    std::env::var(key)
+        .ok()
+        .and_then(|v| match v.to_ascii_lowercase().as_str() {
+            "1" | "true" | "yes" | "on" => Some(true),
+            "0" | "false" | "no" | "off" => Some(false),
+            _ => None,
+        })
+        .unwrap_or(default)
+}
+
+pub fn extract_api_key(headers: &HeaderMap) -> Option<String> {
+    if let Some(raw) = headers.get("x-api-key").and_then(|v| v.to_str().ok()) {
+        let key = raw.trim();
+        if !key.is_empty() {
+            return Some(key.to_string());
+        }
+    }
+
+    headers
+        .get("authorization")
+        .and_then(|v| v.to_str().ok())
+        .and_then(|raw| {
+            raw.strip_prefix("Bearer ")
+                .or_else(|| raw.strip_prefix("bearer "))
+        })
+        .map(str::trim)
+        .filter(|v| !v.is_empty())
+        .map(str::to_string)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use axum::http::HeaderValue;
+
+    #[tokio::test]
+    async fn rate_limit_blocks_when_capacity_reached() {
+        let state = SecurityState {
+            require_auth: false,
+            api_keys: HashSet::new(),
+            rate_limit_per_minute: 2,
+            limiter: Mutex::new(HashMap::new()),
+        };
+
+        assert!(state.check_rate_limit("anon").await.is_ok());
+        assert!(state.check_rate_limit("anon").await.is_ok());
+        assert!(matches!(
+            state.check_rate_limit("anon").await,
+            Err(SecurityError::RateLimited(_))
+        ));
+    }
+
+    #[test]
+    fn extracts_key_from_headers() {
+        let mut headers = HeaderMap::new();
+        headers.insert("x-api-key", HeaderValue::from_static("test-key"));
+        assert_eq!(extract_api_key(&headers).as_deref(), Some("test-key"));
+
+        headers.clear();
+        headers.insert(
+            "authorization",
+            HeaderValue::from_static("Bearer another-key"),
+        );
+        assert_eq!(extract_api_key(&headers).as_deref(), Some("another-key"));
+    }
+}

src/server/mod.rs

@@ -19,6 +19,7 @@ use crate::chunker;
 use crate::ingestor::{self, Encoding, FileFormat};
 use crate::reranker::Reranker;
 use crate::retriever::Document;
+use crate::security::SecurityError;
 use crate::AppState;
 
 #[derive(Serialize, Deserialize)]
@@ -55,6 +56,21 @@ fn extract_tenant_id(headers: &HeaderMap) -> String {
         .unwrap_or_else(|| "default".to_string())
 }
 
+fn security_error_response(err: SecurityError) -> axum::response::Response {
+    match err {
+        SecurityError::Unauthorized(msg) => (
+            axum::http::StatusCode::UNAUTHORIZED,
+            Json(serde_json::to_value(ErrorResponse { error: msg }).unwrap()),
+        )
+            .into_response(),
+        SecurityError::RateLimited(msg) => (
+            axum::http::StatusCode::TOO_MANY_REQUESTS,
+            Json(serde_json::to_value(ErrorResponse { error: msg }).unwrap()),
+        )
+            .into_response(),
+    }
+}
+
 #[derive(Clone, Serialize, Deserialize)]
 pub struct QuerySource {
     pub text: String,
@@ -97,6 +113,13 @@ async fn ingest_handler(
     Json(req): Json<IngestRequest>,
 ) -> impl IntoResponse {
     let tenant_id = extract_tenant_id(&headers);
+    if let Err(err) = state
+        .security
+        .authorize_and_check_rate_limit(&headers, &tenant_id)
+        .await
+    {
+        return security_error_response(err);
+    }
     tracing::info!(
         "Ingesting text ({} chars) for tenant={}",
         req.text.len(),
@@ -115,7 +138,8 @@ async fn ingest_handler(
                 })
                 .unwrap(),
             ),
-        );
+        )
+            .into_response();
     }
 
     let metadata = req.metadata.unwrap_or(serde_json::json!({}));
@@ -132,7 +156,8 @@ async fn ingest_handler(
                     })
                     .unwrap(),
                 ),
-            );
+            )
+                .into_response();
         }
     };
 
@@ -168,6 +193,7 @@ async fn ingest_handler(
                     .unwrap(),
                 ),
             )
+                .into_response()
         }
         Err(e) => {
             tracing::error!("Qdrant upsert failed: {}", e);
@@ -180,6 +206,7 @@ async fn ingest_handler(
                     .unwrap(),
                 ),
             )
+                .into_response()
         }
     }
 }
@@ -221,6 +248,13 @@ async fn batch_ingest_handler(
     Json(req): Json<Vec<BatchFileEntry>>,
 ) -> impl IntoResponse {
     let tenant_id = extract_tenant_id(&headers);
+    if let Err(err) = state
+        .security
+        .authorize_and_check_rate_limit(&headers, &tenant_id)
+        .await
+    {
+        return security_error_response(err);
+    }
     let total = req.len();
     tracing::info!("Batch ingest: {} files for tenant={}", total, tenant_id);
 
@@ -362,6 +396,7 @@ async fn batch_ingest_handler(
             .unwrap(),
         ),
     )
+        .into_response()
 }
 
 async fn query_handler(
@@ -370,6 +405,13 @@ async fn query_handler(
     Json(req): Json<QueryRequest>,
 ) -> impl IntoResponse {
     let tenant_id = extract_tenant_id(&headers);
+    if let Err(err) = state
+        .security
+        .authorize_and_check_rate_limit(&headers, &tenant_id)
+        .await
+    {
+        return security_error_response(err);
+    }
     tracing::info!("Query: {} (tenant={})", req.question, tenant_id);
 
     let query_embeddings = match state.embedder.embed(std::slice::from_ref(&req.question)) {
@@ -502,6 +544,17 @@ async fn query_stream_handler(
     Json(req): Json<QueryRequest>,
 ) -> Result<Sse<SseStream>, (StatusCode, Json<ErrorResponse>)> {
     let tenant_id = extract_tenant_id(&headers);
+    if let Err(err) = state
+        .security
+        .authorize_and_check_rate_limit(&headers, &tenant_id)
+        .await
+    {
+        let (status, message) = match err {
+            SecurityError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg),
+            SecurityError::RateLimited(msg) => (StatusCode::TOO_MANY_REQUESTS, msg),
+        };
+        return Err((status, Json(ErrorResponse { error: message })));
+    }
     tracing::info!("Streaming query: {} (tenant={})", req.question, tenant_id);
 
     let query_embeddings = state
@@ -637,6 +690,7 @@ mod tests {
     use crate::llm::LlmClient;
     use crate::reranker::Reranker;
     use crate::retriever::Retriever;
+    use crate::security::SecurityState;
     use axum::body::Body;
     use axum::http::{Request, StatusCode};
     use std::sync::Arc;
@@ -667,6 +721,7 @@ mod tests {
             llm_client: Arc::new(llm_client),
             reranker: Arc::new(reranker),
             chunker_config,
+            security: Arc::new(SecurityState::from_env()),
         }
     }