KiloStar/subprojects/stardomain/src/docker/mod.rs

/*
 * Copyright 2026 zhaoxi826
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 */

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: i64,
    pub cpu_period: i64,
    pub cpu_quota: i64,
}

impl Default for DockerBackend {
    fn default() -> Self {
        DockerBackend {
            image: "python:3.13-slim".to_string(),
            network_disabled: false,
            memory_limit: 512 * 1024 * 1024,
            cpu_period: 100_000,
            cpu_quota: 50_000,
        }
    }
}

impl DockerBackend {
    pub fn with_policy(policy_name: &str) -> Self {
        let mut backend = Self::default();
        if policy_name == "untrusted" {
            backend.network_disabled = true;
            backend.memory_limit = 128 * 1024 * 1024;
            backend.cpu_quota = 25_000;
        }
        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::<StartContainerOptions<String>>,
            )
            .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::<bollard::container::KillContainerOptions<String>>,
                    )
                    .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::<String> {
        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())
}