421 lines
18 KiB
Python
421 lines
18 KiB
Python
"""Docker 沙箱执行器单测：命令构造、安全挂载、超时与取消。
|
||
|
||
运行（在 ~/workspace/codex-app-server 下）：
|
||
python3 -m unittest discover -s cloud/tests -p "test_docker_sandbox.py" -v
|
||
|
||
说明：全部用注入的假 docker runner，不调用真实 docker、不联网、
|
||
不跑真实模型回合；env 映射里的值只是占位假名，绝不是真实密钥。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.sandbox import DockerExecutor
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.sandbox import DockerExecutor
|
||
from cloud.server import init_db
|
||
from cloud.worker import Worker
|
||
|
||
|
||
def _ok_runner_factory(seen):
|
||
def fake(cmd, timeout=None):
|
||
seen.append(list(cmd))
|
||
return SimpleNamespace(
|
||
returncode=0,
|
||
stdout=json.dumps({"status": "completed", "text": "done"}),
|
||
stderr="")
|
||
return fake
|
||
|
||
|
||
def _insert_task(db_path, task_id, prompt="do work"):
|
||
now = time.time()
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" idempotency_key, depends_on, priority, created_by,"
|
||
" created_at, updated_at, last_heartbeat)"
|
||
" VALUES(?, 'p1', ?, 'queued', NULL, '[]', 0, 'tester', ?, ?, NULL)",
|
||
(task_id, prompt, now, now),
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _task_status(db_path, task_id):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT status, error_text FROM tasks WHERE id = ?",
|
||
(task_id,)).fetchone()
|
||
finally:
|
||
conn.close()
|
||
return row
|
||
|
||
|
||
class DockerCommandTestCase(unittest.TestCase):
|
||
def _executor(self, seen, **kwargs):
|
||
kwargs.setdefault("runner", _ok_runner_factory(seen))
|
||
kwargs.setdefault("env", {"XCM_API_KEY": "placeholder-not-a-real-key"})
|
||
return DockerExecutor("/fake/host.py", **kwargs)
|
||
|
||
def test_command_mounts_only_task_worktree_and_ro_code(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt-t1"},
|
||
["python3", "/app/host.py", "status", "--task", "t1"])
|
||
mounts = [cmd[i + 1] for i, a in enumerate(cmd) if a == "-v"]
|
||
self.assertEqual(mounts, ["/tmp/wt-t1:/work", "/fake:/app:ro"])
|
||
# 绝不挂载宿主根目录。
|
||
self.assertNotIn("/:/work", mounts)
|
||
self.assertNotIn("/:/app:ro", mounts)
|
||
self.assertIn("--rm", cmd)
|
||
self.assertIn("/tmp/wt-t1:/work", cmd)
|
||
|
||
def test_command_resource_limits_and_non_root_user(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen, cpus="1", memory="512m",
|
||
user="1000:1000")
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt"},
|
||
["echo", "hi"])
|
||
self.assertEqual(cmd[cmd.index("--cpus") + 1], "1")
|
||
self.assertEqual(cmd[cmd.index("--memory") + 1], "512m")
|
||
self.assertEqual(cmd[cmd.index("--user") + 1], "1000:1000")
|
||
self.assertNotEqual(cmd[cmd.index("--user") + 1], "0")
|
||
self.assertNotEqual(cmd[cmd.index("--user") + 1], "root")
|
||
|
||
def test_command_defaults_are_2cpu_2g(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
|
||
self.assertEqual(cmd[cmd.index("--cpus") + 1], "2")
|
||
self.assertEqual(cmd[cmd.index("--memory") + 1], "2g")
|
||
|
||
def test_container_name_sanitized(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
self.assertEqual(ex.container_name({"id": "a b/c"}),
|
||
"codex-a-b-c")
|
||
cmd = ex.build_run_command(
|
||
{"id": "a b/c", "worktree": "/tmp/wt"}, ["echo", "hi"])
|
||
self.assertIn("codex-a-b-c", cmd)
|
||
self.assertEqual(cmd[cmd.index("--name") + 1], "codex-a-b-c")
|
||
|
||
def test_env_name_only_no_value_in_command(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
|
||
joined = " ".join(cmd)
|
||
self.assertIn("-e", cmd)
|
||
self.assertIn("XCM_API_KEY", cmd)
|
||
# 值绝不出现在命令列表里（-e 只传名，docker 从宿主环境取值）。
|
||
self.assertNotIn("placeholder-not-a-real-key", joined)
|
||
self.assertNotIn("placeholder-not-a-real-key=", joined)
|
||
|
||
def test_missing_key_does_not_start_container(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py",
|
||
runner=_ok_runner_factory(seen), env={})
|
||
with self.assertRaises(RuntimeError) as ctx:
|
||
ex.dispatch_task({"id": "t1", "prompt": "hi",
|
||
"worktree": "/tmp/wt"})
|
||
self.assertIn("Missing environment variable", str(ctx.exception))
|
||
self.assertEqual(seen, [])
|
||
|
||
def test_missing_docker_does_not_call_runner(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which", return_value=None):
|
||
ex = self._executor(seen)
|
||
with self.assertRaises(RuntimeError) as ctx:
|
||
ex.dispatch_task({"id": "t1", "prompt": "hi",
|
||
"worktree": "/tmp/wt"})
|
||
self.assertIn("docker not available", str(ctx.exception))
|
||
self.assertEqual(seen, [])
|
||
|
||
def test_dispatch_and_run_shapes(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
dispatched = ex.dispatch_task(
|
||
{"id": "t1", "prompt": "do it", "worktree": "/tmp/wt"})
|
||
ran = ex.run_task(
|
||
{"id": "t1", "prompt": "again", "worktree": "/tmp/wt"})
|
||
self.assertEqual(dispatched["status"], "completed")
|
||
self.assertEqual(dispatched["executor"], "docker")
|
||
dispatch_cmd = seen[0]
|
||
self.assertIn("dispatch", dispatch_cmd)
|
||
self.assertIn("t1", dispatch_cmd)
|
||
self.assertIn("do it", dispatch_cmd)
|
||
run_cmd = seen[1]
|
||
self.assertIn("run", run_cmd)
|
||
self.assertIn("again", run_cmd)
|
||
|
||
def test_execute_is_readonly_status_not_dispatch(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
result = ex.execute({"id": "t1", "worktree": "/tmp/wt"})
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertIn("status", seen[0])
|
||
self.assertNotIn("dispatch", seen[0])
|
||
# run 在命令里会出现 "--name" 等，但内层命令不应是 run：
|
||
# 内层 host.py 命令的第一个动词是 status。
|
||
inner = seen[0][seen[0].index("codex-app-server:latest") + 1:]
|
||
self.assertEqual(inner[2], "status")
|
||
|
||
def test_timeout_kills_and_removes_container(self):
|
||
seen = []
|
||
|
||
def fake(cmd, timeout=None):
|
||
seen.append(list(cmd))
|
||
if len(cmd) > 1 and cmd[1] == "run":
|
||
raise subprocess.TimeoutExpired(cmd, timeout)
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py", runner=fake,
|
||
env={"XCM_API_KEY": "placeholder"},
|
||
timeout=5)
|
||
result = ex.dispatch_task(
|
||
{"id": "t9", "prompt": "slow", "worktree": "/tmp/wt"})
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("timeout", result["stdout"])
|
||
kinds = [(c[1], c[2]) for c in seen if len(c) > 2]
|
||
self.assertIn(("kill", "codex-t9"), kinds)
|
||
self.assertIn(("rm", "-f"), [(c[1], c[2]) for c in seen
|
||
if len(c) > 2])
|
||
|
||
def test_timeout_keeps_partial_stdout_and_container_logs(self):
|
||
# 旧超时路径只剩一句 'docker run timeout'：TimeoutExpired
|
||
# 自带的 partial stdout 与 kill 后抓的 docker logs 尾都
|
||
# 被丢了。新路径：先 kill 不 rm，抓 logs，再 rm。
|
||
seen = []
|
||
|
||
def fake(cmd, timeout=None):
|
||
seen.append(list(cmd))
|
||
if len(cmd) > 1 and cmd[1] == "run":
|
||
exc = subprocess.TimeoutExpired(cmd, timeout)
|
||
exc.stdout = "partial-host-output"
|
||
exc.stderr = "partial-err"
|
||
raise exc
|
||
if len(cmd) > 1 and cmd[1] == "logs":
|
||
return SimpleNamespace(returncode=0,
|
||
stdout="container-tail-log",
|
||
stderr="")
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py", runner=fake,
|
||
env={"XCM_API_KEY": "placeholder"},
|
||
timeout=5)
|
||
result = ex.dispatch_task(
|
||
{"id": "t10", "prompt": "slow", "worktree": "/tmp/wt"})
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("partial-host-output", result["stdout"])
|
||
# partial 已有内容时 stdout 优先用 partial，logs 另存字段。
|
||
self.assertEqual(result["container_log_tail"],
|
||
"container-tail-log")
|
||
# 顺序：kill → logs → rm，logs 必须在 rm 之前抓到。
|
||
kinds = [c[1] for c in seen if len(c) > 1]
|
||
self.assertLess(kinds.index("logs"), kinds.index("rm"))
|
||
|
||
def test_kill_container_calls_kill_and_rm(self):
|
||
seen = []
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor(seen)
|
||
ok = ex.kill_container({"id": "t2"})
|
||
self.assertTrue(ok)
|
||
self.assertEqual([c[1] for c in seen], ["kill", "rm"])
|
||
self.assertEqual(seen[0][2], "codex-t2")
|
||
self.assertEqual(seen[1][2], "-f")
|
||
self.assertEqual(seen[1][3], "codex-t2")
|
||
|
||
|
||
class WorkerDockerTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
self.env = {"XCM_API_KEY": "placeholder-not-a-real-key"}
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def test_worker_docker_missing_fails_task_with_docker_error(self):
|
||
seen = []
|
||
|
||
def fake(cmd, timeout=None):
|
||
seen.append(list(cmd))
|
||
return SimpleNamespace(returncode=0, stdout="{}", stderr="")
|
||
|
||
_insert_task(self.db_path, "d1", prompt="do it")
|
||
with mock.patch("cloud.sandbox.shutil.which", return_value=None):
|
||
ex = DockerExecutor("/fake/host.py", runner=fake, env=self.env)
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
|
||
result = worker.run_one()
|
||
self.assertTrue(result["dispatched"])
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertIn("docker", result["error"])
|
||
self.assertEqual(seen, [])
|
||
status, error_text = _task_status(self.db_path, "d1")
|
||
self.assertEqual(status, "failed")
|
||
self.assertIn("docker", error_text)
|
||
|
||
def test_worker_docker_success(self):
|
||
seen = []
|
||
_insert_task(self.db_path, "d2", prompt="do it")
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py",
|
||
runner=_ok_runner_factory(seen),
|
||
env=self.env)
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
|
||
result = worker.run_one()
|
||
self.assertTrue(result["dispatched"])
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(result["executor"], "docker")
|
||
self.assertTrue(any(c[1] == "run" for c in seen))
|
||
|
||
def test_worker_cancel_kills_docker_container(self):
|
||
seen = []
|
||
_insert_task(self.db_path, "d3", prompt="do it")
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py",
|
||
runner=_ok_runner_factory(seen),
|
||
env=self.env)
|
||
worker = Worker(self.db_path, "/fake/host.py",
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"}, executor=ex)
|
||
# 模拟 dispatch 阻塞期间：worker 已把容器名记下来。
|
||
worker._running_set("d3", container="codex-d3")
|
||
conn = sqlite3.connect(self.db_path)
|
||
try:
|
||
conn.execute(
|
||
"UPDATE tasks SET status = 'running' WHERE id = 'd3'")
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
self.assertTrue(worker.cancel("d3"))
|
||
self.assertEqual([c[1] for c in seen], ["kill", "rm"])
|
||
status, _ = _task_status(self.db_path, "d3")
|
||
self.assertEqual(status, "interrupted")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|
||
|
||
class TestDockerExtraMounts(DockerCommandTestCase):
|
||
def test_extra_mounts_appended_after_base_mounts(self):
|
||
with tempfile.TemporaryDirectory() as td:
|
||
work = os.path.join(td, "work")
|
||
os.makedirs(work)
|
||
state = os.path.join(td, "host-state.json")
|
||
open(state, "w").write("{}")
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py", runner=_ok_runner_factory([]),
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"},
|
||
extra_mounts=[f"{work}:/app/work", f"{state}:/app/host-state.json"])
|
||
cmd = ex.build_run_command({"id": "t1"}, ["python3", "/app/host.py", "status"])
|
||
# 基础挂载在前，额外挂载在后，且 /app/work 覆盖只读子路径
|
||
self.assertIn(f"{work}:/app/work:rw", cmd)
|
||
self.assertIn(f"{state}:/app/host-state.json:rw", cmd)
|
||
mounts = [cmd[i + 1] for i, a in enumerate(cmd) if a == "-v"]
|
||
self.assertEqual(mounts[1], "/fake:/app:ro")
|
||
self.assertIn(f"{work}:/app/work:rw", mounts)
|
||
self.assertLess(mounts.index("/fake:/app:ro"), mounts.index(f"{work}:/app/work:rw"))
|
||
|
||
def test_tmpfs_tmp_and_readonly_root_flag(self):
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = self._executor([])
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
|
||
self.assertIn("--tmpfs", cmd)
|
||
self.assertIn("/tmp", cmd)
|
||
self.assertIn("/var/tmp", cmd)
|
||
self.assertNotIn("--read-only", cmd)
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
ex = DockerExecutor("/fake/host.py", runner=_ok_runner_factory([]),
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"},
|
||
readonly_root=True)
|
||
cmd = ex.build_run_command(
|
||
{"id": "t1", "worktree": "/tmp/wt"}, ["echo", "hi"])
|
||
self.assertIn("--read-only", cmd)
|
||
with self.assertRaises(ValueError):
|
||
DockerExecutor("/fake/host.py",
|
||
extra_mounts=["/tmp/host-tmp:/tmp"])
|
||
|
||
def test_extra_mounts_reject_dangerous_paths(self):
|
||
with self.assertRaises(ValueError):
|
||
DockerExecutor("/fake/host.py", extra_mounts=["/: /x"])
|
||
with self.assertRaises(ValueError):
|
||
DockerExecutor("/fake/host.py", extra_mounts=["/etc/codex-app-server/env:/x"])
|
||
with self.assertRaises(ValueError):
|
||
DockerExecutor("/fake/host.py", extra_mounts=["/tmp:/app"])
|
||
with self.assertRaises(ValueError):
|
||
DockerExecutor("/fake/host.py", extra_mounts=["/nonexistent-xyz:/x"])
|
||
|
||
|
||
class TestDockerDefaultRunnerEnv(unittest.TestCase):
|
||
def test_env_merged_not_replaced(self):
|
||
from cloud.sandbox import _docker_default_runner
|
||
out = _docker_default_runner(
|
||
["python3", "-c",
|
||
"import os,sys;sys.stdout.write("
|
||
"os.environ.get('PROBE_X','')+':'+"
|
||
"('1' if os.environ.get('PATH') else '0'))"],
|
||
env={"PROBE_X": "42"})
|
||
self.assertEqual(out.returncode, 0)
|
||
self.assertEqual(out.stdout, "42:1")
|
||
|
||
def test_call_runner_passes_env_to_default_runner(self):
|
||
seen = []
|
||
|
||
def recording(cmd, timeout=None, env=None):
|
||
seen.append(env)
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
|
||
ex = DockerExecutor("/fake/host.py", runner=recording,
|
||
env={"XCM_API_KEY": "placeholder-not-a-real-key"})
|
||
ex._call_runner(["docker", "version"], env={"SUB2API_API_KEY": "x"})
|
||
self.assertEqual(seen, [{"SUB2API_API_KEY": "x"}])
|