Files
codex-app-server/cloud/tests/test_p1_security_worker.py
T
host 9bb9a50e4e 残留批次A:成员删除、PR needs_local_worker 落库、推送 DNS TOCTOU
- 成员删除:DELETE /projects/{id}/members/{user_id},限 owner 或
  全局管理员,禁删唯一 owner,记 remove_member 审计
  (AuthStore.remove_member + server 分发,沿用 PATCH 守卫口径)
- PR needs_local_worker 落库:非 local 扫到 queued 任务时 CAS
  写行状态为 needs_local_worker;claim/find_open 同步认该状态,
  标记过的行仍可被 local worker 认领、不被搁置;另修
  claim/find_open 对裸 sqlite3 tuple 行的 dict() 兜底
  (TypeError/ValueError)
- 推送 DNS TOCTOU:发送改走建连时解析+判段的连接类,解析
  一次、只拨判过的 IP,Host/SNI 仍为原域名;opener 直接构造
  以排除 ProxyHandler 改写(build_opener 默认会塞代理)
2026-10-27 19:10:49 +00:00

541 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""P1 执行面针对性单测(只测本路改动,不跑全量三套)。
运行(在 ~/workspace/codex-app-server-p1security 下):
python3 -m unittest discover -s cloud/tests -p "test_p1_security_worker.py" -v
说明:只用注入的假 runner 与测试假值,不跑真实模型回合、不联网、
不登录任何服务;env 里的值只是占位用的假名,绝不是真实密钥。
本文件只归本路(执行面)所有,不碰 test_p1_security_server.py。
"""
import json
import os
import sqlite3
import sys
import tempfile
import threading
import time
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock
try:
import cloud.sandbox as sandbox_mod
from cloud.cost import CostStore
from cloud.pr_exec import enqueue_pr_job
from cloud.sandbox import DockerExecutor, LocalExecutor
from cloud.server import init_db
from cloud.worker import Worker, _code_signature
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import cloud.sandbox as sandbox_mod
from cloud.cost import CostStore
from cloud.pr_exec import enqueue_pr_job
from cloud.sandbox import DockerExecutor, LocalExecutor
from cloud.server import init_db
from cloud.worker import Worker, _code_signature
FAKE_ENV = {"XCM_API_KEY": "test-placeholder"}
def _completed_runner(cmd):
if "dispatch" in cmd:
name = cmd[cmd.index("--name") + 1]
else:
name = cmd[cmd.index("--task") + 1]
return SimpleNamespace(
returncode=0,
stdout=json.dumps({"task": name, "status": "completed",
"text": "done"}),
stderr="")
def _usage_runner(usage, calls):
def fake(cmd):
calls.append(list(cmd))
return SimpleNamespace(
returncode=0,
stdout=json.dumps({"status": "completed", "text": "done",
"usage": usage}),
stderr="")
return fake
def _insert_task(db_path, task_id, created_by="tester", status="queued",
project_id="p1", **cols):
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(?, ?, ?, ?, NULL, '[]', 0, ?, ?, ?, NULL)",
(task_id, project_id, "do work", status, created_by, now, now),
)
for key, value in cols.items():
conn.execute(
"UPDATE tasks SET %s = ? WHERE id = ?" % key,
(value, task_id))
conn.commit()
finally:
conn.close()
def _insert_quota(db_path, user_id, micros):
conn = sqlite3.connect(db_path)
try:
conn.execute(
"CREATE TABLE IF NOT EXISTS user_quotas("
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT,"
" updated_at DOUBLE PRECISION)")
conn.execute(
"INSERT INTO user_quotas(user_id, quota_usd_micros, updated_at)"
" VALUES(?, ?, ?)",
(user_id, micros, time.time()))
conn.commit()
finally:
conn.close()
def _insert_cost(db_path, task_id, model, input_tokens, output_tokens,
micros, source="actual"):
_insert_task(db_path, task_id, status="completed")
store = CostStore(db_path)
try:
store.record(task_id, "p1", model, input_tokens=input_tokens,
output_tokens=output_tokens, usd_micros=micros,
source=source)
finally:
try:
store._conn.close()
except Exception:
pass
def _status(db_path, task_id):
conn = sqlite3.connect(db_path)
try:
row = conn.execute(
"SELECT status FROM tasks WHERE id = ?", (task_id,)).fetchone()
finally:
conn.close()
return row[0] if row else None
def _error_text(db_path, task_id):
conn = sqlite3.connect(db_path)
try:
row = conn.execute(
"SELECT error_text FROM tasks WHERE id = ?",
(task_id,)).fetchone()
finally:
conn.close()
return row[0] if row else None
class P1WorkerTestCase(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.calls = []
def tearDown(self):
self.tmp.cleanup()
def _worker(self, runner=None, **kwargs):
kwargs.setdefault("env", FAKE_ENV)
return Worker(self.db_path, "/fake/host.py",
runner=runner or _completed_runner, **kwargs)
# 1. 无配额行按默认 $10 额度(不再是不限额)
def test_missing_quota_row_uses_default_quota(self):
_insert_cost(self.db_path, "donor-a", "gpt-6.1-sol",
1, 0, 11000000)
_insert_cost(self.db_path, "donor-b", "gpt-6.1-sol",
1, 0, 9000000)
worker = self._worker()
over = worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1"}, None)
# donor-a 归 tester:已花 $11 > 默认 $10,挡住。
self.assertIsNotNone(over)
self.assertIn("quota_exceeded", over)
# donor-b 同样归 tester 会把总额推得更高,单独测另一用户:
_insert_task(self.db_path, "donor-c", created_by="rich",
status="completed")
store = CostStore(self.db_path)
try:
store.record("donor-c", "p1", "gpt-6.1-sol", input_tokens=1,
output_tokens=0, usd_micros=9000000,
source="actual")
finally:
try:
store._conn.close()
except Exception:
pass
under = worker._quota_exceeded(
{"created_by": "rich", "project_id": "p1"}, None)
self.assertIsNone(under)
# 2. 显式 NULL 配额行仍是不限额
def test_explicit_null_quota_row_stays_unlimited(self):
_insert_quota(self.db_path, "tester", None)
_insert_cost(self.db_path, "donor", "gpt-6.1-sol",
1, 0, 999000000)
worker = self._worker()
self.assertIsNone(worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1"}, None))
# 3. 属主为空拒绝新回合(不再返回 None 放行)
def test_ownerless_task_refused(self):
worker = self._worker()
for owner in (None, ""):
err = worker._quota_exceeded(
{"created_by": owner, "project_id": "p1"}, None)
self.assertIsNotNone(err)
self.assertIn("quota_exceeded", err)
# 4. 协调员子任务属主跟父任务走
def test_child_inherits_parent_owner(self):
_insert_quota(self.db_path, "tester", 100)
_insert_cost(self.db_path, "donor", "gpt-6.1-sol",
1, 0, 200)
_insert_task(self.db_path, "parent", created_by="tester")
_insert_task(self.db_path, "parent2", created_by="fresh")
worker = self._worker()
blocked = worker._quota_exceeded(
{"created_by": None, "parent_task_id": "parent",
"project_id": "p1"}, None)
self.assertIsNotNone(blocked)
self.assertIn("quota_exceeded", blocked)
# fresh 用户无配额行、未花钱:默认额度内放行。
self.assertIsNone(worker._quota_exceeded(
{"created_by": None, "parent_task_id": "parent2",
"project_id": "p1"}, None))
# 5. 查库失败 fail-closed(新回合拒绝,不自动放行)
def test_quota_db_failure_is_fail_closed(self):
worker = self._worker()
with mock.patch("cloud.worker._db_connect",
side_effect=RuntimeError("test db down")):
err = worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1"}, None)
self.assertIsNotNone(err)
self.assertTrue(err.startswith("quota_check_failed:"))
# 6. Docker 配额看门狗:命中杀容器、未命中不杀
def test_docker_quota_watchdog(self):
killed = []
def fake(cmd, timeout=None):
killed.append(list(cmd))
return SimpleNamespace(returncode=0, stdout="", stderr="")
executor = DockerExecutor(runner=fake)
worker = Worker(self.db_path, "/fake/host.py",
env=FAKE_ENV, executor=executor)
# 未命中:0.2 秒后停表,kill 不应被调。
stop = threading.Event()
threading.Timer(0.2, stop.set).start()
worker._docker_quota_watchdog(
"T1", "codex-test", stop, lambda: False, interval=0.05)
self.assertEqual(killed, [])
# 命中:立即杀容器并记命中标记。
stop2 = threading.Event()
worker._docker_quota_watchdog(
"T2", "codex-test", stop2, lambda: True, interval=0.05)
self.assertTrue(any("kill" in c for c in killed))
self.assertTrue(worker._take_quota_hit("T2"))
# 命中已消费,第二次取不到(不污染下一轮)。
self.assertFalse(worker._take_quota_hit("T2"))
# 7. usage_cost_usd 按实际成本累加,max_cost_usd 不再是死参数
def test_usage_cost_accumulates_and_budget_trips(self):
_insert_task(self.db_path, "T1")
worker = self._worker(
runner=_usage_runner(
{"input_tokens": 10000, "cached_input_tokens": 0,
"output_tokens": 1000, "total_tokens": 11000},
self.calls),
model="gpt-6.1-sol", max_cost_usd=0.01)
self.assertFalse(worker._budget_exceeded())
result = worker.run_one()
self.assertEqual(result["status"], "completed")
# 10000*$2 + 1000*$10 = $0.03,超过 $0.01 上限。
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
self.assertTrue(worker._budget_exceeded())
conn = sqlite3.connect(self.db_path)
try:
row = conn.execute(
"SELECT usd_micros, source FROM cost_records"
" WHERE task_id = 'T1'").fetchone()
finally:
conn.close()
self.assertEqual(row, (30000, "actual"))
# 8. 取消分支先记成本再返回(与配额中断收尾同款记账点)
def test_cancel_branch_records_cost(self):
worker = self._worker()
task = {"id": "T1", "project_id": "p1", "prompt": "do work"}
with worker._db_lock:
cost_id = worker._record_usage_cost(
task, "gpt-6.1-sol",
{"input_tokens": 10000, "cached_input_tokens": 0,
"output_tokens": 1000, "total_tokens": 11000},
"done", 0)
self.assertIsNotNone(cost_id)
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
# 无用量时回落估算,不累加 usage_cost_usd。
with worker._db_lock:
worker._record_usage_cost(
{"id": "T2", "project_id": "p1", "prompt": "do work"},
"gpt-6.1-sol", None, "done", 0)
self.assertAlmostEqual(worker.usage_cost_usd, 0.03, places=6)
# 9. 估算行不计入无价 token 上限
def test_estimated_rows_excluded_from_unpriced_cap(self):
_insert_cost(self.db_path, "donor-est", "gpt-5.6-terra",
100000, 0, None, source="estimated")
worker = self._worker()
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
try:
# 10 万估算 token 在上限 100 下仍放行(估算不是真花费)。
self.assertIsNone(worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1"}, None))
_insert_cost(self.db_path, "donor-act", "gpt-5.6-terra",
75, 75, None, source="actual")
err = worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1"}, None)
self.assertIsNotNone(err)
self.assertIn("无价模型 token", err)
finally:
if old is None:
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
else:
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
# 10. cache_write 不计入无价上限(只 input+output 计)
def test_cache_write_excluded_from_unpriced_cap(self):
worker = self._worker()
old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP")
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "100"
try:
usage = {"input_tokens": 10, "output_tokens": 5,
"cache_write_input_tokens": 9999999,
"total_tokens": 15}
self.assertIsNone(worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1",
"model": "gpt-5.6-terra"}, usage))
usage2 = {"input_tokens": 200, "output_tokens": 0,
"cache_write_input_tokens": 0,
"total_tokens": 200}
err = worker._quota_exceeded(
{"created_by": "tester", "project_id": "p1",
"model": "gpt-5.6-terra"}, usage2)
self.assertIsNotNone(err)
finally:
if old is None:
os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None)
else:
os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old
# 11. 协调员子任务继承父任务模型/强度,分工显式值优先
def test_coordinator_child_inherits_model_effort(self):
worker = self._worker()
task = {"id": "coord", "project_id": "p1", "prompt": "coord",
"priority": 0, "created_by": "tester", "model": "gpt-6.1-sol",
"effort": "high", "executor": None, "environment_id": None}
children = worker._coord_sync_children(task, [
{"name": "a", "prompt": "do a", "model": None,
"effort": None, "executor": None},
{"name": "b", "prompt": "do b", "model": "gpt-5.6-terra",
"effort": "low", "executor": None},
])
self.assertEqual(len(children), 2)
conn = sqlite3.connect(self.db_path)
try:
rows = {r[0]: (r[1], r[2]) for r in conn.execute(
"SELECT coordinator_part, model, effort FROM tasks"
" WHERE parent_task_id = 'coord'").fetchall()}
finally:
conn.close()
self.assertEqual(rows["a"], ("gpt-6.1-sol", "high"))
self.assertEqual(rows["b"], ("gpt-5.6-terra", "low"))
# 12. 断点续跑:Local/Docker/popen 三路都不传 --goal-text
def test_goal_resume_breakpoint_omits_goal_text(self):
calls = []
def fake(cmd):
calls.append(list(cmd))
return SimpleNamespace(
returncode=0,
stdout="goal met after rounds\n"
"usage input=10 cached=0 cache_write=0 output=5 "
"reasoning=0 total=15 completions=1\n",
stderr="")
task = {"id": "T1", "project_id": "p1", "prompt": "orig goal",
"verify_command": "true", "goal_max_rounds": 3}
local = LocalExecutor("/fake/host.py", runner=fake)
local.goal_resume_task(task, None)
self.assertIn("goal-resume", calls[-1])
self.assertNotIn("--goal-text", calls[-1])
docker_calls = []
def dfake(cmd, timeout=None):
docker_calls.append(list(cmd))
return SimpleNamespace(
returncode=0,
stdout="goal met after rounds\n"
"usage input=10 cached=0 cache_write=0 output=5 "
"reasoning=0 total=15 completions=1\n",
stderr="")
docker = DockerExecutor(runner=dfake, env=FAKE_ENV)
with mock.patch("cloud.sandbox.shutil.which",
return_value="/usr/bin/docker"):
docker.goal_resume_task(task, None)
self.assertIn("goal-resume", docker_calls[-1])
self.assertNotIn("--goal-text", docker_calls[-1])
with mock.patch.object(sandbox_mod.subprocess, "Popen") as popen:
popen.return_value = SimpleNamespace(pid=9999)
local.popen_task(
{"id": "T1", "prompt": "resume msg",
"_resume_kind": "breakpoint"})
# popen_task 默认 kind=dispatch;显式测 goal-resume:
local.popen_task(
{"id": "T1", "prompt": "resume msg",
"_resume_kind": "breakpoint"}, kind="goal-resume")
cmd = popen.call_args[0][0]
self.assertIn("goal-resume", cmd)
self.assertNotIn("--goal-text", cmd)
# 13. 开新目标:goal-resume 带 --goal-text(旧行为保留)
def test_goal_resume_new_goal_keeps_goal_text(self):
calls = []
def fake(cmd):
calls.append(list(cmd))
return SimpleNamespace(
returncode=0,
stdout="goal met after rounds\n"
"usage input=10 cached=0 cache_write=0 output=5 "
"reasoning=0 total=15 completions=1\n",
stderr="")
task = {"id": "T1", "project_id": "p1", "prompt": "new goal"}
local = LocalExecutor("/fake/host.py", runner=fake)
local.goal_resume_task(task, "new goal")
self.assertIn("--goal-text", calls[-1])
self.assertIn("new goal", calls[-1])
# 14. goal 轮次上限:sandbox loud fail + worker 认领后 failed
def test_goal_max_rounds_cap(self):
task = {"id": "T1", "project_id": "p1", "prompt": "g",
"verify_command": "true", "goal_max_rounds": 101}
local = LocalExecutor("/fake/host.py", runner=_completed_runner)
with self.assertRaises(ValueError):
local.goal_task(task)
_insert_task(self.db_path, "T1", verify_command="true",
goal_max_rounds=101)
worker = self._worker()
result = worker.run_one()
self.assertEqual(result["status"], "failed")
self.assertIn("上限", result["error"])
self.assertEqual(_status(self.db_path, "T1"), "failed")
self.assertEqual(self.calls, [])
# 15. PR 非 local:显式 needs_local_worker,行状态持久化,
# 任务本身不动,标记过的行仍可被 local worker 认领
def test_pr_non_local_needs_local_worker(self):
conn = sqlite3.connect(self.db_path)
try:
job_id = enqueue_pr_job(conn, "T1", "p1", "test-host")
finally:
conn.close()
executor = DockerExecutor(runner=lambda cmd, timeout=None:
SimpleNamespace(returncode=0, stdout="",
stderr=""))
worker = Worker(self.db_path, "/fake/host.py", env=FAKE_ENV,
executor=executor, host_name="test-host")
results = worker._run_pr_jobs()
self.assertEqual(len(results), 1)
self.assertEqual(results[0]["status"], "needs_local_worker")
self.assertEqual(results[0]["pr_job_id"], job_id)
conn = sqlite3.connect(self.db_path)
try:
row = conn.execute(
"SELECT status FROM pr_jobs WHERE id = ?",
(job_id,)).fetchone()
self.assertEqual(row[0], "needs_local_worker")
# 标记过的行仍可被 local worker 认领,不会被搁置。
from cloud.pr_exec import claim_pr_job
claimed = claim_pr_job(conn, "test-host")
self.assertIsNotNone(claimed)
self.assertEqual(claimed["id"], job_id)
finally:
conn.close()
# 16. 多文件签名 + 签名缺失/已变时认领 fail-closed
def test_code_signature_multi_file_and_claim_refusal(self):
sig = _code_signature("/fake/host.py")
self.assertIsNotNone(sig)
self.assertEqual(len(sig), 6)
for entry in sig:
self.assertEqual(len(entry), 2)
_insert_task(self.db_path, "T1")
worker = self._worker()
worker._code_signature = ("stale-signature",)
result = worker.run_one()
self.assertFalse(result["dispatched"])
self.assertEqual(result["reason"], "already_claimed")
self.assertEqual(_status(self.db_path, "T1"), "queued")
worker._code_signature = None
result = worker.run_one()
self.assertFalse(result["dispatched"])
self.assertEqual(_status(self.db_path, "T1"), "queued")
# 17. goal 首轮撞已存在 → 断点续跑一次,不直接 failed
def test_goal_already_exists_falls_back_to_breakpoint(self):
calls = []
def fake(cmd):
calls.append(list(cmd))
if "goal-resume" in cmd:
return SimpleNamespace(
returncode=0,
stdout="round 1: turn=completed verify_rc=0 "
"verify_out=\ngoal met after rounds\n"
"usage input=10 cached=0 cache_write=0 output=5 "
"reasoning=0 total=15 completions=1\n",
stderr="")
if "goal" in cmd:
return SimpleNamespace(
returncode=0,
stdout="host 任务已存在:T1 已存在",
stderr="")
return _completed_runner(cmd)
_insert_task(self.db_path, "T1", verify_command="true",
goal_max_rounds=2)
worker = self._worker(runner=fake)
result = worker.run_one()
self.assertEqual(result["status"], "completed")
# 注意:cmd 是元素列表,"goal" in cmd 只匹配恰好等于
# "goal" 的元素,goal-resume 命令里并没有 "goal" 元素,
# 所以按“任一元素含 goal 子串”过滤。
goal_cmds = [c for c in calls
if any("goal" in e for e in c)]
self.assertTrue(any("goal-resume" in c for c in goal_cmds))
resume_cmd = [c for c in goal_cmds if "goal-resume" in c][0]
self.assertNotIn("--goal-text", resume_cmd)
self.assertEqual(_status(self.db_path, "T1"), "completed")
if __name__ == "__main__":
unittest.main()