codex-app-server/cloud/tests/test_tenant_isolation.py

"""多租户隔离单测。

只用临时目录与假 runner，不跑真实 Docker、不联网。
覆盖：租户目录播种（codex-home 只取 config.toml，绝不
复制 thread_history）、挂载按项目分开、网络名安全化、
执行器 per-call 挂载/网络覆盖、网络创建失败不降级。
"""

import os
import sys
import tempfile
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:
    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="", stderr="")
    return fake


class TenantDirsTestCase(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)
        # 模板 home：config.toml + 绝不能复制的 thread_history
        # 与假密钥文件；模板 repo：sample.py + .git。
        self.template_home = os.path.join(self.tmp.name, "tpl-home")
        os.makedirs(self.template_home)
        open(os.path.join(self.template_home, "config.toml"),
             "w").write('[model_provider.xcmapi]\nbase_url="https://xcmapi.org/v1"\n')
        open(os.path.join(self.template_home, "thread_history_1.sqlite"),
             "w").write("OTHER TENANT HISTORY")
        open(os.path.join(self.template_home, "auth.json"),
             "w").write("SECRET")
        self.template_repo = os.path.join(self.tmp.name, "tpl-repo")
        os.makedirs(os.path.join(self.template_repo, ".git"))
        open(os.path.join(self.template_repo, "sample.py"),
             "w").write("def add(a, b):\n    return a + b\n")

    def tearDown(self):
        self.tmp.cleanup()

    def _worker(self, **kwargs):
        kwargs.setdefault("env", {"XCM_API_KEY": "test-placeholder"})
        kwargs.setdefault("tenants_root",
                          os.path.join(self.tmp.name, "tenants"))
        kwargs.setdefault("tenant_template_home", self.template_home)
        kwargs.setdefault("tenant_template_repo", self.template_repo)
        kwargs.setdefault("tenant_isolation", True)
        return Worker(self.db_path, "/fake/host.py", **kwargs)

    def test_dirs_seed_config_only_not_history(self):
        worker = self._worker()
        dirs = worker._tenant_dirs("projA")
        home = dirs["codex_home"]
        self.assertTrue(os.path.exists(os.path.join(home, "config.toml")))
        self.assertFalse(os.path.exists(
            os.path.join(home, "thread_history_1.sqlite")))
        self.assertFalse(os.path.exists(os.path.join(home, "auth.json")))
        # test-repo 副本有 sample.py，且是复制不是链接。
        self.assertTrue(os.path.exists(
            os.path.join(dirs["test_repo"], "sample.py")))
        # 第二次调用不覆盖租户自己的改动。
        marker = os.path.join(dirs["test_repo"], "mine.txt")
        open(marker, "w").write("mine")
        worker._tenant_dirs("projA")
        self.assertTrue(os.path.exists(marker))

    def test_mounts_are_per_project(self):
        worker = self._worker()
        mounts_a = worker._tenant_mounts("projA")
        mounts_b = worker._tenant_mounts("projB")
        # 容器内路径相同（/app/work 等），宿主源不同项目不同。
        cont_a = [m.split(":")[1] for m in mounts_a]
        cont_b = [m.split(":")[1] for m in mounts_b]
        self.assertEqual(cont_a, cont_b)
        self.assertIn("/app/work", cont_a)
        self.assertIn("/var/lib/codex-app-server/codex-home", cont_a)
        hosts_a = {m.split(":")[0] for m in mounts_a}
        hosts_b = {m.split(":")[0] for m in mounts_b}
        self.assertFalse(hosts_a & hosts_b)
        # 租户 A 的 work 里写文件，B 的挂载里看不到（宿主侧）。
        open(os.path.join(worker._tenant_dirs("projA")["work"],
                          "secret-a.txt"), "w").write("SECRET-A")
        self.assertFalse(os.path.exists(os.path.join(
            worker._tenant_dirs("projB")["work"], "secret-a.txt")))

    def test_network_name_sanitized(self):
        worker = self._worker()
        self.assertEqual(worker._tenant_network("abc123"),
                         "codex-abc123")
        self.assertEqual(worker._tenant_network("a b/c"),
                         "codex-a-b-c")
        with self.assertRaises(ValueError):
            worker._tenant_network("...")

    def test_template_home_without_config_fails_loudly(self):
        empty_home = os.path.join(self.tmp.name, "empty-home")
        os.makedirs(empty_home)
        worker = self._worker(tenant_template_home=empty_home)
        with self.assertRaises(RuntimeError):
            worker._tenant_dirs("projA")


class DockerTenantCommandTestCase(unittest.TestCase):
    def test_per_call_mounts_network_and_drop_caps(self):
        seen = []
        static_dir = tempfile.mkdtemp()
        tenant_dir = tempfile.mkdtemp()
        with mock.patch("cloud.sandbox.shutil.which",
                        return_value="/usr/bin/docker"):
            ex = DockerExecutor(
                "/fake/host.py",
                runner=_ok_runner_factory(seen),
                env={"XCM_API_KEY": "test-placeholder"},
                extra_mounts=["%s:/app/work:rw" % static_dir])
            cmd = ex.build_run_command(
                {"id": "t1", "worktree": "/tmp/wt"},
                ["echo", "hi"],
                extra_mounts=["%s:/app/work:rw" % tenant_dir],
                network="codex-projA")
        mounts = [cmd[i + 1] for i, a in enumerate(cmd) if a == "-v"]
        # per-call 挂载替换静态挂载，不是追加。
        self.assertIn("%s:/app/work:rw" % tenant_dir, mounts)
        self.assertNotIn("%s:/app/work:rw" % static_dir, mounts)
        self.assertEqual(cmd[cmd.index("--network") + 1], "codex-projA")
        self.assertIn("--cap-drop=ALL", cmd)
        self.assertIn("--security-opt=no-new-privileges", cmd)

    def test_ensure_network_already_exists_is_ok(self):
        def exists_runner(cmd, timeout=None):
            return SimpleNamespace(
                returncode=1, stdout="",
                stderr="Error response from daemon: network with name "
                       "codex-p already exists")
        with mock.patch("cloud.sandbox.shutil.which",
                        return_value="/usr/bin/docker"):
            ex = DockerExecutor("/fake/host.py", runner=exists_runner,
                                env={"XCM_API_KEY": "test-placeholder"})
            self.assertTrue(ex.ensure_network("codex-p"))

    def test_ensure_network_real_failure_is_false(self):
        def fail_runner(cmd, timeout=None):
            return SimpleNamespace(returncode=1, stdout="",
                                   stderr="daemon not reachable")
        with mock.patch("cloud.sandbox.shutil.which",
                        return_value="/usr/bin/docker"):
            ex = DockerExecutor("/fake/host.py", runner=fail_runner,
                                env={"XCM_API_KEY": "test-placeholder"})
            self.assertFalse(ex.ensure_network("codex-p"))


if __name__ == "__main__":
    unittest.main()