yggdrasil/infra/
docker.rs

1use futures::StreamExt;
2use std::collections::HashMap;
3#[cfg(test)]
4use std::future::Future;
5use std::sync::LazyLock;
6use std::time::Duration;
7use tokio::sync::oneshot;
8use tokio::time::timeout;
9
10use crate::infra::runner_config::{ResourceLimits, RUNNER_CONFIG};
11use bollard::container::LogOutput;
12use bollard::models::{ContainerCreateBody, HostConfig, ResourcesUlimits};
13use bollard::query_parameters::{
14    AttachContainerOptions, CreateContainerOptions, RemoveContainerOptions, StartContainerOptions,
15    WaitContainerOptions,
16};
17use bollard::Docker;
18
19type CleanupSignal = oneshot::Sender<bool>;
20type ContainerOutput = (Option<i64>, String, String, bool);
21type ContainerResult = Result<ContainerOutput, bollard::errors::Error>;
22
23/// 共享的 Docker 客户端。
24///
25/// `connect_with_unix` 会立即探测 socket 是否存在：缺失（未安装 / 未运行 Docker，
26/// 或 `DOCKER_SOCKET_PATH` 指错）时返回 `SocketNotFoundError`。
27///
28/// 连接失败不 panic：release profile 是 `panic = "abort"`，会让整个服务进程一起挂掉，
29/// 而博客本身并不依赖 Docker。改为记 error 日志后返回 `None`，代码运行器把
30/// `None` 转成普通 bollard 错误向上冒泡，最终在 execute.rs 的错误脱敏层统一映射为
31/// 「系统暂时不可用」，其余功能不受影响。
32pub static DOCKER_CLIENT: LazyLock<Option<Docker>> = LazyLock::new(|| {
33    match Docker::connect_with_unix(
34        &RUNNER_CONFIG.docker_socket_path,
35        120,
36        bollard::API_DEFAULT_VERSION,
37    ) {
38        Ok(d) => Some(d),
39        Err(e) => {
40            tracing::error!(
41                error = ?e,
42                socket = %RUNNER_CONFIG.docker_socket_path,
43                "无法连接 Docker daemon，代码运行功能将不可用。\
44                 请确认 Docker 已安装并运行，或设置正确的 DOCKER_SOCKET_PATH",
45            );
46            None
47        }
48    }
49});
50
51/// 取共享 Docker 客户端；daemon 不可用时返回 IOError（NotFound）。
52///
53/// `NotFound` 既不命中 `TimedOut` 也不命中超时判断，会走 execute.rs 的通用失败路径
54/// （`ExecStatus::Failed` + 「系统暂时不可用」），不会误报成超时。
55fn get_docker() -> Result<&'static Docker, bollard::errors::Error> {
56    DOCKER_CLIENT
57        .as_ref()
58        .ok_or_else(|| bollard::errors::Error::IOError {
59            err: std::io::Error::new(
60                std::io::ErrorKind::NotFound,
61                "Docker daemon 不可用（未安装或未运行）",
62            ),
63        })
64}
65
66/// 跨运行持久化的编译缓存（Docker named volume）。
67///
68/// 语言镜像把编译缓存目录（如 Go 的 GOCACHE）指到 `mount_path`，docker.rs 把
69/// 共享 named volume 挂载到该路径，使每次运行复用上次编译产物。仅编译型语言使用。
70#[derive(Clone, Debug, PartialEq)]
71pub struct CacheVolume {
72    pub name: String,
73    pub mount_path: String,
74}
75
76pub fn build_host_config(
77    limits: &ResourceLimits,
78    cache_volume: Option<&CacheVolume>,
79) -> HostConfig {
80    let mut tmpfs = HashMap::new();
81    // /code 用 mode=1777（sticky + all-rwx）让容器内 1000:1000 用户可写。
82    // 不用 `uid=1000,gid=1000`：那是 Docker 的 tmpfs 扩展选项，Podman 报
83    // `unknown mount option "uid=1000"`。mode=1777 是 POSIX 标准 tmpfs 选项，
84    // Docker 与 Podman 都支持，语义等价（任意 UID 可写 /code）。
85    tmpfs.insert("/code".to_string(), "size=16m,mode=1777".to_string());
86    // /tmp 必须 exec：编译型语言（go/rust）把编译产物落在 /tmp 后再 exec，
87    // Docker tmpfs 默认 noexec 会让执行二进制时报 EACCES（permission denied）。
88    // 解释型语言（python/node）执行根文件系统的解释器，不受影响。
89    tmpfs.insert("/tmp".to_string(), "size=64m,mode=1777,exec".to_string());
90    tmpfs.insert("/run".to_string(), "size=16m,mode=1777".to_string());
91
92    let memory = (limits.memory_mb * 1024 * 1024) as i64;
93
94    HostConfig {
95        cpu_quota: Some((limits.cpu_cores * 100_000.0) as i64),
96        cpu_period: Some(100_000),
97        memory: Some(memory),
98        memory_swap: Some(memory), // = memory, disable swap
99        network_mode: Some(if limits.allow_network {
100            "bridge".to_string()
101        } else {
102            "none".to_string()
103        }),
104        readonly_rootfs: Some(true),
105        tmpfs: Some(tmpfs),
106        // 编译缓存 named volume：裸卷名 `name:/path` 即命名卷挂载，Docker/Podman
107        // 都会自动创建卷。binds 不受 readonly_rootfs 影响，挂载点仍可写；
108        // mount_path（如 /go-cache）是镜像内新路径，与 tmpfs 挂载无叠加。
109        binds: cache_volume.map(|cv| vec![format!("{}:{}", cv.name, cv.mount_path)]),
110        pids_limit: Some(64),
111        // 只保留 nofile（fd 数上限，语义正常）。
112        // 不设 nproc：RLIMIT_NPROC 在 setrlimit 时按 UID 计数，配合 non-root 用户会让
113        // 容器初始 exec /bin/sh 直接 EAGAIN（"exec: resource temporarily unavailable"），
114        // 与容器内实际进程数无关。pids_limit 已在 cgroup 层兜底，nproc 是冗余且有害的双重约束。
115        ulimits: Some(vec![ResourcesUlimits {
116            name: Some("nofile".to_string()),
117            soft: Some(64),
118            hard: Some(64),
119        }]),
120        cap_drop: Some(vec!["ALL".to_string()]),
121        security_opt: Some(vec!["no-new-privileges".to_string()]),
122        auto_remove: Some(false), // must be false to avoid premature removal before getting logs
123        ..Default::default()
124    }
125}
126
127struct ContainerGuard {
128    container_id: String,
129    docker: Docker,
130    cleanup_signal: Option<CleanupSignal>,
131}
132
133async fn remove_container_with_retry(docker: Docker, container_id: String) -> bool {
134    let max_attempts = 3u8;
135    let mut backoff = Duration::from_millis(200);
136    for attempt in 1..=max_attempts {
137        let remove_options = Some(RemoveContainerOptions {
138            force: true,
139            ..Default::default()
140        });
141        match docker.remove_container(&container_id, remove_options).await {
142            Ok(()) => return true,
143            Err(e) if attempt < max_attempts => {
144                tracing::warn!(
145                    attempt,
146                    max_attempts,
147                    "remove_container 失败，稍后重试: {:?}",
148                    e
149                );
150                tokio::time::sleep(backoff).await;
151                backoff *= 2;
152            }
153            Err(e) => {
154                tracing::error!(
155                    container_id = %container_id,
156                    "重试 {} 次后仍无法删除容器，可能泄漏；请手动执行 `docker rm -f {}`: {:?}",
157                    max_attempts,
158                    container_id,
159                    e
160                );
161                return false;
162            }
163        }
164    }
165    false
166}
167
168impl Drop for ContainerGuard {
169    fn drop(&mut self) {
170        let docker = self.docker.clone();
171        let container_id = self.container_id.clone();
172        let cleanup_signal = self.cleanup_signal.take();
173        // 容器清理是 fire-and-forget：调用方已返回，无法把错误回传给业务层。
174        // 因此重试几次以抵抗瞬时故障（daemon 繁忙 / socket 抖动），
175        // 仍失败则记录 error 级日志并带上 container_id，便于运维手动 `docker rm -f` 清理，
176        // 避免容器静默泄漏、长期堆积。
177        tokio::spawn(async move {
178            let cleaned = remove_container_with_retry(docker, container_id).await;
179            if let Some(signal) = cleanup_signal {
180                let _ = signal.send(cleaned);
181            }
182        });
183    }
184}
185
186pub async fn run_in_container(
187    image_name: &str,
188    run_cmd: &str,
189    source: &str,
190    ext: &str,
191    limits: ResourceLimits,
192    cache_volume: Option<&CacheVolume>,
193) -> ContainerResult {
194    run_in_container_inner(image_name, run_cmd, source, ext, limits, cache_volume, None).await
195}
196
197async fn run_in_container_inner(
198    image_name: &str,
199    run_cmd: &str,
200    source: &str,
201    ext: &str,
202    limits: ResourceLimits,
203    cache_volume: Option<&CacheVolume>,
204    cleanup_signal: Option<CleanupSignal>,
205) -> ContainerResult {
206    let docker = get_docker()?;
207    let host_config = build_host_config(&limits, cache_volume);
208
209    // Source injection script: use sh -c to first receive stdin and write to file, then exec the actual command
210    let setup_cmd = format!("cat > /code/main.{} && exec {}", ext, run_cmd);
211    let cmd = vec!["sh".to_string(), "-c".to_string(), setup_cmd];
212
213    let config = ContainerCreateBody {
214        image: Some(image_name.to_string()),
215        cmd: Some(cmd),
216        host_config: Some(host_config),
217        attach_stdin: Some(true),
218        attach_stdout: Some(true),
219        attach_stderr: Some(true),
220        open_stdin: Some(true),
221        stdin_once: Some(true),
222        user: Some("1000:1000".to_string()), // non-root user
223        working_dir: Some("/code".to_string()),
224        ..Default::default()
225    };
226
227    let container = docker
228        .create_container(None::<CreateContainerOptions>, config)
229        .await?;
230
231    let container_id = container.id;
232    let _guard = ContainerGuard {
233        container_id: container_id.clone(),
234        docker: docker.clone(),
235        cleanup_signal,
236    };
237
238    // Attach to container to stream stdin, stdout, and stderr
239    let attach_res = docker
240        .attach_container(
241            &container_id,
242            Some(AttachContainerOptions {
243                stdin: true,
244                stdout: true,
245                stderr: true,
246                stream: true,
247                logs: false,
248                ..Default::default()
249            }),
250        )
251        .await;
252
253    let (mut writer, mut stream) = match attach_res {
254        Ok(res) => (res.input, res.output),
255        Err(e) => return Err(e),
256    };
257
258    // Start container
259    docker
260        .start_container(&container_id, None::<StartContainerOptions>)
261        .await?;
262
263    // Write source code to stdin and drop/close the writer
264    use tokio::io::AsyncWriteExt;
265    let write_fut = async {
266        let _ = writer.write_all(source.as_bytes()).await;
267        let _ = writer.flush().await;
268        let _ = writer.shutdown().await;
269    };
270
271    if timeout(Duration::from_secs(5), write_fut).await.is_err() {
272        return Err(bollard::errors::Error::IOError {
273            err: std::io::Error::new(std::io::ErrorKind::TimedOut, "Writing to stdin timed out"),
274        });
275    }
276    drop(writer);
277
278    // Wait for execution with timeout control
279    let wait_future = async {
280        let mut wait_stream = docker.wait_container(&container_id, None::<WaitContainerOptions>);
281        wait_stream.next().await
282    };
283
284    let wait_res = timeout(Duration::from_secs(limits.timeout_secs), wait_future).await;
285
286    let mut timed_out = false;
287    let mut exit_code = None;
288
289    match wait_res {
290        Ok(Some(Ok(exit_status))) => {
291            exit_code = Some(exit_status.status_code);
292        }
293        Ok(_) => {} // wait error
294        Err(_) => {
295            // timeout, kill container
296            timed_out = true;
297            let _ = docker.kill_container(&container_id, None).await;
298        }
299    }
300
301    // Collect logs
302    let mut stdout_buf = Vec::new();
303    let mut stderr_buf = Vec::new();
304
305    while let Some(item) = stream.next().await {
306        match item {
307            Ok(chunk) => {
308                match chunk {
309                    LogOutput::StdOut { message } => {
310                        let remaining = (limits.output_bytes as usize)
311                            .saturating_sub(stdout_buf.len() + stderr_buf.len());
312                        if remaining > 0 {
313                            let to_add = message.len().min(remaining);
314                            stdout_buf.extend_from_slice(&message[..to_add]);
315                        }
316                    }
317                    LogOutput::StdErr { message } => {
318                        let remaining = (limits.output_bytes as usize)
319                            .saturating_sub(stdout_buf.len() + stderr_buf.len());
320                        if remaining > 0 {
321                            let to_add = message.len().min(remaining);
322                            stderr_buf.extend_from_slice(&message[..to_add]);
323                        }
324                    }
325                    _ => {}
326                }
327                if stdout_buf.len() + stderr_buf.len() >= limits.output_bytes as usize {
328                    break;
329                }
330            }
331            Err(e) => {
332                tracing::error!("Error reading container log stream: {:?}", e);
333                break;
334            }
335        }
336    }
337
338    // Check OOM status
339    let inspect = docker.inspect_container(&container_id, None).await;
340    let oom_killed = inspect
341        .ok()
342        .and_then(|info| info.state.and_then(|s| s.oom_killed))
343        .unwrap_or(false);
344
345    // Truncate output to limits.output_bytes
346    let limit_bytes = limits.output_bytes as usize;
347    let stdout_len = stdout_buf.len().min(limit_bytes);
348    let stderr_len = stderr_buf.len().min(limit_bytes);
349
350    let stdout = String::from_utf8_lossy(&stdout_buf[..stdout_len]).into_owned();
351    let stderr = String::from_utf8_lossy(&stderr_buf[..stderr_len]).into_owned();
352
353    if timed_out {
354        return Err(bollard::errors::Error::IOError {
355            err: std::io::Error::new(std::io::ErrorKind::TimedOut, "Execution timed out"),
356        });
357    }
358
359    Ok((exit_code, stdout, stderr, oom_killed))
360}
361
362#[cfg(test)]
363fn run_in_container_with_cleanup_signal<'a>(
364    image_name: &'a str,
365    run_cmd: &'a str,
366    source: &'a str,
367    ext: &'a str,
368    limits: ResourceLimits,
369    cache_volume: Option<&'a CacheVolume>,
370) -> (
371    impl Future<Output = ContainerResult> + 'a,
372    oneshot::Receiver<bool>,
373) {
374    let (sender, receiver) = oneshot::channel();
375    (
376        run_in_container_inner(
377            image_name,
378            run_cmd,
379            source,
380            ext,
381            limits,
382            cache_volume,
383            Some(sender),
384        ),
385        receiver,
386    )
387}
388
389/// 流式输出 chunk：run_in_container_stream 边读日志边推送给 SSE handler。
390///
391/// 序列化后作为 SSE event data；`Done` 同时携带终态信息（退出码 / OOM / 超时 / 耗时）。
392#[derive(Clone, Debug)]
393pub enum OutputChunk {
394    /// stdout 块（容器逐块产出）。
395    Stdout(String),
396    /// stderr 块（容器逐块产出）。
397    Stderr(String),
398    /// 终态：容器执行结束。exit_code=None 表示拿不到退出码（wait 出错）。
399    /// duration_ms = start_container 到 wait 完成的耗时。
400    Done {
401        exit_code: Option<i64>,
402        oom_killed: bool,
403        timed_out: bool,
404        duration_ms: u64,
405        /// 系统级错误消息（镜像缺失 / daemon 不可达等）。正常执行时为 None；
406        /// 由 spawn_exec_task 的 Err 分支在容器拉起失败时推送，让前端 SSE
407        /// 能拿到可操作消息，而非笼统的「连接异常」。
408        error: Option<String>,
409    },
410}
411
412/// 流式执行：与 [`run_in_container`] 相同的容器生命周期与清理（`ContainerGuard`），
413/// 但边读日志流边推 chunk 到 `tx`，同时保留完整 buffer 供调用方回填 EXEC_TASKS。
414///
415/// 与 `run_in_container` 的差异：
416/// 1. 日志循环里每块 chunk 既 `tx.send` 推流，也 append 到本地 buffer。
417/// 2. 用 `tokio::select!` 在日志读取中并发等待 `tx` 关闭——客户端断开（SSE 关闭）
418///    → `tx` 所有 Sender drop → `rx` 返回 None → 中止读取。
419/// 3. 终态推 `OutputChunk::Done` 后 return。
420///
421/// 返回完整 buffer（exit_code / stdout / stderr / oom / timed_out），供调用方写 EXEC_TASKS，
422/// 让轮询兜底路径（get_exec_result）也能拿到完整结果。
423pub async fn run_in_container_stream(
424    image_name: &str,
425    run_cmd: &str,
426    source: &str,
427    ext: &str,
428    limits: ResourceLimits,
429    cache_volume: Option<&CacheVolume>,
430    tx: tokio::sync::mpsc::Sender<OutputChunk>,
431) -> Result<(Option<i64>, String, String, bool, bool), bollard::errors::Error> {
432    let docker = get_docker()?;
433    let host_config = build_host_config(&limits, cache_volume);
434
435    // 与 run_in_container 相同的 stdin 注入脚本。
436    let setup_cmd = format!("cat > /code/main.{} && exec {}", ext, run_cmd);
437    let cmd = vec!["sh".to_string(), "-c".to_string(), setup_cmd];
438
439    let config = ContainerCreateBody {
440        image: Some(image_name.to_string()),
441        cmd: Some(cmd),
442        host_config: Some(host_config),
443        attach_stdin: Some(true),
444        attach_stdout: Some(true),
445        attach_stderr: Some(true),
446        open_stdin: Some(true),
447        stdin_once: Some(true),
448        user: Some("1000:1000".to_string()),
449        working_dir: Some("/code".to_string()),
450        ..Default::default()
451    };
452
453    let container = docker
454        .create_container(None::<CreateContainerOptions>, config)
455        .await?;
456
457    let container_id = container.id;
458    let _guard = ContainerGuard {
459        container_id: container_id.clone(),
460        docker: docker.clone(),
461        cleanup_signal: None,
462    };
463
464    // Attach 到容器的 stdin/stdout/stderr 流。
465    let attach_res = docker
466        .attach_container(
467            &container_id,
468            Some(AttachContainerOptions {
469                stdin: true,
470                stdout: true,
471                stderr: true,
472                stream: true,
473                logs: false,
474                ..Default::default()
475            }),
476        )
477        .await;
478
479    let (mut writer, mut stream) = match attach_res {
480        Ok(res) => (res.input, res.output),
481        Err(e) => return Err(e),
482    };
483
484    docker
485        .start_container(&container_id, None::<StartContainerOptions>)
486        .await?;
487
488    // 容器开始执行的时刻，用于计算 duration_ms（start_container 返回即视为起点）。
489    let start_time = std::time::Instant::now();
490
491    // 写入源码到 stdin 后关闭 writer。
492    use tokio::io::AsyncWriteExt;
493    let write_fut = async {
494        let _ = writer.write_all(source.as_bytes()).await;
495        let _ = writer.flush().await;
496        let _ = writer.shutdown().await;
497    };
498
499    if timeout(Duration::from_secs(5), write_fut).await.is_err() {
500        return Err(bollard::errors::Error::IOError {
501            err: std::io::Error::new(std::io::ErrorKind::TimedOut, "Writing to stdin timed out"),
502        });
503    }
504    drop(writer);
505
506    // —— 关键：wait_container 与日志读取必须并发，否则流式失效 ——
507    //
508    // 若先 await wait_container（等容器退出）再读日志流，wait 会阻塞到程序结束，
509    // 届时 attach stream 里已缓冲全部输出，stream.next() 一次性快速读完——
510    // 表现为"等完再一次性输出"，流式名存实亡。
511    //
512    // 用 tokio::select! 让两条分支并发：
513    // - log_reader：持续读 attach stream，每块 chunk 立即 tx.send 推流 + 累积 buffer
514    // - wait_with_timeout：等容器退出（带超时），退出后日志流自然结束（stream 返回 None）
515    // 先完成的一方触发 select 返回；若 wait 超时则 kill 容器。
516    let mut stdout_buf = Vec::new();
517    let mut stderr_buf = Vec::new();
518    let limit_bytes = limits.output_bytes as usize;
519    let mut client_disconnected = false;
520    let mut timed_out = false;
521    let mut exit_code = None;
522
523    // 日志读取循环：逐块推流 + 累积 buffer。
524    // 循环正常退出条件：stream 返回 None（容器退出后 Docker 关闭 attach 流），
525    // 或输出超限 break，或 select 被另一分支抢先完成（log_reader 被 drop）。
526    let log_reader = async {
527        while let Some(item) = stream.next().await {
528            match item {
529                Ok(chunk) => match chunk {
530                    LogOutput::StdOut { message } => {
531                        let remaining =
532                            limit_bytes.saturating_sub(stdout_buf.len() + stderr_buf.len());
533                        if remaining > 0 {
534                            let to_add = message.len().min(remaining);
535                            let slice = &message[..to_add];
536                            stdout_buf.extend_from_slice(slice);
537                            if !client_disconnected {
538                                let text = String::from_utf8_lossy(slice).into_owned();
539                                if tx.send(OutputChunk::Stdout(text)).await.is_err() {
540                                    client_disconnected = true;
541                                }
542                            }
543                        }
544                    }
545                    LogOutput::StdErr { message } => {
546                        let remaining =
547                            limit_bytes.saturating_sub(stdout_buf.len() + stderr_buf.len());
548                        if remaining > 0 {
549                            let to_add = message.len().min(remaining);
550                            let slice = &message[..to_add];
551                            stderr_buf.extend_from_slice(slice);
552                            if !client_disconnected {
553                                let text = String::from_utf8_lossy(slice).into_owned();
554                                if tx.send(OutputChunk::Stderr(text)).await.is_err() {
555                                    client_disconnected = true;
556                                }
557                            }
558                        }
559                    }
560                    _ => {}
561                },
562                Err(e) => {
563                    tracing::error!("Error reading container log stream: {:?}", e);
564                    break;
565                }
566            }
567            if stdout_buf.len() + stderr_buf.len() >= limit_bytes {
568                break;
569            }
570        }
571    };
572
573    // 带超时地等待容器退出。
574    let wait_future = async {
575        let mut wait_stream = docker.wait_container(&container_id, None::<WaitContainerOptions>);
576        wait_stream.next().await
577    };
578    let wait_with_timeout = async {
579        match timeout(Duration::from_secs(limits.timeout_secs), wait_future).await {
580            Ok(Some(Ok(exit_status))) => Some(exit_status.status_code),
581            Ok(Some(Err(_))) | Ok(None) => None, // wait error or stream ended
582            Err(_) => {
583                // 超时，杀容器。kill 后 attach stream 会被 Docker 关闭，log_reader 自然结束。
584                timed_out = true;
585                let _ = docker.kill_container(&container_id, None).await;
586                None
587            }
588        }
589    };
590
591    // 并发：哪边先完成就先用其结果。
592    // 通常 wait 先完成（容器退出 → Docker 关闭 attach stream → log_reader 也很快结束），
593    // 但若日志流先因输出超限 break，wait 会被 select drop 掉（容器仍在跑，后续 _guard 清理）。
594    tokio::select! {
595        status = wait_with_timeout => {
596            exit_code = status;
597            // 容器已退出，但 attach stream 可能还有缓冲的尾部日志。
598            // 继续读完日志流（非阻塞：stream 即将返回 None）。
599            while let Some(item) = stream.next().await {
600                if let Ok(chunk) = item {
601                    match chunk {
602                        LogOutput::StdOut { message } => {
603                            let remaining = limit_bytes.saturating_sub(stdout_buf.len() + stderr_buf.len());
604                            let to_add = message.len().min(remaining);
605                            if to_add > 0 {
606                                let slice = &message[..to_add];
607                                stdout_buf.extend_from_slice(slice);
608                                if !client_disconnected {
609                                    let text = String::from_utf8_lossy(slice).into_owned();
610                                    let _ = tx.send(OutputChunk::Stdout(text)).await;
611                                }
612                            }
613                        }
614                        LogOutput::StdErr { message } => {
615                            let remaining = limit_bytes.saturating_sub(stdout_buf.len() + stderr_buf.len());
616                            let to_add = message.len().min(remaining);
617                            if to_add > 0 {
618                                let slice = &message[..to_add];
619                                stderr_buf.extend_from_slice(slice);
620                                if !client_disconnected {
621                                    let text = String::from_utf8_lossy(slice).into_owned();
622                                    let _ = tx.send(OutputChunk::Stderr(text)).await;
623                                }
624                            }
625                        }
626                        _ => {}
627                    }
628                }
629            }
630        }
631        _ = log_reader => {
632            // 日志流先结束（输出超限 break，或 attach 断开），容器可能仍在运行。
633            // C3 修复：旧代码 wait_container().next().await 无超时无 kill，注释还谎称
634            // "短超时"。`while True: print()` 式程序会永久挂起 → RUNNER_SEMAPHORE 许可
635            // 不释放（默认 4 并发 → 4 个即 DoS 全部代码执行器）+ ContainerGuard 不 drop
636            // → 容器永久运行烧 CPU。与 wait_with_timeout 分支对称：带超时地等退出，超时则 kill。
637            let mut wait_stream = docker.wait_container(&container_id, None::<WaitContainerOptions>);
638            match timeout(Duration::from_secs(limits.timeout_secs), wait_stream.next()).await {
639                Ok(Some(Ok(status))) => exit_code = Some(status.status_code),
640                Ok(Some(Err(_))) | Ok(None) => {}
641                Err(_) => {
642                    // 超时：kill 后再等一次回收 exit_code（kill 后 wait 立即返回）。
643                    timed_out = true;
644                    let _ = docker.kill_container(&container_id, None).await;
645                    if let Some(Ok(status)) = wait_stream.next().await {
646                        exit_code = Some(status.status_code);
647                    }
648                }
649            }
650        }
651    }
652
653    // 检查 OOM 状态。
654    let inspect = docker.inspect_container(&container_id, None).await;
655    let oom_killed = inspect
656        .ok()
657        .and_then(|info| info.state.and_then(|s| s.oom_killed))
658        .unwrap_or(false);
659    // 推送终态 chunk（客户端已断开则跳过，send 必然失败）。
660    let duration_ms = start_time.elapsed().as_millis() as u64;
661    if !client_disconnected {
662        let _ = tx
663            .send(OutputChunk::Done {
664                exit_code,
665                oom_killed,
666                timed_out,
667                duration_ms,
668                error: None,
669            })
670            .await;
671    }
672
673    let stdout_len = stdout_buf.len().min(limit_bytes);
674    let stderr_len = stderr_buf.len().min(limit_bytes);
675    let stdout = String::from_utf8_lossy(&stdout_buf[..stdout_len]).into_owned();
676    let stderr = String::from_utf8_lossy(&stderr_buf[..stderr_len]).into_owned();
677
678    if timed_out {
679        return Err(bollard::errors::Error::IOError {
680            err: std::io::Error::new(std::io::ErrorKind::TimedOut, "Execution timed out"),
681        });
682    }
683
684    Ok((exit_code, stdout, stderr, oom_killed, timed_out))
685}
686
687#[cfg(test)]
688mod tests {
689    use super::*;
690    use crate::infra::runner_config::ResourceLimits;
691
692    #[test]
693    fn test_host_config_generation() {
694        let limits = ResourceLimits {
695            cpu_cores: 1.5,
696            memory_mb: 256,
697            timeout_secs: 5,
698            output_bytes: 1024,
699            allow_network: false,
700        };
701        let host_config = build_host_config(&limits, None);
702        assert_eq!(host_config.cpu_quota, Some(150_000));
703        assert_eq!(host_config.memory, Some(256 * 1024 * 1024));
704        assert_eq!(host_config.readonly_rootfs, Some(true));
705        assert_eq!(host_config.network_mode.as_deref(), Some("none"));
706
707        // 隔离不变量：这些配置若被误删，容器隔离会悄悄失效。
708        // 单独断言而非隐含在端到端测试里，确保回归可见。
709        assert_eq!(
710            host_config.cap_drop.as_deref(),
711            Some(&["ALL".to_string()][..])
712        );
713        assert_eq!(
714            host_config.security_opt.as_deref(),
715            Some(&["no-new-privileges".to_string()][..])
716        );
717        assert_eq!(host_config.pids_limit, Some(64));
718        // memory_swap == memory 才是真正禁用 swap；二者不等价于可换出到磁盘
719        assert_eq!(host_config.memory_swap, host_config.memory);
720        assert_eq!(host_config.memory_swap, Some(256 * 1024 * 1024));
721    }
722
723    #[test]
724    fn host_config_cache_volume_binds() {
725        // 编译缓存 named volume：None 时不挂载任何 binds；Some 时以裸卷名
726        // `name:/path` 挂载（Docker/Podman 自动创建卷，不受 readonly_rootfs 影响）。
727        let limits = ResourceLimits {
728            cpu_cores: 1.0,
729            memory_mb: 128,
730            timeout_secs: 5,
731            output_bytes: 1024,
732            allow_network: false,
733        };
734        assert_eq!(build_host_config(&limits, None).binds, None);
735
736        let cache = CacheVolume {
737            name: "yggdrasil-gocache".to_string(),
738            mount_path: "/go-cache".to_string(),
739        };
740        assert_eq!(
741            build_host_config(&limits, Some(&cache)).binds,
742            Some(vec!["yggdrasil-gocache:/go-cache".to_string()])
743        );
744    }
745
746    #[test]
747    fn host_config_ulimits_drops_only_nofile() {
748        // nproc 在容器内对 non-root 用户有害（初始 exec /bin/sh 直接 EAGAIN），
749        // 只保留 nofile=64。若有人加回 nproc，这里会失败。
750        let limits = ResourceLimits {
751            cpu_cores: 1.0,
752            memory_mb: 128,
753            timeout_secs: 5,
754            output_bytes: 1024,
755            allow_network: false,
756        };
757        let host_config = build_host_config(&limits, None);
758        let ulimits = host_config.ulimits.expect("ulimits must be set");
759        assert_eq!(ulimits.len(), 1, "only nofile, no nproc");
760        let nf = &ulimits[0];
761        assert_eq!(nf.name.as_deref(), Some("nofile"));
762        assert_eq!(nf.soft, Some(64));
763        assert_eq!(nf.hard, Some(64));
764    }
765
766    #[test]
767    fn host_config_tmpfs_has_exec_on_tmp_only() {
768        // /tmp 必须 exec（编译型语言把产物落在 /tmp 后再 exec）；
769        // /code、/run 不带 exec（默认 noexec）。若误给 /code 也加 exec，测试失败。
770        let limits = ResourceLimits {
771            cpu_cores: 1.0,
772            memory_mb: 128,
773            timeout_secs: 5,
774            output_bytes: 1024,
775            allow_network: false,
776        };
777        let host_config = build_host_config(&limits, None);
778        let tmpfs = host_config.tmpfs.expect("tmpfs must be set");
779        assert_eq!(tmpfs.len(), 3);
780        assert!(tmpfs["/tmp"].contains("exec"));
781        assert!(!tmpfs["/code"].contains("exec"));
782        assert!(!tmpfs["/run"].contains("exec"));
783    }
784
785    #[test]
786    fn host_config_network_mode_follows_allow_network() {
787        let base = ResourceLimits {
788            cpu_cores: 1.0,
789            memory_mb: 128,
790            timeout_secs: 5,
791            output_bytes: 1024,
792            allow_network: false,
793        };
794        assert_eq!(
795            build_host_config(&base, None).network_mode.as_deref(),
796            Some("none")
797        );
798        let mut net = base;
799        net.allow_network = true;
800        assert_eq!(
801            build_host_config(&net, None).network_mode.as_deref(),
802            Some("bridge")
803        );
804    }
805
806    #[test]
807    fn host_config_cpu_and_memory_scale_with_limits() {
808        // 资源钳制是安全边界，断言换算公式不漂移（cpu_cores * 100_000，memory_mb * 1MiB）。
809        let limits = ResourceLimits {
810            cpu_cores: 2.0,
811            memory_mb: 512,
812            timeout_secs: 5,
813            output_bytes: 1024,
814            allow_network: false,
815        };
816        let host_config = build_host_config(&limits, None);
817        assert_eq!(host_config.cpu_quota, Some(200_000));
818        assert_eq!(host_config.cpu_period, Some(100_000));
819        assert_eq!(host_config.memory, Some(512 * 1024 * 1024));
820        assert_eq!(host_config.memory_swap, Some(512 * 1024 * 1024));
821    }
822
823    /// 探测 Docker daemon 是否可用：socket 缺失（`DOCKER_CLIENT == None`）、
824    /// daemon 无响应、或所需镜像不在本地时返回 `None`。集成测试用它做动态
825    /// 守卫——daemon 与镜像都在则跑，否则显式跳过（eprintln + 提前 return），
826    /// 而非 panic 失败。
827    ///
828    /// 这比 `#[ignore]` 更合适：`#[ignore]` 会在所有环境（含有 Docker 的开发机）
829    /// 一律跳过、且需 `cargo test --ignored` 显式触发；探测守卫让这些测试在
830    /// 有 Docker 时自动运行、在无 Docker 或缺镜像的 CI 上静默放行。
831    ///
832    /// 镜像探测：GitHub runner 有 daemon 但默认不带 alpine:latest，创建容器
833    /// 时会 404。这里用 inspect_image 只读检查本地是否已有镜像，缺失即跳过，
834    /// 不在此 pull（CI 不应依赖外网拉镜像——慢且脆弱）。
835    async fn require_docker_with_image(image: &str) -> Option<&'static Docker> {
836        let docker = DOCKER_CLIENT.as_ref()?;
837        // socket 在但 daemon 挂了的情况靠这一步轻量只读调用兜住。
838        docker.version().await.ok()?;
839        // 镜像不在本地 → 视为环境不可用，跳过而非 panic。
840        docker.inspect_image(image).await.ok()?;
841        Some(docker)
842    }
843
844    #[tokio::test]
845    #[serial_test::serial]
846    async fn test_run_in_container_success() {
847        if require_docker_with_image("alpine:latest").await.is_none() {
848            eprintln!("skip: Docker daemon 不可用或缺少 alpine:latest");
849            return;
850        }
851        let limits = ResourceLimits {
852            cpu_cores: 1.0,
853            memory_mb: 128,
854            timeout_secs: 5,
855            output_bytes: 1024,
856            allow_network: false,
857        };
858        let (exit_code, stdout, stderr, oom_killed) = run_in_container(
859            "alpine:latest",
860            "cat /code/main.txt",
861            "hello world",
862            "txt",
863            limits,
864            None,
865        )
866        .await
867        .unwrap();
868
869        assert_eq!(exit_code, Some(0));
870        assert_eq!(stdout, "hello world");
871        assert!(stderr.is_empty());
872        assert!(!oom_killed);
873    }
874
875    #[tokio::test]
876    #[serial_test::serial]
877    async fn test_run_in_container_output_truncation() {
878        if require_docker_with_image("alpine:latest").await.is_none() {
879            eprintln!("skip: Docker daemon 不可用或缺少 alpine:latest");
880            return;
881        }
882        let limits = ResourceLimits {
883            cpu_cores: 1.0,
884            memory_mb: 128,
885            timeout_secs: 5,
886            output_bytes: 5,
887            allow_network: false,
888        };
889        let (exit_code, stdout, stderr, oom_killed) = run_in_container(
890            "alpine:latest",
891            "cat /code/main.txt",
892            "hello world",
893            "txt",
894            limits,
895            None,
896        )
897        .await
898        .unwrap();
899
900        assert_eq!(exit_code, Some(0));
901        assert_eq!(stdout, "hello");
902        assert!(stderr.is_empty());
903        assert!(!oom_killed);
904    }
905
906    #[tokio::test]
907    #[serial_test::serial]
908    async fn test_run_in_container_timeout() {
909        if require_docker_with_image("alpine:latest").await.is_none() {
910            eprintln!("skip: Docker daemon 不可用或缺少 alpine:latest");
911            return;
912        }
913        let limits = ResourceLimits {
914            cpu_cores: 1.0,
915            memory_mb: 128,
916            timeout_secs: 1,
917            output_bytes: 1024,
918            allow_network: false,
919        };
920        let res = run_in_container("alpine:latest", "sleep 10", "", "txt", limits, None).await;
921
922        assert!(res.is_err());
923        let err = res.unwrap_err();
924        match err {
925            bollard::errors::Error::IOError { err } => {
926                assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
927            }
928            _ => panic!("Expected IOError(TimedOut), got {:?}", err),
929        }
930    }
931
932    #[tokio::test]
933    #[serial_test::serial]
934    async fn test_run_in_container_cancellation() {
935        use bollard::query_parameters::ListContainersOptions;
936        let docker = match require_docker_with_image("alpine:latest").await {
937            Some(d) => d,
938            None => {
939                eprintln!("skip: Docker daemon 不可用或缺少 alpine:latest");
940                return;
941            }
942        };
943
944        let before = docker
945            .list_containers(Some(ListContainersOptions {
946                all: true,
947                ..Default::default()
948            }))
949            .await
950            .unwrap();
951        let before_ids: std::collections::HashSet<String> =
952            before.into_iter().map(|c| c.id.unwrap()).collect();
953
954        let limits = ResourceLimits {
955            cpu_cores: 1.0,
956            memory_mb: 128,
957            timeout_secs: 10,
958            output_bytes: 1024,
959            allow_network: false,
960        };
961
962        let (run_fut, cleanup_done) = run_in_container_with_cleanup_signal(
963            "alpine:latest",
964            "sleep 100",
965            "",
966            "txt",
967            limits,
968            None,
969        );
970
971        tokio::select! {
972            _ = run_fut => {
973                panic!("Should have been cancelled");
974            }
975            _ = tokio::time::sleep(Duration::from_secs(1)) => {
976                // Cancelled!
977            }
978        }
979
980        // 取消发生后，Drop 仍通过后台任务删除容器。等待清理任务明确报告结果，
981        // 不再用固定时间窗口猜测 Docker daemon 是否已经完成删除。
982        match timeout(Duration::from_secs(30), cleanup_done).await {
983            Ok(Ok(true)) => {}
984            Ok(Ok(false)) => panic!("Container cleanup failed after retries"),
985            Ok(Err(_)) => {
986                // Future 在 create_container 完成前就被取消，没有容器需要清理。
987                return;
988            }
989            Err(_) => panic!("Container cleanup did not finish within 30 seconds"),
990        }
991
992        // remove_container 返回成功后，给 daemon 很短的时间同步 list API，避免
993        // 删除成功与列表可见性之间的最终一致性造成误报。
994        let deadline = Duration::from_secs(5);
995        let poll_interval = Duration::from_millis(250);
996        let start = std::time::Instant::now();
997
998        let leaked = loop {
999            tokio::time::sleep(poll_interval).await;
1000
1001            let after = docker
1002                .list_containers(Some(ListContainersOptions {
1003                    all: true,
1004                    ..Default::default()
1005                }))
1006                .await
1007                .unwrap();
1008
1009            let leaked: Vec<String> = after
1010                .into_iter()
1011                .filter(|c| {
1012                    c.id.as_ref().is_some_and(|id| !before_ids.contains(id))
1013                        && c.image.as_deref() == Some("alpine:latest")
1014                })
1015                .map(|c| c.id.unwrap())
1016                .collect();
1017
1018            if leaked.is_empty() || start.elapsed() >= deadline {
1019                break leaked;
1020            }
1021        };
1022
1023        assert_eq!(
1024            leaked.len(),
1025            0,
1026            "Found {} leaked containers after confirmed cleanup",
1027            leaked.len()
1028        );
1029    }
1030}
yggdrasil/infra/docker.rs

yggdrasil/infra/
docker.rs
