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.
288 lines
11 KiB
Python
288 lines
11 KiB
Python
"""多用户生产级缺口批单测：sweep 鉴权、默认执行器、新用户默认配额、
|
|
无价模型 token 上限（服务端准入）。
|
|
|
|
只用本地临时端口 + 临时 SQLite，不跑真实模型，不联网。
|
|
环境变量（默认执行器、默认配额额、token 上限）是服务端每次
|
|
请求现读的，测试里用 try/finally 改回，不影响其他测试。
|
|
"""
|
|
|
|
import http.client
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from cloud.auth import AuthStore
|
|
from cloud.cost import CostStore
|
|
from cloud.server import _apply_default_quota, _default_quota_usd_micros
|
|
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.cost import CostStore
|
|
from cloud.server import _apply_default_quota, _default_quota_usd_micros
|
|
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 _utc_today():
|
|
return datetime.now(timezone.utc).strftime("%Y-%m-%d")
|
|
|
|
|
|
def _setenv(name, value):
|
|
old = os.environ.get(name)
|
|
if value is None:
|
|
os.environ.pop(name, None)
|
|
else:
|
|
os.environ[name] = value
|
|
return old
|
|
|
|
|
|
class MultiuserGapsTestCase(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._old_admins = os.environ.get("CODEX_GLOBAL_ADMINS")
|
|
os.environ["CODEX_GLOBAL_ADMINS"] = self.owner_id
|
|
self.viewer_id, self.viewer_key = self.auth.create_user("viewer")
|
|
self.outsider_id, self.outsider_key = self.auth.create_user("outsider")
|
|
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": "gproj"}, self.owner_key)
|
|
self.assertEqual(s, 201, proj)
|
|
self.project_id = proj["id"]
|
|
self.auth.set_member(
|
|
self.project_id, self.owner_id, self.viewer_id, "viewer")
|
|
|
|
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 getattr(self, "_old_admins", None) is None:
|
|
os.environ.pop("CODEX_GLOBAL_ADMINS", None)
|
|
else:
|
|
os.environ["CODEX_GLOBAL_ADMINS"] = self._old_admins
|
|
|
|
# -- helpers -------------------------------------------------------
|
|
def _create_task(self, key=None, body_extra=None):
|
|
body = {"project_id": self.project_id, "prompt": "gap probe"}
|
|
if body_extra:
|
|
body.update(body_extra)
|
|
return _req(self.port, "POST", "/tasks", body,
|
|
key if key is not None else self.owner_key)
|
|
|
|
def _insert_quota(self, user_id, micros):
|
|
conn = sqlite3.connect(self.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 OR REPLACE INTO user_quotas"
|
|
"(user_id, quota_usd_micros, updated_at) VALUES(?, ?, ?)",
|
|
(user_id, micros, time.time()))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def _insert_unpriced_cost(self, task_id, input_tokens, output_tokens):
|
|
conn = sqlite3.connect(self.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(?, ?, ?, 'completed', NULL, '[]', 0, ?, ?, ?, NULL)",
|
|
(task_id, self.project_id, "donor", self.owner_id,
|
|
time.time(), time.time()))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
# CostStore.record 建表并按 UTC 今天记 day，与服务端查询口径一致。
|
|
store = CostStore(self.db_path)
|
|
try:
|
|
store.record(task_id, self.project_id, "gpt-5.6-terra",
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
usd_micros=None, source="actual")
|
|
finally:
|
|
try:
|
|
store._conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def _make_running_task(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
s, _ = _req(self.port, "POST",
|
|
"/tasks/%s/heartbeat" % task["id"], {},
|
|
self.owner_key)
|
|
self.assertEqual(s, 200)
|
|
return task["id"]
|
|
|
|
# -- sweep 鉴权 -----------------------------------------------------
|
|
def test_sweep_viewer_and_outsider_forbidden(self):
|
|
for key in (self.viewer_key, self.outsider_key):
|
|
s, body = _req(self.port, "POST", "/admin/sweep", {},
|
|
key)
|
|
self.assertEqual(s, 403, body)
|
|
|
|
def test_sweep_owner_maintainer_ok_and_sweeps_stale(self):
|
|
tid = self._make_running_task()
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
conn.execute(
|
|
"UPDATE tasks SET last_heartbeat = ? WHERE id = ?",
|
|
(time.time() - 3600, tid))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
s, body = _req(self.port, "POST", "/admin/sweep",
|
|
{"threshold_seconds": 300}, self.owner_key)
|
|
self.assertEqual(s, 200, body)
|
|
self.assertIn(tid, body["swept"])
|
|
|
|
def test_sweep_threshold_clamped_to_60(self):
|
|
tid = self._make_running_task()
|
|
s, body = _req(self.port, "POST", "/admin/sweep",
|
|
{"threshold_seconds": 0}, self.owner_key)
|
|
self.assertEqual(s, 200, body)
|
|
self.assertNotIn(tid, body["swept"])
|
|
|
|
# -- 默认执行器 ------------------------------------------------------
|
|
def test_default_executor_is_docker(self):
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
self.assertEqual(task["executor"], "docker")
|
|
|
|
def test_explicit_executor_kept(self):
|
|
s, task = self._create_task(body_extra={"executor": "local"})
|
|
self.assertEqual(s, 201, task)
|
|
self.assertEqual(task["executor"], "local")
|
|
|
|
def test_default_executor_env_override(self):
|
|
old = _setenv("CLOUD_DEFAULT_EXECUTOR", "local")
|
|
try:
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
self.assertEqual(task["executor"], "local")
|
|
finally:
|
|
_setenv("CLOUD_DEFAULT_EXECUTOR", old)
|
|
old = _setenv("CLOUD_DEFAULT_EXECUTOR", "")
|
|
try:
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
self.assertIsNone(task["executor"])
|
|
finally:
|
|
_setenv("CLOUD_DEFAULT_EXECUTOR", old)
|
|
|
|
# -- 新用户默认配额 ---------------------------------------------------
|
|
def test_apply_default_quota_idempotent(self):
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
self.assertTrue(_apply_default_quota(conn, "newbie"))
|
|
row = conn.execute(
|
|
"SELECT quota_usd_micros FROM user_quotas"
|
|
" WHERE user_id = 'newbie'").fetchone()
|
|
self.assertEqual(row[0], _default_quota_usd_micros())
|
|
# 第二次不再写，不覆盖管理员后来改过的额度。
|
|
conn.execute(
|
|
"UPDATE user_quotas SET quota_usd_micros = 12345"
|
|
" WHERE user_id = 'newbie'")
|
|
conn.commit()
|
|
self.assertFalse(_apply_default_quota(conn, "newbie"))
|
|
row = conn.execute(
|
|
"SELECT quota_usd_micros FROM user_quotas"
|
|
" WHERE user_id = 'newbie'").fetchone()
|
|
self.assertEqual(row[0], 12345)
|
|
finally:
|
|
conn.close()
|
|
|
|
def test_default_quota_env_override(self):
|
|
old = _setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", "5000")
|
|
try:
|
|
self.assertEqual(_default_quota_usd_micros(), 5000)
|
|
finally:
|
|
_setenv("CLOUD_DEFAULT_QUOTA_USD_MICROS", old)
|
|
|
|
# -- 无价 token 上限（服务端准入）--------------------------------------
|
|
def test_unpriced_token_cap_blocks_create(self):
|
|
self._insert_quota(self.owner_id, 10 ** 12) # USD 很大，不挡
|
|
for i in range(7): # 7 * 150 = 1050 token
|
|
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
|
|
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
|
|
try:
|
|
s, body = self._create_task()
|
|
self.assertEqual(s, 402, body)
|
|
self.assertEqual(body["error"], "quota_exceeded")
|
|
self.assertEqual(body["unpriced_tokens_used"], 1050)
|
|
finally:
|
|
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
|
|
|
|
def test_unpriced_token_cap_under_cap_passes(self):
|
|
self._insert_quota(self.owner_id, 10 ** 12)
|
|
for i in range(6): # 6 * 150 = 900 < 1000
|
|
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
|
|
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
|
|
try:
|
|
s, task = self._create_task()
|
|
self.assertEqual(s, 201, task)
|
|
finally:
|
|
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
|
|
|
|
def test_unpriced_token_cap_blocks_without_quota_row(self):
|
|
# P0 安全批后：无配额行不再等于不限额，无价 cap 照样挡。
|
|
for i in range(7):
|
|
self._insert_unpriced_cost("donor-%d" % i, 100, 50)
|
|
old = _setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", "1000")
|
|
try:
|
|
s, body = self._create_task()
|
|
self.assertEqual(s, 402, body)
|
|
self.assertEqual(body["error"], "quota_exceeded")
|
|
finally:
|
|
_setenv("CLOUD_QUOTA_UNPRICED_TOKEN_CAP", old)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|