pos nomas

yuubox-core/Cargo.toml

@@ -0,0 +1,24 @@
+[package]
+name = "yuubox-core"
+version = "1.0.0"
+edition = "2021"
+license = "Apache-2.0"
+
+[lib]
+name = "yuubox_core"
+crate-type = ["cdylib"]
+
+[dependencies]
+pyo3 = { version = "0.20", features = ["extension-module", "abi3-py39"] }
+tokio = { version = "1.35", features = ["full"] }
+bollard = "0.15"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+chrono = "0.4"
+anyhow = "1.0"
+thiserror = "1.0"
+uuid = { version = "1.6", features = ["v4"] }
+futures-util = "0.3"
+
+[dev-dependencies]
+tokio-test = "0.4"

yuubox-core/src/container.rs

@@ -0,0 +1,222 @@
+use bollard::Docker;
+use bollard::container::{Config, CreateContainerOptions, RemoveContainerOptions, StartContainerOptions, WaitContainerOptions};
+use bollard::models::{HostConfig, RestartPolicy, RestartPolicyNameEnum};
+use bollard::exec::{CreateExecOptions, StartExecResults};
+use anyhow::{Result, anyhow};
+use std::collections::HashMap;
+use tokio::time::{timeout, Duration};
+use futures_util::stream::StreamExt;
+
+use crate::limits::ResourceLimits;
+
+pub struct ContainerResult {
+    pub stdout: String,
+    pub stderr: String,
+    pub exit_code: i64,
+    pub memory_used: u64,
+    pub cpu_time: f64,
+}
+
+pub struct ContainerManager {
+    docker: Docker,
+}
+
+impl ContainerManager {
+    pub async fn new() -> Result<Self> {
+        let docker = Docker::connect_with_socket_defaults()
+            .map_err(|e| anyhow!("Failed to connect to Docker: {}", e))?;
+        
+        Ok(Self { docker })
+    }
+
+    pub async fn execute(
+        &self,
+        code: &str,
+        language: &str,
+        limits: &ResourceLimits,
+    ) -> Result<ContainerResult> {
+        let image = self.get_image(language)?;
+        let command = self.build_command(code, language)?;
+
+        let host_config = Some(HostConfig {
+            memory: Some(limits.memory_bytes as i64),
+            nano_cpus: Some((limits.cpu_quota * 1_000_000_000.0) as i64),
+            network_mode: Some("none".to_string()),
+            read_only_root_fs: Some(true),
+            tmpfs: Some({
+                let mut map = HashMap::new();
+                map.insert("/tmp".to_string(), "size=100m".to_string());
+                map
+            }),
+            cap_drop: Some(vec!["ALL".to_string()]),
+            restart_policy: Some(RestartPolicy {
+                name: Some(RestartPolicyNameEnum::NO),
+                ..Default::default()
+            }),
+            ..Default::default()
+        });
+
+        let config = Config {
+            image: Some(image.clone()),
+            cmd: Some(command),
+            working_dir: Some("/workspace".to_string()),
+            user: Some("1000:1000".to_string()),
+            host_config,
+            attach_stdout: Some(true),
+            attach_stderr: Some(true),
+            tty: Some(false),
+            ..Default::default()
+        };
+
+        let container_name = format!("yuubox-{}", uuid::Uuid::new_v4());
+        
+        let container = self.docker
+            .create_container(
+                Some(CreateContainerOptions {
+                    name: &container_name,
+                    ..Default::default()
+                }),
+                config,
+            )
+            .await?;
+
+        let container_id = container.id.clone();
+
+        let exec_result = timeout(
+            Duration::from_secs(limits.timeout_seconds),
+            self.run_container(&container_id)
+        ).await;
+
+        let result = match exec_result {
+            Ok(Ok(res)) => res,
+            Ok(Err(e)) => {
+                self.cleanup_container(&container_id).await?;
+                return Err(e);
+            }
+            Err(_) => {
+                self.docker.kill_container::<String>(&container_id, None).await.ok();
+                self.cleanup_container(&container_id).await?;
+                return Err(anyhow!("Execution timeout after {} seconds", limits.timeout_seconds));
+            }
+        };
+
+        self.cleanup_container(&container_id).await?;
+
+        Ok(result)
+    }
+
+    async fn run_container(&self, container_id: &str) -> Result<ContainerResult> {
+        self.docker.start_container::<String>(container_id, None).await?;
+
+        let mut wait_stream = self.docker.wait_container(
+            container_id,
+            Some(WaitContainerOptions {
+                condition: "not-running",
+            })
+        );
+
+        let mut exit_code = 0i64;
+        while let Some(wait_result) = wait_stream.next().await {
+            if let Ok(status) = wait_result {
+                exit_code = status.status_code;
+            }
+        }
+
+        let logs = self.docker.logs::<String>(
+            container_id,
+            Some(bollard::container::LogsOptions {
+                stdout: true,
+                stderr: true,
+                ..Default::default()
+            })
+        );
+
+        let mut stdout = String::new();
+        let mut stderr = String::new();
+
+        let log_output = logs.collect::<Vec<_>>().await;
+        for log in log_output {
+            if let Ok(log_line) = log {
+                match log_line {
+                    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));
+                    }
+                    _ => {}
+                }
+            }
+        }
+
+        Ok(ContainerResult {
+            stdout,
+            stderr,
+            exit_code,
+            memory_used: 0,
+            cpu_time: 0.0,
+        })
+    }
+
+    async fn cleanup_container(&self, container_id: &str) -> Result<()> {
+        self.docker.remove_container(
+            container_id,
+            Some(RemoveContainerOptions {
+                force: true,
+                ..Default::default()
+            })
+        ).await?;
+        
+        Ok(())
+    }
+
+    pub async fn cleanup(&self) -> Result<()> {
+        Ok(())
+    }
+
+    fn get_image(&self, language: &str) -> Result<String> {
+        let image = match language.to_lowercase().as_str() {
+            "python" => "python:3.11-slim",
+            "javascript" | "js" | "node" => "node:20-slim",
+            "rust" => "rust:1.75-slim",
+            "go" => "golang:1.21-alpine",
+            "java" => "openjdk:17-slim",
+            _ => return Err(anyhow!("Unsupported language: {}", language)),
+        };
+        
+        Ok(image.to_string())
+    }
+
+    fn build_command(&self, code: &str, language: &str) -> Result<Vec<String>> {
+        let command = match language.to_lowercase().as_str() {
+            "python" => vec![
+                "python".to_string(),
+                "-c".to_string(),
+                code.to_string(),
+            ],
+            "javascript" | "js" | "node" => vec![
+                "node".to_string(),
+                "-e".to_string(),
+                code.to_string(),
+            ],
+            "rust" => vec![
+                "sh".to_string(),
+                "-c".to_string(),
+                format!("echo '{}' > /tmp/main.rs && rustc /tmp/main.rs -o /tmp/main && /tmp/main", code),
+            ],
+            "go" => vec![
+                "sh".to_string(),
+                "-c".to_string(),
+                format!("echo '{}' > /tmp/main.go && go run /tmp/main.go", code),
+            ],
+            "java" => vec![
+                "sh".to_string(),
+                "-c".to_string(),
+                format!("echo '{}' > /tmp/Main.java && javac /tmp/Main.java && java -cp /tmp Main", code),
+            ],
+            _ => return Err(anyhow!("Unsupported language: {}", language)),
+        };
+        
+        Ok(command)
+    }
+}

