222 lines
7.8 KiB
Rust
222 lines
7.8 KiB
Rust
//! 魂镜安全沙箱服务入口。
|
||
//!
|
||
//! 定位见 docs/architecture/service-boundaries.md：Rust 承载的安全与本地可靠能力，
|
||
//! 隶属于「百工 Tool Runtime」的执行隔离层，通过 HTTP 微服务形态供 Python 后端调用。
|
||
//! 跨语言契约与安全边界见 docs/architecture/service-boundaries.md。
|
||
//!
|
||
//! 容器隔离 + 预热池设计见 docs/superpowers/specs/2026-07-05-sandbox-container-warmpool-design.md。
|
||
|
||
mod api;
|
||
mod auth;
|
||
mod codemode;
|
||
mod config;
|
||
mod docker_exec;
|
||
mod error;
|
||
mod executor;
|
||
mod fileio;
|
||
mod metrics;
|
||
mod rlimits;
|
||
mod warm_pool;
|
||
|
||
use std::path::Path;
|
||
use std::sync::Arc;
|
||
|
||
use anyhow::Context;
|
||
use config::Config;
|
||
use executor::Executor;
|
||
use tokio::sync::Semaphore;
|
||
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
|
||
use warm_pool::WarmPool;
|
||
|
||
/// 贯穿 handler 的共享状态。
|
||
pub struct AppState {
|
||
pub config: Arc<Config>,
|
||
pub executor: Arc<Executor>,
|
||
pub semaphore: Arc<Semaphore>,
|
||
pub warm_pool: Option<Arc<WarmPool>>,
|
||
}
|
||
|
||
#[tokio::main]
|
||
async fn main() -> anyhow::Result<()> {
|
||
// 本地 cargo run 时从项目根 .env 注入环境变量；容器内 .dockerignore 已排除 .env，
|
||
// dotenv_iter() 找不到文件返回 Err，if let Ok 静默跳过，直接用 compose environment 注入的值。
|
||
if let Ok(iter) = dotenvy::dotenv_iter() {
|
||
for item in iter.flatten() {
|
||
let (k, v) = item;
|
||
if std::env::var_os(&k).is_none() {
|
||
std::env::set_var(&k, v);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 先读配置：此时 tracing 尚未初始化，失败经 `?` 返回由 Rust runtime 打到 stderr。
|
||
let config = Arc::new(Config::from_env().context("读取沙箱配置失败")?);
|
||
|
||
// 初始化 tracing：stdout（彩色，控制台实时可见）+ 文件（sandbox.log 按日轮转，持久化）。
|
||
if let Err(e) = std::fs::create_dir_all(&config.log_dir) {
|
||
eprintln!("警告：创建日志目录失败 {:?}：{e}（将仅输出到控制台）", config.log_dir);
|
||
}
|
||
init_logging(&config.log_dir);
|
||
|
||
tracing::info!(
|
||
host = %config.host,
|
||
port = config.port,
|
||
base = ?config.base_workdir,
|
||
log_dir = ?config.log_dir,
|
||
max_timeout_ms = config.max_timeout.as_millis(),
|
||
max_concurrent = config.max_concurrent,
|
||
auth_enabled = config.auth_enabled,
|
||
docker_enabled = config.docker_enabled,
|
||
"启动魂镜安全沙箱服务",
|
||
);
|
||
|
||
tokio::fs::create_dir_all(&config.base_workdir)
|
||
.await
|
||
.with_context(|| format!("创建沙箱工作目录失败: {:?}", config.base_workdir))?;
|
||
|
||
// 清理上次异常退出残留的临时工作目录（tempfile::TempDir 命名为 .tmpXXXXXX）
|
||
cleanup_stale_tempdirs(&config.base_workdir).await;
|
||
|
||
// 初始化预热池（Docker 不可用时降级为 None）
|
||
let warm_pool: Option<Arc<WarmPool>> = if config.docker_enabled {
|
||
// 解析工作目录为绝对路径（Docker bind mount 要求绝对路径）
|
||
let workdir_abs = config
|
||
.base_workdir
|
||
.canonicalize()
|
||
.unwrap_or_else(|_| config.base_workdir.clone())
|
||
.to_str()
|
||
.unwrap_or("/app/sandbox-work")
|
||
.to_string();
|
||
|
||
match WarmPool::try_new(config.warm_pool.clone(), workdir_abs)
|
||
.await
|
||
{
|
||
Some(pool) => {
|
||
// 后台预热（不阻塞启动）
|
||
let pool_clone = pool.clone();
|
||
tokio::spawn(async move {
|
||
pool_clone.warmup_all().await;
|
||
});
|
||
|
||
// 后台健康检查
|
||
let pool_clone = pool.clone();
|
||
tokio::spawn(async move {
|
||
pool_clone.health_check_loop().await;
|
||
});
|
||
|
||
// 后台空闲回收
|
||
let pool_clone = pool.clone();
|
||
tokio::spawn(async move {
|
||
pool_clone.idle_eviction_loop().await;
|
||
});
|
||
|
||
tracing::info!("预热池已初始化，容器模式已启用");
|
||
Some(pool)
|
||
}
|
||
None => {
|
||
tracing::warn!("Docker 不可用，容器模式降级为进程模式");
|
||
None
|
||
}
|
||
}
|
||
} else {
|
||
tracing::info!("容器模式未启用（SANDBOX_DOCKER_ENABLED=false），仅使用进程隔离");
|
||
None
|
||
};
|
||
|
||
let warm_pool_for_shutdown = warm_pool.clone();
|
||
let state = Arc::new(AppState {
|
||
warm_pool,
|
||
config: config.clone(),
|
||
executor: Arc::new(Executor::new(config.clone())?),
|
||
semaphore: Arc::new(Semaphore::new(config.max_concurrent)),
|
||
});
|
||
|
||
let app = api::router(state);
|
||
let addr = format!("{}:{}", config.host, config.port);
|
||
let listener = tokio::net::TcpListener::bind(&addr)
|
||
.await
|
||
.with_context(|| format!("监听 {addr} 失败"))?;
|
||
tracing::info!("安全沙箱服务监听 {addr}（GET /healthz，POST /v1/exec）");
|
||
|
||
axum::serve(listener, app)
|
||
.with_graceful_shutdown(shutdown_signal(warm_pool_for_shutdown))
|
||
.await
|
||
.context("axum 服务退出异常")?;
|
||
|
||
tracing::info!("安全沙箱服务已停止");
|
||
Ok(())
|
||
}
|
||
|
||
/// 装配 tracing subscriber：同一 EnvFilter 同时扇出到 stdout（彩色）与日志文件（无颜色）。
|
||
/// 文件按日轮转写入 `<log_dir>/sandbox.log`（当天无后缀，历史文件追加日期后缀）。
|
||
fn init_logging(log_dir: &Path) {
|
||
let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));
|
||
let file_writer = tracing_appender::rolling::daily(log_dir, "sandbox.log");
|
||
tracing_subscriber::registry()
|
||
.with(filter)
|
||
.with(fmt::layer().with_writer(std::io::stdout))
|
||
.with(fmt::layer().with_writer(file_writer).with_ansi(false))
|
||
.init();
|
||
}
|
||
|
||
/// 清理 base 下 tempfile 残留的临时目录（`.tmpXXXXXX`），防御异常退出后的目录泄漏。
|
||
async fn cleanup_stale_tempdirs(base: &Path) {
|
||
let mut entries = match tokio::fs::read_dir(base).await {
|
||
Ok(e) => e,
|
||
Err(e) => {
|
||
tracing::warn!(error = %e, "读取沙箱工作目录失败，跳过残留清理");
|
||
return;
|
||
}
|
||
};
|
||
let mut cleaned = 0u32;
|
||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||
let name = entry.file_name();
|
||
if name.to_string_lossy().starts_with(".tmp") {
|
||
let path = entry.path();
|
||
match tokio::fs::remove_dir_all(&path).await {
|
||
Ok(_) => cleaned += 1,
|
||
Err(e) => tracing::warn!(path = %path.display(), error = %e, "清理残留临时目录失败"),
|
||
}
|
||
}
|
||
}
|
||
if cleaned > 0 {
|
||
tracing::info!(cleaned, "已清理上次运行残留的临时工作目录");
|
||
}
|
||
}
|
||
|
||
/// 优雅关闭：捕获 Ctrl-C（开发）与 SIGTERM（容器编排），回收预热池容器。
|
||
async fn shutdown_signal(warm_pool: Option<Arc<WarmPool>>) {
|
||
let ctrl_c = async {
|
||
let _ = tokio::signal::ctrl_c().await;
|
||
};
|
||
|
||
#[cfg(unix)]
|
||
let terminate = async {
|
||
match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
|
||
Ok(mut s) => {
|
||
s.recv().await;
|
||
}
|
||
Err(_) => std::future::pending::<()>().await,
|
||
}
|
||
};
|
||
#[cfg(not(unix))]
|
||
let terminate = std::future::pending::<()>();
|
||
|
||
tokio::select! {
|
||
_ = ctrl_c => {
|
||
tracing::info!("收到中断信号（Ctrl-C），开始优雅关闭");
|
||
if let Some(pool) = warm_pool {
|
||
pool.shutdown().await;
|
||
tracing::info!("预热池容器已全部回收");
|
||
}
|
||
}
|
||
_ = terminate => {
|
||
tracing::info!("收到终止信号（SIGTERM），开始优雅关闭");
|
||
if let Some(pool) = warm_pool {
|
||
pool.shutdown().await;
|
||
tracing::info!("预热池容器已全部回收");
|
||
}
|
||
}
|
||
}
|
||
}
|