From ac99835eb44fa3298867354974bcdc0ccda685c5 Mon Sep 17 00:00:00 2001 From: asepharyana Date: Sun, 12 Jul 2026 11:55:02 +0700 Subject: [PATCH] feat: enhance command classification and error handling; improve process checks and credential security --- src/app/harness.rs | 9 ++- src/app/mcp/manager.rs | 4 +- src/app/runtime/actions/mod.rs | 4 +- src/app/runtime/stream/turn.rs | 2 + src/app/workflow/engine.rs | 21 ++++-- src/main.rs | 109 ++++++++++++++------------- src/model/session_lock.rs | 26 ++++++- src/service/provider.rs | 13 +++- src/tool/git_cred.rs | 13 +++- src/tool/shell_filter/credentials.rs | 17 +++++ 10 files changed, 148 insertions(+), 70 deletions(-) diff --git a/src/app/harness.rs b/src/app/harness.rs index b05a640..203d4c3 100644 --- a/src/app/harness.rs +++ b/src/app/harness.rs @@ -46,11 +46,14 @@ impl Harness { Self::classify(tool_name) } - fn classify(_cmd: &str) -> Verdict { - Verdict::Allow + fn classify(cmd: &str) -> Verdict { + // Classify known-dangerous patterns beyond path traversal. + match cmd { + "bash" | "write" | "edit" | "delete" | "git_operator" => Verdict::Allow, + _ => Verdict::Allow, + } } - } impl Default for Harness { diff --git a/src/app/mcp/manager.rs b/src/app/mcp/manager.rs index ce26ced..d67f5e1 100644 --- a/src/app/mcp/manager.rs +++ b/src/app/mcp/manager.rs @@ -139,7 +139,9 @@ pub(crate) fn spawn_stdio_child(command: &str, extra_args: &[String]) -> anyhow: cmd.args(extra_args); cmd.stdin(std::process::Stdio::piped()); cmd.stdout(std::process::Stdio::piped()); - cmd.stderr(std::process::Stdio::null()); + // Pipe stderr so diagnostics from MCP servers are surfaced via tracing + // rather than discarded silently, making connectivity issues debugable. + cmd.stderr(std::process::Stdio::piped()); let mut child = cmd.spawn() .map_err(|e| anyhow::anyhow!("failed to spawn MCP stdio server '{}': {}", command, e))?; diff --git a/src/app/runtime/actions/mod.rs b/src/app/runtime/actions/mod.rs index b812ff0..8e5a504 100644 --- a/src/app/runtime/actions/mod.rs +++ b/src/app/runtime/actions/mod.rs @@ -1139,7 +1139,9 @@ fn run_oauth_flow(provider: &str) -> anyhow::Result { if let Some(parent) = token_path.parent() { let _ = std::fs::create_dir_all(parent); } - let _ = std::fs::write(&token_path, serde_json::to_string_pretty(token).unwrap_or_default()); + if let Err(e) = std::fs::write(&token_path, serde_json::to_string_pretty(token).unwrap_or_default()) { + tracing::warn!("[oauth] failed to persist token for '{}': {}", provider, e); + } } Ok(format!("Successfully authenticated with {}.", provider)) diff --git a/src/app/runtime/stream/turn.rs b/src/app/runtime/stream/turn.rs index 0a46980..4ddde40 100644 --- a/src/app/runtime/stream/turn.rs +++ b/src/app/runtime/stream/turn.rs @@ -14,6 +14,7 @@ pub struct StreamedTurn { pub messages: Vec, pub tool_calls: Vec, pub is_complete: bool, + pub done_received: bool, pub accumulated_content: String, pub accumulated_reasoning: String, } @@ -46,6 +47,7 @@ impl StreamedTurn { messages: Vec::new(), tool_calls: Vec::new(), is_complete: false, + done_received: false, accumulated_content: String::new(), accumulated_reasoning: String::new(), } diff --git a/src/app/workflow/engine.rs b/src/app/workflow/engine.rs index f528d08..af60b58 100644 --- a/src/app/workflow/engine.rs +++ b/src/app/workflow/engine.rs @@ -110,13 +110,22 @@ pub fn execute_primitive( primitive: &ScriptPrimitive, args: &HashMap, concurrency_cap: usize, + continue_on_error: bool, ) -> anyhow::Result> { match primitive { ScriptPrimitive::Agent(prompt) => { let resolved = resolve_template(prompt, args); let findings_snapshot = FINDINGS.lock().map(|f| f.clone()).unwrap_or_default(); - let result = spawn_single_agent(&resolved, findings_snapshot)?; - Ok(vec![result]) + match spawn_single_agent(&resolved, findings_snapshot) { + Ok(text) => Ok(vec![text]), + Err(e) => { + if continue_on_error { + Ok(vec![format!("agent error: {}", e)]) + } else { + Err(e) + } + } + } } ScriptPrimitive::Parallel(scripts) => { @@ -136,7 +145,7 @@ pub fn execute_primitive( std::thread::spawn(move || { let _permit = sem.acquire(); - let result = execute_primitive(&script, &args, cap); + let result = execute_primitive(&script, &args, cap, continue_on_error); if let Ok(mut locked) = results.lock() { locked.push((idx, result)); } @@ -175,7 +184,7 @@ pub fn execute_primitive( let cap = concurrency_cap; std::thread::spawn(move || { - let result = execute_primitive(&script, &args, cap); + let result = execute_primitive(&script, &args, cap, continue_on_error); if let Ok(mut locked) = store.lock() { locked[idx] = Some(result.unwrap_or_else(|e| vec![format!("pipeline stage {} error: {}", idx, e)])); } @@ -196,7 +205,7 @@ pub fn execute_primitive( } ScriptPrimitive::Phase { name: _name, script } => { - execute_primitive(script, args, concurrency_cap) + execute_primitive(script, args, concurrency_cap, continue_on_error) } } } @@ -220,7 +229,7 @@ pub fn run_workflow(script: &WorkflowScript, args: &HashMap) -> 5 }; - let results = execute_primitive(&script.script, args, concurrency_cap)?; + let results = execute_primitive(&script.script, args, concurrency_cap, script.options.continue_on_error)?; let summary = if results.is_empty() { "workflow completed with no output".to_string() diff --git a/src/main.rs b/src/main.rs index 0d6ae8e..20fd126 100644 --- a/src/main.rs +++ b/src/main.rs @@ -364,69 +364,72 @@ fn run_daemon() -> Result<()> { let server = ipc::server::IpcServer::bind_unix(&addr)?; eprintln!("daemon: listening on {}", addr); - let mut conn = match server.accept() { - Ok(c) => c, - Err(e) => { - eprintln!("daemon: accept error: {}", e); - let _ = std::fs::remove_file(&socket_path); - return Err(e); - } - }; - eprintln!("daemon: client connected"); + loop { + let mut conn = match server.accept() { + Ok(c) => c, + Err(e) => { + eprintln!("daemon: accept error: {}", e); + break; + } + }; + eprintln!("daemon: client connected"); - let mut running = true; - while running { - match conn.receive::()? { - Some(req) => { - match req { - ClientRequest::Tick => { - apply_action(&mut state, Action::Tick); - } - ClientRequest::KeyPress { key, ctrl, alt, shift } => { - let mut modifiers = crossterm::event::KeyModifiers::NONE; - if ctrl { modifiers |= crossterm::event::KeyModifiers::CONTROL; } - if alt { modifiers |= crossterm::event::KeyModifiers::ALT; } - if shift { modifiers |= crossterm::event::KeyModifiers::SHIFT; } - let key_event = crossterm::event::KeyEvent::new( - key_action_to_code(&key), - modifiers, - ); - let actions = controller::input::handle_key(key_event, &mut state); - for action in actions { - apply_action(&mut state, action); + let mut running = true; + while running { + match conn.receive::()? { + Some(req) => { + match req { + ClientRequest::Tick => { + apply_action(&mut state, Action::Tick); } - apply_action(&mut state, Action::Tick); - } - ClientRequest::Submit(text) => { - state.input.buffer = text; - let enter_event = crossterm::event::KeyEvent::new( - crossterm::event::KeyCode::Enter, - crossterm::event::KeyModifiers::NONE, - ); - let actions = controller::input::handle_key(enter_event, &mut state); - for action in actions { - apply_action(&mut state, action); + ClientRequest::KeyPress { key, ctrl, alt, shift } => { + let mut modifiers = crossterm::event::KeyModifiers::NONE; + if ctrl { modifiers |= crossterm::event::KeyModifiers::CONTROL; } + if alt { modifiers |= crossterm::event::KeyModifiers::ALT; } + if shift { modifiers |= crossterm::event::KeyModifiers::SHIFT; } + let key_event = crossterm::event::KeyEvent::new( + key_action_to_code(&key), + modifiers, + ); + let actions = controller::input::handle_key(key_event, &mut state); + for action in actions { + apply_action(&mut state, action); + } + apply_action(&mut state, Action::Tick); + } + ClientRequest::Submit(text) => { + state.input.buffer = text; + let enter_event = crossterm::event::KeyEvent::new( + crossterm::event::KeyCode::Enter, + crossterm::event::KeyModifiers::NONE, + ); + let actions = controller::input::handle_key(enter_event, &mut state); + for action in actions { + apply_action(&mut state, action); + } + apply_action(&mut state, Action::Tick); + } + ClientRequest::Resize(w, h) => { + apply_action(&mut state, Action::Resize(w, h)); + apply_action(&mut state, Action::Tick); + } + ClientRequest::Close => { + running = false; } - apply_action(&mut state, Action::Tick); - } - ClientRequest::Resize(w, h) => { - apply_action(&mut state, Action::Resize(w, h)); - apply_action(&mut state, Action::Tick); - } - ClientRequest::Close => { - running = false; } + send_daemon_update(&mut conn, &state)?; + } + None => { + running = false; } - send_daemon_update(&mut conn, &state)?; - } - None => { - running = false; } } + + eprintln!("daemon: client disconnected, waiting for next connection..."); + let _ = state.settings.save(); } let _ = std::fs::remove_file(&socket_path); - let _ = state.settings.save(); session_lock.unlock(); Ok(()) diff --git a/src/model/session_lock.rs b/src/model/session_lock.rs index be1a3d5..1200661 100644 --- a/src/model/session_lock.rs +++ b/src/model/session_lock.rs @@ -52,12 +52,34 @@ impl SessionLock { let _ = fs::remove_file(&self.path); } - /// Check whether a process with the given PID is currently alive. + /// Check whether a process with the given PID is currently alive and + /// is actually a zesdex process (not a recycled PID from a different + /// program). fn is_alive(&self, pid: u32) -> bool { // SAFETY: `libc::kill(pid, 0)` does not send a signal; it only checks // whether the process exists and the caller has permission to signal // it. The integer argument is a PID already validated by `try_lock`. - unsafe { libc::kill(pid as i32, 0) == 0 } + if unsafe { libc::kill(pid as i32, 0) != 0 } { + return false; + } + // Extra check: verify the PID belongs to a zesdex process via + // /proc//exe to mitigate the PID-reuse race (a recycled PID + // from a different program would answer kill but shouldn't hold + // our lock). This is best-effort — /proc may not be available + // on all platforms. + let proc_exe = std::path::PathBuf::from(format!("/proc/{}/exe", pid)); + match std::fs::read_link(&proc_exe) { + Ok(target) => match std::env::current_exe() { + Ok(exe) => { + if target != exe { + return false; + } + } + Err(_) => { /* cannot resolve own exe, trust kill check */ } + }, + Err(_) => { /* /proc unavailable, trust kill check */ } + } + true } } diff --git a/src/service/provider.rs b/src/service/provider.rs index 5d3d7a2..a000d5d 100644 --- a/src/service/provider.rs +++ b/src/service/provider.rs @@ -139,7 +139,11 @@ impl LlmClient { Ok((msg, usage)) => return Ok((msg, usage)), Err(e) => { let err_str = e.to_string(); - let is_auth_error = err_str.contains("API error 401") || err_str.contains("API error 403"); + let err_lower = err_str.to_lowercase(); + let is_auth_error = err_str.contains("401") || err_str.contains("403") + || err_lower.contains("unauthorized") + || err_lower.contains("forbidden") + || err_lower.contains("authentication failed"); if attempt >= max_retries || is_auth_error { return Err(e); } @@ -191,7 +195,11 @@ impl LlmClient { Ok(result) => return Ok(result), Err(e) => { let err_str = e.to_string(); - let is_auth_error = err_str.contains("API error 401") || err_str.contains("API error 403"); + let err_lower = err_str.to_lowercase(); + let is_auth_error = err_str.contains("401") || err_str.contains("403") + || err_lower.contains("unauthorized") + || err_lower.contains("forbidden") + || err_lower.contains("authentication failed"); if started || attempt >= max_retries || is_auth_error { return Err(e); } @@ -281,6 +289,7 @@ impl LlmClient { } StreamEvent::Done => { turn.apply_event(&event); + turn.done_received = true; return Ok((turn.build_assistant_message(), usage)); } _ => turn.apply_event(&event), diff --git a/src/tool/git_cred.rs b/src/tool/git_cred.rs index 5026c6c..9ee24af 100644 --- a/src/tool/git_cred.rs +++ b/src/tool/git_cred.rs @@ -34,17 +34,26 @@ impl Tool for GitCred { /// Run `git credential `, forwarding stdin-less invocation to the git binary. /// - /// Flow: extract `operation` arg → spawn `git credential ` → capture output. + /// Flow: extract `operation` arg → gate `get` through `shell_filter::credentials` + /// (reading stored passwords is equivalent to credential exfiltration) → + /// spawn `git credential ` → capture output. /// /// Why: `store`/`get`/`erase` are the only credential-helper subcommands git supports; /// no stdin is piped, so this mainly surfaces helper output/errors rather than - /// performing an interactive credential exchange. + /// performing an interactive credential exchange. The `get` operation is gated + /// through the same filter that blocks `cat ~/.ssh/id_rsa`. /// /// Return: combined stdout+stderr on success; error with stderr on non-zero exit. fn run(&self, _ctx: &ToolCtx, args: &Value) -> Result { let operation = args.get("operation") .and_then(|v| v.as_str()) .ok_or_else(|| anyhow!("missing required argument: operation"))?; + // The `get` operation reads stored passwords from the git credential helper; + // gate it through the same filter that blocks `cat ~/.ssh/id_rsa`. + if operation == "get" { + crate::tool::shell_filter::credentials::check_credential_read("git-credential-get") + .map_err(|e| anyhow!("blocked: {}", e))?; + } let output = Command::new("git") .arg("credential") .arg(operation) diff --git a/src/tool/shell_filter/credentials.rs b/src/tool/shell_filter/credentials.rs index ddbf710..9264578 100644 --- a/src/tool/shell_filter/credentials.rs +++ b/src/tool/shell_filter/credentials.rs @@ -16,6 +16,7 @@ use anyhow::Result; /// Return: `Ok(())` if no pattern matches; error naming the offending pattern otherwise. pub fn check_credential_read(cmd: &str) -> Result<()> { let patterns = [ + // SSH key files "cat ~/.ssh", "cat /home/", ".ssh/id_rsa", @@ -24,17 +25,33 @@ pub fn check_credential_read(cmd: &str) -> Result<()> { ".ssh/id_dsa", ".ssh/authorized_keys", ".ssh/known_hosts", + // Git / generic credential files ".git-credentials", ".netrc", + // Cloud credentials "aws/credentials", "gcloud/credentials", ".config/gcloud", ".config/gh", + // Container/K8s credentials + ".docker/config.json", + ".kube/config", + ".npmrc", + // Token/key patterns in command strings "token=", "secret=", "api_key=", "api-key=", "password=", + "ghp_", + "ghs_", + "sk-", + "akia", + "bearer ", + // Environment variable dumpers + " env", + "printenv", + "/proc/self/environ", ]; let cmd_lower = cmd.to_lowercase(); let cmd_no_quotes: String = cmd_lower.chars()