host
706d5a8ca4
首版共享会话在缓存未命中时各线程各自登录：突发冷缓存下 N 路并发开通会在网关刷出 N 次 bcrypt。改为锁内双重检查 登录（RLock 可重入，_login 成功后写回同一缓存）：第一路 登录，其余线程等锁后直接播种。实例本身仍不跨线程，401 时照常走刷新->重登录自愈。 测试：FakeApi 加锁（假实现自身线程安全，失败指向真代码）； 新增 10 线程同起并发单测，断言 10 路全成功且管理员登录 恰好 1 次；全量 cloud 795（4 跳过）、host 69、deploy 19 全绿；真网关 10 路突发实测 0 错误、墙钟 0.774s、临时号 已清。
434 lines
17 KiB
Python
434 lines
17 KiB
Python
"""sub2api 网关接入：系统用管理 API 为每位用户建号发 key。
|
||
|
||
设计（2026-10-27 独立部署 sub2api 后）：
|
||
- 系统侧只用管理 API 管用户：管理员登录 -> 建 sub2api 账号 ->
|
||
以该账号登录 -> 发 key -> key 存我方库，任务注入该用户的 key。
|
||
- key 创建是用户侧端点 POST /api/v1/keys（不是 /api/v1/api-keys，
|
||
后者只有管理员改组）；管理员不能代用户建/删 key，所以 refresh_token
|
||
也要存下来：以后轮换/吊销 key 时靠它换用户 access token。
|
||
- key 明文存 PG：与 github_app_config 的 App 私钥同一信任边界
|
||
（库只在宿主 loopback 可达，服务端 env 同机）；at-rest 加密是
|
||
后续项，不在本批假装已做。
|
||
- 零第三方依赖，只用标准库；HTTP 可注入 request 可调用对象，
|
||
便于单测（见 cloud/tests/test_sub2api.py）。
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import secrets
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
|
||
GATEWAY_TABLE_SQL = (
|
||
"CREATE TABLE IF NOT EXISTS user_gateway_keys("
|
||
"user_id TEXT PRIMARY KEY, "
|
||
"sub2api_user_id BIGINT NOT NULL, "
|
||
"api_key TEXT NOT NULL, "
|
||
"refresh_token TEXT, "
|
||
"created_at DOUBLE PRECISION, "
|
||
"updated_at DOUBLE PRECISION)")
|
||
|
||
|
||
class Sub2apiError(RuntimeError):
|
||
def __init__(self, message, status=None):
|
||
super().__init__(message)
|
||
self.status = status
|
||
|
||
|
||
def _default_request(method, url, body=None, token=None):
|
||
"""urllib 实现的 request：返回 (status, parsed_json_dict)。"""
|
||
headers = {"Content-Type": "application/json"}
|
||
if token:
|
||
headers["Authorization"] = "Bearer " + token
|
||
req = urllib.request.Request(
|
||
url, method=method,
|
||
data=json.dumps(body).encode() if body is not None else None,
|
||
headers=headers)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=30) as resp:
|
||
raw = resp.read() or b"{}"
|
||
return resp.status, json.loads(raw)
|
||
except urllib.error.HTTPError as e:
|
||
try:
|
||
data = json.loads(e.read() or b"{}")
|
||
except ValueError:
|
||
data = {}
|
||
return e.code, data
|
||
|
||
|
||
def gateway_origin():
|
||
"""sub2api origin（不带 /v1）：未配则返回 None。"""
|
||
raw = os.environ.get("SUB2API_BASE_URL", "").strip()
|
||
return raw.rstrip("/") or None
|
||
|
||
|
||
def gateway_name():
|
||
raw = os.environ.get("CLOUD_GATEWAY", "xcmapi").strip().lower()
|
||
return raw if raw in ("xcmapi", "sub2api") else "xcmapi"
|
||
|
||
|
||
def gateway_group_name():
|
||
return os.environ.get("SUB2API_GROUP_NAME", "codex-system").strip() \
|
||
or "codex-system"
|
||
|
||
|
||
def opening_balance():
|
||
"""开通时授的初始余额（美元）：默认 10，显式值优先。
|
||
|
||
新用户余额为 0 时网关鉴权层直接 403（余额门在任何上游路由
|
||
之前），所以开通必须同时授余额，否则发出的 key 立刻不可用。
|
||
显式配 0 表示不授（调用方自理）；非法/负值回落默认 10。
|
||
"""
|
||
raw = os.environ.get(
|
||
"CLOUD_SUB2API_OPENING_BALANCE_USD", "").strip()
|
||
if not raw:
|
||
return 10.0
|
||
try:
|
||
value = float(raw)
|
||
except ValueError:
|
||
return 10.0
|
||
return value if value >= 0 else 10.0
|
||
|
||
|
||
def ensure_gateway_keys_table(conn):
|
||
conn.execute(GATEWAY_TABLE_SQL)
|
||
|
||
|
||
def get_user_gateway_key(conn, user_id):
|
||
"""只读：返回该用户的网关行 dict 或 None，绝不触发开通。"""
|
||
ensure_gateway_keys_table(conn)
|
||
cur = conn.execute(
|
||
"SELECT user_id, sub2api_user_id, api_key, refresh_token "
|
||
"FROM user_gateway_keys WHERE user_id = ?",
|
||
(user_id,))
|
||
row = cur.fetchone()
|
||
if row is None:
|
||
return None
|
||
try:
|
||
keys = ("user_id", "sub2api_user_id", "api_key", "refresh_token")
|
||
return dict(zip(keys, row))
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
# 进程内共享管理员会话：(origin, email) -> (token, refresh, expires_at)。
|
||
# 只复用 token 元组，不共用实例（实例不跨线程）；注册与首任务开通
|
||
# 每次新建实例都要登录一次（实测约 111ms，占开通链路约 1/3），
|
||
# 共享后同一进程内第二次起免登录。401/刷新失败会回落到重登录，
|
||
# 密码轮换后下一次登录即用新密码刷新缓存，不会卡死。
|
||
_SHARED_ADMIN = {}
|
||
_SHARED_ADMIN_LOCK = threading.RLock()
|
||
|
||
|
||
def _store_shared_session(admin):
|
||
key = getattr(admin, "_shared_key", None)
|
||
if key is None or not admin._token:
|
||
return
|
||
with _SHARED_ADMIN_LOCK:
|
||
_SHARED_ADMIN[key] = (admin._token, admin._refresh,
|
||
admin._expires_at)
|
||
|
||
|
||
def _evict_shared_session(admin):
|
||
key = getattr(admin, "_shared_key", None)
|
||
if key is None:
|
||
return
|
||
with _SHARED_ADMIN_LOCK:
|
||
_SHARED_ADMIN.pop(key, None)
|
||
|
||
|
||
def shared_admin_session(request=None):
|
||
"""带进程内 token 复用的管理员会话（生产开通路径默认用它）。
|
||
|
||
缓存命中且未过期时直接播种 token，免一次管理员登录；
|
||
未命中/已过期时在锁内登录一次再返回：突发冷缓存下
|
||
N 个并发开通也只登录一次（其余线程等锁后直接播种），
|
||
不会各登各的把网关刷出 N 次 bcrypt。调用方显式传
|
||
admin 的单测不受影响（不走这条工厂）。
|
||
"""
|
||
admin = Sub2apiAdmin.from_env(request=request)
|
||
key = (admin.origin, admin.email)
|
||
admin._shared_key = key
|
||
with _SHARED_ADMIN_LOCK:
|
||
cached = _SHARED_ADMIN.get(key)
|
||
if cached and cached[2] > time.time():
|
||
admin._token, admin._refresh, admin._expires_at = cached
|
||
return admin
|
||
# 缓存未命中：在锁内登录（_login 成功后经
|
||
# _store_shared_session 写回同一缓存，RLock 可重入）。
|
||
admin._login()
|
||
return admin
|
||
|
||
|
||
class Sub2apiAdmin:
|
||
"""管理员会话：登录、建/删用户、查组。token 自动刷新一次。
|
||
|
||
每次实例只在进程内复用 token 元组（见 shared_admin_session）：
|
||
实例本身不跨线程共用，只复用 (token, refresh, expires_at)，
|
||
401 时走已有刷新->重登录链路自愈，密码轮换后下一次登录
|
||
即用新密码并刷新缓存，不会卡在旧 token 上。
|
||
"""
|
||
|
||
def __init__(self, origin, email, password, request=None):
|
||
if not origin:
|
||
raise Sub2apiError("SUB2API_BASE_URL 未配置")
|
||
if not email or not password:
|
||
raise Sub2apiError("SUB2API 管理员邮箱/密码未配置")
|
||
self.origin = origin
|
||
self.email = email
|
||
self.password = password
|
||
self.request = request or _default_request
|
||
self._token = None
|
||
self._refresh = None
|
||
self._expires_at = 0
|
||
self._shared_key = None
|
||
|
||
@classmethod
|
||
def from_env(cls, request=None):
|
||
return cls(gateway_origin(),
|
||
os.environ.get("SUB2API_ADMIN_EMAIL", ""),
|
||
os.environ.get("SUB2API_ADMIN_PASSWORD", ""),
|
||
request=request)
|
||
|
||
def _login(self):
|
||
st, data = self.request(
|
||
"POST", self.origin + "/api/v1/auth/login",
|
||
{"email": self.email, "password": self.password})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
_evict_shared_session(self)
|
||
raise Sub2apiError(
|
||
"sub2api 管理员登录失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
self._token = payload["access_token"]
|
||
self._refresh = payload.get("refresh_token")
|
||
try:
|
||
ttl = int(payload.get("expires_in") or 3600)
|
||
except (TypeError, ValueError):
|
||
ttl = 3600
|
||
self._expires_at = time.time() + max(60, ttl - 60)
|
||
_store_shared_session(self)
|
||
|
||
def _refresh_token(self):
|
||
if not self._refresh:
|
||
self._login()
|
||
return
|
||
st, data = self.request(
|
||
"POST", self.origin + "/api/v1/auth/refresh",
|
||
{"refresh_token": self._refresh})
|
||
payload = (data or {}).get("data") or {}
|
||
if st == 200 and payload.get("access_token"):
|
||
self._token = payload["access_token"]
|
||
if payload.get("refresh_token"):
|
||
self._refresh = payload["refresh_token"]
|
||
try:
|
||
ttl = int(payload.get("expires_in") or 3600)
|
||
except (TypeError, ValueError):
|
||
ttl = 3600
|
||
self._expires_at = time.time() + max(60, ttl - 60)
|
||
_store_shared_session(self)
|
||
return
|
||
self._login()
|
||
|
||
def _call(self, method, path, body=None):
|
||
if not self._token or time.time() >= self._expires_at:
|
||
self._login()
|
||
st, data = self.request(method, self.origin + path, body,
|
||
token=self._token)
|
||
if st == 401:
|
||
self._refresh_token()
|
||
st, data = self.request(method, self.origin + path, body,
|
||
token=self._token)
|
||
return st, data
|
||
|
||
def create_user(self, email, password, username):
|
||
st, data = self._call("POST", "/api/v1/admin/users", {
|
||
"email": email, "password": password,
|
||
"username": username, "role": "user"})
|
||
payload = (data or {}).get("data") or {}
|
||
if st not in (200, 201) or payload.get("id") is None:
|
||
raise Sub2apiError(
|
||
"sub2api 建用户失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload["id"]
|
||
|
||
def delete_user(self, sub_user_id):
|
||
st, _data = self._call(
|
||
"DELETE", "/api/v1/admin/users/%s" % sub_user_id)
|
||
return st in (200, 204)
|
||
|
||
def set_balance(self, sub_user_id, amount):
|
||
st, data = self._call(
|
||
"POST", "/api/v1/admin/users/%s/balance" % sub_user_id,
|
||
{"balance": amount, "operation": "set",
|
||
"notes": "codex provisioning opening balance"})
|
||
if st != 200:
|
||
raise Sub2apiError(
|
||
"sub2api 授余额失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
|
||
def list_groups(self):
|
||
st, data = self._call("GET", "/api/v1/admin/groups")
|
||
payload = (data or {}).get("data")
|
||
if st != 200:
|
||
raise Sub2apiError("sub2api 查组失败: %s" % st, status=st)
|
||
if isinstance(payload, list):
|
||
return payload
|
||
if isinstance(payload, dict):
|
||
for key in ("items", "list", "groups"):
|
||
if isinstance(payload.get(key), list):
|
||
return payload[key]
|
||
return []
|
||
|
||
def resolve_group_id(self):
|
||
raw = os.environ.get("SUB2API_GROUP_ID", "").strip()
|
||
if raw:
|
||
try:
|
||
return int(raw)
|
||
except ValueError:
|
||
raise Sub2apiError("SUB2API_GROUP_ID 非法: %r" % raw)
|
||
want = gateway_group_name()
|
||
for grp in self.list_groups():
|
||
if isinstance(grp, dict) and grp.get("name") == want:
|
||
if grp.get("id") is not None:
|
||
return grp["id"]
|
||
raise Sub2apiError("sub2api 组不存在: %s" % want)
|
||
|
||
|
||
def user_login(origin, email, password, request=None):
|
||
"""以 sub2api 用户身份登录：返回 (access_token, refresh_token)。"""
|
||
req = request or _default_request
|
||
st, data = req("POST", origin + "/api/v1/auth/login",
|
||
{"email": email, "password": password})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
raise Sub2apiError(
|
||
"sub2api 用户登录失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload["access_token"], payload.get("refresh_token")
|
||
|
||
|
||
def user_create_key(origin, user_token, name, group_id, request=None):
|
||
"""以用户身份发 key：返回 (key_id, key_value)。"""
|
||
req = request or _default_request
|
||
body = {"name": name}
|
||
if group_id is not None:
|
||
body["group_id"] = group_id
|
||
st, data = req("POST", origin + "/api/v1/keys", body, token=user_token)
|
||
payload = (data or {}).get("data") or {}
|
||
value = payload.get("key") or payload.get("api_key") \
|
||
or payload.get("value")
|
||
if st not in (200, 201) or not value:
|
||
raise Sub2apiError(
|
||
"sub2api 发 key 失败: %s"
|
||
% ((data or {}).get("message") or st), status=st)
|
||
return payload.get("id"), value
|
||
|
||
|
||
def user_delete_key(origin, user_token, key_id, request=None):
|
||
req = request or _default_request
|
||
st, _data = req("DELETE", origin + "/api/v1/keys/%s" % key_id,
|
||
token=user_token)
|
||
return st in (200, 204)
|
||
|
||
|
||
def user_refresh_access(origin, refresh_token, request=None):
|
||
"""refresh_token 换 access_token：返回 (新 access, 新 refresh|None)。"""
|
||
req = request or _default_request
|
||
st, data = req("POST", origin + "/api/v1/auth/refresh",
|
||
{"refresh_token": refresh_token})
|
||
payload = (data or {}).get("data") or {}
|
||
if st != 200 or not payload.get("access_token"):
|
||
raise Sub2apiError("sub2api refresh 失败: %s" % st, status=st)
|
||
return payload["access_token"], payload.get("refresh_token")
|
||
|
||
|
||
def list_gateway_models(origin, api_key, request=None):
|
||
"""网关 /v1/models：返回模型 id 列表（key 无模型时为空列表）。"""
|
||
req = request or _default_request
|
||
st, data = req("GET", origin + "/v1/models", token=api_key)
|
||
if st != 200:
|
||
raise Sub2apiError("sub2api 查模型失败: %s" % st, status=st)
|
||
payload = (data or {}).get("data", data)
|
||
if isinstance(payload, dict):
|
||
payload = payload.get("data") or []
|
||
out = []
|
||
if isinstance(payload, list):
|
||
for item in payload:
|
||
if isinstance(item, dict) and item.get("id"):
|
||
out.append(item["id"])
|
||
elif isinstance(item, str):
|
||
out.append(item)
|
||
return out
|
||
|
||
|
||
def ensure_user_gateway_key(conn, user_id, admin=None, request=None):
|
||
"""幂等开通：已有行直接返回 key；否则建号->登录->发 key->授余额->入库。
|
||
|
||
入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错；
|
||
授余额失败同样走这条清理（余额为 0 的 key 在网关侧直接 403，
|
||
留着没有意义）。同一用户并发开通时建号撞 email 已存在，
|
||
则轮询复用胜出者入库的行。管理员工单默认走进程内共享会话
|
||
（同一进程内第二次起免管理员登录，见 shared_admin_session）。
|
||
调用方负责 conn 的提交语义：本函数在入库后
|
||
commit（开通是独立事务，不跟调用方的业务事务绑在一起）。
|
||
"""
|
||
ensure_gateway_keys_table(conn)
|
||
row = get_user_gateway_key(conn, user_id)
|
||
if row is not None:
|
||
return row["api_key"]
|
||
origin = gateway_origin()
|
||
if not origin:
|
||
raise Sub2apiError("SUB2API_BASE_URL 未配置，无法开通网关账号")
|
||
admin = admin or shared_admin_session(request=request)
|
||
req = request or _default_request
|
||
email = "codex-%s@codex.local" % user_id
|
||
username = "codex-" + str(user_id)[:12]
|
||
password = secrets.token_hex(16)
|
||
try:
|
||
sub_id = admin.create_user(email, password, username)
|
||
except Sub2apiError as exc:
|
||
if "already exists" not in str(exc).lower():
|
||
raise
|
||
# 并发开通竞态：同一用户的另一路已建号（email 唯一），
|
||
# 等它入库后复用那一行；等不到则原错抛出（任务失败后
|
||
# 可重跑，孤儿号场景下重跑仍撞号需管理员介入清理）。
|
||
for _ in range(5):
|
||
time.sleep(2)
|
||
row = get_user_gateway_key(conn, user_id)
|
||
if row is not None:
|
||
return row["api_key"]
|
||
raise
|
||
key_id = None
|
||
try:
|
||
user_token, refresh = user_login(origin, email, password,
|
||
request=req)
|
||
group_id = admin.resolve_group_id()
|
||
key_id, key = user_create_key(origin, user_token,
|
||
"codex-task-key", group_id,
|
||
request=req)
|
||
amount = opening_balance()
|
||
if amount > 0:
|
||
admin.set_balance(sub_id, amount)
|
||
now = time.time()
|
||
conn.execute(
|
||
"INSERT INTO user_gateway_keys(user_id, sub2api_user_id, "
|
||
"api_key, refresh_token, created_at, updated_at) "
|
||
"VALUES(?, ?, ?, ?, ?, ?)",
|
||
(user_id, sub_id, key, refresh, now, now))
|
||
conn.commit()
|
||
except Exception:
|
||
if key_id is not None:
|
||
try:
|
||
user_delete_key(origin, user_token, key_id, request=req)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
admin.delete_user(sub_id)
|
||
except Exception:
|
||
pass
|
||
raise
|
||
return key
|