DerekSurge
191 lines
8.2 KiB
Python
191 lines
8.2 KiB
Python
"""Tests for core.setup — copies uploaded context files into the workdir."""
|
|
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from core import setup as setup_mod # aliased to avoid clash with pytest's xunit setup_module fixture
|
|
from core.tools import sandbox
|
|
|
|
|
|
class SetupTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = Path(tempfile.mkdtemp())
|
|
self.world_root = self.temp_dir / "world"
|
|
self.workdir = self.temp_dir / "workdir"
|
|
self.world_root.mkdir()
|
|
|
|
self._original_workdir = sandbox.WORKDIR
|
|
sandbox.WORKDIR = str(self.workdir)
|
|
|
|
def tearDown(self):
|
|
sandbox.WORKDIR = self._original_workdir
|
|
shutil.rmtree(self.temp_dir, ignore_errors=True)
|
|
|
|
def _run(self, task_id: str | None = None) -> int:
|
|
"""Invoke setup.main() in-process and return its exit code (0 on normal return)."""
|
|
env = {"WORLDBENCH_ROOT": str(self.world_root)}
|
|
if task_id is not None:
|
|
env["WORLDBENCH_TASK_ID"] = task_id
|
|
with mock.patch.dict(os.environ, env, clear=False):
|
|
os.environ.pop("BUNDLEDIR", None)
|
|
if task_id is None:
|
|
os.environ.pop("WORLDBENCH_TASK_ID", None)
|
|
try:
|
|
setup_mod.main()
|
|
return 0
|
|
except SystemExit as e:
|
|
return int(e.code or 0)
|
|
|
|
def test_copies_task_specific_files_into_workdir(self):
|
|
"""Task-specific setup_data/{task_id}/files/ is copied into WORKDIR."""
|
|
task_id = "task-abc"
|
|
files_dir = self.world_root / "tasks" / "setup_data" / task_id / "files"
|
|
files_dir.mkdir(parents=True)
|
|
(files_dir / "james.txt").write_text("james is the CEO")
|
|
(files_dir / "nested").mkdir()
|
|
(files_dir / "nested" / "notes.md").write_text("# notes")
|
|
|
|
assert self._run(task_id=task_id) == 0
|
|
|
|
assert (self.workdir / "james.txt").read_text() == "james is the CEO"
|
|
assert (self.workdir / "nested" / "notes.md").read_text() == "# notes"
|
|
|
|
def test_falls_back_to_generic_setup_data_when_task_dir_missing(self):
|
|
"""If no task-specific dir exists, copies from {world_root}/setup_data/files/."""
|
|
generic_files = self.world_root / "setup_data" / "files"
|
|
generic_files.mkdir(parents=True)
|
|
(generic_files / "shared.txt").write_text("shared context")
|
|
|
|
assert self._run(task_id="unknown-task") == 0
|
|
assert (self.workdir / "shared.txt").read_text() == "shared context"
|
|
|
|
def test_uses_generic_when_no_task_id(self):
|
|
"""With no WORLDBENCH_TASK_ID set, generic setup_data/files/ is used."""
|
|
generic_files = self.world_root / "setup_data" / "files"
|
|
generic_files.mkdir(parents=True)
|
|
(generic_files / "default.txt").write_text("default")
|
|
|
|
assert self._run(task_id=None) == 0
|
|
assert (self.workdir / "default.txt").read_text() == "default"
|
|
|
|
def test_merges_into_existing_workdir(self):
|
|
"""Existing workdir contents are preserved; new files are added alongside."""
|
|
self.workdir.mkdir(parents=True)
|
|
(self.workdir / "preexisting.txt").write_text("keep me")
|
|
|
|
generic_files = self.world_root / "setup_data" / "files"
|
|
generic_files.mkdir(parents=True)
|
|
(generic_files / "new.txt").write_text("new file")
|
|
|
|
assert self._run(task_id=None) == 0
|
|
assert (self.workdir / "preexisting.txt").read_text() == "keep me"
|
|
assert (self.workdir / "new.txt").read_text() == "new file"
|
|
|
|
def test_no_setup_data_exits_cleanly(self):
|
|
"""If neither task-specific nor generic setup_data exists, exit 0 without error."""
|
|
assert self._run(task_id="whatever") == 0
|
|
|
|
def test_setup_data_without_files_subdir_is_noop(self):
|
|
"""setup_data/ with no files/ subdirectory is a no-op (does not error).
|
|
|
|
Workdir IS created (ensure_workdir runs unconditionally so subsequent
|
|
tools have something to chdir into), but no files are copied.
|
|
"""
|
|
(self.world_root / "setup_data").mkdir()
|
|
|
|
assert self._run(task_id=None) == 0
|
|
|
|
def test_rejects_symlink_in_setup_files(self):
|
|
"""Symlinks in the source tree are refused — they'd let a bundle
|
|
materialize an arbitrary host path (e.g. /app/tasks) into /workdir
|
|
and then chown it to the model user."""
|
|
files_dir = self.world_root / "setup_data" / "files"
|
|
files_dir.mkdir(parents=True)
|
|
(files_dir / "real.txt").write_text("ok")
|
|
# Target doesn't need to exist for the check to fire.
|
|
os.symlink("/app/tasks", str(files_dir / "rubrics"))
|
|
|
|
assert self._run(task_id=None) == 1
|
|
assert not (self.workdir / "rubrics").exists()
|
|
assert not (self.workdir / "real.txt").exists() # copytree never ran
|
|
|
|
def test_rejects_symlinked_files_dir_root(self):
|
|
"""A symlinked files/ root is rejected too."""
|
|
real_dir = self.world_root / "setup_data" / "real_files"
|
|
real_dir.mkdir(parents=True)
|
|
(real_dir / "x.txt").write_text("x")
|
|
os.symlink(str(real_dir), str(self.world_root / "setup_data" / "files"))
|
|
|
|
assert self._run(task_id=None) == 1
|
|
assert not (self.workdir / "x.txt").exists()
|
|
|
|
def test_rejects_hardlink_in_setup_files(self):
|
|
"""Hardlinked files in the source tree are refused — they could
|
|
point at a locked-down inode (e.g. /app/packages/grading/*) and
|
|
slip past the symlink guard."""
|
|
files_dir = self.world_root / "setup_data" / "files"
|
|
files_dir.mkdir(parents=True)
|
|
# Two paths sharing one inode: the "target" simulates a protected
|
|
# file outside the bundle; the entry inside files/ is its hardlink.
|
|
protected = self.temp_dir / "fake-protected.py"
|
|
protected.write_text("rubric body")
|
|
os.link(str(protected), str(files_dir / "rubric.py"))
|
|
|
|
assert self._run(task_id=None) == 1
|
|
assert not (self.workdir / "rubric.py").exists()
|
|
|
|
def test_rejects_preexisting_workdir_symlink(self):
|
|
"""A symlink planted in /workdir before setup runs is refused —
|
|
otherwise copytree(dirs_exist_ok=True) would write through it as
|
|
root, into paths the model can't normally write."""
|
|
files_dir = self.world_root / "setup_data" / "files"
|
|
files_dir.mkdir(parents=True)
|
|
(files_dir / "context").mkdir()
|
|
(files_dir / "context" / "doc.txt").write_text("payload")
|
|
|
|
# Plant a model-controlled symlink in /workdir before setup.
|
|
self.workdir.mkdir(parents=True)
|
|
protected_target = self.temp_dir / "fake-protected-dir"
|
|
protected_target.mkdir()
|
|
os.symlink(str(protected_target), str(self.workdir / "context"))
|
|
|
|
assert self._run(task_id=None) == 1
|
|
# The bundle file must not have been written through the symlink.
|
|
assert not (protected_target / "doc.txt").exists()
|
|
|
|
def test_chowns_workdir_to_target_uid_gid_when_root(self):
|
|
"""Under (mocked) root, copied files are chowned back to SETUID:SETGID
|
|
so the unprivileged model user can modify them."""
|
|
generic_files = self.world_root / "setup_data" / "files"
|
|
generic_files.mkdir(parents=True)
|
|
(generic_files / "a.txt").write_text("hello")
|
|
(generic_files / "sub").mkdir()
|
|
(generic_files / "sub" / "b.txt").write_text("world")
|
|
|
|
env = {"WORLDBENCH_ROOT": str(self.world_root), "SETUID": "1000", "SETGID": "1000"}
|
|
with (
|
|
mock.patch.dict(os.environ, env, clear=False),
|
|
mock.patch("core.privilege.os.geteuid", return_value=0),
|
|
mock.patch("core.privilege.os.chown") as chown,
|
|
):
|
|
os.environ.pop("BUNDLEDIR", None)
|
|
os.environ.pop("WORLDBENCH_TASK_ID", None)
|
|
setup_mod.main()
|
|
|
|
chowned = sorted(call.args[0] for call in chown.call_args_list)
|
|
# ensure_workdir chowns WORKDIR; chown_tree_to_target chowns it again plus
|
|
# every entry under it. We only assert the post-copy entries are present.
|
|
assert str(self.workdir / "a.txt") in chowned
|
|
assert str(self.workdir / "sub") in chowned
|
|
assert str(self.workdir / "sub" / "b.txt") in chowned
|
|
for call in chown.call_args_list:
|
|
assert call.args[1:] == (1000, 1000)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|