Skip to main content

tiny_agent/sandbox/
host.rs

1use super::{ExecOutput, OutStream, OutputCallback, Sandbox, SandboxError};
2use async_trait::async_trait;
3use std::path::PathBuf;
4use std::process::Stdio;
5use std::sync::Arc;
6use tokio::io::{AsyncBufReadExt, AsyncRead, BufReader};
7use tokio::process::Command;
8
9/// 直接在宿主机上执行命令、读写真实文件系统的 sandbox。
10///
11/// 这是默认实现,行为等价于"不沙箱",不提供任何隔离。
12#[derive(Default)]
13pub struct HostSandbox {
14    /// 命令执行的工作目录;`None` = 继承当前进程的 cwd(原行为)。
15    root: Option<PathBuf>,
16}
17
18impl HostSandbox {
19    pub fn new() -> Self {
20        Self { root: None }
21    }
22
23    /// 把命令限定在 `root` 目录下执行(以它为工作目录)。
24    ///
25    /// 注意:这是**工作目录**级别的限定——相对路径会落在 `root` 内,但**不阻止**绝对路径
26    /// (如 `/etc/...`)或 `cd ..` 逃逸,因此**不是安全沙箱**,只是"默认在这个目录里干活"。
27    pub fn with_root(root: impl Into<PathBuf>) -> Self {
28        Self {
29            root: Some(root.into()),
30        }
31    }
32}
33
34#[async_trait]
35impl Sandbox for HostSandbox {
36    async fn open(self: Arc<Self>, _session_id: &str) -> Result<Arc<dyn Sandbox>, SandboxError> {
37        Ok(self)
38    }
39
40    async fn execute(
41        &self,
42        command: &str,
43        on_output: Option<OutputCallback>,
44    ) -> Result<ExecOutput, SandboxError> {
45        let mut cmd = Command::new("bash");
46        cmd.arg("-c")
47            .arg(command)
48            .stdin(Stdio::null())
49            .stdout(Stdio::piped())
50            .stderr(Stdio::piped())
51            // 被 drop(取消/超时杀任务)时连带杀掉直接子进程。
52            .kill_on_drop(true);
53        // 限定了 root 就以它为工作目录,命令里的相对路径都落在目录内。
54        if let Some(root) = &self.root {
55            cmd.current_dir(root);
56        }
57        let mut child = cmd.spawn()?;
58
59        let stdout = child.stdout.take().expect("stdout piped");
60        let stderr = child.stderr.take().expect("stderr piped");
61
62        // 两个流并发逐行读:每行既回调(实时)又累积(最终返回)。
63        let (out_acc, err_acc) = tokio::join!(
64            pump(stdout, OutStream::Stdout, on_output.clone()),
65            pump(stderr, OutStream::Stderr, on_output.clone()),
66        );
67
68        let status = child.wait().await?;
69        Ok(ExecOutput {
70            stdout: out_acc,
71            stderr: err_acc,
72            exit_code: status.code(),
73        })
74    }
75}
76
77/// 逐行读一个流:每行调一次回调(若有)并累积,EOF 后返回累积内容。
78async fn pump<R: AsyncRead + Unpin>(
79    reader: R,
80    stream: OutStream,
81    on_output: Option<OutputCallback>,
82) -> String {
83    let mut lines = BufReader::new(reader).lines();
84    let mut acc = String::new();
85    while let Ok(Some(line)) = lines.next_line().await {
86        if let Some(cb) = &on_output {
87            cb(stream, &line);
88        }
89        acc.push_str(&line);
90        acc.push('\n');
91    }
92    acc
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[tokio::test]
100    async fn execute_captures_stdout_and_exit_code() {
101        let sb = HostSandbox::new();
102        let out = sb.execute("echo hello", None).await.unwrap();
103        assert_eq!(out.stdout, "hello\n");
104        assert!(out.is_success());
105    }
106
107    #[tokio::test]
108    async fn execute_reports_nonzero_exit() {
109        let sb = HostSandbox::new();
110        let out = sb.execute("exit 3", None).await.unwrap();
111        assert_eq!(out.exit_code, Some(3));
112        assert!(!out.is_success());
113    }
114
115    #[tokio::test]
116    async fn with_root_runs_in_that_directory() {
117        let dir = std::env::temp_dir();
118        let sb = HostSandbox::with_root(&dir);
119        let out = sb.execute("pwd", None).await.unwrap();
120        // macOS 上 /var/folders 通过软链到 /private/var,两边都 canonicalize 再比。
121        let got = std::fs::canonicalize(out.stdout.trim()).unwrap();
122        let want = std::fs::canonicalize(&dir).unwrap();
123        assert_eq!(got, want);
124        // 相对路径写文件会落在 root 内。
125        let probe = format!("host_root_probe_{}.txt", std::process::id());
126        sb.execute(&format!("echo hi > {probe}"), None).await.unwrap();
127        let landed = dir.join(&probe);
128        assert!(landed.exists(), "相对路径应落在 root: {}", landed.display());
129        let _ = std::fs::remove_file(landed);
130    }
131}