From 976c5bbf3ccaee70d5e7bb774f0bb87d644055bc Mon Sep 17 00:00:00 2001 From: argenis de la rosa Date: Sat, 14 Feb 2026 01:47:08 -0500 Subject: [PATCH] hardening: fix 7 production weaknesses found in codebase scan MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scan findings and fixes: 1. Gateway buffer overflow (8KB → 64KB) - Fixed: Increased request buffer from 8,192 to 65,536 bytes - Large POST bodies (long prompts) were silently truncated 2. Gateway slow-loris attack (no read timeout → 30s) - Fixed: tokio::time::timeout(30s) on stream.read() - Malicious clients could hold connections indefinitely 3. Webhook secret timing attack (== → constant_time_eq) - Fixed: Now uses constant_time_eq() for secret comparison - Prevents timing side-channel on webhook authentication 4. Pairing brute force (no limit → 5 attempts + 5min lockout) - Fixed: PairingGuard tracks failed attempts with lockout - Returns 429 Too Many Requests with retry_after seconds 5. Shell tool hang (no timeout → 60s kill) - Fixed: tokio::time::timeout(60s) on Command::output() - Commands that hang are killed and return error 6. Shell tool OOM (unbounded output → 1MB cap) - Fixed: stdout/stderr truncated at 1MB with warning - Prevents memory exhaustion from verbose commands 7. Provider HTTP timeout (none → 120s request + 10s connect) - Fixed: All 5 providers (OpenRouter, Anthropic, OpenAI, Ollama, Compatible) now have reqwest timeouts - Ollama gets 300s (local models are slower) 949 tests passing, 0 clippy warnings, cargo fmt clean --- src/gateway/mod.rs | 50 +++++++++------ src/providers/anthropic.rs | 6 +- src/providers/compatible.rs | 6 +- src/providers/ollama.rs | 6 +- src/providers/openai.rs | 6 +- src/providers/openrouter.rs | 6 +- src/security/pairing.rs | 118 +++++++++++++++++++++++++++++++++--- src/tools/shell.rs | 70 +++++++++++++++------ 8 files changed, 219 insertions(+), 49 deletions(-) diff --git a/src/gateway/mod.rs b/src/gateway/mod.rs index eac6571..294e199 100644 --- a/src/gateway/mod.rs +++ b/src/gateway/mod.rs @@ -1,9 +1,10 @@ use crate::config::Config; use crate::memory::{self, Memory, MemoryCategory}; use crate::providers::{self, Provider}; -use crate::security::pairing::{is_public_bind, PairingGuard}; +use crate::security::pairing::{constant_time_eq, is_public_bind, PairingGuard}; use anyhow::Result; use std::sync::Arc; +use std::time::Duration; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; @@ -106,9 +107,11 @@ pub async fn run_gateway(host: &str, port: u16, config: Config) -> Result<()> { let pairing = pairing.clone(); tokio::spawn(async move { - let mut buf = vec![0u8; 8192]; - let n = match stream.read(&mut buf).await { - Ok(n) if n > 0 => n, + // Read with 30s timeout to prevent slow-loris attacks + let mut buf = vec![0u8; 65_536]; // 64KB max request + let n = match tokio::time::timeout(Duration::from_secs(30), stream.read(&mut buf)).await + { + Ok(Ok(n)) if n > 0 => n, _ => return, }; @@ -179,18 +182,31 @@ async fn handle_request( // Pairing endpoint — exchange one-time code for bearer token ("POST", "/pair") => { let code = extract_header(request, "X-Pairing-Code").unwrap_or(""); - if let Some(token) = pairing.try_pair(code) { - tracing::info!("🔐 New client paired successfully"); - let body = serde_json::json!({ - "paired": true, - "token": token, - "message": "Save this token — use it as Authorization: Bearer " - }); - let _ = send_json(stream, 200, &body).await; - } else { - tracing::warn!("🔐 Pairing attempt with invalid code"); - let err = serde_json::json!({"error": "Invalid pairing code"}); - let _ = send_json(stream, 403, &err).await; + match pairing.try_pair(code) { + Ok(Some(token)) => { + tracing::info!("🔐 New client paired successfully"); + let body = serde_json::json!({ + "paired": true, + "token": token, + "message": "Save this token — use it as Authorization: Bearer " + }); + let _ = send_json(stream, 200, &body).await; + } + Ok(None) => { + tracing::warn!("🔐 Pairing attempt with invalid code"); + let err = serde_json::json!({"error": "Invalid pairing code"}); + let _ = send_json(stream, 403, &err).await; + } + Err(lockout_secs) => { + tracing::warn!( + "🔐 Pairing locked out — too many failed attempts ({lockout_secs}s remaining)" + ); + let err = serde_json::json!({ + "error": format!("Too many failed attempts. Try again in {lockout_secs}s."), + "retry_after": lockout_secs + }); + let _ = send_json(stream, 429, &err).await; + } } } @@ -213,7 +229,7 @@ async fn handle_request( if let Some(secret) = webhook_secret { let header_val = extract_header(request, "X-Webhook-Secret"); match header_val { - Some(val) if val == secret.as_ref() => {} + Some(val) if constant_time_eq(val, secret.as_ref()) => {} _ => { tracing::warn!( "Webhook: rejected request — invalid or missing X-Webhook-Secret" diff --git a/src/providers/anthropic.rs b/src/providers/anthropic.rs index 1b52826..9cddba1 100644 --- a/src/providers/anthropic.rs +++ b/src/providers/anthropic.rs @@ -38,7 +38,11 @@ impl AnthropicProvider { pub fn new(api_key: Option<&str>) -> Self { Self { api_key: api_key.map(ToString::to_string), - client: Client::new(), + client: Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), } } } diff --git a/src/providers/compatible.rs b/src/providers/compatible.rs index 90e4568..15f7a32 100644 --- a/src/providers/compatible.rs +++ b/src/providers/compatible.rs @@ -36,7 +36,11 @@ impl OpenAiCompatibleProvider { base_url: base_url.trim_end_matches('/').to_string(), api_key: api_key.map(ToString::to_string), auth_header: auth_style, - client: Client::new(), + client: Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), } } } diff --git a/src/providers/ollama.rs b/src/providers/ollama.rs index ee8c070..adc3e6e 100644 --- a/src/providers/ollama.rs +++ b/src/providers/ollama.rs @@ -44,7 +44,11 @@ impl OllamaProvider { .unwrap_or("http://localhost:11434") .trim_end_matches('/') .to_string(), - client: Client::new(), + client: Client::builder() + .timeout(std::time::Duration::from_secs(300)) // Ollama runs locally, may be slow + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), } } } diff --git a/src/providers/openai.rs b/src/providers/openai.rs index dae89fe..3481ce4 100644 --- a/src/providers/openai.rs +++ b/src/providers/openai.rs @@ -40,7 +40,11 @@ impl OpenAiProvider { pub fn new(api_key: Option<&str>) -> Self { Self { api_key: api_key.map(ToString::to_string), - client: Client::new(), + client: Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), } } } diff --git a/src/providers/openrouter.rs b/src/providers/openrouter.rs index 4e5bb78..3d99481 100644 --- a/src/providers/openrouter.rs +++ b/src/providers/openrouter.rs @@ -40,7 +40,11 @@ impl OpenRouterProvider { pub fn new(api_key: Option<&str>) -> Self { Self { api_key: api_key.map(ToString::to_string), - client: Client::new(), + client: Client::builder() + .timeout(std::time::Duration::from_secs(120)) + .connect_timeout(std::time::Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), } } } diff --git a/src/security/pairing.rs b/src/security/pairing.rs index ce338a6..08c99af 100644 --- a/src/security/pairing.rs +++ b/src/security/pairing.rs @@ -10,6 +10,12 @@ use std::collections::HashSet; use std::sync::Mutex; +use std::time::Instant; + +/// Maximum failed pairing attempts before lockout. +const MAX_PAIR_ATTEMPTS: u32 = 5; +/// Lockout duration after too many failed pairing attempts. +const PAIR_LOCKOUT_SECS: u64 = 300; // 5 minutes /// Manages pairing state for the gateway. #[derive(Debug)] @@ -20,6 +26,8 @@ pub struct PairingGuard { pairing_code: Option, /// Set of valid bearer tokens (persisted across restarts). paired_tokens: Mutex>, + /// Brute-force protection: failed attempt counter + lockout time. + failed_attempts: Mutex<(u32, Option)>, } impl PairingGuard { @@ -38,6 +46,7 @@ impl PairingGuard { require_pairing, pairing_code: code, paired_tokens: Mutex::new(tokens), + failed_attempts: Mutex::new((0, None)), } } @@ -52,19 +61,57 @@ impl PairingGuard { } /// Attempt to pair with the given code. Returns a bearer token on success. - pub fn try_pair(&self, code: &str) -> Option { + /// Returns `Err(lockout_seconds)` if locked out due to brute force. + pub fn try_pair(&self, code: &str) -> Result, u64> { + // Check brute force lockout + { + let attempts = self + .failed_attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let (count, Some(locked_at)) = &*attempts { + if *count >= MAX_PAIR_ATTEMPTS { + let elapsed = locked_at.elapsed().as_secs(); + if elapsed < PAIR_LOCKOUT_SECS { + return Err(PAIR_LOCKOUT_SECS - elapsed); + } + } + } + } + if let Some(ref expected) = self.pairing_code { if constant_time_eq(code.trim(), expected.trim()) { + // Reset failed attempts on success + { + let mut attempts = self + .failed_attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *attempts = (0, None); + } let token = generate_token(); let mut tokens = self .paired_tokens .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); tokens.insert(token.clone()); - return Some(token); + return Ok(Some(token)); } } - None + + // Increment failed attempts + { + let mut attempts = self + .failed_attempts + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + attempts.0 += 1; + if attempts.0 >= MAX_PAIR_ATTEMPTS { + attempts.1 = Some(Instant::now()); + } + } + + Ok(None) } /// Check if a bearer token is valid. @@ -117,7 +164,7 @@ fn generate_token() -> String { } /// Constant-time string comparison to prevent timing attacks on pairing code. -fn constant_time_eq(a: &str, b: &str) -> bool { +pub fn constant_time_eq(a: &str, b: &str) -> bool { if a.len() != b.len() { return false; } @@ -165,7 +212,7 @@ mod tests { fn try_pair_correct_code() { let guard = PairingGuard::new(true, &[]); let code = guard.pairing_code().unwrap().to_string(); - let token = guard.try_pair(&code); + let token = guard.try_pair(&code).unwrap(); assert!(token.is_some()); assert!(token.unwrap().starts_with("zc_")); assert!(guard.is_paired()); @@ -174,16 +221,16 @@ mod tests { #[test] fn try_pair_wrong_code() { let guard = PairingGuard::new(true, &[]); - let token = guard.try_pair("000000"); + let result = guard.try_pair("000000").unwrap(); // Might succeed if code happens to be 000000, but extremely unlikely - // Just check it doesn't panic - let _ = token; + // Just check it returns Ok(None) normally + let _ = result; } #[test] fn try_pair_empty_code() { let guard = PairingGuard::new(true, &[]); - assert!(guard.try_pair("").is_none()); + assert!(guard.try_pair("").unwrap().is_none()); } #[test] @@ -217,7 +264,7 @@ mod tests { fn pair_then_authenticate() { let guard = PairingGuard::new(true, &[]); let code = guard.pairing_code().unwrap().to_string(); - let token = guard.try_pair(&code).unwrap(); + let token = guard.try_pair(&code).unwrap().unwrap(); assert!(guard.is_authenticated(&token)); assert!(!guard.is_authenticated("wrong")); } @@ -273,4 +320,55 @@ mod tests { assert!(token.starts_with("zc_")); assert!(token.len() > 10); } + + // ── Brute force protection ─────────────────────────────── + + #[test] + fn brute_force_lockout_after_max_attempts() { + let guard = PairingGuard::new(true, &[]); + // Exhaust all attempts with wrong codes + for i in 0..MAX_PAIR_ATTEMPTS { + let result = guard.try_pair(&format!("wrong_{i}")); + assert!(result.is_ok(), "Attempt {i} should not be locked out yet"); + } + // Next attempt should be locked out + let result = guard.try_pair("another_wrong"); + assert!( + result.is_err(), + "Should be locked out after {MAX_PAIR_ATTEMPTS} attempts" + ); + let lockout_secs = result.unwrap_err(); + assert!(lockout_secs > 0, "Lockout should have remaining seconds"); + assert!( + lockout_secs <= PAIR_LOCKOUT_SECS, + "Lockout should not exceed max" + ); + } + + #[test] + fn correct_code_resets_failed_attempts() { + let guard = PairingGuard::new(true, &[]); + let code = guard.pairing_code().unwrap().to_string(); + // Fail a few times + for _ in 0..3 { + let _ = guard.try_pair("wrong"); + } + // Correct code should still work (under MAX_PAIR_ATTEMPTS) + let result = guard.try_pair(&code).unwrap(); + assert!(result.is_some(), "Correct code should work before lockout"); + } + + #[test] + fn lockout_returns_remaining_seconds() { + let guard = PairingGuard::new(true, &[]); + for _ in 0..MAX_PAIR_ATTEMPTS { + let _ = guard.try_pair("wrong"); + } + let err = guard.try_pair("wrong").unwrap_err(); + // Should be close to PAIR_LOCKOUT_SECS (within a second) + assert!( + err >= PAIR_LOCKOUT_SECS - 1, + "Remaining lockout should be ~{PAIR_LOCKOUT_SECS}s, got {err}s" + ); + } } diff --git a/src/tools/shell.rs b/src/tools/shell.rs index 8d29e5d..92a5582 100644 --- a/src/tools/shell.rs +++ b/src/tools/shell.rs @@ -3,6 +3,12 @@ use crate::security::SecurityPolicy; use async_trait::async_trait; use serde_json::json; use std::sync::Arc; +use std::time::Duration; + +/// Maximum shell command execution time before kill. +const SHELL_TIMEOUT_SECS: u64 = 60; +/// Maximum output size in bytes (1MB). +const MAX_OUTPUT_BYTES: usize = 1_048_576; /// Shell command execution tool with sandboxing pub struct ShellTool { @@ -53,25 +59,55 @@ impl Tool for ShellTool { }); } - let output = tokio::process::Command::new("sh") - .arg("-c") - .arg(command) - .current_dir(&self.security.workspace_dir) - .output() - .await?; + // Execute with timeout to prevent hanging commands + let result = tokio::time::timeout( + Duration::from_secs(SHELL_TIMEOUT_SECS), + tokio::process::Command::new("sh") + .arg("-c") + .arg(command) + .current_dir(&self.security.workspace_dir) + .output(), + ) + .await; - let stdout = String::from_utf8_lossy(&output.stdout).to_string(); - let stderr = String::from_utf8_lossy(&output.stderr).to_string(); + match result { + Ok(Ok(output)) => { + let mut stdout = String::from_utf8_lossy(&output.stdout).to_string(); + let mut stderr = String::from_utf8_lossy(&output.stderr).to_string(); - Ok(ToolResult { - success: output.status.success(), - output: stdout, - error: if stderr.is_empty() { - None - } else { - Some(stderr) - }, - }) + // Truncate output to prevent OOM + if stdout.len() > MAX_OUTPUT_BYTES { + stdout.truncate(MAX_OUTPUT_BYTES); + stdout.push_str("\n... [output truncated at 1MB]"); + } + if stderr.len() > MAX_OUTPUT_BYTES { + stderr.truncate(MAX_OUTPUT_BYTES); + stderr.push_str("\n... [stderr truncated at 1MB]"); + } + + Ok(ToolResult { + success: output.status.success(), + output: stdout, + error: if stderr.is_empty() { + None + } else { + Some(stderr) + }, + }) + } + Ok(Err(e)) => Ok(ToolResult { + success: false, + output: String::new(), + error: Some(format!("Failed to execute command: {e}")), + }), + Err(_) => Ok(ToolResult { + success: false, + output: String::new(), + error: Some(format!( + "Command timed out after {SHELL_TIMEOUT_SECS}s and was killed" + )), + }), + } } }