"""P0 安全批单测（基线 99efa93 后的硬缺口修法）。 不跑真实模型回合、不联网、不碰生产：server 侧走本地临时 端口 + 临时 SQLite；用量直接插 cost_records 行；webhook 走 本地签名；GitHub App 回调走注入的 fake 换 code 客户端。 测试用的 key、secret、state 全是假值。 """ import hashlib import hmac import http.client import json import os import sqlite3 import sys import tempfile import threading import unittest from datetime import datetime, timezone from pathlib import Path from unittest import mock try: from cloud.auth import AuthStore from cloud.github_app import SetupStateStore from cloud.server import create_server from cloud.webhooks import is_trusted_fix_review_author except ImportError: sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from cloud.auth import AuthStore from cloud.github_app import SetupStateStore from cloud.server import create_server from cloud.webhooks import is_trusted_fix_review_author def _req(port, method, path, body=None, key=None, headers=None, raw_body=None): conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10) hdrs = dict(headers or {}) if raw_body is not None: payload = raw_body elif body is not None: payload = json.dumps(body).encode("utf-8") hdrs.setdefault("Content-Type", "application/json") else: payload = None if key is not None: hdrs["X-API-Key"] = key try: conn.request(method, path, body=payload, headers=hdrs) resp = conn.getresponse() data = resp.read() finally: conn.close() try: parsed = json.loads(data.decode("utf-8")) if data else None except ValueError: parsed = data.decode("utf-8", "replace") return resp.status, parsed def _sig(secret, raw): return "sha256=" + hmac.new( secret.encode(), raw, hashlib.sha256).hexdigest() def _utc_today(): return datetime.now(timezone.utc).strftime("%Y-%m-%d") class P0SecurityTestCase(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.db_path = os.path.join(self.tmp.name, "control.db") self.auth = AuthStore(self.db_path) self.owner_id, self.owner_key = self.auth.create_user("owner") self.member_id, self.member_key = self.auth.create_user("member") self.outsider_id, self.outsider_key = self.auth.create_user( "outsider") self._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS") os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id self.server = create_server(self.db_path, "127.0.0.1", 0) self.port = self.server.server_address[1] self.thread = threading.Thread( target=self.server.serve_forever, daemon=True) self.thread.start() s, proj = _req(self.port, "POST", "/projects", {"name": "owner/repo"}, self.owner_key) self.assertEqual(s, 201, proj) self.project_id = proj["id"] self.auth.set_member( self.project_id, self.owner_id, self.member_id, "maintainer") def tearDown(self): self.server.shutdown() self.server.server_close() self.thread.join(timeout=5) try: self.auth.conn.close() except Exception: pass self.tmp.cleanup() if self._old_admins is None: os.environ.pop("CODEX_GLOBAL_ADMINS", None) else: os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins os.environ.pop("GITHUB_APP_PUBLIC_BASE", None) os.environ.pop("GITHUB_WEBHOOK_SECRET", None) # -- helpers --------------------------------------------------- def _create_task(self, key=None, prompt="p0 probe"): return _req(self.port, "POST", "/tasks", {"project_id": self.project_id, "prompt": prompt}, key if key is not None else self.owner_key) def _set_status(self, task_id, status): conn = sqlite3.connect(self.db_path) try: conn.execute("UPDATE tasks SET status = ? WHERE id = ?", (status, task_id)) conn.commit() finally: conn.close() def _insert_cost(self, task_id, usd_micros, day=None): conn = sqlite3.connect(self.db_path) try: conn.execute( "INSERT INTO cost_records(task_id, project_id, model," " input_tokens, output_tokens, total_tokens, usd_micros," " source, day, created_at) VALUES(?, ?, ?, ?, ?, ?, ?," " ?, ?, ?)", (task_id, self.project_id, "gpt-5.6-terra", 100, 50, 150, usd_micros, "actual", day or _utc_today(), 0)) conn.commit() finally: conn.close() def _insert_unpriced_cost(self, task_id, in_tok, out_tok): conn = sqlite3.connect(self.db_path) try: conn.execute( "INSERT INTO cost_records(task_id, project_id, model," " input_tokens, output_tokens, total_tokens, usd_micros," " source, day, created_at) VALUES(?, ?, ?, ?, ?, ?," " NULL, ?, ?, ?)", (task_id, self.project_id, "gpt-5.6-terra", in_tok, out_tok, in_tok + out_tok, "actual", _utc_today(), 0)) conn.commit() finally: conn.close() def _set_quota(self, target, quota): return _req(self.port, "POST", "/admin/quotas", {"user_id": target, "quota_usd_micros": quota}, self.owner_key) def _audit_actions(self): conn = sqlite3.connect(self.db_path) try: return [r[0] for r in conn.execute( "SELECT action FROM audit_log").fetchall()] finally: conn.close() def _webhook(self, payload, delivery, secret="test-secret"): raw = json.dumps(payload).encode("utf-8") headers = { "Content-Type": "application/json", "X-Hub-Signature-256": _sig(secret, raw), "X-GitHub-Event": "issue_comment", "X-GitHub-Delivery": delivery, } with mock.patch.dict( os.environ, {"GITHUB_WEBHOOK_SECRET": secret}): return _req(self.port, "POST", "/webhooks/github", headers=headers, raw_body=raw) def _fix_payload(self, author="alice", association="COLLABORATOR", comment="请 /fix-review 看一下"): return { "action": "created", "comment": {"id": 999, "body": comment, "user": {"login": author, "author_association": association}}, "repository": {"full_name": "owner/repo"}, "issue": {"number": 42, "user": {"login": "alice"}}, } def _tasks(self): conn = sqlite3.connect(self.db_path) try: return conn.execute( "SELECT id, prompt, status, executor, created_by" " FROM tasks").fetchall() finally: conn.close() # -- 1. 真全局管理员 ------------------------------------------- def test_admin_fail_closed_without_env(self): with mock.patch.dict(os.environ, {}, clear=False): os.environ.pop("CODEX_GLOBAL_ADMINS", None) for method, path, body in ( ("POST", "/admin/sweep", {}), ("GET", "/admin/quotas?user_id=" + self.owner_id, None), ("POST", "/admin/backup", {"out_dir": "x"}), ("POST", "/admin/push/dispatch", {"event_type": "failed", "project_id": self.project_id})): s, _ = _req(self.port, method, path, body, self.owner_key) self.assertEqual(s, 403, (method, path)) # 配了 env 的管理员仍可调 sweep。 s, body = _req(self.port, "POST", "/admin/sweep", {}, self.owner_key) self.assertEqual(s, 200, body) self.assertIn("admin_sweep", self._audit_actions()) def test_admin_denied_is_audited(self): with mock.patch.dict(os.environ, {}, clear=False): os.environ.pop("CODEX_GLOBAL_ADMINS", None) _req(self.port, "POST", "/admin/sweep", {}, self.member_key) self.assertIn("admin_sweep_denied", self._audit_actions()) # -- 2. 所有花钱路径统一配额门 ---------------------------------- def test_continue_blocked_when_creator_over_quota(self): # 任务由 owner 创建，member（maintainer）来续跑： # member 自己没超，但 owner 超了，仍应 402。 s, task = self._create_task(key=self.owner_key) self.assertEqual(s, 201, task) self._set_status(task["id"], "completed") self._insert_cost(task["id"], 5000000) s, _ = self._set_quota(self.owner_id, 100) self.assertEqual(s, 200) s, body = _req(self.port, "POST", "/tasks/%s/continue" % task["id"], {"message": "继续"}, self.member_key) self.assertEqual(s, 402, body) self.assertEqual(body["error"], "quota_exceeded") def test_followup_blocked_when_creator_over_quota(self): s, task = self._create_task(key=self.owner_key) self.assertEqual(s, 201, task) self._set_status(task["id"], "running") self._insert_cost(task["id"], 5000000) s, _ = self._set_quota(self.owner_id, 100) self.assertEqual(s, 200) s, body = _req(self.port, "POST", "/tasks/%s/followups" % task["id"], {"body": "后续", "mode": "queued"}, self.member_key) self.assertEqual(s, 402, body) def test_peer_message_blocked_when_target_creator_over_quota(self): s, target = self._create_task(key=self.owner_key) self.assertEqual(s, 201, target) self._set_status(target["id"], "running") s, sender = self._create_task(key=self.member_key) self.assertEqual(s, 201, sender) self._insert_cost(target["id"], 5000000) s, _ = self._set_quota(self.owner_id, 100) self.assertEqual(s, 200) s, body = _req(self.port, "POST", "/tasks/%s/messages" % target["id"], {"from_task_id": sender["id"], "body": "hi"}, self.member_key) self.assertEqual(s, 402, body) def test_retry_blocked_when_creator_over_quota(self): s, task = self._create_task(key=self.owner_key) self.assertEqual(s, 201, task) self._set_status(task["id"], "failed") self._insert_cost(task["id"], 5000000) s, _ = self._set_quota(self.owner_id, 100) self.assertEqual(s, 200) s, body = _req(self.port, "POST", "/tasks/%s/retry" % task["id"], {}, self.member_key) self.assertEqual(s, 402, body) # -- 3. 无配额行不再等于不限额 ---------------------------------- def test_no_quota_row_uses_default_usd_quota(self): # member 是 maintainer 但配额行是缺的：先建任务、再把 # 当月费用推到默认额度之上，下一次建任务应 402。 s, task = self._create_task(key=self.member_key) self.assertEqual(s, 201, task) self._insert_cost(task["id"], 11000000) s, body = self._create_task(key=self.member_key) self.assertEqual(s, 402, body) self.assertEqual(body["quota_usd_micros"], 10000000) def test_no_quota_row_unpriced_cap_blocks(self): old = os.environ.get("CLOUD_QUOTA_UNPRICED_TOKEN_CAP") os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = "1000" try: s, task = self._create_task(key=self.member_key) self.assertEqual(s, 201, task) self._insert_unpriced_cost(task["id"], 700, 350) s, body = self._create_task(key=self.member_key) self.assertEqual(s, 402, body) self.assertEqual(body["unpriced_tokens_used"], 1050) finally: if old is None: os.environ.pop("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", None) else: os.environ["CLOUD_QUOTA_UNPRICED_TOKEN_CAP"] = old # -- 4. webhook /fix-review 收紧 -------------------------------- def test_webhook_untrusted_author_rejected(self): s, body = self._webhook( self._fix_payload(author="bob", association="CONTRIBUTOR"), delivery="d-p0-untrusted") self.assertEqual(s, 200, body) self.assertIsNone(body["task_id"]) self.assertEqual(body["error"], "untrusted fix-review author") self.assertEqual(self._tasks(), []) self.assertIn("webhook_fix_review_rejected", self._audit_actions()) def test_webhook_trusted_forces_docker_and_owner_and_truncates(self): long_comment = "请 /fix-review 看一下 " + "X" * 3000 + "TAILMARK" s, body = self._webhook( self._fix_payload(comment=long_comment), delivery="d-p0-trusted") self.assertEqual(s, 200, body) self.assertIsNotNone(body["task_id"]) rows = self._tasks() self.assertEqual(len(rows), 1) task_id, prompt, status, executor, created_by = rows[0] self.assertEqual(task_id, body["task_id"]) self.assertEqual(status, "queued") self.assertEqual(executor, "docker") self.assertEqual(created_by, self.owner_id) self.assertNotIn("TAILMARK", prompt) self.assertLessEqual(len(prompt), 2200) def test_webhook_owner_quota_blocks(self): s, task = self._create_task(key=self.owner_key) self.assertEqual(s, 201, task) self._insert_cost(task["id"], 5000000) s, _ = self._set_quota(self.owner_id, 100) self.assertEqual(s, 200) s, body = self._webhook(self._fix_payload(), delivery="d-p0-quota") self.assertEqual(s, 402, body) # 未建新任务：只有之前手动建的那一个任务。 self.assertEqual(len(self._tasks()), 1) def test_is_trusted_fix_review_author_unit(self): self.assertTrue(is_trusted_fix_review_author( {"comment": {"user": {"login": "bob", "author_association": "MEMBER"}}})) self.assertTrue(is_trusted_fix_review_author( {"comment": {"user": {"login": "alice", "author_association": "FIRST_TIMER"}}, "issue": {"user": {"login": "alice"}}})) self.assertFalse(is_trusted_fix_review_author( {"comment": {"user": {"login": "bob", "author_association": "CONTRIBUTOR"}}, "issue": {"user": {"login": "alice"}}})) self.assertFalse(is_trusted_fix_review_author( {"comment": {"user": {"login": "bob"}}})) # -- 5. GitHub App 配置覆盖限管理员 ------------------------------- def test_github_app_setup_requires_admin(self): os.environ["GITHUB_APP_PUBLIC_BASE"] = "https://example.invalid" s, _ = _req(self.port, "POST", "/github-app/setup/state", {}, self.member_key) self.assertEqual(s, 403) s, body = _req(self.port, "POST", "/github-app/setup/state", {}, self.owner_key) self.assertEqual(s, 200, body) self.assertIn("state", body) def test_setup_state_issuer_mismatch_rejected(self): store = SetupStateStore(self.db_path) try: state = store.create_state(issued_by="someone-else") self.assertFalse(store.consume_state( state, allowed_issuers={self.owner_id})) # 未被消耗：管理员清单里换成签发人后仍可消耗。 self.assertTrue(store.consume_state( state, allowed_issuers={"someone-else"})) self.assertFalse(store.consume_state(state)) finally: store.conn.close() def test_github_app_callback_issuer_mismatch(self): self.server.github_app_exchange_client = lambda url, headers: { "id": 424242, "slug": "fake-codex-app", "client_id": "Iv1.x", "webhook_secret": "fake-app-secret", "pem": "fake-pem", "html_url": "https://github.com/apps/fake-codex-app"} store = SetupStateStore(self.db_path) try: state = store.create_state(issued_by="someone-else") finally: store.conn.close() s, page = _req(self.port, "GET", "/github-app/callback?code=x&state=%s" % state) self.assertEqual(s, 400, page) class P0WorkerCancelTestCase(unittest.TestCase): """取消写回守卫：迟到写回不覆盖新结果（mock DB 级验证）。""" def test_write_result_guarded_by_status(self): from cloud.worker import Worker tmp = tempfile.TemporaryDirectory() try: db_path = os.path.join(tmp.name, "control.db") conn = sqlite3.connect(db_path) try: conn.execute( "CREATE TABLE tasks(id TEXT PRIMARY KEY," " status TEXT, result_text TEXT, error_text TEXT," " actual_model TEXT, started_at REAL," " finished_at REAL, turn_id TEXT," " root_turn_id TEXT, updated_at REAL)") conn.execute( "INSERT INTO tasks(id, status) VALUES('t1', 'queued')") conn.commit() finally: conn.close() worker = Worker.__new__(Worker) worker.db_path = db_path worker._ensure_result_columns = lambda c: None # 状态已是 queued，不是 completed：带守卫的写回应被拦。 ok = worker._write_execution_result( "t1", result_text="late", only_if_status="completed") self.assertFalse(ok) conn = sqlite3.connect(db_path) try: row = conn.execute( "SELECT result_text FROM tasks WHERE id = 't1'" ).fetchone() finally: conn.close() self.assertIsNone(row[0]) # 状态对上时写回成功。 ok = worker._write_execution_result( "t1", result_text="ok", only_if_status="queued") self.assertTrue(ok) finally: tmp.cleanup() def test_shell_env_exclude_covers_tokens(self): import sys as _sys _sys.path.insert(0, str(Path(__file__).resolve().parents[2])) import host as _host excl = _host.SHELL_ENV_EXCLUDE for pat in ("*_TOKEN", "*_AUTH*", "*_CREDENTIAL*", "*API_KEY*"): self.assertIn(pat, excl) # gh 走文件登录：host.py 里不应有 GH_TOKEN 环境变量依赖。 import pathlib text = pathlib.Path( Path(__file__).resolve().parents[2] / "host.py" ).read_text() self.assertNotIn("GH_TOKEN", text) self.assertNotIn("GITHUB_TOKEN", text) if __name__ == "__main__": unittest.main()