1230 lines
43 KiB
Rust
1230 lines
43 KiB
Rust
//! 沙箱执行内核：以受限子进程方式运行外部命令，提供超时、取消、流式输出与进程组清理。
|
||
//!
|
||
//! 当前阶段隔离措施（跨 Linux / macOS 可用）：
|
||
//! - 独立进程组：超时或取消时对整个进程组发送 SIGKILL，清理其派生的子进程。
|
||
//! - 环境隔离：默认清空继承环境，仅注入白名单 + 请求指定变量。
|
||
//! - 工作目录隔离：默认每次执行使用 base 下的临时目录，执行后回收；可指定 base 内子目录（拒绝路径逃逸）。
|
||
//! - 进程级资源限额：fork 后、exec 前 setrlimit 限制 CPU / 内存 / 文件 / 进程数 / fd（见 rlimits 模块）。
|
||
//! - 超时强制：达到上限（min(请求, 全局)）即 SIGKILL 进程组。
|
||
//! - 取消：调用方持 trace_id 经 ``cancel`` 中止运行中任务（SIGKILL 进程组），产出标记 ``cancelled``。
|
||
//! - 流式输出：``run_stream`` 经 mpsc channel 实时推送 stdout/stderr 增量，供 SSE 端点转发。
|
||
//! - 输出截断：stdout / stderr 各自截断至上限，超出标记 truncated 并继续 drain 至 EOF 以防管道阻塞。
|
||
//! - 可观测：每次执行记录 Prometheus 指标（计数/耗时/并发/流式）。
|
||
//!
|
||
//! 容器化部署时可在 Docker 侧叠加 seccomp / cgroup v2 资源限制（见 docker/sandbox/Dockerfile 与 sandbox/README.md）。
|
||
|
||
use std::collections::HashMap;
|
||
use std::path::PathBuf;
|
||
use std::process::Stdio;
|
||
use std::sync::Arc;
|
||
use std::time::{Duration, Instant};
|
||
|
||
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt};
|
||
use tokio::process::Command;
|
||
use tokio::sync::{mpsc, oneshot, Mutex};
|
||
|
||
use crate::config::Config;
|
||
use crate::docker_exec::{self, detect_language_from_argv};
|
||
use crate::metrics::Metrics;
|
||
use crate::rlimits::EffectiveRlimit;
|
||
use crate::warm_pool::WarmPool;
|
||
|
||
/// 执行错误：`InvalidArgs` 映射 HTTP 400（客户端可修正），`Internal` 映射 HTTP 500。
|
||
#[derive(Debug)]
|
||
pub enum ExecError {
|
||
/// 请求参数不合法，例如 argv 为空、workdir 路径逃逸。
|
||
InvalidArgs(String),
|
||
/// 内部错误：`user_msg` 返回客户端（通用、脱敏，不含路径/errno），
|
||
/// `detail` 含底层细节只进 tracing 日志（带 trace_id 关联）。
|
||
Internal { user_msg: &'static str, detail: String },
|
||
}
|
||
|
||
impl std::fmt::Display for ExecError {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
ExecError::InvalidArgs(m) => write!(f, "参数错误: {m}"),
|
||
ExecError::Internal { user_msg, .. } => write!(f, "{user_msg}"),
|
||
}
|
||
}
|
||
}
|
||
|
||
impl std::error::Error for ExecError {}
|
||
|
||
/// 运行中任务的句柄：持有 pid 与取消信号发送端，供 ``cancel`` 中止。
|
||
struct RunningHandle {
|
||
#[allow(dead_code)]
|
||
pid: u32,
|
||
cancel_tx: Option<oneshot::Sender<()>>,
|
||
}
|
||
|
||
/// 已 spawn 的子进程及其管道、取消信号与临时工作目录持有者。
|
||
struct SpawnedChild {
|
||
child: tokio::process::Child,
|
||
pid: Option<u32>,
|
||
stdout: tokio::process::ChildStdout,
|
||
stderr: tokio::process::ChildStderr,
|
||
cancel_rx: oneshot::Receiver<()>,
|
||
start: Instant,
|
||
/// 临时工作目录：持有到子进程退出后才回收（None 表示用 base 内持久子目录）。
|
||
_temp_keep: Option<tempfile::TempDir>,
|
||
}
|
||
|
||
/// 流式执行事件：stdout / stderr 增量，或执行结束的汇总。
|
||
pub enum StreamEvent {
|
||
Stdout(Vec<u8>),
|
||
Stderr(Vec<u8>),
|
||
Done(ExecOutcome),
|
||
}
|
||
|
||
#[derive(Clone, Copy)]
|
||
enum OutputKind {
|
||
Stdout,
|
||
Stderr,
|
||
}
|
||
|
||
pub struct Executor {
|
||
config: Arc<Config>,
|
||
base_canon: PathBuf,
|
||
/// 运行中任务注册表：trace_id → 句柄。cancel 端点据此定位并 SIGKILL 进程组。
|
||
running: Arc<Mutex<HashMap<String, RunningHandle>>>,
|
||
metrics: Arc<Metrics>,
|
||
}
|
||
|
||
/// 单次执行的产出。
|
||
#[derive(Debug)]
|
||
pub struct ExecOutcome {
|
||
pub exit_code: Option<i32>,
|
||
pub stdout: String,
|
||
pub stderr: String,
|
||
pub duration_ms: u128,
|
||
pub timeout: bool,
|
||
pub cancelled: bool,
|
||
pub stdout_truncated: bool,
|
||
pub stderr_truncated: bool,
|
||
}
|
||
|
||
/// 把执行结果分类为 Prometheus 标签值。
|
||
fn classify_outcome(o: &ExecOutcome) -> &'static str {
|
||
if o.cancelled {
|
||
"cancelled"
|
||
} else if o.timeout {
|
||
"timeout"
|
||
} else if o.exit_code == Some(0) {
|
||
"completed"
|
||
} else {
|
||
"failed"
|
||
}
|
||
}
|
||
|
||
impl Executor {
|
||
pub fn new(config: Arc<Config>) -> std::io::Result<Self> {
|
||
let base_canon = config
|
||
.base_workdir
|
||
.canonicalize()
|
||
.unwrap_or_else(|_| config.base_workdir.clone());
|
||
let metrics = Arc::new(Metrics::new(config.max_concurrent, config.auth_enabled));
|
||
Ok(Self {
|
||
config,
|
||
base_canon,
|
||
running: Arc::new(Mutex::new(HashMap::new())),
|
||
metrics,
|
||
})
|
||
}
|
||
|
||
/// 暴露指标集合（供 /metrics 端点 gather）。
|
||
pub fn metrics(&self) -> &Arc<Metrics> {
|
||
&self.metrics
|
||
}
|
||
|
||
/// 沙箱工作根目录（供文件 IO 端点复用路径校验与读写）。
|
||
pub fn base(&self) -> &std::path::Path {
|
||
&self.config.base_workdir
|
||
}
|
||
|
||
/// 沙箱工作根目录的 canonical 形式（路径逃逸校验基准）。
|
||
pub fn base_canon(&self) -> &std::path::Path {
|
||
&self.base_canon
|
||
}
|
||
|
||
/// 取消运行中的任务。返回是否命中（未命中表示任务不存在或已完成）。
|
||
/// 命中后向 run()/run_stream() 的 select! 发送取消信号，由其 SIGKILL 进程组并标记 cancelled。
|
||
pub async fn cancel(&self, trace_id: &str) -> bool {
|
||
let handle = self.running.lock().await.remove(trace_id);
|
||
if let Some(h) = handle {
|
||
if let Some(tx) = h.cancel_tx {
|
||
let _ = tx.send(());
|
||
}
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
|
||
/// 共享的 spawn 前置：工作目录解析、子进程组装、stdin 写入、注册 cancel。
|
||
/// 返回的 SpawnedChild 持有 child 与管道；调用方负责 select! 等待与注销 trace_id。
|
||
async fn spawn_tracked(
|
||
&self,
|
||
argv: Vec<String>,
|
||
stdin: Option<String>,
|
||
workdir: Option<String>,
|
||
env: Vec<(String, String)>,
|
||
trace_id: String,
|
||
effective: Vec<EffectiveRlimit>,
|
||
) -> Result<SpawnedChild, ExecError> {
|
||
if argv.is_empty() {
|
||
return Err(ExecError::InvalidArgs("argv 不能为空".into()));
|
||
}
|
||
|
||
// 1) 工作目录：临时目录或 base 内子目录（拒绝逃逸）
|
||
let (workdir, temp_keep) = match workdir.as_deref() {
|
||
None => {
|
||
let td = tempfile::TempDir::new_in(&self.config.base_workdir).map_err(|e| {
|
||
ExecError::Internal {
|
||
user_msg: "工作目录不可用",
|
||
detail: format!("创建临时工作目录失败: {e}"),
|
||
}
|
||
})?;
|
||
(td.path().to_path_buf(), Some(td))
|
||
}
|
||
Some(rel) => {
|
||
let candidate = self.config.base_workdir.join(rel);
|
||
tokio::fs::create_dir_all(&candidate).await.map_err(|e| {
|
||
ExecError::Internal {
|
||
user_msg: "工作目录不可用",
|
||
detail: format!("创建工作目录失败: {e}"),
|
||
}
|
||
})?;
|
||
let canon = tokio::fs::canonicalize(&candidate)
|
||
.await
|
||
.map_err(|e| ExecError::Internal {
|
||
user_msg: "工作目录不可用",
|
||
detail: format!("解析工作目录失败: {e}"),
|
||
})?;
|
||
if !canon.starts_with(&self.base_canon) {
|
||
return Err(ExecError::InvalidArgs(format!(
|
||
"workdir 逃逸出沙箱根目录: {rel}"
|
||
)));
|
||
}
|
||
(canon, None)
|
||
}
|
||
};
|
||
|
||
// 2) 组装子进程
|
||
let mut cmd = Command::new(&argv[0]);
|
||
cmd.args(&argv[1..]);
|
||
cmd.env_clear();
|
||
// env_clear 后始终注入兜底 PATH，避免子进程找不到外部命令。
|
||
// allowed_env 若有 PATH 会覆盖此默认值。
|
||
cmd.env("PATH", "/usr/local/bin:/usr/bin:/bin");
|
||
for k in &self.config.allowed_env {
|
||
if let Ok(v) = std::env::var(k) {
|
||
cmd.env(k, v);
|
||
}
|
||
}
|
||
for (k, v) in &env {
|
||
cmd.env(k, v);
|
||
}
|
||
cmd.current_dir(&workdir);
|
||
cmd.stdin(Stdio::piped());
|
||
cmd.stdout(Stdio::piped());
|
||
cmd.stderr(Stdio::piped());
|
||
cmd.process_group(0); // 子进程成为新进程组 leader，pgid == pid
|
||
|
||
// 进程级资源限额：fork 后、exec 前 setrlimit（容器级隔离之外的纵深防御）。
|
||
// 非 unix 平台无 setrlimit 等价物，跳过（rlimits::apply 已是空操作）。
|
||
#[cfg(unix)]
|
||
{
|
||
let is_root = unsafe { libc::getuid() } == 0;
|
||
let eff = effective.clone();
|
||
// pre_exec 闭包在 fork 后、exec 前的子进程运行，仅调用 async-signal-safe 的 setrlimit，
|
||
// 不分配内存、不持锁；tokio 将其标记 unsafe 以强制调用方确认这些约束。
|
||
unsafe {
|
||
cmd.pre_exec(move || crate::rlimits::apply(&eff, is_root));
|
||
}
|
||
}
|
||
#[cfg(not(unix))]
|
||
{
|
||
let _ = &effective;
|
||
}
|
||
|
||
let start = Instant::now();
|
||
// 重试逻辑：EAGAIN / ENOMEM / EMFILE 等瞬态错误做指数退避重试
|
||
let retryable_errnos = [libc::EAGAIN, libc::ENOMEM, libc::EMFILE];
|
||
let retry_delays = [
|
||
Duration::from_millis(200),
|
||
Duration::from_millis(400),
|
||
Duration::from_millis(800),
|
||
];
|
||
let mut child = {
|
||
let mut result = cmd.spawn();
|
||
let mut attempt = 0usize;
|
||
loop {
|
||
match result {
|
||
Ok(child) => break child,
|
||
Err(e) => {
|
||
let is_retryable = e
|
||
.raw_os_error()
|
||
.map(|errno| retryable_errnos.contains(&errno))
|
||
.unwrap_or(false);
|
||
if !is_retryable {
|
||
return Err(ExecError::Internal {
|
||
user_msg: "无法启动进程",
|
||
detail: format!("spawn 失败: {e}"),
|
||
});
|
||
}
|
||
if attempt >= retry_delays.len() {
|
||
return Err(ExecError::Internal {
|
||
user_msg: "无法启动进程",
|
||
detail: format!(
|
||
"spawn 失败（已重试 {} 次）: {e}",
|
||
retry_delays.len()
|
||
),
|
||
});
|
||
}
|
||
tracing::warn!(
|
||
attempt = attempt + 1,
|
||
max_retries = retry_delays.len(),
|
||
error = %e,
|
||
"spawn 重试（进程/文件/内存资源不足）"
|
||
);
|
||
tokio::time::sleep(retry_delays[attempt]).await;
|
||
result = cmd.spawn();
|
||
attempt += 1;
|
||
}
|
||
}
|
||
}
|
||
};
|
||
|
||
// 3) 写入 stdin（若有）
|
||
if let Some(input) = stdin {
|
||
if let Some(mut child_stdin) = child.stdin.take() {
|
||
tokio::spawn(async move {
|
||
let _ = child_stdin.write_all(input.as_bytes()).await;
|
||
let _ = child_stdin.shutdown().await;
|
||
});
|
||
}
|
||
} else {
|
||
drop(child.stdin.take());
|
||
}
|
||
|
||
let pid = child.id();
|
||
let stdout = child.stdout.take().expect("stdout piped");
|
||
let stderr = child.stderr.take().expect("stderr piped");
|
||
|
||
// 4) 注册运行中任务 + 取消信号
|
||
let (cancel_tx, cancel_rx) = oneshot::channel::<()>();
|
||
{
|
||
let mut map = self.running.lock().await;
|
||
map.insert(
|
||
trace_id.clone(),
|
||
RunningHandle {
|
||
pid: pid.unwrap_or(0),
|
||
cancel_tx: Some(cancel_tx),
|
||
},
|
||
);
|
||
}
|
||
|
||
Ok(SpawnedChild {
|
||
child,
|
||
pid,
|
||
stdout,
|
||
stderr,
|
||
cancel_rx,
|
||
start,
|
||
_temp_keep: temp_keep,
|
||
})
|
||
}
|
||
|
||
/// 同步执行：收集全部 stdout/stderr 后返回。
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub async fn run(
|
||
&self,
|
||
argv: Vec<String>,
|
||
stdin: Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
workdir: Option<String>,
|
||
env: Vec<(String, String)>,
|
||
max_output_bytes: Option<usize>,
|
||
trace_id: String,
|
||
effective: Vec<EffectiveRlimit>,
|
||
) -> Result<ExecOutcome, ExecError> {
|
||
let global_ms = self.config.max_timeout.as_millis() as u64;
|
||
let timeout = Duration::from_millis(timeout_ms.unwrap_or(global_ms).min(global_ms));
|
||
let max_out = max_output_bytes.unwrap_or(self.config.max_output_bytes);
|
||
|
||
let SpawnedChild {
|
||
mut child,
|
||
pid,
|
||
stdout,
|
||
stderr,
|
||
cancel_rx,
|
||
start,
|
||
_temp_keep,
|
||
} = self
|
||
.spawn_tracked(argv, stdin, workdir, env, trace_id.clone(), effective)
|
||
.await?;
|
||
// spawn 成功：计入并发
|
||
self.metrics.concurrent_inc();
|
||
|
||
let out_task = tokio::spawn(capture_limited(stdout, max_out));
|
||
let err_task = tokio::spawn(capture_limited(stderr, max_out));
|
||
|
||
let mut timeout_hit = false;
|
||
let mut cancelled = false;
|
||
let status: Option<std::process::ExitStatus> = tokio::select! {
|
||
res = child.wait() => match res {
|
||
Ok(s) => Some(s),
|
||
Err(e) => {
|
||
self.running.lock().await.remove(&trace_id);
|
||
self.metrics.concurrent_dec();
|
||
return Err(ExecError::Internal {
|
||
user_msg: "执行失败",
|
||
detail: format!("wait 子进程失败: {e}"),
|
||
});
|
||
}
|
||
},
|
||
_ = tokio::time::sleep(timeout) => {
|
||
timeout_hit = true;
|
||
if let Some(pid) = pid {
|
||
unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
|
||
}
|
||
child.wait().await.ok()
|
||
}
|
||
_ = cancel_rx => {
|
||
cancelled = true;
|
||
if let Some(pid) = pid {
|
||
unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
|
||
}
|
||
child.wait().await.ok()
|
||
}
|
||
};
|
||
|
||
self.running.lock().await.remove(&trace_id);
|
||
drop(_temp_keep);
|
||
|
||
let (out_bytes, out_trunc) = out_task.await.unwrap_or_default();
|
||
let (err_bytes, err_trunc) = err_task.await.unwrap_or_default();
|
||
|
||
let outcome = ExecOutcome {
|
||
exit_code: status.and_then(|s| s.code()),
|
||
stdout: String::from_utf8_lossy(&out_bytes).into_owned(),
|
||
stderr: String::from_utf8_lossy(&err_bytes).into_owned(),
|
||
duration_ms: start.elapsed().as_millis(),
|
||
timeout: timeout_hit,
|
||
cancelled,
|
||
stdout_truncated: out_trunc,
|
||
stderr_truncated: err_trunc,
|
||
};
|
||
self.metrics.concurrent_dec();
|
||
self.metrics
|
||
.record_exec(classify_outcome(&outcome), outcome.duration_ms);
|
||
Ok(outcome)
|
||
}
|
||
|
||
/// 流式执行：返回事件 Receiver，stdout/stderr 增量实时推送，末尾一个 Done。
|
||
/// `permit` 持有到驱动 task 结束（并发闸），由调用方从 Semaphore acquire 后传入。
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub async fn run_stream(
|
||
&self,
|
||
argv: Vec<String>,
|
||
stdin: Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
workdir: Option<String>,
|
||
env: Vec<(String, String)>,
|
||
max_output_bytes: Option<usize>,
|
||
trace_id: String,
|
||
permit: Option<tokio::sync::OwnedSemaphorePermit>,
|
||
effective: Vec<EffectiveRlimit>,
|
||
) -> Result<mpsc::Receiver<StreamEvent>, ExecError> {
|
||
let global_ms = self.config.max_timeout.as_millis() as u64;
|
||
let timeout = Duration::from_millis(timeout_ms.unwrap_or(global_ms).min(global_ms));
|
||
let max_out = max_output_bytes.unwrap_or(self.config.max_output_bytes);
|
||
|
||
let SpawnedChild {
|
||
child,
|
||
pid,
|
||
stdout,
|
||
stderr,
|
||
cancel_rx,
|
||
start,
|
||
_temp_keep,
|
||
} = self
|
||
.spawn_tracked(argv, stdin, workdir, env, trace_id.clone(), effective)
|
||
.await?;
|
||
// spawn 成功：计入并发
|
||
self.metrics.concurrent_inc();
|
||
|
||
let (tx, rx) = mpsc::channel::<StreamEvent>(64);
|
||
let out_task =
|
||
tokio::spawn(capture_streaming(stdout, max_out, OutputKind::Stdout, tx.clone()));
|
||
let err_task =
|
||
tokio::spawn(capture_streaming(stderr, max_out, OutputKind::Stderr, tx.clone()));
|
||
|
||
let running = self.running.clone();
|
||
let metrics = self.metrics.clone();
|
||
let done_trace_id = trace_id.clone();
|
||
tokio::spawn(async move {
|
||
let _permit = permit;
|
||
let _temp_keep = _temp_keep;
|
||
let mut child = child;
|
||
|
||
// 客户端断开探测：`tx.closed()` 在接收端（rx）drop 时完成（== SSE 断连）。
|
||
// 对有输出命令，capture_streaming 的 send 也会先一步感知；这里覆盖无输出
|
||
// 长命令（如 sleep 30）——单纯靠 send err 检测会漏掉。单次 select! 即可，无需轮询。
|
||
|
||
let mut timeout_hit = false;
|
||
let mut cancelled = false;
|
||
let status: Option<std::process::ExitStatus> = tokio::select! {
|
||
res = child.wait() => res.ok(),
|
||
_ = tokio::time::sleep(timeout) => {
|
||
timeout_hit = true;
|
||
if let Some(pid) = pid {
|
||
unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
|
||
}
|
||
child.wait().await.ok()
|
||
}
|
||
_ = cancel_rx => {
|
||
cancelled = true;
|
||
if let Some(pid) = pid {
|
||
unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
|
||
}
|
||
child.wait().await.ok()
|
||
}
|
||
_ = tx.closed() => {
|
||
// 接收端已关闭：客户端断开，取消子进程并收尸
|
||
cancelled = true;
|
||
if let Some(pid) = pid {
|
||
unsafe { libc::kill(-(pid as i32), libc::SIGKILL) };
|
||
}
|
||
child.wait().await.ok()
|
||
}
|
||
};
|
||
|
||
running.lock().await.remove(&done_trace_id);
|
||
let out_trunc = out_task.await.unwrap_or(false);
|
||
let err_trunc = err_task.await.unwrap_or(false);
|
||
|
||
let outcome = ExecOutcome {
|
||
exit_code: status.and_then(|s| s.code()),
|
||
stdout: String::new(),
|
||
stderr: String::new(),
|
||
duration_ms: start.elapsed().as_millis(),
|
||
timeout: timeout_hit,
|
||
cancelled,
|
||
stdout_truncated: out_trunc,
|
||
stderr_truncated: err_trunc,
|
||
};
|
||
metrics.concurrent_dec();
|
||
metrics.record_stream(classify_outcome(&outcome));
|
||
|
||
let _ = tx.send(StreamEvent::Done(outcome)).await;
|
||
});
|
||
|
||
Ok(rx)
|
||
}
|
||
|
||
/// 统一入口：根据隔离模式路由到进程或容器执行路径。
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub async fn run_routed(
|
||
&self,
|
||
argv: Vec<String>,
|
||
stdin: Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
workdir: Option<String>,
|
||
env: Vec<(String, String)>,
|
||
max_output_bytes: Option<usize>,
|
||
trace_id: String,
|
||
effective: Vec<EffectiveRlimit>,
|
||
isolation_mode: &str,
|
||
warm_pool: Option<Arc<WarmPool>>,
|
||
allow_fallback: bool,
|
||
) -> Result<ExecOutcome, ExecError> {
|
||
match (isolation_mode, warm_pool) {
|
||
("container", Some(pool)) => {
|
||
match self.run_container(pool, &argv, &stdin, timeout_ms, &env, max_output_bytes, &trace_id).await {
|
||
Ok(outcome) => Ok(outcome),
|
||
Err(e) if allow_fallback => {
|
||
tracing::warn!(
|
||
error = %e,
|
||
trace_id = %trace_id,
|
||
"容器执行失败，降级为进程模式"
|
||
);
|
||
self.metrics.record_downgrade("container_fallback");
|
||
self.run(argv, stdin, timeout_ms, workdir, env, max_output_bytes, trace_id, effective).await
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
error = %e,
|
||
trace_id = %trace_id,
|
||
"容器执行失败（hard isolation policy，不降级）"
|
||
);
|
||
self.metrics.record_downgrade("container_hard_fail");
|
||
Err(e)
|
||
}
|
||
}
|
||
}
|
||
_ => {
|
||
self.run(argv, stdin, timeout_ms, workdir, env, max_output_bytes, trace_id, effective).await
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 容器执行路径（支持取消）。
|
||
async fn run_container(
|
||
&self,
|
||
pool: Arc<WarmPool>,
|
||
argv: &[String],
|
||
stdin: &Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
env: &[(String, String)],
|
||
max_output_bytes: Option<usize>,
|
||
trace_id: &str,
|
||
) -> Result<ExecOutcome, ExecError> {
|
||
let global_ms = self.config.max_timeout.as_millis() as u64;
|
||
let timeout = Duration::from_millis(timeout_ms.unwrap_or(global_ms).min(global_ms));
|
||
let max_out = max_output_bytes.unwrap_or(self.config.max_output_bytes);
|
||
let lang = detect_language_from_argv(argv);
|
||
|
||
let handle = pool
|
||
.clone()
|
||
.acquire(lang)
|
||
.await
|
||
.map_err(|e| ExecError::Internal {
|
||
user_msg: "容器资源不可用",
|
||
detail: format!("warm_pool acquire 失败: {e}"),
|
||
})?;
|
||
|
||
let start = Instant::now();
|
||
|
||
// 注册到 running map 支持取消
|
||
let (cancel_tx, cancel_rx) = oneshot::channel::<()>();
|
||
{
|
||
let mut map = self.running.lock().await;
|
||
map.remove(trace_id);
|
||
map.insert(
|
||
trace_id.to_string(),
|
||
RunningHandle {
|
||
pid: 0,
|
||
cancel_tx: Some(cancel_tx),
|
||
},
|
||
);
|
||
}
|
||
|
||
let outcome = tokio::select! {
|
||
result = docker_exec::exec_in_container(
|
||
pool.docker(),
|
||
&handle.id,
|
||
argv,
|
||
stdin.clone(),
|
||
env,
|
||
timeout,
|
||
max_out,
|
||
) => {
|
||
match result {
|
||
Ok(o) => o,
|
||
Err(e) => {
|
||
self.running.lock().await.remove(trace_id);
|
||
pool.release(handle).await;
|
||
return Err(e);
|
||
}
|
||
}
|
||
}
|
||
_ = cancel_rx => {
|
||
// 用户取消：杀死容器（exec 随之终止）
|
||
let _ = pool.docker()
|
||
.kill_container(&handle.id, None::<bollard::container::KillContainerOptions<String>>)
|
||
.await;
|
||
ExecOutcome {
|
||
exit_code: None,
|
||
stdout: String::new(),
|
||
stderr: String::new(),
|
||
duration_ms: start.elapsed().as_millis(),
|
||
timeout: false,
|
||
cancelled: true,
|
||
stdout_truncated: false,
|
||
stderr_truncated: false,
|
||
}
|
||
}
|
||
};
|
||
|
||
self.metrics
|
||
.record_container_exec(lang, outcome.duration_ms);
|
||
|
||
self.running.lock().await.remove(trace_id);
|
||
pool.release(handle).await;
|
||
Ok(outcome)
|
||
}
|
||
|
||
/// 流式容器执行路径。
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub async fn run_stream_routed(
|
||
&self,
|
||
argv: Vec<String>,
|
||
stdin: Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
workdir: Option<String>,
|
||
env: Vec<(String, String)>,
|
||
max_output_bytes: Option<usize>,
|
||
trace_id: String,
|
||
effective: Vec<EffectiveRlimit>,
|
||
isolation_mode: &str,
|
||
warm_pool: Option<Arc<WarmPool>>,
|
||
permit: Option<tokio::sync::OwnedSemaphorePermit>,
|
||
allow_fallback: bool,
|
||
) -> Result<mpsc::Receiver<StreamEvent>, ExecError> {
|
||
match (isolation_mode, warm_pool) {
|
||
("container", Some(pool)) => {
|
||
match self.run_stream_container(pool, &argv, &stdin, timeout_ms, &env, max_output_bytes, &trace_id).await {
|
||
Ok(rx) => Ok(rx),
|
||
Err(e) if allow_fallback => {
|
||
tracing::warn!(
|
||
error = %e,
|
||
trace_id = %trace_id,
|
||
"流式容器执行失败，降级为进程模式"
|
||
);
|
||
self.metrics.record_downgrade("container_stream_fallback");
|
||
self.run_stream(argv, stdin, timeout_ms, workdir, env, max_output_bytes, trace_id, permit, effective).await
|
||
}
|
||
Err(e) => {
|
||
tracing::warn!(
|
||
error = %e,
|
||
trace_id = %trace_id,
|
||
"流式容器执行失败（hard isolation policy，不降级）"
|
||
);
|
||
self.metrics.record_downgrade("container_stream_hard_fail");
|
||
Err(e)
|
||
}
|
||
}
|
||
}
|
||
_ => {
|
||
self.run_stream(argv, stdin, timeout_ms, workdir, env, max_output_bytes, trace_id, permit, effective).await
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 流式容器执行。
|
||
async fn run_stream_container(
|
||
&self,
|
||
pool: Arc<WarmPool>,
|
||
argv: &[String],
|
||
stdin: &Option<String>,
|
||
timeout_ms: Option<u64>,
|
||
env: &[(String, String)],
|
||
max_output_bytes: Option<usize>,
|
||
_trace_id: &str,
|
||
) -> Result<mpsc::Receiver<StreamEvent>, ExecError> {
|
||
let global_ms = self.config.max_timeout.as_millis() as u64;
|
||
let timeout = Duration::from_millis(timeout_ms.unwrap_or(global_ms).min(global_ms));
|
||
let max_out = max_output_bytes.unwrap_or(self.config.max_output_bytes);
|
||
let lang = detect_language_from_argv(argv);
|
||
|
||
let handle = pool
|
||
.clone()
|
||
.acquire(lang)
|
||
.await
|
||
.map_err(|e| ExecError::Internal {
|
||
user_msg: "容器资源不可用",
|
||
detail: format!("warm_pool acquire 失败: {e}"),
|
||
})?;
|
||
|
||
let rx = docker_exec::exec_in_container_stream(
|
||
pool.docker(),
|
||
&handle.id,
|
||
argv,
|
||
stdin.clone(),
|
||
env,
|
||
timeout,
|
||
max_out,
|
||
)
|
||
.await?;
|
||
|
||
pool.release(handle).await;
|
||
Ok(rx)
|
||
}
|
||
}
|
||
|
||
/// 读取流到 `max` 字节即停止存储，但继续 drain 至 EOF 以防管道阻塞导致子进程挂死。
|
||
async fn capture_limited<R: AsyncRead + Unpin>(reader: R, max: usize) -> (Vec<u8>, bool) {
|
||
let mut reader = reader;
|
||
let mut buf = Vec::new();
|
||
let mut truncated = false;
|
||
let mut chunk = [0u8; 8192];
|
||
loop {
|
||
match reader.read(&mut chunk).await {
|
||
Ok(0) => break,
|
||
Ok(n) => {
|
||
if buf.len() < max {
|
||
let room = max - buf.len();
|
||
if n <= room {
|
||
buf.extend_from_slice(&chunk[..n]);
|
||
} else {
|
||
buf.extend_from_slice(&chunk[..room]);
|
||
truncated = true;
|
||
}
|
||
} else {
|
||
truncated = true;
|
||
}
|
||
}
|
||
Err(_) => break,
|
||
}
|
||
}
|
||
(buf, truncated)
|
||
}
|
||
|
||
/// 流式读取：把不超过 `max` 的增量发往 channel，超出后继续 drain 到 EOF（防管道阻塞），返回是否截断。
|
||
async fn capture_streaming<R: AsyncRead + Unpin>(
|
||
reader: R,
|
||
max: usize,
|
||
kind: OutputKind,
|
||
tx: mpsc::Sender<StreamEvent>,
|
||
) -> bool {
|
||
let mut reader = reader;
|
||
let mut written = 0usize;
|
||
let mut truncated = false;
|
||
let mut chunk = [0u8; 8192];
|
||
loop {
|
||
match reader.read(&mut chunk).await {
|
||
Ok(0) => break,
|
||
Ok(n) => {
|
||
if written < max {
|
||
let room = max - written;
|
||
let take = n.min(room);
|
||
if take > 0 {
|
||
let event = match kind {
|
||
OutputKind::Stdout => StreamEvent::Stdout(chunk[..take].to_vec()),
|
||
OutputKind::Stderr => StreamEvent::Stderr(chunk[..take].to_vec()),
|
||
};
|
||
if tx.send(event).await.is_err() {
|
||
break; // 接收端已关闭（客户端断开）
|
||
}
|
||
}
|
||
written += take;
|
||
if n > room {
|
||
truncated = true;
|
||
}
|
||
} else {
|
||
truncated = true;
|
||
}
|
||
}
|
||
Err(_) => break,
|
||
}
|
||
}
|
||
truncated
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn make_cfg(base: PathBuf) -> Config {
|
||
Config {
|
||
host: "127.0.0.1".into(),
|
||
port: 0,
|
||
base_workdir: base,
|
||
max_timeout: Duration::from_millis(2000),
|
||
max_concurrent: 4,
|
||
max_output_bytes: 1024,
|
||
max_file_bytes: 4096,
|
||
allowed_env: vec!["PATH".into()],
|
||
auth_enabled: false,
|
||
auth_secret: String::new(),
|
||
auth_secret_previous: vec![],
|
||
rlimits: crate::rlimits::RlimitsConfig::new(),
|
||
log_dir: PathBuf::from("../logs"),
|
||
isolation_default: "process".into(),
|
||
docker_enabled: false,
|
||
warm_pool: crate::config::WarmPoolConfig::default(),
|
||
}
|
||
}
|
||
|
||
fn sh() -> String {
|
||
if std::path::Path::new("/bin/sh").exists() {
|
||
"/bin/sh".into()
|
||
} else {
|
||
"sh".into()
|
||
}
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn runs_echo_and_captures_stdout() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![sh(), "-c".into(), "echo hello".into()],
|
||
None,
|
||
None,
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_echo".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert!(!out.timeout);
|
||
assert!(!out.cancelled);
|
||
assert_eq!(out.exit_code, Some(0));
|
||
assert_eq!(out.stdout.trim(), "hello");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn timeout_kills_process_group() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![sh(), "-c".into(), "sleep 30".into()],
|
||
None,
|
||
Some(200),
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_timeout".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert!(out.timeout);
|
||
assert!(out.duration_ms < 5000, "应在 200ms 内被强制结束");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn cancel_aborts_running_task() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Arc::new(Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap());
|
||
let trace = "trace_cancel".to_string();
|
||
|
||
let ex2 = ex.clone();
|
||
let trace2 = trace.clone();
|
||
let handle = tokio::spawn(async move {
|
||
ex2.run(
|
||
vec![sh(), "-c".into(), "sleep 30".into()],
|
||
None,
|
||
Some(30_000),
|
||
None,
|
||
vec![],
|
||
None,
|
||
trace2,
|
||
vec![],
|
||
)
|
||
.await
|
||
});
|
||
|
||
let deadline = Instant::now() + Duration::from_secs(2);
|
||
loop {
|
||
if ex.running.lock().await.contains_key(&trace) {
|
||
break;
|
||
}
|
||
if Instant::now() >= deadline {
|
||
panic!("任务未在预期时间内注册");
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||
}
|
||
|
||
assert!(ex.cancel(&trace).await, "应命中运行中任务");
|
||
assert!(!ex.cancel(&trace).await, "重复取消不应命中");
|
||
|
||
let outcome = handle.await.unwrap().unwrap();
|
||
assert!(outcome.cancelled);
|
||
assert!(outcome.duration_ms < 5000, "取消后应快速返回");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn cancel_unknown_trace_returns_false() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
assert!(!ex.cancel("never_existed").await);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn truncates_large_output() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let mut cfg = make_cfg(dir.path().to_path_buf());
|
||
cfg.max_output_bytes = 64;
|
||
let ex = Executor::new(Arc::new(cfg)).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![sh(), "-c".into(), "yes A | head -c 100000".into()],
|
||
None,
|
||
Some(3000),
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_trunc".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert!(out.stdout_truncated);
|
||
assert_eq!(out.stdout.len(), 64);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn rejects_path_traversal_workdir() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let err = ex
|
||
.run(
|
||
vec![sh(), "-c".into(), "echo x".into()],
|
||
None,
|
||
None,
|
||
Some("../../etc".into()),
|
||
vec![],
|
||
None,
|
||
"trace_traversal".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap_err();
|
||
assert!(matches!(err, ExecError::InvalidArgs(_)), "逃逸应为参数错误");
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn passes_stdin_and_env() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![sh(), "-c".into(), "cat; echo $GREETING".into()],
|
||
Some("injected\n".into()),
|
||
None,
|
||
None,
|
||
vec![("GREETING".into(), "hi".into())],
|
||
None,
|
||
"trace_stdin".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(out.exit_code, Some(0));
|
||
assert!(out.stdout.contains("injected"));
|
||
assert!(out.stdout.contains("hi"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn run_stream_emits_stdout_then_done() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let mut rx = ex
|
||
.run_stream(
|
||
vec![sh(), "-c".into(), "echo hello-stream".into()],
|
||
None,
|
||
None,
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_stream".into(),
|
||
None,
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
let mut stdout = String::new();
|
||
let mut done = false;
|
||
while let Some(event) = rx.recv().await {
|
||
match event {
|
||
StreamEvent::Stdout(b) => stdout.push_str(&String::from_utf8_lossy(&b)),
|
||
StreamEvent::Stderr(_) => {}
|
||
StreamEvent::Done(o) => {
|
||
assert_eq!(o.exit_code, Some(0));
|
||
assert!(!o.timeout && !o.cancelled);
|
||
done = true;
|
||
}
|
||
}
|
||
}
|
||
assert!(done, "应以 Done 结束");
|
||
assert!(stdout.contains("hello-stream"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn run_stream_reports_timeout() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let mut rx = ex
|
||
.run_stream(
|
||
vec![sh(), "-c".into(), "sleep 30".into()],
|
||
None,
|
||
Some(200),
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_stream_to".into(),
|
||
None,
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
let mut timeout = false;
|
||
while let Some(event) = rx.recv().await {
|
||
if let StreamEvent::Done(o) = event {
|
||
assert!(o.timeout);
|
||
assert!(o.duration_ms < 5000);
|
||
timeout = true;
|
||
}
|
||
}
|
||
assert!(timeout);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn run_stream_cancels_when_receiver_dropped() {
|
||
// 客户端断开 SSE（drop Receiver）后，后台任务应主动取消子进程并从 running 注销，
|
||
// 而非让命令跑到自然结束或超时——否则并发槽被已离开的用户长任务占满。
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Arc::new(Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap());
|
||
let trace = "trace_stream_drop".to_string();
|
||
|
||
let rx = ex
|
||
.run_stream(
|
||
vec![sh(), "-c".into(), "sleep 30".into()],
|
||
None,
|
||
Some(30_000),
|
||
None,
|
||
vec![],
|
||
None,
|
||
trace.clone(),
|
||
None,
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
|
||
// 等待任务注册到 running map（确认它确实在跑）
|
||
let deadline = Instant::now() + Duration::from_secs(2);
|
||
loop {
|
||
if ex.running.lock().await.contains_key(&trace) {
|
||
break;
|
||
}
|
||
if Instant::now() >= deadline {
|
||
panic!("任务未在预期时间内注册");
|
||
}
|
||
tokio::time::sleep(Duration::from_millis(20)).await;
|
||
}
|
||
|
||
// 模拟客户端断开：drop Receiver
|
||
drop(rx);
|
||
|
||
// 给后台自动取消留出时间（远小于 30s，证明不是自然结束/超时）
|
||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||
|
||
// trace_id 应已从 running 注销 → 手动 cancel 必然返回 false
|
||
assert!(
|
||
!ex.running.lock().await.contains_key(&trace),
|
||
"Receiver drop 后任务应已自动从 running 注销"
|
||
);
|
||
assert!(
|
||
!ex.cancel(&trace).await,
|
||
"Receiver drop 后任务应已被自动取消，手动 cancel 不应再命中"
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn metrics_record_exec_outcome() {
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
ex.run(
|
||
vec![sh(), "-c".into(), "echo ok".into()],
|
||
None,
|
||
None,
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_metrics".into(),
|
||
vec![],
|
||
)
|
||
.await
|
||
.unwrap();
|
||
let text = ex.metrics().gather();
|
||
assert!(text.contains("sandbox_exec_total"), "应有执行计数指标");
|
||
assert!(
|
||
text.contains("outcome=\"completed\""),
|
||
"应记录 completed 结果"
|
||
);
|
||
assert!(text.contains("sandbox_exec_duration_seconds"));
|
||
assert!(text.contains("sandbox_concurrent_limit"));
|
||
}
|
||
|
||
/// 构造一个 FSIZE 限额（写入字节数上限），用于验证 setrlimit 实际生效。
|
||
fn fsize_limit(bytes: u64) -> Vec<EffectiveRlimit> {
|
||
vec![EffectiveRlimit {
|
||
kind: crate::rlimits::RlimitKind::Fsize,
|
||
value: bytes,
|
||
}]
|
||
}
|
||
|
||
#[cfg(unix)]
|
||
#[tokio::test]
|
||
async fn rlimit_fsize_blocks_large_write() {
|
||
// FSIZE=4KB 时，写超过限额的文件应被拦（write 触发 SIGXFSZ 或 EFBIG → 进程非零退出）。
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![
|
||
sh(),
|
||
"-c".into(),
|
||
"dd if=/dev/zero of=big.dat bs=1024 count=64 2>/dev/null".into(),
|
||
],
|
||
None,
|
||
Some(5000),
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_fsize".into(),
|
||
fsize_limit(4096),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
// dd 写 64KB 远超 4KB 限额 → 被中断，exit_code 非 0 或被信号终止
|
||
assert!(
|
||
out.exit_code != Some(0),
|
||
"FSIZE 限额应阻止大文件写入（exit={:?}）",
|
||
out.exit_code
|
||
);
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn rlimit_fsize_allows_small_write() {
|
||
// FSIZE 足够大时，正常写入应成功（验证限额不误杀正常命令）。
|
||
let dir = tempfile::tempdir().unwrap();
|
||
let ex = Executor::new(Arc::new(make_cfg(dir.path().to_path_buf()))).unwrap();
|
||
let out = ex
|
||
.run(
|
||
vec![
|
||
sh(),
|
||
"-c".into(),
|
||
"echo hello > small.dat && cat small.dat".into(),
|
||
],
|
||
None,
|
||
Some(5000),
|
||
None,
|
||
vec![],
|
||
None,
|
||
"trace_fsize_ok".into(),
|
||
fsize_limit(1_048_576),
|
||
)
|
||
.await
|
||
.unwrap();
|
||
assert_eq!(out.exit_code, Some(0));
|
||
assert!(out.stdout.contains("hello"));
|
||
}
|
||
|
||
#[tokio::test]
|
||
async fn internal_error_message_is_sanitized() {
|
||
// Internal 错误的 Display 只含 user_msg，不含 detail（路径/errno）。
|
||
let err = ExecError::Internal {
|
||
user_msg: "无法启动进程",
|
||
detail: "spawn 失败: No such file or directory (os error 2): /secret/path".into(),
|
||
};
|
||
let display = format!("{err}");
|
||
assert_eq!(display, "无法启动进程");
|
||
assert!(
|
||
!display.contains("/secret/path"),
|
||
"user_msg 不应泄露 detail 路径"
|
||
);
|
||
assert!(!display.contains("os error"), "user_msg 不应泄露 errno");
|
||
}
|
||
}
|