yuubox-core/src/lib.rs

@@ -0,0 +1,96 @@
+use pyo3::prelude::*;
+use pyo3::exceptions::PyRuntimeError;
+use tokio::runtime::Runtime;
+use std::time::Instant;
+
+mod container;
+mod limits;
+mod monitor;
+
+use container::ContainerManager;
+use limits::ResourceLimits;
+
+#[pyclass]
+#[derive(Clone)]
+pub struct ExecutionResult {
+    #[pyo3(get)]
+    pub stdout: String,
+    #[pyo3(get)]
+    pub stderr: String,
+    #[pyo3(get)]
+    pub exit_code: i64,
+    #[pyo3(get)]
+    pub execution_time: f64,
+    #[pyo3(get)]
+    pub memory_used: u64,
+    #[pyo3(get)]
+    pub cpu_time: f64,
+}
+
+#[pyclass]
+pub struct ContainerExecutor {
+    runtime: Runtime,
+    manager: ContainerManager,
+}
+
+#[pymethods]
+impl ContainerExecutor {
+    #[new]
+    pub fn new() -> PyResult<Self> {
+        let runtime = Runtime::new()
+            .map_err(|e| PyRuntimeError::new_err(format!("Failed to create runtime: {}", e)))?;
+        
+        let manager = runtime.block_on(async {
+            ContainerManager::new().await
+        }).map_err(|e| PyRuntimeError::new_err(format!("Failed to create manager: {}", e)))?;
+
+        Ok(Self { runtime, manager })
+    }
+
+    pub fn execute(
+        &self,
+        code: String,
+        language: String,
+        memory_mb: Option<u64>,
+        cpu_quota: Option<f64>,
+        timeout_seconds: Option<u64>,
+    ) -> PyResult<ExecutionResult> {
+        let limits = ResourceLimits {
+            memory_bytes: memory_mb.unwrap_or(256) * 1024 * 1024,
+            cpu_quota: cpu_quota.unwrap_or(1.0),
+            timeout_seconds: timeout_seconds.unwrap_or(60),
+        };
+
+        let start = Instant::now();
+        
+        let result = self.runtime.block_on(async {
+            self.manager.execute(&code, &language, &limits).await
+        }).map_err(|e| PyRuntimeError::new_err(format!("Execution failed: {}", e)))?;
+
+        let execution_time = start.elapsed().as_secs_f64();
+
+        Ok(ExecutionResult {
+            stdout: result.stdout,
+            stderr: result.stderr,
+            exit_code: result.exit_code,
+            execution_time,
+            memory_used: result.memory_used,
+            cpu_time: result.cpu_time,
+        })
+    }
+
+    pub fn cleanup(&self) -> PyResult<()> {
+        self.runtime.block_on(async {
+            self.manager.cleanup().await
+        }).map_err(|e| PyRuntimeError::new_err(format!("Cleanup failed: {}", e)))?;
+        
+        Ok(())
+    }
+}
+
+#[pymodule]
+fn yuubox_core(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_class::<ContainerExecutor>()?;
+    m.add_class::<ExecutionResult>()?;
+    Ok(())
+}

yuubox-core/src/limits.rs

@@ -0,0 +1,16 @@
+#[derive(Debug, Clone)]
+pub struct ResourceLimits {
+    pub memory_bytes: u64,
+    pub cpu_quota: f64,
+    pub timeout_seconds: u64,
+}
+
+impl Default for ResourceLimits {
+    fn default() -> Self {
+        Self {
+            memory_bytes: 256 * 1024 * 1024,
+            cpu_quota: 1.0,
+            timeout_seconds: 60,
+        }
+    }
+}

yuubox-core/src/monitor.rs

@@ -0,0 +1,17 @@
+use std::time::Instant;
+
+pub struct ResourceMonitor {
+    start_time: Instant,
+}
+
+impl ResourceMonitor {
+    pub fn new() -> Self {
+        Self {
+            start_time: Instant::now(),
+        }
+    }
+
+    pub fn elapsed(&self) -> f64 {
+        self.start_time.elapsed().as_secs_f64()
+    }
+}

yuubox-core/tests/test_container.rs

@@ -0,0 +1,7 @@
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_basic() {
+        assert_eq!(2 + 2, 4);
+    }
+}