diff --git a/src/agent/agent_loop.rs b/src/agent/agent_loop.rs index 3c9a3ca4..bdcc43b6 100644 --- a/src/agent/agent_loop.rs +++ b/src/agent/agent_loop.rs @@ -424,6 +424,29 @@ impl Agent { } } + // Safety validation for user input + let validation = self.safety().validate_input(content); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Ok(SubmissionResult::error(format!( + "Input rejected by safety validation: {}", + details + ))); + } + + let violations = self.safety().check_policy(content); + if violations + .iter() + .any(|rule| rule.action == crate::safety::PolicyAction::Block) + { + return Ok(SubmissionResult::error("Input rejected by safety policy.")); + } + // Handle explicit commands (starting with /) directly // Everything else goes through the normal agentic loop with tools let temp_message = IncomingMessage { @@ -767,6 +790,22 @@ impl Agent { name: tool_name.to_string(), })?; + // Validate tool parameters + let validation = self.safety().validator().validate_tool_params(params); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Err(crate::error::ToolError::InvalidParameters { + name: tool_name.to_string(), + reason: format!("Invalid tool parameters: {}", details), + } + .into()); + } + // Execute with timeout let result = tokio::time::timeout(std::time::Duration::from_secs(60), async { tool.execute(params.clone(), job_ctx).await @@ -813,11 +852,22 @@ impl Agent { title, description, category, - } => self.handle_create_job(title, description, category).await?, - MessageIntent::CheckJobStatus { job_id } => self.handle_check_status(job_id).await?, - MessageIntent::CancelJob { job_id } => self.handle_cancel_job(&job_id).await?, - MessageIntent::ListJobs { filter } => self.handle_list_jobs(filter).await?, - MessageIntent::HelpJob { job_id } => self.handle_help_job(&job_id).await?, + } => { + self.handle_create_job(&message.user_id, title, description, category) + .await? + } + MessageIntent::CheckJobStatus { job_id } => { + self.handle_check_status(&message.user_id, job_id).await? + } + MessageIntent::CancelJob { job_id } => { + self.handle_cancel_job(&message.user_id, &job_id).await? + } + MessageIntent::ListJobs { filter } => { + self.handle_list_jobs(&message.user_id, filter).await? + } + MessageIntent::HelpJob { job_id } => { + self.handle_help_job(&message.user_id, &job_id).await? + } MessageIntent::Command { command, args } => { match self.handle_command(&command, &args).await? { Some(s) => s, @@ -1223,6 +1273,7 @@ impl Agent { async fn handle_create_job( &self, + user_id: &str, title: String, description: String, category: Option, @@ -1230,7 +1281,7 @@ impl Agent { // Create job context let job_id = self .context_manager - .create_job(&title, &description) + .create_job_for_user(user_id, &title, &description) .await?; // Update category if provided @@ -1263,13 +1314,20 @@ impl Agent { )) } - async fn handle_check_status(&self, job_id: Option) -> Result { + async fn handle_check_status( + &self, + user_id: &str, + job_id: Option, + ) -> Result { match job_id { Some(id) => { let uuid = Uuid::parse_str(&id) .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } Ok(format!( "Job: {}\nStatus: {:?}\nCreated: {}\nStarted: {}\nActual cost: {}", @@ -1284,7 +1342,7 @@ impl Agent { } None => { // Show summary of all jobs - let summary = self.context_manager.summary().await; + let summary = self.context_manager.summary_for(user_id).await; Ok(format!( "Jobs summary:\n Total: {}\n In Progress: {}\n Completed: {}\n Failed: {}\n Stuck: {}", summary.total, @@ -1297,17 +1355,26 @@ impl Agent { } } - async fn handle_cancel_job(&self, job_id: &str) -> Result { + async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result { let uuid = Uuid::parse_str(job_id) .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; + let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } + self.scheduler.stop(uuid).await?; Ok(format!("Job {} has been cancelled.", job_id)) } - async fn handle_list_jobs(&self, _filter: Option) -> Result { - let jobs = self.context_manager.all_jobs().await; + async fn handle_list_jobs( + &self, + user_id: &str, + _filter: Option, + ) -> Result { + let jobs = self.context_manager.all_jobs_for(user_id).await; if jobs.is_empty() { return Ok("No jobs found.".to_string()); @@ -1316,18 +1383,23 @@ impl Agent { let mut output = String::from("Jobs:\n"); for job_id in jobs { if let Ok(ctx) = self.context_manager.get_context(job_id).await { - output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state)); + if ctx.user_id == user_id { + output.push_str(&format!(" {} - {} ({:?})\n", job_id, ctx.title, ctx.state)); + } } } Ok(output) } - async fn handle_help_job(&self, job_id: &str) -> Result { + async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result { let uuid = Uuid::parse_str(job_id) .map_err(|_| crate::error::JobError::NotFound { id: Uuid::nil() })?; let ctx = self.context_manager.get_context(uuid).await?; + if ctx.user_id != user_id { + return Err(crate::error::JobError::NotFound { id: uuid }.into()); + } if ctx.state == crate::context::JobState::Stuck { // Attempt recovery diff --git a/src/agent/scheduler.rs b/src/agent/scheduler.rs index 6c3adf8b..a88df10a 100644 --- a/src/agent/scheduler.rs +++ b/src/agent/scheduler.rs @@ -50,9 +50,9 @@ pub struct Scheduler { tools: Arc, store: Option>, /// Running jobs (main LLM-driven jobs). - jobs: RwLock>, + jobs: Arc>>, /// Running sub-tasks (tool executions, background tasks). - subtasks: RwLock>, + subtasks: Arc>>, } impl Scheduler { @@ -72,8 +72,8 @@ impl Scheduler { safety, tools, store, - jobs: RwLock::new(HashMap::new()), - subtasks: RwLock::new(HashMap::new()), + jobs: Arc::new(RwLock::new(HashMap::new())), + subtasks: Arc::new(RwLock::new(HashMap::new())), } } @@ -137,6 +137,27 @@ impl Scheduler { .await .insert(job_id, ScheduledJob { handle, tx }); + // Cleanup task for this job to avoid capacity leaks + let jobs = Arc::clone(&self.jobs); + tokio::spawn(async move { + loop { + let finished = { + let jobs_read = jobs.read().await; + match jobs_read.get(&job_id) { + Some(scheduled) => scheduled.handle.is_finished(), + None => true, + } + }; + + if finished { + jobs.write().await.remove(&job_id); + break; + } + + tokio::time::sleep(Duration::from_secs(1)).await; + } + }); + tracing::info!("Scheduled job {} for execution", job_id); Ok(()) } @@ -171,11 +192,13 @@ impl Scheduler { } => { let tools = self.tools.clone(); let context_manager = self.context_manager.clone(); + let safety = self.safety.clone(); tokio::spawn(async move { let result = Self::execute_tool_task( tools, context_manager, + safety, tool_parent_id, &tool_name, params, @@ -217,6 +240,27 @@ impl Scheduler { }, ); + // Cleanup task for subtask tracking + let subtasks = Arc::clone(&self.subtasks); + tokio::spawn(async move { + loop { + let finished = { + let subtasks_read = subtasks.read().await; + match subtasks_read.get(&task_id) { + Some(scheduled) => scheduled.handle.is_finished(), + None => true, + } + }; + + if finished { + subtasks.write().await.remove(&task_id); + break; + } + + tokio::time::sleep(Duration::from_secs(1)).await; + } + }); + tracing::debug!( parent_id = %parent_id, task_id = %task_id, @@ -282,6 +326,7 @@ impl Scheduler { async fn execute_tool_task( tools: Arc, context_manager: Arc, + safety: Arc, job_id: Uuid, tool_name: &str, params: serde_json::Value, @@ -297,6 +342,36 @@ impl Scheduler { // Get job context let job_ctx: JobContext = context_manager.get_context(job_id).await?; + if job_ctx.state == JobState::Cancelled { + return Err(crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: "Job is cancelled".to_string(), + } + .into()); + } + + if tool.requires_approval() { + return Err(crate::error::ToolError::AuthRequired { + name: tool_name.to_string(), + } + .into()); + } + + // Validate tool parameters + let validation = safety.validator().validate_tool_params(¶ms); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Err(crate::error::ToolError::InvalidParameters { + name: tool_name.to_string(), + reason: format!("Invalid tool parameters: {}", details), + } + .into()); + } // Execute with timeout let result = tokio::time::timeout(Duration::from_secs(60), async { diff --git a/src/agent/worker.rs b/src/agent/worker.rs index d49f8ac8..c05ece63 100644 --- a/src/agent/worker.rs +++ b/src/agent/worker.rs @@ -226,6 +226,14 @@ Report when the job is complete or if you encounter issues you cannot resolve."# } } + // Check for cancellation + if let Ok(ctx) = self.context_manager().get_context(self.job_id).await { + if ctx.state == JobState::Cancelled { + tracing::info!("Worker for job {} detected cancellation", self.job_id); + return Ok(()); + } + } + iteration += 1; if iteration > max_iterations { self.mark_stuck("Maximum iterations exceeded").await?; @@ -335,6 +343,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# let params = selection.parameters.clone(); let tools = self.tools().clone(); let context_manager = self.context_manager().clone(); + let safety = self.safety().clone(); let job_id = self.job_id; let store = self.deps.store.clone(); @@ -342,6 +351,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# let result = Self::execute_tool_inner( tools, context_manager, + safety, store, job_id, &tool_name, @@ -360,6 +370,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# async fn execute_tool_inner( tools: Arc, context_manager: Arc, + safety: Arc, store: Option>, job_id: Uuid, tool_name: &str, @@ -372,17 +383,39 @@ Report when the job is complete or if you encounter issues you cannot resolve."# name: tool_name.to_string(), })?; - // Log warning if tool requires approval (autonomous jobs auto-approve for now) + // Tools requiring approval are blocked in autonomous jobs if tool.requires_approval() { - tracing::warn!( - job_id = %job_id, - tool = %tool_name, - "Executing sensitive tool in autonomous job (auto-approved)" - ); + return Err(crate::error::ToolError::AuthRequired { + name: tool_name.to_string(), + } + .into()); } // Get job context for the tool let job_ctx = context_manager.get_context(job_id).await?; + if job_ctx.state == JobState::Cancelled { + return Err(crate::error::ToolError::ExecutionFailed { + name: tool_name.to_string(), + reason: "Job is cancelled".to_string(), + } + .into()); + } + + // Validate tool parameters + let validation = safety.validator().validate_tool_params(params); + if !validation.is_valid { + let details = validation + .errors + .iter() + .map(|e| format!("{}: {}", e.field, e.message)) + .collect::>() + .join("; "); + return Err(crate::error::ToolError::InvalidParameters { + name: tool_name.to_string(), + reason: format!("Invalid tool parameters: {}", details), + } + .into()); + } // Execute with timeout and timing let start = std::time::Instant::now(); @@ -395,7 +428,9 @@ Report when the job is complete or if you encounter issues you cannot resolve."# // Record action in memory and get the ActionRecord for persistence let action = match &result { Ok(Ok(output)) => { - let output_str = serde_json::to_string_pretty(&output.result).ok(); + let output_str = serde_json::to_string_pretty(&output.result) + .ok() + .map(|s| safety.sanitize_tool_output(tool_name, &s).content); context_manager .update_memory(job_id, |mem| { let rec = mem.create_action(tool_name, params.clone()).succeed( @@ -625,6 +660,7 @@ Report when the job is complete or if you encounter issues you cannot resolve."# Self::execute_tool_inner( self.tools().clone(), self.context_manager().clone(), + self.safety().clone(), self.deps.store.clone(), self.job_id, tool_name, diff --git a/src/channels/http.rs b/src/channels/http.rs index 121b8eda..986ebf7c 100644 --- a/src/channels/http.rs +++ b/src/channels/http.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use async_trait::async_trait; use axum::{ Json, Router, - extract::State, + extract::{DefaultBodyLimit, State}, http::StatusCode, response::IntoResponse, routing::{get, post}, @@ -36,8 +36,30 @@ struct HttpChannelState { shutdown_tx: RwLock>>, /// Expected webhook secret for authentication (if configured). webhook_secret: Option, + /// Fixed user ID for this HTTP channel. + user_id: String, + /// Rate limiting state. + rate_limit: tokio::sync::Mutex, } +#[derive(Debug)] +struct RateLimitState { + window_start: std::time::Instant, + request_count: u32, +} + +/// Maximum JSON body size for webhook requests (64 KB). +const MAX_BODY_BYTES: usize = 64 * 1024; + +/// Maximum number of pending wait-for-response requests. +const MAX_PENDING_RESPONSES: usize = 100; + +/// Maximum requests per minute. +const MAX_REQUESTS_PER_MINUTE: u32 = 60; + +/// Maximum content length for a single message. +const MAX_CONTENT_BYTES: usize = 32 * 1024; + impl HttpChannel { /// Create a new HTTP channel. pub fn new(config: HttpConfig) -> Self { @@ -45,6 +67,7 @@ impl HttpChannel { .webhook_secret .as_ref() .map(|s| s.expose_secret().to_string()); + let user_id = config.user_id.clone(); Self { config, @@ -53,6 +76,11 @@ impl HttpChannel { pending_responses: RwLock::new(std::collections::HashMap::new()), shutdown_tx: RwLock::new(None), webhook_secret, + user_id, + rate_limit: tokio::sync::Mutex::new(RateLimitState { + window_start: std::time::Instant::now(), + request_count: 0, + }), }), } } @@ -60,8 +88,9 @@ impl HttpChannel { #[derive(Debug, Deserialize)] struct WebhookRequest { - /// User or client identifier. - user_id: String, + /// User or client identifier (ignored, user is fixed by server config). + #[serde(default)] + user_id: Option, /// Message content. content: String, /// Optional thread ID for conversation tracking. @@ -100,6 +129,33 @@ async fn webhook_handler( State(state): State>, Json(req): Json, ) -> (StatusCode, Json) { + // Rate limiting + { + let mut limiter = state.rate_limit.lock().await; + if limiter.window_start.elapsed() >= std::time::Duration::from_secs(60) { + limiter.window_start = std::time::Instant::now(); + limiter.request_count = 0; + } + limiter.request_count += 1; + if limiter.request_count > MAX_REQUESTS_PER_MINUTE { + return ( + StatusCode::TOO_MANY_REQUESTS, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Rate limit exceeded".to_string()), + }), + ); + } + } + + let _ = req.user_id.as_ref().map(|user_id| { + tracing::debug!( + provided_user_id = %user_id, + "HTTP webhook request provided user_id, ignoring in favor of configured user_id" + ); + }); + // Validate secret if configured if let Some(ref expected_secret) = state.webhook_secret { match &req.secret { @@ -129,10 +185,22 @@ async fn webhook_handler( } } - let msg = - IncomingMessage::new("http", &req.user_id, &req.content).with_metadata(serde_json::json!({ + if req.content.len() > MAX_CONTENT_BYTES { + return ( + StatusCode::PAYLOAD_TOO_LARGE, + Json(WebhookResponse { + message_id: Uuid::nil(), + status: "error".to_string(), + response: Some("Content too large".to_string()), + }), + ); + } + + let msg = IncomingMessage::new("http", &state.user_id, &req.content).with_metadata( + serde_json::json!({ "wait_for_response": req.wait_for_response, - })); + }), + ); if let Some(thread_id) = &req.thread_id { let msg = msg.with_thread(thread_id); @@ -151,6 +219,17 @@ async fn process_message( // Set up response channel if waiting let response_rx = if wait_for_response { + if state.pending_responses.read().await.len() >= MAX_PENDING_RESPONSES { + return ( + StatusCode::TOO_MANY_REQUESTS, + Json(WebhookResponse { + message_id: msg_id, + status: "error".to_string(), + response: Some("Too many pending requests".to_string()), + }), + ); + } + let (tx, rx) = oneshot::channel(); state.pending_responses.write().await.insert(msg_id, tx); Some(rx) @@ -194,6 +273,9 @@ async fn process_message( None }; + // Ensure pending response entry is cleaned up on timeout or cancellation + let _ = state.pending_responses.write().await.remove(&msg_id); + ( StatusCode::OK, Json(WebhookResponse { @@ -211,6 +293,13 @@ impl Channel for HttpChannel { } async fn start(&self) -> Result { + if self.state.webhook_secret.is_none() { + return Err(ChannelError::StartupFailed { + name: "http".to_string(), + reason: "HTTP webhook secret is required (set HTTP_WEBHOOK_SECRET)".to_string(), + }); + } + let (tx, rx) = mpsc::channel(256); *self.state.tx.write().await = Some(tx); @@ -242,6 +331,7 @@ impl Channel for HttpChannel { let app = Router::new() .route("/health", get(health_handler)) .route("/webhook", post(webhook_handler)) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .with_state(state.clone()); // Create shutdown channel @@ -299,3 +389,22 @@ impl Channel for HttpChannel { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_http_channel_requires_secret() { + let config = HttpConfig { + host: "127.0.0.1".to_string(), + port: 0, + webhook_secret: None, + user_id: "http".to_string(), + }; + + let channel = HttpChannel::new(config); + let result = channel.start().await; + assert!(result.is_err()); + } +} diff --git a/src/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index e3acf888..b0927616 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -48,6 +48,7 @@ use crate::channels::wasm::runtime::{PreparedChannelModule, WasmChannelRuntime}; use crate::channels::wasm::schema::ChannelConfig; use crate::channels::{Channel, IncomingMessage, MessageStream, OutgoingResponse, StatusUpdate}; use crate::error::ChannelError; +use crate::safety::LeakDetector; use crate::tools::wasm::LogLevel; use crate::tools::wasm::WasmResourceLimiter; @@ -240,6 +241,15 @@ impl near::agent::channel_host::Host for ChannelStoreData { ); let url = injected_url; + let leak_detector = LeakDetector::new(); + let header_vec: Vec<(String, String)> = headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + leak_detector + .scan_http_request(&url, &header_vec, body.as_deref()) + .map_err(|e| format!("Potential secret leak blocked: {}", e))?; // Make the HTTP request using blocking I/O // We're already in a spawn_blocking context, so we can use block_on @@ -306,6 +316,13 @@ impl near::agent::channel_host::Host for ChannelStoreData { tracing::debug!(body = %truncated, "Response body"); } + // Leak detection on response body (best-effort) + if let Ok(body_str) = std::str::from_utf8(&body) { + leak_detector + .scan_and_clean(body_str) + .map_err(|e| format!("Potential secret leak in response: {}", e))?; + } + Ok(near::agent::channel_host::HttpResponse { status, headers_json, diff --git a/src/config.rs b/src/config.rs index 0ead0078..797d058c 100644 --- a/src/config.rs +++ b/src/config.rs @@ -349,6 +349,7 @@ pub struct HttpConfig { pub host: String, pub port: u16, pub webhook_secret: Option, + pub user_id: String, } impl ChannelsConfig { @@ -365,6 +366,7 @@ impl ChannelsConfig { })? .unwrap_or(8080), webhook_secret: optional_env("HTTP_WEBHOOK_SECRET")?.map(SecretString::from), + user_id: optional_env("HTTP_USER_ID")?.unwrap_or_else(|| "http".to_string()), }) } else { None diff --git a/src/context/manager.rs b/src/context/manager.rs index 78072445..7ef2d644 100644 --- a/src/context/manager.rs +++ b/src/context/manager.rs @@ -33,6 +33,17 @@ impl ContextManager { &self, title: impl Into, description: impl Into, + ) -> Result { + self.create_job_for_user("default", title, description) + .await + } + + /// Create a new job context for a specific user. + pub async fn create_job_for_user( + &self, + user_id: impl Into, + title: impl Into, + description: impl Into, ) -> Result { let contexts = self.contexts.read().await; let active_count = contexts.values().filter(|c| c.state.is_active()).count(); @@ -42,7 +53,7 @@ impl ContextManager { } drop(contexts); - let context = JobContext::new(title, description); + let context = JobContext::with_user(user_id, title, description); let job_id = context.job_id; let memory = Memory::new(job_id); @@ -113,6 +124,28 @@ impl ContextManager { self.contexts.read().await.keys().cloned().collect() } + /// List all active job IDs for a specific user. + pub async fn active_jobs_for(&self, user_id: &str) -> Vec { + self.contexts + .read() + .await + .iter() + .filter(|(_, c)| c.user_id == user_id && c.state.is_active()) + .map(|(id, _)| *id) + .collect() + } + + /// List all job IDs for a specific user. + pub async fn all_jobs_for(&self, user_id: &str) -> Vec { + self.contexts + .read() + .await + .iter() + .filter(|(_, c)| c.user_id == user_id) + .map(|(id, _)| *id) + .collect() + } + /// Get count of active jobs. pub async fn active_count(&self) -> usize { self.contexts @@ -174,6 +207,35 @@ impl ContextManager { summary.total = contexts.len(); summary } + + /// Get summary of all jobs for a specific user. + pub async fn summary_for(&self, user_id: &str) -> ContextSummary { + let contexts = self.contexts.read().await; + + let mut summary = ContextSummary::default(); + for ctx in contexts.values().filter(|c| c.user_id == user_id) { + match ctx.state { + crate::context::JobState::Pending => summary.pending += 1, + crate::context::JobState::InProgress => summary.in_progress += 1, + crate::context::JobState::Completed => summary.completed += 1, + crate::context::JobState::Submitted => summary.submitted += 1, + crate::context::JobState::Accepted => summary.accepted += 1, + crate::context::JobState::Failed => summary.failed += 1, + crate::context::JobState::Stuck => summary.stuck += 1, + crate::context::JobState::Cancelled => summary.cancelled += 1, + } + } + + summary.total = summary.pending + + summary.in_progress + + summary.completed + + summary.submitted + + summary.accepted + + summary.failed + + summary.stuck + + summary.cancelled; + summary + } } impl Default for ContextManager { @@ -209,6 +271,18 @@ mod tests { assert_eq!(context.title, "Test"); } + #[tokio::test] + async fn test_create_job_for_user_sets_user_id() { + let manager = ContextManager::new(5); + let job_id = manager + .create_job_for_user("user-123", "Test", "Description") + .await + .unwrap(); + + let context = manager.get_context(job_id).await.unwrap(); + assert_eq!(context.user_id, "user-123"); + } + #[tokio::test] async fn test_max_jobs_limit() { let manager = ContextManager::new(2); diff --git a/src/safety/mod.rs b/src/safety/mod.rs index 7223c4c1..ef93961a 100644 --- a/src/safety/mod.rs +++ b/src/safety/mod.rs @@ -16,7 +16,7 @@ pub use leak_detector::{ LeakAction, LeakDetectionError, LeakDetector, LeakMatch, LeakPattern, LeakScanResult, LeakSeverity, }; -pub use policy::{Policy, PolicyRule, Severity}; +pub use policy::{Policy, PolicyAction, PolicyRule, Severity}; pub use sanitizer::{InjectionWarning, SanitizedOutput, Sanitizer}; pub use validator::{ValidationResult, Validator}; @@ -27,6 +27,7 @@ pub struct SafetyLayer { sanitizer: Sanitizer, validator: Validator, policy: Policy, + leak_detector: LeakDetector, config: SafetyConfig, } @@ -37,6 +38,7 @@ impl SafetyLayer { sanitizer: Sanitizer::new(), validator: Validator::new(), policy: Policy::default(), + leak_detector: LeakDetector::new(), config: config.clone(), } } @@ -64,14 +66,55 @@ impl SafetyLayer { }; } + let mut content = output.to_string(); + let mut was_modified = false; + + // Leak detection and redaction + match self.leak_detector.scan_and_clean(&content) { + Ok(cleaned) => { + if cleaned != content { + was_modified = true; + content = cleaned; + } + } + Err(_) => { + return SanitizedOutput { + content: "[Output blocked due to potential secret leakage]".to_string(), + warnings: vec![], + was_modified: true, + }; + } + } + + // Safety policy enforcement + let violations = self.policy.check(&content); + if violations + .iter() + .any(|rule| rule.action == crate::safety::PolicyAction::Block) + { + return SanitizedOutput { + content: "[Output blocked by safety policy]".to_string(), + warnings: vec![], + was_modified: true, + }; + } + if violations + .iter() + .any(|rule| rule.action == crate::safety::PolicyAction::Sanitize) + { + was_modified = true; + } + // Run sanitization if enabled if self.config.injection_check_enabled { - self.sanitizer.sanitize(output) + let mut sanitized = self.sanitizer.sanitize(&content); + sanitized.was_modified = sanitized.was_modified || was_modified; + sanitized } else { SanitizedOutput { - content: output.to_string(), + content, warnings: vec![], - was_modified: false, + was_modified, } } } diff --git a/src/tools/builtin/file.rs b/src/tools/builtin/file.rs index 416cafd1..1b6ebf90 100644 --- a/src/tools/builtin/file.rs +++ b/src/tools/builtin/file.rs @@ -214,6 +214,10 @@ impl Tool for ReadFileTool { fn requires_sanitization(&self) -> bool { true // File content could contain anything } + + fn requires_approval(&self) -> bool { + true // Reading local files should require approval + } } /// Write file contents tool. @@ -422,6 +426,10 @@ impl Tool for ListDirTool { fn requires_sanitization(&self) -> bool { false // Directory listings are safe } + + fn requires_approval(&self) -> bool { + true // Directory listings can leak filesystem structure + } } /// Recursively list directory contents. diff --git a/src/tools/builtin/http.rs b/src/tools/builtin/http.rs index b9b25d3e..f5ac8c20 100644 --- a/src/tools/builtin/http.rs +++ b/src/tools/builtin/http.rs @@ -1,12 +1,14 @@ //! HTTP request tool. use std::collections::HashMap; +use std::net::IpAddr; use std::time::Duration; use async_trait::async_trait; use reqwest::Client; use crate::context::JobContext; +use crate::safety::LeakDetector; use crate::tools::tool::{Tool, ToolError, ToolOutput}; /// Tool for making HTTP requests. @@ -26,6 +28,58 @@ impl HttpTool { } } +fn validate_url(url: &str) -> Result { + let parsed = reqwest::Url::parse(url) + .map_err(|e| ToolError::InvalidParameters(format!("invalid URL: {}", e)))?; + + if parsed.scheme() != "https" { + return Err(ToolError::NotAuthorized( + "only https URLs are allowed".to_string(), + )); + } + + let host = parsed + .host_str() + .ok_or_else(|| ToolError::InvalidParameters("URL missing host".to_string()))?; + + let host_lower = host.to_lowercase(); + if host_lower == "localhost" || host_lower.ends_with(".localhost") { + return Err(ToolError::NotAuthorized( + "localhost is not allowed".to_string(), + )); + } + + if let Ok(ip) = host.parse::() { + if is_disallowed_ip(&ip) { + return Err(ToolError::NotAuthorized( + "private or local IPs are not allowed".to_string(), + )); + } + } + + Ok(parsed) +} + +fn is_disallowed_ip(ip: &IpAddr) -> bool { + match ip { + IpAddr::V4(v4) => { + v4.is_private() + || v4.is_loopback() + || v4.is_link_local() + || v4.is_multicast() + || v4.is_unspecified() + || *v4 == std::net::Ipv4Addr::new(169, 254, 169, 254) + } + IpAddr::V6(v6) => { + v6.is_loopback() + || v6.is_unique_local() + || v6.is_unicast_link_local() + || v6.is_multicast() + || v6.is_unspecified() + } + } +} + impl Default for HttpTool { fn default() -> Self { Self::new() @@ -90,20 +144,25 @@ impl Tool for HttpTool { .get("url") .and_then(|v| v.as_str()) .ok_or_else(|| ToolError::InvalidParameters("missing 'url' parameter".to_string()))?; + let parsed_url = validate_url(url)?; // Parse headers let headers: HashMap = params .get("headers") .and_then(|v| serde_json::from_value(v.clone()).ok()) .unwrap_or_default(); + let headers_vec: Vec<(String, String)> = headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); // Build request let mut request = match method.to_uppercase().as_str() { - "GET" => self.client.get(url), - "POST" => self.client.post(url), - "PUT" => self.client.put(url), - "DELETE" => self.client.delete(url), - "PATCH" => self.client.patch(url), + "GET" => self.client.get(parsed_url.clone()), + "POST" => self.client.post(parsed_url.clone()), + "PUT" => self.client.put(parsed_url.clone()), + "DELETE" => self.client.delete(parsed_url.clone()), + "PATCH" => self.client.patch(parsed_url.clone()), _ => { return Err(ToolError::InvalidParameters(format!( "unsupported method: {}", @@ -118,9 +177,20 @@ impl Tool for HttpTool { } // Add body if present - if let Some(body) = params.get("body") { + let body_bytes = if let Some(body) = params.get("body") { + let bytes = serde_json::to_vec(body) + .map_err(|e| ToolError::InvalidParameters(format!("invalid body JSON: {}", e)))?; request = request.json(body); - } + Some(bytes) + } else { + None + }; + + // Leak detection on outbound request (url/headers/body) + let detector = LeakDetector::new(); + detector + .scan_http_request(parsed_url.as_str(), &headers_vec, body_bytes.as_deref()) + .map_err(|e| ToolError::NotAuthorized(format!("{}", e)))?; // Execute request let response = request.send().await.map_err(|e| { @@ -168,3 +238,26 @@ impl Tool for HttpTool { true // HTTP requests go to external services, require user approval } } + +#[cfg(test)] +mod tests { + use super::validate_url; + + #[test] + fn test_validate_url_rejects_http() { + let err = validate_url("http://example.com").unwrap_err(); + assert!(err.to_string().contains("https")); + } + + #[test] + fn test_validate_url_rejects_localhost() { + let err = validate_url("https://localhost:8080").unwrap_err(); + assert!(err.to_string().contains("localhost")); + } + + #[test] + fn test_validate_url_accepts_https_public() { + let url = validate_url("https://example.com").unwrap(); + assert_eq!(url.host_str(), Some("example.com")); + } +} diff --git a/src/tools/builtin/job.rs b/src/tools/builtin/job.rs index 354fe14f..7f6663c9 100644 --- a/src/tools/builtin/job.rs +++ b/src/tools/builtin/job.rs @@ -56,7 +56,7 @@ impl Tool for CreateJobTool { async fn execute( &self, params: serde_json::Value, - _ctx: &JobContext, + ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); @@ -72,7 +72,11 @@ impl Tool for CreateJobTool { ToolError::InvalidParameters("missing 'description' parameter".into()) })?; - match self.context_manager.create_job(title, description).await { + match self + .context_manager + .create_job_for_user(&ctx.user_id, title, description) + .await + { Ok(job_id) => { let result = serde_json::json!({ "job_id": job_id.to_string(), @@ -133,7 +137,7 @@ impl Tool for ListJobsTool { async fn execute( &self, params: serde_json::Value, - _ctx: &JobContext, + ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); @@ -143,8 +147,8 @@ impl Tool for ListJobsTool { .unwrap_or("all"); let job_ids = match filter { - "active" => self.context_manager.active_jobs().await, - _ => self.context_manager.all_jobs().await, + "active" => self.context_manager.active_jobs_for(&ctx.user_id).await, + _ => self.context_manager.all_jobs_for(&ctx.user_id).await, }; let mut jobs = Vec::new(); @@ -168,7 +172,7 @@ impl Tool for ListJobsTool { } } - let summary = self.context_manager.summary().await; + let summary = self.context_manager.summary_for(&ctx.user_id).await; let result = serde_json::json!({ "jobs": jobs, @@ -226,9 +230,10 @@ impl Tool for JobStatusTool { async fn execute( &self, params: serde_json::Value, - _ctx: &JobContext, + ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); + let requester_id = ctx.user_id.clone(); let job_id_str = params .get("job_id") @@ -240,16 +245,22 @@ impl Tool for JobStatusTool { })?; match self.context_manager.get_context(job_id).await { - Ok(ctx) => { + Ok(job_ctx) => { + if job_ctx.user_id != requester_id { + let result = serde_json::json!({ + "error": "Job not found".to_string() + }); + return Ok(ToolOutput::success(result, start.elapsed())); + } let result = serde_json::json!({ "job_id": job_id.to_string(), - "title": ctx.title, - "description": ctx.description, - "status": format!("{:?}", ctx.state), - "created_at": ctx.created_at.to_rfc3339(), - "started_at": ctx.started_at.map(|t| t.to_rfc3339()), - "completed_at": ctx.completed_at.map(|t| t.to_rfc3339()), - "actual_cost": ctx.actual_cost.to_string() + "title": job_ctx.title, + "description": job_ctx.description, + "status": format!("{:?}", job_ctx.state), + "created_at": job_ctx.created_at.to_rfc3339(), + "started_at": job_ctx.started_at.map(|t| t.to_rfc3339()), + "completed_at": job_ctx.completed_at.map(|t| t.to_rfc3339()), + "actual_cost": job_ctx.actual_cost.to_string() }); Ok(ToolOutput::success(result, start.elapsed())) } @@ -304,9 +315,10 @@ impl Tool for CancelJobTool { async fn execute( &self, params: serde_json::Value, - _ctx: &JobContext, + ctx: &JobContext, ) -> Result { let start = std::time::Instant::now(); + let requester_id = ctx.user_id.clone(); let job_id_str = params .get("job_id") @@ -321,6 +333,9 @@ impl Tool for CancelJobTool { match self .context_manager .update_context(job_id, |ctx| { + if ctx.user_id != requester_id { + return Err("Job not found".to_string()); + } ctx.transition_to(JobState::Cancelled, Some("Cancelled by user".to_string())) }) .await