diff --git a/src/daemon.rs b/src/daemon.rs index 9293b3d..8276b1e 100644 --- a/src/daemon.rs +++ b/src/daemon.rs @@ -113,10 +113,11 @@ pub async fn run(database: &Arc, memory_store: &Arc) { )))); let registry = Arc::new(RegistryManager::load("tools")); - // 沙箱容器(如果启用):daemon 启动时创建长期容器,失败则回退本地执行 + // 沙箱容器(如果启用):daemon 启动时创建长期容器,失败则设置降级标志 + // executor 会据此回退本地执行,避免不受信任工具持续 SpawnFailed if crate::tools::sandbox::enabled() { if let Err(e) = crate::tools::sandbox::ensure_container("tools").await { - tracing::warn!(target: "ias::sandbox", "沙箱容器启动失败,回退本地执行: {e}"); + tracing::warn!(target: "ias::sandbox", "沙箱容器启动失败: {e}"); } } diff --git a/src/tools/approval.rs b/src/tools/approval.rs index 27dbc59..fcab5af 100644 --- a/src/tools/approval.rs +++ b/src/tools/approval.rs @@ -129,6 +129,33 @@ impl ApprovalManager { Ok((code_str, rx)) } + /// 测试专用:创建指定 TTL(秒)的审批请求 + #[cfg(test)] + async fn create_with_ttl( + &self, + user_id: &str, + skill_name: &str, + ttl_secs: u64, + ) -> Result<(String, oneshot::Receiver), String> { + let code: u32 = rand::rng().random_range(100000..999999); + let code_str = code.to_string(); + let code_hash = format!("{:x}", Sha256::digest(code_str.as_bytes())); + let (tx, rx) = oneshot::channel(); + + let mut map = self.pending.lock().await; + map.entry(user_id.to_string()) + .or_default() + .push(PendingEntry { + code_hash, + skill_name: skill_name.to_string(), + attempts_left: 3, + expires_at: Instant::now() + std::time::Duration::from_secs(ttl_secs), + sender: tx, + }); + + Ok((code_str, rx)) + } + /// 处理用户回复(由消息循环调用) /// 返回 (技能名, 决策);匹配失败返回 None pub async fn handle_reply( @@ -241,3 +268,66 @@ impl ApprovalManager { map.retain(|_, v| !v.is_empty()); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn clean_expired_only_removes_expired_entries() { + // 无 DB,纯内存 + let mgr = ApprovalManager::new(None); + + // 两条审批:一条立即过期(ttl=0),一条 300s 不过期 + let (_expired_code, expired_rx) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap(); + let (live_code, live_rx) = mgr.create_with_ttl("u1", "tool_b", 300).await.unwrap(); + + // 立即过期的条目 expires_at 已成过去,sleep 确保时间推进 + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + mgr.clean_expired().await; + + // 过期条目应收到 Expired + assert_eq!(expired_rx.await.unwrap(), ApprovalDecision::Expired); + // 未过期条目不应被误删:handle_reply 用正确确认码应能匹配 + let (name, decision) = mgr.handle_reply("u1", &live_code).await.unwrap(); + assert_eq!(name, "tool_b"); + assert_eq!(decision, ApprovalDecision::Approved); + // live_rx 应收到 Approved(未被 clean_expired 误标记) + assert_eq!(live_rx.await.unwrap(), ApprovalDecision::Approved); + } + + #[tokio::test] + async fn clean_expired_safe_when_no_pending() { + // 空 pending 不应 panic + let mgr = ApprovalManager::new(None); + mgr.clean_expired().await; + // 重复调用也安全 + mgr.clean_expired().await; + } + + #[tokio::test] + async fn clean_expired_skips_entries_already_removed_by_reply() { + // 回归保护:clean_expired 第一阶段收集过期 hash 后、第二阶段删除前, + // 若该条目已被 handle_reply 提前处理,clean_expired 应安全跳过,不误删其他条目。 + // (原 bug:按快照 index 删除会导致误删;现按 code_hash 删除应安全跳过) + let mgr = ApprovalManager::new(None); + + // 两条同用户审批,都设为立即过期 + let (_code_a, rx_a) = mgr.create_with_ttl("u1", "tool_a", 0).await.unwrap(); + let (_code_b, rx_b) = mgr.create_with_ttl("u1", "tool_b", 0).await.unwrap(); + + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + // 先用 handle_reply “0” 取消最早的条目(tool_a),模拟并发回复 + let (name, decision) = mgr.handle_reply("u1", "0").await.unwrap(); + assert_eq!(name, "tool_a"); + assert_eq!(decision, ApprovalDecision::Rejected); + assert_eq!(rx_a.await.unwrap(), ApprovalDecision::Rejected); + + // 再 clean_expired:tool_a 已被 remove,按 hash 查找应跳过; + // tool_b 仍在且过期,应被正确标记 Expired(不会被误当成 tool_a 跳过) + mgr.clean_expired().await; + assert_eq!(rx_b.await.unwrap(), ApprovalDecision::Expired); + } +} diff --git a/src/tools/executor.rs b/src/tools/executor.rs index 06361f5..6bee010 100644 --- a/src/tools/executor.rs +++ b/src/tools/executor.rs @@ -146,20 +146,15 @@ pub async fn execute_loop( // ── spawn ── // 执行隔离边界:启用沙箱且工具 profile 非 local 时,通过 docker exec - // 在长期容器内执行;否则本地 spawn(env_clear + 白名单注入) - let use_docker = - crate::tools::sandbox::enabled() && spec.sandbox.profile != "local"; - let mut cmd = if use_docker { - crate::tools::sandbox::build_exec_cmd(spec) + // 在长期容器内执行;否则本地 spawn(env_clear + 白名单注入)。 + // 沙箱降级(ensure 失败或容器崩溃未恢复)时自动回退本地执行。 + let use_docker = if crate::tools::sandbox::should_use_docker(&spec.sandbox.profile) { + // 执行前确保容器就绪:未运行则重建,重建失败则降级 + crate::tools::sandbox::ensure_ready_or_degrade("tools").await } else { - build_local_cmd(spec, binary_path) + false }; - cmd.stdin(std::process::Stdio::piped()) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .kill_on_drop(true); - - let mut child = match cmd.spawn() { + let mut child = match spawn_child(spec, binary_path, use_docker).await { Ok(c) => c, Err(e) => { error!(target: "ias::err","[spawn] {e}"); @@ -370,3 +365,48 @@ fn build_local_cmd(spec: &ToolSpec, binary_path: &str) -> Command { } cmd } + +/// spawn 工具进程。docker exec 模式下若首次 spawn 失败,尝试检测容器崩溃并重建后重试一次; +/// 重建失败则降级为本地执行(保留可用性,但仍受审批约束)。 +async fn spawn_child( + spec: &ToolSpec, + binary_path: &str, + use_docker: bool, +) -> Result { + let make_cmd = || { + let mut cmd = if use_docker { + crate::tools::sandbox::build_exec_cmd(spec) + } else { + build_local_cmd(spec, binary_path) + }; + cmd.stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true); + cmd + }; + + if let Ok(child) = make_cmd().spawn() { + return Ok(child); + } + + // 首次 spawn 失败。docker 模式下可能是容器崩溃,尝试重建后重试一次 + if use_docker && crate::tools::sandbox::recover_if_crashed("tools").await { + tracing::info!(target: "ias::sandbox", "沙箱容器已重建,重试执行"); + if let Ok(child) = make_cmd().spawn() { + return Ok(child); + } + // 重建后仍失败,降级本地执行 + tracing::warn!(target: "ias::sandbox", "重建后 exec 仍失败,降级本地执行"); + } else if use_docker { + tracing::warn!(target: "ias::sandbox", "容器不可用,降级本地执行"); + } + + // 降级:本地 spawn(审批已在入口完成,安全边界不受影响) + build_local_cmd(spec, binary_path) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .kill_on_drop(true) + .spawn() +} diff --git a/src/tools/sandbox.rs b/src/tools/sandbox.rs index 15c16f9..aab2433 100644 --- a/src/tools/sandbox.rs +++ b/src/tools/sandbox.rs @@ -9,7 +9,8 @@ //! 5 秒后仍不退出发 SIGKILL,确保容器内不累积孤儿进程。executor 的 tokio timeout //! 作为外层保险,kill `docker exec` 客户端释放宿主资源。 -use std::path::{Path, PathBuf}; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; use tokio::process::Command; use crate::tools::spec::ToolSpec; @@ -19,6 +20,10 @@ pub const CONTAINER_NAME: &str = "ias-sandbox"; /// 容器内产物挂载点 const MOUNT_TARGET: &str = "/tools"; +/// 降级标志:ensure_container 失败或容器崩溃后置 true, +/// executor 据此回退本地执行,避免不受信任工具持续 SpawnFailed +static DEGRADED: AtomicBool = AtomicBool::new(false); + /// 全局沙箱开关:`IAS_DOCKER_SANDBOX=1` 启用 pub fn enabled() -> bool { std::env::var("IAS_DOCKER_SANDBOX") @@ -27,13 +32,37 @@ pub fn enabled() -> bool { == Some("1") } +/// 当前是否处于降级状态(ensure 失败或容器崩溃) +pub fn degraded() -> bool { + DEGRADED.load(Ordering::SeqCst) +} + +/// 执行前检查:沙箱启用且未降级时才走 docker exec +pub fn should_use_docker(profile: &str) -> bool { + enabled() && profile != "local" && !degraded() +} + /// 确保 ias-sandbox 长期容器在运行。daemon 启动时调用。 /// /// - 容器已运行:直接复用 /// - 容器不存在/已停止:清理后重新创建 /// -/// 失败返回 Err,调用方应回退本地执行。 +/// 失败时设置降级标志,executor 据此回退本地执行。 pub async fn ensure_container(dir: &str) -> Result<(), String> { + match ensure_container_inner(dir).await { + Ok(()) => { + DEGRADED.store(false, Ordering::SeqCst); + Ok(()) + } + Err(e) => { + DEGRADED.store(true, Ordering::SeqCst); + tracing::warn!(target: "ias::sandbox", "沙箱降级为本地执行: {e}"); + Err(e) + } + } +} + +async fn ensure_container_inner(dir: &str) -> Result<(), String> { let root = locate_tools_root(dir)?; let bin_dir = root.join("bin"); std::fs::create_dir_all(&bin_dir).map_err(|e| format!("创建产物目录失败: {e}"))?; @@ -104,17 +133,57 @@ pub async fn ensure_container(dir: &str) -> Result<(), String> { Ok(()) } +/// exec 失败时调用:检测容器是否崩溃,崩溃则尝试重建。 +/// 重建成功返回 true(调用方可重试 exec);否则置降级标志返回 false。 +pub async fn recover_if_crashed(dir: &str) -> bool { + let running = Command::new("docker") + .args(["inspect", "-f", "{{.State.Running}}", CONTAINER_NAME]) + .output() + .await + .map(|o| String::from_utf8_lossy(&o.stdout).trim() == "true") + .unwrap_or(false); + + if running { + // 容器在运行,exec 失败是其他原因(如工具二进制不存在),不降级 + return false; + } + + tracing::warn!(target: "ias::sandbox", "沙箱容器已崩溃,尝试重建"); + ensure_container(dir).await.is_ok() +} + +/// 执行前调用:沙箱启用且未降级时,确保容器在运行。 +/// 容器不在运行则重建;重建失败置降级标志并返回 false(调用方降级本地执行)。 +pub async fn ensure_ready_or_degrade(dir: &str) -> bool { + if !enabled() || degraded() { + return false; + } + let running = Command::new("docker") + .args(["inspect", "-f", "{{.State.Running}}", CONTAINER_NAME]) + .output() + .await + .map(|o| String::from_utf8_lossy(&o.stdout).trim() == "true") + .unwrap_or(false); + if running { + return true; + } + // 容器不在运行,尝试重建 + if ensure_container(dir).await.is_ok() { + true + } else { + tracing::warn!(target: "ias::sandbox", "容器不可用,降级本地执行"); + false + } +} + /// 构建 `docker exec` 命令,在长期容器内执行工具。 /// /// 容器内路径 = MOUNT_TARGET + 二进制文件名。 /// 用 `timeout -k 5` 包裹:超时 SIGTERM,5 秒后 SIGKILL,避免孤儿进程。 /// env 按 spec 白名单注入(`-e` 必须在容器名前)。 pub fn build_exec_cmd(spec: &ToolSpec) -> Command { - let binary_name = Path::new(spec.path.as_deref().unwrap_or(&spec.tool)) - .file_name() - .map(|n| n.to_string_lossy().to_string()) - .unwrap_or_else(|| spec.tool.clone()); - let container_path = format!("{MOUNT_TARGET}/{binary_name}"); + // 直接用 spec.tool 作为容器内文件名,语义清晰,不依赖 spec.path 格式 + let container_path = format!("{MOUNT_TARGET}/{}", spec.tool); let mut cmd = Command::new("docker"); cmd.arg("exec").arg("-i");