diff --git a/subprojects/stardomain/Cargo.lock b/subprojects/stardomain/Cargo.lock index 31aa5ba..c0f5a26 100644 --- a/subprojects/stardomain/Cargo.lock +++ b/subprojects/stardomain/Cargo.lock @@ -1081,6 +1081,7 @@ dependencies = [ "anyhow", "bollard", "clap", + "futures-util", "pyo3", "pyo3-async-runtimes", "regex", diff --git a/subprojects/stardomain/Cargo.toml b/subprojects/stardomain/Cargo.toml index 6ba31be..2aebb3f 100644 --- a/subprojects/stardomain/Cargo.toml +++ b/subprojects/stardomain/Cargo.toml @@ -19,6 +19,7 @@ pyo3 = { version = "0.28", features = ["extension-module"] } pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] } tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros", "process", "time"] } bollard = "0.18" +futures-util = "0.3" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" serde_yaml = "0.9" diff --git a/subprojects/stardomain/src/docker/mod.rs b/subprojects/stardomain/src/docker/mod.rs index 5946138..f86aa09 100644 --- a/subprojects/stardomain/src/docker/mod.rs +++ b/subprojects/stardomain/src/docker/mod.rs @@ -8,25 +8,34 @@ * http://www.apache.org/licenses/LICENSE-2.0 */ -/// Docker integration layer. -/// Uses bollard to manage container lifecycle for sandbox execution. +use std::time::Duration; +use bollard::Docker; +use bollard::container::{ + Config, CreateContainerOptions, StartContainerOptions, + WaitContainerOptions, LogsOptions, RemoveContainerOptions, +}; +use bollard::models::HostConfig; +use futures_util::StreamExt; +use tokio::time::timeout; + +use crate::sandbox::SandboxResult; pub struct DockerBackend { pub image: String, pub network_disabled: bool, - pub memory_limit: u64, - pub cpu_period: u64, - pub cpu_quota: u64, + pub memory_limit: i64, + pub cpu_period: i64, + pub cpu_quota: i64, } impl Default for DockerBackend { fn default() -> Self { DockerBackend { - image: "stardomain-runtime:latest".to_string(), + image: "python:3.13-slim".to_string(), network_disabled: false, - memory_limit: 512 * 1024 * 1024, // 512MB + memory_limit: 512 * 1024 * 1024, cpu_period: 100_000, - cpu_quota: 50_000, // 50% of one core + cpu_quota: 50_000, } } } @@ -41,4 +50,162 @@ impl DockerBackend { } backend } + + pub async fn run_command( + &self, + command: &str, + workspace: &str, + timeout_secs: u64, + ) -> SandboxResult { + let docker = match Docker::connect_with_local_defaults() { + Ok(d) => d, + Err(e) => return SandboxResult::error(&format!("Docker 连接失败: {}", e)), + }; + + let container_name = format!("stardomain_{}", uuid_short()); + + let host_config = HostConfig { + memory: Some(self.memory_limit), + cpu_period: Some(self.cpu_period), + cpu_quota: Some(self.cpu_quota), + network_mode: if self.network_disabled { + Some("none".to_string()) + } else { + None + }, + binds: Some(vec![format!("{}:/workspace", workspace)]), + ..Default::default() + }; + + let config = Config { + image: Some(self.image.clone()), + cmd: Some(vec![ + "sh".to_string(), + "-c".to_string(), + command.to_string(), + ]), + working_dir: Some("/workspace".to_string()), + host_config: Some(host_config), + ..Default::default() + }; + + let create_opts = CreateContainerOptions { + name: &container_name, + platform: None, + }; + if let Err(e) = docker.create_container(Some(create_opts), config).await { + return SandboxResult::error(&format!("创建容器失败: {}", e)); + } + + if let Err(e) = docker + .start_container( + &container_name, + None::>, + ) + .await + { + let _ = cleanup(&docker, &container_name).await; + return SandboxResult::error(&format!("启动容器失败: {}", e)); + } + + let wait_fut = wait_for_container(&docker, &container_name); + let result = match timeout(Duration::from_secs(timeout_secs), wait_fut).await + { + Ok(r) => r, + Err(_) => { + let _ = docker + .kill_container( + &container_name, + None::>, + ) + .await; + let logs = collect_logs(&docker, &container_name).await; + let _ = cleanup(&docker, &container_name).await; + return SandboxResult { + stdout: logs.0, + stderr: logs.1, + exit_code: -1, + killed_by_timeout: true, + }; + } + }; + + let logs = collect_logs(&docker, &container_name).await; + let _ = cleanup(&docker, &container_name).await; + + SandboxResult { + stdout: logs.0, + stderr: logs.1, + exit_code: result, + killed_by_timeout: false, + } + } + + pub async fn run_python( + &self, + code: &str, + workspace: &str, + timeout_secs: u64, + ) -> SandboxResult { + let escaped = code.replace('\'', "'\\''"); + let cmd = format!("python3 -c '{}'", escaped); + self.run_command(&cmd, workspace, timeout_secs).await + } } + +async fn wait_for_container(docker: &Docker, name: &str) -> i32 { + let opts = WaitContainerOptions { condition: "not-running" }; + let mut stream = docker.wait_container(name, Some(opts)); + if let Some(Ok(response)) = stream.next().await { + response.status_code as i32 + } else { + -1 + } +} + +async fn collect_logs(docker: &Docker, name: &str) -> (String, String) { + let opts = LogsOptions:: { + stdout: true, + stderr: true, + ..Default::default() + }; + let mut stream = docker.logs(name, Some(opts)); + let mut stdout = String::new(); + let mut stderr = String::new(); + + while let Some(Ok(output)) = stream.next().await { + match output { + bollard::container::LogOutput::StdOut { message } => { + stdout.push_str(&String::from_utf8_lossy(&message)); + } + bollard::container::LogOutput::StdErr { message } => { + stderr.push_str(&String::from_utf8_lossy(&message)); + } + _ => {} + } + } + (stdout, stderr) +} + +async fn cleanup( + docker: &Docker, + name: &str, +) -> Result<(), bollard::errors::Error> { + docker + .remove_container( + name, + Some(RemoveContainerOptions { + force: true, + ..Default::default() + }), + ) + .await +} + +fn uuid_short() -> String { + use std::time::{SystemTime, UNIX_EPOCH}; + let t = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default(); + format!("{:x}{:x}", t.as_secs(), t.subsec_nanos()) +} \ No newline at end of file diff --git a/subprojects/stardomain/src/lib.rs b/subprojects/stardomain/src/lib.rs index 15322d0..f581bb3 100644 --- a/subprojects/stardomain/src/lib.rs +++ b/subprojects/stardomain/src/lib.rs @@ -19,14 +19,25 @@ pub mod policy; pub mod sandbox; use pyo3::prelude::*; - use sandbox::SandboxResult; +fn get_runtime() -> &'static tokio::runtime::Runtime { + use std::sync::OnceLock; + static RT: OnceLock = OnceLock::new(); + RT.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("Failed to create tokio runtime") + }) +} + #[pyclass] struct Sandbox { mode: String, workspace: String, timeout: u64, + #[allow(dead_code)] memory_limit_mb: u64, policy_name: String, } @@ -34,26 +45,102 @@ struct Sandbox { #[pymethods] impl Sandbox { #[new] - #[pyo3(signature = (mode="sandbox", workspace="/tmp/stardomain_ws", timeout=30, memory_limit_mb=512, policy="agent_exec"))] - fn new(mode: &str, workspace: &str, timeout: u64, memory_limit_mb: u64, policy: &str) -> Self { - Sandbox { + #[pyo3(signature = (mode="local", workspace="/tmp/stardomain_ws", timeout=30, memory_limit_mb=512, policy="agent_exec"))] + fn new(mode: &str, workspace: &str, timeout: u64, memory_limit_mb: u64, policy: &str) -> PyResult { + std::fs::create_dir_all(workspace).map_err(|e| { + PyErr::new::(format!("无法创建工作目录: {}", e)) + })?; + Ok(Sandbox { mode: mode.to_string(), workspace: workspace.to_string(), timeout, memory_limit_mb, policy_name: policy.to_string(), - } + }) } fn run_command(&self, command: &str) -> PyResult { - let policy = policy::get_policy(&self.policy_name); - policy::filter::validate_command(command, &policy)?; - Ok(SandboxResult::stub(command)) + let pol = policy::get_policy(&self.policy_name); + policy::filter::validate_command(command, &pol)?; + + let rt = get_runtime(); + match self.mode.as_str() { + "sandbox" | "docker" => { + let backend = docker::DockerBackend::with_policy(&self.policy_name); + Ok(rt.block_on(backend.run_command(command, &self.workspace, self.timeout))) + } + _ => { + Ok(rt.block_on(sandbox::executor::run_local(command, &self.workspace, self.timeout))) + } + } } fn run_python(&self, code: &str) -> PyResult { - let policy = policy::get_policy(&self.policy_name); - policy::filter::validate_python(code, &policy)?; - Ok(SandboxResult::stub(code)) + let pol = policy::get_policy(&self.policy_name); + policy::filter::validate_python(code, &pol)?; + + let rt = get_runtime(); + match self.mode.as_str() { + "sandbox" | "docker" => { + let backend = docker::DockerBackend::with_policy(&self.policy_name); + Ok(rt.block_on(backend.run_python(code, &self.workspace, self.timeout))) + } + _ => { + Ok(rt.block_on(sandbox::executor::run_python_local(code, &self.workspace, self.timeout))) + } + } + } + + fn run_command_async<'py>(&self, py: Python<'py>, command: String) -> PyResult> { + let pol = policy::get_policy(&self.policy_name); + policy::filter::validate_command(&command, &pol)?; + + let mode = self.mode.clone(); + let workspace = self.workspace.clone(); + let timeout = self.timeout; + let policy_name = self.policy_name.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let result = match mode.as_str() { + "sandbox" | "docker" => { + let backend = docker::DockerBackend::with_policy(&policy_name); + backend.run_command(&command, &workspace, timeout).await + } + _ => { + sandbox::executor::run_local(&command, &workspace, timeout).await + } + }; + Ok(result) + }) + } + + fn run_python_async<'py>(&self, py: Python<'py>, code: String) -> PyResult> { + let pol = policy::get_policy(&self.policy_name); + policy::filter::validate_python(&code, &pol)?; + + let mode = self.mode.clone(); + let workspace = self.workspace.clone(); + let timeout = self.timeout; + let policy_name = self.policy_name.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + let result = match mode.as_str() { + "sandbox" | "docker" => { + let backend = docker::DockerBackend::with_policy(&policy_name); + backend.run_python(&code, &workspace, timeout).await + } + _ => { + sandbox::executor::run_python_local(&code, &workspace, timeout).await + } + }; + Ok(result) + }) } } + +#[pymodule] +fn stardomain(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/subprojects/stardomain/src/sandbox/executor.rs b/subprojects/stardomain/src/sandbox/executor.rs index 265e495..1ccc31f 100644 --- a/subprojects/stardomain/src/sandbox/executor.rs +++ b/subprojects/stardomain/src/sandbox/executor.rs @@ -8,16 +8,54 @@ * http://www.apache.org/licenses/LICENSE-2.0 */ -/// Executor: responsible for running commands either locally or in Docker. -/// This is a stub — actual Docker execution will be implemented later. -pub struct Executor; +use std::time::Duration; +use tokio::process::Command; +use tokio::time::timeout; -impl Executor { - pub fn run_local(_command: &str) -> (String, String, i32) { - ("".to_string(), "".to_string(), 0) - } +use crate::sandbox::SandboxResult; - pub fn run_docker(_command: &str) -> (String, String, i32) { - ("".to_string(), "[stardomain] Docker execution not yet implemented".to_string(), 1) +pub async fn run_local( + command: &str, + workspace: &str, + timeout_secs: u64, +) -> SandboxResult { + let fut = Command::new("sh") + .arg("-c") + .arg(command) + .current_dir(workspace) + .output(); + + match timeout(Duration::from_secs(timeout_secs), fut).await { + Ok(Ok(output)) => SandboxResult { + stdout: String::from_utf8_lossy(&output.stdout).to_string(), + stderr: String::from_utf8_lossy(&output.stderr).to_string(), + exit_code: output.status.code().unwrap_or(-1), + killed_by_timeout: false, + }, + Ok(Err(e)) => SandboxResult { + stdout: String::new(), + stderr: format!("执行失败: {}", e), + exit_code: -1, + killed_by_timeout: false, + }, + Err(_) => SandboxResult { + stdout: String::new(), + stderr: "执行超时".to_string(), + exit_code: -1, + killed_by_timeout: true, + }, } } + +pub async fn run_python_local( + code: &str, + workspace: &str, + timeout_secs: u64, +) -> SandboxResult { + let cmd = format!("python3 -c {}", shell_escape(code)); + run_local(&cmd, workspace, timeout_secs).await +} + +fn shell_escape(s: &str) -> String { + format!("'{}'", s.replace('\'', "'\\''")) +} diff --git a/subprojects/stardomain/src/sandbox/mod.rs b/subprojects/stardomain/src/sandbox/mod.rs index fbaf86a..105b0a2 100644 --- a/subprojects/stardomain/src/sandbox/mod.rs +++ b/subprojects/stardomain/src/sandbox/mod.rs @@ -35,4 +35,13 @@ impl SandboxResult { killed_by_timeout: false, } } + + pub fn error(msg: &str) -> Self { + SandboxResult { + stdout: String::new(), + stderr: msg.to_string(), + exit_code: -1, + killed_by_timeout: false, + } + } }