初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,647 @@
|
||||
//! HTTP 路由、请求 / 响应模型与 handler。
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::extract::{Extension, State};
|
||||
use axum::http::{header, HeaderValue};
|
||||
use axum::response::IntoResponse;
|
||||
use axum::response::sse::{Event, KeepAlive, Sse};
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine as _};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio_stream::StreamExt;
|
||||
use tokio_stream::wrappers::ReceiverStream;
|
||||
|
||||
use crate::auth::AuthPayload;
|
||||
use crate::error::ApiError;
|
||||
use crate::executor::{ExecError, ExecOutcome, StreamEvent};
|
||||
use crate::AppState;
|
||||
|
||||
/// 装配路由：GET /healthz、GET /metrics、POST /v1/exec、POST /v1/exec/stream、POST /v1/cancel。
|
||||
/// 受保护端点经司命授权中间件；`/healthz` 与 `/metrics` 不受保护。
|
||||
pub fn router(state: Arc<AppState>) -> Router {
|
||||
let protected = Router::new()
|
||||
.route("/v1/exec", post(exec))
|
||||
.route("/v1/exec/stream", post(exec_stream))
|
||||
.route("/v1/cancel", post(cancel_exec))
|
||||
.route("/v1/file/read", post(file_read))
|
||||
.route("/v1/file/write", post(file_write))
|
||||
.route("/v1/file/list", post(file_list))
|
||||
.route("/v1/file/stat", post(file_stat))
|
||||
.route_layer(axum::middleware::from_fn_with_state(
|
||||
state.clone(),
|
||||
crate::auth::require_authorization,
|
||||
));
|
||||
Router::new()
|
||||
.route("/healthz", get(healthz))
|
||||
.route("/metrics", get(metrics_handler))
|
||||
.merge(protected)
|
||||
.layer(tower_http::trace::TraceLayer::new_for_http())
|
||||
.layer(tower_http::limit::RequestBodyLimitLayer::new(2 * 1024 * 1024))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
pub async fn healthz(State(state): State<Arc<AppState>>) -> impl IntoResponse {
|
||||
let concurrent = state.executor.metrics().concurrent_value();
|
||||
let max_concurrent = state.config.max_concurrent;
|
||||
|
||||
// 收集预热池统计（如果有）
|
||||
let pool_stats = if let Some(pool) = &state.warm_pool {
|
||||
let metrics = pool.collect_metrics().await;
|
||||
let total: usize = metrics.values().map(|(idle, running, _)| (idle + running) as usize).sum();
|
||||
Some(serde_json::json!({
|
||||
"total_containers": total,
|
||||
"by_language": metrics.iter().map(|(lang, (idle, running, waiters))| {
|
||||
(lang.clone(), serde_json::json!({
|
||||
"idle": idle,
|
||||
"running": running,
|
||||
"waiters": waiters,
|
||||
}))
|
||||
}).collect::<serde_json::Map<_, _>>(),
|
||||
}))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
Json(serde_json::json!({
|
||||
"status": "ok",
|
||||
"service": "hunjing-sandbox",
|
||||
"version": env!("CARGO_PKG_VERSION"),
|
||||
"load": {
|
||||
"concurrent": concurrent,
|
||||
"max_concurrent": max_concurrent,
|
||||
"utilization_pct": if max_concurrent > 0 {
|
||||
(concurrent as f64 / max_concurrent as f64 * 100.0).round()
|
||||
} else { 0.0 },
|
||||
},
|
||||
"config": {
|
||||
"auth_enabled": state.config.auth_enabled,
|
||||
"docker_enabled": state.config.docker_enabled,
|
||||
"default_isolation": &state.config.isolation_default,
|
||||
},
|
||||
"warm_pool": pool_stats,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Prometheus 指标（text/plain, version=0.0.4）。不鉴权，便于抓取。
|
||||
pub async fn metrics_handler(State(state): State<Arc<AppState>>) -> impl IntoResponse {
|
||||
let body = state.executor.metrics().gather();
|
||||
(
|
||||
[(
|
||||
header::CONTENT_TYPE,
|
||||
HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"),
|
||||
)],
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
/// 执行请求体。`argv`（命令式）与 `code`+`language`（代码片段式）二选一，同时或都不给出返回 400。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ExecRequest {
|
||||
/// 命令参数数组，首个元素为可执行文件（命令模式）。
|
||||
#[serde(default)]
|
||||
pub argv: Option<Vec<String>>,
|
||||
/// 代码片段（code 模式）；与 `argv` 互斥。沙箱转译为 argv 后走执行路径。
|
||||
#[serde(default)]
|
||||
pub code: Option<String>,
|
||||
/// code 模式语言，默认 `python`。P0 仅支持 python。
|
||||
#[serde(default)]
|
||||
pub language: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stdin: Option<String>,
|
||||
#[serde(default)]
|
||||
pub timeout_ms: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub workdir: Option<String>,
|
||||
#[serde(default)]
|
||||
pub env: HashMap<String, String>,
|
||||
#[serde(default)]
|
||||
pub max_output_bytes: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub trace_id: Option<String>,
|
||||
/// 资源限额收紧（key: fsize/as/cpu/nproc/nofile → u64）；只能比凭证授予/默认更严，不能放宽。
|
||||
#[serde(default)]
|
||||
pub rlimits: Option<BTreeMap<String, u64>>,
|
||||
/// 隔离策略：soft（默认，容器失败时降级 process）或 hard（容器失败直接报错，不降级）。
|
||||
#[serde(default)]
|
||||
pub isolation_policy: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct ExecResponse {
|
||||
pub ok: bool,
|
||||
pub exit_code: Option<i32>,
|
||||
pub stdout: String,
|
||||
pub stderr: String,
|
||||
pub duration_ms: u128,
|
||||
pub timeout: bool,
|
||||
pub cancelled: bool,
|
||||
pub truncated: bool,
|
||||
pub trace_id: String,
|
||||
/// 实际生效的隔离模式（process/container）。P0 期间 container 降级为 process。
|
||||
pub isolation_mode: String,
|
||||
}
|
||||
|
||||
/// 取消请求体：按 trace_id 定位运行中任务。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CancelRequest {
|
||||
pub trace_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct CancelResponse {
|
||||
pub cancelled: bool,
|
||||
pub trace_id: String,
|
||||
}
|
||||
|
||||
/// 文件读取请求。`path` 为 base 内相对路径；`offset`/`limit` 字节级分页（可选）。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct FileReadRequest {
|
||||
pub path: String,
|
||||
#[serde(default)]
|
||||
pub offset: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub limit: Option<u64>,
|
||||
}
|
||||
|
||||
/// 文件写入请求。`content_b64` 为 base64url-no-pad 编码（二进制安全，不经 shell）。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct FileWriteRequest {
|
||||
pub path: String,
|
||||
pub content_b64: String,
|
||||
}
|
||||
|
||||
/// 目录列出请求。`path` 默认 base 根（`.`）。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct FileListRequest {
|
||||
#[serde(default = "default_list_path")]
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
fn default_list_path() -> String {
|
||||
".".into()
|
||||
}
|
||||
|
||||
/// 文件元信息请求。
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct FileStatRequest {
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
pub async fn exec(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<ExecRequest>,
|
||||
) -> Result<Json<ExecResponse>, ApiError> {
|
||||
validate_mutex(&req)?;
|
||||
let isolation = resolve_isolation(&state, &auth, req.isolation_policy.as_deref())?;
|
||||
let effective_mode = isolation.mode.clone();
|
||||
let allow_fallback = isolation.allow_fallback;
|
||||
|
||||
// 授权绑定校验：body_hash 按 code/argv 模式分别重算，必须与凭证一致（防 TTL 内重放）。
|
||||
if let Some(expected) = &auth.body_hash {
|
||||
let actual = compute_request_body_hash(&req);
|
||||
if &actual != expected {
|
||||
return Err(ApiError::Unauthorized(
|
||||
"请求内容与授权凭证不匹配（body_hash 校验失败）".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// code → argv 转译。translation 持有 TempPath，在本函数作用域内存活至 run().await 结束
|
||||
// （同步执行：run 返回时子进程已 open 脚本，drop 安全）。
|
||||
let translation = translate_if_code(&req).await?;
|
||||
let effective_argv: Vec<String> = match &translation {
|
||||
Some(t) => t.argv.clone(),
|
||||
None => req.argv.clone().unwrap(),
|
||||
};
|
||||
if effective_argv[0].trim().is_empty() {
|
||||
return Err(ApiError::BadRequest("argv[0] 不能为空".into()));
|
||||
}
|
||||
|
||||
let trace_id = req.trace_id.clone().unwrap_or_else(generate_trace_id);
|
||||
|
||||
// 资源限额：凭证授予（司命 tool policy）+ 请求收紧，clamp 到全局天花板。
|
||||
let effective = crate::rlimits::compute_effective(
|
||||
&state.config.rlimits,
|
||||
auth.rlimit_overrides.as_ref(),
|
||||
req.rlimits.as_ref(),
|
||||
);
|
||||
|
||||
// 并发闸：超出上限时排队，避免资源耗尽。
|
||||
let _permit = state
|
||||
.semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| ApiError::Unavailable("沙箱并发闸已关闭".into()))?;
|
||||
|
||||
let env_pairs: Vec<(String, String)> = req.env.into_iter().collect();
|
||||
|
||||
let outcome: ExecOutcome = state
|
||||
.executor
|
||||
.run_routed(
|
||||
effective_argv,
|
||||
req.stdin,
|
||||
req.timeout_ms,
|
||||
req.workdir,
|
||||
env_pairs,
|
||||
req.max_output_bytes,
|
||||
trace_id.clone(),
|
||||
effective,
|
||||
&effective_mode,
|
||||
state.warm_pool.clone(),
|
||||
allow_fallback,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
ExecError::InvalidArgs(m) => ApiError::BadRequest(m),
|
||||
ExecError::Internal { user_msg, detail } => {
|
||||
tracing::error!(error = %detail, "沙箱执行内部错误");
|
||||
ApiError::Internal(user_msg.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
Ok(Json(ExecResponse {
|
||||
ok: !outcome.timeout && !outcome.cancelled && outcome.exit_code == Some(0),
|
||||
exit_code: outcome.exit_code,
|
||||
stdout: outcome.stdout,
|
||||
stderr: outcome.stderr,
|
||||
duration_ms: outcome.duration_ms,
|
||||
timeout: outcome.timeout,
|
||||
cancelled: outcome.cancelled,
|
||||
truncated: outcome.stdout_truncated || outcome.stderr_truncated,
|
||||
trace_id,
|
||||
isolation_mode: effective_mode,
|
||||
}))
|
||||
}
|
||||
|
||||
/// 流式执行：返回 SSE，实时推送 stdout / stderr 增量（base64 保真），末尾一个 done 事件。
|
||||
pub async fn exec_stream(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<ExecRequest>,
|
||||
) -> Result<Sse<impl tokio_stream::Stream<Item = Result<Event, axum::Error>>>, ApiError> {
|
||||
validate_mutex(&req)?;
|
||||
let isolation = resolve_isolation(&state, &auth, req.isolation_policy.as_deref())?;
|
||||
let effective_mode = isolation.mode.clone();
|
||||
let allow_fallback = isolation.allow_fallback;
|
||||
if let Some(expected) = &auth.body_hash {
|
||||
let actual = compute_request_body_hash(&req);
|
||||
if &actual != expected {
|
||||
return Err(ApiError::Unauthorized(
|
||||
"请求内容与授权凭证不匹配（body_hash 校验失败）".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
let translation = translate_if_code(&req).await?;
|
||||
let effective_argv: Vec<String> = match &translation {
|
||||
Some(t) => t.argv.clone(),
|
||||
None => req.argv.clone().unwrap(),
|
||||
};
|
||||
if effective_argv[0].trim().is_empty() {
|
||||
return Err(ApiError::BadRequest("argv[0] 不能为空".into()));
|
||||
}
|
||||
|
||||
let trace_id = req.trace_id.clone().unwrap_or_else(generate_trace_id);
|
||||
|
||||
// 资源限额：凭证授予（司命 tool policy）+ 请求收紧，clamp 到全局天花板。
|
||||
let effective = crate::rlimits::compute_effective(
|
||||
&state.config.rlimits,
|
||||
auth.rlimit_overrides.as_ref(),
|
||||
req.rlimits.as_ref(),
|
||||
);
|
||||
|
||||
let permit = state
|
||||
.semaphore
|
||||
.clone()
|
||||
.acquire_owned()
|
||||
.await
|
||||
.map_err(|_| ApiError::Unavailable("沙箱并发闸已关闭".into()))?;
|
||||
|
||||
let env_pairs: Vec<(String, String)> = req.env.into_iter().collect();
|
||||
// Clone effective_mode before moving into SSE closure; run_stream_routed borrows the original.
|
||||
let mode_for_sse = effective_mode.clone();
|
||||
let rx = state
|
||||
.executor
|
||||
.run_stream_routed(
|
||||
effective_argv,
|
||||
req.stdin,
|
||||
req.timeout_ms,
|
||||
req.workdir,
|
||||
env_pairs,
|
||||
req.max_output_bytes,
|
||||
trace_id.clone(),
|
||||
effective,
|
||||
&effective_mode,
|
||||
state.warm_pool.clone(),
|
||||
Some(permit),
|
||||
allow_fallback,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
ExecError::InvalidArgs(m) => ApiError::BadRequest(m),
|
||||
ExecError::Internal { user_msg, detail } => {
|
||||
tracing::error!(error = %detail, "沙箱执行内部错误");
|
||||
ApiError::Internal(user_msg.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
let stream = ReceiverStream::new(rx).map(move |event| {
|
||||
let _hold_script = &translation; // 持有 TempPath 至流结束
|
||||
let ev = match event {
|
||||
StreamEvent::Stdout(b) => Event::default()
|
||||
.event("stdout")
|
||||
.data(BASE64_URL_SAFE_NO_PAD.encode(&b)),
|
||||
StreamEvent::Stderr(b) => Event::default()
|
||||
.event("stderr")
|
||||
.data(BASE64_URL_SAFE_NO_PAD.encode(&b)),
|
||||
StreamEvent::Done(o) => {
|
||||
let payload = serde_json::json!({
|
||||
"exit_code": o.exit_code,
|
||||
"duration_ms": o.duration_ms,
|
||||
"timeout": o.timeout,
|
||||
"cancelled": o.cancelled,
|
||||
"truncated": o.stdout_truncated || o.stderr_truncated,
|
||||
"trace_id": trace_id,
|
||||
"isolation_mode": mode_for_sse,
|
||||
});
|
||||
Event::default().event("done").data(payload.to_string())
|
||||
}
|
||||
};
|
||||
Ok::<_, axum::Error>(ev)
|
||||
});
|
||||
|
||||
Ok(Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))))
|
||||
}
|
||||
|
||||
/// 取消运行中的任务（按 trace_id）。
|
||||
///
|
||||
/// 安全权衡：cancel 仅中止任务、不执行新命令，危害有限，故凭证不绑定 trace_id
|
||||
/// （验签 + 过期即可），与 exec 的 body_hash 强绑定不同。
|
||||
pub async fn cancel_exec(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(_auth): Extension<AuthPayload>,
|
||||
Json(req): Json<CancelRequest>,
|
||||
) -> Result<Json<CancelResponse>, ApiError> {
|
||||
if req.trace_id.trim().is_empty() {
|
||||
return Err(ApiError::BadRequest("trace_id 不能为空".into()));
|
||||
}
|
||||
let cancelled = state.executor.cancel(&req.trace_id).await;
|
||||
if cancelled {
|
||||
Ok(Json(CancelResponse {
|
||||
cancelled: true,
|
||||
trace_id: req.trace_id,
|
||||
}))
|
||||
} else {
|
||||
Err(ApiError::NotFound(format!(
|
||||
"任务不存在或已完成: {}",
|
||||
req.trace_id
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
// ==============================
|
||||
// 文件 IO 端点（不经 shell，二进制安全）
|
||||
// ==============================
|
||||
|
||||
/// `POST /v1/file/read`：读取 base 内文件，返回 base64 内容 + 元信息。
|
||||
pub async fn file_read(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<FileReadRequest>,
|
||||
) -> Result<Json<crate::fileio::FileReadOutcome>, ApiError> {
|
||||
verify_file_op_hash(&auth, "read", &req.path)?;
|
||||
let outcome = crate::fileio::file_read(
|
||||
state.executor.base(),
|
||||
state.executor.base_canon(),
|
||||
&req.path,
|
||||
req.offset,
|
||||
req.limit,
|
||||
state.config.max_file_bytes,
|
||||
)
|
||||
.await
|
||||
.map_err(file_error_to_api)?;
|
||||
Ok(Json(outcome))
|
||||
}
|
||||
|
||||
/// `POST /v1/file/write`：写入 base 内文件（覆盖），content 为 base64url-no-pad。
|
||||
pub async fn file_write(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<FileWriteRequest>,
|
||||
) -> Result<Json<crate::fileio::FileWriteOutcome>, ApiError> {
|
||||
verify_file_op_hash(&auth, "write", &req.path)?;
|
||||
let content = BASE64_URL_SAFE_NO_PAD
|
||||
.decode(req.content_b64.as_bytes())
|
||||
.map_err(|_| ApiError::BadRequest("content_b64 不是合法 base64url-no-pad".into()))?;
|
||||
let outcome = crate::fileio::file_write(
|
||||
state.executor.base(),
|
||||
state.executor.base_canon(),
|
||||
&req.path,
|
||||
&content,
|
||||
state.config.max_file_bytes,
|
||||
)
|
||||
.await
|
||||
.map_err(file_error_to_api)?;
|
||||
Ok(Json(outcome))
|
||||
}
|
||||
|
||||
/// `POST /v1/file/list`：列出 base 内目录条目（按名字排序）。
|
||||
pub async fn file_list(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<FileListRequest>,
|
||||
) -> Result<Json<crate::fileio::FileListOutcome>, ApiError> {
|
||||
verify_file_op_hash(&auth, "list", &req.path)?;
|
||||
let outcome = crate::fileio::file_list(
|
||||
state.executor.base(),
|
||||
state.executor.base_canon(),
|
||||
&req.path,
|
||||
)
|
||||
.await
|
||||
.map_err(file_error_to_api)?;
|
||||
Ok(Json(outcome))
|
||||
}
|
||||
|
||||
/// `POST /v1/file/stat`：返回 base 内文件/目录元信息。
|
||||
pub async fn file_stat(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Extension(auth): Extension<AuthPayload>,
|
||||
Json(req): Json<FileStatRequest>,
|
||||
) -> Result<Json<crate::fileio::FileStatOutcome>, ApiError> {
|
||||
verify_file_op_hash(&auth, "stat", &req.path)?;
|
||||
let outcome = crate::fileio::file_stat(
|
||||
state.executor.base(),
|
||||
state.executor.base_canon(),
|
||||
&req.path,
|
||||
)
|
||||
.await
|
||||
.map_err(file_error_to_api)?;
|
||||
Ok(Json(outcome))
|
||||
}
|
||||
|
||||
/// 校验文件操作凭证绑定：凭证 file_op_hash 必须与 {op, path} 摘要一致，防重放任意文件操作。
|
||||
fn verify_file_op_hash(auth: &AuthPayload, op: &str, path: &str) -> Result<(), ApiError> {
|
||||
if let Some(expected) = &auth.file_op_hash {
|
||||
let actual = crate::auth::compute_file_op_hash(op, path);
|
||||
if &actual != expected {
|
||||
return Err(ApiError::Unauthorized(
|
||||
"文件操作与授权凭证不匹配（file_op_hash 校验失败）".into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 文件操作错误 → API 错误映射（路径问题 4xx，IO 内部错误脱敏后 500）。
|
||||
fn file_error_to_api(e: crate::fileio::FileError) -> ApiError {
|
||||
use crate::fileio::FileError;
|
||||
match e {
|
||||
FileError::InvalidPath(m) => ApiError::BadRequest(m),
|
||||
FileError::NotFound => ApiError::NotFound("文件或目录不存在".into()),
|
||||
FileError::WrongType { expected, actual } => {
|
||||
ApiError::BadRequest(format!("类型不符：期望 {expected}，实际 {actual}"))
|
||||
}
|
||||
FileError::TooLarge { limit, actual } => {
|
||||
ApiError::BadRequest(format!("内容超过大小上限 {limit} 字节（实际 {actual}）"))
|
||||
}
|
||||
FileError::Io(e) => {
|
||||
tracing::error!(error = %e, "文件 IO 内部错误");
|
||||
ApiError::Internal("文件操作失败".into())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn generate_trace_id() -> String {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
format!("sbx_{nanos:x}")
|
||||
}
|
||||
|
||||
/// argv / code 互斥：同时或都不给出 → 400。
|
||||
fn validate_mutex(req: &ExecRequest) -> Result<(), ApiError> {
|
||||
match (req.argv.as_ref(), req.code.as_ref()) {
|
||||
(Some(_), Some(_)) => Err(ApiError::BadRequest("argv 与 code 不能同时指定".into())),
|
||||
(None, None) => Err(ApiError::BadRequest("必须指定 argv 或 code 之一".into())),
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// 计算请求内容摘要：code 模式按 {language, code, stdin}，argv 模式按 {argv, stdin}。
|
||||
fn compute_request_body_hash(req: &ExecRequest) -> String {
|
||||
if let Some(code) = req.code.as_deref() {
|
||||
let lang = req.language.as_deref().unwrap_or("python");
|
||||
crate::auth::compute_body_hash_code(lang, code, &req.stdin)
|
||||
} else {
|
||||
let argv = req.argv.as_deref().unwrap_or_default();
|
||||
crate::auth::compute_body_hash(argv, &req.stdin)
|
||||
}
|
||||
}
|
||||
|
||||
/// code 模式转译为 argv；argv 模式返回 None。返回的 `CodeTranslation` 持有 `TempPath`，
|
||||
/// 调用方负责存活至执行结束（同步 exec 由作用域持有；流式 exec_stream 由 SSE 闭包持有）。
|
||||
async fn translate_if_code(
|
||||
req: &ExecRequest,
|
||||
) -> Result<Option<crate::codemode::CodeTranslation>, ApiError> {
|
||||
if let Some(code) = req.code.as_deref() {
|
||||
let lang = req.language.as_deref().unwrap_or("python");
|
||||
crate::codemode::translate(lang, code)
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(ApiError::BadRequest)
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
/// 生效隔离模式与降级策略。
|
||||
struct IsolationConfig {
|
||||
/// 实际执行模式："process" 或 "container"。
|
||||
mode: String,
|
||||
/// 容器执行失败时是否允许降级为进程模式。
|
||||
allow_fallback: bool,
|
||||
}
|
||||
|
||||
/// 解析生效隔离模式与降级策略。
|
||||
///
|
||||
/// `auth.isolation_mode` 为凭证授予的隔离要求；`policy` 来自请求体 `isolation_policy`
|
||||
/// （"hard"=不降级 / "soft"=降级 / None=默认降级）。
|
||||
fn resolve_isolation(
|
||||
state: &Arc<AppState>,
|
||||
auth: &AuthPayload,
|
||||
policy: Option<&str>,
|
||||
) -> Result<IsolationConfig, ApiError> {
|
||||
let hard = policy == Some("hard");
|
||||
let fallback = |mode: &str, downgrade: bool| {
|
||||
let m = mode.to_string();
|
||||
IsolationConfig {
|
||||
mode: m,
|
||||
allow_fallback: if downgrade { true } else { !hard },
|
||||
}
|
||||
};
|
||||
|
||||
match auth.isolation_mode.as_deref() {
|
||||
Some("container") => {
|
||||
if state.warm_pool.is_some() {
|
||||
Ok(fallback("container", false))
|
||||
} else if hard {
|
||||
Err(ApiError::Unavailable(
|
||||
"凭证要求 container 隔离但预热池未就绪，且 isolation_policy=hard 拒绝降级".into(),
|
||||
))
|
||||
} else {
|
||||
tracing::warn!(
|
||||
tool = ?auth.tool,
|
||||
"凭证要求 container 隔离但预热池不可用，降级 process"
|
||||
);
|
||||
Ok(fallback("process", true))
|
||||
}
|
||||
}
|
||||
Some("process") => Ok(fallback("process", true)),
|
||||
Some(other) => {
|
||||
tracing::warn!(mode = other, "未知 isolation_mode，降级 process");
|
||||
Ok(fallback("process", true))
|
||||
}
|
||||
None => {
|
||||
if state.config.isolation_default == "container" && state.warm_pool.is_some() {
|
||||
Ok(fallback("container", false))
|
||||
} else {
|
||||
Ok(fallback("process", true))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn exec_request_code_mode_deserializes() {
|
||||
let json = r#"{"language":"python","code":"print(1)"}"#;
|
||||
let req: ExecRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.code.as_deref(), Some("print(1)"));
|
||||
assert_eq!(req.language.as_deref(), Some("python"));
|
||||
assert!(req.argv.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exec_request_argv_mode_still_works() {
|
||||
let json = r#"{"argv":["python3","-c","print(1)"]}"#;
|
||||
let req: ExecRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(
|
||||
req.argv.as_deref().unwrap(),
|
||||
&[
|
||||
"python3".to_string(),
|
||||
"-c".to_string(),
|
||||
"print(1)".to_string()
|
||||
]
|
||||
);
|
||||
assert!(req.code.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,460 @@
|
||||
//! 司命授权：HMAC-SHA256 签名的短期凭证。
|
||||
//!
|
||||
//! 凭证格式：`<payload_b64url>.<sig_b64url>`
|
||||
//! - payload: JSON `{exp, trace_id, risk_level, tool, character_id, run_id, body_hash, rlimit_overrides}`，
|
||||
//! 字段顺序与后端 Python 客户端一致。
|
||||
//! - sig: HMAC-SHA256(secret, payload_b64url) 的 base64url（无填充）。
|
||||
//!
|
||||
//! 安全语义：沙箱不做风险判定，只校验"调用方是否持有由司命签发的有效凭证"。
|
||||
//! `body_hash` 把凭证绑死到具体执行内容（argv + stdin 摘要），防 TTL 内重放执行任意命令；
|
||||
//! `rlimit_overrides` 携带司命授予的资源配额上限（授权属性，不进 body_hash），验签后信任并 clamp。
|
||||
//! verify 支持多 secret 并存（主 + previous），支持无中断轮换。
|
||||
//! 验签失败、过期或 body_hash 不匹配一律 401。
|
||||
//! 设计依据见 docs/architecture/service-boundaries.md 与 cross-language-integration-design.md。
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use axum::extract::{Request, State};
|
||||
use axum::http::{header, StatusCode};
|
||||
use axum::middleware::Next;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::Json;
|
||||
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine as _};
|
||||
use hmac::{Hmac, Mac};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::AppState;
|
||||
|
||||
type HmacSha256 = Hmac<Sha256>;
|
||||
|
||||
/// 授权凭证 payload（验签通过后插入请求扩展，供 handler 关联 trace、body_hash 校验与资源配额）。
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AuthPayload {
|
||||
/// 过期时间（Unix 秒）。
|
||||
pub exp: u64,
|
||||
#[serde(default)]
|
||||
pub trace_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub risk_level: Option<String>,
|
||||
#[serde(default)]
|
||||
pub tool: Option<String>,
|
||||
#[serde(default)]
|
||||
pub character_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub run_id: Option<String>,
|
||||
/// 请求内容摘要 SHA256(canonical JSON `{argv, stdin}`)，绑定凭证与具体执行内容，
|
||||
/// 防止凭证在 TTL 内被重放去执行任意命令；None 表示不绑定（仅 auth_enabled=false 本地联调）。
|
||||
#[serde(default)]
|
||||
pub body_hash: Option<String>,
|
||||
/// 司命授予该上下文的资源配额上限（key 见 `RlimitKind::as_str`：fsize/as/cpu/nproc/nofile）。
|
||||
/// 随 HMAC 凭证下发，sandbox 验签后信任并 clamp 到全局天花板；属授权属性，不进 body_hash。
|
||||
#[serde(default)]
|
||||
pub rlimit_overrides: Option<BTreeMap<String, u64>>,
|
||||
/// 文件操作摘要 SHA256(canonical JSON `{op, path}`)，绑定凭证与具体文件操作
|
||||
/// （read/write/list/stat），防 TTL 内重放执行任意文件操作。exec 凭证为 None、file 凭证
|
||||
/// 为 Some；与 body_hash 互斥使用，不参与 body_hash 计算。
|
||||
#[serde(default)]
|
||||
pub file_op_hash: Option<String>,
|
||||
/// 隔离模式授权（L0-L2 → process，L3-L4 → container）。随 HMAC 凭证下发，sandbox 验签后
|
||||
/// 据此路由执行路径；P0 期间 container 降级 process。属授权属性，不进 body_hash，与
|
||||
/// rlimit_overrides 同链路。None 表示司命未指定，按 SANDBOX_ISOLATION_DEFAULT 兜底。
|
||||
#[serde(default)]
|
||||
pub isolation_mode: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum AuthError {
|
||||
Missing,
|
||||
Malformed,
|
||||
BadSignature,
|
||||
Expired,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for AuthError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let s = match self {
|
||||
AuthError::Missing => "缺少 Authorization 头",
|
||||
AuthError::Malformed => "授权凭证格式不合法",
|
||||
AuthError::BadSignature => "授权凭证签名校验失败",
|
||||
AuthError::Expired => "授权凭证已过期",
|
||||
};
|
||||
f.write_str(s)
|
||||
}
|
||||
}
|
||||
|
||||
/// 签发授权凭证（与后端 Python `sign_authorization` 共用同一算法与密钥）。
|
||||
#[allow(dead_code)] // 沙箱端仅校验（verify）；sign 供对称测试与未来本地签发场景保留
|
||||
pub fn sign(secret: &str, payload: &AuthPayload) -> anyhow::Result<String> {
|
||||
let json = serde_json::to_vec(payload)?;
|
||||
let payload_b64 = BASE64_URL_SAFE_NO_PAD.encode(&json);
|
||||
let mut mac =
|
||||
HmacSha256::new_from_slice(secret.as_bytes()).expect("HMAC 接受任意长度密钥");
|
||||
mac.update(payload_b64.as_bytes());
|
||||
let sig_b64 = BASE64_URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
|
||||
Ok(format!("{payload_b64}.{sig_b64}"))
|
||||
}
|
||||
|
||||
/// 校验授权凭证（多 secret）：逐个尝试 secrets，任一验签通过即放行。
|
||||
///
|
||||
/// 返回 `(payload, 命中 secret 的索引)`：0 为主 secret（`SANDBOX_AUTH_SECRET`），
|
||||
/// \>0 为旧 secret（`SANDBOX_AUTH_SECRET_PREVIOUS` 第几个）。索引供 metric 区分轮换进度。
|
||||
/// 格式 / 过期错误与 secret 无关，只判一次；仅验签阶段遍历多 secret（支持无中断轮换的并存窗口）。
|
||||
pub fn verify_any(secrets: &[String], token: &str) -> Result<(AuthPayload, usize), AuthError> {
|
||||
let (payload_b64, sig_b64) = token.rsplit_once('.').ok_or(AuthError::Malformed)?;
|
||||
let sig_bytes = BASE64_URL_SAFE_NO_PAD
|
||||
.decode(sig_b64.as_bytes())
|
||||
.map_err(|_| AuthError::Malformed)?;
|
||||
|
||||
// 验签：任一 secret 通过即可（常量时间比较由 hmac verify_slice 提供）。
|
||||
let mut slot = None;
|
||||
for (i, secret) in secrets.iter().enumerate() {
|
||||
let mut mac = HmacSha256::new_from_slice(secret.as_bytes())
|
||||
.map_err(|_| AuthError::Malformed)?;
|
||||
mac.update(payload_b64.as_bytes());
|
||||
if mac.verify_slice(&sig_bytes).is_ok() {
|
||||
slot = Some(i);
|
||||
break;
|
||||
}
|
||||
}
|
||||
let slot = slot.ok_or(AuthError::BadSignature)?;
|
||||
|
||||
let payload_bytes = BASE64_URL_SAFE_NO_PAD
|
||||
.decode(payload_b64.as_bytes())
|
||||
.map_err(|_| AuthError::Malformed)?;
|
||||
let payload: AuthPayload =
|
||||
serde_json::from_slice(&payload_bytes).map_err(|_| AuthError::Malformed)?;
|
||||
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
if payload.exp <= now {
|
||||
return Err(AuthError::Expired);
|
||||
}
|
||||
Ok((payload, slot))
|
||||
}
|
||||
|
||||
/// 单 secret 校验（测试便捷包装：内部转 `verify_any`）。
|
||||
/// 仅测试用——生产路径 `require_authorization` 直连 `verify_any` 以支持多 secret 轮换。
|
||||
#[cfg(test)]
|
||||
pub fn verify(secret: &str, token: &str) -> Result<AuthPayload, AuthError> {
|
||||
verify_any(&[secret.to_string()], token).map(|(p, _)| p)
|
||||
}
|
||||
|
||||
/// 计算请求内容摘要：SHA256(canonical JSON of `{argv, stdin}`)。
|
||||
///
|
||||
/// 与后端 Python `compute_body_hash` 字节级一致：
|
||||
/// - JSON 紧凑（无空格），字段顺序固定 `argv, stdin`；`serde_json` 默认紧凑输出 + Option None → null，
|
||||
/// 对齐 Python `json.dumps(..., separators=(",",":"), ensure_ascii=False)`
|
||||
/// - 返回小写十六进制（64 字符），对齐 `hashlib.sha256().hexdigest()`
|
||||
/// - 注意：`rlimit_overrides` 不参与摘要（属授权属性，非执行内容）
|
||||
pub fn compute_body_hash(argv: &[String], stdin: &Option<String>) -> String {
|
||||
#[derive(Serialize)]
|
||||
struct BodyHashInput<'a> {
|
||||
argv: &'a [String],
|
||||
stdin: &'a Option<String>,
|
||||
}
|
||||
let bytes = serde_json::to_vec(&BodyHashInput { argv, stdin })
|
||||
.expect("argv/stdin 序列化不会失败");
|
||||
let digest = Sha256::digest(&bytes);
|
||||
digest.iter().map(|b| format!("{:02x}", b)).collect()
|
||||
}
|
||||
|
||||
/// 计算 code 模式请求内容摘要：SHA256(canonical JSON `{language, code, stdin}`)。
|
||||
///
|
||||
/// 与后端 Python `compute_body_hash_code` 字节级一致：
|
||||
/// - JSON 紧凑、字段顺序固定 `language, code, stdin`；Option None → null
|
||||
/// - 返回小写十六进制（64 字符）
|
||||
/// - code 模式凭证 body_hash 用此函数（argv 模式仍用 compute_body_hash），把凭证绑死到
|
||||
/// 具体代码内容，防 TTL 内重放执行任意代码。
|
||||
pub fn compute_body_hash_code(language: &str, code: &str, stdin: &Option<String>) -> String {
|
||||
#[derive(Serialize)]
|
||||
struct CodeHashInput<'a> {
|
||||
language: &'a str,
|
||||
code: &'a str,
|
||||
stdin: &'a Option<String>,
|
||||
}
|
||||
let bytes = serde_json::to_vec(&CodeHashInput {
|
||||
language,
|
||||
code,
|
||||
stdin,
|
||||
})
|
||||
.expect("language/code/stdin 序列化不会失败");
|
||||
let digest = Sha256::digest(&bytes);
|
||||
digest.iter().map(|b| format!("{:02x}", b)).collect()
|
||||
}
|
||||
|
||||
/// 计算文件操作摘要：SHA256(canonical JSON `{op, path}`)。
|
||||
///
|
||||
/// 与后端 Python `compute_file_op_hash` 字节级一致：紧凑 JSON、字段顺序固定 `op, path`、
|
||||
/// 小写十六进制。把凭证绑死到具体文件操作，防 TTL 内重放执行任意文件读写。
|
||||
pub fn compute_file_op_hash(op: &str, path: &str) -> String {
|
||||
#[derive(Serialize)]
|
||||
struct FileOpHashInput<'a> {
|
||||
op: &'a str,
|
||||
path: &'a str,
|
||||
}
|
||||
let bytes = serde_json::to_vec(&FileOpHashInput { op, path })
|
||||
.expect("op/path 序列化不会失败");
|
||||
let digest = Sha256::digest(&bytes);
|
||||
digest.iter().map(|b| format!("{:02x}", b)).collect()
|
||||
}
|
||||
|
||||
/// axum 中间件：校验 `Authorization: Bearer <token>`；`auth_enabled=false` 时直接放行。
|
||||
pub async fn require_authorization(
|
||||
State(state): State<Arc<AppState>>,
|
||||
mut req: Request,
|
||||
next: Next,
|
||||
) -> Response {
|
||||
if !state.config.auth_enabled {
|
||||
// 本地联调：跳过验签，但仍插入占位 payload，使 handler 能统一从扩展取 AuthPayload。
|
||||
req.extensions_mut().insert(AuthPayload {
|
||||
exp: 0,
|
||||
trace_id: None,
|
||||
risk_level: None,
|
||||
tool: None,
|
||||
character_id: None,
|
||||
run_id: None,
|
||||
body_hash: None,
|
||||
rlimit_overrides: None,
|
||||
file_op_hash: None,
|
||||
isolation_mode: None,
|
||||
});
|
||||
return next.run(req).await;
|
||||
}
|
||||
|
||||
let token = req
|
||||
.headers()
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|s| s.strip_prefix("Bearer ").map(str::trim));
|
||||
|
||||
match token {
|
||||
Some(t) => {
|
||||
// 主 secret + previous secrets 一并尝试（无中断轮换的并存窗口）。
|
||||
let mut secrets = vec![state.config.auth_secret.clone()];
|
||||
secrets.extend(state.config.auth_secret_previous.iter().cloned());
|
||||
match verify_any(&secrets, t) {
|
||||
Ok((payload, slot)) => {
|
||||
state.executor.metrics().record_auth_secret(slot);
|
||||
req.extensions_mut().insert(payload);
|
||||
next.run(req).await
|
||||
}
|
||||
Err(e) => unauthorized(e),
|
||||
}
|
||||
}
|
||||
None => unauthorized(AuthError::Missing),
|
||||
}
|
||||
}
|
||||
|
||||
fn unauthorized(err: AuthError) -> Response {
|
||||
let body = Json(serde_json::json!({
|
||||
"error": { "code": "sandbox_unauthorized", "message": err.to_string() }
|
||||
}));
|
||||
(StatusCode::UNAUTHORIZED, body).into_response()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const SECRET: &str = "test-shared-secret";
|
||||
|
||||
fn payload(exp_offset_secs: i64) -> AuthPayload {
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
AuthPayload {
|
||||
exp: (now as i64 + exp_offset_secs) as u64,
|
||||
trace_id: Some("trace_abc".into()),
|
||||
risk_level: Some("L1".into()),
|
||||
tool: Some("shell".into()),
|
||||
character_id: Some("character_01H".into()),
|
||||
run_id: Some("run_01H".into()),
|
||||
body_hash: Some("deadbeef".into()),
|
||||
rlimit_overrides: None,
|
||||
file_op_hash: None,
|
||||
isolation_mode: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sign_then_verify_roundtrip() {
|
||||
let token = sign(SECRET, &payload(60)).unwrap();
|
||||
let claims = verify(SECRET, &token).expect("合法凭证应校验通过");
|
||||
assert_eq!(claims.trace_id.as_deref(), Some("trace_abc"));
|
||||
assert_eq!(claims.risk_level.as_deref(), Some("L1"));
|
||||
assert_eq!(claims.body_hash.as_deref(), Some("deadbeef"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_wrong_secret() {
|
||||
let token = sign(SECRET, &payload(60)).unwrap();
|
||||
assert!(matches!(
|
||||
verify("wrong-secret", &token),
|
||||
Err(AuthError::BadSignature)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_tampered_token() {
|
||||
let mut token = sign(SECRET, &payload(60)).unwrap();
|
||||
// 篡改签名末位
|
||||
let last = token.len() - 1;
|
||||
token.replace_range(last.., if &token[last..] == "A" { "B" } else { "A" });
|
||||
assert!(verify(SECRET, &token).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_expired() {
|
||||
let token = sign(SECRET, &payload(-10)).unwrap();
|
||||
assert!(matches!(verify(SECRET, &token), Err(AuthError::Expired)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_malformed() {
|
||||
assert!(matches!(
|
||||
verify(SECRET, "not-a-valid-token"),
|
||||
Err(AuthError::Malformed)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_hash_known_vector_matches_python() {
|
||||
// 与后端 Python compute_body_hash 共享同一固定向量，保证跨端字节级一致。
|
||||
// 向量：{"argv":["sh","-c","echo hi"],"stdin":null}
|
||||
let argv = vec!["sh".to_string(), "-c".to_string(), "echo hi".to_string()];
|
||||
let hash = compute_body_hash(&argv, &None);
|
||||
assert_eq!(
|
||||
hash,
|
||||
"f5fa97ea2ee92c505174672da52c59c952e49c32fa56a1db5bd983c10a829540"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_hash_distinguishes_none_and_empty_stdin() {
|
||||
let argv = vec!["sh".to_string()];
|
||||
// None → "stdin":null；Some("") → "stdin":""，两者字节不同
|
||||
assert_ne!(
|
||||
compute_body_hash(&argv, &None),
|
||||
compute_body_hash(&argv, &Some(String::new()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_op_hash_known_vector_matches_python() {
|
||||
// 与后端 Python compute_file_op_hash 共享同一固定向量，保证跨端字节级一致。
|
||||
// 向量：{"op":"read","path":"data/note.txt"}
|
||||
let hash = compute_file_op_hash("read", "data/note.txt");
|
||||
assert_eq!(
|
||||
hash,
|
||||
"717f67252c61cf1681cb3c6018d3801b79a511564664567b57d12d775bb8fa53"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn file_op_hash_differs_per_op_and_path() {
|
||||
// 凭证绑死具体文件操作：op 或 path 不同则摘要不同，杜绝重放任意文件读写。
|
||||
let read_a = compute_file_op_hash("read", "a.txt");
|
||||
let write_a = compute_file_op_hash("write", "a.txt");
|
||||
let read_b = compute_file_op_hash("read", "b.txt");
|
||||
assert_ne!(read_a, write_a); // 同 path 不同 op
|
||||
assert_ne!(read_a, read_b); // 同 op 不同 path
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_any_accepts_previous_secret() {
|
||||
// 轮换并存窗口：主 secret + 旧 secret 都应被接受，且命中 slot 不同。
|
||||
let token_old = sign("old-secret", &payload(60)).unwrap();
|
||||
let token_new = sign("new-secret", &payload(60)).unwrap();
|
||||
let secrets = vec!["new-secret".to_string(), "old-secret".to_string()];
|
||||
// 新 secret 签发 → 命中 slot 0（primary）
|
||||
let (p, slot) = verify_any(&secrets, &token_new).unwrap();
|
||||
assert_eq!(slot, 0);
|
||||
assert_eq!(p.trace_id.as_deref(), Some("trace_abc"));
|
||||
// 旧 secret 签发 → 命中 slot 1（previous）
|
||||
let (_, slot) = verify_any(&secrets, &token_old).unwrap();
|
||||
assert_eq!(slot, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn verify_any_rejects_when_no_secret_matches() {
|
||||
let token = sign("actual-secret", &payload(60)).unwrap();
|
||||
let secrets = vec!["wrong1".to_string(), "wrong2".to_string()];
|
||||
assert!(matches!(
|
||||
verify_any(&secrets, &token),
|
||||
Err(AuthError::BadSignature)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rlimit_overrides_roundtrip_through_payload() {
|
||||
// rlimit_overrides 随凭证序列化/反序列化，且不进 body_hash
|
||||
let mut p = payload(60);
|
||||
let mut overrides = BTreeMap::new();
|
||||
overrides.insert("fsize".to_string(), 200_000_000u64);
|
||||
p.rlimit_overrides = Some(overrides);
|
||||
let token = sign(SECRET, &p).unwrap();
|
||||
let claims = verify(SECRET, &token).unwrap();
|
||||
assert_eq!(
|
||||
claims.rlimit_overrides.as_ref().unwrap().get("fsize"),
|
||||
Some(&200_000_000u64)
|
||||
);
|
||||
// body_hash 不受 rlimit_overrides 影响（测试用占位值，未绑定真实 argv/stdin）
|
||||
assert_eq!(claims.body_hash.as_deref(), Some("deadbeef"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_hash_code_known_vector_matches_python() {
|
||||
// 与后端 Python compute_body_hash_code 共享固定向量，保证跨端字节级一致。
|
||||
// 向量：{"language":"python","code":"print(1+1)","stdin":null}
|
||||
let hash = compute_body_hash_code("python", "print(1+1)", &None);
|
||||
assert_eq!(
|
||||
hash,
|
||||
"e7d633605676c3b19c642ee2437d91f9fb22562aaffad40f217d84c587322b69"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_hash_code_distinguishes_language_and_stdin() {
|
||||
let base = compute_body_hash_code("python", "print(1+1)", &None);
|
||||
// 不同 language → 不同摘要
|
||||
assert_ne!(base, compute_body_hash_code("node", "print(1+1)", &None));
|
||||
// 不同 stdin（None vs Some）→ 不同摘要
|
||||
assert_ne!(
|
||||
base,
|
||||
compute_body_hash_code("python", "print(1+1)", &Some("42".into()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn isolation_mode_roundtrip_and_backward_compat() {
|
||||
// 新字段随凭证序列化/反序列化
|
||||
let mut p = payload(60);
|
||||
p.isolation_mode = Some("container".into());
|
||||
let token = sign(SECRET, &p).unwrap();
|
||||
let claims = verify(SECRET, &token).unwrap();
|
||||
assert_eq!(claims.isolation_mode.as_deref(), Some("container"));
|
||||
|
||||
// 向后兼容：手工构造不含 isolation_mode 的旧 payload JSON，仍可验签且字段为 None。
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
let legacy_json = format!(
|
||||
"{{\"exp\":{},\"trace_id\":\"trace_abc\",\"risk_level\":\"L1\",\"tool\":\"shell\",\"character_id\":\"character_01H\",\"run_id\":\"run_01H\",\"body_hash\":\"deadbeef\",\"rlimit_overrides\":null,\"file_op_hash\":null}}",
|
||||
now as i64 + 60
|
||||
);
|
||||
let payload_b64 = BASE64_URL_SAFE_NO_PAD.encode(legacy_json.as_bytes());
|
||||
let mut mac = HmacSha256::new_from_slice(SECRET.as_bytes()).unwrap();
|
||||
mac.update(payload_b64.as_bytes());
|
||||
let sig_b64 = BASE64_URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes());
|
||||
let token = format!("{payload_b64}.{sig_b64}");
|
||||
let claims = verify(SECRET, &token).expect("旧凭证（无 isolation_mode）应可验签");
|
||||
assert!(claims.isolation_mode.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
//! code 模式转译：language + code → argv + 临时脚本文件。
|
||||
//!
|
||||
//! code 模式（`/v1/exec` 的 `code`/`language` 字段）由本模块转译为 argv，复用现有进程执行路径。
|
||||
//! 走脚本文件而非 `-c`，避免多行 / 模块导入 / shebang 问题；脚本经 `tempfile` 写入 OS 临时目录，
|
||||
//! 返回的 `CodeTranslation.script_path`（`TempPath`）由调用方持有至执行结束——drop 时自动删除。
|
||||
//! 流式执行场景须把 `CodeTranslation` move 进 SSE 闭包，确保子进程 open 脚本前文件不被删。
|
||||
|
||||
use std::io::Write;
|
||||
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
/// code 模式转译结果：argv 用于执行，script_path 须由调用方持有至执行结束。
|
||||
pub struct CodeTranslation {
|
||||
pub argv: Vec<String>,
|
||||
/// 临时脚本路径，持有至执行结束（RAII：drop 时自动删除）
|
||||
#[allow(dead_code)]
|
||||
pub script_path: tempfile::TempPath,
|
||||
}
|
||||
|
||||
/// 语言 → (interpreter, 脚本扩展名)。P0 仅支持 python；node/java 留 P2。
|
||||
pub fn interpreter_for(language: &str) -> Result<(&'static str, &'static str), String> {
|
||||
match language {
|
||||
"python" => Ok(("python3", "py")),
|
||||
other => Err(format!("不支持的语言: {other}（P0 仅支持 python）")),
|
||||
}
|
||||
}
|
||||
|
||||
/// 写 code 到临时脚本文件，返回 argv 与 TempPath（调用方持有至执行结束）。
|
||||
pub async fn translate(language: &str, code: &str) -> Result<CodeTranslation, String> {
|
||||
let (interp, ext) = interpreter_for(language)?;
|
||||
// NamedTempFile 是同步 IO，文件小、写一次，无需 tokio 化。
|
||||
let mut tmp = NamedTempFile::with_suffix(format!(".{ext}"))
|
||||
.map_err(|e| format!("创建临时脚本失败: {e}"))?;
|
||||
tmp.write_all(code.as_bytes())
|
||||
.map_err(|e| format!("写入临时脚本失败: {e}"))?;
|
||||
// detach 成 TempPath：文件留在磁盘，TempPath drop 时删除。
|
||||
let script_path = tmp.into_temp_path();
|
||||
let argv = vec![
|
||||
interp.to_string(),
|
||||
script_path.to_string_lossy().into_owned(),
|
||||
];
|
||||
Ok(CodeTranslation { argv, script_path })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn interpreter_for_python_ok_others_rejected() {
|
||||
assert_eq!(interpreter_for("python").unwrap(), ("python3", "py"));
|
||||
assert!(interpreter_for("node").is_err());
|
||||
assert!(interpreter_for("").is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn translate_python_writes_script_and_returns_argv() {
|
||||
let code = "print(1+1)\nimport sys\nprint(sys.argv)";
|
||||
let trans = translate("python", code).await.unwrap();
|
||||
assert_eq!(trans.argv[0], "python3");
|
||||
assert!(trans.argv[1].ends_with(".py"));
|
||||
let on_disk = std::fs::read_to_string(&trans.argv[1]).unwrap();
|
||||
assert_eq!(on_disk, code);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn translate_unsupported_language_errors() {
|
||||
assert!(translate("ruby", "puts 1").await.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn translate_temp_path_keeps_file_until_drop() {
|
||||
let trans = translate("python", "x = 1").await.unwrap();
|
||||
let path = trans.argv[1].clone();
|
||||
assert!(std::path::Path::new(&path).exists(), "持有 TempPath 时文件应存在");
|
||||
drop(trans);
|
||||
assert!(!std::path::Path::new(&path).exists(), "TempPath drop 后文件应删除");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
//! 沙箱服务运行配置，全部来自环境变量（与项目根 `.env` 统一）。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::env;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
|
||||
use anyhow::bail;
|
||||
|
||||
use crate::rlimits::{RlimitKind, RlimitsConfig};
|
||||
|
||||
/// 沙箱服务运行配置。字段公开以支持测试构造。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Config {
|
||||
/// 监听地址。
|
||||
pub host: String,
|
||||
/// 监听端口。
|
||||
pub port: u16,
|
||||
/// 沙箱工作根目录：所有执行的工作目录都约束在其之下。
|
||||
pub base_workdir: PathBuf,
|
||||
/// 单次执行的全局超时上限（请求 timeout_ms 不可超过该值）。
|
||||
pub max_timeout: Duration,
|
||||
/// 最大并发执行数。
|
||||
pub max_concurrent: usize,
|
||||
/// 单流（stdout / stderr 各自）最大捕获字节数。
|
||||
pub max_output_bytes: usize,
|
||||
/// 单次文件读写大小上限（字节）。文件操作在沙箱主进程执行（非 exec 子进程），
|
||||
/// 不受 RLIMIT_* 约束，故用此配置上限防止内存放大。
|
||||
pub max_file_bytes: usize,
|
||||
/// 默认注入子进程的环境变量白名单（在 env_clear 后注入）。
|
||||
pub allowed_env: Vec<String>,
|
||||
/// 是否校验调用方授权凭证。默认开启；关闭后跳过校验（仅限本地测试）。
|
||||
pub auth_enabled: bool,
|
||||
/// 授权 HMAC 共享密钥（与后端司命签发端共用）。
|
||||
pub auth_secret: String,
|
||||
/// 日志目录：tracing 文件 appender 写入位置。与后端 log_dir="../logs" 对齐
|
||||
/// （相对沙箱 cwd，即项目根 logs/）；容器内由 SANDBOX_LOG_DIR=/app/logs 覆盖。
|
||||
pub log_dir: PathBuf,
|
||||
/// 进程级资源限额配置（各维度 default/max + 禁用清单），由 env `SANDBOX_RLIMIT_*` 解析。
|
||||
pub rlimits: RlimitsConfig,
|
||||
/// 凭证轮换：旧 secret 列表（`SANDBOX_AUTH_SECRET_PREVIOUS` 逗号分隔），verify 时与主 secret 一并尝试。
|
||||
pub auth_secret_previous: Vec<String>,
|
||||
/// 凭证未下发 isolation_mode 时的兜底隔离模式（process/container）。
|
||||
pub isolation_default: String,
|
||||
/// 是否启用 Docker 容器模式。
|
||||
pub docker_enabled: bool,
|
||||
/// 预热池配置。
|
||||
pub warm_pool: WarmPoolConfig,
|
||||
}
|
||||
|
||||
/// 预热池配置。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WarmPoolConfig {
|
||||
/// 每种语言最大空闲容器数。
|
||||
pub max_idle_per_lang: usize,
|
||||
/// 全局最大容器数（各分桶总和）。
|
||||
pub global_max_containers: usize,
|
||||
/// 容器闲置超时，超过则销毁回收。
|
||||
pub idle_ttl: Duration,
|
||||
/// 创建容器的超时时间。
|
||||
pub create_timeout: Duration,
|
||||
/// 每容器内存上限（字节）。
|
||||
pub memory_per_container: i64,
|
||||
/// 每容器 CPU 上限（纳秒核）。
|
||||
pub nano_cpus_per_container: i64,
|
||||
/// 每容器进程数上限。
|
||||
pub pids_limit: i64,
|
||||
/// 语言 → 镜像名映射。
|
||||
pub images: HashMap<String, String>,
|
||||
/// Docker socket 路径。
|
||||
pub docker_socket: String,
|
||||
/// 容器网络模式。非空时设置 network_mode，空字符串使用 Docker 默认 bridge。
|
||||
pub container_network: String,
|
||||
}
|
||||
|
||||
impl Default for WarmPoolConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_idle_per_lang: 3,
|
||||
global_max_containers: 20,
|
||||
idle_ttl: Duration::from_secs(300),
|
||||
create_timeout: Duration::from_secs(60),
|
||||
memory_per_container: 512 * 1024 * 1024,
|
||||
nano_cpus_per_container: 1_000_000_000,
|
||||
pids_limit: 64,
|
||||
images: HashMap::new(),
|
||||
docker_socket: "unix:///var/run/docker.sock".into(),
|
||||
container_network: String::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn from_env() -> anyhow::Result<Self> {
|
||||
let auth_enabled = parse_or("SANDBOX_AUTH_ENABLED", true);
|
||||
let auth_secret = env::var("SANDBOX_AUTH_SECRET").unwrap_or_default();
|
||||
if auth_enabled && auth_secret.is_empty() {
|
||||
bail!("SANDBOX_AUTH_ENABLED=true 时必须设置 SANDBOX_AUTH_SECRET（与后端共享的 HMAC 密钥）");
|
||||
}
|
||||
|
||||
let base_workdir = PathBuf::from(
|
||||
env::var("SANDBOX_BASE_WORKDIR").unwrap_or_else(|_| "../sandbox-work".into()),
|
||||
);
|
||||
let log_dir = PathBuf::from(
|
||||
env::var("SANDBOX_LOG_DIR").unwrap_or_else(|_| "../logs".into()),
|
||||
);
|
||||
let max_timeout = Duration::from_millis(parse_or("SANDBOX_MAX_TIMEOUT_MS", 30_000_u64));
|
||||
|
||||
// 进程级资源限额：每维度 default（普通工具默认）+ max（全局硬天花板，未设则=default）。
|
||||
// CPU 默认对齐全局超时上限（一个命令最多烧这么多 CPU 秒）。
|
||||
let cpu_default = max_timeout.as_secs().max(1);
|
||||
let mut rlimits = RlimitsConfig::new();
|
||||
rlimits.set(
|
||||
RlimitKind::Fsize,
|
||||
parse_or("SANDBOX_RLIMIT_FSIZE_BYTES", 67_108_864_u64),
|
||||
parse_or(
|
||||
"SANDBOX_RLIMIT_FSIZE_MAX",
|
||||
parse_or("SANDBOX_RLIMIT_FSIZE_BYTES", 67_108_864_u64),
|
||||
),
|
||||
);
|
||||
rlimits.set(
|
||||
RlimitKind::As,
|
||||
parse_or("SANDBOX_RLIMIT_AS_BYTES", 536_870_912_u64),
|
||||
parse_or(
|
||||
"SANDBOX_RLIMIT_AS_MAX",
|
||||
parse_or("SANDBOX_RLIMIT_AS_BYTES", 536_870_912_u64),
|
||||
),
|
||||
);
|
||||
rlimits.set(
|
||||
RlimitKind::Cpu,
|
||||
parse_or("SANDBOX_RLIMIT_CPU_SECS", cpu_default),
|
||||
parse_or("SANDBOX_RLIMIT_CPU_MAX", cpu_default),
|
||||
);
|
||||
rlimits.set(
|
||||
RlimitKind::Nproc,
|
||||
parse_or("SANDBOX_RLIMIT_NPROC", 64_u64),
|
||||
parse_or(
|
||||
"SANDBOX_RLIMIT_NPROC_MAX",
|
||||
parse_or("SANDBOX_RLIMIT_NPROC", 64_u64),
|
||||
),
|
||||
);
|
||||
rlimits.set(
|
||||
RlimitKind::Nofile,
|
||||
parse_or("SANDBOX_RLIMIT_NOFILE", 256_u64),
|
||||
parse_or(
|
||||
"SANDBOX_RLIMIT_NOFILE_MAX",
|
||||
parse_or("SANDBOX_RLIMIT_NOFILE", 256_u64),
|
||||
),
|
||||
);
|
||||
// 显式禁用维度（逗号分隔，如 nproc,nofile）；未知维度名静默忽略（from_env 早于 tracing 初始化）。
|
||||
for s in env::var("SANDBOX_RLIMIT_DISABLE")
|
||||
.unwrap_or_default()
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
{
|
||||
if let Some(kind) = RlimitKind::from_key(s) {
|
||||
rlimits.disable(kind);
|
||||
}
|
||||
}
|
||||
|
||||
// 凭证轮换：旧 secret 列表，verify 时与主 secret 一并尝试（无中断轮换的并存窗口）。
|
||||
let auth_secret_previous = env::var("SANDBOX_AUTH_SECRET_PREVIOUS")
|
||||
.unwrap_or_default()
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)
|
||||
.collect();
|
||||
|
||||
Ok(Self {
|
||||
host: env::var("SANDBOX_HOST").unwrap_or_else(|_| "0.0.0.0".into()),
|
||||
port: parse_or("SANDBOX_PORT", 8090),
|
||||
base_workdir,
|
||||
max_timeout,
|
||||
max_concurrent: parse_or("SANDBOX_MAX_CONCURRENT", 8),
|
||||
max_output_bytes: parse_or("SANDBOX_MAX_OUTPUT_BYTES", 1_048_576),
|
||||
max_file_bytes: parse_or("SANDBOX_MAX_FILE_BYTES", 1_048_576),
|
||||
allowed_env: env::var("SANDBOX_ALLOWED_ENV")
|
||||
.unwrap_or_else(|_| "PATH,HOME,TZ,LANG".into())
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|s| !s.is_empty())
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
auth_enabled,
|
||||
auth_secret,
|
||||
auth_secret_previous,
|
||||
rlimits,
|
||||
log_dir,
|
||||
isolation_default: env::var("SANDBOX_ISOLATION_DEFAULT")
|
||||
.unwrap_or_else(|_| "process".into()),
|
||||
docker_enabled: parse_or("SANDBOX_DOCKER_ENABLED", false),
|
||||
warm_pool: WarmPoolConfig {
|
||||
max_idle_per_lang: parse_or("SANDBOX_WARM_POOL_MAX_IDLE", 3),
|
||||
global_max_containers: parse_or("SANDBOX_WARM_POOL_GLOBAL_MAX", 20),
|
||||
idle_ttl: Duration::from_secs(parse_or("SANDBOX_CONTAINER_IDLE_TTL_SECS", 300)),
|
||||
create_timeout: Duration::from_millis(parse_or("SANDBOX_CONTAINER_CREATE_TIMEOUT_MS", 60_000)),
|
||||
memory_per_container: parse_or("SANDBOX_CONTAINER_MEMORY_BYTES", 536_870_912i64),
|
||||
nano_cpus_per_container: parse_or("SANDBOX_CONTAINER_CPU_NANOCORES", 1_000_000_000i64),
|
||||
pids_limit: parse_or("SANDBOX_CONTAINER_PIDS_LIMIT", 64i64),
|
||||
docker_socket: env::var("SANDBOX_DOCKER_SOCKET")
|
||||
.unwrap_or_else(|_| "unix:///var/run/docker.sock".into()),
|
||||
container_network: env::var("SANDBOX_CONTAINER_NETWORK")
|
||||
.unwrap_or_default(),
|
||||
images: {
|
||||
let mut m = HashMap::new();
|
||||
if let Some(img) = env::var("SANDBOX_CONTAINER_IMAGE_PYTHON").ok().filter(|s| !s.is_empty()) {
|
||||
m.insert("python".into(), img);
|
||||
}
|
||||
if let Some(img) = env::var("SANDBOX_CONTAINER_IMAGE_NODE").ok().filter(|s| !s.is_empty()) {
|
||||
m.insert("node".into(), img);
|
||||
}
|
||||
if let Some(img) = env::var("SANDBOX_CONTAINER_IMAGE_JAVA").ok().filter(|s| !s.is_empty()) {
|
||||
m.insert("java".into(), img);
|
||||
}
|
||||
if let Some(img) = env::var("SANDBOX_CONTAINER_IMAGE_GENERIC").ok().filter(|s| !s.is_empty()) {
|
||||
m.insert("generic".into(), img);
|
||||
}
|
||||
m
|
||||
},
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_or<T: std::str::FromStr>(key: &str, default: T) -> T {
|
||||
env::var(key)
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(default)
|
||||
}
|
||||
@@ -0,0 +1,707 @@
|
||||
//! Docker 容器执行封装：在预热池容器中执行命令、构建容器配置、镜像管理。
|
||||
//!
|
||||
//! 与 executor.rs 共享 ExecOutcome / ExecError 类型。
|
||||
//! 调用方（warm_pool）先 acquire 容器句柄，再调用 exec_in_container 执行。
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use bollard::container::Config;
|
||||
use bollard::exec::{CreateExecOptions, StartExecResults};
|
||||
use bollard::image::CreateImageOptions;
|
||||
use bollard::container::LogOutput;
|
||||
use bollard::models::HostConfig;
|
||||
use bollard::{Docker, API_DEFAULT_VERSION};
|
||||
use futures_util::StreamExt;
|
||||
|
||||
use crate::config::WarmPoolConfig;
|
||||
use crate::executor::{ExecError, ExecOutcome, StreamEvent};
|
||||
|
||||
/// 在已运行的容器中执行一条命令，返回 ExecOutcome（与进程模式相同的接口）。
|
||||
///
|
||||
/// # 参数
|
||||
/// - `argv`：命令及参数数组
|
||||
/// - `stdin`：可选标准输入
|
||||
/// - `env`：环境变量键值对
|
||||
/// - `timeout`：执行超时
|
||||
/// - `max_output`：单流最大捕获字节数
|
||||
pub async fn exec_in_container(
|
||||
docker: &Docker,
|
||||
container_id: &str,
|
||||
argv: &[String],
|
||||
stdin: Option<String>,
|
||||
env: &[(String, String)],
|
||||
timeout: Duration,
|
||||
max_output: usize,
|
||||
) -> Result<ExecOutcome, ExecError> {
|
||||
let start = Instant::now();
|
||||
|
||||
// 1) 创建 exec 实例
|
||||
let exec = docker
|
||||
.create_exec(
|
||||
container_id,
|
||||
CreateExecOptions {
|
||||
cmd: Some(argv.to_vec()),
|
||||
env: Some(
|
||||
env.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
attach_stdout: Some(true),
|
||||
attach_stderr: Some(true),
|
||||
attach_stdin: Some(stdin.is_some()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| ExecError::Internal {
|
||||
user_msg: "创建容器 exec 失败",
|
||||
detail: format!("create_exec 失败: {e}"),
|
||||
})?;
|
||||
|
||||
// 2) 启动 exec 并捕获输出
|
||||
let mut stdout = Vec::new();
|
||||
let mut stderr = Vec::new();
|
||||
let mut stdout_truncated = false;
|
||||
let mut stderr_truncated = false;
|
||||
let mut timeout_hit = false;
|
||||
let cancelled = false;
|
||||
|
||||
let exec_result = tokio::time::timeout(timeout, async {
|
||||
match docker.start_exec(&exec.id, None).await {
|
||||
Ok(StartExecResults::Attached {
|
||||
mut output,
|
||||
mut input,
|
||||
}) => {
|
||||
// 先写入 stdin（若有），然后关闭输入管道
|
||||
if let Some(data) = &stdin {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let _ = input.write_all(data.as_bytes()).await;
|
||||
let _ = input.shutdown().await;
|
||||
}
|
||||
drop(input); // 关闭 stdin 触发命令的 EOF
|
||||
|
||||
while let Some(chunk) = output.next().await {
|
||||
match chunk {
|
||||
Ok(LogOutput::StdOut { message }) => {
|
||||
let data = message.as_ref();
|
||||
if stdout.len() < max_output {
|
||||
let room = max_output - stdout.len();
|
||||
let take = data.len().min(room);
|
||||
stdout.extend_from_slice(&data[..take]);
|
||||
if take < data.len() {
|
||||
stdout_truncated = true;
|
||||
}
|
||||
} else {
|
||||
stdout_truncated = true;
|
||||
}
|
||||
}
|
||||
Ok(LogOutput::StdErr { message }) => {
|
||||
let data = message.as_ref();
|
||||
if stderr.len() < max_output {
|
||||
let room = max_output - stderr.len();
|
||||
let take = data.len().min(room);
|
||||
stderr.extend_from_slice(&data[..take]);
|
||||
if take < data.len() {
|
||||
stderr_truncated = true;
|
||||
}
|
||||
} else {
|
||||
stderr_truncated = true;
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, exec_id = %exec.id, "容器 exec 流错误");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(_) => {
|
||||
// Detached（无输出者适用），直接检查退出码
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, exec_id = %exec.id, "启动容器 exec 失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
if exec_result.is_err() {
|
||||
timeout_hit = true;
|
||||
// Docker exec 无 kill 端点；exec 会在容器中继续运行，
|
||||
// 但容器会在 release 时被销毁，无需额外清理。
|
||||
}
|
||||
|
||||
// 4) 获取退出码
|
||||
let exit_code = docker
|
||||
.inspect_exec(&exec.id)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|i| i.exit_code)
|
||||
.map(|c| c as i32);
|
||||
|
||||
Ok(ExecOutcome {
|
||||
exit_code,
|
||||
stdout: String::from_utf8_lossy(&stdout).into_owned(),
|
||||
stderr: String::from_utf8_lossy(&stderr).into_owned(),
|
||||
duration_ms: start.elapsed().as_millis(),
|
||||
timeout: timeout_hit,
|
||||
cancelled,
|
||||
stdout_truncated,
|
||||
stderr_truncated,
|
||||
})
|
||||
}
|
||||
|
||||
/// 流式执行：在容器中执行命令并通过 channel 实时推送输出事件。
|
||||
pub async fn exec_in_container_stream(
|
||||
docker: &Docker,
|
||||
container_id: &str,
|
||||
argv: &[String],
|
||||
stdin: Option<String>,
|
||||
env: &[(String, String)],
|
||||
timeout: Duration,
|
||||
max_output: usize,
|
||||
) -> Result<tokio::sync::mpsc::Receiver<StreamEvent>, ExecError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let exec = docker
|
||||
.create_exec(
|
||||
container_id,
|
||||
CreateExecOptions {
|
||||
cmd: Some(argv.to_vec()),
|
||||
env: Some(
|
||||
env.iter()
|
||||
.map(|(k, v)| format!("{k}={v}"))
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
attach_stdout: Some(true),
|
||||
attach_stderr: Some(true),
|
||||
attach_stdin: Some(stdin.is_some()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(|e| ExecError::Internal {
|
||||
user_msg: "创建容器 exec 失败",
|
||||
detail: format!("create_exec 失败: {e}"),
|
||||
})?;
|
||||
|
||||
let (tx, rx) = tokio::sync::mpsc::channel::<StreamEvent>(64);
|
||||
|
||||
// 消费 exec 输出并转发到 channel
|
||||
let docker_inner = docker.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut out_written = 0usize;
|
||||
let mut err_written = 0usize;
|
||||
let mut out_trunc = false;
|
||||
let mut err_trunc = false;
|
||||
|
||||
let exec_result = tokio::time::timeout(timeout, async {
|
||||
match docker_inner.start_exec(&exec.id, None).await {
|
||||
Ok(StartExecResults::Attached {
|
||||
mut output,
|
||||
mut input,
|
||||
}) => {
|
||||
// 先写入 stdin（若有），然后关闭输入管道
|
||||
if let Some(data) = &stdin {
|
||||
use tokio::io::AsyncWriteExt;
|
||||
let _ = input.write_all(data.as_bytes()).await;
|
||||
let _ = input.shutdown().await;
|
||||
}
|
||||
drop(input);
|
||||
|
||||
while let Some(chunk) = output.next().await {
|
||||
match chunk {
|
||||
Ok(log_output) => {
|
||||
match log_output {
|
||||
LogOutput::StdOut { message } => {
|
||||
let data = message.as_ref();
|
||||
if out_written < max_output {
|
||||
let room = max_output - out_written;
|
||||
let take = data.len().min(room);
|
||||
if take > 0 {
|
||||
let _ = tx
|
||||
.send(StreamEvent::Stdout(data[..take].to_vec()))
|
||||
.await;
|
||||
}
|
||||
out_written += take;
|
||||
if take < data.len() {
|
||||
out_trunc = true;
|
||||
}
|
||||
} else {
|
||||
out_trunc = true;
|
||||
}
|
||||
}
|
||||
LogOutput::StdErr { message } => {
|
||||
let data = message.as_ref();
|
||||
if err_written < max_output {
|
||||
let room = max_output - err_written;
|
||||
let take = data.len().min(room);
|
||||
if take > 0 {
|
||||
let _ = tx
|
||||
.send(StreamEvent::Stderr(data[..take].to_vec()))
|
||||
.await;
|
||||
}
|
||||
err_written += take;
|
||||
if take < data.len() {
|
||||
err_trunc = true;
|
||||
}
|
||||
} else {
|
||||
err_trunc = true;
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "容器 exec 流错误");
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "启动容器 exec 失败");
|
||||
}
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
let timeout_hit = exec_result.is_err();
|
||||
if timeout_hit {
|
||||
// Docker exec 无 kill 端点；exec 在容器内继续运行，
|
||||
// 容器会在 release 时被销毁，无需额外清理。
|
||||
}
|
||||
|
||||
let exit_code = docker_inner
|
||||
.inspect_exec(&exec.id)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|i| i.exit_code)
|
||||
.map(|c| c as i32);
|
||||
|
||||
let _ = tx
|
||||
.send(StreamEvent::Done(ExecOutcome {
|
||||
exit_code,
|
||||
stdout: String::new(),
|
||||
stderr: String::new(),
|
||||
duration_ms: start.elapsed().as_millis(),
|
||||
timeout: timeout_hit,
|
||||
cancelled: false,
|
||||
stdout_truncated: out_trunc,
|
||||
stderr_truncated: err_trunc,
|
||||
}))
|
||||
.await;
|
||||
});
|
||||
|
||||
Ok(rx)
|
||||
}
|
||||
|
||||
/// 为预热容器构建 Docker Config。
|
||||
///
|
||||
/// 预热容器使用 `sleep infinity` 待命，实际执行通过 `exec_in_container` 注入命令。
|
||||
pub fn build_container_config(
|
||||
config: &WarmPoolConfig,
|
||||
image: &str,
|
||||
base_workdir: &str,
|
||||
) -> Config<String> {
|
||||
// Docker bind mount 要求绝对路径；相对路径返回 400 错误。
|
||||
if !base_workdir.starts_with('/') {
|
||||
tracing::warn!(
|
||||
base_workdir,
|
||||
"base_workdir 不是绝对路径，Docker bind mount 将失败"
|
||||
);
|
||||
}
|
||||
let network_mode = if config.container_network.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(config.container_network.clone())
|
||||
};
|
||||
Config {
|
||||
image: Some(image.to_string()),
|
||||
cmd: Some(vec!["sleep".into(), "infinity".into()]),
|
||||
host_config: Some(HostConfig {
|
||||
network_mode,
|
||||
readonly_rootfs: Some(true),
|
||||
cap_drop: Some(vec!["ALL".into()]),
|
||||
security_opt: Some(vec!["no-new-privileges:true".into()]),
|
||||
memory: Some(config.memory_per_container),
|
||||
memory_swap: Some(config.memory_per_container),
|
||||
pids_limit: Some(config.pids_limit),
|
||||
nano_cpus: Some(config.nano_cpus_per_container),
|
||||
tmpfs: Some(HashMap::from([(
|
||||
"/tmp".into(),
|
||||
"size=67108864,mode=1777".into(),
|
||||
)])),
|
||||
binds: Some(vec![format!("{base_workdir}:/app/sandbox-work:rw")]),
|
||||
..Default::default()
|
||||
}),
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
/// 检查镜像是否存在，不存在则拉取。
|
||||
pub async fn ensure_image(docker: &Docker, image: &str) -> Result<(), String> {
|
||||
let images = docker
|
||||
.list_images::<String>(None)
|
||||
.await
|
||||
.map_err(|e| format!("列出镜像失败: {e}"))?;
|
||||
|
||||
if images
|
||||
.iter()
|
||||
.any(|i| i.repo_tags.contains(&image.to_string()))
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
tracing::info!(image, "拉取容器镜像（首次使用）");
|
||||
docker
|
||||
.create_image(
|
||||
Some(CreateImageOptions {
|
||||
from_image: image,
|
||||
..Default::default()
|
||||
}),
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.for_each(|_| async {}) // 静默消费拉取流
|
||||
.await;
|
||||
|
||||
tracing::info!(image, "镜像拉取完成");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 连接 Docker 守护进程。
|
||||
pub fn connect_docker(socket: &str) -> Result<Docker, String> {
|
||||
if socket.starts_with("unix://") {
|
||||
Docker::connect_with_unix(socket.trim_start_matches("unix://"), 120, API_DEFAULT_VERSION)
|
||||
.map_err(|e| format!("连接 Docker Unix socket 失败: {e}"))
|
||||
} else if socket.starts_with("tcp://") {
|
||||
Docker::connect_with_http(socket, 120, API_DEFAULT_VERSION)
|
||||
.map_err(|e| format!("连接 Docker TCP 失败: {e}"))
|
||||
} else {
|
||||
Docker::connect_with_local_defaults()
|
||||
.map_err(|e| format!("连接本地 Docker 失败: {e}"))
|
||||
}
|
||||
}
|
||||
|
||||
/// 从 argv 首元素推断语言（用于选择预热池分桶）。
|
||||
pub fn detect_language_from_argv(argv: &[String]) -> &str {
|
||||
let cmd = argv.first().map(|s| s.as_str()).unwrap_or("");
|
||||
if cmd.ends_with("python3") || cmd == "python" || cmd.contains("python") {
|
||||
"python"
|
||||
} else if cmd.ends_with("node") || cmd.contains("node") {
|
||||
"node"
|
||||
} else if cmd.ends_with("java") || cmd.contains("java") {
|
||||
"java"
|
||||
} else {
|
||||
"generic"
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detect_language_from_argv_python() {
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["python3".into(), "script.py".into()]),
|
||||
"python"
|
||||
);
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["python".into(), "-c".into(), "print(1)".into()]),
|
||||
"python"
|
||||
);
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["/usr/bin/python3".into(), "script.py".into()]),
|
||||
"python"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_language_from_argv_node() {
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["node".into(), "index.js".into()]),
|
||||
"node"
|
||||
);
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["/usr/local/bin/node".into(), "app.js".into()]),
|
||||
"node"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_language_from_argv_java() {
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["java".into(), "-jar".into(), "app.jar".into()]),
|
||||
"java"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_language_from_argv_fallback() {
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["cat".into(), "/etc/passwd".into()]),
|
||||
"generic"
|
||||
);
|
||||
assert_eq!(
|
||||
detect_language_from_argv(&["sh".into(), "-c".into(), "echo hi".into()]),
|
||||
"generic"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_language_from_argv_empty() {
|
||||
assert_eq!(detect_language_from_argv(&[]), "generic");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_container_config_sets_basic_fields() {
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images.insert(
|
||||
"python".into(),
|
||||
"python:3.12-slim".into(),
|
||||
);
|
||||
let image = cfg.images.get("python").unwrap();
|
||||
let config = build_container_config(&cfg, image, "/app/sandbox-work");
|
||||
assert_eq!(config.image.as_deref(), Some("python:3.12-slim"));
|
||||
assert!(config.host_config.is_some());
|
||||
let hc = config.host_config.unwrap();
|
||||
assert_eq!(hc.network_mode, None, "默认容器网络应为 None（bridge）");
|
||||
assert_eq!(hc.readonly_rootfs, Some(true));
|
||||
assert_eq!(hc.cap_drop, Some(vec!["ALL".into()]));
|
||||
assert_eq!(hc.memory, Some(512 * 1024 * 1024));
|
||||
}
|
||||
}
|
||||
|
||||
/// Docker 集成测试（需要在运行中的 Docker 守护进程）。
|
||||
/// cargo test --features docker-tests
|
||||
#[cfg(feature = "docker-tests")]
|
||||
#[cfg(test)]
|
||||
mod docker_integration_tests {
|
||||
use super::*;
|
||||
use crate::warm_pool::WarmPool;
|
||||
|
||||
fn test_image() -> &'static str {
|
||||
"alpine:3.20"
|
||||
}
|
||||
|
||||
async fn setup_docker() -> Option<(Docker, String)> {
|
||||
let docker = connect_docker("unix:///var/run/docker.sock").ok()?;
|
||||
ensure_image(&docker, test_image()).await.ok()?;
|
||||
Some((docker, test_image().to_string()))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exec_echo_captures_stdout() {
|
||||
let (docker, image) = match setup_docker().await {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
eprintln!("跳过：Docker 不可用");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images.insert("generic".into(), image);
|
||||
let pool = WarmPool::new(cfg, "/tmp".into(), docker.clone());
|
||||
|
||||
let config = build_container_config(pool.config(), &image, "/tmp");
|
||||
let container = docker
|
||||
.create_container::<String, String>(None, config)
|
||||
.await
|
||||
.expect("创建容器");
|
||||
docker
|
||||
.start_container::<String>(&container.id, None)
|
||||
.await
|
||||
.expect("启动容器");
|
||||
|
||||
let outcome = exec_in_container(
|
||||
&docker,
|
||||
&container.id,
|
||||
&["echo".into(), "hello-docker".into()],
|
||||
None,
|
||||
&[],
|
||||
Duration::from_secs(10),
|
||||
4096,
|
||||
)
|
||||
.await
|
||||
.expect("exec_in_container");
|
||||
|
||||
assert_eq!(outcome.exit_code, Some(0));
|
||||
assert!(outcome.stdout.contains("hello-docker"));
|
||||
assert!(!outcome.timeout);
|
||||
assert!(!outcome.cancelled);
|
||||
|
||||
docker
|
||||
.remove_container(
|
||||
&container.id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
v: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exec_timeout_kills() {
|
||||
let (docker, image) = match setup_docker().await {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
eprintln!("跳过：Docker 不可用");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images.insert("generic".into(), image);
|
||||
let pool = WarmPool::new(cfg, "/tmp".into(), docker.clone());
|
||||
|
||||
let config = build_container_config(pool.config(), &image, "/tmp");
|
||||
let container = docker
|
||||
.create_container::<String, String>(None, config)
|
||||
.await
|
||||
.expect("创建容器");
|
||||
docker
|
||||
.start_container::<String>(&container.id, None)
|
||||
.await
|
||||
.expect("启动容器");
|
||||
|
||||
let outcome = exec_in_container(
|
||||
&docker,
|
||||
&container.id,
|
||||
&["sleep".into(), "30".into()],
|
||||
None,
|
||||
&[],
|
||||
Duration::from_millis(200),
|
||||
4096,
|
||||
)
|
||||
.await
|
||||
.expect("exec_in_container");
|
||||
|
||||
assert!(outcome.timeout, "应超时返回");
|
||||
assert!(outcome.duration_ms < 5000);
|
||||
|
||||
docker
|
||||
.remove_container(
|
||||
&container.id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
v: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn exec_in_container_stream_emits_events() {
|
||||
let (docker, image) = match setup_docker().await {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
eprintln!("跳过：Docker 不可用");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images.insert("generic".into(), image);
|
||||
let pool = WarmPool::new(cfg, "/tmp".into(), docker.clone());
|
||||
|
||||
let config = build_container_config(pool.config(), &image, "/tmp");
|
||||
let container = docker
|
||||
.create_container::<String, String>(None, config)
|
||||
.await
|
||||
.expect("创建容器");
|
||||
docker
|
||||
.start_container::<String>(&container.id, None)
|
||||
.await
|
||||
.expect("启动容器");
|
||||
|
||||
let mut rx = exec_in_container_stream(
|
||||
&docker,
|
||||
&container.id,
|
||||
&["echo".into(), "stream-test".into()],
|
||||
None,
|
||||
&[],
|
||||
Duration::from_secs(10),
|
||||
4096,
|
||||
)
|
||||
.await
|
||||
.expect("exec_in_container_stream");
|
||||
|
||||
use futures_util::StreamExt;
|
||||
let mut found_stdout = false;
|
||||
let mut found_done = false;
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
StreamEvent::Stdout(data) => {
|
||||
if String::from_utf8_lossy(&data).contains("stream-test") {
|
||||
found_stdout = true;
|
||||
}
|
||||
}
|
||||
StreamEvent::Done(o) => {
|
||||
assert_eq!(o.exit_code, Some(0));
|
||||
found_done = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
assert!(found_stdout, "应收到 stdout 事件");
|
||||
assert!(found_done, "应以 Done 结束");
|
||||
|
||||
docker
|
||||
.remove_container(
|
||||
&container.id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
v: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.ok();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn warm_pool_acquire_release_cycle() {
|
||||
let (docker, image) = match setup_docker().await {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
eprintln!("跳过：Docker 不可用");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images.insert("generic".into(), image);
|
||||
cfg.max_idle_per_lang = 1;
|
||||
cfg.global_max_containers = 5;
|
||||
|
||||
let pool = WarmPool::new(cfg, "/tmp".into(), docker);
|
||||
pool.warmup_all().await;
|
||||
|
||||
let stats = pool.collect_metrics().await;
|
||||
let (idle, ..) = stats.get("generic").copied().unwrap_or((0, 0, 0));
|
||||
assert!(idle >= 1, "预热后应有至少 1 个空闲容器");
|
||||
|
||||
let handle = pool.clone().acquire("generic").await.expect("acquire");
|
||||
assert_eq!(handle.language, "generic");
|
||||
|
||||
pool.release(handle).await;
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
let stats2 = pool.collect_metrics().await;
|
||||
let (idle2, ..) = stats2.get("generic").copied().unwrap_or((0, 0, 0));
|
||||
assert!(idle2 >= 1, "release 后应有替补容器");
|
||||
|
||||
pool.shutdown().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
//! 统一错误类型，序列化为跨语言契约约定的结构。
|
||||
//! 约定见 docs/architecture/service-boundaries.md 第 7 节。
|
||||
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::Json;
|
||||
use serde_json::json;
|
||||
|
||||
pub enum ApiError {
|
||||
/// 请求参数不合法（4xx）。
|
||||
BadRequest(String),
|
||||
/// 未授权：授权凭证缺失 / 过期 / 签名无效，或请求内容与凭证 body_hash 不匹配（401）。
|
||||
Unauthorized(String),
|
||||
/// 资源不存在，例如取消一个未知或已完成的任务（404）。
|
||||
NotFound(String),
|
||||
/// 服务暂不可用，例如并发闸耗尽或容器资源不足（503）。
|
||||
Unavailable(String),
|
||||
/// 执行失败或内部错误（500）。
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
impl IntoResponse for ApiError {
|
||||
fn into_response(self) -> Response {
|
||||
let (code, message, status) = match self {
|
||||
ApiError::BadRequest(m) => ("bad_request", m, StatusCode::BAD_REQUEST),
|
||||
ApiError::Unauthorized(m) => {
|
||||
("sandbox_unauthorized", m, StatusCode::UNAUTHORIZED)
|
||||
}
|
||||
ApiError::NotFound(m) => ("not_found", m, StatusCode::NOT_FOUND),
|
||||
ApiError::Unavailable(m) => {
|
||||
("service_unavailable", m, StatusCode::SERVICE_UNAVAILABLE)
|
||||
}
|
||||
ApiError::Internal(m) => ("exec_failed", m, StatusCode::INTERNAL_SERVER_ERROR),
|
||||
};
|
||||
let body = Json(json!({ "error": { "code": code, "message": message } }));
|
||||
(status, body).into_response()
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,474 @@
|
||||
//! 文件 IO：在沙箱 base 工作目录内提供不经 shell 的二进制安全文件读写。
|
||||
//!
|
||||
//! 与 builtins 经 shell（cat/sed/ls）的文件操作相比，原生端点不经 shell 解释
|
||||
//! （无注入面）、二进制安全（base64 保真）、大小精确可控、含 stat 元信息。
|
||||
//!
|
||||
//! 路径安全：所有路径经 `resolve_within_base` 校验——
|
||||
//! - read/list/stat（目标已存在）：`canonicalize` 解析 `..` 与符号链接后校验 `starts_with(base_canon)`。
|
||||
//! - write（目标可不存在）：先词汇预检（`..` 不可使深度为负，避免 `create_dir_all` 在 base 外
|
||||
//! 产生副作用），再 `canonicalize` 父目录兜底符号链接。
|
||||
//!
|
||||
//! 与 executor.rs 的 workdir 校验同一保证。
|
||||
//!
|
||||
//! 设计权衡：文件操作在沙箱主进程执行（非 exec 子进程），不受 RLIMIT_* 约束，
|
||||
//! 改由配置上限 `SANDBOX_MAX_FILE_BYTES` 限制单次读写大小（防内存放大）。
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use base64::prelude::{BASE64_URL_SAFE_NO_PAD, Engine as _};
|
||||
use serde::Serialize;
|
||||
|
||||
/// 文件操作错误。
|
||||
#[derive(Debug)]
|
||||
pub enum FileError {
|
||||
/// 路径非法或逃逸出沙箱根目录。
|
||||
InvalidPath(String),
|
||||
/// 目标不存在。
|
||||
NotFound,
|
||||
/// 期望文件但目标是目录（或反之）。
|
||||
WrongType {
|
||||
expected: &'static str,
|
||||
actual: &'static str,
|
||||
},
|
||||
/// 内容超过大小上限。
|
||||
TooLarge {
|
||||
limit: u64,
|
||||
actual: u64,
|
||||
},
|
||||
/// 底层 IO 错误。
|
||||
Io(std::io::Error),
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for FileError {
|
||||
fn from(e: std::io::Error) -> Self {
|
||||
if e.kind() == std::io::ErrorKind::NotFound {
|
||||
FileError::NotFound
|
||||
} else {
|
||||
FileError::Io(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_within(path: &Path, base: &Path) -> Result<(), FileError> {
|
||||
if path.starts_with(base) {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(FileError::InvalidPath("路径逃逸出沙箱根目录".into()))
|
||||
}
|
||||
}
|
||||
|
||||
/// 词汇层面判断 `rel` 规范化后不向上逃逸：`..` 不可使深度为负。
|
||||
/// 用于 write 前预检（write 不 canonicalize 目标，需先确保 `create_dir_all` 不在 base 外建目录）。
|
||||
fn is_lexically_within(rel: &str) -> bool {
|
||||
let mut depth: i32 = 0;
|
||||
for comp in Path::new(rel).components() {
|
||||
match comp {
|
||||
std::path::Component::Normal(_) => depth += 1,
|
||||
std::path::Component::ParentDir => {
|
||||
depth -= 1;
|
||||
if depth < 0 {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
std::path::Component::CurDir => {}
|
||||
std::path::Component::RootDir | std::path::Component::Prefix(_) => return false,
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// 解析相对路径为 base 内的绝对路径，拒绝逃逸。
|
||||
///
|
||||
/// - `exist_required=true`（read/list/stat）：`canonicalize` 已存在路径后校验 `starts_with`。
|
||||
/// - `exist_required=false`（write）：词汇预检防 `..` 逃逸，再 `canonicalize` 父目录兜底符号链接。
|
||||
pub async fn resolve_within_base(
|
||||
base: &Path,
|
||||
base_canon: &Path,
|
||||
rel: &str,
|
||||
exist_required: bool,
|
||||
) -> Result<PathBuf, FileError> {
|
||||
if rel.is_empty() || Path::new(rel).is_absolute() {
|
||||
return Err(FileError::InvalidPath("相对路径非法（空或绝对路径）".into()));
|
||||
}
|
||||
// 词汇预检：.. 不可使深度为负。read/list/stat 也先判，使逃逸路径报 InvalidPath 而非 NotFound；
|
||||
// write 则避免 create_dir_all 在 base 外产生副作用。
|
||||
if !is_lexically_within(rel) {
|
||||
return Err(FileError::InvalidPath("路径逃逸出沙箱根目录".into()));
|
||||
}
|
||||
let candidate = base.join(rel);
|
||||
|
||||
if exist_required {
|
||||
let canon = tokio::fs::canonicalize(&candidate).await?;
|
||||
ensure_within(&canon, base_canon)?;
|
||||
Ok(canon)
|
||||
} else {
|
||||
// write：目标可不存在。词汇预检已确保 .. 不逃逸（见上），此处 canonicalize 父目录
|
||||
// 兜底符号链接跳出，再 join 文件名。
|
||||
let parent = candidate.parent().unwrap_or(base);
|
||||
if parent != base && !parent.as_os_str().is_empty() {
|
||||
tokio::fs::create_dir_all(parent).await?;
|
||||
}
|
||||
let parent_canon = tokio::fs::canonicalize(parent).await?;
|
||||
ensure_within(&parent_canon, base_canon)?;
|
||||
let name = candidate
|
||||
.file_name()
|
||||
.ok_or_else(|| FileError::InvalidPath("路径缺少文件名".into()))?;
|
||||
Ok(parent_canon.join(name))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct FileReadOutcome {
|
||||
/// 文件内容（base64url-no-pad，二进制安全；调用方按需解码为 UTF-8）。
|
||||
pub content_b64: String,
|
||||
/// 文件总字节数。
|
||||
pub size: u64,
|
||||
/// 本次读取起始字节偏移。
|
||||
pub offset: u64,
|
||||
/// 本次返回字节数。
|
||||
pub returned: u64,
|
||||
/// 是否因超过 max_bytes 而截断。
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// 读取文件（二进制安全，字节级 offset/limit 分页，超 max_bytes 截断并标记）。
|
||||
///
|
||||
/// 实现走流式读取（`File` + `AsyncReadExt`），只读 `min(请求范围, max_bytes)` 字节，
|
||||
/// **不把全文读进内存**——文件远大于 max_bytes 时内存峰值受 max_bytes 钳制，
|
||||
/// 而非文件大小。`size` 字段仍报全文大小（取自 `metadata`），与原行为一致。
|
||||
pub async fn file_read(
|
||||
base: &Path,
|
||||
base_canon: &Path,
|
||||
rel: &str,
|
||||
offset: Option<u64>,
|
||||
limit: Option<u64>,
|
||||
max_bytes: usize,
|
||||
) -> Result<FileReadOutcome, FileError> {
|
||||
let path = resolve_within_base(base, base_canon, rel, true).await?;
|
||||
let meta = tokio::fs::metadata(&path).await?;
|
||||
if !meta.is_file() {
|
||||
return Err(FileError::WrongType {
|
||||
expected: "file",
|
||||
actual: if meta.is_dir() { "dir" } else { "other" },
|
||||
});
|
||||
}
|
||||
let total = meta.len() as usize;
|
||||
|
||||
// 计算请求范围 [start, end)（与原逻辑一致：offset/limit 经 total 钳制）。
|
||||
let start = (offset.unwrap_or(0) as usize).min(total);
|
||||
let end = limit
|
||||
.map(|l| (start + (l as usize)).min(total))
|
||||
.unwrap_or(total);
|
||||
let desired = end - start; // 期望读取的字节数
|
||||
let truncated = desired > max_bytes;
|
||||
let read_len = desired.min(max_bytes);
|
||||
|
||||
// 流式读取：seek 到 start，只读 read_len 字节（不读全文）。
|
||||
let mut file = tokio::fs::File::open(&path).await?;
|
||||
if start > 0 {
|
||||
use tokio::io::AsyncSeekExt;
|
||||
file.seek(std::io::SeekFrom::Start(start as u64)).await?;
|
||||
}
|
||||
let mut buf = vec![0u8; read_len];
|
||||
let mut filled = 0usize;
|
||||
use tokio::io::AsyncReadExt;
|
||||
while filled < read_len {
|
||||
let n = file.read(&mut buf[filled..]).await?;
|
||||
if n == 0 {
|
||||
break; // EOF（理论上 start+read_len <= total，不会提前 EOF，但防御性处理）
|
||||
}
|
||||
filled += n;
|
||||
}
|
||||
buf.truncate(filled);
|
||||
|
||||
Ok(FileReadOutcome {
|
||||
content_b64: BASE64_URL_SAFE_NO_PAD.encode(&buf),
|
||||
size: total as u64,
|
||||
offset: start as u64,
|
||||
returned: filled as u64,
|
||||
truncated,
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct FileWriteOutcome {
|
||||
pub bytes_written: u64,
|
||||
}
|
||||
|
||||
/// 写入文件（覆盖）。content 为原始字节。超 max_bytes 拒绝。
|
||||
pub async fn file_write(
|
||||
base: &Path,
|
||||
base_canon: &Path,
|
||||
rel: &str,
|
||||
content: &[u8],
|
||||
max_bytes: usize,
|
||||
) -> Result<FileWriteOutcome, FileError> {
|
||||
if content.len() > max_bytes {
|
||||
return Err(FileError::TooLarge {
|
||||
limit: max_bytes as u64,
|
||||
actual: content.len() as u64,
|
||||
});
|
||||
}
|
||||
let path = resolve_within_base(base, base_canon, rel, false).await?;
|
||||
tokio::fs::write(&path, content).await?;
|
||||
Ok(FileWriteOutcome {
|
||||
bytes_written: content.len() as u64,
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct DirEntry {
|
||||
pub name: String,
|
||||
pub kind: String, // "file" | "dir" | "other"
|
||||
pub size: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct FileListOutcome {
|
||||
pub entries: Vec<DirEntry>,
|
||||
}
|
||||
|
||||
/// 列出目录条目（按名字排序，确定性）。
|
||||
pub async fn file_list(
|
||||
base: &Path,
|
||||
base_canon: &Path,
|
||||
rel: &str,
|
||||
) -> Result<FileListOutcome, FileError> {
|
||||
let path = resolve_within_base(base, base_canon, rel, true).await?;
|
||||
let meta = tokio::fs::metadata(&path).await?;
|
||||
if !meta.is_dir() {
|
||||
return Err(FileError::WrongType {
|
||||
expected: "dir",
|
||||
actual: if meta.is_file() { "file" } else { "other" },
|
||||
});
|
||||
}
|
||||
let mut rd = tokio::fs::read_dir(&path).await?;
|
||||
let mut entries = Vec::new();
|
||||
while let Some(entry) = rd.next_entry().await? {
|
||||
let name = entry.file_name().to_string_lossy().into_owned();
|
||||
let ft = entry.file_type().await?;
|
||||
let kind = if ft.is_dir() {
|
||||
"dir"
|
||||
} else if ft.is_file() {
|
||||
"file"
|
||||
} else {
|
||||
"other"
|
||||
};
|
||||
let size = entry.metadata().await?.len();
|
||||
entries.push(DirEntry {
|
||||
name,
|
||||
kind: kind.into(),
|
||||
size,
|
||||
});
|
||||
}
|
||||
entries.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
Ok(FileListOutcome { entries })
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct FileStatOutcome {
|
||||
pub kind: String,
|
||||
pub size: u64,
|
||||
/// 最后修改时间（UNIX 秒）。
|
||||
pub modified_unix: u64,
|
||||
}
|
||||
|
||||
/// 文件元信息（类型 / 大小 / 修改时间）。
|
||||
pub async fn file_stat(
|
||||
base: &Path,
|
||||
base_canon: &Path,
|
||||
rel: &str,
|
||||
) -> Result<FileStatOutcome, FileError> {
|
||||
let path = resolve_within_base(base, base_canon, rel, true).await?;
|
||||
let meta = tokio::fs::metadata(&path).await?;
|
||||
let kind = if meta.is_dir() {
|
||||
"dir"
|
||||
} else if meta.is_file() {
|
||||
"file"
|
||||
} else {
|
||||
"other"
|
||||
};
|
||||
let modified_unix = meta
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|t| t.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(0);
|
||||
Ok(FileStatOutcome {
|
||||
kind: kind.into(),
|
||||
size: meta.len(),
|
||||
modified_unix,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::TempDir;
|
||||
|
||||
struct Sandbox {
|
||||
tmp: TempDir,
|
||||
base_canon: PathBuf,
|
||||
}
|
||||
|
||||
impl Sandbox {
|
||||
fn new() -> Self {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let base_canon = tmp.path().canonicalize().unwrap();
|
||||
Self { tmp, base_canon }
|
||||
}
|
||||
fn base(&self) -> &Path {
|
||||
self.tmp.path()
|
||||
}
|
||||
fn base_canon(&self) -> &Path {
|
||||
&self.base_canon
|
||||
}
|
||||
fn write(&self, rel: &str, content: &[u8]) {
|
||||
let p = self.base().join(rel);
|
||||
if let Some(parent) = p.parent() {
|
||||
std::fs::create_dir_all(parent).unwrap();
|
||||
}
|
||||
std::fs::write(p, content).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_write_roundtrip_binary_safe() {
|
||||
let sb = Sandbox::new();
|
||||
let content = b"hello \xe4\xb8\x96\xe7\x95\x8c"; // 二进制安全（UTF-8 "世界"）
|
||||
file_write(sb.base(), sb.base_canon(), "a/b.txt", content, 1 << 20)
|
||||
.await
|
||||
.unwrap();
|
||||
let out = file_read(sb.base(), sb.base_canon(), "a/b.txt", None, None, 1 << 20)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(out.size, content.len() as u64);
|
||||
let decoded = BASE64_URL_SAFE_NO_PAD.decode(&out.content_b64).unwrap();
|
||||
assert_eq!(decoded, content);
|
||||
assert!(!out.truncated);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_rejects_path_traversal() {
|
||||
let sb = Sandbox::new();
|
||||
sb.write("safe.txt", b"x");
|
||||
let err = file_read(sb.base(), sb.base_canon(), "../../etc/passwd", None, None, 1 << 20).await;
|
||||
assert!(matches!(err, Err(FileError::InvalidPath(_))));
|
||||
let err2 = file_read(sb.base(), sb.base_canon(), "/etc/passwd", None, None, 1 << 20).await;
|
||||
assert!(matches!(err2, Err(FileError::InvalidPath(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_rejects_path_traversal() {
|
||||
let sb = Sandbox::new();
|
||||
let err = file_write(sb.base(), sb.base_canon(), "../../../escape.txt", b"x", 1 << 20).await;
|
||||
assert!(matches!(err, Err(FileError::InvalidPath(_))));
|
||||
// 绝对路径
|
||||
let err2 = file_write(sb.base(), sb.base_canon(), "/tmp/escape.txt", b"x", 1 << 20).await;
|
||||
assert!(matches!(err2, Err(FileError::InvalidPath(_))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_allows_dotdot_within_base() {
|
||||
// sub/../valid 词汇上在 base 内（抵消后 = valid），应允许
|
||||
let sb = Sandbox::new();
|
||||
std::fs::create_dir_all(sb.base().join("sub")).unwrap();
|
||||
file_write(sb.base(), sb.base_canon(), "sub/../valid.txt", b"x", 1 << 20)
|
||||
.await
|
||||
.unwrap();
|
||||
let out = file_read(sb.base(), sb.base_canon(), "valid.txt", None, None, 1 << 20)
|
||||
.await
|
||||
.unwrap();
|
||||
let decoded = BASE64_URL_SAFE_NO_PAD.decode(&out.content_b64).unwrap();
|
||||
assert_eq!(decoded, b"x");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_offset_limit_pagination() {
|
||||
let sb = Sandbox::new();
|
||||
sb.write("f.txt", b"0123456789");
|
||||
let out = file_read(sb.base(), sb.base_canon(), "f.txt", Some(3), Some(4), 1 << 20)
|
||||
.await
|
||||
.unwrap();
|
||||
let decoded = BASE64_URL_SAFE_NO_PAD.decode(&out.content_b64).unwrap();
|
||||
assert_eq!(decoded, b"3456");
|
||||
assert_eq!(out.offset, 3);
|
||||
assert_eq!(out.returned, 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_truncates_at_max_bytes() {
|
||||
let sb = Sandbox::new();
|
||||
sb.write("big.txt", b"0123456789ABCDEF"); // 16 bytes
|
||||
let out = file_read(sb.base(), sb.base_canon(), "big.txt", None, None, 8)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.truncated);
|
||||
assert_eq!(out.returned, 8);
|
||||
assert_eq!(out.size, 16);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_large_file_truncates_to_max_bytes_content() {
|
||||
// 锁定契约（重构保护）：文件远大于 max_bytes 且 limit=None 时，
|
||||
// size 报全文大小、returned=max_bytes、truncated=true、内容为前 max_bytes 字节。
|
||||
let sb = Sandbox::new();
|
||||
let mut content = Vec::with_capacity(65536);
|
||||
for i in 0..65536u32 {
|
||||
content.push((i % 251) as u8); // 非平凡内容，验证切片正确性
|
||||
}
|
||||
sb.write("large.bin", &content);
|
||||
let out = file_read(sb.base(), sb.base_canon(), "large.bin", None, None, 1024)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(out.truncated);
|
||||
assert_eq!(out.size, 65536);
|
||||
assert_eq!(out.returned, 1024);
|
||||
let decoded = BASE64_URL_SAFE_NO_PAD.decode(&out.content_b64).unwrap();
|
||||
assert_eq!(decoded, &content[..1024]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn write_rejects_too_large() {
|
||||
let sb = Sandbox::new();
|
||||
let err = file_write(sb.base(), sb.base_canon(), "big.txt", &[0u8; 100], 10).await;
|
||||
assert!(matches!(err, Err(FileError::TooLarge { limit: 10, actual: 100 })));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_returns_sorted_entries() {
|
||||
let sb = Sandbox::new();
|
||||
sb.write("z.txt", b"1");
|
||||
sb.write("a.txt", b"1");
|
||||
std::fs::create_dir_all(sb.base().join("mdir")).unwrap();
|
||||
let out = file_list(sb.base(), sb.base_canon(), ".").await.unwrap();
|
||||
let names: Vec<_> = out.entries.iter().map(|e| e.name.as_str()).collect();
|
||||
assert_eq!(names, vec!["a.txt", "mdir", "z.txt"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stat_returns_kind_and_size() {
|
||||
let sb = Sandbox::new();
|
||||
sb.write("f.txt", b"hello");
|
||||
let out = file_stat(sb.base(), sb.base_canon(), "f.txt").await.unwrap();
|
||||
assert_eq!(out.kind, "file");
|
||||
assert_eq!(out.size, 5);
|
||||
assert!(out.modified_unix > 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_missing_file_not_found() {
|
||||
let sb = Sandbox::new();
|
||||
let err = file_read(sb.base(), sb.base_canon(), "nope.txt", None, None, 1 << 20).await;
|
||||
assert!(matches!(err, Err(FileError::NotFound)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_on_dir_wrong_type() {
|
||||
let sb = Sandbox::new();
|
||||
std::fs::create_dir_all(sb.base().join("d")).unwrap();
|
||||
let err = file_read(sb.base(), sb.base_canon(), "d", None, None, 1 << 20).await;
|
||||
assert!(matches!(err, Err(FileError::WrongType { .. })));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,221 @@
|
||||
//! 魂镜安全沙箱服务入口。
|
||||
//!
|
||||
//! 定位见 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!("预热池容器已全部回收");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
//! Prometheus 指标：执行计数（按结果）、耗时分布、并发占用、流式计数、授权 secret 分布。
|
||||
//!
|
||||
//! 指标在 executor 内部记录，与执行生命周期绑定；``/metrics`` 端点调用 ``gather`` 输出文本格式。
|
||||
|
||||
use prometheus::{
|
||||
Encoder, Gauge, Histogram, HistogramOpts, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec,
|
||||
Opts, Registry, TextEncoder,
|
||||
};
|
||||
|
||||
/// 沙箱可观测指标集合。
|
||||
pub struct Metrics {
|
||||
exec_total: IntCounterVec,
|
||||
exec_duration: Histogram,
|
||||
concurrent: Gauge,
|
||||
stream_total: IntCounterVec,
|
||||
auth_secret_used: IntCounterVec,
|
||||
// 预热池与容器指标
|
||||
warm_pool_idle: IntGaugeVec,
|
||||
warm_pool_running: IntGaugeVec,
|
||||
warm_pool_global_total: IntGauge,
|
||||
warm_pool_waiters: IntGaugeVec,
|
||||
container_cold_start: Histogram,
|
||||
pool_hit_total: IntCounterVec,
|
||||
pool_miss_total: IntCounterVec,
|
||||
container_exec_duration: HistogramVec,
|
||||
container_downgrade_total: IntCounterVec,
|
||||
registry: Registry,
|
||||
}
|
||||
|
||||
impl Metrics {
|
||||
/// 构造指标集合并注册到独立 Registry；``max_concurrent`` 写入 ``sandbox_concurrent_limit``，
|
||||
/// ``auth_enabled`` 写入 ``sandbox_auth_enabled``（0/1），便于监控告警"生产环境鉴权被关"。
|
||||
pub fn new(max_concurrent: usize, auth_enabled: bool) -> Self {
|
||||
let registry = Registry::new();
|
||||
|
||||
let exec_total = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_exec_total",
|
||||
"同步执行次数（outcome=completed/timeout/cancelled/failed）",
|
||||
),
|
||||
&["outcome"],
|
||||
)
|
||||
.expect("exec_total 指标构造");
|
||||
let exec_duration = Histogram::with_opts(
|
||||
HistogramOpts::new("sandbox_exec_duration_seconds", "同步执行耗时分布（秒）")
|
||||
// 沙箱命令多为秒级；超时上限 30s
|
||||
.buckets(vec![0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0]),
|
||||
)
|
||||
.expect("exec_duration 指标构造");
|
||||
let concurrent = Gauge::new(
|
||||
"sandbox_concurrent",
|
||||
"当前并发执行数（含同步与流式，上限见 sandbox_concurrent_limit）",
|
||||
)
|
||||
.expect("concurrent 指标构造");
|
||||
let concurrent_limit = Gauge::new(
|
||||
"sandbox_concurrent_limit",
|
||||
"并发执行上限（SANDBOX_MAX_CONCURRENT）",
|
||||
)
|
||||
.expect("concurrent_limit 指标构造");
|
||||
concurrent_limit.set(max_concurrent as f64);
|
||||
let stream_total = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_stream_total",
|
||||
"流式执行次数（outcome=completed/timeout/cancelled/failed）",
|
||||
),
|
||||
&["outcome"],
|
||||
)
|
||||
.expect("stream_total 指标构造");
|
||||
let auth_secret_used = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_auth_secret_used",
|
||||
"授权凭证验签命中 secret 分布（slot=primary|previous），用于判断轮换何时完成",
|
||||
),
|
||||
&["slot"],
|
||||
)
|
||||
.expect("auth_secret_used 指标构造");
|
||||
let auth_enabled_gauge = Gauge::new(
|
||||
"sandbox_auth_enabled",
|
||||
"鉴权是否开启（1=SANDBOX_AUTH_ENABLED true，0=false）；0 时应触发告警",
|
||||
)
|
||||
.expect("auth_enabled 指标构造");
|
||||
auth_enabled_gauge.set(if auth_enabled { 1.0 } else { 0.0 });
|
||||
|
||||
// ==== 预热池指标 ====
|
||||
let warm_pool_idle = IntGaugeVec::new(
|
||||
Opts::new("sandbox_warm_pool_idle", "各语言空闲容器数"),
|
||||
&["language"],
|
||||
)
|
||||
.expect("warm_pool_idle 指标构造");
|
||||
let warm_pool_running = IntGaugeVec::new(
|
||||
Opts::new("sandbox_warm_pool_running", "各语言运行中容器数"),
|
||||
&["language"],
|
||||
)
|
||||
.expect("warm_pool_running 指标构造");
|
||||
let warm_pool_global_total = IntGauge::new(
|
||||
"sandbox_warm_pool_global_total",
|
||||
"全局容器总数（所有语言）",
|
||||
)
|
||||
.expect("warm_pool_global_total 指标构造");
|
||||
let warm_pool_global_max = IntGauge::new(
|
||||
"sandbox_warm_pool_global_max",
|
||||
"全局容器上限（SANDBOX_WARM_POOL_GLOBAL_MAX）",
|
||||
)
|
||||
.expect("warm_pool_global_max 指标构造");
|
||||
warm_pool_global_max.set(max_concurrent as i64); // 复用 max_concurrent 的占位，实际由 warm_pool 设置
|
||||
let warm_pool_waiters = IntGaugeVec::new(
|
||||
Opts::new("sandbox_warm_pool_waiters", "排队等待容器数"),
|
||||
&["language"],
|
||||
)
|
||||
.expect("warm_pool_waiters 指标构造");
|
||||
let container_cold_start = Histogram::with_opts(
|
||||
HistogramOpts::new("sandbox_container_cold_start_seconds", "容器冷启动耗时分布（秒）")
|
||||
.buckets(vec![0.5, 1.0, 2.0, 3.0, 5.0, 10.0, 30.0]),
|
||||
)
|
||||
.expect("container_cold_start 指标构造");
|
||||
let pool_hit_total = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_container_pool_hit_total",
|
||||
"预热命中计数（按语言）",
|
||||
),
|
||||
&["language"],
|
||||
)
|
||||
.expect("pool_hit_total 指标构造");
|
||||
let pool_miss_total = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_container_pool_miss_total",
|
||||
"冷启动计数（按语言）",
|
||||
),
|
||||
&["language"],
|
||||
)
|
||||
.expect("pool_miss_total 指标构造");
|
||||
let container_exec_duration = HistogramVec::new(
|
||||
HistogramOpts::new(
|
||||
"sandbox_container_exec_duration_seconds",
|
||||
"容器内执行耗时分布（秒）",
|
||||
)
|
||||
.buckets(vec![0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0]),
|
||||
&["language"],
|
||||
)
|
||||
.expect("container_exec_duration 指标构造");
|
||||
let container_downgrade_total = IntCounterVec::new(
|
||||
Opts::new(
|
||||
"sandbox_container_downgrade_total",
|
||||
"容器降级为进程计数（按原因）",
|
||||
),
|
||||
&["reason"],
|
||||
)
|
||||
.expect("container_downgrade_total 指标构造");
|
||||
|
||||
registry
|
||||
.register(Box::new(exec_total.clone()))
|
||||
.expect("register exec_total");
|
||||
registry
|
||||
.register(Box::new(exec_duration.clone()))
|
||||
.expect("register exec_duration");
|
||||
registry
|
||||
.register(Box::new(concurrent.clone()))
|
||||
.expect("register concurrent");
|
||||
registry
|
||||
.register(Box::new(concurrent_limit))
|
||||
.expect("register concurrent_limit");
|
||||
registry
|
||||
.register(Box::new(stream_total.clone()))
|
||||
.expect("register stream_total");
|
||||
registry
|
||||
.register(Box::new(auth_secret_used.clone()))
|
||||
.expect("register auth_secret_used");
|
||||
registry
|
||||
.register(Box::new(auth_enabled_gauge))
|
||||
.expect("register auth_enabled_gauge");
|
||||
registry
|
||||
.register(Box::new(warm_pool_idle.clone()))
|
||||
.expect("register warm_pool_idle");
|
||||
registry
|
||||
.register(Box::new(warm_pool_running.clone()))
|
||||
.expect("register warm_pool_running");
|
||||
registry
|
||||
.register(Box::new(warm_pool_global_total.clone()))
|
||||
.expect("register warm_pool_global_total");
|
||||
registry
|
||||
.register(Box::new(warm_pool_global_max))
|
||||
.expect("register warm_pool_global_max");
|
||||
registry
|
||||
.register(Box::new(warm_pool_waiters.clone()))
|
||||
.expect("register warm_pool_waiters");
|
||||
registry
|
||||
.register(Box::new(container_cold_start.clone()))
|
||||
.expect("register container_cold_start");
|
||||
registry
|
||||
.register(Box::new(pool_hit_total.clone()))
|
||||
.expect("register pool_hit_total");
|
||||
registry
|
||||
.register(Box::new(pool_miss_total.clone()))
|
||||
.expect("register pool_miss_total");
|
||||
registry
|
||||
.register(Box::new(container_exec_duration.clone()))
|
||||
.expect("register container_exec_duration");
|
||||
registry
|
||||
.register(Box::new(container_downgrade_total.clone()))
|
||||
.expect("register container_downgrade_total");
|
||||
|
||||
Self {
|
||||
exec_total,
|
||||
exec_duration,
|
||||
concurrent,
|
||||
stream_total,
|
||||
auth_secret_used,
|
||||
warm_pool_idle,
|
||||
warm_pool_running,
|
||||
warm_pool_global_total,
|
||||
warm_pool_waiters,
|
||||
container_cold_start,
|
||||
pool_hit_total,
|
||||
pool_miss_total,
|
||||
container_exec_duration,
|
||||
container_downgrade_total,
|
||||
registry,
|
||||
}
|
||||
}
|
||||
|
||||
/// 记录一次同步执行：结果分类 + 耗时。
|
||||
pub fn record_exec(&self, outcome: &str, duration_ms: u128) {
|
||||
self.exec_total.with_label_values(&[outcome]).inc();
|
||||
self.exec_duration.observe(duration_ms as f64 / 1000.0);
|
||||
}
|
||||
|
||||
/// 记录一次流式执行：结果分类。
|
||||
pub fn record_stream(&self, outcome: &str) {
|
||||
self.stream_total.with_label_values(&[outcome]).inc();
|
||||
}
|
||||
|
||||
/// 记录授权验签命中的 secret：slot 0 为主（primary），>0 为旧（previous）。
|
||||
pub fn record_auth_secret(&self, slot: usize) {
|
||||
let label = if slot == 0 { "primary" } else { "previous" };
|
||||
self.auth_secret_used.with_label_values(&[label]).inc();
|
||||
}
|
||||
|
||||
/// 当前并发执行数（用于 /healthz 暴露系统负载）。
|
||||
pub fn concurrent_value(&self) -> i64 {
|
||||
self.concurrent.get() as i64
|
||||
}
|
||||
|
||||
pub fn concurrent_inc(&self) {
|
||||
self.concurrent.inc();
|
||||
}
|
||||
|
||||
pub fn concurrent_dec(&self) {
|
||||
self.concurrent.dec();
|
||||
}
|
||||
|
||||
// ==== 预热池与容器指标 ====
|
||||
|
||||
pub fn set_warm_pool_idle(&self, language: &str, count: i64) {
|
||||
self.warm_pool_idle
|
||||
.with_label_values(&[language])
|
||||
.set(count);
|
||||
}
|
||||
|
||||
pub fn set_warm_pool_running(&self, language: &str, count: i64) {
|
||||
self.warm_pool_running
|
||||
.with_label_values(&[language])
|
||||
.set(count);
|
||||
}
|
||||
|
||||
pub fn set_warm_pool_global_total(&self, count: i64) {
|
||||
self.warm_pool_global_total.set(count);
|
||||
}
|
||||
|
||||
pub fn set_warm_pool_waiters(&self, language: &str, count: i64) {
|
||||
self.warm_pool_waiters
|
||||
.with_label_values(&[language])
|
||||
.set(count);
|
||||
}
|
||||
|
||||
pub fn record_container_cold_start(&self, duration_ms: u128) {
|
||||
self.container_cold_start
|
||||
.observe(duration_ms as f64 / 1000.0);
|
||||
}
|
||||
|
||||
pub fn record_pool_hit(&self, language: &str) {
|
||||
self.pool_hit_total
|
||||
.with_label_values(&[language])
|
||||
.inc();
|
||||
}
|
||||
|
||||
pub fn record_pool_miss(&self, language: &str) {
|
||||
self.pool_miss_total
|
||||
.with_label_values(&[language])
|
||||
.inc();
|
||||
}
|
||||
|
||||
pub fn record_container_exec(&self, language: &str, duration_ms: u128) {
|
||||
self.container_exec_duration
|
||||
.with_label_values(&[language])
|
||||
.observe(duration_ms as f64 / 1000.0);
|
||||
}
|
||||
|
||||
pub fn record_downgrade(&self, reason: &str) {
|
||||
self.container_downgrade_total
|
||||
.with_label_values(&[reason])
|
||||
.inc();
|
||||
}
|
||||
|
||||
pub fn warm_pool_metrics_ref(
|
||||
&self,
|
||||
) -> (
|
||||
&IntGaugeVec,
|
||||
&IntGaugeVec,
|
||||
&IntGauge,
|
||||
&IntGaugeVec,
|
||||
&Histogram,
|
||||
&IntCounterVec,
|
||||
&IntCounterVec,
|
||||
&HistogramVec,
|
||||
&IntCounterVec,
|
||||
) {
|
||||
(
|
||||
&self.warm_pool_idle,
|
||||
&self.warm_pool_running,
|
||||
&self.warm_pool_global_total,
|
||||
&self.warm_pool_waiters,
|
||||
&self.container_cold_start,
|
||||
&self.pool_hit_total,
|
||||
&self.pool_miss_total,
|
||||
&self.container_exec_duration,
|
||||
&self.container_downgrade_total,
|
||||
)
|
||||
}
|
||||
|
||||
/// 输出 Prometheus 文本格式（供 /metrics 端点）。
|
||||
pub fn gather(&self) -> String {
|
||||
let encoder = TextEncoder::new();
|
||||
let mut buf = Vec::new();
|
||||
encoder
|
||||
.encode(&self.registry.gather(), &mut buf)
|
||||
.expect("编码指标不会失败");
|
||||
String::from_utf8(buf).expect("指标文本为 UTF-8")
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn auth_enabled_gauge_reflects_config() {
|
||||
// 鉴权开启：sandbox_auth_enabled 应为 1，便于监控告警"生产环境鉴权被关"。
|
||||
let m_on = Metrics::new(8, true);
|
||||
let text_on = m_on.gather();
|
||||
assert!(
|
||||
text_on.contains("sandbox_auth_enabled"),
|
||||
"应暴露 sandbox_auth_enabled 指标"
|
||||
);
|
||||
assert!(
|
||||
text_on.contains("sandbox_auth_enabled 1"),
|
||||
"auth_enabled=true 时指标值应为 1，实际：{}",
|
||||
text_on.lines().find(|l| l.contains("sandbox_auth_enabled")).unwrap_or("")
|
||||
);
|
||||
|
||||
// 鉴权关闭：指标值应为 0，监控可据此告警。
|
||||
let m_off = Metrics::new(8, false);
|
||||
let text_off = m_off.gather();
|
||||
assert!(
|
||||
text_off.contains("sandbox_auth_enabled 0"),
|
||||
"auth_enabled=false 时指标值应为 0，实际：{}",
|
||||
text_off.lines().find(|l| l.contains("sandbox_auth_enabled")).unwrap_or("")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
//! 进程级资源限额：在子进程 fork 后、exec 前通过 setrlimit 限制
|
||||
//! CPU / 内存 / 文件 / 进程数 / 打开文件数，作为容器级隔离之外的进程级纵深防御。
|
||||
//!
|
||||
//! 配额分层（自上而下 clamp，越宽越受上层钳制）：
|
||||
//! - 全局天花板 `*_MAX`：env 硬上限，运维设的安全阀，任何执行不可逾越。
|
||||
//! - 凭证授予 granted：司命按 tool policy 写入 `rlimit_overrides`（特定工具可放宽到此）。
|
||||
//! - 请求值 request：调用方传 `rlimits`，只能收紧、不能放宽。
|
||||
//! - 默认值 default：env `SANDBOX_RLIMIT_*`（普通工具的默认配额）。
|
||||
//!
|
||||
//! 生效规则：`baseline = granted.unwrap_or(default)`；
|
||||
//! `effective = min(request.unwrap_or(baseline), baseline)`；再 clamp 到 `*_MAX`。
|
||||
//! 沙箱不做策略判定——granted 由经司命审核的签发端决定、随 HMAC 凭证下发，此处只 clamp。
|
||||
//! 设计依据见 docs/superpowers/specs/2026-06-28-sandbox-batch-a-safety-hardening-design.md。
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
/// 资源限额维度。`as_str` 返回跨语言一致的字典 key（与 backend `auth.py` 对齐）。
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub enum RlimitKind {
|
||||
/// 单文件最大写入字节数（RLIMIT_FSIZE）。
|
||||
Fsize,
|
||||
/// 虚拟内存最大字节数（RLIMIT_AS）。
|
||||
As,
|
||||
/// 最大 CPU 秒数（RLIMIT_CPU）。
|
||||
Cpu,
|
||||
/// real uid 最大进程数（RLIMIT_NPROC）。
|
||||
Nproc,
|
||||
/// 最大打开文件描述符数（RLIMIT_NOFILE）。
|
||||
Nofile,
|
||||
}
|
||||
|
||||
impl RlimitKind {
|
||||
pub const ALL: [RlimitKind; 5] = [
|
||||
RlimitKind::Fsize,
|
||||
RlimitKind::As,
|
||||
RlimitKind::Cpu,
|
||||
RlimitKind::Nproc,
|
||||
RlimitKind::Nofile,
|
||||
];
|
||||
|
||||
/// 跨语言一致的 key 字符串（backend `rlimit_overrides` 字典 key 与此对齐）。
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
RlimitKind::Fsize => "fsize",
|
||||
RlimitKind::As => "as",
|
||||
RlimitKind::Cpu => "cpu",
|
||||
RlimitKind::Nproc => "nproc",
|
||||
RlimitKind::Nofile => "nofile",
|
||||
}
|
||||
}
|
||||
|
||||
/// 从跨语言 key 字符串解析；非法返回 None（调用方忽略未知 key，不报错）。
|
||||
pub fn from_key(s: &str) -> Option<RlimitKind> {
|
||||
match s {
|
||||
"fsize" => Some(RlimitKind::Fsize),
|
||||
"as" => Some(RlimitKind::As),
|
||||
"cpu" => Some(RlimitKind::Cpu),
|
||||
"nproc" => Some(RlimitKind::Nproc),
|
||||
"nofile" => Some(RlimitKind::Nofile),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 某维度的配额边界：默认值（普通工具用）+ 全局硬天花板（granted / request 不可逾越）。
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct RlimitBound {
|
||||
pub default: u64,
|
||||
pub max: u64,
|
||||
}
|
||||
|
||||
/// 全部维度的配额配置（边界 + 显式禁用清单），由 `Config` 从 env 解析后持有。
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RlimitsConfig {
|
||||
bounds: HashMap<RlimitKind, RlimitBound>,
|
||||
disabled: Vec<RlimitKind>,
|
||||
}
|
||||
|
||||
impl RlimitsConfig {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// 设置某维度的边界。`max` 是硬天花板；若配置 `max < default`，以 `max` 收紧默认（保守）。
|
||||
pub fn set(&mut self, kind: RlimitKind, default: u64, max: u64) -> &mut Self {
|
||||
let effective_default = default.min(max);
|
||||
self.bounds.insert(
|
||||
kind,
|
||||
RlimitBound {
|
||||
default: effective_default,
|
||||
max,
|
||||
},
|
||||
);
|
||||
self
|
||||
}
|
||||
|
||||
/// 显式禁用某维度（不对其 setrlimit）。幂等。
|
||||
pub fn disable(&mut self, kind: RlimitKind) -> &mut Self {
|
||||
if !self.disabled.contains(&kind) {
|
||||
self.disabled.push(kind);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
pub fn bound(&self, kind: RlimitKind) -> Option<RlimitBound> {
|
||||
self.bounds.get(&kind).copied()
|
||||
}
|
||||
|
||||
pub fn is_disabled(&self, kind: RlimitKind) -> bool {
|
||||
self.disabled.contains(&kind)
|
||||
}
|
||||
}
|
||||
|
||||
/// 一个维度最终生效的限额（setrlimit 时软 = 硬）。
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct EffectiveRlimit {
|
||||
pub kind: RlimitKind,
|
||||
pub value: u64,
|
||||
}
|
||||
|
||||
/// 计算各维度最终生效限额。
|
||||
///
|
||||
/// - `granted`：来自凭证的 `rlimit_overrides`（司命授予该上下文的上限）。
|
||||
/// - `request`：来自请求体的 `rlimits`（调用方主动收紧）。
|
||||
///
|
||||
/// 两者都「只能收紧」：baseline 取 granted 或 default；request 在 baseline 之上取较小；
|
||||
/// 最后统一 clamp 到全局天花板 `*_MAX`。被禁用的维度不产出。
|
||||
pub fn compute_effective(
|
||||
config: &RlimitsConfig,
|
||||
granted: Option<&BTreeMap<String, u64>>,
|
||||
request: Option<&BTreeMap<String, u64>>,
|
||||
) -> Vec<EffectiveRlimit> {
|
||||
let mut out = Vec::new();
|
||||
for kind in RlimitKind::ALL {
|
||||
if config.is_disabled(kind) {
|
||||
continue;
|
||||
}
|
||||
let Some(bound) = config.bound(kind) else {
|
||||
continue;
|
||||
};
|
||||
let granted_val = granted.and_then(|g| g.get(kind.as_str()).copied());
|
||||
let baseline = granted_val.unwrap_or(bound.default);
|
||||
let mut effective = request
|
||||
.and_then(|r| r.get(kind.as_str()).copied())
|
||||
.map(|r| r.min(baseline))
|
||||
.unwrap_or(baseline);
|
||||
if effective > bound.max {
|
||||
tracing::warn!(
|
||||
kind = kind.as_str(),
|
||||
requested = effective,
|
||||
max = bound.max,
|
||||
"配额超过全局天花板，已 clamp（检查司命 tool policy 或 *_MAX 配置）"
|
||||
);
|
||||
effective = bound.max;
|
||||
}
|
||||
out.push(EffectiveRlimit { kind, value: effective });
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// 在 fork 后、exec 前对子进程设置资源限额（软 = 硬）。
|
||||
///
|
||||
/// 供 `tokio::Command::pre_exec` 闭包调用；setrlimit 失败将导致 spawn 失败
|
||||
/// （保守：宁拒绝执行也不放过无限制的命令）。
|
||||
///
|
||||
/// `is_root` 为真时跳过 `Nproc`——root 的 `RLIMIT_NPROC` 按 uid 计会误伤系统进程。
|
||||
#[cfg(unix)]
|
||||
pub fn apply(rlimits: &[EffectiveRlimit], is_root: bool) -> std::io::Result<()> {
|
||||
for rl in rlimits {
|
||||
if rl.kind == RlimitKind::Nproc && is_root {
|
||||
continue;
|
||||
}
|
||||
// macOS 的 RLIMIT_AS 实为 RLIMIT_RSS 别名（资源号同为 5），不真正限制虚拟地址空间；
|
||||
// 且沙箱 server 自身 AS 常已超过配置上限，fork 后子进程 setrlimit(AS) 会返回 EINVAL
|
||||
// 致 spawn 失败。故 macOS 上跳过 AS——Linux 才真正强制 AS 限额。
|
||||
#[cfg(target_os = "macos")]
|
||||
if rl.kind == RlimitKind::As {
|
||||
continue;
|
||||
}
|
||||
// macOS 的 RLIMIT_NPROC 检查 uid 级总进程数（非子进程树内计数），与系统 ulimit -u
|
||||
// 重叠且更严格——用户进程数稍高即致子进程 spawn 失败。已有 ulimit -u 做系统级防护，
|
||||
// 沙箱内重复 setrlimit 会误杀合法命令，故 macOS 上跳过 NPROC。
|
||||
#[cfg(target_os = "macos")]
|
||||
if rl.kind == RlimitKind::Nproc {
|
||||
continue;
|
||||
}
|
||||
let resource = match rl.kind {
|
||||
RlimitKind::Fsize => libc::RLIMIT_FSIZE,
|
||||
RlimitKind::As => libc::RLIMIT_AS,
|
||||
RlimitKind::Cpu => libc::RLIMIT_CPU,
|
||||
RlimitKind::Nproc => libc::RLIMIT_NPROC,
|
||||
RlimitKind::Nofile => libc::RLIMIT_NOFILE,
|
||||
};
|
||||
let rlim = libc::rlimit {
|
||||
rlim_cur: rl.value,
|
||||
rlim_max: rl.value,
|
||||
};
|
||||
// setrlimit 在 fork 后、exec 前的子进程上下文调用；失败多为 EPERM（超过当前 hard limit）。
|
||||
let rc = unsafe { libc::setrlimit(resource, &rlim) };
|
||||
if rc != 0 {
|
||||
return Err(std::io::Error::last_os_error());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
pub fn apply(_rlimits: &[EffectiveRlimit], _is_root: bool) -> std::io::Result<()> {
|
||||
// 非 unix 平台无 setrlimit 等价物；资源限额仅作记录，不实际施加。
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// 每个维度 default=max=同值 的配置（等价于无分级天花板）。
|
||||
fn flat_cfg(value: u64) -> RlimitsConfig {
|
||||
let mut c = RlimitsConfig::new();
|
||||
for kind in RlimitKind::ALL {
|
||||
c.set(kind, value, value);
|
||||
}
|
||||
c
|
||||
}
|
||||
|
||||
/// 每个维度 default/max 不同的配置。
|
||||
fn tiered_cfg(default: u64, max: u64) -> RlimitsConfig {
|
||||
let mut c = RlimitsConfig::new();
|
||||
for kind in RlimitKind::ALL {
|
||||
c.set(kind, default, max);
|
||||
}
|
||||
c
|
||||
}
|
||||
|
||||
fn val(eff: &[EffectiveRlimit], kind: RlimitKind) -> u64 {
|
||||
eff.iter()
|
||||
.find(|e| e.kind == kind)
|
||||
.map(|e| e.value)
|
||||
.unwrap_or_else(|| panic!("缺少 {:?} 维度", kind))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_when_no_grant_no_request() {
|
||||
let c = tiered_cfg(100, 200);
|
||||
let eff = compute_effective(&c, None, None);
|
||||
for kind in RlimitKind::ALL {
|
||||
assert_eq!(val(&eff, kind), 100, "{:?} 应使用默认值", kind);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn granted_relaxes_within_max() {
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("fsize".to_string(), 150u64);
|
||||
let eff = compute_effective(&c, Some(&g), None);
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 150);
|
||||
// 未授予的维度仍用默认
|
||||
assert_eq!(val(&eff, RlimitKind::As), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn granted_above_max_clamped() {
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("fsize".to_string(), 300u64); // 超过 max
|
||||
let eff = compute_effective(&c, Some(&g), None);
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 200);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_only_tightens() {
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut r = BTreeMap::new();
|
||||
r.insert("fsize".to_string(), 50u64);
|
||||
let eff = compute_effective(&c, None, Some(&r));
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 50);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_cannot_widen_past_baseline() {
|
||||
// granted=150，request 想放宽到 180 → 应被 baseline(150) 钳制（不能放宽）
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("fsize".to_string(), 150u64);
|
||||
let mut r = BTreeMap::new();
|
||||
r.insert("fsize".to_string(), 180u64);
|
||||
let eff = compute_effective(&c, Some(&g), Some(&r));
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 150); // min(180, baseline=150)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_tightens_below_granted() {
|
||||
// granted=150，request 收紧到 80 → 生效 80
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("fsize".to_string(), 150u64);
|
||||
let mut r = BTreeMap::new();
|
||||
r.insert("fsize".to_string(), 80u64);
|
||||
let eff = compute_effective(&c, Some(&g), Some(&r));
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 80);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_kind_omitted() {
|
||||
let mut c = tiered_cfg(100, 200);
|
||||
c.disable(RlimitKind::Nproc);
|
||||
let eff = compute_effective(&c, None, None);
|
||||
assert!(eff.iter().all(|e| e.kind != RlimitKind::Nproc));
|
||||
assert_eq!(eff.len(), RlimitKind::ALL.len() - 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unknown_grant_key_ignored() {
|
||||
let c = tiered_cfg(100, 200);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("unknown_key".to_string(), 999u64);
|
||||
let eff = compute_effective(&c, Some(&g), None);
|
||||
for kind in RlimitKind::ALL {
|
||||
assert_eq!(val(&eff, kind), 100, "{:?} 不受未知 key 影响", kind);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn max_below_default_tightens_default() {
|
||||
// 配置 max < default 时，default 被收紧到 max（保守）
|
||||
let mut c = RlimitsConfig::new();
|
||||
c.set(RlimitKind::Fsize, 100, 40);
|
||||
assert_eq!(c.bound(RlimitKind::Fsize).unwrap().default, 40);
|
||||
assert_eq!(c.bound(RlimitKind::Fsize).unwrap().max, 40);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_config_equivalent_to_single_ceiling() {
|
||||
// default=max 时，granted 任何值都被 clamp 到该值（无分级空间）
|
||||
let c = flat_cfg(100);
|
||||
let mut g = BTreeMap::new();
|
||||
g.insert("fsize".to_string(), 150u64);
|
||||
let eff = compute_effective(&c, Some(&g), None);
|
||||
assert_eq!(val(&eff, RlimitKind::Fsize), 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rlimit_kind_roundtrip_keys() {
|
||||
for kind in RlimitKind::ALL {
|
||||
assert_eq!(RlimitKind::from_key(kind.as_str()), Some(kind));
|
||||
}
|
||||
assert!(RlimitKind::from_key("bogus").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,609 @@
|
||||
//! 懒预热 + 双缓冲预热池：管理每语言 Docker 容器的生命周期。
|
||||
//!
|
||||
//! 核心策略：
|
||||
//! - acquire() 从空闲队列取容器，通过通道通知后台创建替补（双缓冲）
|
||||
//! - release() 销毁旧容器，通过通道通知后台创建新容器（不复用脏环境）
|
||||
//! - 后台 Creator 任务处理所有容器创建，避免 Arc<Self> 自引用问题
|
||||
//!
|
||||
//! 设计规格见 docs/superpowers/specs/2026-07-05-sandbox-container-warmpool-design.md
|
||||
|
||||
use std::collections::{HashMap, HashSet, VecDeque};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use bollard::container::RemoveContainerOptions;
|
||||
use bollard::Docker;
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
|
||||
use crate::config::WarmPoolConfig;
|
||||
use crate::docker_exec::{build_container_config, connect_docker, ensure_image};
|
||||
|
||||
/// 预热池错误类型。
|
||||
#[derive(Debug)]
|
||||
pub enum WarmPoolError {
|
||||
PoolExhausted,
|
||||
DockerDaemonDown(String),
|
||||
CreateFailed(String),
|
||||
InvalidOperation(String),
|
||||
}
|
||||
|
||||
impl std::fmt::Display for WarmPoolError {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
WarmPoolError::PoolExhausted => write!(f, "预热池容器已满"),
|
||||
WarmPoolError::DockerDaemonDown(m) => write!(f, "Docker 不可用: {m}"),
|
||||
WarmPoolError::CreateFailed(m) => write!(f, "容器创建失败: {m}"),
|
||||
WarmPoolError::InvalidOperation(m) => write!(f, "非法操作: {m}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for WarmPoolError {}
|
||||
|
||||
/// 容器实例句柄。
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerHandle {
|
||||
pub id: String,
|
||||
pub language: String,
|
||||
pub created_at: Instant,
|
||||
}
|
||||
|
||||
/// 单语言容器分桶的共享状态。
|
||||
struct LanguagePool {
|
||||
idle: VecDeque<ContainerHandle>,
|
||||
creating: HashSet<String>,
|
||||
max_idle: usize,
|
||||
#[allow(dead_code)]
|
||||
image: String,
|
||||
}
|
||||
|
||||
/// 预热池共享状态（包装在 Mutex 中）。
|
||||
struct PoolInner {
|
||||
/// 语言 → 分桶
|
||||
pools: HashMap<String, LanguagePool>,
|
||||
/// 排队等待者（池满时挂起请求）
|
||||
waiters: HashMap<String, VecDeque<oneshot::Sender<ContainerHandle>>>,
|
||||
}
|
||||
|
||||
/// 预热池核心逻辑（通过 Arc 在后台任务间共享）。
|
||||
pub struct WarmPool {
|
||||
/// 共享状态
|
||||
inner: Mutex<PoolInner>,
|
||||
/// Docker 客户端
|
||||
docker: Docker,
|
||||
/// 全局容器计数
|
||||
global_count: AtomicUsize,
|
||||
/// 全局最大容器数
|
||||
global_max: usize,
|
||||
/// 预热池配置
|
||||
config: WarmPoolConfig,
|
||||
/// 工作基目录（bind mount 用）
|
||||
base_workdir: String,
|
||||
/// 容器创建通道发送端（后台 Creator 消费）
|
||||
create_tx: mpsc::UnboundedSender<String>,
|
||||
}
|
||||
|
||||
impl WarmPool {
|
||||
/// 创建预热池并启动后台 Creator。
|
||||
pub fn new(config: WarmPoolConfig, base_workdir: String, docker: Docker) -> Arc<Self> {
|
||||
let (create_tx, create_rx) = mpsc::unbounded_channel::<String>();
|
||||
|
||||
let pool = Arc::new(Self {
|
||||
inner: Mutex::new(PoolInner {
|
||||
pools: HashMap::new(),
|
||||
waiters: HashMap::new(),
|
||||
}),
|
||||
docker,
|
||||
global_count: AtomicUsize::new(0),
|
||||
global_max: config.global_max_containers,
|
||||
config,
|
||||
base_workdir,
|
||||
create_tx,
|
||||
});
|
||||
|
||||
// 启动后台 Creator
|
||||
let pool_clone = pool.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut rx = create_rx;
|
||||
while let Some(lang) = rx.recv().await {
|
||||
pool_clone.clone().handle_create_request(&lang).await;
|
||||
}
|
||||
});
|
||||
|
||||
pool
|
||||
}
|
||||
|
||||
/// 连接 Docker 并创建预热池。Docker 不可用时返回 None。
|
||||
pub async fn try_new(config: WarmPoolConfig, base_workdir: String) -> Option<Arc<Self>> {
|
||||
match connect_docker(&config.docker_socket) {
|
||||
Ok(docker) => {
|
||||
// 确保所有镜像已拉取（后台进行；失败跳过该语言）
|
||||
for (lang, image) in &config.images {
|
||||
if let Err(e) = ensure_image(&docker, image).await {
|
||||
tracing::warn!(lang, image, error = %e, "拉取镜像失败，跳过预热");
|
||||
}
|
||||
}
|
||||
Some(Self::new(config, base_workdir, docker))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "Docker 连接失败，容器模式不可用");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 预热池配置引用。
|
||||
pub fn config(&self) -> &WarmPoolConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Docker 客户端引用。
|
||||
pub fn docker(&self) -> &Docker {
|
||||
&self.docker
|
||||
}
|
||||
|
||||
/// 预热所有已配置语言的容器。
|
||||
pub async fn warmup_all(self: Arc<Self>) {
|
||||
for (lang, image) in &self.config.images {
|
||||
// 初始化分桶
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
if !inner.pools.contains_key(lang) {
|
||||
inner.pools.insert(
|
||||
lang.clone(),
|
||||
LanguagePool {
|
||||
idle: VecDeque::new(),
|
||||
creating: HashSet::new(),
|
||||
max_idle: self.config.max_idle_per_lang,
|
||||
image: image.clone(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// 逐个预热（失败不阻断）
|
||||
for i in 0..self.config.max_idle_per_lang {
|
||||
tracing::info!(lang, n = i + 1, "预热容器");
|
||||
if let Err(e) = create_one_container_now(
|
||||
&self.docker,
|
||||
&self.inner,
|
||||
&self.config,
|
||||
&self.base_workdir,
|
||||
&self.global_count,
|
||||
self.global_max,
|
||||
lang,
|
||||
image,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::warn!(lang, error = %e, "预热容器失败");
|
||||
break;
|
||||
}
|
||||
}
|
||||
tracing::info!(lang, count = self.config.max_idle_per_lang, "预热完成");
|
||||
}
|
||||
}
|
||||
|
||||
/// 从池中获取一个容器。
|
||||
pub async fn acquire(self: Arc<Self>, language: &str) -> Result<ContainerHandle, WarmPoolError> {
|
||||
let lang = language.to_string();
|
||||
|
||||
// 确保分桶已初始化
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner.pools.entry(lang.clone()).or_insert_with(|| {
|
||||
let image = self
|
||||
.config
|
||||
.images
|
||||
.get(&lang)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
LanguagePool {
|
||||
idle: VecDeque::new(),
|
||||
creating: HashSet::new(),
|
||||
max_idle: self.config.max_idle_per_lang,
|
||||
image,
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// 1) 尝试从空闲队列取
|
||||
let _hit = {
|
||||
let mut inner = self.inner.lock().await;
|
||||
if let Some(pool) = inner.pools.get_mut(&lang) {
|
||||
if let Some(handle) = pool.idle.pop_front() {
|
||||
// 双缓冲：通知后台创建替补
|
||||
let _ = self.create_tx.send(lang.clone());
|
||||
drop(inner);
|
||||
return Ok(handle);
|
||||
}
|
||||
}
|
||||
false
|
||||
};
|
||||
|
||||
// 2) 空闲为空，尝试同步创建（冷启动）
|
||||
if self.global_count.load(Ordering::Acquire) < self.global_max {
|
||||
tracing::info!(lang, "冷启动容器（空闲池为空）");
|
||||
let image = self
|
||||
.config
|
||||
.images
|
||||
.get(&lang)
|
||||
.cloned()
|
||||
.ok_or_else(|| WarmPoolError::InvalidOperation(format!("语言 '{lang}' 未配置镜像")))?;
|
||||
let handle = create_one_container_now(
|
||||
&self.docker,
|
||||
&self.inner,
|
||||
&self.config,
|
||||
&self.base_workdir,
|
||||
&self.global_count,
|
||||
self.global_max,
|
||||
&lang,
|
||||
&image,
|
||||
)
|
||||
.await?;
|
||||
// 再通知后台创建一个作为双缓冲
|
||||
let _ = self.create_tx.send(lang.clone());
|
||||
return Ok(handle);
|
||||
}
|
||||
|
||||
// 3) 已达上限：排队等待
|
||||
self.enqueue_or_timeout(&lang).await
|
||||
}
|
||||
|
||||
/// 释放容器：销毁旧容器 + 通知后台创建替补。
|
||||
pub async fn release(&self, handle: ContainerHandle) {
|
||||
let lang = handle.language.clone();
|
||||
destroy_one_container(&self.docker, &handle.id, &self.global_count).await;
|
||||
let _ = self.create_tx.send(lang);
|
||||
}
|
||||
|
||||
/// 非阻塞尝试取一个空闲容器。
|
||||
pub async fn try_acquire(&self, language: &str) -> Option<ContainerHandle> {
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner
|
||||
.pools
|
||||
.get_mut(language)
|
||||
.and_then(|pool| pool.idle.pop_front())
|
||||
}
|
||||
|
||||
/// 健康检查循环：定期 ping 空闲容器，剔除异常退出的。
|
||||
pub async fn health_check_loop(self: Arc<Self>) {
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(30));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let to_check: Vec<(String, String)> = {
|
||||
let inner = self.inner.lock().await;
|
||||
inner
|
||||
.pools
|
||||
.iter()
|
||||
.flat_map(|(lang, pool)| {
|
||||
pool.idle
|
||||
.iter()
|
||||
.map(move |h| (lang.clone(), h.id.clone()))
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
|
||||
for (lang, id) in to_check {
|
||||
let ok = tokio::time::timeout(Duration::from_secs(3), self.docker.inspect_container(&id, None::<bollard::container::InspectContainerOptions>))
|
||||
.await
|
||||
.is_ok();
|
||||
if !ok {
|
||||
tracing::warn!(container_id = %id, language = %lang, "健康检查失败，移除异常容器");
|
||||
destroy_one_container(&self.docker, &id, &self.global_count).await;
|
||||
let _ = self.create_tx.send(lang);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 空闲回收循环：超过 idle_ttl 的容器销毁回收（不补充，主动缩容）。
|
||||
pub async fn idle_eviction_loop(self: Arc<Self>) {
|
||||
let mut interval = tokio::time::interval(Duration::from_secs(60));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
let mut evict: Vec<(String, String)> = Vec::new();
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
for (lang, pool) in inner.pools.iter_mut() {
|
||||
while let Some(handle) = pool.idle.front() {
|
||||
if handle.created_at.elapsed() > self.config.idle_ttl {
|
||||
let handle = pool.idle.pop_front().unwrap();
|
||||
evict.push((lang.clone(), handle.id.clone()));
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
for (_lang, id) in evict {
|
||||
destroy_one_container(&self.docker, &id, &self.global_count).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 优雅关闭：销毁所有容器。
|
||||
pub async fn shutdown(&self) {
|
||||
let inner = self.inner.lock().await;
|
||||
for (_lang, pool) in inner.pools.iter() {
|
||||
for handle in &pool.idle {
|
||||
destroy_one_container(&self.docker, &handle.id, &self.global_count).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 获取当前池统计信息。
|
||||
pub async fn collect_metrics(&self) -> HashMap<String, (i64, i64, i64)> {
|
||||
let mut stats = HashMap::new();
|
||||
let inner = self.inner.lock().await;
|
||||
for (lang, pool) in inner.pools.iter() {
|
||||
stats.insert(
|
||||
lang.clone(),
|
||||
(
|
||||
pool.idle.len() as i64,
|
||||
pool.creating.len() as i64,
|
||||
self.global_count.load(Ordering::Acquire) as i64,
|
||||
),
|
||||
);
|
||||
}
|
||||
stats
|
||||
}
|
||||
|
||||
/// 排队等待空闲容器。
|
||||
async fn enqueue_or_timeout(
|
||||
&self,
|
||||
language: &str,
|
||||
) -> Result<ContainerHandle, WarmPoolError> {
|
||||
let (tx, rx) = oneshot::channel::<ContainerHandle>();
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
inner
|
||||
.waiters
|
||||
.entry(language.to_string())
|
||||
.or_default()
|
||||
.push_back(tx);
|
||||
}
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), rx)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|r| r.ok())
|
||||
.ok_or(WarmPoolError::PoolExhausted)
|
||||
}
|
||||
/// 处理后台创建请求（由 Creator 任务调用）。
|
||||
async fn handle_create_request(self: Arc<Self>, language: &str) {
|
||||
if self.global_count.load(Ordering::Acquire) >= self.global_max {
|
||||
return;
|
||||
}
|
||||
|
||||
let image = match self.config.images.get(language) {
|
||||
Some(img) => img.clone(),
|
||||
None => return,
|
||||
};
|
||||
|
||||
// 检查 creating 去重和 idle 上限
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
if let Some(pool) = inner.pools.get_mut(language) {
|
||||
if pool.creating.contains(language) || pool.idle.len() >= pool.max_idle {
|
||||
return;
|
||||
}
|
||||
pool.creating.insert(language.to_string());
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let result = create_one_container_now(
|
||||
&self.docker,
|
||||
&self.inner,
|
||||
&self.config,
|
||||
&self.base_workdir,
|
||||
&self.global_count,
|
||||
self.global_max,
|
||||
language,
|
||||
&image,
|
||||
)
|
||||
.await;
|
||||
|
||||
// 清除 creating 标记
|
||||
{
|
||||
let mut inner = self.inner.lock().await;
|
||||
if let Some(pool) = inner.pools.get_mut(language) {
|
||||
pool.creating.remove(language);
|
||||
}
|
||||
}
|
||||
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(lang = language, error = %e, "后台创建容器失败");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// 自由函数
|
||||
|
||||
// ============================================================
|
||||
// 自由函数：容器创建和销毁
|
||||
// ============================================================
|
||||
|
||||
/// 同步创建一个容器并放入空闲队列。
|
||||
async fn create_one_container_now(
|
||||
docker: &Docker,
|
||||
inner: &Mutex<PoolInner>,
|
||||
config: &WarmPoolConfig,
|
||||
base_workdir: &str,
|
||||
global_count: &AtomicUsize,
|
||||
_global_max: usize,
|
||||
language: &str,
|
||||
image: &str,
|
||||
) -> Result<ContainerHandle, WarmPoolError> {
|
||||
let start = Instant::now();
|
||||
|
||||
let container_config = build_container_config(config, image, base_workdir);
|
||||
|
||||
let container = tokio::time::timeout(config.create_timeout, async {
|
||||
docker
|
||||
.create_container::<String, String>(None, container_config)
|
||||
.await
|
||||
.map_err(|e| WarmPoolError::CreateFailed(format!("create_container: {e}")))
|
||||
})
|
||||
.await
|
||||
.map_err(|_| WarmPoolError::CreateFailed("容器创建超时".into()))?
|
||||
.map_err(|e| {
|
||||
global_count.fetch_sub(1, Ordering::Release);
|
||||
e
|
||||
})?;
|
||||
|
||||
docker
|
||||
.start_container::<String>(&container.id, None)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
global_count.fetch_sub(1, Ordering::Release);
|
||||
WarmPoolError::CreateFailed(format!("start_container: {e}"))
|
||||
})?;
|
||||
|
||||
global_count.fetch_add(1, Ordering::Release);
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
tracing::debug!(
|
||||
container_id = %container.id,
|
||||
language,
|
||||
elapsed_ms = elapsed.as_millis(),
|
||||
"容器创建完成"
|
||||
);
|
||||
|
||||
let handle = ContainerHandle {
|
||||
id: container.id,
|
||||
language: language.to_string(),
|
||||
created_at: Instant::now(),
|
||||
};
|
||||
|
||||
// 放入空闲队列
|
||||
{
|
||||
let mut inner_lock = inner.lock().await;
|
||||
if let Some(pool) = inner_lock.pools.get_mut(language) {
|
||||
pool.idle.push_back(handle.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// 尝试唤醒等待者
|
||||
notify_waiter(inner, language, &handle).await;
|
||||
|
||||
Ok(handle)
|
||||
}
|
||||
|
||||
/// 销毁容器。
|
||||
async fn destroy_one_container(docker: &Docker, container_id: &str, global_count: &AtomicUsize) {
|
||||
let result = docker
|
||||
.remove_container(
|
||||
container_id,
|
||||
Some(RemoveContainerOptions {
|
||||
force: true,
|
||||
v: true,
|
||||
..Default::default()
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
global_count.fetch_sub(1, Ordering::Release);
|
||||
|
||||
if let Err(e) = result {
|
||||
tracing::warn!(container_id, error = %e, "销毁容器失败（可能已被外部杀死）");
|
||||
}
|
||||
}
|
||||
|
||||
/// 唤醒一个等待者。
|
||||
async fn notify_waiter(
|
||||
inner: &Mutex<PoolInner>,
|
||||
language: &str,
|
||||
handle: &ContainerHandle,
|
||||
) {
|
||||
let mut inner_lock = inner.lock().await;
|
||||
if let Some(queue) = inner_lock.waiters.get_mut(language) {
|
||||
while let Some(tx) = queue.pop_front() {
|
||||
if tx.send(handle.clone()).is_ok() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_config() -> WarmPoolConfig {
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images
|
||||
.insert("python".into(), "python:3.12-slim".into());
|
||||
cfg.images
|
||||
.insert("generic".into(), "debian:bookworm-slim".into());
|
||||
cfg.max_idle_per_lang = 2;
|
||||
cfg.global_max_containers = 10;
|
||||
cfg
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warm_pool_config_defaults() {
|
||||
let cfg = WarmPoolConfig::default();
|
||||
assert_eq!(cfg.max_idle_per_lang, 3);
|
||||
assert_eq!(cfg.global_max_containers, 20);
|
||||
assert_eq!(cfg.memory_per_container, 512 * 1024 * 1024);
|
||||
assert_eq!(cfg.pids_limit, 64);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn container_handle_basic() {
|
||||
let h = ContainerHandle {
|
||||
id: "abc123".into(),
|
||||
language: "python".into(),
|
||||
created_at: Instant::now(),
|
||||
};
|
||||
assert_eq!(h.id, "abc123");
|
||||
assert_eq!(h.language, "python");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warm_pool_config_image_mapping() {
|
||||
let mut cfg = WarmPoolConfig::default();
|
||||
cfg.images
|
||||
.insert("python".into(), "python:3.12-slim".into());
|
||||
cfg.images.insert("node".into(), "node:22-slim".into());
|
||||
assert_eq!(cfg.images.get("python").unwrap(), "python:3.12-slim");
|
||||
assert_eq!(cfg.images.get("node").unwrap(), "node:22-slim");
|
||||
assert!(cfg.images.get("java").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_initialized_with_empty_idle() {
|
||||
let docker = connect_docker("unix:///nonexistent").ok();
|
||||
if docker.is_none() {
|
||||
eprintln!("跳过：Docker 不可用");
|
||||
return;
|
||||
}
|
||||
let pool = WarmPool::new(make_config(), "/tmp".into(), docker.unwrap());
|
||||
let stats = pool.collect_metrics().await;
|
||||
assert!(stats.is_empty() || stats.values().all(|(idle, ..)| *idle == 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn global_count_atomic() {
|
||||
let count = AtomicUsize::new(0);
|
||||
assert_eq!(count.load(Ordering::Acquire), 0);
|
||||
count.fetch_add(1, Ordering::Release);
|
||||
assert_eq!(count.load(Ordering::Acquire), 1);
|
||||
count.fetch_sub(1, Ordering::Release);
|
||||
assert_eq!(count.load(Ordering::Acquire), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn warm_pool_error_display() {
|
||||
assert_eq!(format!("{}", WarmPoolError::PoolExhausted), "预热池容器已满");
|
||||
assert_eq!(
|
||||
format!("{}", WarmPoolError::CreateFailed("oom".into())),
|
||||
"容器创建失败: oom"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user