From de4acb35ca395f4ca1584c1b6d79a23e21cb7c24 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 04:37:53 +0000 Subject: [PATCH 1/5] Initial plan From 23f1ad64fbca95e18046a97d6527d59b31c013c3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 04:57:22 +0000 Subject: [PATCH 2/5] Add rate limiting and max entries limit for whitelist endpoints Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/Cargo.toml | 2 ++ server/src/config.rs | 2 ++ server/src/handlers/project.rs | 46 ++++++++++++++++++++++++++++++++++ server/src/state.rs | 17 +++++++++++++ 4 files changed, 67 insertions(+) diff --git a/server/Cargo.toml b/server/Cargo.toml index 197c7c0..26ee5d5 100644 --- a/server/Cargo.toml +++ b/server/Cargo.toml @@ -24,3 +24,5 @@ time = "0.3.46" openssl = { version = "0.10", features = ["vendored"] } openssl-sys = { version = "0.9", features = ["vendored"] } url = "2.5.8" +dashmap = "6.1" +governor = "0.10" diff --git a/server/src/config.rs b/server/src/config.rs index d2bcdd3..f23f552 100644 --- a/server/src/config.rs +++ b/server/src/config.rs @@ -3,6 +3,7 @@ use sqlx::postgres::PgPoolOptions; use std::sync::{Arc, Mutex}; use std::collections::HashMap; use crate::state::AppState; +use dashmap::DashMap; pub async fn setup_database_and_docker() -> Result> { // 1. Docker setup @@ -22,6 +23,7 @@ pub async fn setup_database_and_docker() -> Result = sqlx::query_as( "SELECT id FROM projects WHERE owner_id = $1 AND LOWER(slug) = LOWER($2)", @@ -506,6 +518,15 @@ pub async fn add_to_whitelist( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Session Error: {}", e)))?; let user = user.ok_or((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))?; + // Apply rate limiting: 20 requests per minute per user + let rate_limiter = state.get_whitelist_rate_limiter(user.id); + if rate_limiter.check().is_err() { + return Err(( + StatusCode::TOO_MANY_REQUESTS, + "Rate limit exceeded. Please try again later.".to_string() + )); + } + let trimmed_url = payload.allowed_url.trim(); if trimmed_url.is_empty() { return Err((StatusCode::BAD_REQUEST, "allowed_url is required".to_string())); @@ -525,6 +546,22 @@ pub async fn add_to_whitelist( None => return Err((StatusCode::NOT_FOUND, "Project not found".to_string())), }; + // Check current whitelist entry count to prevent storage exhaustion + let count_row: (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM project_whitelists WHERE project_id = $1", + ) + .bind(project_id) + .fetch_one(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; + + if count_row.0 >= MAX_WHITELIST_ENTRIES { + return Err(( + StatusCode::FORBIDDEN, + format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) + )); + } + // Unique(project_id, allowed_url) is enforced by the DB; ignore conflicts let result = sqlx::query( "INSERT INTO project_whitelists (project_id, allowed_url) \ @@ -556,6 +593,15 @@ pub async fn remove_from_whitelist( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Session Error: {}", e)))?; let user = user.ok_or((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))?; + // Apply rate limiting: 20 requests per minute per user + let rate_limiter = state.get_whitelist_rate_limiter(user.id); + if rate_limiter.check().is_err() { + return Err(( + StatusCode::TOO_MANY_REQUESTS, + "Rate limit exceeded. Please try again later.".to_string() + )); + } + let trimmed_url = payload.allowed_url.trim(); if trimmed_url.is_empty() { return Err((StatusCode::BAD_REQUEST, "allowed_url is required".to_string())); diff --git a/server/src/state.rs b/server/src/state.rs index 7dcac10..474b10a 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -2,6 +2,9 @@ use bollard::Docker; use std::sync::{Arc, Mutex, MutexGuard}; use std::collections::HashMap; use std::time::Instant; +use dashmap::DashMap; +use governor::{Quota, RateLimiter, clock::DefaultClock, state::{InMemoryState, NotKeyed}}; +use std::num::NonZeroU32; // New Struct to track container details AND ownership #[derive(Clone, Debug)] @@ -25,6 +28,8 @@ pub struct AppState { pub github_id: String, pub github_secret: String, pub sessions: SessionMap, + // Rate limiter for whitelist operations: per-user tracking + pub whitelist_rate_limiters: Arc>>>, } impl AppState { @@ -37,4 +42,16 @@ impl AppState { } } } + + /// Get or create a rate limiter for a user (20 requests per minute for whitelist operations) + pub fn get_whitelist_rate_limiter(&self, user_id: i64) -> Arc> { + self.whitelist_rate_limiters + .entry(user_id) + .or_insert_with(|| { + // 20 requests per 60 seconds + let quota = Quota::per_minute(NonZeroU32::new(20).unwrap()); + Arc::new(RateLimiter::direct(quota)) + }) + .clone() + } } \ No newline at end of file From be66fb43b49c1d4c027be6cd30ba58e282d1e2e1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 04:59:07 +0000 Subject: [PATCH 3/5] Address code review feedback: fix race condition and improve code quality Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/handlers/project.rs | 59 ++++++++++++++++++++-------------- server/src/state.rs | 8 ++--- 2 files changed, 39 insertions(+), 28 deletions(-) diff --git a/server/src/handlers/project.rs b/server/src/handlers/project.rs index dd6df7b..26600b5 100644 --- a/server/src/handlers/project.rs +++ b/server/src/handlers/project.rs @@ -21,6 +21,9 @@ use std::collections::HashMap; // Maximum number of whitelist entries allowed per project const MAX_WHITELIST_ENTRIES: i64 = 100; +// Rate limit for whitelist operations (requests per minute per user) +pub const WHITELIST_RATE_LIMIT_PER_MINUTE: u32 = 20; + #[derive(Deserialize)] pub struct SearchQuery { q: String, @@ -469,7 +472,7 @@ pub async fn get_whitelist( let user = user.ok_or((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))?; // Apply rate limiting: 20 requests per minute per user - let rate_limiter = state.get_whitelist_rate_limiter(user.id); + let rate_limiter = state.get_or_create_whitelist_rate_limiter(user.id); if rate_limiter.check().is_err() { return Err(( StatusCode::TOO_MANY_REQUESTS, @@ -519,7 +522,7 @@ pub async fn add_to_whitelist( let user = user.ok_or((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))?; // Apply rate limiting: 20 requests per minute per user - let rate_limiter = state.get_whitelist_rate_limiter(user.id); + let rate_limiter = state.get_or_create_whitelist_rate_limiter(user.id); if rate_limiter.check().is_err() { return Err(( StatusCode::TOO_MANY_REQUESTS, @@ -546,35 +549,43 @@ pub async fn add_to_whitelist( None => return Err((StatusCode::NOT_FOUND, "Project not found".to_string())), }; - // Check current whitelist entry count to prevent storage exhaustion - let count_row: (i64,) = sqlx::query_as( - "SELECT COUNT(*) FROM project_whitelists WHERE project_id = $1", - ) - .bind(project_id) - .fetch_one(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; - - if count_row.0 >= MAX_WHITELIST_ENTRIES { - return Err(( - StatusCode::FORBIDDEN, - format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) - )); - } - - // Unique(project_id, allowed_url) is enforced by the DB; ignore conflicts + // Use a single query to check count and insert atomically to prevent race conditions + // This INSERT will only succeed if the count is below the limit let result = sqlx::query( "INSERT INTO project_whitelists (project_id, allowed_url) \ - VALUES ($1, $2) \ + SELECT $1, $2 \ + WHERE (SELECT COUNT(*) FROM project_whitelists WHERE project_id = $1) < $3 \ ON CONFLICT (project_id, allowed_url) DO NOTHING", ) .bind(project_id) .bind(trimmed_url) + .bind(MAX_WHITELIST_ENTRIES) .execute(&state.db) - .await; + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; - if let Err(e) = result { - return Err((StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e))); + // Check if the insert was successful or if limit was reached + if result.rows_affected() == 0 { + // Check if it's because of the limit or because it already exists + let exists: (bool,) = sqlx::query_as( + "SELECT EXISTS(SELECT 1 FROM project_whitelists WHERE project_id = $1 AND allowed_url = $2)", + ) + .bind(project_id) + .bind(trimmed_url) + .fetch_one(&state.db) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; + + if exists.0 { + // Entry already exists, this is acceptable + return Ok(StatusCode::CREATED); + } else { + // Must have hit the limit + return Err(( + StatusCode::FORBIDDEN, + format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) + )); + } } Ok(StatusCode::CREATED) @@ -594,7 +605,7 @@ pub async fn remove_from_whitelist( let user = user.ok_or((StatusCode::UNAUTHORIZED, "Unauthorized".to_string()))?; // Apply rate limiting: 20 requests per minute per user - let rate_limiter = state.get_whitelist_rate_limiter(user.id); + let rate_limiter = state.get_or_create_whitelist_rate_limiter(user.id); if rate_limiter.check().is_err() { return Err(( StatusCode::TOO_MANY_REQUESTS, diff --git a/server/src/state.rs b/server/src/state.rs index 474b10a..4969a87 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -5,6 +5,7 @@ use std::time::Instant; use dashmap::DashMap; use governor::{Quota, RateLimiter, clock::DefaultClock, state::{InMemoryState, NotKeyed}}; use std::num::NonZeroU32; +use crate::handlers::project::WHITELIST_RATE_LIMIT_PER_MINUTE; // New Struct to track container details AND ownership #[derive(Clone, Debug)] @@ -43,13 +44,12 @@ impl AppState { } } - /// Get or create a rate limiter for a user (20 requests per minute for whitelist operations) - pub fn get_whitelist_rate_limiter(&self, user_id: i64) -> Arc> { + /// Get or create a rate limiter for a user's whitelist operations + pub fn get_or_create_whitelist_rate_limiter(&self, user_id: i64) -> Arc> { self.whitelist_rate_limiters .entry(user_id) .or_insert_with(|| { - // 20 requests per 60 seconds - let quota = Quota::per_minute(NonZeroU32::new(20).unwrap()); + let quota = Quota::per_minute(NonZeroU32::new(WHITELIST_RATE_LIMIT_PER_MINUTE).unwrap()); Arc::new(RateLimiter::direct(quota)) }) .clone() From fdb09a32237da7b69d520b7e2b80732b05ce9700 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:00:08 +0000 Subject: [PATCH 4/5] Fix race condition using PostgreSQL advisory locks in transaction Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/handlers/project.rs | 82 +++++++++++++++++++++------------- 1 file changed, 52 insertions(+), 30 deletions(-) diff --git a/server/src/handlers/project.rs b/server/src/handlers/project.rs index 26600b5..e8272ee 100644 --- a/server/src/handlers/project.rs +++ b/server/src/handlers/project.rs @@ -549,45 +549,67 @@ pub async fn add_to_whitelist( None => return Err((StatusCode::NOT_FOUND, "Project not found".to_string())), }; - // Use a single query to check count and insert atomically to prevent race conditions - // This INSERT will only succeed if the count is below the limit - let result = sqlx::query( - "INSERT INTO project_whitelists (project_id, allowed_url) \ - SELECT $1, $2 \ - WHERE (SELECT COUNT(*) FROM project_whitelists WHERE project_id = $1) < $3 \ - ON CONFLICT (project_id, allowed_url) DO NOTHING", + // Use a database transaction with table-level advisory lock to prevent race conditions + let mut tx = state.db.begin() + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Transaction Error: {}", e)))?; + + // Get an advisory lock for this project's whitelist to prevent concurrent modifications + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(project_id) + .execute(&mut *tx) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Lock Error: {}", e)))?; + + // Check if entry already exists + let exists: (bool,) = sqlx::query_as( + "SELECT EXISTS(SELECT 1 FROM project_whitelists WHERE project_id = $1 AND allowed_url = $2)", ) .bind(project_id) .bind(trimmed_url) - .bind(MAX_WHITELIST_ENTRIES) - .execute(&state.db) + .fetch_one(&mut *tx) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; - // Check if the insert was successful or if limit was reached - if result.rows_affected() == 0 { - // Check if it's because of the limit or because it already exists - let exists: (bool,) = sqlx::query_as( - "SELECT EXISTS(SELECT 1 FROM project_whitelists WHERE project_id = $1 AND allowed_url = $2)", - ) - .bind(project_id) - .bind(trimmed_url) - .fetch_one(&state.db) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; + if exists.0 { + // Entry already exists, return success + tx.commit().await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Commit Error: {}", e)))?; + return Ok(StatusCode::CREATED); + } - if exists.0 { - // Entry already exists, this is acceptable - return Ok(StatusCode::CREATED); - } else { - // Must have hit the limit - return Err(( - StatusCode::FORBIDDEN, - format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) - )); - } + // Check current count + let count_row: (i64,) = sqlx::query_as( + "SELECT COUNT(*) FROM project_whitelists WHERE project_id = $1", + ) + .bind(project_id) + .fetch_one(&mut *tx) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; + + if count_row.0 >= MAX_WHITELIST_ENTRIES { + // Rollback transaction and return error + tx.rollback().await.ok(); + return Err(( + StatusCode::FORBIDDEN, + format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) + )); } + // Insert the new entry + sqlx::query( + "INSERT INTO project_whitelists (project_id, allowed_url) VALUES ($1, $2)", + ) + .bind(project_id) + .bind(trimmed_url) + .execute(&mut *tx) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Insert Error: {}", e)))?; + + // Commit transaction + tx.commit() + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Commit Error: {}", e)))?; + Ok(StatusCode::CREATED) } From 8298271e9d34c9e0478cd3c8a60fb9070561aaba Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 11 Feb 2026 05:01:04 +0000 Subject: [PATCH 5/5] Address final code review comments Co-authored-by: Yashb404 <139128977+Yashb404@users.noreply.github.com> --- server/src/handlers/project.rs | 7 +++---- server/src/state.rs | 5 ++++- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/server/src/handlers/project.rs b/server/src/handlers/project.rs index e8272ee..3c7ccb7 100644 --- a/server/src/handlers/project.rs +++ b/server/src/handlers/project.rs @@ -572,9 +572,9 @@ pub async fn add_to_whitelist( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; if exists.0 { - // Entry already exists, return success + // Entry already exists, commit transaction and return success tx.commit().await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Commit Error: {}", e)))?; - return Ok(StatusCode::CREATED); + return Ok(StatusCode::OK); // 200 OK - idempotent operation, entry already exists } // Check current count @@ -587,8 +587,7 @@ pub async fn add_to_whitelist( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB Error: {}", e)))?; if count_row.0 >= MAX_WHITELIST_ENTRIES { - // Rollback transaction and return error - tx.rollback().await.ok(); + // Transaction will automatically rollback when dropped return Err(( StatusCode::FORBIDDEN, format!("Maximum whitelist entries ({}) reached for this project", MAX_WHITELIST_ENTRIES) diff --git a/server/src/state.rs b/server/src/state.rs index 4969a87..e93230a 100644 --- a/server/src/state.rs +++ b/server/src/state.rs @@ -49,7 +49,10 @@ impl AppState { self.whitelist_rate_limiters .entry(user_id) .or_insert_with(|| { - let quota = Quota::per_minute(NonZeroU32::new(WHITELIST_RATE_LIMIT_PER_MINUTE).unwrap()); + let quota = Quota::per_minute( + NonZeroU32::new(WHITELIST_RATE_LIMIT_PER_MINUTE) + .expect("WHITELIST_RATE_LIMIT_PER_MINUTE must be non-zero") + ); Arc::new(RateLimiter::direct(quota)) }) .clone()