host
7091a2d0d2
- 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.
361 lines
15 KiB
Python
361 lines
15 KiB
Python
"""租户配额与计费单测：月度 USD 上限 + 准入 402 + 管理端点。
|
||
|
||
不跑真实模型回合、不联网：server 侧走本地临时端口 + 临时
|
||
SQLite；用量直接插 cost_records 行（按任务 created_by 归属、
|
||
UTC day），验证 JOIN 归属、自然月窗口与三路准入（建任务 /
|
||
分支 / 重试）。配额单位是微美元（1 美元 = 1_000_000），
|
||
与 cost_records 一致。
|
||
"""
|
||
|
||
import http.client
|
||
import json
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import unittest
|
||
from datetime import date, datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import create_server
|
||
except ImportError: # discover -s cloud/tests 时把项目根目录补进 sys.path
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.auth import AuthStore
|
||
from cloud.server import create_server
|
||
|
||
|
||
def _req(port, method, path, body=None, key=None):
|
||
conn = http.client.HTTPConnection("127.0.0.1", port, timeout=10)
|
||
hdrs = {}
|
||
payload = None
|
||
if body is not None:
|
||
payload = json.dumps(body).encode("utf-8")
|
||
hdrs["Content-Type"] = "application/json"
|
||
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 = None
|
||
return resp.status, parsed
|
||
|
||
|
||
def _set_task_status(db_path, task_id, status):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute("UPDATE tasks SET status = ? WHERE id = ?",
|
||
(status, task_id))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _insert_cost(db_path, task_id, project_id, usd_micros, day):
|
||
conn = sqlite3.connect(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, project_id, "gpt-5.6-terra", 100, 50, 150,
|
||
usd_micros, "actual", day, 0))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _utc_today():
|
||
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
||
|
||
|
||
def _prev_month_day():
|
||
today = date.today()
|
||
first = today.replace(day=1)
|
||
return (first - timedelta(days=1)).isoformat()
|
||
|
||
|
||
class ServerQuotasTestCase(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.other_id, self.other_key = self.auth.create_user("other")
|
||
# P0 安全批后：admin 端点只认 CODEX_GLOBAL_ADMINS，
|
||
# 本文件里 owner 即管理员，测前写 env、测后还原。
|
||
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": "qproj"}, self.owner_key)
|
||
self.assertEqual(s, 201, proj)
|
||
self.project_id = proj["id"]
|
||
|
||
def tearDown(self):
|
||
self.server.shutdown()
|
||
self.server.server_close()
|
||
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
|
||
|
||
def _create_task(self, key=None):
|
||
return _req(self.port, "POST", "/tasks",
|
||
{"project_id": self.project_id,
|
||
"prompt": "quota probe"},
|
||
key if key is not None else self.owner_key)
|
||
|
||
def _set_quota(self, target, quota, key=None):
|
||
return _req(self.port, "POST", "/admin/quotas",
|
||
{"user_id": target, "quota_usd_micros": quota},
|
||
key if key is not None else self.owner_key)
|
||
|
||
def test_self_quota_default_is_default_quota_not_unlimited(self):
|
||
# P0 安全批后：无配额行按默认额度处理，不再是 None。
|
||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["quota_usd_micros"], 10000000)
|
||
self.assertEqual(q["used_usd_micros"], 0)
|
||
self.assertEqual(q["remaining_usd_micros"], 10000000)
|
||
self.assertEqual(q["month"],
|
||
datetime.now(timezone.utc).strftime("%Y-%m"))
|
||
|
||
def test_admin_set_and_get(self):
|
||
s, q = self._set_quota(self.owner_id, 5000000)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||
s, q = _req(self.port, "GET",
|
||
"/admin/quotas?user_id=" + self.owner_id,
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||
key=self.owner_key)
|
||
self.assertEqual(q["quota_usd_micros"], 5000000)
|
||
# null = 恢复不限额
|
||
s, q = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertIsNone(q["quota_usd_micros"])
|
||
|
||
def test_admin_set_forbidden_and_validation(self):
|
||
# other 没有任何项目角色，不是管理员
|
||
s, q = self._set_quota(self.owner_id, 100, self.other_key)
|
||
self.assertEqual(s, 403, q)
|
||
s, q = _req(self.port, "GET",
|
||
"/admin/quotas?user_id=" + self.owner_id,
|
||
key=self.other_key)
|
||
self.assertEqual(s, 403, q)
|
||
s, q = self._set_quota(self.owner_id, -1)
|
||
self.assertEqual(s, 400, q)
|
||
s, q = self._set_quota("no-such-user", 100)
|
||
self.assertEqual(s, 404, q)
|
||
s, q = _req(self.port, "POST", "/admin/quotas",
|
||
{"quota_usd_micros": 100}, self.owner_key)
|
||
self.assertEqual(s, 400, q)
|
||
s, q = _req(self.port, "GET", "/admin/quotas",
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 400, q)
|
||
|
||
def test_usage_counts_current_month_only(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
5000000, _utc_today())
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
100000000, _prev_month_day())
|
||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["used_usd_micros"], 5000000)
|
||
|
||
def test_create_blocked_at_quota(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
5000000, _utc_today())
|
||
s, _ = self._set_quota(self.owner_id, 5000000)
|
||
self.assertEqual(s, 200)
|
||
s, body = self._create_task()
|
||
self.assertEqual(s, 402, body)
|
||
self.assertEqual(body["error"], "quota_exceeded")
|
||
self.assertEqual(body["quota_usd_micros"], 5000000)
|
||
self.assertEqual(body["used_usd_micros"], 5000000)
|
||
# 额度调高后放行
|
||
s, _ = self._set_quota(self.owner_id, 6000000)
|
||
self.assertEqual(s, 200)
|
||
s, task2 = self._create_task()
|
||
self.assertEqual(s, 201, task2)
|
||
|
||
def test_fork_blocked_at_quota(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
3000000, _utc_today())
|
||
_set_task_status(self.db_path, task["id"], "completed")
|
||
s, _ = self._set_quota(self.owner_id, 3000000)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/fork" % task["id"],
|
||
{"message": "branch probe"}, self.owner_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, forked = _req(self.port, "POST",
|
||
"/tasks/%s/fork" % task["id"],
|
||
{"message": "branch probe"}, self.owner_key)
|
||
self.assertEqual(s, 201, forked)
|
||
|
||
def test_retry_blocked_at_quota(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
2000000, _utc_today())
|
||
_set_task_status(self.db_path, task["id"], "failed")
|
||
s, _ = self._set_quota(self.owner_id, 2000000)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.owner_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(row["status"], "failed")
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, retried = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.owner_key)
|
||
self.assertEqual(s, 200, retried)
|
||
self.assertEqual(retried["status"], "queued")
|
||
|
||
def test_retry_blocked_when_operator_or_creator_over_quota(self):
|
||
# P0 安全批后：重试按双人门查配额（触发人 + 创建人）。
|
||
# #90 的“只查创建者、操作人超限不拦”已被 P0 有意收紧：
|
||
# 操作人用自己的额度过门、费用记创建者名下，仍是绕限，
|
||
# 所以操作人超限也必须拦。other 须是 maintainer 才能过
|
||
# approve_merge 权限门，否则 403 会掩盖配额判断。
|
||
self.auth.set_member(self.project_id, self.owner_id,
|
||
self.other_id, "maintainer")
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
2000000, _utc_today())
|
||
_set_task_status(self.db_path, task["id"], "failed")
|
||
# 创建者超限：维护者替他重试也应拦。
|
||
s, _ = self._set_quota(self.owner_id, 2000000)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.other_key)
|
||
self.assertEqual(s, 402, body)
|
||
# 创建者不限额、操作人额度为 0：P0 双人门应拦（#90
|
||
# 只查创建者的旧口径会放行，那正是 P0 要堵的绕限）。
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, _ = self._set_quota(self.other_id, 0)
|
||
self.assertEqual(s, 200)
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.other_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(row["status"], "failed")
|
||
# 两人都不超限才放行。
|
||
s, _ = self._set_quota(self.other_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, retried = _req(self.port, "POST",
|
||
"/tasks/%s/retry" % task["id"], {},
|
||
self.other_key)
|
||
self.assertEqual(s, 200, retried)
|
||
self.assertEqual(retried["status"], "queued")
|
||
|
||
def test_continue_blocked_at_creator_quota(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
5000000, _utc_today())
|
||
_set_task_status(self.db_path, task["id"], "completed")
|
||
s, _ = self._set_quota(self.owner_id, 5000000)
|
||
self.assertEqual(s, 200)
|
||
# 以前：先回 200 排队、随后认领门悄悄标 interrupted。
|
||
# 现在：API 侧直接 402，状态保持 completed 不变。
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task["id"],
|
||
{"message": "again"}, self.owner_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, row = _req(self.port, "GET", "/tasks/" + task["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(row["status"], "completed")
|
||
# 额度调高后放行。
|
||
s, _ = self._set_quota(self.owner_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, continued = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task["id"],
|
||
{"message": "again"}, self.owner_key)
|
||
self.assertEqual(s, 200, continued)
|
||
self.assertEqual(continued["status"], "queued")
|
||
self.auth.set_member(self.project_id, self.owner_id,
|
||
self.other_id, "maintainer")
|
||
s, task2 = self._create_task()
|
||
self.assertEqual(s, 201, task2)
|
||
_set_task_status(self.db_path, task2["id"], "completed")
|
||
s, _ = self._set_quota(self.other_id, 0)
|
||
self.assertEqual(s, 200)
|
||
# P0 安全批后：续跑按双人门查，操作人额度为 0、
|
||
# 创建者不限额时也应拦（#90 只查创建者的旧口径会
|
||
# 放行，那正是 P0 要堵的绕限）；两人都不超限才放行。
|
||
s, body = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task2["id"],
|
||
{"message": "again"}, self.other_key)
|
||
self.assertEqual(s, 402, body)
|
||
s, row = _req(self.port, "GET", "/tasks/" + task2["id"],
|
||
key=self.owner_key)
|
||
self.assertEqual(row["status"], "completed")
|
||
s, _ = self._set_quota(self.other_id, None)
|
||
self.assertEqual(s, 200)
|
||
s, continued2 = _req(self.port, "POST",
|
||
"/tasks/%s/continue" % task2["id"],
|
||
{"message": "again"}, self.other_key)
|
||
self.assertEqual(s, 200, continued2)
|
||
self.assertEqual(continued2["status"], "queued")
|
||
|
||
def test_quota_reports_unpriced_tokens(self):
|
||
s, task = self._create_task()
|
||
self.assertEqual(s, 201, task)
|
||
# 有价 5000000 计入 used；无价 actual 150 只进
|
||
# unpriced_actual_tokens，不计价、不拦任务。
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
5000000, _utc_today())
|
||
_insert_cost(self.db_path, task["id"], self.project_id,
|
||
None, _utc_today())
|
||
s, q = _req(self.port, "GET", "/users/me/quota",
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["used_usd_micros"], 5000000)
|
||
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||
s, q = _req(self.port, "GET",
|
||
"/admin/quotas?user_id=" + self.owner_id,
|
||
key=self.owner_key)
|
||
self.assertEqual(s, 200, q)
|
||
self.assertEqual(q["unpriced_actual_tokens"], 150)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|