- CODEX_GLOBAL_ADMINS replaces maintainer-derived global admin; admin endpoints (quotas, sweep, backup, push dispatch) fail-closed 403 when unset, with audit on allow and deny. - continue/followups/messages/retry check both caller and original creator quotas (USD + unpriced token cap), 402 on either overrun. - Quota-less users now get the #86 default quota instead of unlimited; init_db backfills and OAuth create/bind/existing-login branches INSERT OR IGNORE a default row; explicit NULL row stays unlimited. user_quotas added to migration tables. - /fix-review: untrusted authors rejected before dedup claim, no task; trusted tasks force executor=docker, created_by=project owner, owner quota gate, prompt is a fixed template plus a 2000-char excerpt. - GitHub App setup state only issued by global admins, state rows bind issued_by (ALTER-compatible for old DBs), callback rejects issuers no longer in the admin list without consuming the state. - Cancel: popen poll checks DB status every second, Docker path runs a cancel watchdog that kills by container name, write-back is CAS-guarded so late results never overwrite a new status. - SHELL_ENV_EXCLUDE adds *_TOKEN, *_AUTH*, *_CREDENTIAL*, *API_KEY*; comment states the residual risk: same-UID /proc environ remains readable, true isolation needs UID separation or hidepid. - Tests: new cloud/tests/test_p0_security.py (17 cases); old tests updated to the new semantics (admin env, default quota, trusted fix-review author) without weakening any gate.
467 lines
19 KiB
Python
467 lines
19 KiB
Python
"""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()
|