gitops/optimclaw
1
0
Fork 0
mirror of https://github.com/outbackdingo/optimclaw.git synced 2026-08-25 14:53:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Addressing vareity of security issues

Browse Source
This commit is contained in:
Illia Polosukhin
2026-02-05 10:40:24 -08:00
parent 0ab9643843
commit 3e6dfb8409
11 changed files with 602 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agent/agent_loop.rs
+85 -13
View File
@@ -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::<Vec<_>>()
.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::<Vec<_>>()
.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<String>,
@@ -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<String>) -> Result<String, Error> {
async fn handle_check_status(
&self,
user_id: &str,
job_id: Option<String>,
) -> Result<String, Error> {
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<String, Error> {
async fn handle_cancel_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
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<String>) -> Result<String, Error> {
let jobs = self.context_manager.all_jobs().await;
async fn handle_list_jobs(
&self,
user_id: &str,
_filter: Option<String>,
) -> Result<String, Error> {
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<String, Error> {
async fn handle_help_job(&self, user_id: &str, job_id: &str) -> Result<String, Error> {
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
src/agent/scheduler.rs
+79 -4
View File
@@ -50,9 +50,9 @@ pub struct Scheduler {
tools: Arc<ToolRegistry>,
store: Option<Arc<Store>>,
/// Running jobs (main LLM-driven jobs).
jobs: RwLock<HashMap<Uuid, ScheduledJob>>,
jobs: Arc<RwLock<HashMap<Uuid, ScheduledJob>>>,
/// Running sub-tasks (tool executions, background tasks).
subtasks: RwLock<HashMap<Uuid, ScheduledSubtask>>,
subtasks: Arc<RwLock<HashMap<Uuid, ScheduledSubtask>>>,
}
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<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
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(&params);
if !validation.is_valid {
let details = validation
.errors
.iter()
.map(|e| format!("{}: {}", e.field, e.message))
.collect::<Vec<_>>()
.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 {
src/agent/worker.rs
+43 -7
View File
@@ -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<ToolRegistry>,
context_manager: Arc<ContextManager>,
safety: Arc<SafetyLayer>,
store: Option<Arc<Store>>,
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::<Vec<_>>()
.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,
src/channels/http.rs
+115 -6
View File
@@ -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<Option<oneshot::Sender<()>>>,
/// Expected webhook secret for authentication (if configured).
webhook_secret: Option<String>,
/// Fixed user ID for this HTTP channel.
user_id: String,
/// Rate limiting state.
rate_limit: tokio::sync::Mutex<RateLimitState>,
}
#[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<String>,
/// Message content.
content: String,
/// Optional thread ID for conversation tracking.
@@ -100,6 +129,33 @@ async fn webhook_handler(
State(state): State<Arc<HttpChannelState>>,
Json(req): Json<WebhookRequest>,
) -> (StatusCode, Json<WebhookResponse>) {
// 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<MessageStream, ChannelError> {
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());
}
}
src/channels/wasm/wrapper.rs
+17
View File
@@ -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,
src/config.rs
+2
View File
@@ -349,6 +349,7 @@ pub struct HttpConfig {
pub host: String,
pub port: u16,
pub webhook_secret: Option<SecretString>,
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
src/context/manager.rs
+75 -1
View File
@@ -33,6 +33,17 @@ impl ContextManager {
&self,
title: impl Into<String>,
description: impl Into<String>,
) -> Result<Uuid, JobError> {
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<String>,
title: impl Into<String>,
description: impl Into<String>,
) -> Result<Uuid, JobError> {
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<Uuid> {
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<Uuid> {
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);
src/safety/mod.rs
+47 -4
View File
@@ -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,
}
}
}
src/tools/builtin/file.rs
+8
View File
@@ -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.
src/tools/builtin/http.rs
+100 -7
View File
@@ -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<reqwest::Url, ToolError> {
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::<IpAddr>() {
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<String, String> = 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"));
}
}
src/tools/builtin/job.rs
+31 -16
View File
@@ -56,7 +56,7 @@ impl Tool for CreateJobTool {
async fn execute(
&self,
params: serde_json::Value,
_ctx: &JobContext,
ctx: &JobContext,
) -> Result<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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<ToolOutput, ToolError> {
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