mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 20:12:39 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
27f094e096 | ||
|
|
fa402b909a | ||
|
|
2e5973b165 | ||
|
|
ecf199d29d | ||
|
|
7cbb0d8552 | ||
|
|
c4a8878f1f | ||
|
|
95c8b2401e | ||
|
|
e669e51b92 | ||
|
|
ce9e70d5cf | ||
|
|
deb5e8756a | ||
|
|
2ee66c401d | ||
|
|
3071d41443 | ||
|
|
58ee0c8d75 | ||
|
|
78f85583ea | ||
|
|
7a7f51a534 | ||
|
|
db4295d646 | ||
|
|
3c41984733 | ||
|
|
04617cc0ff | ||
|
|
dabca7ca79 | ||
|
|
a90c18cbea | ||
|
|
e502766170 | ||
|
|
b8e844f4ee | ||
|
|
80a533c890 | ||
|
|
6d152893ff | ||
|
|
97afb4089c | ||
|
|
48fad26433 | ||
|
|
a05b873215 | ||
|
|
1b4029b050 | ||
|
|
139a3d4346 | ||
|
|
8c9cda1803 | ||
|
|
0a8831777b | ||
|
|
34c8dbd604 | ||
|
|
e311d368d3 | ||
|
|
f133fde643 | ||
|
|
c0b8f6b0a8 | ||
|
|
9ccc9077cc | ||
|
|
f5e484aa70 | ||
|
|
aa615c2677 | ||
|
|
032001598a | ||
|
|
6def36a171 | ||
|
|
d2ef0cb151 | ||
|
|
25bd4956f0 | ||
|
|
2518691dbc | ||
|
|
bfabfe60c8 | ||
|
|
f3312cf7d9 | ||
|
|
30202f8266 | ||
|
|
a2acbf553b | ||
|
|
dd0cbe91c9 | ||
|
|
0cc382e331 | ||
|
|
6becd11e39 | ||
|
|
80dd2929a3 | ||
|
|
4990ca0b2c | ||
|
|
344f303919 | ||
|
|
ff5b0e6254 | ||
|
|
fedeba2568 | ||
|
|
3c86e249f4 | ||
|
|
f72958d530 | ||
|
|
8b8b6b3132 | ||
|
|
9bc7d10c08 | ||
|
|
496005d688 | ||
|
|
041db5d824 | ||
|
|
3379603022 | ||
|
|
4a97491be6 | ||
|
|
35b5edb24c | ||
|
|
01062c3244 | ||
|
|
604b764cf5 | ||
|
|
f69042ca6e | ||
|
|
41b58b640a | ||
|
|
5da3325a30 | ||
|
|
93717394b2 | ||
|
|
6aa6b3a968 | ||
|
|
39107ca457 | ||
|
|
6e2ba0e4cd | ||
|
|
8a147fcc51 | ||
|
|
d0fa8c63f1 | ||
|
|
3a4966f786 | ||
|
|
3ed2873f14 | ||
|
|
c044d5c1f1 | ||
|
|
732aeef880 | ||
|
|
3dcf412b65 | ||
|
|
c8ff478876 | ||
|
|
f7979ae8ff | ||
|
|
8fd6f91b7e | ||
|
|
7840eb1c49 | ||
|
|
6102d64274 | ||
|
|
eea02c38fc | ||
|
|
0343dba916 | ||
|
|
b08cab91a9 | ||
|
|
ef88cf3f89 | ||
|
|
9d79b42729 | ||
|
|
7a05a7cdcc | ||
|
|
752b2eaa9b | ||
|
|
026d2a8083 | ||
|
|
5af38ccf41 | ||
|
|
340dc99e02 | ||
|
|
8a2d7c5bd3 | ||
|
|
ba2b0044e3 | ||
|
|
d2753dc2e2 | ||
|
|
c99e479bcc | ||
|
|
b8c640d218 | ||
|
|
1433fbc44d | ||
|
|
2cff4bf169 | ||
|
|
361983cf4d | ||
|
|
510ee451bd | ||
|
|
7bbe170405 | ||
|
|
3401a971a8 | ||
|
|
ed280b81fc | ||
|
|
841544345f | ||
|
|
4c73149af8 | ||
|
|
a1e1997632 | ||
|
|
a67351caec | ||
|
|
62846bb074 | ||
|
|
d4b9797ff7 | ||
|
|
19149ca196 | ||
|
|
831812b39d | ||
|
|
d98ac48eb9 | ||
|
|
72b29a7d4e | ||
|
|
14608dc6d4 | ||
|
|
daa4989124 | ||
|
|
8a95a46a69 | ||
|
|
9d498c2474 | ||
|
|
a2f802d409 | ||
|
|
b01196a7a8 | ||
|
|
e14fb2b6ff | ||
|
|
dd7a2b22f0 | ||
|
|
4e8ea7d568 | ||
|
|
a3f2b8fd87 | ||
|
|
23cdd20606 | ||
|
|
a62b821b5f | ||
|
|
774ff5d8c8 | ||
|
|
d2667393b3 | ||
|
|
e002fbb349 | ||
|
|
005bc5c8d4 | ||
|
|
bc4bd118d8 | ||
|
|
f3925db11b | ||
|
|
f9a467a4c0 | ||
|
|
8ce9288a6a | ||
|
|
c1e702be5e | ||
|
|
bc6b69289c | ||
|
|
5a0492bf88 | ||
|
|
fdc9d92fdb | ||
|
|
8cdd372b7e | ||
|
|
080a52121a | ||
|
|
37db8d031b | ||
|
|
f0e0b7e6d8 | ||
|
|
4c845ca697 | ||
|
|
95d27f2dce | ||
|
|
f2e7a55703 | ||
|
|
1ed9f59599 | ||
|
|
b4ad7fd363 | ||
|
|
ef2a22be70 | ||
|
|
c831bb979f | ||
|
|
3edcbfd120 | ||
|
|
1db10dc559 | ||
|
|
3f6b5c7bd7 | ||
|
|
2946281a43 | ||
|
|
335edde9c8 | ||
|
|
1a855d3e0e | ||
|
|
b1a6b80267 | ||
|
|
c3f759d13a | ||
|
|
539bfc8bad | ||
|
|
8b75dd5576 | ||
|
|
bd0f2d6405 | ||
|
|
71afe24710 | ||
|
|
500f241692 | ||
|
|
b92bbf0299 | ||
|
|
2594950993 | ||
|
|
9386ce8a13 | ||
|
|
be74deae73 | ||
|
|
2264a33085 | ||
|
|
d2585097f3 | ||
|
|
e86063155f | ||
|
|
1f7b7b9132 | ||
|
|
7043982ae3 | ||
|
|
57914967cb | ||
|
|
26b5e87457 |
+19
-5
@@ -18,7 +18,9 @@ ARG NPM_CONFIG_REGISTRY=
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 1: Frontend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM ${NODE_IMAGE} AS frontend-builder
|
||||
# --platform=$BUILDPLATFORM: the frontend output is JS (arch-neutral), so build
|
||||
# it on the native host arch instead of under QEMU emulation for the target.
|
||||
FROM --platform=${BUILDPLATFORM} ${NODE_IMAGE} AS frontend-builder
|
||||
ARG NPM_CONFIG_REGISTRY
|
||||
|
||||
WORKDIR /app/frontend
|
||||
@@ -44,7 +46,11 @@ RUN pnpm run build
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 2: Backend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM ${GOLANG_IMAGE} AS backend-builder
|
||||
# --platform=$BUILDPLATFORM: run the Go toolchain on the native host arch and
|
||||
# cross-compile to the target arch below. The binary is CGO_ENABLED=0, so this
|
||||
# is a clean pure-Go cross-compile — no QEMU emulation of go mod download / go
|
||||
# build (emulated networking here was dropping module fetches with EOF).
|
||||
FROM --platform=${BUILDPLATFORM} ${GOLANG_IMAGE} AS backend-builder
|
||||
|
||||
# Build arguments for version info (set by CI)
|
||||
ARG VERSION=
|
||||
@@ -52,6 +58,9 @@ ARG COMMIT=docker
|
||||
ARG DATE
|
||||
ARG GOPROXY
|
||||
ARG GOSUMDB
|
||||
# Populated by buildx from the --platform target (e.g. linux/amd64).
|
||||
ARG TARGETOS
|
||||
ARG TARGETARCH
|
||||
|
||||
ENV GOPROXY=${GOPROXY}
|
||||
ENV GOSUMDB=${GOSUMDB}
|
||||
@@ -63,7 +72,10 @@ WORKDIR /app/backend
|
||||
|
||||
# Copy go mod files first (better caching)
|
||||
COPY backend/go.mod backend/go.sum ./
|
||||
RUN go mod download
|
||||
# Cache mount keeps the module cache across builds so a transient CDN blip on
|
||||
# retry resumes instead of re-fetching every zip from scratch.
|
||||
RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \
|
||||
go mod download
|
||||
|
||||
# Copy backend source first
|
||||
COPY backend/ ./
|
||||
@@ -73,10 +85,12 @@ COPY --from=frontend-builder /app/backend/internal/web/dist ./internal/web/dist
|
||||
|
||||
# Build the binary (BuildType=release for CI builds, embed frontend)
|
||||
# Version precedence: build arg VERSION > exact git tag > cmd/server/VERSION
|
||||
RUN VERSION_VALUE="${VERSION}" && \
|
||||
RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \
|
||||
--mount=type=cache,id=sub2api-gobuild,target=/root/.cache/go-build \
|
||||
VERSION_VALUE="${VERSION}" && \
|
||||
if [ -z "${VERSION_VALUE}" ]; then VERSION_VALUE="$(./scripts/resolve-version.sh)"; fi && \
|
||||
DATE_VALUE="${DATE:-$(date -u +%Y-%m-%dT%H:%M:%SZ)}" && \
|
||||
CGO_ENABLED=0 GOOS=linux go build \
|
||||
CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build \
|
||||
-tags embed \
|
||||
-ldflags="-s -w -X main.Version=${VERSION_VALUE} -X main.Commit=${COMMIT} -X main.Date=${DATE_VALUE} -X main.BuildType=release" \
|
||||
-trimpath \
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
# Sub2API
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="assets/logo.svg" alt="Sub2API Logo" width="128" />
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
@@ -77,8 +79,8 @@ Please read the following carefully before using this project:
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>Thanks to PPToken.org for sponsoring this project! <a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> specializes in GPT model API relay services, supporting Codex, Claude Code, OpenAI-compatible clients, and Gemini CLI integration. Top-ups are 1:1 (¥1 = $1 credit); GPT models start at 0.16x rate multiplier, with overall cost at roughly 2.2% of official pricing and first-token latency around 1 second — ideal for developers seeking low-cost, high-speed access to GPT model capabilities. Technical support: 24/7 real human responses (no bots), @tech in the group chat and get a reply within 10 minutes. Sponsor benefit: the first 200 users who register via the <a href="https://api.pptoken.org/register?promo=SUB2API">exclusive registration link</a> and enter promo code `SUB2API` can claim free Codex / Claude Code trial credits — no minimum spend, no card required.
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>Thanks to PPToken.cc for sponsoring this project! <a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> specializes in GPT model API relay services, supporting Codex, Claude Code, OpenAI-compatible clients, and Gemini CLI integration. Top-ups are 1:1 (¥1 = $1 credit); GPT models start at 0.16x rate multiplier, with overall cost at roughly 2.2% of official pricing and first-token latency around 1 second — ideal for developers seeking low-cost, high-speed access to GPT model capabilities. Technical support: 24/7 real human responses (no bots), @tech in the group chat and get a reply within 10 minutes. Sponsor benefit: the first 200 users who register via the <a href="https://api.pptoken.cc/register?promo=SUB2API">exclusive registration link</a> and enter promo code `SUB2API` can claim free Codex / Claude Code trial credits — no minimum spend, no card required.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -571,9 +573,18 @@ Additional security-related options are available in `config.yaml`:
|
||||
- `security.response_headers.enabled` to enable configurable response header filtering (disabled uses default allowlist)
|
||||
- `security.csp` to control Content-Security-Policy headers
|
||||
- `billing.circuit_breaker` to fail closed on billing errors
|
||||
- `server.trusted_proxies` to enable X-Forwarded-For parsing
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` enables legacy raw forwarded-header takeover (enabled by default for upgrade compatibility); disable it to enforce `server.trusted_proxies`, which should contain only the exact proxy CIDRs that connect directly to Sub2API
|
||||
- `security.forwarded_client_ip_headers` configures up to 16 third-party CDN client-IP header names; they are checked in order before the built-in headers only while legacy takeover is enabled
|
||||
- `turnstile.required` to require Turnstile in release mode
|
||||
|
||||
Custom client-IP headers can be set in YAML or as a comma-separated environment variable:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
Header names are validated, canonicalized, and de-duplicated. The admin security settings can update the list without a restart; new installations persist YAML/environment defaults and existing installations backfill a missing database value. When legacy takeover is disabled, all custom and built-in raw forwarding headers are ignored and Gin uses only `server.trusted_proxies`. While takeover is enabled, firewall the origin to CDN/proxy addresses and make the edge overwrite every trusted client-IP header. See [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md) for the complete migration and trust-boundary rules.
|
||||
|
||||
**⚠️ Security Warning: HTTP URL Configuration**
|
||||
|
||||
When `security.url_allowlist.enabled=false`, the system performs minimal URL validation and **allows HTTP URLs by default** (dev-friendly mode; Docker Compose deployments use the same default). For production, explicitly tighten this to HTTPS-only:
|
||||
@@ -639,8 +650,20 @@ override this limit.
|
||||
The connection cap is coordinated through Redis using a 60-second lease that
|
||||
is refreshed every 20 seconds. A process that cannot confirm a lease for a
|
||||
full lease lifetime closes its local WebSocket rather than continuing outside
|
||||
the global cap. Use `http_bridge` for client-WebSocket/upstream-HTTP operation
|
||||
when rolling out or mitigating upstream WebSocket issues.
|
||||
the global cap.
|
||||
|
||||
Enable the v2 mode router before selecting an account-level WS mode such as
|
||||
`http_bridge`:
|
||||
|
||||
```yaml
|
||||
gateway:
|
||||
openai_ws:
|
||||
mode_router_v2_enabled: true
|
||||
```
|
||||
|
||||
Or set `GATEWAY_OPENAI_WS_MODE_ROUTER_V2_ENABLED=true` in the environment.
|
||||
Use `http_bridge` for client-WebSocket/upstream-HTTP operation when rolling out
|
||||
or mitigating upstream WebSocket issues.
|
||||
|
||||
#### ⚠️ Important: Creating the Admin Account
|
||||
|
||||
@@ -788,9 +811,9 @@ xAI quota is passive. Sub2API does not invent subscription quota values; it reco
|
||||
|
||||
`401` responses temporarily remove accounts with invalid credentials from scheduling. `403` responses are treated as access or entitlement failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling.
|
||||
|
||||
New Grok image and video generation requests use a media-specific eligibility check. An OAuth account is excluded from new media generation when its recorded weekly or monthly billing probe returns `403`; chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account` instead of forwarding the request to a known-ineligible account.
|
||||
New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts require positive paid-entitlement evidence from the xAI billing probe; Free, forbidden, missing, malformed, and inconclusive billing observations are excluded from new media generation. Unobserved OAuth accounts are probed before the first media request is forwarded, and imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`.
|
||||
|
||||
Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A missing billing observation does not block legacy routing, and a weekly allowance period by itself is not treated as evidence that the account is ineligible.
|
||||
Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A weekly allowance period alone is not treated as a paid tier signal. Successful image responses must contain at least one actual image output; empty HTTP `200` responses trigger account failover instead of being counted and returned as successful generations.
|
||||
|
||||
---
|
||||
|
||||
|
||||
+16
-5
@@ -1,7 +1,9 @@
|
||||
# Sub2API
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="assets/logo.svg" alt="Sub2API Logo" width="128" />
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
@@ -80,8 +82,8 @@
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>感谢 PPToken.org 赞助本项目! <a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> 主打 GPT 系列模型 API 中转服务,支持 Codex、Claude Code、OpenAI 兼容客户端及 Gemini CLI 等工具接入。充值 1:1,1 元=1 美元额度;GPT 模型最低 0.16 倍倍率,综合成本约为官方价格的 0.22 折,最快首字 Token 约 1 秒,适合开发者低成本、高响应速度接入 GPT 模型能力。技术支持: 7×24 小时真人响应(不是机器人),群内@技术,10 分钟内有回复 。赞助商福利:前 200 名用户通过 <a href="https://api.pptoken.org/register?promo=SUB2API">[专属注册链接]</a> 注册,输入优惠码 `SUB2API`,可领取 Codex / Claude Code 免费试用额度,无门槛、不绑卡。
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>感谢 PPToken.cc 赞助本项目! <a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> 主打 GPT 系列模型 API 中转服务,支持 Codex、Claude Code、OpenAI 兼容客户端及 Gemini CLI 等工具接入。充值 1:1,1 元=1 美元额度;GPT 模型最低 0.16 倍倍率,综合成本约为官方价格的 0.22 折,最快首字 Token 约 1 秒,适合开发者低成本、高响应速度接入 GPT 模型能力。技术支持: 7×24 小时真人响应(不是机器人),群内@技术,10 分钟内有回复 。赞助商福利:前 200 名用户通过 <a href="https://api.pptoken.cc/register?promo=SUB2API">[专属注册链接]</a> 注册,输入优惠码 `SUB2API`,可领取 Codex / Claude Code 免费试用额度,无门槛、不绑卡。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -607,9 +609,18 @@ gateway:
|
||||
- `security.response_headers.enabled` 可启用可配置响应头过滤(关闭时使用默认白名单)
|
||||
- `security.csp` 配置 Content-Security-Policy
|
||||
- `billing.circuit_breaker` 计费异常时 fail-closed
|
||||
- `server.trusted_proxies` 启用可信代理解析 X-Forwarded-For
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` 控制旧版原始转发头接管(为升级兼容默认开启);关闭后严格使用 `server.trusted_proxies`,其中只应填写直接连接 Sub2API 的精确代理 CIDR
|
||||
- `security.forwarded_client_ip_headers` 最多配置 16 个第三方 CDN 客户端 IP 请求头;仅在旧版接管开启时按顺序优先于内置请求头解析
|
||||
- `turnstile.required` 在 release 模式强制启用 Turnstile
|
||||
|
||||
自定义客户端 IP 请求头可通过 YAML 配置,也可使用逗号分隔的环境变量:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
请求头名称会经过合法性校验、规范化和大小写无关去重。管理员可在安全设置中动态更新列表,无需重启;新安装会持久化 YAML/环境变量默认值,旧安装缺少数据库字段时会自动回填。关闭旧版接管后,自定义头和内置原始转发头均被忽略,只使用 `server.trusted_proxies`。开启接管时必须限制源站仅允许 CDN/代理访问,并确保边缘代理覆盖所有受信客户端 IP 请求头。完整迁移规则和信任边界见 [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md)。
|
||||
|
||||
**网关防御纵深建议(重点)**
|
||||
|
||||
- `gateway.upstream_response_read_max_bytes`:限制非流式上游响应读取大小(默认 `8MB`),用于防止异常响应导致内存放大。
|
||||
|
||||
+16
-5
@@ -1,7 +1,9 @@
|
||||
# Sub2API
|
||||
|
||||
<div align="center">
|
||||
|
||||
<img src="assets/logo.svg" alt="Sub2API Logo" width="128" />
|
||||
|
||||
# Sub2API
|
||||
|
||||
[](https://golang.org/)
|
||||
[](https://vuejs.org/)
|
||||
[](https://www.postgresql.org/)
|
||||
@@ -79,8 +81,8 @@
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://api.pptoken.org/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>PPToken.org のご支援に感謝します!<a href="https://api.pptoken.org/register?promo=SUB2API">PPToken.org</a> は GPT シリーズモデルの API 中継サービスを専門としており、Codex、Claude Code、OpenAI 互換クライアント、Gemini CLI などのツール接続をサポートしています。チャージは 1:1(1元=1ドル分のクレジット)、GPT モデルは最低 0.16 倍のレート倍率で、総合コストは公式価格の約 2.2% 、最速ファーストトークンは約1秒 — 開発者が低コスト・高速レスポンスで GPT モデル機能にアクセスするのに最適です。テクニカルサポート:24時間365日リアルな人間が対応(ボットではありません)、グループ内で @技術 すれば 10 分以内に返信。スポンサー特典:先着 200 名のユーザーが<a href="https://api.pptoken.org/register?promo=SUB2API">専用登録リンク</a>から登録し、プロモコード `SUB2API` を入力すると、Codex / Claude Code の無料トライアルクレジットを獲得できます — 最低利用額なし、カード登録不要。
|
||||
<td width="180"><a href="https://api.pptoken.cc/register?promo=SUB2API"><img src="assets/partners/logos/pptoken.png" alt="pptoken" width="150"></a></td>
|
||||
<td>PPToken.cc のご支援に感謝します!<a href="https://api.pptoken.cc/register?promo=SUB2API">PPToken.cc</a> は GPT シリーズモデルの API 中継サービスを専門としており、Codex、Claude Code、OpenAI 互換クライアント、Gemini CLI などのツール接続をサポートしています。チャージは 1:1(1元=1ドル分のクレジット)、GPT モデルは最低 0.16 倍のレート倍率で、総合コストは公式価格の約 2.2% 、最速ファーストトークンは約1秒 — 開発者が低コスト・高速レスポンスで GPT モデル機能にアクセスするのに最適です。テクニカルサポート:24時間365日リアルな人間が対応(ボットではありません)、グループ内で @技術 すれば 10 分以内に返信。スポンサー特典:先着 200 名のユーザーが<a href="https://api.pptoken.cc/register?promo=SUB2API">専用登録リンク</a>から登録し、プロモコード `SUB2API` を入力すると、Codex / Claude Code の無料トライアルクレジットを獲得できます — 最低利用額なし、カード登録不要。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
@@ -569,9 +571,18 @@ default:
|
||||
- `security.response_headers.enabled` - 設定可能なレスポンスヘッダーフィルタリングを有効化(無効時はデフォルトの許可リストを使用)
|
||||
- `security.csp` - Content-Security-Policy ヘッダーの制御
|
||||
- `billing.circuit_breaker` - 課金エラー時にフェイルクローズ
|
||||
- `server.trusted_proxies` - X-Forwarded-For パースの有効化
|
||||
- `security.trust_forwarded_ip_for_api_key_acl` - 従来の生転送ヘッダーによる上書きを制御(アップグレード互換性のため既定で有効)。無効にすると `server.trusted_proxies` を厳格に使用し、Sub2API に直接接続するプロキシの正確な CIDR のみを指定
|
||||
- `security.forwarded_client_ip_headers` - サードパーティ CDN のクライアント IP ヘッダーを最大 16 個指定。従来モードが有効な場合のみ、設定順で組み込みヘッダーより先に評価
|
||||
- `turnstile.required` - リリースモードでの Turnstile 必須化
|
||||
|
||||
カスタムクライアント IP ヘッダーは YAML またはカンマ区切りの環境変数で設定できます:
|
||||
|
||||
```bash
|
||||
SECURITY_FORWARDED_CLIENT_IP_HEADERS=True-Client-IP,X-CDN-Client-IP
|
||||
```
|
||||
|
||||
ヘッダー名は検証、正規化、大小文字を区別しない重複排除が行われます。管理画面のセキュリティ設定から再起動せずに更新でき、新規インストールでは YAML/環境変数の既定値を保存し、既存環境ではデータベース値がない場合に補完します。従来モードを無効にするとカスタムおよび組み込みの生転送ヘッダーはすべて無視され、`server.trusted_proxies` のみを使用します。有効にする場合はオリジンへの接続元を CDN/プロキシに制限し、エッジで信頼する全クライアント IP ヘッダーを上書きしてください。移行規則と信頼境界の詳細は [`deploy/EDGE_SECURITY.md`](deploy/EDGE_SECURITY.md) を参照してください。
|
||||
|
||||
**⚠️ セキュリティ警告: HTTP URL 設定**
|
||||
|
||||
`security.url_allowlist.enabled=false` の場合、システムは最小限の URL バリデーションのみを行い、**デフォルトで HTTP URL を許可**します(開発フレンドリーモード。Docker Compose デプロイのデフォルトも同じです)。本番環境では、以下のように明示的に HTTPS のみに制限することを推奨します:
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512" role="img" aria-labelledby="title desc">
|
||||
<title id="title">Sub2API</title>
|
||||
<desc id="desc">An interlocking S symbol representing subscription routing into APIs.</desc>
|
||||
<defs>
|
||||
<linearGradient id="s2a-bg" x1="72" y1="44" x2="442" y2="478" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#142B56"/>
|
||||
<stop offset=".52" stop-color="#0A1A39"/>
|
||||
<stop offset="1" stop-color="#061127"/>
|
||||
</linearGradient>
|
||||
<radialGradient id="s2a-ambient" cx="0" cy="0" r="1" gradientUnits="userSpaceOnUse" gradientTransform="translate(168 92) rotate(51) scale(342 382)">
|
||||
<stop stop-color="#3E68B0" stop-opacity=".28"/>
|
||||
<stop offset="1" stop-color="#3E68B0" stop-opacity="0"/>
|
||||
</radialGradient>
|
||||
<linearGradient id="s2a-brand" x1="4" y1="4" x2="20" y2="20" gradientUnits="userSpaceOnUse">
|
||||
<stop stop-color="#79F4BD"/>
|
||||
<stop offset=".48" stop-color="#39D9E7"/>
|
||||
<stop offset="1" stop-color="#3875F6"/>
|
||||
</linearGradient>
|
||||
<filter id="s2a-shadow" x="-25%" y="-25%" width="150%" height="165%" color-interpolation-filters="sRGB">
|
||||
<feDropShadow dx="0" dy=".55" stdDeviation=".65" flood-color="#020817" flood-opacity=".34"/>
|
||||
</filter>
|
||||
</defs>
|
||||
|
||||
<rect x="16" y="16" width="480" height="480" rx="120" fill="url(#s2a-bg)"/>
|
||||
<rect x="16" y="16" width="480" height="480" rx="120" fill="url(#s2a-ambient)"/>
|
||||
<rect x="16.75" y="16.75" width="478.5" height="478.5" rx="119.25" fill="none" stroke="#A5BFFF" stroke-opacity=".16" stroke-width="1.5"/>
|
||||
|
||||
<g transform="translate(52 52) scale(17)" fill="none" stroke="url(#s2a-brand)" stroke-linecap="round" stroke-linejoin="round" stroke-width="2.7" filter="url(#s2a-shadow)">
|
||||
<path d="m19.25 7.65-2.55-3.2a1.33 1.33 0 0 0-1.03-.5H8.58c-.34 0-.67.13-.91.37L4.15 7.65c-.93.88-.6 1.52.65 2.3l9.55 5.97"/>
|
||||
<path d="m4.75 16.35 2.55 3.2c.25.31.63.5 1.03.5h7.09c.34 0 .67-.13.91-.37l3.52-3.33c.93-.88.6-1.52-.65-2.3L9.65 8.08"/>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 2.0 KiB |
@@ -0,0 +1,20 @@
|
||||
# Ingress rejection log cleanup
|
||||
|
||||
This maintenance command removes historical admission rejections from
|
||||
`ops_error_logs` without matching unrelated authentication or upstream errors.
|
||||
It is a dry run unless `--execute` is supplied, and always requires an explicit
|
||||
RFC3339 cutoff.
|
||||
|
||||
```sh
|
||||
go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z
|
||||
go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z --execute
|
||||
```
|
||||
|
||||
Run the execute form only after every application instance has been upgraded so
|
||||
older instances cannot add new ingress rejection rows below the chosen cutoff.
|
||||
The classifier intentionally retains invariant failures such as
|
||||
`USER_NOT_FOUND`, database errors, quota/billing errors, and upstream failures.
|
||||
|
||||
After the rollout and cleanup are verified, run
|
||||
`backend/scripts/finalize-ingress-reject-cleanup.sql` in a maintenance window to
|
||||
remove the deprecated plaintext-key audit table and attribution columns.
|
||||
@@ -0,0 +1,218 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
_ "github.com/Wei-Shaw/sub2api/ent/runtime"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/repository"
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
const classifierVersion = "ingress-reject-v1"
|
||||
|
||||
type candidate struct {
|
||||
id int64
|
||||
statusCode int
|
||||
message string
|
||||
body string
|
||||
}
|
||||
|
||||
func main() {
|
||||
beforeRaw := flag.String("before", "", "required RFC3339 cutoff; only older rows are considered")
|
||||
execute := flag.Bool("execute", false, "delete matched rows (default is dry-run)")
|
||||
batchSize := flag.Int("batch-size", 5000, "scan/delete batch size (1-5000)")
|
||||
flag.Parse()
|
||||
|
||||
if *beforeRaw == "" {
|
||||
log.Fatal("--before is required")
|
||||
}
|
||||
before, err := time.Parse(time.RFC3339, *beforeRaw)
|
||||
if err != nil {
|
||||
log.Fatalf("invalid --before: %v", err)
|
||||
}
|
||||
if *batchSize < 1 || *batchSize > 5000 {
|
||||
log.Fatal("--batch-size must be between 1 and 5000")
|
||||
}
|
||||
|
||||
cfg, err := config.LoadForBootstrap()
|
||||
if err != nil {
|
||||
log.Fatalf("load config: %v", err)
|
||||
}
|
||||
client, db, err := repository.InitEnt(cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("initialize database: %v", err)
|
||||
}
|
||||
defer func() { _ = client.Close() }()
|
||||
|
||||
ctx := context.Background()
|
||||
counts, scanned, matched, deleted, err := cleanup(ctx, db, before, *batchSize, *execute)
|
||||
if err != nil {
|
||||
log.Fatalf("cleanup failed: %v", err)
|
||||
}
|
||||
|
||||
digest := sha256.Sum256([]byte(classifierVersion))
|
||||
mode := "dry-run"
|
||||
if *execute {
|
||||
mode = "execute"
|
||||
}
|
||||
fmt.Printf("mode=%s before=%s classifier=%s scanned=%d matched=%d deleted=%d\n",
|
||||
mode, before.UTC().Format(time.RFC3339), hex.EncodeToString(digest[:]), scanned, matched, deleted)
|
||||
reasons := make([]string, 0, len(counts))
|
||||
for reason := range counts {
|
||||
reasons = append(reasons, reason)
|
||||
}
|
||||
sort.Strings(reasons)
|
||||
for _, reason := range reasons {
|
||||
fmt.Printf("reason=%s count=%d\n", reason, counts[reason])
|
||||
}
|
||||
if *execute && deleted > 0 {
|
||||
fmt.Println("cleanup complete; schedule VACUUM (ANALYZE) ops_error_logs during normal maintenance")
|
||||
}
|
||||
}
|
||||
|
||||
func cleanup(ctx context.Context, db *sql.DB, before time.Time, batchSize int, execute bool) (map[string]int64, int64, int64, int64, error) {
|
||||
counts := make(map[string]int64)
|
||||
var cursor, scanned, matched, deleted int64
|
||||
for {
|
||||
rows, err := db.QueryContext(ctx, `
|
||||
SELECT id, COALESCE(status_code, 0), COALESCE(error_message, ''), COALESCE(error_body, '')
|
||||
FROM ops_error_logs
|
||||
WHERE id > $1
|
||||
AND created_at < $2
|
||||
AND error_phase = 'auth'
|
||||
AND account_id IS NULL
|
||||
AND upstream_status_code IS NULL
|
||||
AND COALESCE(upstream_error_message, '') = ''
|
||||
AND COALESCE(upstream_error_detail, '') = ''
|
||||
ORDER BY id ASC
|
||||
LIMIT $3`, cursor, before, batchSize)
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
|
||||
batch := make([]candidate, 0, batchSize)
|
||||
for rows.Next() {
|
||||
var item candidate
|
||||
if err := rows.Scan(&item.id, &item.statusCode, &item.message, &item.body); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
batch = append(batch, item)
|
||||
cursor = item.id
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
_ = rows.Close()
|
||||
if len(batch) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
ids := make([]int64, 0, len(batch))
|
||||
for _, item := range batch {
|
||||
scanned++
|
||||
if reason, ok := historicalIngressRejectReason(item); ok {
|
||||
matched++
|
||||
counts[reason]++
|
||||
ids = append(ids, item.id)
|
||||
}
|
||||
}
|
||||
if execute && len(ids) > 0 {
|
||||
result, err := db.ExecContext(ctx,
|
||||
`DELETE FROM ops_error_logs WHERE id = ANY($1) AND created_at < $2`, pq.Array(ids), before)
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
n, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
deleted += n
|
||||
}
|
||||
}
|
||||
return counts, scanned, matched, deleted, nil
|
||||
}
|
||||
|
||||
func historicalIngressRejectReason(item candidate) (string, bool) {
|
||||
code, message := parseErrorIdentity(item.body, item.message)
|
||||
switch code {
|
||||
case "API_KEY_REQUIRED":
|
||||
return "missing_key", true
|
||||
case "INVALID_API_KEY":
|
||||
return "invalid_key", true
|
||||
case "API_KEY_DISABLED":
|
||||
return "key_disabled", true
|
||||
case "USER_INACTIVE":
|
||||
return "user_inactive", true
|
||||
case "GROUP_DELETED":
|
||||
return "group_deleted", true
|
||||
case "GROUP_DISABLED":
|
||||
return "group_disabled", true
|
||||
case "GROUP_NOT_ALLOWED":
|
||||
return "group_forbidden", true
|
||||
case "ACCESS_DENIED":
|
||||
return "ip_acl_denied", true
|
||||
case "api_key_in_query_deprecated":
|
||||
return "query_key_deprecated", true
|
||||
}
|
||||
|
||||
normalized := strings.TrimSpace(message)
|
||||
switch {
|
||||
case normalized == "API key is required":
|
||||
return "missing_key", true
|
||||
case normalized == "Invalid API key":
|
||||
return "invalid_key", true
|
||||
case normalized == "API key is disabled":
|
||||
return "key_disabled", true
|
||||
case normalized == "User account is not active":
|
||||
return "user_inactive", true
|
||||
case normalized == "API Key 所属分组已删除":
|
||||
return "group_deleted", true
|
||||
case normalized == "API Key 所属分组已停用":
|
||||
return "group_disabled", true
|
||||
case normalized == "API Key 所属专属分组不再允许当前用户使用":
|
||||
return "group_forbidden", true
|
||||
case normalized == "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.":
|
||||
return "group_unassigned", true
|
||||
case strings.HasPrefix(normalized, "Access denied. Your IP is "):
|
||||
return "ip_acl_denied", true
|
||||
case normalized == "Query parameter api_key is deprecated. Use Authorization header or key instead.":
|
||||
return "query_key_deprecated", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func parseErrorIdentity(body, fallbackMessage string) (string, string) {
|
||||
var payload struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Error struct {
|
||||
Code json.RawMessage `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(body), &payload); err != nil {
|
||||
return "", fallbackMessage
|
||||
}
|
||||
message := payload.Message
|
||||
if message == "" {
|
||||
message = payload.Error.Message
|
||||
}
|
||||
if message == "" {
|
||||
message = fallbackMessage
|
||||
}
|
||||
return strings.TrimSpace(payload.Code), message
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHistoricalIngressRejectReason(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
item candidate
|
||||
reason string
|
||||
match bool
|
||||
}{
|
||||
{name: "standard invalid key", item: candidate{body: `{"code":"INVALID_API_KEY","message":"Invalid API key"}`}, reason: "invalid_key", match: true},
|
||||
{name: "google missing key", item: candidate{body: `{"error":{"code":401,"message":"API key is required","status":"UNAUTHENTICATED"}}`}, reason: "missing_key", match: true},
|
||||
{name: "google group deleted", item: candidate{body: `{"error":{"code":403,"message":"API Key 所属分组已删除","status":"PERMISSION_DENIED"}}`}, reason: "group_deleted", match: true},
|
||||
{name: "ip acl", item: candidate{body: `{"code":"ACCESS_DENIED","message":"Access denied. Your IP is 192.0.2.1"}`}, reason: "ip_acl_denied", match: true},
|
||||
{name: "user not found remains", item: candidate{body: `{"code":"USER_NOT_FOUND","message":"User associated with API key not found"}`}, match: false},
|
||||
{name: "quota remains", item: candidate{body: `{"code":"API_KEY_QUOTA_EXHAUSTED","message":"quota"}`}, match: false},
|
||||
{name: "database failure remains", item: candidate{statusCode: 500, message: "Failed to validate API key", body: `{"code":"INTERNAL_ERROR","message":"Failed to validate API key"}`}, match: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
reason, ok := historicalIngressRejectReason(tt.item)
|
||||
if ok != tt.match || reason != tt.reason {
|
||||
t.Fatalf("got (%q, %v), want (%q, %v)", reason, ok, tt.reason, tt.match)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1 +1 @@
|
||||
0.1.159
|
||||
0.1.161
|
||||
|
||||
@@ -155,10 +155,11 @@ func runMainServer() {
|
||||
defer app.Cleanup()
|
||||
if app.PromptAudit != nil {
|
||||
if err := app.PromptAudit.Start(context.Background()); err != nil {
|
||||
// Startup continues so unrelated APIs stay up, but Prompt Audit itself
|
||||
// fails closed (unavailable) until a later reload installs a trusted
|
||||
// snapshot—avoiding a silent ModeOff bypass of persisted blocking policy.
|
||||
log.Printf("Prompt Audit started in degraded fail-closed state: %v", err)
|
||||
// Startup continues so unrelated APIs stay up. Fail-closed (unavailable)
|
||||
// applies only when a persisted blocking policy was observed; without
|
||||
// blocking intent, Prompt Audit stays ModeOff so the gateway remains
|
||||
// usable and administrators can still disable the feature (#4560).
|
||||
log.Printf("Prompt Audit started in degraded state: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -81,6 +81,10 @@ func provideCleanup(
|
||||
opsCleanup *service.OpsCleanupService,
|
||||
opsScheduledReport *service.OpsScheduledReportService,
|
||||
opsSystemLogSink *service.OpsSystemLogSink,
|
||||
opsService *service.OpsService,
|
||||
opsIngressReject *service.OpsIngressRejectAggregator,
|
||||
apiKeyService *service.APIKeyService,
|
||||
authCacheInvalidationWorker *service.AuthCacheInvalidationWorker,
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
@@ -121,6 +125,30 @@ func provideCleanup(
|
||||
|
||||
// 应用层清理步骤可并行执行,基础设施资源(Redis/Ent)最后按顺序关闭。
|
||||
parallelSteps := []cleanupStep{
|
||||
{"OpsIngressRejectAggregator", func() error {
|
||||
if opsIngressReject != nil {
|
||||
opsIngressReject.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationWorker", func() error {
|
||||
if authCacheInvalidationWorker != nil {
|
||||
authCacheInvalidationWorker.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationSubscriber", func() error {
|
||||
if apiKeyService != nil {
|
||||
apiKeyService.StopAuthCacheInvalidationSubscriber()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsRuntimeSettingsRefresh", func() error {
|
||||
if opsService != nil {
|
||||
opsService.StopRuntimeSettingsRefresh()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"PromptAuditService", func() error {
|
||||
if promptAudit != nil {
|
||||
return promptAudit.Shutdown(ctx)
|
||||
|
||||
@@ -157,7 +157,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService, internal500CounterCache)
|
||||
geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig)
|
||||
opsSystemLogSink := service.ProvideOpsSystemLogSink(opsRepository)
|
||||
opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService)
|
||||
authCacheInvalidationOutboxRepository := repository.NewAuthCacheInvalidationOutboxRepository(db)
|
||||
authCacheInvalidationWorker := service.ProvideAuthCacheInvalidationWorker(authCacheInvalidationOutboxRepository, apiKeyCache, apiKeyService)
|
||||
opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService, authCacheInvalidationWorker, apiKeyService)
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService, opsService, settingService)
|
||||
redeemHandler := handler.NewRedeemHandler(redeemService)
|
||||
subscriptionHandler := handler.NewSubscriptionHandler(subscriptionService)
|
||||
@@ -179,7 +181,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
|
||||
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
|
||||
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
|
||||
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
|
||||
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
|
||||
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
|
||||
@@ -198,7 +200,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
|
||||
dbDumper := repository.NewPgDumper(configConfig)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper)
|
||||
backupHandler := admin.NewBackupHandler(backupService, userService)
|
||||
imageStorageFactory := repository.ProvideImageStorageFactory()
|
||||
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
|
||||
backupHandler := admin.NewBackupHandler(backupService, userService, imageStorageSettingService)
|
||||
oAuthHandler := admin.NewOAuthHandler(oAuthService)
|
||||
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
|
||||
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
|
||||
@@ -216,7 +220,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
registry := payment.ProvideRegistry()
|
||||
defaultLoadBalancer := payment.ProvideDefaultLoadBalancer(client, encryptionKey)
|
||||
paymentService := service.ProvidePaymentService(client, registry, defaultLoadBalancer, redeemService, subscriptionService, paymentConfigService, userRepository, groupRepository, affiliateService, notificationEmailService)
|
||||
settingHandler := handler.ProvideAdminSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService, notificationEmailService)
|
||||
settingHandler := handler.ProvideAdminSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService, notificationEmailService, totpService, userService)
|
||||
opsHandler := admin.NewOpsHandler(opsService)
|
||||
updateCache := repository.NewUpdateCache(redisClient)
|
||||
gitHubReleaseClient := repository.ProvideGitHubReleaseClient(configConfig)
|
||||
@@ -270,18 +274,14 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
|
||||
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
|
||||
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
|
||||
imageTaskStore := repository.NewImageTaskStore(redisClient)
|
||||
imageStorage, err := repository.ProvideImageStorage(configConfig)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
imageTaskService := service.ProvideImageTaskService(imageTaskStore, imageStorage, configConfig)
|
||||
imageTaskService := service.ProvideImageTaskService(imageTaskStore, imageStorageSettingService)
|
||||
asyncImageHandler := handler.NewAsyncImageHandler(imageTaskService, openAIGatewayHandler)
|
||||
batchImageRepository := repository.NewBatchImageRepository(db)
|
||||
batchImageQueue := repository.NewBatchImageQueue(redisClient, configConfig)
|
||||
@@ -298,7 +298,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService, auditLogService)
|
||||
apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig)
|
||||
auditLogMiddleware := middleware.NewAuditLogMiddleware(auditLogService)
|
||||
stepUpAuthMiddleware := middleware.NewStepUpAuthMiddleware(totpService, userService)
|
||||
stepUpAuthMiddleware := middleware.NewStepUpAuthMiddleware(totpService, userService, settingService)
|
||||
engine := server.ProvideRouter(configConfig, handlers, jwtAuthMiddleware, adminAuthMiddleware, apiKeyAuthMiddleware, auditLogMiddleware, stepUpAuthMiddleware, apiKeyService, subscriptionService, opsService, settingService, redisClient)
|
||||
httpServer := server.ProvideHTTPServer(configConfig, engine)
|
||||
opsMetricsCollector := service.ProvideOpsMetricsCollector(opsRepository, settingRepository, accountRepository, concurrencyService, db, redisClient, configConfig)
|
||||
@@ -306,6 +306,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
opsAlertEvaluatorService := service.ProvideOpsAlertEvaluatorService(opsService, opsRepository, emailService, redisClient, configConfig, proxyRepository)
|
||||
opsCleanupService := service.ProvideOpsCleanupService(opsRepository, db, redisClient, configConfig, channelMonitorService, settingRepository, opsService)
|
||||
opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig)
|
||||
opsIngressRejectAggregator := service.ProvideOpsIngressRejectAggregator(opsRepository, opsService)
|
||||
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
|
||||
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
|
||||
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
|
||||
@@ -314,7 +315,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db)
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
|
||||
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService, promptService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService, promptService)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
PromptAudit: promptService,
|
||||
@@ -351,6 +352,10 @@ func provideCleanup(
|
||||
opsCleanup *service.OpsCleanupService,
|
||||
opsScheduledReport *service.OpsScheduledReportService,
|
||||
opsSystemLogSink *service.OpsSystemLogSink,
|
||||
opsService *service.OpsService,
|
||||
opsIngressReject *service.OpsIngressRejectAggregator,
|
||||
apiKeyService *service.APIKeyService,
|
||||
authCacheInvalidationWorker *service.AuthCacheInvalidationWorker,
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
@@ -390,6 +395,30 @@ func provideCleanup(
|
||||
}
|
||||
|
||||
parallelSteps := []cleanupStep{
|
||||
{"OpsIngressRejectAggregator", func() error {
|
||||
if opsIngressReject != nil {
|
||||
opsIngressReject.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationWorker", func() error {
|
||||
if authCacheInvalidationWorker != nil {
|
||||
authCacheInvalidationWorker.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationSubscriber", func() error {
|
||||
if apiKeyService != nil {
|
||||
apiKeyService.StopAuthCacheInvalidationSubscriber()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsRuntimeSettingsRefresh", func() error {
|
||||
if opsService != nil {
|
||||
opsService.StopRuntimeSettingsRefresh()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"PromptAuditService", func() error {
|
||||
if promptAudit != nil {
|
||||
return promptAudit.Shutdown(ctx)
|
||||
|
||||
@@ -58,6 +58,10 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
&service.OpsCleanupService{},
|
||||
&service.OpsScheduledReportService{},
|
||||
opsSystemLogSinkSvc,
|
||||
nil, // opsService
|
||||
nil, // opsIngressRejectAggregator
|
||||
nil, // apiKeyService
|
||||
nil, // authCacheInvalidationWorker
|
||||
schedulerSnapshotSvc,
|
||||
tokenRefreshSvc,
|
||||
accountExpirySvc,
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"math"
|
||||
"net/textproto"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
@@ -14,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
"golang.org/x/net/http/httpguts"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -255,6 +257,22 @@ func (c *ImageStorageConfig) Active() bool {
|
||||
return c.Enabled && c.IsConfigured()
|
||||
}
|
||||
|
||||
// MissingCredentialKeys 返回 IsConfigured 所缺的配置键名。
|
||||
// 用于启动日志:只说"凭证不完整"会让运维以为自己漏填了,而实际可能是值填了却没被读到。
|
||||
func (c *ImageStorageConfig) MissingCredentialKeys() []string {
|
||||
var missing []string
|
||||
if c.Bucket == "" {
|
||||
missing = append(missing, "image_storage.bucket")
|
||||
}
|
||||
if c.AccessKeyID == "" {
|
||||
missing = append(missing, "image_storage.access_key_id")
|
||||
}
|
||||
if c.SecretAccessKey == "" {
|
||||
missing = append(missing, "image_storage.secret_access_key")
|
||||
}
|
||||
return missing
|
||||
}
|
||||
|
||||
type LinuxDoConnectConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
ClientID string `mapstructure:"client_id"`
|
||||
@@ -639,16 +657,18 @@ type PricingConfig struct {
|
||||
}
|
||||
|
||||
type ServerConfig struct {
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
Mode string `mapstructure:"mode"` // debug/release
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP)
|
||||
MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制
|
||||
H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置
|
||||
Host string `mapstructure:"host"`
|
||||
Port int `mapstructure:"port"`
|
||||
Mode string `mapstructure:"mode"` // debug/release
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
MaxHeaderBytes int `mapstructure:"max_header_bytes"` // 请求头最大字节数(HTTP/2 映射为 header-list 上限)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP)
|
||||
TrustedProxiesConfigured bool `mapstructure:"-" json:"-" yaml:"-"` // 是否显式配置了可信代理列表
|
||||
MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制
|
||||
H2C H2CConfig `mapstructure:"h2c"` // HTTP/2 Cleartext 配置
|
||||
}
|
||||
|
||||
// H2CConfig HTTP/2 Cleartext 配置
|
||||
@@ -666,36 +686,103 @@ type CORSConfig struct {
|
||||
AllowCredentials bool `mapstructure:"allow_credentials"`
|
||||
}
|
||||
|
||||
const MaxForwardedClientIPHeaders = 16
|
||||
|
||||
type ForwardedClientIPSettings struct {
|
||||
TrustForwardedIP bool
|
||||
Headers []string
|
||||
}
|
||||
|
||||
type SecurityConfig struct {
|
||||
URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"`
|
||||
ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"`
|
||||
CSP CSPConfig `mapstructure:"csp"`
|
||||
ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"`
|
||||
ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"`
|
||||
TrustForwardedIPForAPIKeyACL bool `mapstructure:"trust_forwarded_ip_for_api_key_acl"`
|
||||
trustForwardedIPForAPIKeyACLLive *atomic.Bool `mapstructure:"-"`
|
||||
URLAllowlist URLAllowlistConfig `mapstructure:"url_allowlist"`
|
||||
ResponseHeaders ResponseHeaderConfig `mapstructure:"response_headers"`
|
||||
CSP CSPConfig `mapstructure:"csp"`
|
||||
ProxyFallback ProxyFallbackConfig `mapstructure:"proxy_fallback"`
|
||||
ProxyProbe ProxyProbeConfig `mapstructure:"proxy_probe"`
|
||||
// TrustForwardedIPForAPIKeyACL enables legacy raw forwarded-header takeover.
|
||||
// When disabled, server.trusted_proxies is authoritative for all client-IP consumers.
|
||||
TrustForwardedIPForAPIKeyACL bool `mapstructure:"trust_forwarded_ip_for_api_key_acl"`
|
||||
ForwardedClientIPHeaders []string `mapstructure:"forwarded_client_ip_headers" json:"forwarded_client_ip_headers" yaml:"forwarded_client_ip_headers"`
|
||||
forwardedClientIPSettingsLive *atomic.Pointer[ForwardedClientIPSettings] `mapstructure:"-" json:"-" yaml:"-"`
|
||||
}
|
||||
|
||||
func NormalizeForwardedClientIPHeaders(headers []string) ([]string, error) {
|
||||
normalized := make([]string, 0, len(headers))
|
||||
seen := make(map[string]struct{}, len(headers))
|
||||
for _, header := range headers {
|
||||
header = strings.TrimSpace(header)
|
||||
if !httpguts.ValidHeaderFieldName(header) {
|
||||
return nil, fmt.Errorf("invalid HTTP header field name %q", header)
|
||||
}
|
||||
canonical := textproto.CanonicalMIMEHeaderKey(header)
|
||||
key := strings.ToLower(canonical)
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
if len(normalized) == MaxForwardedClientIPHeaders {
|
||||
return nil, fmt.Errorf("forwarded client IP headers must contain at most %d unique names", MaxForwardedClientIPHeaders)
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
normalized = append(normalized, canonical)
|
||||
}
|
||||
return normalized, nil
|
||||
}
|
||||
|
||||
func cloneForwardedClientIPHeaders(headers []string) []string {
|
||||
if len(headers) == 0 {
|
||||
return []string{}
|
||||
}
|
||||
return append([]string(nil), headers...)
|
||||
}
|
||||
|
||||
func (c *Config) ForwardedClientIPSettings() ForwardedClientIPSettings {
|
||||
if c == nil {
|
||||
return ForwardedClientIPSettings{Headers: []string{}}
|
||||
}
|
||||
live := c.Security.forwardedClientIPSettingsLive
|
||||
if live != nil {
|
||||
if snapshot := live.Load(); snapshot != nil {
|
||||
return ForwardedClientIPSettings{
|
||||
TrustForwardedIP: snapshot.TrustForwardedIP,
|
||||
Headers: cloneForwardedClientIPHeaders(snapshot.Headers),
|
||||
}
|
||||
}
|
||||
}
|
||||
return ForwardedClientIPSettings{
|
||||
TrustForwardedIP: c.Security.TrustForwardedIPForAPIKeyACL,
|
||||
Headers: cloneForwardedClientIPHeaders(c.Security.ForwardedClientIPHeaders),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Config) TrustForwardedIPForAPIKeyACL() bool {
|
||||
return c.ForwardedClientIPSettings().TrustForwardedIP
|
||||
}
|
||||
|
||||
// ForwardedClientIPTrustEnabled reports whether the legacy forwarded-header
|
||||
// compatibility mode currently overrides server.trusted_proxies.
|
||||
func (c *Config) ForwardedClientIPTrustEnabled() bool {
|
||||
return c != nil && c.TrustForwardedIPForAPIKeyACL()
|
||||
}
|
||||
|
||||
func (c *Config) SetForwardedClientIPSettings(enabled bool, headers []string) {
|
||||
if c == nil {
|
||||
return false
|
||||
return
|
||||
}
|
||||
live := c.Security.trustForwardedIPForAPIKeyACLLive
|
||||
if live == nil {
|
||||
return c.Security.TrustForwardedIPForAPIKeyACL
|
||||
headers = cloneForwardedClientIPHeaders(headers)
|
||||
if c.Security.forwardedClientIPSettingsLive == nil {
|
||||
c.Security.forwardedClientIPSettingsLive = &atomic.Pointer[ForwardedClientIPSettings]{}
|
||||
}
|
||||
return live.Load()
|
||||
c.Security.forwardedClientIPSettingsLive.Store(&ForwardedClientIPSettings{
|
||||
TrustForwardedIP: enabled,
|
||||
Headers: headers,
|
||||
})
|
||||
}
|
||||
|
||||
func (c *Config) SetTrustForwardedIPForAPIKeyACL(enabled bool) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.Security.TrustForwardedIPForAPIKeyACL = enabled
|
||||
if c.Security.trustForwardedIPForAPIKeyACLLive == nil {
|
||||
c.Security.trustForwardedIPForAPIKeyACLLive = &atomic.Bool{}
|
||||
}
|
||||
c.Security.trustForwardedIPForAPIKeyACLLive.Store(enabled)
|
||||
c.SetForwardedClientIPSettings(enabled, c.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
type URLAllowlistConfig struct {
|
||||
@@ -796,6 +883,8 @@ type GatewayConfig struct {
|
||||
OpenAIHighEffortFirstOutputTimeoutSeconds int `mapstructure:"openai_high_effort_first_output_timeout_seconds"`
|
||||
// 请求体最大字节数,用于网关请求体大小限制
|
||||
MaxBodySize int64 `mapstructure:"max_body_size"`
|
||||
// TextMaxBodySize limits endpoints that cannot carry inline image/video payloads.
|
||||
TextMaxBodySize int64 `mapstructure:"text_max_body_size"`
|
||||
// 非流式上游响应体读取上限(字节),用于防止无界读取导致内存放大
|
||||
UpstreamResponseReadMaxBytes int64 `mapstructure:"upstream_response_read_max_bytes"`
|
||||
// 代理探测响应体读取上限(字节)
|
||||
@@ -1419,12 +1508,22 @@ type RateLimitConfig struct {
|
||||
|
||||
// APIKeyAuthCacheConfig API Key 认证缓存配置
|
||||
type APIKeyAuthCacheConfig struct {
|
||||
L1Size int `mapstructure:"l1_size"`
|
||||
L1TTLSeconds int `mapstructure:"l1_ttl_seconds"`
|
||||
L2TTLSeconds int `mapstructure:"l2_ttl_seconds"`
|
||||
NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"`
|
||||
JitterPercent int `mapstructure:"jitter_percent"`
|
||||
Singleflight bool `mapstructure:"singleflight"`
|
||||
L1Size int `mapstructure:"l1_size"`
|
||||
L1TTLSeconds int `mapstructure:"l1_ttl_seconds"`
|
||||
L2TTLSeconds int `mapstructure:"l2_ttl_seconds"`
|
||||
NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"`
|
||||
JitterPercent int `mapstructure:"jitter_percent"`
|
||||
Singleflight bool `mapstructure:"singleflight"`
|
||||
LookupConcurrency int `mapstructure:"lookup_concurrency"`
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
WindowSeconds int `mapstructure:"window_seconds"`
|
||||
BlockSeconds int `mapstructure:"block_seconds"`
|
||||
Capacity int `mapstructure:"capacity"`
|
||||
}
|
||||
|
||||
// SubscriptionCacheConfig 订阅认证 L1 缓存配置
|
||||
@@ -1551,17 +1650,31 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
}
|
||||
// 配置文件不存在时使用默认值
|
||||
}
|
||||
trustedProxiesEnv, trustedProxiesEnvConfigured := os.LookupEnv("SERVER_TRUSTED_PROXIES")
|
||||
forwardedClientIPHeadersEnv, forwardedClientIPHeadersEnvConfigured := os.LookupEnv("SECURITY_FORWARDED_CLIENT_IP_HEADERS")
|
||||
trustedProxiesConfigured := viper.InConfig("server.trusted_proxies") ||
|
||||
viper.IsSet("server.trusted_proxies") || trustedProxiesEnvConfigured
|
||||
|
||||
var cfg Config
|
||||
if err := viper.Unmarshal(&cfg); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal config error: %w", err)
|
||||
}
|
||||
if trustedProxiesEnvConfigured {
|
||||
cfg.Server.TrustedProxies = normalizeStringSlice(strings.Split(trustedProxiesEnv, ","))
|
||||
}
|
||||
if forwardedClientIPHeadersEnvConfigured {
|
||||
cfg.Security.ForwardedClientIPHeaders = normalizeStringSlice(strings.Split(forwardedClientIPHeadersEnv, ","))
|
||||
}
|
||||
cfg.Server.TrustedProxiesConfigured = trustedProxiesConfigured
|
||||
if cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs == 0 {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
|
||||
}
|
||||
if cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate == 0 {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
|
||||
}
|
||||
// Kept as a backstop: setEnvReachableDefaults now registers this key with its
|
||||
// effective default (true), so IsSet always reports true and this branch no
|
||||
// longer fires. It still guards the default if that registration is dropped.
|
||||
if !cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled && !viper.IsSet("gateway.openai_scheduler.sticky_escape_enabled") {
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
|
||||
}
|
||||
@@ -1611,7 +1724,12 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
|
||||
cfg.Security.ResponseHeaders.AdditionalAllowed = normalizeStringSlice(cfg.Security.ResponseHeaders.AdditionalAllowed)
|
||||
cfg.Security.ResponseHeaders.ForceRemove = normalizeStringSlice(cfg.Security.ResponseHeaders.ForceRemove)
|
||||
cfg.Security.CSP.Policy = strings.TrimSpace(cfg.Security.CSP.Policy)
|
||||
cfg.SetTrustForwardedIPForAPIKeyACL(cfg.Security.TrustForwardedIPForAPIKeyACL)
|
||||
forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(cfg.Security.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("security.forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
cfg.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders
|
||||
cfg.SetForwardedClientIPSettings(cfg.Security.TrustForwardedIPForAPIKeyACL, forwardedClientIPHeaders)
|
||||
cfg.Log.Level = strings.ToLower(strings.TrimSpace(cfg.Log.Level))
|
||||
cfg.Log.Format = strings.ToLower(strings.TrimSpace(cfg.Log.Format))
|
||||
cfg.Log.ServiceName = strings.TrimSpace(cfg.Log.ServiceName)
|
||||
@@ -1698,9 +1816,9 @@ func setDefaults() {
|
||||
viper.SetDefault("server.mode", "release")
|
||||
viper.SetDefault("server.enable_server_timing", false)
|
||||
viper.SetDefault("server.frontend_url", "")
|
||||
viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头
|
||||
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
|
||||
viper.SetDefault("server.trusted_proxies", []string{})
|
||||
viper.SetDefault("server.read_header_timeout", 10) // 10秒读取请求头
|
||||
viper.SetDefault("server.max_header_bytes", 64*1024)
|
||||
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
|
||||
viper.SetDefault("server.max_request_body_size", int64(256*1024*1024))
|
||||
// H2C 默认配置
|
||||
viper.SetDefault("server.h2c.enabled", false)
|
||||
@@ -1757,7 +1875,7 @@ func setDefaults() {
|
||||
viper.SetDefault("security.csp.enabled", true)
|
||||
viper.SetDefault("security.csp.policy", DefaultCSPPolicy)
|
||||
viper.SetDefault("security.proxy_probe.insecure_skip_verify", false)
|
||||
viper.SetDefault("security.trust_forwarded_ip_for_api_key_acl", false)
|
||||
viper.SetDefault("security.trust_forwarded_ip_for_api_key_acl", true)
|
||||
|
||||
// Security - disable direct fallback on proxy error
|
||||
viper.SetDefault("security.proxy_fallback.allow_direct_on_error", false)
|
||||
@@ -1926,6 +2044,15 @@ func setDefaults() {
|
||||
viper.SetDefault("image_storage.force_path_style", false)
|
||||
viper.SetDefault("image_storage.presign_expiry_hours", 24)
|
||||
viper.SetDefault("image_storage.max_download_bytes", 33554432)
|
||||
// Registered with empty defaults so AutomaticEnv can reach them: viper only
|
||||
// decodes keys present in AllKeys(), so a credential that is supplied purely
|
||||
// via IMAGE_STORAGE_* and never appears in config.yaml would be dropped and
|
||||
// silently disable the whole async image feature.
|
||||
viper.SetDefault("image_storage.endpoint", "")
|
||||
viper.SetDefault("image_storage.bucket", "")
|
||||
viper.SetDefault("image_storage.access_key_id", "")
|
||||
viper.SetDefault("image_storage.secret_access_key", "")
|
||||
viper.SetDefault("image_storage.public_base_url", "")
|
||||
|
||||
// Ops (vNext)
|
||||
viper.SetDefault("ops.enabled", true)
|
||||
@@ -1983,6 +2110,12 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.negative_ttl_seconds", 30)
|
||||
viper.SetDefault("api_key_auth_cache.jitter_percent", 10)
|
||||
viper.SetDefault("api_key_auth_cache.singleflight", true)
|
||||
viper.SetDefault("api_key_auth_cache.lookup_concurrency", 64)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.enabled", true)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.threshold", 120)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2111,6 +2244,7 @@ func setDefaults() {
|
||||
viper.SetDefault("gateway.antigravity_fallback_cooldown_minutes", 1)
|
||||
viper.SetDefault("gateway.antigravity_extra_retries", 10)
|
||||
viper.SetDefault("gateway.max_body_size", int64(256*1024*1024))
|
||||
viper.SetDefault("gateway.text_max_body_size", int64(32*1024*1024))
|
||||
viper.SetDefault("gateway.upstream_response_read_max_bytes", DefaultUpstreamResponseReadMaxBytes)
|
||||
viper.SetDefault("gateway.proxy_probe_response_read_max_bytes", int64(1024*1024))
|
||||
viper.SetDefault("gateway.gemini_debug_response_headers", false)
|
||||
@@ -2205,9 +2339,140 @@ func setDefaults() {
|
||||
viper.SetDefault("subscription_maintenance.worker_count", 2)
|
||||
viper.SetDefault("subscription_maintenance.queue_size", 1024)
|
||||
|
||||
setEnvReachableDefaults()
|
||||
}
|
||||
|
||||
// setEnvReachableDefaults registers zero-valued defaults for keys that are
|
||||
// documented in deploy/config.example.yaml but had no default of their own.
|
||||
//
|
||||
// viper.Unmarshal only decodes the keys returned by AllKeys(), which unions
|
||||
// SetDefault keys, config-file keys and explicitly bound BindEnv keys.
|
||||
// AutomaticEnv can override a key already in that union, but it never adds one,
|
||||
// and the viper_bind_struct escape hatch is compiled out (we build with
|
||||
// -tags embed). So a key that lives only in the example file was unreachable by
|
||||
// environment variable: the value was read from the process environment and
|
||||
// then silently dropped. Deployments driven purely by env — which is what
|
||||
// deploy/docker-compose.yml does — got the zero value with no warning.
|
||||
//
|
||||
// The values below are deliberately zero rather than the documented example
|
||||
// values: an absent key already unmarshalled to the zero value, so registering
|
||||
// zero keeps behavior identical while making the key addressable from the
|
||||
// environment. Any subsystem that wants a richer default still applies it after
|
||||
// unmarshal, exactly as before.
|
||||
func setEnvReachableDefaults() {
|
||||
viper.SetDefault("gateway.forced_codex_instructions_template_file", "")
|
||||
viper.SetDefault("gateway.session_idle_timeout_minutes", 0)
|
||||
viper.SetDefault("gateway.user_message_queue.mode", "")
|
||||
viper.SetDefault("update.proxy_url", "")
|
||||
|
||||
// sticky_escape_enabled is the one exception to the zero-value rule: its
|
||||
// effective default is true, applied post-unmarshal via a viper.IsSet guard.
|
||||
// Registering false would make IsSet always report true and permanently
|
||||
// disable sticky escape, so register the effective default instead. An
|
||||
// explicit false in config or env still wins.
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_enabled", true)
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_error_rate", 0.0)
|
||||
viper.SetDefault("gateway.openai_scheduler.sticky_escape_ttft_ms", 0)
|
||||
|
||||
// server.trusted_proxies and security.forwarded_client_ip_headers are the
|
||||
// other exception: load() distinguishes explicit configuration from absence
|
||||
// (issue #4600), and viper.IsSet also reports registered defaults, so a
|
||||
// SetDefault would make trusted proxies look permanently configured. Both
|
||||
// environment variables are parsed by hand in load() via os.LookupEnv and
|
||||
// were never silently dropped; binding them here records that reachability
|
||||
// where AllKeys() — and the env-reachability guard — can see it, without
|
||||
// affecting IsSet while the variables are absent. BindEnv only errors when
|
||||
// called without arguments.
|
||||
_ = viper.BindEnv("server.trusted_proxies", "SERVER_TRUSTED_PROXIES")
|
||||
_ = viper.BindEnv("security.forwarded_client_ip_headers", "SECURITY_FORWARDED_CLIENT_IP_HEADERS")
|
||||
|
||||
// Third-party login providers. These carry client secrets and are exactly
|
||||
// the settings an operator expects to inject via the environment, but every
|
||||
// key here was previously unreachable that way.
|
||||
for _, provider := range []string{"github_oauth", "google_oauth"} {
|
||||
viper.SetDefault(provider+".enabled", false)
|
||||
viper.SetDefault(provider+".client_id", "")
|
||||
viper.SetDefault(provider+".client_secret", "")
|
||||
viper.SetDefault(provider+".authorize_url", "")
|
||||
viper.SetDefault(provider+".token_url", "")
|
||||
viper.SetDefault(provider+".userinfo_url", "")
|
||||
viper.SetDefault(provider+".emails_url", "")
|
||||
viper.SetDefault(provider+".scopes", "")
|
||||
viper.SetDefault(provider+".redirect_url", "")
|
||||
viper.SetDefault(provider+".frontend_redirect_url", "")
|
||||
}
|
||||
|
||||
viper.SetDefault("dingtalk_connect.client_id", "")
|
||||
viper.SetDefault("dingtalk_connect.client_secret", "")
|
||||
viper.SetDefault("dingtalk_connect.internal_corp_id", "")
|
||||
viper.SetDefault("dingtalk_connect.redirect_url", "")
|
||||
viper.SetDefault("dingtalk_connect.bypass_registration", false)
|
||||
viper.SetDefault("dingtalk_connect.username_attribute_key", "")
|
||||
viper.SetDefault("dingtalk_connect.enable_attribute_matching", false)
|
||||
viper.SetDefault("dingtalk_connect.enable_attribute_sync", false)
|
||||
viper.SetDefault("dingtalk_connect.attribute_sync_fields", []string{})
|
||||
viper.SetDefault("dingtalk_connect.attribute_sync_overwrite_policy", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_display_name_attr_name", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_dept", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_dept_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_dept_attr_name", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email", false)
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email_attr_key", "")
|
||||
viper.SetDefault("dingtalk_connect.sync_corp_email_attr_name", "")
|
||||
}
|
||||
|
||||
func (c *Config) Validate() error {
|
||||
forwardedClientIPHeaders, err := NormalizeForwardedClientIPHeaders(c.Security.ForwardedClientIPHeaders)
|
||||
if err != nil {
|
||||
return fmt.Errorf("security.forwarded_client_ip_headers: %w", err)
|
||||
}
|
||||
c.Security.ForwardedClientIPHeaders = forwardedClientIPHeaders
|
||||
c.SetForwardedClientIPSettings(c.Security.TrustForwardedIPForAPIKeyACL, forwardedClientIPHeaders)
|
||||
if c.Server.ReadHeaderTimeout < 1 || c.Server.ReadHeaderTimeout > 60 {
|
||||
return fmt.Errorf("server.read_header_timeout must be between 1 and 60 seconds")
|
||||
}
|
||||
if c.Server.MaxHeaderBytes < 8*1024 || c.Server.MaxHeaderBytes > 1024*1024 {
|
||||
return fmt.Errorf("server.max_header_bytes must be between 8192 and 1048576 bytes")
|
||||
}
|
||||
if c.Server.IdleTimeout <= 0 {
|
||||
return fmt.Errorf("server.idle_timeout must be positive")
|
||||
}
|
||||
if c.Server.MaxRequestBodySize < 0 {
|
||||
return fmt.Errorf("server.max_request_body_size must be non-negative")
|
||||
}
|
||||
if c.Server.H2C.Enabled {
|
||||
if c.Server.H2C.MaxConcurrentStreams == 0 {
|
||||
return fmt.Errorf("server.h2c.max_concurrent_streams must be positive")
|
||||
}
|
||||
if c.Server.H2C.IdleTimeout <= 0 {
|
||||
return fmt.Errorf("server.h2c.idle_timeout must be positive")
|
||||
}
|
||||
if c.Server.H2C.MaxReadFrameSize < 16*1024 || c.Server.H2C.MaxReadFrameSize > 16*1024*1024-1 {
|
||||
return fmt.Errorf("server.h2c.max_read_frame_size must be between 16384 and 16777215 bytes")
|
||||
}
|
||||
if c.Server.H2C.MaxUploadBufferPerConnection < 65535 {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_connection must be at least 65535 bytes")
|
||||
}
|
||||
if c.Server.H2C.MaxUploadBufferPerStream <= 0 {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.WindowSeconds < 1 || c.APIKeyAuth.InvalidAbuse.WindowSeconds > 3600 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.window_seconds must be between 1 and 3600")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.BlockSeconds < 1 || c.APIKeyAuth.InvalidAbuse.BlockSeconds > 3600 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.block_seconds must be between 1 and 3600")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Capacity < 256 || c.APIKeyAuth.InvalidAbuse.Capacity > 1_000_000 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.capacity must be between 256 and 1000000")
|
||||
}
|
||||
}
|
||||
jwtSecret := strings.TrimSpace(c.JWT.Secret)
|
||||
if jwtSecret == "" {
|
||||
return fmt.Errorf("jwt.secret is required")
|
||||
@@ -2733,6 +2998,9 @@ func (c *Config) Validate() error {
|
||||
if c.Gateway.MaxBodySize <= 0 {
|
||||
return fmt.Errorf("gateway.max_body_size must be positive")
|
||||
}
|
||||
if c.Gateway.TextMaxBodySize <= 0 || c.Gateway.TextMaxBodySize > c.Gateway.MaxBodySize {
|
||||
return fmt.Errorf("gateway.text_max_body_size must be positive and no greater than gateway.max_body_size")
|
||||
}
|
||||
if c.Gateway.UpstreamResponseReadMaxBytes <= 0 {
|
||||
return fmt.Errorf("gateway.upstream_response_read_max_bytes must be positive")
|
||||
}
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -35,6 +37,213 @@ func TestLoadServerTimingConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadHTTPIngressSafetyDefaults(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 10, cfg.Server.ReadHeaderTimeout)
|
||||
require.Equal(t, 64*1024, cfg.Server.MaxHeaderBytes)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.False(t, cfg.Server.TrustedProxiesConfigured)
|
||||
require.True(t, cfg.TrustForwardedIPForAPIKeyACL())
|
||||
require.Equal(t, int64(32*1024*1024), cfg.Gateway.TextMaxBodySize)
|
||||
require.True(t, cfg.APIKeyAuth.InvalidAbuse.Enabled)
|
||||
require.Equal(t, 120, cfg.APIKeyAuth.InvalidAbuse.Threshold)
|
||||
require.Equal(t, 16384, cfg.APIKeyAuth.InvalidAbuse.Capacity)
|
||||
}
|
||||
|
||||
func TestNormalizeForwardedClientIPHeaders(t *testing.T) {
|
||||
headers, err := NormalizeForwardedClientIPHeaders([]string{
|
||||
" x-cdn-client-ip ",
|
||||
"X-CDN-CLIENT-IP",
|
||||
"true-client-ip",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"X-Cdn-Client-Ip", "True-Client-Ip"}, headers)
|
||||
|
||||
_, err = NormalizeForwardedClientIPHeaders([]string{"X Invalid"})
|
||||
require.ErrorContains(t, err, "invalid HTTP header field name")
|
||||
}
|
||||
|
||||
func TestNormalizeForwardedClientIPHeadersLimit(t *testing.T) {
|
||||
headers := make([]string, 0, MaxForwardedClientIPHeaders+1)
|
||||
for i := 0; i <= MaxForwardedClientIPHeaders; i++ {
|
||||
headers = append(headers, fmt.Sprintf("X-CDN-IP-%d", i))
|
||||
}
|
||||
|
||||
_, err := NormalizeForwardedClientIPHeaders(headers)
|
||||
require.ErrorContains(t, err, "at most 16 unique names")
|
||||
}
|
||||
|
||||
func TestLoadForwardedClientIPHeadersNormalizesAndSnapshots(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{" x-cdn-ip ", "X-CDN-IP", "true-client-ip"})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
snapshot := cfg.ForwardedClientIPSettings()
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, snapshot.Headers)
|
||||
|
||||
snapshot.Headers[0] = "X-Mutated"
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestForwardedClientIPSettingsConcurrentPublication(t *testing.T) {
|
||||
cfg := &Config{}
|
||||
cfg.SetForwardedClientIPSettings(true, []string{"X-Public-A"})
|
||||
|
||||
const iterations = 2000
|
||||
start := make(chan struct{})
|
||||
errCh := make(chan error, 8)
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, settings := range []ForwardedClientIPSettings{
|
||||
{TrustForwardedIP: true, Headers: []string{"X-Public-A"}},
|
||||
{TrustForwardedIP: false, Headers: []string{"X-Public-B"}},
|
||||
} {
|
||||
settings := settings
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for i := 0; i < iterations; i++ {
|
||||
cfg.SetForwardedClientIPSettings(settings.TrustForwardedIP, settings.Headers)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
for i := 0; i < cap(errCh); i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for j := 0; j < iterations; j++ {
|
||||
snapshot := cfg.ForwardedClientIPSettings()
|
||||
validA := snapshot.TrustForwardedIP && len(snapshot.Headers) == 1 && snapshot.Headers[0] == "X-Public-A"
|
||||
validB := !snapshot.TrustForwardedIP && len(snapshot.Headers) == 1 && snapshot.Headers[0] == "X-Public-B"
|
||||
if !validA && !validB {
|
||||
errCh <- fmt.Errorf("observed inconsistent forwarded IP settings: %+v", snapshot)
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
for err := range errCh {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadForwardedClientIPHeadersFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", " x-cdn-ip , X-CDN-IP, true-client-ip ")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"X-Cdn-Ip", "True-Client-Ip"}, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyForwardedClientIPHeadersFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{"X-Yaml-IP"})
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", "")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.ForwardedClientIPSettings().Headers)
|
||||
}
|
||||
|
||||
func TestLoadRejectsInvalidForwardedClientIPHeaderFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SECURITY_FORWARDED_CLIENT_IP_HEADERS", "X-Valid-IP, X Invalid")
|
||||
|
||||
_, err := Load()
|
||||
require.ErrorContains(t, err, "security.forwarded_client_ip_headers")
|
||||
}
|
||||
|
||||
func TestLoadRejectsInvalidForwardedClientIPHeader(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("security.forwarded_client_ip_headers", []string{"X Invalid"})
|
||||
|
||||
_, err := Load()
|
||||
require.ErrorContains(t, err, "security.forwarded_client_ip_headers")
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyTrustedProxiesEnablesConfiguredMode(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("server.trusted_proxies", []string{})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadExplicitTrustedProxiesEnablesConfiguredMode(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
viper.Set("server.trusted_proxies", []string{"127.0.0.1/32"})
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"127.0.0.1/32"}, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadTrustedProxiesFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SERVER_TRUSTED_PROXIES", "127.0.0.1/32, ::1/128")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []string{"127.0.0.1/32", "::1/128"}, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadExplicitEmptyTrustedProxiesFromEnvironment(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("SERVER_TRUSTED_PROXIES", "")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, cfg.Server.TrustedProxies)
|
||||
require.True(t, cfg.Server.TrustedProxiesConfigured)
|
||||
}
|
||||
|
||||
func TestLoadTrustedProxiesPresenceFromYAML(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
yaml string
|
||||
want []string
|
||||
configured bool
|
||||
}{
|
||||
{name: "absent", yaml: "server:\n mode: debug\n", configured: false},
|
||||
{name: "explicit empty", yaml: "server:\n trusted_proxies: []\n", want: []string{}, configured: true},
|
||||
{
|
||||
name: "populated",
|
||||
yaml: "server:\n trusted_proxies:\n - 127.0.0.1/32\n - ::1/128\n",
|
||||
want: []string{"127.0.0.1/32", "::1/128"},
|
||||
configured: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
configDir := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(configDir, "config.yaml"), []byte(test.yaml), 0o600))
|
||||
t.Setenv("DATA_DIR", configDir)
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, test.want, cfg.Server.TrustedProxies)
|
||||
require.Equal(t, test.configured, cfg.Server.TrustedProxiesConfigured)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
|
||||
viper.Reset()
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
@@ -1185,6 +1394,51 @@ func TestValidateConfigErrors(t *testing.T) {
|
||||
mutate func(*Config)
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "server read header timeout",
|
||||
mutate: func(c *Config) { c.Server.ReadHeaderTimeout = 0 },
|
||||
wantErr: "server.read_header_timeout",
|
||||
},
|
||||
{
|
||||
name: "server max header bytes too small",
|
||||
mutate: func(c *Config) { c.Server.MaxHeaderBytes = 4096 },
|
||||
wantErr: "server.max_header_bytes",
|
||||
},
|
||||
{
|
||||
name: "server max request body size",
|
||||
mutate: func(c *Config) { c.Server.MaxRequestBodySize = -1 },
|
||||
wantErr: "server.max_request_body_size",
|
||||
},
|
||||
{
|
||||
name: "h2c zero concurrent streams",
|
||||
mutate: func(c *Config) {
|
||||
c.Server.H2C.Enabled = true
|
||||
c.Server.H2C.MaxConcurrentStreams = 0
|
||||
},
|
||||
wantErr: "server.h2c.max_concurrent_streams",
|
||||
},
|
||||
{
|
||||
name: "h2c oversized read frame",
|
||||
mutate: func(c *Config) {
|
||||
c.Server.H2C.Enabled = true
|
||||
c.Server.H2C.MaxReadFrameSize = 16 * 1024 * 1024
|
||||
},
|
||||
wantErr: "server.h2c.max_read_frame_size",
|
||||
},
|
||||
{
|
||||
name: "invalid auth abuse threshold too small",
|
||||
mutate: func(c *Config) {
|
||||
c.APIKeyAuth.InvalidAbuse.Threshold = 9
|
||||
},
|
||||
wantErr: "api_key_auth_cache.invalid_abuse.threshold",
|
||||
},
|
||||
{
|
||||
name: "invalid auth abuse capacity too small",
|
||||
mutate: func(c *Config) {
|
||||
c.APIKeyAuth.InvalidAbuse.Capacity = 255
|
||||
},
|
||||
wantErr: "api_key_auth_cache.invalid_abuse.capacity",
|
||||
},
|
||||
{
|
||||
name: "jwt secret required",
|
||||
mutate: func(c *Config) { c.JWT.Secret = "" },
|
||||
@@ -1386,6 +1640,11 @@ func TestValidateConfigErrors(t *testing.T) {
|
||||
mutate: func(c *Config) { c.Gateway.MaxBodySize = 0 },
|
||||
wantErr: "gateway.max_body_size",
|
||||
},
|
||||
{
|
||||
name: "gateway text body exceeds media body",
|
||||
mutate: func(c *Config) { c.Gateway.TextMaxBodySize = c.Gateway.MaxBodySize + 1 },
|
||||
wantErr: "gateway.text_max_body_size",
|
||||
},
|
||||
{
|
||||
name: "gateway response header timeout",
|
||||
mutate: func(c *Config) { c.Gateway.ResponseHeaderTimeout = -1 },
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
//go:build unit
|
||||
|
||||
package config
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"sort"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
// collectMapstructureKeys walks a config struct and returns every dotted key
|
||||
// viper would need in order to populate it.
|
||||
func collectMapstructureKeys(t reflect.Type, prefix string, out map[string]string) {
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
field := t.Field(i)
|
||||
if field.PkgPath != "" {
|
||||
continue // unexported
|
||||
}
|
||||
tag := field.Tag.Get("mapstructure")
|
||||
name, _, _ := strings.Cut(tag, ",")
|
||||
if name == "-" {
|
||||
continue
|
||||
}
|
||||
if name == "" {
|
||||
name = strings.ToLower(field.Name)
|
||||
}
|
||||
key := name
|
||||
if prefix != "" {
|
||||
key = prefix + "." + name
|
||||
}
|
||||
|
||||
ft := field.Type
|
||||
for ft.Kind() == reflect.Ptr {
|
||||
ft = ft.Elem()
|
||||
}
|
||||
if ft.Kind() == reflect.Struct {
|
||||
collectMapstructureKeys(ft, key, out)
|
||||
continue
|
||||
}
|
||||
if ft.Kind() == reflect.Map {
|
||||
// A map cannot be expressed in a single environment variable, so it
|
||||
// is out of scope here — such settings need a config file either way.
|
||||
continue
|
||||
}
|
||||
out[strings.ToLower(key)] = ft.String()
|
||||
}
|
||||
}
|
||||
|
||||
// TestConfigKeysAreEnvReachable is the systemic guard behind the image_storage
|
||||
// bug: viper.Unmarshal only decodes keys returned by AllKeys(), which unions
|
||||
// SetDefault keys, config-file keys and explicit BindEnv keys. AutomaticEnv can
|
||||
// override a key already in that union but never introduces one, and the
|
||||
// viper_bind_struct escape hatch is compiled out (we build with -tags embed).
|
||||
//
|
||||
// So a Config field with no registered default is unreachable by environment
|
||||
// variable whenever the deployment has no config.yaml containing it — the
|
||||
// operator sets the variable, the loader discards it, and the feature behaves
|
||||
// as if it were never configured. That is exactly how image_storage credentials
|
||||
// were lost, silently disabling async image tasks for env-driven deployments.
|
||||
//
|
||||
// When this fails, register a zero-valued default in setEnvReachableDefaults
|
||||
// for each reported key.
|
||||
func TestConfigKeysAreEnvReachable(t *testing.T) {
|
||||
bound := map[string]string{}
|
||||
collectMapstructureKeys(reflect.TypeOf(Config{}), "", bound)
|
||||
|
||||
viper.Reset()
|
||||
t.Cleanup(viper.Reset)
|
||||
setDefaults()
|
||||
registered := map[string]struct{}{}
|
||||
for _, key := range viper.AllKeys() {
|
||||
registered[key] = struct{}{}
|
||||
}
|
||||
|
||||
var unreachable []string
|
||||
for key, kind := range bound {
|
||||
if _, ok := registered[key]; !ok {
|
||||
unreachable = append(unreachable, key+" ("+kind+")")
|
||||
}
|
||||
}
|
||||
sort.Strings(unreachable)
|
||||
|
||||
if len(unreachable) > 0 {
|
||||
t.Fatalf("%d config keys have no default registered, so their environment variables are silently ignored:\n %s",
|
||||
len(unreachable), strings.Join(unreachable, "\n "))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//go:build unit
|
||||
|
||||
package config
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestLoadImageStorageFromEnv guards against a viper trap that silently disabled
|
||||
// asynchronous image tasks for every environment-variable-only deployment.
|
||||
//
|
||||
// viper only decodes keys returned by AllKeys(), which unions SetDefault keys,
|
||||
// config-file keys and explicit BindEnv keys. AutomaticEnv can override a key
|
||||
// that is already in that list, but it never introduces a new one. Credentials
|
||||
// such as image_storage.bucket therefore need an (empty) default registered, or
|
||||
// IMAGE_STORAGE_BUCKET is dropped on the floor and Active() stays false while
|
||||
// image_storage.enabled reads true — the endpoints 404 with no useful signal.
|
||||
func TestLoadImageStorageFromEnv(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
t.Setenv("IMAGE_STORAGE_ENABLED", "true")
|
||||
t.Setenv("IMAGE_STORAGE_ENDPOINT", "https://acct.r2.cloudflarestorage.com")
|
||||
t.Setenv("IMAGE_STORAGE_BUCKET", "my-images")
|
||||
t.Setenv("IMAGE_STORAGE_ACCESS_KEY_ID", "ak")
|
||||
t.Setenv("IMAGE_STORAGE_SECRET_ACCESS_KEY", "sk")
|
||||
t.Setenv("IMAGE_STORAGE_PUBLIC_BASE_URL", "https://cdn.example.com")
|
||||
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
|
||||
require.True(t, cfg.ImageStorage.Enabled)
|
||||
require.Equal(t, "https://acct.r2.cloudflarestorage.com", cfg.ImageStorage.Endpoint)
|
||||
require.Equal(t, "my-images", cfg.ImageStorage.Bucket)
|
||||
require.Equal(t, "ak", cfg.ImageStorage.AccessKeyID)
|
||||
require.Equal(t, "sk", cfg.ImageStorage.SecretAccessKey)
|
||||
require.Equal(t, "https://cdn.example.com", cfg.ImageStorage.PublicBaseURL)
|
||||
|
||||
require.True(t, cfg.ImageStorage.IsConfigured())
|
||||
require.True(t, cfg.ImageStorage.Active(), "async image tasks must be active when every credential is supplied via env")
|
||||
}
|
||||
@@ -47,6 +47,100 @@ func TestNormalizeCodexImportEntryAcceptsAgentIdentityAuthJSON(t *testing.T) {
|
||||
require.NotEmpty(t, item.WarningTexts)
|
||||
}
|
||||
|
||||
func TestBuildCodexAgentIdentityKeysUseChatGPTAccountOnly(t *testing.T) {
|
||||
keys := buildCodexAgentIdentityKeys("team-a")
|
||||
require.Equal(t, []string{"account:team-a"}, keys)
|
||||
}
|
||||
|
||||
func TestCodexAgentIdentityIndexSeparatesTeamsForSameUser(t *testing.T) {
|
||||
existing := service.Account{
|
||||
ID: 1,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": service.OpenAIAuthModeAgentIdentity,
|
||||
"chatgpt_account_id": "team-a",
|
||||
"chatgpt_user_id": "same-user",
|
||||
"agent_runtime_id": "runtime-a",
|
||||
},
|
||||
}
|
||||
index := buildCodexAccountIndex([]service.Account{existing})
|
||||
|
||||
teamBKeys := buildCodexAgentIdentityKeys("team-b")
|
||||
matched, _ := index.Find(teamBKeys, "same-user")
|
||||
require.Nil(t, matched)
|
||||
|
||||
teamAKeys := buildCodexAgentIdentityKeys("team-a")
|
||||
matched, matchedKey := index.Find(teamAKeys, "same-user")
|
||||
require.NotNil(t, matched)
|
||||
require.Equal(t, int64(1), matched.ID)
|
||||
require.Equal(t, "account:team-a", matchedKey)
|
||||
}
|
||||
|
||||
func TestImportCodexSessionsKeepsAgentIdentityTeamsSeparate(t *testing.T) {
|
||||
first := buildAgentIdentityImportValue(t, "runtime-a", "team-a", "same-user", "task-a")
|
||||
second := buildAgentIdentityImportValue(t, "runtime-b", "team-b", "same-user", "task-b")
|
||||
svc := newCodexImportMemoryAdminService(nil)
|
||||
handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
result, err := handler.importCodexSessions(context.Background(), CodexSessionImportRequest{
|
||||
SkipDefaultGroupBind: boolPtr(true),
|
||||
}, []codexImportEntry{{Index: 1, Value: first}, {Index: 2, Value: second}})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, result.Created)
|
||||
require.Zero(t, result.Updated)
|
||||
require.Zero(t, result.Skipped)
|
||||
require.Len(t, svc.createdAccounts, 2)
|
||||
}
|
||||
|
||||
func TestImportCodexSessionsMergesAgentIdentityRuntimesForSameTeam(t *testing.T) {
|
||||
first := buildAgentIdentityImportValue(t, "runtime-a", "team-a", "same-user", "task-a")
|
||||
second := buildAgentIdentityImportValue(t, "runtime-b", "team-a", "same-user", "task-b")
|
||||
firstIdentity, ok := first["agent_identity"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
existing := service.Account{
|
||||
ID: 41,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"auth_mode": service.OpenAIAuthModeAgentIdentity,
|
||||
"agent_runtime_id": firstIdentity["agent_runtime_id"],
|
||||
"agent_private_key": firstIdentity["agent_private_key"],
|
||||
"task_id": firstIdentity["task_id"],
|
||||
"chatgpt_account_id": firstIdentity["account_id"],
|
||||
"chatgpt_user_id": firstIdentity["chatgpt_user_id"],
|
||||
},
|
||||
}
|
||||
svc := newCodexImportMemoryAdminService([]service.Account{existing})
|
||||
handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
result, err := handler.importCodexSessions(context.Background(), CodexSessionImportRequest{
|
||||
SkipDefaultGroupBind: boolPtr(true),
|
||||
}, []codexImportEntry{{Index: 1, Value: second}})
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, result.Created)
|
||||
require.Equal(t, 1, result.Updated)
|
||||
require.Len(t, svc.updatedAccounts, 1)
|
||||
require.Equal(t, "runtime-b", svc.updatedAccounts[0].input.Credentials["agent_runtime_id"])
|
||||
require.Equal(t, "task-b", svc.updatedAccounts[0].input.Credentials["task_id"])
|
||||
}
|
||||
|
||||
func buildAgentIdentityImportValue(t *testing.T, runtimeID, accountID, userID, taskID string) map[string]any {
|
||||
t.Helper()
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
der, err := x509.MarshalPKCS8PrivateKey(privateKey)
|
||||
require.NoError(t, err)
|
||||
return map[string]any{
|
||||
"auth_mode": "agentIdentity",
|
||||
"agent_identity": map[string]any{
|
||||
"agent_runtime_id": runtimeID,
|
||||
"agent_private_key": base64.StdEncoding.EncodeToString(der),
|
||||
"task_id": taskID,
|
||||
"account_id": accountID,
|
||||
"chatgpt_user_id": userID,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func TestImportCodexSessionsCreatesAgentIdentityWithoutOAuthExpiry(t *testing.T) {
|
||||
_, privateKey, err := ed25519.GenerateKey(rand.Reader)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -532,7 +532,7 @@ func normalizeCodexImportEntry(entry codexImportEntry) (*codexImportAccount, err
|
||||
if item.AgentTaskID == "" {
|
||||
item.WarningTexts = append(item.WarningTexts, "未包含 task_id,首次请求会使用现有 runtime 注册新 task")
|
||||
}
|
||||
item.IdentityKeys = buildCodexAgentIdentityKeys(item.AccountID, item.UserID, item.Email, item.AgentRuntimeID)
|
||||
item.IdentityKeys = buildCodexAgentIdentityKeys(item.AccountID)
|
||||
item.Name = buildCodexImportAccountName(item, entry.Index)
|
||||
return item, nil
|
||||
}
|
||||
@@ -894,12 +894,17 @@ func buildCodexImportIdentityKeys(accountID, userID, email, accessToken, refresh
|
||||
return buildCodexStoredIdentityKeys(accountID, userID, email, accessToken)
|
||||
}
|
||||
|
||||
func buildCodexAgentIdentityKeys(accountID, userID, email, runtimeID string) []string {
|
||||
keys := buildCodexStoredIdentityKeys(accountID, userID, email, "")
|
||||
if runtimeID = strings.TrimSpace(runtimeID); runtimeID != "" {
|
||||
keys = append([]string{"agent:" + runtimeID}, keys...)
|
||||
func buildCodexAgentIdentityKeys(accountID string) []string {
|
||||
// Agent Identity credentials belonging to the same ChatGPT account are
|
||||
// intentionally merged, while the same user may own multiple accounts.
|
||||
// Do not use user/email/runtime as fallback keys here: user_id is shared
|
||||
// across Team workspaces and runtime_id changes when a new runtime is
|
||||
// registered for the same account.
|
||||
accountID = strings.TrimSpace(accountID)
|
||||
if accountID == "" {
|
||||
return nil
|
||||
}
|
||||
return keys
|
||||
return []string{"account:" + accountID}
|
||||
}
|
||||
|
||||
// buildCodexStoredIdentityKeys 生成存量账号索引键,保留 user/account 维度,
|
||||
|
||||
@@ -61,7 +61,7 @@ type AccountHandler struct {
|
||||
sessionLimitCache service.SessionLimitCache
|
||||
rpmCache service.RPMCache
|
||||
tokenCacheInvalidator service.TokenCacheInvalidator
|
||||
grokImportProber grokUsageProber
|
||||
grokImportProber grokImportProber
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService
|
||||
}
|
||||
|
||||
@@ -121,6 +121,7 @@ type CreateAccountRequest struct {
|
||||
GroupIDs []int64 `json:"group_ids"`
|
||||
ExpiresAt *int64 `json:"expires_at"`
|
||||
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
|
||||
ProbeEnabled *bool `json:"upstream_billing_probe_enabled"`
|
||||
ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险
|
||||
}
|
||||
|
||||
@@ -159,6 +160,7 @@ type BulkUpdateAccountsRequest struct {
|
||||
GroupIDs *[]int64 `json:"group_ids"`
|
||||
Credentials map[string]any `json:"credentials"`
|
||||
Extra map[string]any `json:"extra"`
|
||||
ProbeEnabled *bool `json:"upstream_billing_probe_enabled"`
|
||||
ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险
|
||||
}
|
||||
|
||||
@@ -828,6 +830,7 @@ func (h *AccountHandler) Create(c *gin.Context) {
|
||||
GroupIDs: req.GroupIDs,
|
||||
ExpiresAt: req.ExpiresAt,
|
||||
AutoPauseOnExpired: req.AutoPauseOnExpired,
|
||||
ProbeEnabled: req.ProbeEnabled,
|
||||
SkipMixedChannelCheck: skipCheck,
|
||||
})
|
||||
if execErr != nil {
|
||||
@@ -1907,7 +1910,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
req.Schedulable != nil ||
|
||||
req.GroupIDs != nil ||
|
||||
len(req.Credentials) > 0 ||
|
||||
len(req.Extra) > 0
|
||||
len(req.Extra) > 0 ||
|
||||
req.ProbeEnabled != nil
|
||||
|
||||
if !hasUpdates {
|
||||
response.BadRequest(c, "No updates provided")
|
||||
@@ -1928,6 +1932,7 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
GroupIDs: req.GroupIDs,
|
||||
Credentials: req.Credentials,
|
||||
Extra: req.Extra,
|
||||
ProbeEnabled: req.ProbeEnabled,
|
||||
SkipMixedChannelCheck: skipCheck,
|
||||
})
|
||||
if err != nil {
|
||||
|
||||
@@ -222,3 +222,22 @@ func TestBulkUpdateAcceptsFilterTargetRequest(t *testing.T) {
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, float64(0), resp["code"])
|
||||
}
|
||||
|
||||
func TestBulkUpdateAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) {
|
||||
adminSvc := newStubAdminService()
|
||||
router := setupAccountMixedChannelRouter(adminSvc)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"account_ids": []int64{1, 2},
|
||||
"upstream_billing_probe_enabled": false,
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.NotNil(t, adminSvc.lastBulkUpdateAccountInput)
|
||||
require.NotNil(t, adminSvc.lastBulkUpdateAccountInput.ProbeEnabled)
|
||||
require.False(t, *adminSvc.lastBulkUpdateAccountInput.ProbeEnabled)
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
|
||||
router := gin.New()
|
||||
adminSvc := newStubAdminService()
|
||||
|
||||
userHandler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
userHandler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
groupHandler := NewGroupHandler(adminSvc, nil, nil)
|
||||
proxyHandler := NewProxyHandler(adminSvc)
|
||||
redeemHandler := NewRedeemHandler(adminSvc, nil)
|
||||
|
||||
@@ -33,6 +33,7 @@ type stubAdminService struct {
|
||||
createSparkShadowErr error
|
||||
updateAccountErr error
|
||||
bulkUpdateAccountErr error
|
||||
lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput
|
||||
getAccountResult *service.Account
|
||||
updateAccountCalls int
|
||||
updateAccountExtraCalls int
|
||||
@@ -478,6 +479,7 @@ func (s *stubAdminService) SetAccountSchedulable(ctx context.Context, id int64,
|
||||
}
|
||||
|
||||
func (s *stubAdminService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) {
|
||||
s.lastBulkUpdateAccountInput = input
|
||||
if s.bulkUpdateAccountErr != nil {
|
||||
return nil, s.bulkUpdateAccountErr
|
||||
}
|
||||
|
||||
@@ -10,12 +10,14 @@ import (
|
||||
type BackupHandler struct {
|
||||
backupService *service.BackupService
|
||||
userService *service.UserService
|
||||
imageStorage *service.ImageStorageSettingService
|
||||
}
|
||||
|
||||
func NewBackupHandler(backupService *service.BackupService, userService *service.UserService) *BackupHandler {
|
||||
func NewBackupHandler(backupService *service.BackupService, userService *service.UserService, imageStorage *service.ImageStorageSettingService) *BackupHandler {
|
||||
return &BackupHandler{
|
||||
backupService: backupService,
|
||||
userService: userService,
|
||||
imageStorage: imageStorage,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,3 +205,48 @@ func (h *BackupHandler) RestoreBackup(c *gin.Context) {
|
||||
}
|
||||
response.Accepted(c, record)
|
||||
}
|
||||
|
||||
// ─── 异步生图对象存储配置 ───
|
||||
//
|
||||
// 与备份共用一套 S3 客户端构造,因此放在同一个页面下:勾选"复用备份 S3"即可直接
|
||||
// 借用备份已配置的端点与密钥,只用不同的前缀区分对象(备份走 backups/,图片走 images/)。
|
||||
|
||||
func (h *BackupHandler) GetImageStorageConfig(c *gin.Context) {
|
||||
ctx := c.Request.Context()
|
||||
cfg, err := h.imageStorage.Get(ctx)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{
|
||||
"config": cfg,
|
||||
"secret_configured": h.imageStorage.SecretConfigured(ctx),
|
||||
})
|
||||
}
|
||||
|
||||
func (h *BackupHandler) UpdateImageStorageConfig(c *gin.Context) {
|
||||
var req service.ImageStorageSettings
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
cfg, err := h.imageStorage.Update(c.Request.Context(), req)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, cfg)
|
||||
}
|
||||
|
||||
func (h *BackupHandler) TestImageStorageConnection(c *gin.Context) {
|
||||
var req service.ImageStorageSettings
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := h.imageStorage.TestConnection(c.Request.Context(), req); err != nil {
|
||||
response.Success(c, gin.H{"ok": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{"ok": true, "message": "connection successful"})
|
||||
}
|
||||
|
||||
@@ -15,12 +15,12 @@ const (
|
||||
grokImportProbeTimeout = 25 * time.Second
|
||||
)
|
||||
|
||||
type grokUsageProber interface {
|
||||
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
|
||||
type grokImportProber interface {
|
||||
QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
|
||||
}
|
||||
|
||||
type grokImportProbeTask struct {
|
||||
prober grokUsageProber
|
||||
prober grokImportProber
|
||||
accountID int64
|
||||
}
|
||||
|
||||
@@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm
|
||||
}
|
||||
}
|
||||
|
||||
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
|
||||
func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) {
|
||||
if s == nil || prober == nil || account == nil || account.ID <= 0 {
|
||||
return
|
||||
}
|
||||
@@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
|
||||
return task, true
|
||||
}
|
||||
|
||||
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
|
||||
func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
slog.Error(
|
||||
@@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64)
|
||||
// while this timeout only bounds the actual upstream probe execution.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
|
||||
defer cancel()
|
||||
result, err := prober.ProbeUsage(ctx, accountID)
|
||||
result, err := prober.QueryQuota(ctx, accountID)
|
||||
if err != nil {
|
||||
slog.Warn(
|
||||
"grok_import_active_probe_failed",
|
||||
|
||||
@@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
|
||||
func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
|
||||
_, deadlineSeen := ctx.Deadline()
|
||||
s.mu.Lock()
|
||||
s.calls[accountID]++
|
||||
@@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (
|
||||
return nil, failure
|
||||
}
|
||||
return &service.GrokQuotaProbeResult{
|
||||
Source: "active_probe",
|
||||
Source: "hybrid_probe",
|
||||
Model: "grok-4.5",
|
||||
StatusCode: 200,
|
||||
ResetSupported: false,
|
||||
|
||||
@@ -23,7 +23,7 @@ type GrokOAuthHandler struct {
|
||||
grokOAuthService *service.GrokOAuthService
|
||||
adminService service.AdminService
|
||||
quotaService *service.GrokQuotaService
|
||||
importProber grokUsageProber
|
||||
importProber grokImportProber
|
||||
reconciler service.GrokOAuthReconciler
|
||||
}
|
||||
|
||||
|
||||
@@ -136,8 +136,12 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
|
||||
for i, upstreamReq := range requests {
|
||||
require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization"))
|
||||
if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" {
|
||||
require.Equal(t, "application/json, text/event-stream", upstreamReq.Header.Get("Accept"))
|
||||
require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`)
|
||||
require.Contains(t, string(bodies[i]), `"store":false`)
|
||||
require.Contains(t, string(bodies[i]), `"input":"hi"`)
|
||||
require.Contains(t, string(bodies[i]), `"stream":true`)
|
||||
require.NotContains(t, string(bodies[i]), `"max_output_tokens"`)
|
||||
require.NotContains(t, string(bodies[i]), `"store"`)
|
||||
}
|
||||
}
|
||||
require.NotNil(t, repo.updates[42])
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetAuthCacheInvalidationHealth exposes durable outbox lag and subscriber health.
|
||||
func (h *OpsHandler) GetAuthCacheInvalidationHealth(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, h.opsService.GetAuthCacheInvalidationHealth(c.Request.Context()))
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var ingressRejectReasons = map[string]struct{}{
|
||||
"query_api_key_deprecated": {}, "api_key_required": {}, "invalid_api_key": {},
|
||||
"invalid_auth_rate_limited": {},
|
||||
"api_key_auth_overloaded": {},
|
||||
"api_key_disabled": {}, "ip_restricted": {}, "user_inactive": {}, "group_deleted": {},
|
||||
"group_disabled": {}, "group_not_allowed": {}, "group_unassigned": {}, "other": {},
|
||||
}
|
||||
|
||||
var ingressRejectRouteFamilies = map[string]struct{}{
|
||||
"antigravity": {}, "gemini": {}, "codex": {}, "messages": {}, "responses": {},
|
||||
"chat_completions": {}, "images": {}, "videos": {}, "embeddings": {}, "models": {}, "other": {},
|
||||
}
|
||||
|
||||
var ingressRejectProtocols = map[string]struct{}{
|
||||
"google": {}, "anthropic": {}, "openai": {}, "gateway": {}, "other": {},
|
||||
}
|
||||
|
||||
// ListIngressRejects returns bounded security aggregates, never raw credentials or request bodies.
|
||||
func (h *OpsHandler) ListIngressRejects(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
if pageSize > 200 {
|
||||
pageSize = 200
|
||||
}
|
||||
startTime, endTime, err := parseOpsTimeRange(c, "1h")
|
||||
if err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
filter := &service.OpsIngressRejectFilter{Page: page, PageSize: pageSize}
|
||||
if !startTime.IsZero() {
|
||||
filter.StartTime = &startTime
|
||||
}
|
||||
if !endTime.IsZero() {
|
||||
filter.EndTime = &endTime
|
||||
}
|
||||
if filter.RejectReason, err = parseIngressRejectEnum(c, "reason", ingressRejectReasons); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.RouteFamily, err = parseIngressRejectEnum(c, "route_family", ingressRejectRouteFamilies); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.Protocol, err = parseIngressRejectEnum(c, "protocol", ingressRejectProtocols); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if raw := strings.TrimSpace(c.Query("client_ip")); raw != "" {
|
||||
addr, parseErr := netip.ParseAddr(raw)
|
||||
if parseErr != nil {
|
||||
response.BadRequest(c, "Invalid client_ip")
|
||||
return
|
||||
}
|
||||
addr = addr.Unmap()
|
||||
if addr.Is6() {
|
||||
addr = netip.PrefixFrom(addr, 64).Masked().Addr()
|
||||
}
|
||||
filter.ClientIP = addr.String()
|
||||
}
|
||||
if filter.UserID, err = parseOptionalPositiveID(c, "user_id"); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.APIKeyID, err = parseOptionalPositiveID(c, "api_key_id"); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.opsService.ListIngressRejects(c.Request.Context(), filter)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
func (h *OpsHandler) GetIngressRejectHealth(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, h.opsService.GetIngressRejectHealth())
|
||||
}
|
||||
|
||||
func parseIngressRejectEnum(c *gin.Context, name string, allowed map[string]struct{}) (string, error) {
|
||||
value := strings.TrimSpace(c.Query(name))
|
||||
if value == "" {
|
||||
return "", nil
|
||||
}
|
||||
if _, ok := allowed[value]; !ok {
|
||||
return "", &ingressRejectQueryError{message: "Invalid " + name}
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func parseOptionalPositiveID(c *gin.Context, name string) (*int64, error) {
|
||||
raw := strings.TrimSpace(c.Query(name))
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
value, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || value <= 0 {
|
||||
return nil, &ingressRejectQueryError{message: "Invalid " + name}
|
||||
}
|
||||
return &value, nil
|
||||
}
|
||||
|
||||
type ingressRejectQueryError struct{ message string }
|
||||
|
||||
func (e *ingressRejectQueryError) Error() string { return e.message }
|
||||
@@ -0,0 +1,41 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newIngressRejectHandlerForTest() *OpsHandler {
|
||||
return NewOpsHandler(service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil))
|
||||
}
|
||||
|
||||
func TestListIngressRejectsValidatesFilters(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, path := range []string{
|
||||
"/api/v1/admin/ops/ingress-rejections?reason=not-valid",
|
||||
"/api/v1/admin/ops/ingress-rejections?client_ip=not-an-ip",
|
||||
"/api/v1/admin/ops/ingress-rejections?user_id=0",
|
||||
"/api/v1/admin/ops/ingress-rejections?api_key_id=-1",
|
||||
} {
|
||||
t.Run(path, func(t *testing.T) {
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = httptest.NewRequest(http.MethodGet, path, nil)
|
||||
newIngressRejectHandlerForTest().ListIngressRejects(context)
|
||||
require.Equal(t, http.StatusBadRequest, context.Writer.Status())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetIngressRejectHealthAvailable(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ingress-rejections/health", nil)
|
||||
newIngressRejectHandlerForTest().GetIngressRejectHealth(context)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), `"capacity":8192`)
|
||||
}
|
||||
@@ -59,6 +59,8 @@ type SettingHandler struct {
|
||||
paymentService *service.PaymentService
|
||||
userAttributeService *service.UserAttributeService
|
||||
notificationEmailService *service.NotificationEmailService
|
||||
totpService *service.TotpService
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
// NewSettingHandler 创建系统设置处理器
|
||||
@@ -80,6 +82,15 @@ func (h *SettingHandler) SetNotificationEmailService(notificationEmailService *s
|
||||
h.notificationEmailService = notificationEmailService
|
||||
}
|
||||
|
||||
// SetStepUpDeps attaches the services backing the step-up switch preconditions
|
||||
// (enable requires the acting admin to have TOTP enabled; disable is itself a
|
||||
// step-up gated operation), without changing the constructor signature used by
|
||||
// existing unit tests.
|
||||
func (h *SettingHandler) SetStepUpDeps(totpService *service.TotpService, userService *service.UserService) {
|
||||
h.totpService = totpService
|
||||
h.userService = userService
|
||||
}
|
||||
|
||||
// GetSettings 获取所有系统设置
|
||||
// GET /api/v1/admin/settings
|
||||
func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
@@ -124,6 +135,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
TotpEnabled: settings.TotpEnabled,
|
||||
TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(),
|
||||
SessionBindingEnabled: settings.SessionBindingEnabled,
|
||||
StepUpEnabled: settings.StepUpEnabled,
|
||||
AuditLogRetentionDays: settings.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: settings.LoginAgreementEnabled,
|
||||
LoginAgreementMode: settings.LoginAgreementMode,
|
||||
@@ -140,6 +152,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
TurnstileSiteKey: settings.TurnstileSiteKey,
|
||||
TurnstileSecretKeyConfigured: settings.TurnstileSecretKeyConfigured,
|
||||
APIKeyACLTrustForwardedIP: settings.APIKeyACLTrustForwardedIP,
|
||||
ForwardedClientIPHeaders: settings.ForwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: settings.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: settings.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecretConfigured: settings.LinuxDoConnectClientSecretConfigured,
|
||||
|
||||
@@ -56,6 +56,12 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.TotpEnabled != after.TotpEnabled {
|
||||
changed = append(changed, "totp_enabled")
|
||||
}
|
||||
if before.SessionBindingEnabled != after.SessionBindingEnabled {
|
||||
changed = append(changed, "session_binding_enabled")
|
||||
}
|
||||
if before.StepUpEnabled != after.StepUpEnabled {
|
||||
changed = append(changed, "step_up_enabled")
|
||||
}
|
||||
if before.LoginAgreementEnabled != after.LoginAgreementEnabled {
|
||||
changed = append(changed, "login_agreement_enabled")
|
||||
}
|
||||
@@ -101,6 +107,9 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.APIKeyACLTrustForwardedIP != after.APIKeyACLTrustForwardedIP {
|
||||
changed = append(changed, "api_key_acl_trust_forwarded_ip")
|
||||
}
|
||||
if !equalStringSlice(before.ForwardedClientIPHeaders, after.ForwardedClientIPHeaders) {
|
||||
changed = append(changed, "forwarded_client_ip_headers")
|
||||
}
|
||||
if before.LinuxDoConnectEnabled != after.LinuxDoConnectEnabled {
|
||||
changed = append(changed, "linuxdo_connect_enabled")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// step-up 开关转换的门控测试。
|
||||
// 测试环境不注入认证上下文/userService,因此一旦触发校验会以 401/403/500 中止;
|
||||
// 借此区分「触发了转换校验」与「直接放行到常规保存(200)」。
|
||||
|
||||
func newStepUpSwitchTestHandler(t *testing.T, stored map[string]string) (*SettingHandler, *settingHandlerRepoStub) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
repo := &settingHandlerRepoStub{values: stored}
|
||||
svc := service.NewSettingService(repo, &config.Config{Default: config.DefaultConfig{UserConcurrency: 5}})
|
||||
return NewSettingHandler(svc, nil, nil, nil, nil, nil, nil), repo
|
||||
}
|
||||
|
||||
func doUpdateSettings(t *testing.T, h *SettingHandler, body map[string]any, prepare func(c *gin.Context)) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rawBody, err := json.Marshal(body)
|
||||
require.NoError(t, err)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPut, "/api/v1/admin/settings", bytes.NewReader(rawBody))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
if prepare != nil {
|
||||
prepare(c)
|
||||
}
|
||||
|
||||
h.UpdateSettings(c)
|
||||
return rec
|
||||
}
|
||||
|
||||
// 开启开关(false→true):无认证上下文时拒绝,且带专用错误标记。
|
||||
func TestUpdateSettingsEnableStepUpRejectsWithoutSession(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ENABLE_REQUIRES_TOTP")
|
||||
require.NotEqual(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 开启开关:admin API key(机器凭证)一律拒绝,reason 与门控保持一致便于前端分流。
|
||||
func TestUpdateSettingsEnableStepUpRejectsAdminAPIKey(t *testing.T) {
|
||||
h, _ := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, func(c *gin.Context) {
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
|
||||
}
|
||||
|
||||
// 开启开关:有认证会话但 userService 未注入时 fail-closed(500),不得放行。
|
||||
func TestUpdateSettingsEnableStepUpFailsClosedWithoutUserService(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1})
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, rec.Code)
|
||||
require.NotEqual(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 关闭开关(true→false)本身是敏感操作:无认证上下文时被 step-up 门控以 401 拦截。
|
||||
func TestUpdateSettingsDisableStepUpRequiresStepUp(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, nil)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 关闭开关:admin API key 被 step-up 门控以 403 拦截。
|
||||
func TestUpdateSettingsDisableStepUpRejectsAdminAPIKey(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, func(c *gin.Context) {
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 无状态转换(false→false):不触发任何转换校验,常规保存成功且默认持久化为 false。
|
||||
func TestUpdateSettingsStepUpNoTransitionSkipsGate(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyStepUpEnabled])
|
||||
// 会话 IP/UA 绑定默认关闭:未显式提交时持久化 false。
|
||||
require.Equal(t, "false", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
// 保持开启(true→true):不触发转换校验,常规保存不被打断。
|
||||
func TestUpdateSettingsStepUpKeepEnabledSkipsGate(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 省略字段=保持现值:不含 step_up_enabled/session_binding_enabled 的旧客户端全量保存
|
||||
// 不得把已开启的安全开关静默重置,也不触发任何转换门控。
|
||||
func TestUpdateSettingsOmittedSecuritySwitchesKeepStoredValues(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
service.SettingKeySessionBindingEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
require.Equal(t, "true", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
// 省略字段在开关本就关闭时同样保持关闭(默认值路径)。
|
||||
func TestUpdateSettingsOmittedSecuritySwitchesKeepDisabled(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyStepUpEnabled])
|
||||
require.Equal(t, "false", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
func TestUpdateSettingsForwardedClientIPHeadersOmittedPreservesAndEmptyClears(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyForwardedClientIPHeaders: `["X-Cdn-Ip","True-Client-Ip"]`,
|
||||
})
|
||||
|
||||
preserved := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
require.Equal(t, http.StatusOK, preserved.Code)
|
||||
require.JSONEq(t, `["X-Cdn-Ip","True-Client-Ip"]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
require.Contains(t, preserved.Body.String(), `"forwarded_client_ip_headers":["X-Cdn-Ip","True-Client-Ip"]`)
|
||||
|
||||
cleared := doUpdateSettings(t, h, map[string]any{"forwarded_client_ip_headers": []string{}}, nil)
|
||||
require.Equal(t, http.StatusOK, cleared.Code)
|
||||
require.JSONEq(t, `[]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
require.Contains(t, cleared.Body.String(), `"forwarded_client_ip_headers":[]`)
|
||||
}
|
||||
|
||||
func TestUpdateSettingsMalformedForwardedClientIPHeadersRemainFailClosedWhenOmitted(t *testing.T) {
|
||||
cfg := &config.Config{Default: config.DefaultConfig{UserConcurrency: 5}}
|
||||
repo := &settingHandlerRepoStub{values: map[string]string{
|
||||
service.SettingKeyAPIKeyACLTrustForwardedIP: "true",
|
||||
service.SettingKeyForwardedClientIPHeaders: `{"not":"an array"}`,
|
||||
}}
|
||||
svc := service.NewSettingService(repo, cfg)
|
||||
require.ErrorContains(t, svc.LoadForwardedClientIPSettings(context.Background()), "load forwarded client ip headers")
|
||||
require.False(t, cfg.ForwardedClientIPSettings().TrustForwardedIP)
|
||||
h := NewSettingHandler(svc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyAPIKeyACLTrustForwardedIP])
|
||||
require.JSONEq(t, `[]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
runtimeSettings := cfg.ForwardedClientIPSettings()
|
||||
require.False(t, runtimeSettings.TrustForwardedIP)
|
||||
require.Empty(t, runtimeSettings.Headers)
|
||||
require.Contains(t, rec.Body.String(), `"api_key_acl_trust_forwarded_ip":false`)
|
||||
require.Contains(t, rec.Body.String(), `"forwarded_client_ip_headers":[]`)
|
||||
}
|
||||
|
||||
func TestUpdateSettingsRejectsInvalidForwardedClientIPHeader(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyForwardedClientIPHeaders: `["X-Existing-IP"]`,
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{
|
||||
"forwarded_client_ip_headers": []string{"X Invalid"},
|
||||
}, nil)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
require.JSONEq(t, `["X-Existing-IP"]`, repo.values[service.SettingKeyForwardedClientIPHeaders])
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -27,7 +28,8 @@ type UpdateSettingsRequest struct {
|
||||
FrontendURL string `json:"frontend_url"`
|
||||
InvitationCodeEnabled bool `json:"invitation_code_enabled"`
|
||||
TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证
|
||||
SessionBindingEnabled bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定
|
||||
SessionBindingEnabled *bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定(省略=保持现值)
|
||||
StepUpEnabled *bool `json:"step_up_enabled"` // 敏感操作 step-up 2FA(省略=保持现值)
|
||||
AuditLogRetentionDays int `json:"audit_log_retention_days"` // 审计日志保留天数
|
||||
LoginAgreementEnabled bool `json:"login_agreement_enabled"`
|
||||
LoginAgreementMode string `json:"login_agreement_mode"`
|
||||
@@ -49,7 +51,8 @@ type UpdateSettingsRequest struct {
|
||||
TurnstileSecretKey string `json:"turnstile_secret_key"`
|
||||
|
||||
// API Key IP 访问控制设置
|
||||
APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
APIKeyACLTrustForwardedIP *bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
ForwardedClientIPHeaders *[]string `json:"forwarded_client_ip_headers"`
|
||||
|
||||
// LinuxDo Connect OAuth 登录
|
||||
LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"`
|
||||
@@ -335,6 +338,41 @@ type UpdateSettingsRequest struct {
|
||||
|
||||
// UpdateSettings 更新系统设置
|
||||
// PUT /api/v1/admin/settings
|
||||
// ensureActorTotpForStepUp 校验当前操作者具备开启 step-up 门控的条件:
|
||||
// 必须是真人管理员会话(admin API key 无法完成 TOTP step-up,拒绝)且本人已启用 TOTP。
|
||||
// 校验失败时写入错误响应并返回 false。
|
||||
func (h *SettingHandler) ensureActorTotpForStepUp(c *gin.Context) bool {
|
||||
if c.GetString("auth_method") == service.AuditAuthMethodAdminAPIKey {
|
||||
response.ErrorWithDetails(c, http.StatusForbidden,
|
||||
"Admin API key cannot enable step-up verification; use an admin session with TOTP enabled",
|
||||
"STEP_UP_ADMIN_API_KEY_FORBIDDEN", nil)
|
||||
return false
|
||||
}
|
||||
subject, ok := middleware.GetAuthSubjectFromContext(c)
|
||||
if !ok || subject.UserID <= 0 {
|
||||
response.ErrorWithDetails(c, http.StatusForbidden,
|
||||
"Enabling step-up verification requires an authenticated admin session",
|
||||
"STEP_UP_ENABLE_REQUIRES_TOTP", nil)
|
||||
return false
|
||||
}
|
||||
if h.userService == nil {
|
||||
response.InternalError(c, "Step-up precondition check unavailable")
|
||||
return false
|
||||
}
|
||||
user, err := h.userService.GetByID(c.Request.Context(), subject.UserID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return false
|
||||
}
|
||||
if !user.TotpEnabled {
|
||||
response.ErrorWithDetails(c, http.StatusBadRequest,
|
||||
"Enable two-factor authentication (TOTP) for your account before turning on step-up verification",
|
||||
"STEP_UP_ENABLE_REQUIRES_TOTP", nil)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
var req UpdateSettingsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -353,6 +391,38 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 两个安全开关的请求字段为指针:省略字段=保持现值,避免旧客户端/脚本
|
||||
// 用不含新字段的全量 payload 保存设置时把安全开关静默重置。
|
||||
sessionBindingEnabled := previousSettings.SessionBindingEnabled
|
||||
if req.SessionBindingEnabled != nil {
|
||||
sessionBindingEnabled = *req.SessionBindingEnabled
|
||||
}
|
||||
stepUpEnabled := previousSettings.StepUpEnabled
|
||||
if req.StepUpEnabled != nil {
|
||||
stepUpEnabled = *req.StepUpEnabled
|
||||
}
|
||||
forwardedClientIPHeaders := append([]string(nil), previousSettings.ForwardedClientIPHeaders...)
|
||||
if req.ForwardedClientIPHeaders != nil {
|
||||
forwardedClientIPHeaders = append([]string(nil), (*req.ForwardedClientIPHeaders)...)
|
||||
}
|
||||
|
||||
// 开启敏感操作 step-up 门控属自锁风险操作:仅允许本人已启用 TOTP 的管理员会话开启,
|
||||
// 否则开启后操作者立即被挡在所有敏感操作之外。仅在 false→true 的开启瞬间校验,
|
||||
// 保持开启状态的常规设置保存不受影响。
|
||||
if stepUpEnabled && !previousSettings.StepUpEnabled {
|
||||
if !h.ensureActorTotpForStepUp(c) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// 关闭 step-up 门控本身就是敏感操作:防止拿到管理员会话的攻击者先关闸再执行导出/备份。
|
||||
// previousSettings 已证实开关处于开启状态,使用无条件门控变体,
|
||||
// 避免门控内部二次读取开关时因存储故障 fail-open(前端捕获 STEP_UP_REQUIRED 弹码重试)。
|
||||
if !stepUpEnabled && previousSettings.StepUpEnabled {
|
||||
if !middleware.EnforceStepUpAlways(c, h.totpService, h.userService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 验证参数
|
||||
if req.DefaultConcurrency < 1 {
|
||||
req.DefaultConcurrency = 1
|
||||
@@ -1181,7 +1251,8 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
FrontendURL: req.FrontendURL,
|
||||
InvitationCodeEnabled: req.InvitationCodeEnabled,
|
||||
TotpEnabled: req.TotpEnabled,
|
||||
SessionBindingEnabled: req.SessionBindingEnabled,
|
||||
SessionBindingEnabled: sessionBindingEnabled,
|
||||
StepUpEnabled: stepUpEnabled,
|
||||
AuditLogRetentionDays: req.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: req.LoginAgreementEnabled,
|
||||
LoginAgreementMode: loginAgreementMode,
|
||||
@@ -1203,6 +1274,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
}
|
||||
return previousSettings.APIKeyACLTrustForwardedIP
|
||||
}(),
|
||||
ForwardedClientIPHeaders: forwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: req.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: req.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecret: req.LinuxDoConnectClientSecret,
|
||||
@@ -1622,6 +1694,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if h.opsService != nil {
|
||||
h.opsService.SetMonitoringEnabled(settings.OpsMonitoringEnabled)
|
||||
}
|
||||
|
||||
// Update OpenAI fast policy (stored under dedicated key, only when provided).
|
||||
if req.OpenAIFastPolicySettings != nil {
|
||||
@@ -1710,6 +1785,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
TotpEnabled: updatedSettings.TotpEnabled,
|
||||
TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(),
|
||||
SessionBindingEnabled: updatedSettings.SessionBindingEnabled,
|
||||
StepUpEnabled: updatedSettings.StepUpEnabled,
|
||||
AuditLogRetentionDays: updatedSettings.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: updatedSettings.LoginAgreementEnabled,
|
||||
LoginAgreementMode: updatedSettings.LoginAgreementMode,
|
||||
@@ -1726,6 +1802,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
TurnstileSiteKey: updatedSettings.TurnstileSiteKey,
|
||||
TurnstileSecretKeyConfigured: updatedSettings.TurnstileSecretKeyConfigured,
|
||||
APIKeyACLTrustForwardedIP: updatedSettings.APIKeyACLTrustForwardedIP,
|
||||
ForwardedClientIPHeaders: updatedSettings.ForwardedClientIPHeaders,
|
||||
LinuxDoConnectEnabled: updatedSettings.LinuxDoConnectEnabled,
|
||||
LinuxDoConnectClientID: updatedSettings.LinuxDoConnectClientID,
|
||||
LinuxDoConnectClientSecretConfigured: updatedSettings.LinuxDoConnectClientSecretConfigured,
|
||||
|
||||
@@ -22,6 +22,27 @@ type SystemHandler struct {
|
||||
lockSvc *service.SystemOperationLockService
|
||||
}
|
||||
|
||||
// systemUpdateTimeout bounds a full in-place update or rollback: the release
|
||||
// manifest fetch plus a large binary download over slow links. It must stay
|
||||
// above the GitHub download client timeout (10 minutes) so the download owns
|
||||
// its own deadline.
|
||||
const systemUpdateTimeout = 15 * time.Minute
|
||||
|
||||
// systemUpdateContext detaches a long-running update/rollback from the HTTP
|
||||
// request lifetime. Browsers and reverse proxies commonly abort idle requests
|
||||
// after 30-60s (axios default, nginx proxy_read_timeout), which canceled
|
||||
// c.Request.Context() mid-download and killed the update with
|
||||
// "download failed: context canceled" (#4504). The swap keeps running after a
|
||||
// client disconnect; a later retry then hits the system operation lock or
|
||||
// reports "Already up to date".
|
||||
func systemUpdateContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
||||
base := context.Background()
|
||||
if ctx != nil {
|
||||
base = context.WithoutCancel(ctx)
|
||||
}
|
||||
return context.WithTimeout(base, systemUpdateTimeout)
|
||||
}
|
||||
|
||||
type systemUpdateService interface {
|
||||
CheckUpdate(ctx context.Context, force bool) (*service.UpdateInfo, error)
|
||||
PerformUpdate(ctx context.Context) error
|
||||
@@ -75,9 +96,12 @@ func (h *SystemHandler) PerformUpdate(c *gin.Context) {
|
||||
release(releaseReason, succeeded)
|
||||
}()
|
||||
|
||||
if err := h.updateSvc.PerformUpdate(ctx); err != nil {
|
||||
updateCtx, cancel := systemUpdateContext(ctx)
|
||||
defer cancel()
|
||||
|
||||
if err := h.updateSvc.PerformUpdate(updateCtx); err != nil {
|
||||
if errors.Is(err, service.ErrNoUpdateAvailable) {
|
||||
info, checkErr := h.updateSvc.CheckUpdate(ctx, false)
|
||||
info, checkErr := h.updateSvc.CheckUpdate(updateCtx, false)
|
||||
if checkErr != nil {
|
||||
releaseReason = "SYSTEM_UPDATE_FAILED"
|
||||
return nil, checkErr
|
||||
@@ -152,7 +176,10 @@ func (h *SystemHandler) Rollback(c *gin.Context) {
|
||||
}()
|
||||
|
||||
if targetVersion != "" {
|
||||
err = h.updateSvc.RollbackToVersion(ctx, targetVersion)
|
||||
// 指定版本回退同样要下载完整二进制,与更新一样和请求生命周期解耦。
|
||||
rollbackCtx, cancel := systemUpdateContext(ctx)
|
||||
defer cancel()
|
||||
err = h.updateSvc.RollbackToVersion(rollbackCtx, targetVersion)
|
||||
} else {
|
||||
err = h.updateSvc.Rollback()
|
||||
}
|
||||
|
||||
@@ -18,18 +18,22 @@ import (
|
||||
)
|
||||
|
||||
type systemHandlerUpdateServiceStub struct {
|
||||
performErr error
|
||||
updateInfo *service.UpdateInfo
|
||||
checkErr error
|
||||
checkForces []bool
|
||||
performCall int
|
||||
rollbackCall int
|
||||
rollbackToCall int
|
||||
rollbackToVersions []string
|
||||
rollbackToErr error
|
||||
rollbackVersions []service.RollbackVersion
|
||||
rollbackVersionsErr error
|
||||
rollbackVersionsCall int
|
||||
performErr error
|
||||
updateInfo *service.UpdateInfo
|
||||
checkErr error
|
||||
checkForces []bool
|
||||
performCall int
|
||||
performCtxErr error
|
||||
performHasDeadline bool
|
||||
rollbackCall int
|
||||
rollbackToCall int
|
||||
rollbackToCtxErr error
|
||||
rollbackToHasDeadline bool
|
||||
rollbackToVersions []string
|
||||
rollbackToErr error
|
||||
rollbackVersions []service.RollbackVersion
|
||||
rollbackVersionsErr error
|
||||
rollbackVersionsCall int
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) CheckUpdate(_ context.Context, force bool) (*service.UpdateInfo, error) {
|
||||
@@ -37,8 +41,10 @@ func (s *systemHandlerUpdateServiceStub) CheckUpdate(_ context.Context, force bo
|
||||
return s.updateInfo, s.checkErr
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) PerformUpdate(context.Context) error {
|
||||
func (s *systemHandlerUpdateServiceStub) PerformUpdate(ctx context.Context) error {
|
||||
s.performCall++
|
||||
s.performCtxErr = ctx.Err()
|
||||
_, s.performHasDeadline = ctx.Deadline()
|
||||
return s.performErr
|
||||
}
|
||||
|
||||
@@ -52,8 +58,10 @@ func (s *systemHandlerUpdateServiceStub) ListRollbackVersions(context.Context) (
|
||||
return s.rollbackVersions, s.rollbackVersionsErr
|
||||
}
|
||||
|
||||
func (s *systemHandlerUpdateServiceStub) RollbackToVersion(_ context.Context, version string) error {
|
||||
func (s *systemHandlerUpdateServiceStub) RollbackToVersion(ctx context.Context, version string) error {
|
||||
s.rollbackToCall++
|
||||
s.rollbackToCtxErr = ctx.Err()
|
||||
_, s.rollbackToHasDeadline = ctx.Deadline()
|
||||
s.rollbackToVersions = append(s.rollbackToVersions, version)
|
||||
return s.rollbackToErr
|
||||
}
|
||||
@@ -165,6 +173,55 @@ func TestSystemHandlerPerformUpdateFailureStillReturnsInternalError(t *testing.T
|
||||
require.Equal(t, "internal error", body.Message)
|
||||
}
|
||||
|
||||
// TestSystemHandlerPerformUpdateSurvivesClientDisconnect reproduces #4504:
|
||||
// the browser or a reverse proxy (axios 30s default, nginx proxy_read_timeout
|
||||
// 60s) aborts the long-running update request and cancels the request
|
||||
// context. The download must keep running on a detached, bounded context
|
||||
// instead of dying with "download failed: context canceled".
|
||||
func TestSystemHandlerPerformUpdateSurvivesClientDisconnect(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
router := newSystemHandlerTestRouter(t, updateSvc, repo)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/update", nil)
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
req = req.WithContext(canceledCtx)
|
||||
req.Header.Set("Idempotency-Key", "disconnected-update")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, 1, updateSvc.performCall)
|
||||
require.NoError(t, updateSvc.performCtxErr,
|
||||
"update must not observe the canceled request context")
|
||||
require.True(t, updateSvc.performHasDeadline,
|
||||
"detached update context must still be bounded by a deadline")
|
||||
requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded)
|
||||
}
|
||||
|
||||
func TestSystemHandlerRollbackToVersionSurvivesClientDisconnect(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
router := newSystemHandlerTestRouter(t, updateSvc, repo)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/system/rollback",
|
||||
strings.NewReader(`{"version":"0.1.146"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
req = req.WithContext(canceledCtx)
|
||||
req.Header.Set("Idempotency-Key", "disconnected-rollback")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, 1, updateSvc.rollbackToCall)
|
||||
require.NoError(t, updateSvc.rollbackToCtxErr,
|
||||
"versioned rollback must not observe the canceled request context")
|
||||
require.True(t, updateSvc.rollbackToHasDeadline,
|
||||
"detached rollback context must still be bounded by a deadline")
|
||||
requireSystemLockStatus(t, repo, service.IdempotencyStatusSucceeded)
|
||||
}
|
||||
|
||||
func TestSystemHandlerRollbackWithoutBodyUsesLegacyBackup(t *testing.T) {
|
||||
updateSvc := &systemHandlerUpdateServiceStub{}
|
||||
repo := newMemoryIdempotencyRepoStub()
|
||||
|
||||
@@ -33,6 +33,7 @@ type UserHandler struct {
|
||||
billingCache service.BillingCache // T17/T18 缓存失效(PUT/POST 路径)
|
||||
totpService *service.TotpService // 角色提升为管理员的 step-up 门控
|
||||
userService *service.UserService
|
||||
settingService *service.SettingService // step-up 功能开关
|
||||
}
|
||||
|
||||
// NewUserHandler creates a new admin user handler
|
||||
@@ -43,6 +44,7 @@ func NewUserHandler(
|
||||
billingCache service.BillingCache,
|
||||
totpService *service.TotpService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
) *UserHandler {
|
||||
return &UserHandler{
|
||||
adminService: adminService,
|
||||
@@ -51,6 +53,7 @@ func NewUserHandler(
|
||||
billingCache: billingCache,
|
||||
totpService: totpService,
|
||||
userService: userService,
|
||||
settingService: settingService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -275,7 +278,7 @@ func (h *UserHandler) Create(c *gin.Context) {
|
||||
|
||||
// 创建管理员账号属权限敏感操作:需最近完成 step-up 2FA 验证。
|
||||
if req.Role == service.RoleAdmin {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService) {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService, h.settingService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -331,7 +334,7 @@ func (h *UserHandler) Update(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
if target.Role != service.RoleAdmin {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService) {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService, h.settingService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
@@ -35,7 +35,7 @@ func TestUserHandlerListIncludesActivityFieldsAndSortParams(t *testing.T) {
|
||||
UpdatedAt: lastLoginAt,
|
||||
},
|
||||
}
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
@@ -89,7 +89,7 @@ func TestUserHandlerGetByIDIncludesActivityFields(t *testing.T) {
|
||||
UpdatedAt: lastLoginAt,
|
||||
},
|
||||
}
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
@@ -44,7 +44,7 @@ func (s *batchLimitsAdminServiceStub) BatchUpdateLimits(_ context.Context, userI
|
||||
func setupBatchLimitsRouter(serviceStub service.AdminService) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handler := NewUserHandler(serviceStub, nil, nil, nil, nil, nil)
|
||||
handler := NewUserHandler(serviceStub, nil, nil, nil, nil, nil, nil)
|
||||
router.POST("/api/v1/admin/users/batch-limits", handler.BatchUpdateLimits)
|
||||
return router
|
||||
}
|
||||
|
||||
@@ -26,7 +26,7 @@ func (s *getByIDAdminStub) GetUserIncludeDeleted(_ context.Context, id int64) (*
|
||||
func setupGetByIDRouter(svc service.AdminService) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
h := NewUserHandler(svc, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(svc, nil, nil, nil, nil, nil, nil)
|
||||
r.GET("/admin/users/:id", h.GetByID)
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ func TestAdminUserList_ParsesAPIKeyGroupID(t *testing.T) {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
stub := &listUsersFilterStub{AdminService: newStubAdminService()}
|
||||
r := gin.New()
|
||||
h := NewUserHandler(stub, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(stub, nil, nil, nil, nil, nil, nil)
|
||||
r.GET("/admin/users", h.List)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
@@ -29,7 +29,7 @@ func setupRoleStepUpRouter(t *testing.T) (*gin.Engine, *stubAdminService) {
|
||||
Status: service.StatusActive,
|
||||
})
|
||||
|
||||
h := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
router.POST("/api/v1/admin/users", h.Create)
|
||||
router.PUT("/api/v1/admin/users/:id", h.Update)
|
||||
return router, adminSvc
|
||||
|
||||
@@ -46,14 +46,14 @@ type CreateAPIKeyRequest struct {
|
||||
|
||||
// UpdateAPIKeyRequest represents the update API key request payload
|
||||
type UpdateAPIKeyRequest struct {
|
||||
Name string `json:"name"`
|
||||
GroupID *int64 `json:"group_id"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
IPWhitelist []string `json:"ip_whitelist"` // IP 白名单
|
||||
IPBlacklist []string `json:"ip_blacklist"` // IP 黑名单
|
||||
Quota *float64 `json:"quota"` // 配额限制 (USD), 0=无限制
|
||||
ExpiresAt *string `json:"expires_at"` // 过期时间 (ISO 8601)
|
||||
ResetQuota *bool `json:"reset_quota"` // 重置已用配额
|
||||
Name string `json:"name"`
|
||||
GroupID *int64 `json:"group_id"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
IPWhitelist *[]string `json:"ip_whitelist"` // IP 白名单(nil 不修改,空数组清空)
|
||||
IPBlacklist *[]string `json:"ip_blacklist"` // IP 黑名单(nil 不修改,空数组清空)
|
||||
Quota *float64 `json:"quota"` // 配额限制 (USD), 0=无限制
|
||||
ExpiresAt *string `json:"expires_at"` // 过期时间 (ISO 8601)
|
||||
ResetQuota *bool `json:"reset_quota"` // 重置已用配额
|
||||
|
||||
// Rate limit fields (nil = no change, 0 = unlimited)
|
||||
RateLimit5h *float64 `json:"rate_limit_5h"`
|
||||
|
||||
@@ -37,6 +37,7 @@ type SystemSettings struct {
|
||||
TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证
|
||||
TotpEncryptionKeyConfigured bool `json:"totp_encryption_key_configured"` // TOTP 加密密钥是否已配置
|
||||
SessionBindingEnabled bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定
|
||||
StepUpEnabled bool `json:"step_up_enabled"` // 敏感操作 step-up 2FA
|
||||
AuditLogRetentionDays int `json:"audit_log_retention_days"` // 审计日志保留天数
|
||||
LoginAgreementEnabled bool `json:"login_agreement_enabled"`
|
||||
LoginAgreementMode string `json:"login_agreement_mode"`
|
||||
@@ -51,10 +52,11 @@ type SystemSettings struct {
|
||||
SMTPFromName string `json:"smtp_from_name"`
|
||||
SMTPUseTLS bool `json:"smtp_use_tls"`
|
||||
|
||||
TurnstileEnabled bool `json:"turnstile_enabled"`
|
||||
TurnstileSiteKey string `json:"turnstile_site_key"`
|
||||
TurnstileSecretKeyConfigured bool `json:"turnstile_secret_key_configured"`
|
||||
APIKeyACLTrustForwardedIP bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
TurnstileEnabled bool `json:"turnstile_enabled"`
|
||||
TurnstileSiteKey string `json:"turnstile_site_key"`
|
||||
TurnstileSecretKeyConfigured bool `json:"turnstile_secret_key_configured"`
|
||||
APIKeyACLTrustForwardedIP bool `json:"api_key_acl_trust_forwarded_ip"`
|
||||
ForwardedClientIPHeaders []string `json:"forwarded_client_ip_headers"`
|
||||
|
||||
LinuxDoConnectEnabled bool `json:"linuxdo_connect_enabled"`
|
||||
LinuxDoConnectClientID string `json:"linuxdo_connect_client_id"`
|
||||
|
||||
@@ -83,8 +83,9 @@ func (s *FailoverState) HandleFailoverError(
|
||||
return FailoverExhausted
|
||||
}
|
||||
|
||||
// 缓存计费判断
|
||||
if needForceCacheBilling(s.hasBoundSession, failoverErr) {
|
||||
// 同账号重试不算切换账号,粘性会话仅在实际切换时强制缓存计费。
|
||||
sameAccountRetry := failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit
|
||||
if needForceCacheBilling(s.hasBoundSession, failoverErr, sameAccountRetry) {
|
||||
s.ForceCacheBilling = true
|
||||
}
|
||||
|
||||
@@ -174,9 +175,9 @@ func (s *FailoverState) HandleSelectionExhausted(ctx context.Context) FailoverAc
|
||||
}
|
||||
|
||||
// needForceCacheBilling 判断 failover 时是否需要强制缓存计费。
|
||||
// 粘性会话切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费。
|
||||
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError) bool {
|
||||
return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
||||
// 粘性会话实际切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费。
|
||||
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError, sameAccountRetry bool) bool {
|
||||
return (hasBoundSession && !sameAccountRetry) || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
||||
}
|
||||
|
||||
// failoverClientGone 判断下游客户端是否已断开(请求 context 已取消)。
|
||||
|
||||
@@ -271,7 +271,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
t.Run("hasBoundSession为true时设置ForceCacheBilling", func(t *testing.T) {
|
||||
t.Run("hasBoundSession为true且实际切换时设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true) // hasBoundSession=true
|
||||
err := newTestFailoverErr(500, false, false)
|
||||
@@ -280,6 +280,32 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
})
|
||||
|
||||
t.Run("同账号重试时仅凭hasBoundSession不设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, false)
|
||||
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.False(t, fs.ForceCacheBilling)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("同账号重试耗尽并实际切换时设置ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, false)
|
||||
|
||||
for i := 0; i < maxSameAccountRetries; i++ {
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
require.False(t, fs.ForceCacheBilling)
|
||||
}
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
require.Equal(t, 1, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("failoverErr.ForceCacheBilling为true时设置", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
@@ -289,6 +315,17 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
})
|
||||
|
||||
t.Run("同账号重试保留显式ForceCacheBilling", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, true)
|
||||
err := newTestFailoverErr(400, true, true)
|
||||
|
||||
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
|
||||
|
||||
require.True(t, fs.ForceCacheBilling)
|
||||
require.Zero(t, fs.SwitchCount)
|
||||
})
|
||||
|
||||
t.Run("两者均为false时不设置", func(t *testing.T) {
|
||||
mock := &mockTempUnscheduler{}
|
||||
fs := NewFailoverState(3, false)
|
||||
@@ -633,13 +670,14 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
|
||||
for i := 0; i < maxSameAccountRetries; i++ {
|
||||
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
require.False(t, fs.ForceCacheBilling, "同账号重试期间不应仅因绑定会话强制缓存计费")
|
||||
}
|
||||
require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling")
|
||||
|
||||
// 2. 账号 100 超过重试上限 → TempUnschedule + 切换
|
||||
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
|
||||
require.Equal(t, FailoverContinue, action)
|
||||
require.Equal(t, 1, fs.SwitchCount)
|
||||
require.True(t, fs.ForceCacheBilling, "实际切换账号时应设置 ForceCacheBilling")
|
||||
require.Len(t, mock.calls, 1)
|
||||
|
||||
// 3. 账号 200 遇到不可重试错误 → 直接切换
|
||||
|
||||
@@ -166,6 +166,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
reqModel := parsedReq.Model
|
||||
reqStream := parsedReq.Stream
|
||||
reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
|
||||
@@ -1882,6 +1883,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
// count_tokens 走 messages 严格校验时,复用已解析请求,避免二次反序列化。
|
||||
SetClaudeCodeClientContext(c, body, parsedReq)
|
||||
reqLog = reqLog.With(zap.String("model", parsedReq.Model), zap.Bool("stream", parsedReq.Stream))
|
||||
|
||||
@@ -46,6 +46,11 @@ func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
|
||||
}
|
||||
|
||||
// GrokVideoContent proxies downloadable video content through the task's upstream account.
|
||||
func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id"))
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) {
|
||||
streamStarted := false
|
||||
defer h.recoverResponsesPanic(c, &streamStarted)
|
||||
@@ -95,11 +100,12 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
requestInfo := service.ParseGrokMediaRequest(contentType, body)
|
||||
requestModel := requestInfo.Model
|
||||
routingModel := service.NormalizeGrokMediaModelForEndpoint(endpoint, requestModel, requestInfo.HasInputImage())
|
||||
if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" {
|
||||
if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required")
|
||||
return
|
||||
}
|
||||
@@ -159,14 +165,24 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
sessionSeed = []byte(requestID)
|
||||
}
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed)
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus {
|
||||
sessionHash = service.GrokMediaVideoRequestSessionHash(requestID)
|
||||
boundLookupAccountID := int64(0)
|
||||
if endpoint.IsVideoLookupRequest() {
|
||||
sessionHash = service.GrokMediaVideoRequestSessionHash(requestID, subject.UserID, apiKey.ID)
|
||||
boundLookupAccountID, err = h.gatewayService.ResolveGrokMediaVideoRequestAccount(
|
||||
c.Request.Context(), apiKey.GroupID, requestID, subject.UserID, apiKey.ID,
|
||||
)
|
||||
if err != nil || boundLookupAccountID <= 0 {
|
||||
reqLog.Info("grok_media.video_lookup_owner_binding_missing", zap.Error(err))
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
|
||||
return
|
||||
}
|
||||
}
|
||||
requestCtx := c.Request.Context()
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
mediaEligibilityRejected := false
|
||||
switchCount := 0
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
if maxAccountSwitches <= 0 {
|
||||
@@ -184,7 +200,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
apiKey.GroupID,
|
||||
"",
|
||||
sessionHash,
|
||||
requestModel,
|
||||
routingModel,
|
||||
failedAccountIDs,
|
||||
service.OpenAIUpstreamTransportHTTPSSE,
|
||||
requiredCapability,
|
||||
@@ -202,13 +218,14 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
zap.Error(err),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
)
|
||||
if endpoint.IsGenerationRequest() && len(failedAccountIDs) == 0 && errors.Is(err, service.ErrNoAvailableAccounts) {
|
||||
if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) &&
|
||||
(len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
if len(failedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok)
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
@@ -228,13 +245,21 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok)
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
}
|
||||
h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
|
||||
return
|
||||
}
|
||||
if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID {
|
||||
reqLog.Warn("grok_media.video_lookup_bound_account_unavailable",
|
||||
zap.Int64("bound_account_id", boundLookupAccountID),
|
||||
zap.Int64("selected_account_id", selection.Account.ID),
|
||||
)
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
|
||||
return
|
||||
}
|
||||
|
||||
reqLog.Debug("grok_media.account_schedule_decision",
|
||||
zap.String("layer", scheduleDecision.Layer),
|
||||
@@ -246,6 +271,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
)
|
||||
|
||||
account := selection.Account
|
||||
if endpoint.IsGenerationRequest() {
|
||||
eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
|
||||
if !eligible {
|
||||
mediaEligibilityRejected = true
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
reqLog.Warn("grok_media.account_eligibility_rejected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("reason", eligibilityReason),
|
||||
zap.Bool("probe_failed", eligibilityErr != nil),
|
||||
)
|
||||
if switchCount >= maxAccountSwitches {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
continue
|
||||
}
|
||||
}
|
||||
sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account)
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
@@ -285,7 +329,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil)
|
||||
}
|
||||
if c.Writer.Size() != writerSizeBeforeForward {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
@@ -295,6 +339,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
if endpoint.IsVideoLookupRequest() {
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
@@ -333,8 +381,8 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
)
|
||||
continue
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
|
||||
if c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil)
|
||||
if !service.IsResponseCommitted(c) && c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
|
||||
}
|
||||
reqLog.Warn("grok_media.forward_failed",
|
||||
@@ -344,9 +392,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil)
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, result), true, nil)
|
||||
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
|
||||
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
|
||||
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(
|
||||
requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID,
|
||||
); err != nil {
|
||||
reqLog.Warn("grok_media.bind_video_request_account_failed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("request_id", result.ResponseID),
|
||||
@@ -365,6 +415,20 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
}
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) {
|
||||
if account == nil {
|
||||
return false, "missing_account", errors.New("grok media account is required")
|
||||
}
|
||||
eligible, reason := account.GrokMediaGenerationEligibility()
|
||||
if eligible || reason != "billing_unobserved" {
|
||||
return eligible, reason, nil
|
||||
}
|
||||
if h == nil || h.grokMediaEligibilityProber == nil {
|
||||
return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured")
|
||||
}
|
||||
return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID)
|
||||
}
|
||||
|
||||
func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability {
|
||||
if endpoint.IsGenerationRequest() {
|
||||
return service.OpenAIEndpointCapabilityGrokMediaGeneration
|
||||
@@ -372,6 +436,16 @@ func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.Ope
|
||||
return ""
|
||||
}
|
||||
|
||||
func grokMediaScheduleModel(account *service.Account, routingModel string, result *service.OpenAIForwardResult) string {
|
||||
if result != nil && strings.TrimSpace(result.UpstreamModel) != "" {
|
||||
return result.UpstreamModel
|
||||
}
|
||||
if account == nil {
|
||||
return strings.TrimSpace(routingModel)
|
||||
}
|
||||
return account.GetMappedModel(routingModel)
|
||||
}
|
||||
|
||||
func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool {
|
||||
return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != ""
|
||||
}
|
||||
|
||||
@@ -1,12 +1,26 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type grokMediaEligibilityProberStub struct {
|
||||
eligible bool
|
||||
reason string
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *grokMediaEligibilityProberStub) ProbeMediaEligibility(context.Context, int64) (bool, string, error) {
|
||||
s.calls++
|
||||
return s.eligible, s.reason, s.err
|
||||
}
|
||||
|
||||
func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -38,6 +52,12 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "video content skips usage",
|
||||
endpoint: service.GrokMediaEndpointVideoContent,
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "generation skips usage without model",
|
||||
endpoint: service.GrokMediaEndpointImagesGenerations,
|
||||
@@ -65,6 +85,7 @@ func TestGrokMediaRequiredCapability(t *testing.T) {
|
||||
{name: "video edit", endpoint: service.GrokMediaEndpointVideosEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video extension", endpoint: service.GrokMediaEndpointVideosExtensions, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video status preserves lookup", endpoint: service.GrokMediaEndpointVideoStatus, want: ""},
|
||||
{name: "video content preserves lookup", endpoint: service.GrokMediaEndpointVideoContent, want: ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -73,3 +94,74 @@ func TestGrokMediaRequiredCapability(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokMediaScheduleModelUsesNormalizedMappedUpstream(t *testing.T) {
|
||||
account := &service.Account{
|
||||
Platform: service.PlatformGrok,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"grok-imagine-video-1.5": "wrong-raw-model",
|
||||
"grok-imagine-video": "mapped-video-model",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", nil))
|
||||
require.Equal(t, "actual-upstream-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{
|
||||
UpstreamModel: "actual-upstream-model",
|
||||
}))
|
||||
require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{}))
|
||||
require.Equal(t, "grok-imagine-video", grokMediaScheduleModel(nil, " grok-imagine-video ", nil))
|
||||
}
|
||||
|
||||
func TestEnsureGrokMediaAccountEligibility(t *testing.T) {
|
||||
t.Run("non oauth account does not probe", func(t *testing.T) {
|
||||
prober := &grokMediaEligibilityProberStub{}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, eligible)
|
||||
require.Equal(t, "non_oauth", reason)
|
||||
require.Zero(t, prober.calls)
|
||||
})
|
||||
|
||||
t.Run("unobserved oauth is probed before forwarding", func(t *testing.T) {
|
||||
prober := &grokMediaEligibilityProberStub{eligible: true, reason: "eligible"}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{ID: 7, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, eligible)
|
||||
require.Equal(t, "eligible", reason)
|
||||
require.Equal(t, 1, prober.calls)
|
||||
})
|
||||
|
||||
t.Run("missing prober fails closed", func(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
account := &service.Account{ID: 8, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.Error(t, err)
|
||||
require.False(t, eligible)
|
||||
require.Equal(t, "billing_probe_unavailable", reason)
|
||||
})
|
||||
|
||||
t.Run("probe failure fails closed", func(t *testing.T) {
|
||||
probeErr := errors.New("probe failed")
|
||||
prober := &grokMediaEligibilityProberStub{reason: "billing_unobserved", err: probeErr}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{ID: 9, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.ErrorIs(t, err, probeErr)
|
||||
require.False(t, eligible)
|
||||
require.Equal(t, "billing_unobserved", reason)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type toggleSettingRepo struct {
|
||||
mu sync.Mutex
|
||||
values map[string]string
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) Get(context.Context, string) (*service.Setting, error) { return nil, nil }
|
||||
func (r *toggleSettingRepo) GetValue(_ context.Context, key string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return r.values[key], nil
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) Set(_ context.Context, key, value string) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.values[key] = value
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *toggleSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *toggleSettingRepo) SetMultiple(context.Context, map[string]string) error { return nil }
|
||||
func (r *toggleSettingRepo) GetAll(context.Context) (map[string]string, error) {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
func (r *toggleSettingRepo) Delete(context.Context, string) error { return nil }
|
||||
|
||||
type passthroughEncryptor struct{}
|
||||
|
||||
func (passthroughEncryptor) Encrypt(plaintext string) (string, error) { return plaintext, nil }
|
||||
func (passthroughEncryptor) Decrypt(ciphertext string) (string, error) { return ciphertext, nil }
|
||||
|
||||
type noopImageStorage struct{}
|
||||
|
||||
func (noopImageStorage) Save(context.Context, string, string, []byte) (string, error) {
|
||||
return "https://cdn.example.test/object.png", nil
|
||||
}
|
||||
|
||||
// TestAsyncImageEnablesWithoutRestart drives the actual HTTP path for the bug behind
|
||||
// #4458 and #4542: with object storage unconfigured the async endpoint 404s, and the
|
||||
// only way to turn it on used to be editing config.yaml and restarting the container.
|
||||
// Flipping the admin setting must flip the endpoint over in the same process.
|
||||
func TestAsyncImageEnablesWithoutRestart(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
repo := &toggleSettingRepo{values: map[string]string{}}
|
||||
// A fixed encryption key is required to persist a new S3 secret (#4524).
|
||||
backup := service.NewBackupService(repo, &config.Config{
|
||||
Totp: config.TotpConfig{EncryptionKeyConfigured: true},
|
||||
}, passthroughEncryptor{}, nil, nil)
|
||||
factory := func(context.Context, *config.ImageStorageConfig) (service.ImageStorage, error) {
|
||||
return noopImageStorage{}, nil
|
||||
}
|
||||
settings := service.NewImageStorageSettingService(repo, passthroughEncryptor{}, backup, factory, config.ImageStorageConfig{})
|
||||
|
||||
store := &asyncImageMemoryStore{tasks: make(map[string]*service.ImageTaskRecord)}
|
||||
tasks := service.NewImageTaskServiceWithResolver(store, settings.Resolver(), time.Hour, time.Minute)
|
||||
|
||||
h := &AsyncImageHandler{tasks: tasks}
|
||||
h.execute = func(_ string, c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"created": 1, "data": []gin.H{{"url": "https://upstream.test/i.png"}}})
|
||||
}
|
||||
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
groupID := int64(3)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 9, UserID: 7, GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, AllowImageGeneration: true},
|
||||
})
|
||||
c.Next()
|
||||
})
|
||||
router.POST("/v1/images/generations/async", h.Submit)
|
||||
router.GET("/v1/images/tasks/:task_id", h.Get)
|
||||
|
||||
submit := func() *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async",
|
||||
strings.NewReader(`{"model":"gpt-image-1","prompt":"a lighthouse"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
rec := submit()
|
||||
require.Equal(t, http.StatusNotFound, rec.Code, "disabled until an admin configures object storage")
|
||||
require.Contains(t, rec.Body.String(), "async image tasks are not enabled")
|
||||
|
||||
// The admin saves the setting — no restart, same process.
|
||||
_, err := settings.Update(context.Background(), service.ImageStorageSettings{
|
||||
Enabled: true, Bucket: "my-images",
|
||||
Endpoint: "https://acct.r2.cloudflarestorage.com", AccessKeyID: "ak", SecretAccessKey: "sk",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
rec = submit()
|
||||
require.Equal(t, http.StatusAccepted, rec.Code, "the endpoint must go live as soon as the setting is saved")
|
||||
|
||||
var accepted struct {
|
||||
TaskID string `json:"task_id"`
|
||||
PollURL string `json:"poll_url"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &accepted))
|
||||
require.NotEmpty(t, accepted.TaskID)
|
||||
|
||||
// Turning the feature back off must not strand a task that was already accepted.
|
||||
_, err = settings.Update(context.Background(), service.ImageStorageSettings{Enabled: false})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Equal(t, http.StatusNotFound, submit().Code, "new submissions are refused again")
|
||||
|
||||
pollRec := httptest.NewRecorder()
|
||||
router.ServeHTTP(pollRec, httptest.NewRequest(http.MethodGet, accepted.PollURL, nil))
|
||||
require.Equal(t, http.StatusOK, pollRec.Code, "an already-accepted task stays pollable after the switch is turned off")
|
||||
}
|
||||
@@ -39,6 +39,13 @@ func (h *AsyncImageHandler) enabled() bool {
|
||||
return h != nil && h.tasks != nil && h.tasks.Enabled()
|
||||
}
|
||||
|
||||
// pollable reports whether task lookups can be served. It is deliberately weaker
|
||||
// than enabled(): results already written to Redis stay readable after the
|
||||
// feature is switched off, so an in-flight task is never stranded.
|
||||
func (h *AsyncImageHandler) pollable() bool {
|
||||
return h != nil && h.tasks != nil && h.tasks.Pollable()
|
||||
}
|
||||
|
||||
// Submit accepts the same payload as the synchronous Images endpoint and
|
||||
// returns before the upstream image generation begins.
|
||||
func (h *AsyncImageHandler) Submit(c *gin.Context) {
|
||||
@@ -155,7 +162,11 @@ func (h *AsyncImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, apiKe
|
||||
}
|
||||
|
||||
func (h *AsyncImageHandler) Get(c *gin.Context) {
|
||||
if !h.enabled() {
|
||||
// Polling deliberately does not require the feature to be enabled, only that
|
||||
// the task store is reachable. Turning the switch off in the admin UI must not
|
||||
// strand tasks that were already accepted — their results are still in Redis
|
||||
// and their submitters are still polling.
|
||||
if !h.pollable() {
|
||||
imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "async image tasks are not enabled")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -38,11 +38,12 @@ type noAccountErrorClassification struct {
|
||||
// The classifier intentionally does not consume the original error: the
|
||||
// selection layer never tells us *why* the pool came up empty (rate-limited
|
||||
// vs. unsupported model are both wrapped as ErrNoAvailableAccounts). Instead
|
||||
// we re-check pool composition through DiagnoseModelAvailabilityForPlatform,
|
||||
// which only inspects model_mapping configuration and ignores transient
|
||||
// state. That guarantees a 404 is only returned when no operator action
|
||||
// short of editing the account's model_mapping could make this request
|
||||
// succeed.
|
||||
// we re-check pool composition through DiagnoseModelAvailabilityForPlatform.
|
||||
// Its dedicated database query considers only persistent eligibility
|
||||
// (active status + schedulable setting) and model_mapping, bypassing scheduler
|
||||
// snapshots and transient filters. That guarantees a 404 is only returned
|
||||
// when persistent account/group/model configuration must change before the
|
||||
// request can succeed.
|
||||
//
|
||||
// routingModel is the model name that account selection actually compared
|
||||
// against (i.e. after group-level dispatch mapping). displayModel is the
|
||||
|
||||
@@ -154,6 +154,20 @@ func TestClassifyNoAccountError_HasModelSupport_KeepsRoutingMessageGenerationToC
|
||||
require.False(t, cls.ModelNotFound)
|
||||
}
|
||||
|
||||
func TestClassifyNoAccountError_ModelSupportedOnlyByRateLimitedAccount_Returns503(t *testing.T) {
|
||||
c := newTestGinContextWithRequest()
|
||||
// The diagnoser's configured-state lookup still sees the model-supporting
|
||||
// account even though normal scheduling has excluded it during cooldown.
|
||||
fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}}
|
||||
apiKey := &service.APIKey{GroupID: ptrInt64(7)}
|
||||
|
||||
cls := classifyNoAccountErrorFromGin(c, fd, apiKey, "claude-opus-4-8", "claude-opus-4-8", service.PlatformAnthropic)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, cls.Status)
|
||||
require.Equal(t, "api_error", cls.ErrType)
|
||||
require.False(t, cls.ModelNotFound, "temporary account cooldown must remain retryable")
|
||||
}
|
||||
|
||||
func TestClassifyNoAccountError_NoAccountsInPool_Stays503(t *testing.T) {
|
||||
c := newTestGinContextWithRequest()
|
||||
fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: false, HasModelSupport: false}}
|
||||
|
||||
@@ -305,7 +305,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
|
||||
wroteFallback := false
|
||||
if !upstreamErrorAlreadyCommunicated {
|
||||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||||
wroteFallback = h.ensureOpenAIStreamReadErrorResponse(c, err, streamStarted)
|
||||
if !wroteFallback {
|
||||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||||
}
|
||||
}
|
||||
reqLog.Warn("openai_chat_completions.forward_failed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
|
||||
@@ -53,6 +53,8 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
|
||||
return
|
||||
}
|
||||
// 让 ops 错误日志携带实际选中的上游账号,便于定位失效账号(#4544)。
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
|
||||
if err != nil {
|
||||
|
||||
@@ -13,6 +13,48 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// GrokCountTokens handles Anthropic-compatible count_tokens requests locally.
|
||||
// The route middleware already authenticates the API key and resolves the
|
||||
// group; this handler intentionally does not select an account or check billing.
|
||||
func (h *OpenAIGatewayHandler) GrokCountTokens(c *gin.Context) {
|
||||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
h.anthropicErrorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return
|
||||
}
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||||
return
|
||||
}
|
||||
if len(body) == 0 {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||||
return
|
||||
}
|
||||
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
|
||||
if err != nil {
|
||||
logRequestBodyParseFailure(requestLogger(c, "handler.openai_gateway.grok_count_tokens"), body, err)
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
if parsedReq.Model == "" {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
|
||||
estimated, err := service.EstimateGrokCountTokens(parsedReq.Body.Bytes())
|
||||
if err != nil {
|
||||
requestLogger(c, "handler.openai_gateway.grok_count_tokens").Warn("grok_count_tokens.local_estimate_failed", zap.Error(err))
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, parsedReq.Model, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(false, false)))
|
||||
c.JSON(http.StatusOK, gin.H{"input_tokens": estimated})
|
||||
}
|
||||
|
||||
// CountTokens handles Anthropic-compatible POST /v1/messages/count_tokens for OpenAI groups.
|
||||
// It validates billing and routes to an OpenAI token-count bridge without taking concurrency slots
|
||||
// or recording usage.
|
||||
@@ -67,6 +109,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
if parsedReq.Model == "" {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
|
||||
@@ -645,7 +645,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -657,7 +658,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -678,7 +680,7 @@ func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) {
|
||||
{name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`},
|
||||
{name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`},
|
||||
{name: "chat completions raw fallback", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stop":["END"],"stream":false}`},
|
||||
{name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"},
|
||||
{name: "grok media", method: http.MethodPost, path: "/openai/v1/videos/generations", body: `{"model":"grok-imagine-video","prompt":"waves"}`},
|
||||
}
|
||||
|
||||
for _, endpoint := range endpoints {
|
||||
@@ -822,6 +824,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "expired", "refresh_token": "revoked-refresh",
|
||||
"expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
{
|
||||
ID: 802, Name: "healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
@@ -830,6 +833,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "healthy-access", "refresh_token": "healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
}
|
||||
if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" {
|
||||
@@ -843,6 +847,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
})
|
||||
}
|
||||
if mode == "oauth_429_apikey_500" {
|
||||
@@ -904,7 +909,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
apiKey := &service.APIKey{
|
||||
ID: 902, GroupID: &groupID,
|
||||
User: &service.User{ID: 903, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive, AllowImageGeneration: true},
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
@@ -916,6 +921,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
router.POST("/openai/v1/messages", h.Messages)
|
||||
router.POST("/openai/v1/chat/completions", h.ChatCompletions)
|
||||
router.POST("/openai/v1/videos/generations", h.GrokVideoGeneration)
|
||||
router.GET("/openai/v1/videos/:request_id", h.GrokVideoStatus)
|
||||
handlerRefresherStarted.Store(router, refresher.started)
|
||||
cleanup := func() {
|
||||
|
||||
@@ -206,6 +206,10 @@ func TestOpsRecoveredCredentialFailoverUsesAccountAuthAttribution(t *testing.T)
|
||||
require.NotContains(t, job.entry.ErrorMessage, "earlier inference failure")
|
||||
require.NotNil(t, job.entry.UpstreamStatusCode)
|
||||
require.Zero(t, *job.entry.UpstreamStatusCode)
|
||||
require.Len(t, job.entry.UpstreamErrors, 2)
|
||||
require.Equal(t, http.StatusForbidden, job.entry.UpstreamErrors[0].UpstreamStatusCode)
|
||||
require.Nil(t, job.entry.UpstreamErrors)
|
||||
require.NotNil(t, job.entry.UpstreamErrorsJSON)
|
||||
events, err := service.ParseOpsUpstreamErrors(*job.entry.UpstreamErrorsJSON)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, events, 2)
|
||||
require.Equal(t, http.StatusForbidden, events[0].UpstreamStatusCode)
|
||||
}
|
||||
|
||||
@@ -28,18 +28,23 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
securityAuditCoordinator *securityaudit.Coordinator
|
||||
opsService *service.OpsService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
imageLimiter *imageConcurrencyLimiter
|
||||
maxAccountSwitches int
|
||||
cfg *config.Config
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
securityAuditCoordinator *securityaudit.Coordinator
|
||||
grokMediaEligibilityProber grokMediaEligibilityProber
|
||||
opsService *service.OpsService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
imageLimiter *imageConcurrencyLimiter
|
||||
maxAccountSwitches int
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
type grokMediaEligibilityProber interface {
|
||||
ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error)
|
||||
}
|
||||
|
||||
const maxOpenAIFirstOutputTimeoutSwitches = 1
|
||||
@@ -119,6 +124,13 @@ func openAICompatibleRequestPlatform(apiKey *service.APIKey) string {
|
||||
return service.PlatformOpenAI
|
||||
}
|
||||
|
||||
func openAIResponsesRequiredCapability(imageIntent bool, platform string) service.OpenAIEndpointCapability {
|
||||
if imageIntent && platform == service.PlatformOpenAI {
|
||||
return service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
return service.OpenAIEndpointCapabilityChatCompletions
|
||||
}
|
||||
|
||||
func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
|
||||
if apiKey == nil || apiKey.Group == nil {
|
||||
return true
|
||||
@@ -358,12 +370,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 生图意图的 /v1/responses 请求必须调度到确实支持 Responses API 的账号,否则
|
||||
// 会在 forward 阶段被静默降级为无法生图的 Chat Completions 直转(#4417)。
|
||||
// 仅对 OpenAI 平台生效:Grok 生图走独立的 forwardGrokResponses 路径,不应被过滤。
|
||||
// 使用 IsExplicitImageGenerationIntent 排除被动 image_gen namespace 声明,
|
||||
// 避免 Codex 的被动工具目录使 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, body) && requestPlatform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
// 复用前置权限与并发阶段在未修改 body 上确认的显式生图意图,避免大 tools 请求重复扫描。
|
||||
// 该判断已排除 Codex 被动 image_gen namespace,避免 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := openAIResponsesRequiredCapability(imageIntent, requestPlatform)
|
||||
|
||||
for {
|
||||
// Streaming Forward intentionally detaches the upstream request so usage can
|
||||
@@ -2231,12 +2240,29 @@ func (h *OpenAIGatewayHandler) mapUpstreamError(statusCode int) (int, string, st
|
||||
|
||||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||||
func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||||
h.handleStreamingAwareErrorWithCode(c, status, errType, "", message, streamStarted, false)
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleStreamingAwareErrorWithCode(
|
||||
c *gin.Context,
|
||||
status int,
|
||||
errType string,
|
||||
code string,
|
||||
message string,
|
||||
streamStarted bool,
|
||||
countTowardsSLA bool,
|
||||
) {
|
||||
// body-signal compact 心跳可能已把响应头提交为 200:先停心跳(建立
|
||||
// happens-before,接管 ResponseWriter),并升级为流内错误处理。
|
||||
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
|
||||
streamStarted = true
|
||||
}
|
||||
if streamStarted {
|
||||
if countTowardsSLA {
|
||||
service.MarkOpsStreamFailure(c, errType, code, message, status)
|
||||
} else {
|
||||
service.MarkOpsStreamError(c, errType, message, status)
|
||||
}
|
||||
// /v1/responses 的严格 SDK(Codex CLI)要求终止事件必须属于
|
||||
// response.completed/failed/incomplete/cancelled 集合。
|
||||
// 通用 `event: error` 帧不被识别为终止事件,会导致
|
||||
@@ -2249,8 +2275,15 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status
|
||||
// Stream already started, send error as SSE event then close
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if ok {
|
||||
// SSE 错误事件固定 schema,使用 Quote 直拼可避免额外 Marshal 分配。
|
||||
errorEvent := "event: error\ndata: " + `{"error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n"
|
||||
errorObject := gin.H{"type": errType, "message": message}
|
||||
if code != "" {
|
||||
errorObject["code"] = code
|
||||
}
|
||||
payload, err := json.Marshal(gin.H{"error": errorObject})
|
||||
if err != nil {
|
||||
payload = []byte(`{"error":{"type":"upstream_error","message":"Upstream request failed"}}`)
|
||||
}
|
||||
errorEvent := "event: error\ndata: " + string(payload) + "\n\n"
|
||||
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
||||
_ = c.Error(err)
|
||||
}
|
||||
@@ -2260,7 +2293,27 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status
|
||||
}
|
||||
|
||||
// Normal case: return JSON response with proper status code
|
||||
h.errorResponse(c, status, errType, message)
|
||||
if code == "" {
|
||||
h.errorResponse(c, status, errType, message)
|
||||
return
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{
|
||||
"type": errType, "code": code, "message": message,
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) ensureOpenAIStreamReadErrorResponse(c *gin.Context, err error, streamStarted bool) bool {
|
||||
code, message, ok := service.OpenAIUpstreamStreamReadErrorDetails(err)
|
||||
if !ok || c == nil || c.Writer == nil || service.IsResponseCommitted(c) {
|
||||
return false
|
||||
}
|
||||
if c.Writer.Written() {
|
||||
streamStarted = true
|
||||
}
|
||||
h.handleStreamingAwareErrorWithCode(
|
||||
c, http.StatusBadGateway, "upstream_error", code, message, streamStarted, true,
|
||||
)
|
||||
return true
|
||||
}
|
||||
|
||||
// ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。
|
||||
|
||||
@@ -96,6 +96,35 @@ func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIHandleStreamingAwareErrorWithCode_EmitsStableClassification(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
h := &OpenAIGatewayHandler{}
|
||||
h.handleStreamingAwareErrorWithCode(
|
||||
c,
|
||||
http.StatusBadGateway,
|
||||
"upstream_error",
|
||||
service.OpenAIUpstreamHTTP2StreamErrorCode,
|
||||
"Upstream HTTP/2 stream failed",
|
||||
true,
|
||||
true,
|
||||
)
|
||||
|
||||
body := w.Body.String()
|
||||
require.Contains(t, body, "event: error\n")
|
||||
require.Equal(t, "upstream_error", gjson.Get(body[strings.Index(body, "{"):], "error.type").String())
|
||||
require.Equal(t, service.OpenAIUpstreamHTTP2StreamErrorCode, gjson.Get(body[strings.Index(body, "{"):], "error.code").String())
|
||||
require.NotContains(t, body, "stream ID")
|
||||
|
||||
streamErr, ok := service.GetOpsStreamError(c)
|
||||
require.True(t, ok)
|
||||
require.True(t, streamErr.CountTowardsSLA)
|
||||
require.Equal(t, http.StatusBadGateway, streamErr.IntendedStatus)
|
||||
}
|
||||
|
||||
func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
|
||||
require.True(t, openAIForwardSucceededForScheduling(nil))
|
||||
require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{}))
|
||||
@@ -109,6 +138,39 @@ func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
|
||||
}))
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesRequiredCapability(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
imageIntent bool
|
||||
platform string
|
||||
want service.OpenAIEndpointCapability
|
||||
}{
|
||||
{
|
||||
name: "OpenAI explicit image intent requires Responses",
|
||||
imageIntent: true,
|
||||
platform: service.PlatformOpenAI,
|
||||
want: service.OpenAIEndpointCapabilityResponses,
|
||||
},
|
||||
{
|
||||
name: "Grok explicit image intent keeps chat capability",
|
||||
imageIntent: true,
|
||||
platform: service.PlatformGrok,
|
||||
want: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
},
|
||||
{
|
||||
name: "non-image intent keeps chat capability",
|
||||
platform: service.PlatformOpenAI,
|
||||
want: service.OpenAIEndpointCapabilityChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, openAIResponsesRequiredCapability(tt.imageIntent, tt.platform))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveOpenAIMessagesMetadataSession_DoesNotDerivePromptCacheKey(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-sonnet-4-5","metadata":{"user_id":"claude-code-session"},"messages":[{"role":"user","content":"hello"}]}`)
|
||||
|
||||
@@ -1875,6 +1937,212 @@ func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T
|
||||
require.Equal(t, []int64{int64(9902)}, accountRepo.rateLimitedIDs)
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_FirstOutputTimeoutWithoutDownstreamReusesClientForOneFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
firstHitCh := make(chan []byte, 1)
|
||||
secondHitCh := make(chan []byte, 1)
|
||||
var firstConnections atomic.Int32
|
||||
var secondConnections atomic.Int32
|
||||
|
||||
firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
firstConnections.Add(1)
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr == nil {
|
||||
firstHitCh <- payload
|
||||
}
|
||||
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
case <-time.After(3 * time.Second):
|
||||
}
|
||||
}))
|
||||
defer firstUpstream.Close()
|
||||
|
||||
secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
secondConnections.Add(1)
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr == nil {
|
||||
secondHitCh <- payload
|
||||
}
|
||||
|
||||
for _, event := range []string{
|
||||
`{"type":"response.created","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1"}}`,
|
||||
`{"type":"response.output_text.delta","response_id":"resp_ws_timeout_b","delta":"recovered"}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
} {
|
||||
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(event))
|
||||
cancelWrite()
|
||||
if writeErr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
readCtx, cancelRead = context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, _, _ = conn.Read(readCtx)
|
||||
cancelRead()
|
||||
}))
|
||||
defer secondUpstream.Close()
|
||||
|
||||
groupID := int64(4212)
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 9912,
|
||||
Name: "openai-ws-first-semantic-timeout",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
Credentials: map[string]any{"api_key": "sk-first", "base_url": firstUpstream.URL},
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: 9913,
|
||||
Name: "openai-ws-failover-healthy",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 2,
|
||||
Credentials: map[string]any{"api_key": "sk-second", "base_url": secondUpstream.URL},
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.RunMode = config.RunModeSimple
|
||||
cfg.Default.RateMultiplier = 1
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1
|
||||
cfg.Gateway.OpenAIWS.Enabled = true
|
||||
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 3
|
||||
cfg.Gateway.MaxAccountSwitches = 3
|
||||
|
||||
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
|
||||
rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
gatewaySvc := service.NewOpenAIGatewayService(
|
||||
accountRepo, nil, nil, nil, nil, nil, nil, cfg, nil, nil,
|
||||
service.NewBillingService(cfg, nil), rateLimitSvc, billingCacheSvc,
|
||||
nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cache := &concurrencyCacheMock{
|
||||
acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil },
|
||||
acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) {
|
||||
return true, nil
|
||||
},
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
|
||||
maxAccountSwitches: 3,
|
||||
}
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 1812,
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1712, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
|
||||
}
|
||||
handlerDone := make(chan struct{})
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
|
||||
c.Next()
|
||||
})
|
||||
router.GET("/openai/v1/responses", func(c *gin.Context) {
|
||||
h.ResponsesWebSocket(c)
|
||||
close(handlerDone)
|
||||
})
|
||||
handlerServer := httptest.NewServer(router)
|
||||
defer handlerServer.Close()
|
||||
|
||||
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
clientConn, _, err := coderws.Dial(
|
||||
dialCtx,
|
||||
"ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
|
||||
&coderws.DialOptions{CompressionMode: coderws.CompressionContextTakeover},
|
||||
)
|
||||
cancelDial()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = clientConn.CloseNow() }()
|
||||
|
||||
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
var eventTypes []string
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 6*time.Second)
|
||||
for {
|
||||
_, event, readErr := clientConn.Read(readCtx)
|
||||
require.NoError(t, readErr)
|
||||
eventType := gjson.GetBytes(event, "type").String()
|
||||
eventTypes = append(eventTypes, eventType)
|
||||
if eventType == "response.completed" {
|
||||
require.Equal(t, "resp_ws_timeout_b", gjson.GetBytes(event, "response.id").String())
|
||||
break
|
||||
}
|
||||
}
|
||||
cancelRead()
|
||||
require.Contains(t, eventTypes, "response.output_text.delta")
|
||||
require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
|
||||
|
||||
select {
|
||||
case <-handlerDone:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("websocket handler did not finish after healthy failover turn")
|
||||
}
|
||||
select {
|
||||
case <-firstHitCh:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("first upstream did not receive replayable request")
|
||||
}
|
||||
select {
|
||||
case <-secondHitCh:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("second upstream did not receive replayed request")
|
||||
}
|
||||
require.Equal(t, int32(1), firstConnections.Load())
|
||||
require.Equal(t, int32(1), secondConnections.Load())
|
||||
require.NotContains(t, accountRepo.rateLimitedIDs, int64(9913), "healthy failover account must not be penalized")
|
||||
}
|
||||
|
||||
func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
var openAIResponsesImageIntentRoutingBenchmarkSink service.OpenAIEndpointCapability
|
||||
|
||||
func BenchmarkOpenAIResponsesImageIntentRouting_LargeToolsBody(b *testing.B) {
|
||||
body := buildLargeOpenAIResponsesToolsBody(32 << 20)
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body) {
|
||||
b.Fatal("large tools body must not have explicit image intent")
|
||||
}
|
||||
platform := service.PlatformOpenAI
|
||||
|
||||
b.Run("reuse_once", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(body)))
|
||||
b.ReportAllocs()
|
||||
for range b.N {
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body)
|
||||
openAIResponsesImageIntentRoutingBenchmarkSink = openAIResponsesRequiredCapability(imageIntent, platform)
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("rescan_twice", func(b *testing.B) {
|
||||
b.SetBytes(int64(len(body)))
|
||||
b.ReportAllocs()
|
||||
for range b.N {
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body)
|
||||
// 对照优化前路径:路由阶段会再次扫描同一份未修改的 body。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", "gpt-5.4", body) && platform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
if imageIntent && requiredCapability != service.OpenAIEndpointCapabilityResponses {
|
||||
b.Fatal("explicit image intent must require Responses")
|
||||
}
|
||||
openAIResponsesImageIntentRoutingBenchmarkSink = requiredCapability
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func buildLargeOpenAIResponsesToolsBody(targetBytes int) []byte {
|
||||
var builder strings.Builder
|
||||
builder.Grow(targetBytes + 256)
|
||||
_, _ = builder.WriteString(`{"model":"gpt-5.4","tools":[{"type":"function","name":"search","description":"`)
|
||||
_, _ = builder.WriteString(strings.Repeat("x", targetBytes))
|
||||
_, _ = builder.WriteString(`"}],"tool_choice":"auto","input":"write code"}`)
|
||||
return []byte(builder.String())
|
||||
}
|
||||
@@ -72,24 +72,10 @@ const (
|
||||
opsErrorLogMinQueueSize = 256
|
||||
opsErrorLogMaxQueueSize = 8192
|
||||
opsErrorLogBatchSize = 32
|
||||
opsErrorLogMaxQueueBytes = 32 * 1024 * 1024
|
||||
opsErrorLogMaxUserAgentBytes = 512
|
||||
)
|
||||
|
||||
// looksLikeSystemKey 粗筛"形似本系统 key"的输入:长度 16-128 且仅含 [a-zA-Z0-9_-]。
|
||||
// 不用前缀匹配(APIKeyPrefix 可配置)。用于反查审计表前挡掉随机扫描的乱码输入。
|
||||
func looksLikeSystemKey(key string) bool {
|
||||
if len(key) < 16 || len(key) > 128 {
|
||||
return false
|
||||
}
|
||||
for _, c := range key {
|
||||
allowed := (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
|
||||
(c >= '0' && c <= '9') || c == '_' || c == '-'
|
||||
if !allowed {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// keyPrefix 返回脱敏前缀(前 n 个字符);不足 n 则原样返回。
|
||||
func keyPrefix(key string, n int) string {
|
||||
if len(key) <= n {
|
||||
@@ -98,43 +84,26 @@ func keyPrefix(key string, n int) string {
|
||||
return key[:n]
|
||||
}
|
||||
|
||||
// extractAttemptedKey 按认证中间件同样的顺序从请求头提取提交的 key 明文。
|
||||
// 与 api_key_auth.go:43-59 一致:Authorization 仅取 Bearer scheme,非 Bearer 则忽略并继续 x-api-key → x-goog-api-key。
|
||||
func extractAttemptedKey(c *gin.Context) string {
|
||||
if h := c.GetHeader("Authorization"); h != "" {
|
||||
parts := strings.SplitN(h, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
return strings.TrimSpace(parts[1])
|
||||
}
|
||||
// 非 Bearer:与中间件一致,忽略 Authorization,继续尝试其它 header(不在此 return)。
|
||||
}
|
||||
if k := c.GetHeader("x-api-key"); k != "" {
|
||||
return strings.TrimSpace(k)
|
||||
}
|
||||
if k := c.GetHeader("x-goog-api-key"); k != "" {
|
||||
return strings.TrimSpace(k)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type opsErrorLogJob struct {
|
||||
ops *service.OpsService
|
||||
entry *service.OpsInsertErrorLogInput
|
||||
ops *service.OpsService
|
||||
entry *service.OpsInsertErrorLogInput
|
||||
queuedBytes int64
|
||||
}
|
||||
|
||||
var (
|
||||
opsErrorLogOnce sync.Once
|
||||
opsErrorLogQueue chan opsErrorLogJob
|
||||
|
||||
opsErrorLogStopOnce sync.Once
|
||||
opsErrorLogWorkersWg sync.WaitGroup
|
||||
opsErrorLogMu sync.RWMutex
|
||||
opsErrorLogStopping bool
|
||||
opsErrorLogQueueLen atomic.Int64
|
||||
opsErrorLogEnqueued atomic.Int64
|
||||
opsErrorLogDropped atomic.Int64
|
||||
opsErrorLogProcessed atomic.Int64
|
||||
opsErrorLogSanitized atomic.Int64
|
||||
opsErrorLogStopOnce sync.Once
|
||||
opsErrorLogWorkersWg sync.WaitGroup
|
||||
opsErrorLogMu sync.RWMutex
|
||||
opsErrorLogStopping bool
|
||||
opsErrorLogQueueLen atomic.Int64
|
||||
opsErrorLogQueueBytes atomic.Int64
|
||||
opsErrorLogEnqueued atomic.Int64
|
||||
opsErrorLogDropped atomic.Int64
|
||||
opsErrorLogProcessed atomic.Int64
|
||||
opsErrorLogSanitized atomic.Int64
|
||||
|
||||
opsErrorLogLastDropLogAt atomic.Int64
|
||||
|
||||
@@ -154,6 +123,7 @@ func startOpsErrorLogWorkers() {
|
||||
workerCount, queueSize := opsErrorLogConfig()
|
||||
opsErrorLogQueue = make(chan opsErrorLogJob, queueSize)
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
|
||||
opsErrorLogWorkersWg.Add(workerCount)
|
||||
for i := 0; i < workerCount; i++ {
|
||||
@@ -165,6 +135,7 @@ func startOpsErrorLogWorkers() {
|
||||
return
|
||||
}
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-job.queuedBytes)
|
||||
batch := make([]opsErrorLogJob, 0, opsErrorLogBatchSize)
|
||||
batch = append(batch, job)
|
||||
|
||||
@@ -184,6 +155,7 @@ func startOpsErrorLogWorkers() {
|
||||
return
|
||||
}
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-nextJob.queuedBytes)
|
||||
batch = append(batch, nextJob)
|
||||
case <-timer.C:
|
||||
break batchLoop
|
||||
@@ -239,6 +211,20 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo
|
||||
if ops == nil || entry == nil {
|
||||
return
|
||||
}
|
||||
entry.UserAgent = normalizeOpsPersistentUserAgent(entry.UserAgent)
|
||||
if entry.ErrorBody != "" {
|
||||
originalBody := entry.ErrorBody
|
||||
body, truncated := service.SanitizeOpsErrorBodyForQueue(originalBody)
|
||||
entry.ErrorBody = body
|
||||
if truncated || body != originalBody {
|
||||
opsErrorLogSanitized.Add(1)
|
||||
}
|
||||
}
|
||||
if err := service.SanitizeOpsUpstreamErrorsForQueue(entry); err != nil {
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-opsErrorLogShutdownCh:
|
||||
return
|
||||
@@ -259,18 +245,29 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo
|
||||
if opsErrorLogStopping || opsErrorLogQueue == nil {
|
||||
return
|
||||
}
|
||||
queuedBytes := estimateOpsErrorLogJobBytes(entry)
|
||||
if !reserveOpsErrorLogQueueBytes(queuedBytes) {
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry}:
|
||||
opsErrorLogQueueLen.Add(1)
|
||||
case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry, queuedBytes: queuedBytes}:
|
||||
opsErrorLogEnqueued.Add(1)
|
||||
default:
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-queuedBytes)
|
||||
// Queue is full; drop to avoid blocking request handling.
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeOpsPersistentUserAgent(value string) string {
|
||||
return truncateString(strings.TrimSpace(strings.ToValidUTF8(value, "")), opsErrorLogMaxUserAgentBytes)
|
||||
}
|
||||
|
||||
func StopOpsErrorLogWorkers() bool {
|
||||
opsErrorLogStopOnce.Do(func() {
|
||||
opsErrorLogShutdownOnce.Do(func() {
|
||||
@@ -293,6 +290,7 @@ func stopOpsErrorLogWorkers() bool {
|
||||
|
||||
if ch == nil {
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -305,6 +303,7 @@ func stopOpsErrorLogWorkers() bool {
|
||||
select {
|
||||
case <-done:
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
return true
|
||||
case <-time.After(opsErrorLogDrainTimeout):
|
||||
return false
|
||||
@@ -315,6 +314,14 @@ func OpsErrorLogQueueLength() int64 {
|
||||
return opsErrorLogQueueLen.Load()
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueBytes() int64 {
|
||||
return opsErrorLogQueueBytes.Load()
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueBytesCapacity() int64 {
|
||||
return opsErrorLogMaxQueueBytes
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueCapacity() int {
|
||||
opsErrorLogMu.RLock()
|
||||
ch := opsErrorLogQueue
|
||||
@@ -355,12 +362,15 @@ func maybeLogOpsErrorLogDrop() {
|
||||
}
|
||||
|
||||
queued := opsErrorLogQueueLen.Load()
|
||||
queuedBytes := opsErrorLogQueueBytes.Load()
|
||||
queueCap := OpsErrorLogQueueCapacity()
|
||||
|
||||
log.Printf(
|
||||
"[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)",
|
||||
"[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d queued_bytes=%d bytes_cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)",
|
||||
queued,
|
||||
queueCap,
|
||||
queuedBytes,
|
||||
opsErrorLogMaxQueueBytes,
|
||||
opsErrorLogEnqueued.Load(),
|
||||
opsErrorLogDropped.Load(),
|
||||
opsErrorLogProcessed.Load(),
|
||||
@@ -368,6 +378,46 @@ func maybeLogOpsErrorLogDrop() {
|
||||
)
|
||||
}
|
||||
|
||||
func reserveOpsErrorLogQueueBytes(size int64) bool {
|
||||
if size < 1 {
|
||||
size = 1
|
||||
}
|
||||
for {
|
||||
current := opsErrorLogQueueBytes.Load()
|
||||
if current > opsErrorLogMaxQueueBytes-size {
|
||||
return false
|
||||
}
|
||||
if opsErrorLogQueueBytes.CompareAndSwap(current, current+size) {
|
||||
opsErrorLogQueueLen.Add(1)
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func estimateOpsErrorLogJobBytes(entry *service.OpsInsertErrorLogInput) int64 {
|
||||
if entry == nil {
|
||||
return 1
|
||||
}
|
||||
const fixedOverhead = 512
|
||||
size := fixedOverhead + len(entry.RequestID) + len(entry.ClientRequestID) +
|
||||
len(entry.Platform) + len(entry.Model) + len(entry.RequestPath) +
|
||||
len(entry.InboundEndpoint) + len(entry.UpstreamEndpoint) +
|
||||
len(entry.RequestedModel) + len(entry.UpstreamModel) + len(entry.UserAgent) +
|
||||
len(entry.ErrorPhase) + len(entry.ErrorType) + len(entry.Severity) +
|
||||
len(entry.ErrorMessage) + len(entry.ErrorBody) + len(entry.ErrorSource) +
|
||||
len(entry.ErrorOwner) + len(entry.APIKeyPrefix)
|
||||
if entry.UpstreamErrorMessage != nil {
|
||||
size += len(*entry.UpstreamErrorMessage)
|
||||
}
|
||||
if entry.UpstreamErrorDetail != nil {
|
||||
size += len(*entry.UpstreamErrorDetail)
|
||||
}
|
||||
if entry.UpstreamErrorsJSON != nil {
|
||||
size += len(*entry.UpstreamErrorsJSON)
|
||||
}
|
||||
return int64(size)
|
||||
}
|
||||
|
||||
func opsErrorLogConfig() (workerCount int, queueSize int) {
|
||||
workerCount = runtime.GOMAXPROCS(0) * 2
|
||||
if workerCount < opsErrorLogMinWorkerCount {
|
||||
@@ -470,9 +520,12 @@ type opsCaptureWriter struct {
|
||||
gin.ResponseWriter
|
||||
limit int
|
||||
buf bytes.Buffer
|
||||
ctx *gin.Context
|
||||
}
|
||||
|
||||
const opsCaptureWriterLimit = 64 * 1024
|
||||
const opsCaptureWriterLimit = service.OpsErrorLogQueueBodyMaxBytes
|
||||
|
||||
const opsCaptureWriterPoolMaxRetainedCapacity = service.OpsErrorLogQueueBodyMaxBytes
|
||||
|
||||
var opsCaptureWriterPool = sync.Pool{
|
||||
New: func() any {
|
||||
@@ -496,11 +549,19 @@ func releaseOpsCaptureWriter(w *opsCaptureWriter) {
|
||||
return
|
||||
}
|
||||
w.ResponseWriter = nil
|
||||
w.ctx = nil
|
||||
w.limit = opsCaptureWriterLimit
|
||||
if !shouldPoolOpsCaptureWriter(w) {
|
||||
return
|
||||
}
|
||||
w.buf.Reset()
|
||||
opsCaptureWriterPool.Put(w)
|
||||
}
|
||||
|
||||
func shouldPoolOpsCaptureWriter(w *opsCaptureWriter) bool {
|
||||
return w != nil && w.buf.Cap() <= opsCaptureWriterPoolMaxRetainedCapacity
|
||||
}
|
||||
|
||||
func (w *opsCaptureWriter) Status() int {
|
||||
if w.ResponseWriter == nil {
|
||||
return 0
|
||||
@@ -577,7 +638,7 @@ func (w *opsCaptureWriter) Write(b []byte) (int, error) {
|
||||
if w.ResponseWriter == nil {
|
||||
return 0, nil
|
||||
}
|
||||
if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
remaining := w.limit - w.buf.Len()
|
||||
if len(b) > remaining {
|
||||
_, _ = w.buf.Write(b[:remaining])
|
||||
@@ -592,7 +653,7 @@ func (w *opsCaptureWriter) WriteString(s string) (int, error) {
|
||||
if w.ResponseWriter == nil {
|
||||
return 0, nil
|
||||
}
|
||||
if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
remaining := w.limit - w.buf.Len()
|
||||
if len(s) > remaining {
|
||||
_, _ = w.buf.WriteString(s[:remaining])
|
||||
@@ -603,6 +664,14 @@ func (w *opsCaptureWriter) WriteString(s string) (int, error) {
|
||||
return w.ResponseWriter.WriteString(s)
|
||||
}
|
||||
|
||||
func (w *opsCaptureWriter) shouldCapture() bool {
|
||||
if w.ctx == nil {
|
||||
return true
|
||||
}
|
||||
_, rejected := middleware2.GetIngressRejectReason(w.ctx)
|
||||
return !rejected
|
||||
}
|
||||
|
||||
// OpsErrorLoggerMiddleware records error responses (status >= 400) into ops_error_logs.
|
||||
//
|
||||
// Notes:
|
||||
@@ -612,6 +681,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
originalWriter := c.Writer
|
||||
w := acquireOpsCaptureWriter(originalWriter)
|
||||
w.ctx = c
|
||||
defer func() {
|
||||
// Restore the original writer before returning so outer middlewares
|
||||
// don't observe a pooled wrapper that has been released.
|
||||
@@ -623,6 +693,10 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
c.Writer = w
|
||||
c.Next()
|
||||
|
||||
if _, rejected := middleware2.GetIngressRejectReason(c); rejected {
|
||||
return
|
||||
}
|
||||
|
||||
if ops == nil {
|
||||
return
|
||||
}
|
||||
@@ -989,7 +1063,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
IsCountTokens: isCountTokensRequest(c),
|
||||
|
||||
ErrorMessage: parsed.Message,
|
||||
// Keep the full captured error body (capture is already capped at 64KB) so the
|
||||
// Keep the captured error body (already capped at the queue-safe limit) so the
|
||||
// service layer can sanitize JSON before truncating for storage.
|
||||
ErrorBody: string(body),
|
||||
ErrorSource: errorSource,
|
||||
@@ -1002,7 +1076,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
|
||||
if apiKey != nil {
|
||||
entry.APIKeyID = &apiKey.ID
|
||||
// 有效(未删除)key 报错时快照前缀,key 之后被删也保留;与 INVALID_API_KEY 的 attempted_key_prefix 互斥。
|
||||
// 有效 key 报错时快照前缀,key 之后被删也保留。
|
||||
entry.APIKeyPrefix = keyPrefix(apiKey.Key, 8)
|
||||
if apiKey.User != nil {
|
||||
entry.UserID = &apiKey.User.ID
|
||||
@@ -1022,22 +1096,6 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
entry.ClientIP = &clientIP
|
||||
}
|
||||
|
||||
// 已删除 key 归因:仅 INVALID_API_KEY 才尝试。响应已写出,此处不阻塞客户端。
|
||||
if parsed.Code == opsCodeInvalidAPIKey {
|
||||
if attemptedKey := extractAttemptedKey(c); attemptedKey != "" {
|
||||
entry.AttemptedKeyPrefix = keyPrefix(attemptedKey, 8)
|
||||
if looksLikeSystemKey(attemptedKey) {
|
||||
if res, lookupErr := ops.LookupDeletedKeyAudit(c.Request.Context(), attemptedKey); lookupErr != nil {
|
||||
log.Printf("[OpsErrorLogger] LookupDeletedKeyAudit failed: %v", lookupErr)
|
||||
} else if res != nil {
|
||||
owner := res.UserID
|
||||
entry.DeletedKeyOwnerUserID = &owner
|
||||
entry.DeletedKeyName = res.KeyName
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enqueueOpsErrorLog(ops, entry)
|
||||
}
|
||||
}
|
||||
@@ -1072,8 +1130,20 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int)
|
||||
if classifyStatus <= 0 {
|
||||
classifyStatus = wireStatus
|
||||
}
|
||||
normalizedType := normalizeOpsErrorType(streamErr.ErrType, "")
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, "", classifyStatus)
|
||||
normalizedType := normalizeOpsErrorType(streamErr.ErrType, streamErr.Code)
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, streamErr.Code, classifyStatus)
|
||||
recordedStatus := wireStatus
|
||||
if streamErr.CountTowardsSLA && streamErr.IntendedStatus >= 400 {
|
||||
recordedStatus = streamErr.IntendedStatus
|
||||
}
|
||||
errorBody := ""
|
||||
if streamErr.Code != "" {
|
||||
if payload, err := json.Marshal(gin.H{"error": gin.H{
|
||||
"type": normalizedType, "code": streamErr.Code, "message": streamErr.Message,
|
||||
}}); err == nil {
|
||||
errorBody = string(payload)
|
||||
}
|
||||
}
|
||||
|
||||
apiKey := getOpsAPIKey(c)
|
||||
clientRequestID, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
@@ -1140,12 +1210,12 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int)
|
||||
ErrorPhase: phase,
|
||||
ErrorType: normalizedType,
|
||||
Severity: classifyOpsSeverity(normalizedType, classifyStatus),
|
||||
StatusCode: wireStatus,
|
||||
StatusCode: recordedStatus,
|
||||
IsBusinessLimited: isBusinessLimited,
|
||||
IsCountTokens: isCountTokensRequest(c),
|
||||
|
||||
ErrorMessage: streamErr.Message,
|
||||
ErrorBody: "",
|
||||
ErrorBody: errorBody,
|
||||
ErrorSource: errorSource,
|
||||
ErrorOwner: errorOwner,
|
||||
|
||||
@@ -1694,11 +1764,8 @@ func shouldSkipOpsErrorLog(ctx context.Context, ops *service.OpsService, message
|
||||
}
|
||||
|
||||
// Get advanced settings to check filter configuration
|
||||
settings, err := ops.GetOpsAdvancedSettings(ctx)
|
||||
if err != nil || settings == nil {
|
||||
// If we can't get settings, don't skip (fail open)
|
||||
return false
|
||||
}
|
||||
_ = ctx
|
||||
settings := ops.OpsAdvancedSettingsSnapshot()
|
||||
|
||||
msgLower := strings.ToLower(message)
|
||||
bodyLower := strings.ToLower(body)
|
||||
|
||||
@@ -1,39 +1,9 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestLooksLikeSystemKey(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want bool
|
||||
}{
|
||||
{"sk-abcdef0123456789", true},
|
||||
{"ABCdef_-0123456789", true},
|
||||
{"short", false},
|
||||
{"with space xxxxxxxxxx", false},
|
||||
{"汉字key1234567890", false},
|
||||
{"", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := looksLikeSystemKey(c.in); got != c.want {
|
||||
t.Errorf("looksLikeSystemKey(%q)=%v want %v", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
long := make([]byte, 129)
|
||||
for i := range long {
|
||||
long[i] = 'a'
|
||||
}
|
||||
if looksLikeSystemKey(string(long)) {
|
||||
t.Errorf("129-char key should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPrefix(t *testing.T) {
|
||||
if got := keyPrefix("sk-3f2a9c7e", 8); got != "sk-3f2a9" {
|
||||
t.Errorf("keyPrefix=%q want %q", got, "sk-3f2a9")
|
||||
@@ -42,77 +12,3 @@ func TestKeyPrefix(t *testing.T) {
|
||||
t.Errorf("short key should be returned as-is, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractAttemptedKey(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
headers map[string]string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "Bearer in Authorization",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-testkey0123456789"},
|
||||
want: "sk-testkey0123456789",
|
||||
},
|
||||
{
|
||||
name: "Bearer case-insensitive",
|
||||
headers: map[string]string{"Authorization": "BEARER sk-testkey0123456789"},
|
||||
want: "sk-testkey0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-api-key header",
|
||||
headers: map[string]string{"x-api-key": "sk-xapikey0123456789"},
|
||||
want: "sk-xapikey0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-goog-api-key header",
|
||||
headers: map[string]string{"x-goog-api-key": "sk-goog0123456789"},
|
||||
want: "sk-goog0123456789",
|
||||
},
|
||||
{
|
||||
name: "Authorization takes priority over x-api-key",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-auth0123456789", "x-api-key": "sk-xapi0123456789"},
|
||||
want: "sk-auth0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-api-key takes priority over x-goog-api-key",
|
||||
headers: map[string]string{"x-api-key": "sk-xapi0123456789", "x-goog-api-key": "sk-goog0123456789"},
|
||||
want: "sk-xapi0123456789",
|
||||
},
|
||||
{
|
||||
name: "no key headers",
|
||||
headers: map[string]string{},
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "Bearer with leading/trailing spaces trimmed",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-trimmed0123456789 "},
|
||||
want: "sk-trimmed0123456789",
|
||||
},
|
||||
{
|
||||
// 非 Bearer Authorization 应被忽略,继续 fall-through 到 x-api-key(与认证中间件一致)
|
||||
name: "non-Bearer Authorization falls through to x-api-key",
|
||||
headers: map[string]string{"Authorization": "junk-not-bearer", "x-api-key": "sk-realkey1234567"},
|
||||
want: "sk-realkey1234567",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
for k, v := range tc.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
c.Request = req
|
||||
|
||||
got := extractAttemptedKey(c)
|
||||
if got != tc.want {
|
||||
t.Errorf("extractAttemptedKey(%v) = %q, want %q", tc.headers, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -12,6 +15,82 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type ingressRejectSettingRepo struct {
|
||||
service.SettingRepository
|
||||
getValueCalls int
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) GetValue(context.Context, string) (string, error) {
|
||||
r.getValueCalls++
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
r.getValueCalls++
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) Set(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type ingressRejectOpsRepo struct {
|
||||
service.OpsRepository
|
||||
insertCalls int
|
||||
}
|
||||
|
||||
func (r *ingressRejectOpsRepo) InsertErrorLog(context.Context, *service.OpsInsertErrorLogInput) (int64, error) {
|
||||
r.insertCalls++
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (r *ingressRejectOpsRepo) BatchInsertErrorLogs(context.Context, []*service.OpsInsertErrorLogInput) (int64, error) {
|
||||
r.insertCalls++
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func TestOpsErrorLogQueueByteBudget(t *testing.T) {
|
||||
previousBytes := opsErrorLogQueueBytes.Load()
|
||||
previousLen := opsErrorLogQueueLen.Load()
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
t.Cleanup(func() {
|
||||
opsErrorLogQueueBytes.Store(previousBytes)
|
||||
opsErrorLogQueueLen.Store(previousLen)
|
||||
})
|
||||
|
||||
if !reserveOpsErrorLogQueueBytes(opsErrorLogMaxQueueBytes - 1) {
|
||||
t.Fatal("first reservation within byte budget should succeed")
|
||||
}
|
||||
if reserveOpsErrorLogQueueBytes(2) {
|
||||
t.Fatal("reservation beyond byte budget should be rejected")
|
||||
}
|
||||
if got := OpsErrorLogQueueBytes(); got != opsErrorLogMaxQueueBytes-1 {
|
||||
t.Fatalf("queued bytes = %d, want %d", got, opsErrorLogMaxQueueBytes-1)
|
||||
}
|
||||
if got := OpsErrorLogQueueLength(); got != 1 {
|
||||
t.Fatalf("queue length = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateOpsErrorLogJobBytesIncludesVariablePayloads(t *testing.T) {
|
||||
base := estimateOpsErrorLogJobBytes(&service.OpsInsertErrorLogInput{})
|
||||
message := "upstream message"
|
||||
detail := "upstream detail"
|
||||
events := `[{"error":"x"}]`
|
||||
entry := &service.OpsInsertErrorLogInput{
|
||||
ErrorBody: strings.Repeat("x", 1024),
|
||||
ErrorMessage: "client error",
|
||||
UserAgent: "test-agent",
|
||||
UpstreamErrorMessage: &message,
|
||||
UpstreamErrorDetail: &detail,
|
||||
UpstreamErrorsJSON: &events,
|
||||
}
|
||||
if got := estimateOpsErrorLogJobBytes(entry); got <= base+1024 {
|
||||
t.Fatalf("estimated bytes = %d, expected variable payloads above %d", got, base+1024)
|
||||
}
|
||||
}
|
||||
|
||||
func resetOpsErrorLoggerStateForTest(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
@@ -119,6 +198,29 @@ func TestOpsCaptureWriterPool_ResetOnRelease(t *testing.T) {
|
||||
require.Zero(t, reused.buf.Len(), "writer should be reset before reuse")
|
||||
}
|
||||
|
||||
func TestOpsCaptureWriterPool_DropsLargeBuffers(t *testing.T) {
|
||||
w := &opsCaptureWriter{}
|
||||
w.buf.Grow(opsCaptureWriterPoolMaxRetainedCapacity + 1)
|
||||
require.False(t, shouldPoolOpsCaptureWriter(w))
|
||||
}
|
||||
|
||||
func TestEnqueueOpsErrorLog_SanitizesAndBoundsBodyBeforeQueue(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 1)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
secret := strings.Repeat("s", service.OpsErrorLogQueueBodyMaxBytes)
|
||||
entry := &service.OpsInsertErrorLogInput{
|
||||
ErrorPhase: "request",
|
||||
ErrorType: "api_error",
|
||||
ErrorBody: `{"authorization":"Bearer ` + secret + `","message":"failed"}`,
|
||||
}
|
||||
|
||||
enqueueOpsErrorLog(ops, entry)
|
||||
job := <-opsErrorLogQueue
|
||||
require.LessOrEqual(t, len(job.entry.ErrorBody), service.OpsErrorLogQueueBodyMaxBytes)
|
||||
require.NotContains(t, job.entry.ErrorBody, secret)
|
||||
require.Equal(t, int64(1), OpsErrorLogSanitizedTotal())
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_DoesNotBreakOuterMiddlewares(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -149,6 +251,42 @@ func setupOpsErrorLogTestQueue(t *testing.T, size int) {
|
||||
opsErrorLogMu.Unlock()
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_HardSkipsIngressRejection(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
settings := &ingressRejectSettingRepo{}
|
||||
repo := &ingressRejectOpsRepo{}
|
||||
ops := service.NewOpsService(repo, settings, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
// Construction may read unrelated runtime settings; only request-path reads matter here.
|
||||
settings.getValueCalls = 0
|
||||
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.GET("/v1/messages", func(c *gin.Context) {
|
||||
middleware2.MarkIngressRejected(c, middleware2.IngressRejectInvalidAPIKey)
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"code": "INVALID_API_KEY", "message": "Invalid API key"})
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.JSONEq(t, `{"code":"INVALID_API_KEY","message":"Invalid API key"}`, w.Body.String())
|
||||
require.Zero(t, settings.getValueCalls, "ingress rejection must bypass monitoring settings reads")
|
||||
require.Zero(t, repo.insertCalls, "ingress rejection must bypass inserts")
|
||||
require.Zero(t, OpsErrorLogEnqueuedTotal(), "ingress rejection must not enter the error queue")
|
||||
}
|
||||
|
||||
func TestNormalizeOpsPersistentUserAgentBoundsAndPreservesUTF8(t *testing.T) {
|
||||
value := strings.Repeat("a", opsErrorLogMaxUserAgentBytes-1) + "你" + strings.Repeat("b", 32)
|
||||
got := normalizeOpsPersistentUserAgent(" " + value + " ")
|
||||
require.LessOrEqual(t, len(got), opsErrorLogMaxUserAgentBytes)
|
||||
require.True(t, utf8.ValidString(got))
|
||||
require.NotContains(t, got, "b")
|
||||
}
|
||||
|
||||
// 就地(in-band) SSE 错误挂在已固化的 HTTP 200 流上:wire 状态码为 200,
|
||||
// 常规 status>=400 采集路径不会触发。logOpsStreamError 必须据 MarkOpsStreamError
|
||||
// 补记一条错误日志,且用 IntendedStatus(429) 分级、StatusCode 仍记 wire 的 200。
|
||||
@@ -182,6 +320,36 @@ func TestLogOpsStreamError_RecordsInBandConcurrencyLimit(t *testing.T) {
|
||||
require.Equal(t, "Concurrency limit exceeded for account, please retry later", job.entry.ErrorMessage)
|
||||
}
|
||||
|
||||
func TestLogOpsStreamError_UpstreamFailureCountsTowardsSLA(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
c.Set(opsModelKey, "gpt-5.6-sol")
|
||||
|
||||
service.MarkOpsStreamFailure(
|
||||
c,
|
||||
"upstream_error",
|
||||
service.OpenAIUpstreamHTTP2StreamErrorCode,
|
||||
"Upstream HTTP/2 stream failed",
|
||||
http.StatusBadGateway,
|
||||
)
|
||||
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
logOpsStreamError(c, ops, http.StatusOK)
|
||||
|
||||
job := <-opsErrorLogQueue
|
||||
require.NotNil(t, job.entry)
|
||||
require.Equal(t, http.StatusBadGateway, job.entry.StatusCode)
|
||||
require.Equal(t, "upstream_error", job.entry.ErrorType)
|
||||
require.Equal(t, "upstream", job.entry.ErrorPhase)
|
||||
require.Equal(t, "provider", job.entry.ErrorOwner)
|
||||
require.False(t, job.entry.IsBusinessLimited)
|
||||
require.Contains(t, job.entry.ErrorBody, service.OpenAIUpstreamHTTP2StreamErrorCode)
|
||||
}
|
||||
|
||||
// 未标记流内错误时 logOpsStreamError 必须是 no-op(不误记正常的 200 流)。
|
||||
func TestLogOpsStreamError_NoopWhenNotMarked(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpsCaptureWriterDoesNotCopyIngressRejectBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
writer := acquireOpsCaptureWriter(context.Writer)
|
||||
defer releaseOpsCaptureWriter(writer)
|
||||
writer.ctx = context
|
||||
context.Writer = writer
|
||||
middleware2.MarkIngressRejected(context, middleware2.IngressRejectInvalidAPIKey)
|
||||
context.Status(http.StatusUnauthorized)
|
||||
_, err := context.Writer.WriteString(`{"code":"INVALID_API_KEY","message":"Invalid API key"}`)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, writer.buf.Len())
|
||||
}
|
||||
@@ -120,12 +120,14 @@ func ProvideOpenAIGatewayHandler(
|
||||
errorPassthroughService *service.ErrorPassthroughService,
|
||||
contentModerationService *service.ContentModerationService,
|
||||
opsService *service.OpsService,
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
) *OpenAIGatewayHandler {
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
}
|
||||
|
||||
@@ -153,9 +155,10 @@ func ProvideSettingHandler(settingService *service.SettingService, buildInfo Bui
|
||||
}
|
||||
|
||||
// ProvideAdminSettingHandler creates admin.SettingHandler with notification template APIs.
|
||||
func ProvideAdminSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService, paymentConfigService *service.PaymentConfigService, paymentService *service.PaymentService, userAttributeService *service.UserAttributeService, notificationEmailService *service.NotificationEmailService) *admin.SettingHandler {
|
||||
func ProvideAdminSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService, paymentConfigService *service.PaymentConfigService, paymentService *service.PaymentService, userAttributeService *service.UserAttributeService, notificationEmailService *service.NotificationEmailService, totpService *service.TotpService, userService *service.UserService) *admin.SettingHandler {
|
||||
h := admin.NewSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService)
|
||||
h.SetNotificationEmailService(notificationEmailService)
|
||||
h.SetStepUpDeps(totpService, userService)
|
||||
return h
|
||||
}
|
||||
|
||||
|
||||
@@ -146,7 +146,7 @@ func TestAnthropicToResponses_ToolUse(t *testing.T) {
|
||||
assert.Equal(t, "Sunny, 72°F", items[3].Output)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ThinkingIgnored(t *testing.T) {
|
||||
func TestAnthropicToResponses_ThinkingWithoutSignatureIgnored(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.2",
|
||||
MaxTokens: 1024,
|
||||
@@ -162,7 +162,7 @@ func TestAnthropicToResponses_ThinkingIgnored(t *testing.T) {
|
||||
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(resp.Input, &items))
|
||||
// user + assistant(text only, thinking ignored) + user = 3
|
||||
// user + assistant(text only, thinking without signature ignored) + user = 3
|
||||
require.Len(t, items, 3)
|
||||
assert.Equal(t, "assistant", items[1].Role)
|
||||
// Assistant content should only have text, not thinking.
|
||||
@@ -173,6 +173,30 @@ func TestAnthropicToResponses_ThinkingIgnored(t *testing.T) {
|
||||
assert.Equal(t, "Hi!", parts[0].Text)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_ThinkingSignatureBecomesReasoning(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "grok-4.5",
|
||||
MaxTokens: 1024,
|
||||
Messages: []AnthropicMessage{
|
||||
{Role: "user", Content: json.RawMessage(`"Hello"`)},
|
||||
{Role: "assistant", Content: json.RawMessage(`[{"type":"thinking","thinking":"plan","signature":"enc-rs-1"},{"type":"text","text":"Hi!"},{"type":"tool_use","id":"toolu_1","name":"Bash","input":{"command":"ls"}}]`)},
|
||||
{Role: "user", Content: json.RawMessage(`[{"type":"tool_result","tool_use_id":"toolu_1","content":"ok"}]`)},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := AnthropicToResponses(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
var items []ResponsesInputItem
|
||||
require.NoError(t, json.Unmarshal(resp.Input, &items))
|
||||
// user + reasoning + assistant text + function_call + function_call_output
|
||||
require.GreaterOrEqual(t, len(items), 4)
|
||||
assert.Equal(t, "reasoning", items[1].Type)
|
||||
assert.Equal(t, "enc-rs-1", items[1].EncryptedContent)
|
||||
assert.Equal(t, "assistant", items[2].Role)
|
||||
assert.Equal(t, "function_call", items[3].Type)
|
||||
}
|
||||
|
||||
func TestAnthropicToResponses_MaxTokensFloor(t *testing.T) {
|
||||
req := &AnthropicRequest{
|
||||
Model: "gpt-5.2",
|
||||
@@ -208,7 +232,7 @@ func TestResponsesToAnthropic_TextOnly(t *testing.T) {
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "resp_123", anth.ID)
|
||||
assert.Equal(t, "claude-opus-4-6", anth.Model)
|
||||
assert.Equal(t, "end_turn", anth.StopReason)
|
||||
assert.Equal(t, "end_turn", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 1)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
assert.Equal(t, "Hello there!", anth.Content[0].Text)
|
||||
@@ -287,7 +311,7 @@ func TestResponsesToAnthropic_ToolUse(t *testing.T) {
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "tool_use", anth.StopReason)
|
||||
assert.Equal(t, "tool_use", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 2)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
assert.Equal(t, "tool_use", anth.Content[1].Type)
|
||||
@@ -318,7 +342,7 @@ func TestResponsesToAnthropic_ToolUseStopReasonDoesNotDependOnLastBlock(t *testi
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "tool_use", anth.StopReason)
|
||||
assert.Equal(t, "tool_use", AnthropicStopReasonString(anth.StopReason))
|
||||
require.Len(t, anth.Content, 2)
|
||||
assert.Equal(t, "tool_use", anth.Content[0].Type)
|
||||
assert.Equal(t, "text", anth.Content[1].Type)
|
||||
@@ -372,7 +396,8 @@ func TestResponsesToAnthropic_Reasoning(t *testing.T) {
|
||||
Status: "completed",
|
||||
Output: []ResponsesOutput{
|
||||
{
|
||||
Type: "reasoning",
|
||||
Type: "reasoning",
|
||||
EncryptedContent: "enc-rs-roundtrip",
|
||||
Summary: []ResponsesSummary{
|
||||
{Type: "summary_text", Text: "Thinking about the answer..."},
|
||||
},
|
||||
@@ -390,10 +415,57 @@ func TestResponsesToAnthropic_Reasoning(t *testing.T) {
|
||||
require.Len(t, anth.Content, 2)
|
||||
assert.Equal(t, "thinking", anth.Content[0].Type)
|
||||
assert.Equal(t, "Thinking about the answer...", anth.Content[0].Thinking)
|
||||
assert.Equal(t, "enc-rs-roundtrip", anth.Content[0].Signature)
|
||||
assert.Equal(t, "text", anth.Content[1].Type)
|
||||
assert.Equal(t, "42", anth.Content[1].Text)
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_StreamEmitsThinkingSignature(t *testing.T) {
|
||||
state := NewResponsesEventToAnthropicState()
|
||||
var all []AnthropicStreamEvent
|
||||
|
||||
appendAll := func(events []AnthropicStreamEvent) {
|
||||
all = append(all, events...)
|
||||
}
|
||||
|
||||
appendAll(ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.added",
|
||||
OutputIndex: 0,
|
||||
Item: &ResponsesOutput{Type: "reasoning", ID: "rs_1"},
|
||||
}, state))
|
||||
appendAll(ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.reasoning_summary_text.delta",
|
||||
OutputIndex: 0,
|
||||
Delta: "thinking...",
|
||||
SummaryIndex: 0,
|
||||
}, state))
|
||||
// summary.done must not close the thinking block before encrypted_content arrives
|
||||
appendAll(ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.reasoning_summary_text.done",
|
||||
OutputIndex: 0,
|
||||
SummaryIndex: 0,
|
||||
}, state))
|
||||
appendAll(ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.done",
|
||||
OutputIndex: 0,
|
||||
Item: &ResponsesOutput{
|
||||
Type: "reasoning",
|
||||
ID: "rs_1",
|
||||
EncryptedContent: "enc-stream-1",
|
||||
Status: "completed",
|
||||
},
|
||||
}, state))
|
||||
|
||||
var sawSignature bool
|
||||
for _, ev := range all {
|
||||
if ev.Type == "content_block_delta" && ev.Delta != nil && ev.Delta.Type == "signature_delta" {
|
||||
assert.Equal(t, "enc-stream-1", ev.Delta.Signature)
|
||||
sawSignature = true
|
||||
}
|
||||
}
|
||||
require.True(t, sawSignature, "expected signature_delta with encrypted_content")
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_Incomplete(t *testing.T) {
|
||||
resp := &ResponsesResponse{
|
||||
ID: "resp_inc",
|
||||
@@ -411,7 +483,7 @@ func TestResponsesToAnthropic_Incomplete(t *testing.T) {
|
||||
}
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
assert.Equal(t, "max_tokens", anth.StopReason)
|
||||
assert.Equal(t, "max_tokens", AnthropicStopReasonString(anth.StopReason))
|
||||
}
|
||||
|
||||
func TestResponsesToAnthropic_EmptyOutput(t *testing.T) {
|
||||
@@ -785,12 +857,26 @@ func TestStreamingReasoning(t *testing.T) {
|
||||
assert.Equal(t, "thinking_delta", events[0].Delta.Type)
|
||||
assert.Equal(t, "Let me think...", events[0].Delta.Thinking)
|
||||
|
||||
// reasoning done
|
||||
// summary.done keeps thinking open until output_item.done (for signature)
|
||||
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.reasoning_summary_text.done",
|
||||
}, state)
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, "content_block_stop", events[0].Type)
|
||||
require.Len(t, events, 0)
|
||||
|
||||
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.done",
|
||||
OutputIndex: 0,
|
||||
Item: &ResponsesOutput{
|
||||
Type: "reasoning",
|
||||
EncryptedContent: "enc-rs-stream",
|
||||
Status: "completed",
|
||||
},
|
||||
}, state)
|
||||
require.Len(t, events, 2)
|
||||
assert.Equal(t, "content_block_delta", events[0].Type)
|
||||
assert.Equal(t, "signature_delta", events[0].Delta.Type)
|
||||
assert.Equal(t, "enc-rs-stream", events[0].Delta.Signature)
|
||||
assert.Equal(t, "content_block_stop", events[1].Type)
|
||||
}
|
||||
|
||||
func TestStreamingIncomplete(t *testing.T) {
|
||||
@@ -995,7 +1081,7 @@ func TestResponsesToAnthropic_Failed(t *testing.T) {
|
||||
|
||||
anth := ResponsesToAnthropic(resp, "claude-opus-4-6")
|
||||
// Failed status defaults to "end_turn" stop reason
|
||||
assert.Equal(t, "end_turn", anth.StopReason)
|
||||
assert.Equal(t, "end_turn", AnthropicStopReasonString(anth.StopReason))
|
||||
// Should have at least an empty text block
|
||||
require.Len(t, anth.Content, 1)
|
||||
assert.Equal(t, "text", anth.Content[0].Type)
|
||||
@@ -1609,7 +1695,7 @@ func TestAnthropicToResponsesResponse_CacheTokensUseOpenAIInputSemantics(t *test
|
||||
Content: []AnthropicContentBlock{
|
||||
{Type: "text", Text: "ok"},
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 3318,
|
||||
OutputTokens: 123,
|
||||
@@ -1635,7 +1721,7 @@ func TestAnthropicToResponsesResponse_NoCacheTokens(t *testing.T) {
|
||||
Content: []AnthropicContentBlock{
|
||||
{Type: "text", Text: "ok"},
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
@@ -1732,3 +1818,25 @@ func TestAnthropicEventToResponses_CacheTokensFromMessageDelta(t *testing.T) {
|
||||
require.NotNil(t, completed.Response.Usage.InputTokensDetails)
|
||||
assert.Equal(t, 11, completed.Response.Usage.InputTokensDetails.CachedTokens)
|
||||
}
|
||||
|
||||
func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
state := NewResponsesEventToAnthropicState()
|
||||
state.Model = "grok-4.5"
|
||||
events := ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
|
||||
Type: "response.created",
|
||||
Response: &ResponsesResponse{
|
||||
ID: "resp_test",
|
||||
Model: "grok-4.5",
|
||||
},
|
||||
}, state)
|
||||
require.Len(t, events, 1)
|
||||
require.Equal(t, "message_start", events[0].Type)
|
||||
require.NotNil(t, events[0].Message)
|
||||
require.Nil(t, events[0].Message.StopReason, "message_start must use null stop_reason, not empty string")
|
||||
|
||||
sse, err := ResponsesAnthropicEventToSSE(events[0])
|
||||
require.NoError(t, err)
|
||||
// Official Anthropic wire: "stop_reason":null
|
||||
require.Contains(t, sse, `"stop_reason":null`)
|
||||
require.NotContains(t, sse, `"stop_reason":""`)
|
||||
}
|
||||
|
||||
@@ -257,7 +257,9 @@ func anthropicUserToResponses(raw json.RawMessage) ([]ResponsesInputItem, error)
|
||||
// anthropicAssistantToResponses handles an Anthropic assistant message.
|
||||
// Text content → assistant message with output_text parts.
|
||||
// tool_use blocks → function_call items.
|
||||
// thinking blocks → ignored (OpenAI doesn't accept them as input).
|
||||
// thinking blocks with signature → reasoning items (encrypted_content) so
|
||||
// multi-turn Grok/Codex prompt cache can reuse prior reasoning prefixes.
|
||||
// thinking without signature remains ignored (not accepted as plain text input).
|
||||
func anthropicAssistantToResponses(raw json.RawMessage) ([]ResponsesInputItem, error) {
|
||||
// Try plain string.
|
||||
var s string
|
||||
@@ -277,6 +279,25 @@ func anthropicAssistantToResponses(raw json.RawMessage) ([]ResponsesInputItem, e
|
||||
|
||||
var items []ResponsesInputItem
|
||||
|
||||
// Preserve turn order: reasoning → assistant text → tool calls. xAI/Codex
|
||||
// multi-turn cache and tool continuations expect reasoning before the
|
||||
// assistant message that followed it.
|
||||
for _, b := range blocks {
|
||||
if b.Type != "thinking" {
|
||||
continue
|
||||
}
|
||||
sig := strings.TrimSpace(b.Signature)
|
||||
// Only replay provider ciphertext. Skip GPT/Codex-style gAAAA blobs and
|
||||
// empty placeholders — xAI returns 400 on decrypt for foreign signatures.
|
||||
if sig == "" || strings.HasPrefix(sig, "gAAAA") {
|
||||
continue
|
||||
}
|
||||
items = append(items, ResponsesInputItem{
|
||||
Type: "reasoning",
|
||||
EncryptedContent: sig,
|
||||
})
|
||||
}
|
||||
|
||||
// Text content → assistant message with output_text content parts.
|
||||
text := extractAnthropicTextFromBlocks(blocks)
|
||||
if text != "" {
|
||||
|
||||
@@ -89,7 +89,7 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
|
||||
out.Output = outputs
|
||||
|
||||
// Map stop_reason → status
|
||||
out.Status = anthropicStopReasonToResponsesStatus(resp.StopReason, resp.Content)
|
||||
out.Status = anthropicStopReasonToResponsesStatus(AnthropicStopReasonString(resp.StopReason), resp.Content)
|
||||
if out.Status == "incomplete" {
|
||||
out.IncompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
|
||||
}
|
||||
@@ -152,11 +152,25 @@ type AnthropicEventToResponsesState struct {
|
||||
|
||||
// For message output: accumulate text parts
|
||||
ContentIndex int
|
||||
// TextAccum accumulates the current text part so that output_text.done and
|
||||
// content_part.done can carry the full text (deltas carry increments only).
|
||||
TextAccum string
|
||||
|
||||
// For function_call: track per-output info
|
||||
CurrentCallID string
|
||||
CurrentName string
|
||||
|
||||
// Content of the currently open item, folded into Outputs when it closes.
|
||||
CurrentContent []ResponsesContentPart // message
|
||||
CurrentArgs string // function_call
|
||||
CurrentSummary string // reasoning
|
||||
|
||||
// Outputs accumulates every closed output item so that response.completed
|
||||
// can carry the full output list. The OpenAI SDK's get_final_response()
|
||||
// parses the terminal event's response directly; without this, clients see
|
||||
// an empty output_text.
|
||||
Outputs []ResponsesOutput
|
||||
|
||||
// Usage from message_start / message_delta. InputTokens here follows
|
||||
// Anthropic semantics (excludes cached tokens); they are added back when
|
||||
// emitting the OpenAI Responses usage.
|
||||
@@ -293,6 +307,22 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
}))
|
||||
}
|
||||
|
||||
// response.content_part.added must precede the output_text.delta events
|
||||
// for that part. The message item is added with content: [], and the
|
||||
// OpenAI SDK's accumulating stream helper (client.responses.stream) only
|
||||
// appends a content part when it sees content_part.added. Without it the
|
||||
// following output_text.delta indexes output.content[content_index] and
|
||||
// raises IndexError. Raw event iteration
|
||||
// (responses.create(stream=True)) does not accumulate, which is why this
|
||||
// went unnoticed.
|
||||
events = append(events, makeResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Part: &ResponsesContentPart{Type: "output_text", Text: ""},
|
||||
}))
|
||||
state.TextAccum = ""
|
||||
|
||||
case "tool_use":
|
||||
// Close previous item if any
|
||||
events = append(events, closeCurrentResponsesItem(state)...)
|
||||
@@ -327,6 +357,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.Text == "" {
|
||||
return nil
|
||||
}
|
||||
state.TextAccum += evt.Delta.Text
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
@@ -338,6 +369,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.Thinking == "" {
|
||||
return nil
|
||||
}
|
||||
state.CurrentSummary += evt.Delta.Thinking
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
SummaryIndex: 0,
|
||||
@@ -349,6 +381,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.PartialJSON == "" {
|
||||
return nil
|
||||
}
|
||||
state.CurrentArgs += evt.Delta.PartialJSON
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
Delta: evt.Delta.PartialJSON,
|
||||
@@ -393,12 +426,24 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
|
||||
return events
|
||||
|
||||
case "message":
|
||||
// Emit output_text.done (text block is done, but message item stays open for potential more blocks)
|
||||
// Text block is done: emit output_text.done then content_part.done (the
|
||||
// order OpenAI uses), both carrying the part's full text. The message
|
||||
// item itself stays open since more blocks may follow.
|
||||
text := state.TextAccum
|
||||
state.TextAccum = ""
|
||||
state.CurrentContent = append(state.CurrentContent, ResponsesContentPart{Type: "output_text", Text: text})
|
||||
return []ResponsesStreamEvent{
|
||||
makeResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Text: text,
|
||||
}),
|
||||
makeResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Part: &ResponsesContentPart{Type: "output_text", Text: text},
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -454,24 +499,48 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
return nil
|
||||
}
|
||||
|
||||
itemType := state.CurrentItemType
|
||||
itemID := state.CurrentItemID
|
||||
// Assemble the full item: both output_item.done and response.completed must
|
||||
// carry its content. Emitting only {type,id,status} makes SDK-side
|
||||
// accumulation produce an empty output.
|
||||
item := ResponsesOutput{
|
||||
Type: state.CurrentItemType,
|
||||
ID: state.CurrentItemID,
|
||||
Status: "completed",
|
||||
}
|
||||
switch state.CurrentItemType {
|
||||
case "message":
|
||||
item.Role = "assistant"
|
||||
item.Content = state.CurrentContent
|
||||
case "function_call":
|
||||
item.CallID = state.CurrentCallID
|
||||
item.Name = state.CurrentName
|
||||
args := state.CurrentArgs
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
item.Arguments = args
|
||||
case "reasoning":
|
||||
if state.CurrentSummary != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: state.CurrentSummary}}
|
||||
}
|
||||
}
|
||||
state.Outputs = append(state.Outputs, item)
|
||||
|
||||
// Reset
|
||||
state.CurrentItemType = ""
|
||||
state.CurrentItemID = ""
|
||||
state.CurrentCallID = ""
|
||||
state.CurrentName = ""
|
||||
state.CurrentContent = nil
|
||||
state.CurrentArgs = ""
|
||||
state.CurrentSummary = ""
|
||||
state.TextAccum = ""
|
||||
state.OutputIndex++
|
||||
state.ContentIndex = 0
|
||||
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex - 1, // Use the index before increment
|
||||
Item: &ResponsesOutput{
|
||||
Type: itemType,
|
||||
ID: itemID,
|
||||
Status: "completed",
|
||||
},
|
||||
Item: &item,
|
||||
})}
|
||||
}
|
||||
|
||||
@@ -519,6 +588,14 @@ func makeResponsesCompletedEvent(
|
||||
eventType = "response.incomplete"
|
||||
}
|
||||
|
||||
// Carry the output items accumulated over the stream. The SDK's
|
||||
// get_final_response() reads them straight from the terminal event, so an
|
||||
// empty list leaves clients with an empty result.
|
||||
outputs := state.Outputs
|
||||
if outputs == nil {
|
||||
outputs = []ResponsesOutput{}
|
||||
}
|
||||
|
||||
return ResponsesStreamEvent{
|
||||
Type: eventType,
|
||||
SequenceNumber: seq,
|
||||
@@ -527,7 +604,7 @@ func makeResponsesCompletedEvent(
|
||||
Object: "response",
|
||||
Model: state.Model,
|
||||
Status: status,
|
||||
Output: []ResponsesOutput{},
|
||||
Output: outputs,
|
||||
Usage: usage,
|
||||
IncompleteDetails: incompleteDetails,
|
||||
},
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
package apicompat
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestAnthropicEventToResponses_TextEmitsContentPart pins that a message text
|
||||
// stream emits response.content_part.added, and that it precedes the first
|
||||
// output_text.delta for that part.
|
||||
//
|
||||
// Why: the OpenAI SDK's accumulating stream helper (client.responses.stream)
|
||||
// only appends a content part to the message item when it sees
|
||||
// content_part.added. The item is added with content: [], so a missing event
|
||||
// makes the following output_text.delta index output.content[content_index] and
|
||||
// raise IndexError. Raw event iteration does not accumulate, so a regression
|
||||
// here is easy to miss.
|
||||
func TestAnthropicEventToResponses_TextEmitsContentPart(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var types []string
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
for _, out := range AnthropicEventToResponsesEvents(evt, state) {
|
||||
types = append(types, out.Type)
|
||||
}
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1", Model: "claude-sonnet-4-5"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hel"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "lo"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
posOf := func(target string) int {
|
||||
for i, ty := range types {
|
||||
if ty == target {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
partAdded := posOf("response.content_part.added")
|
||||
firstDelta := posOf("response.output_text.delta")
|
||||
|
||||
if partAdded < 0 {
|
||||
t.Fatalf("response.content_part.added was not emitted; got %v", types)
|
||||
}
|
||||
if firstDelta < 0 {
|
||||
t.Fatalf("response.output_text.delta was not emitted; got %v", types)
|
||||
}
|
||||
if partAdded > firstDelta {
|
||||
t.Errorf("content_part.added must precede the first output_text.delta; got %v", types)
|
||||
}
|
||||
if posOf("response.content_part.done") < 0 {
|
||||
t.Errorf("response.content_part.done was not emitted; got %v", types)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_DoneEventsCarryFullText pins that done events
|
||||
// carry the part's full text (deltas carry increments only).
|
||||
func TestAnthropicEventToResponses_DoneEventsCarryFullText(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hello "}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "world"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
|
||||
const want = "Hello world"
|
||||
var sawTextDone, sawPartDone bool
|
||||
for _, e := range events {
|
||||
switch e.Type {
|
||||
case "response.output_text.done":
|
||||
sawTextDone = true
|
||||
if e.Text != want {
|
||||
t.Errorf("output_text.done text = %q, want %q", e.Text, want)
|
||||
}
|
||||
case "response.content_part.done":
|
||||
sawPartDone = true
|
||||
if e.Part == nil || e.Part.Text != want {
|
||||
t.Errorf("content_part.done part = %+v, want text %q", e.Part, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawTextDone || !sawPartDone {
|
||||
t.Errorf("missing done events: output_text.done=%v content_part.done=%v", sawTextDone, sawPartDone)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_CompletedCarriesOutput pins that
|
||||
// response.completed carries the full output list. The SDK's
|
||||
// get_final_response() and tracing integrations parse the terminal event's
|
||||
// response directly; an empty output leaves them with nothing (the text still
|
||||
// renders from deltas, which is why this is invisible when only watching the
|
||||
// stream).
|
||||
func TestAnthropicEventToResponses_CompletedCarriesOutput(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "4826"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
var completed *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.completed" {
|
||||
completed = &events[i]
|
||||
}
|
||||
}
|
||||
if completed == nil || completed.Response == nil {
|
||||
t.Fatalf("response.completed was not emitted")
|
||||
}
|
||||
if len(completed.Response.Output) == 0 {
|
||||
t.Fatalf("response.completed carries an empty output; clients would see no result")
|
||||
}
|
||||
msg := completed.Response.Output[0]
|
||||
if msg.Type != "message" || len(msg.Content) == 0 {
|
||||
t.Fatalf("output[0] = %+v, want a message with content", msg)
|
||||
}
|
||||
if msg.Content[0].Text != "4826" {
|
||||
t.Errorf("output[0].content[0].text = %q, want %q", msg.Content[0].Text, "4826")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments pins that a
|
||||
// function call's accumulated arguments survive into output_item.done and
|
||||
// response.completed.
|
||||
func TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use", ID: "toolu_1", Name: "get_weather",
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `{"city":`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `"SH"}`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
var completed *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.completed" {
|
||||
completed = &events[i]
|
||||
}
|
||||
}
|
||||
if completed == nil || completed.Response == nil || len(completed.Response.Output) == 0 {
|
||||
t.Fatalf("response.completed carries no output")
|
||||
}
|
||||
fc := completed.Response.Output[0]
|
||||
if fc.Type != "function_call" {
|
||||
t.Fatalf("output[0].type = %q, want function_call", fc.Type)
|
||||
}
|
||||
if fc.Arguments != `{"city":"SH"}` {
|
||||
t.Errorf("arguments = %q, want %q", fc.Arguments, `{"city":"SH"}`)
|
||||
}
|
||||
if fc.Name != "get_weather" {
|
||||
t.Errorf("name = %q, want get_weather", fc.Name)
|
||||
}
|
||||
}
|
||||
@@ -390,7 +390,7 @@ func ChatCompletionsResponseToAnthropic(resp *ChatCompletionsResponse, model str
|
||||
if len(resp.Choices) > 0 {
|
||||
choice := resp.Choices[0]
|
||||
out.Content = chatMessageToAnthropicBlocks(choice.Message)
|
||||
out.StopReason = chatFinishReasonToAnthropicStopReason(choice.FinishReason, out.Content)
|
||||
out.StopReason = AnthropicStopReasonPtr(chatFinishReasonToAnthropicStopReason(choice.FinishReason, out.Content))
|
||||
// "length" → "max_tokens" is handled by chatFinishReasonToAnthropicStopReason;
|
||||
// Anthropic conveys max-tokens via stop_reason only, no incomplete_details field.
|
||||
}
|
||||
@@ -404,9 +404,9 @@ func ChatCompletionsResponseToAnthropic(resp *ChatCompletionsResponse, model str
|
||||
}
|
||||
// Empty choices / nil response never enter the choices branch above; the
|
||||
// double-conversion path still reports a completed turn ("end_turn"), and
|
||||
// stop_reason "" is invalid for strict Anthropic clients.
|
||||
if out.StopReason == "" {
|
||||
out.StopReason = chatFinishReasonToAnthropicStopReason("", out.Content)
|
||||
// stop_reason null/"" is invalid for completed non-stream responses.
|
||||
if AnthropicStopReasonString(out.StopReason) == "" {
|
||||
out.StopReason = AnthropicStopReasonPtr(chatFinishReasonToAnthropicStopReason("", out.Content))
|
||||
}
|
||||
// The double-conversion path generates a response id when the upstream
|
||||
// omits one (ChatCompletionsResponseToResponses); clients treat it as required.
|
||||
@@ -701,12 +701,13 @@ func ensureCCAnthropicMessageStart(state *ChatCompletionsToAnthropicStreamState)
|
||||
return []AnthropicStreamEvent{{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
Usage: AnthropicUsage{InputTokens: 0, OutputTokens: 0},
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
StopReason: nil, // JSON null; never ""
|
||||
Usage: AnthropicUsage{InputTokens: 0, OutputTokens: 0},
|
||||
},
|
||||
}}
|
||||
}
|
||||
|
||||
@@ -314,7 +314,7 @@ func TestChatCompletionsResponseToAnthropic_TextOnly(t *testing.T) {
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "hello world", out.Content[0].Text)
|
||||
require.Equal(t, "end_turn", out.StopReason)
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason))
|
||||
require.Equal(t, 5, out.Usage.InputTokens)
|
||||
require.Equal(t, 2, out.Usage.OutputTokens)
|
||||
}
|
||||
@@ -346,7 +346,7 @@ func TestChatCompletionsResponseToAnthropic_ToolUse(t *testing.T) {
|
||||
require.Equal(t, "call_1", out.Content[0].ID)
|
||||
require.Equal(t, "get_weather", out.Content[0].Name)
|
||||
require.Equal(t, `{"city":"SF"}`, string(out.Content[0].Input))
|
||||
require.Equal(t, "tool_use", out.StopReason)
|
||||
require.Equal(t, "tool_use", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_ReasoningOnlyFallback(t *testing.T) {
|
||||
@@ -384,7 +384,7 @@ func TestChatCompletionsResponseToAnthropic_FinishReasonLength(t *testing.T) {
|
||||
}
|
||||
|
||||
out := ChatCompletionsResponseToAnthropic(resp, "claude-sonnet-4-20250514")
|
||||
require.Equal(t, "max_tokens", out.StopReason)
|
||||
require.Equal(t, "max_tokens", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_EmptyChoices(t *testing.T) {
|
||||
@@ -398,7 +398,7 @@ func TestChatCompletionsResponseToAnthropic_EmptyChoices(t *testing.T) {
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "", out.Content[0].Text)
|
||||
require.Equal(t, "end_turn", out.StopReason, "empty choices must not produce an empty stop_reason")
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason), "empty choices must not produce an empty stop_reason")
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_CacheTokens(t *testing.T) {
|
||||
@@ -433,7 +433,7 @@ func TestChatCompletionsResponseToAnthropic_NilResponse(t *testing.T) {
|
||||
out := ChatCompletionsResponseToAnthropic(nil, "claude-sonnet-4-20250514")
|
||||
require.Len(t, out.Content, 1)
|
||||
require.Equal(t, "text", out.Content[0].Type)
|
||||
require.Equal(t, "end_turn", out.StopReason, "nil response must not produce an empty stop_reason")
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(out.StopReason), "nil response must not produce an empty stop_reason")
|
||||
require.NotEmpty(t, out.ID)
|
||||
}
|
||||
|
||||
@@ -678,7 +678,7 @@ func TestDirectBridge_NonStreamingMatchesDoubleConversion(t *testing.T) {
|
||||
double := ResponsesToAnthropic(responsesResp, "claude-sonnet-4-20250514")
|
||||
|
||||
// Compare key fields
|
||||
require.Equal(t, direct.StopReason, double.StopReason)
|
||||
require.Equal(t, AnthropicStopReasonString(direct.StopReason), AnthropicStopReasonString(double.StopReason))
|
||||
require.Equal(t, direct.Model, double.Model)
|
||||
require.Len(t, direct.Content, len(double.Content))
|
||||
for i := range direct.Content {
|
||||
@@ -1061,8 +1061,8 @@ func TestDirectBridge_NonStreamingMatchesDoubleConversion_EmptyChoices(t *testin
|
||||
responsesResp := ChatCompletionsResponseToResponses(resp, "claude-sonnet-4-20250514", nil, false, nil)
|
||||
double := ResponsesToAnthropic(responsesResp, "claude-sonnet-4-20250514")
|
||||
|
||||
require.Equal(t, double.StopReason, direct.StopReason)
|
||||
require.Equal(t, "end_turn", direct.StopReason)
|
||||
require.Equal(t, AnthropicStopReasonString(double.StopReason), AnthropicStopReasonString(direct.StopReason))
|
||||
require.Equal(t, "end_turn", AnthropicStopReasonString(direct.StopReason))
|
||||
}
|
||||
|
||||
func TestChatCompletionsResponseToAnthropic_ContentFilterWithToolUse(t *testing.T) {
|
||||
@@ -1086,5 +1086,5 @@ func TestChatCompletionsResponseToAnthropic_ContentFilterWithToolUse(t *testing.
|
||||
}
|
||||
|
||||
out := ChatCompletionsResponseToAnthropic(resp, "claude-sonnet-4-20250514")
|
||||
require.Equal(t, "tool_use", out.StopReason)
|
||||
require.Equal(t, "tool_use", AnthropicStopReasonString(out.StopReason))
|
||||
}
|
||||
|
||||
@@ -88,7 +88,7 @@ func TestAnthropicToResponsesResponse_CacheCreation(t *testing.T) {
|
||||
CacheReadInputTokens: 4,
|
||||
CacheCreationInputTokens: 6,
|
||||
},
|
||||
StopReason: "end_turn",
|
||||
StopReason: AnthropicStopReasonPtr("end_turn"),
|
||||
}
|
||||
|
||||
out := AnthropicToResponsesResponse(&resp)
|
||||
|
||||
@@ -3,6 +3,7 @@ package apicompat
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
@@ -32,10 +33,14 @@ func ResponsesToAnthropic(resp *ResponsesResponse, model string) *AnthropicRespo
|
||||
summaryText += s.Text
|
||||
}
|
||||
}
|
||||
if summaryText != "" {
|
||||
// Always surface encrypted_content as thinking.signature so Claude
|
||||
// Code / multi-turn clients can send it back. Signature-only
|
||||
// thinking blocks are valid when the model omits a visible summary.
|
||||
if summaryText != "" || strings.TrimSpace(item.EncryptedContent) != "" {
|
||||
blocks = append(blocks, AnthropicContentBlock{
|
||||
Type: "thinking",
|
||||
Thinking: summaryText,
|
||||
Type: "thinking",
|
||||
Thinking: summaryText,
|
||||
Signature: item.EncryptedContent,
|
||||
})
|
||||
}
|
||||
case "message":
|
||||
@@ -81,7 +86,7 @@ func ResponsesToAnthropic(resp *ResponsesResponse, model string) *AnthropicRespo
|
||||
}
|
||||
out.Content = blocks
|
||||
|
||||
out.StopReason = responsesStatusToAnthropicStopReason(resp.Status, resp.IncompleteDetails, blocks)
|
||||
out.StopReason = AnthropicStopReasonPtr(responsesStatusToAnthropicStopReason(resp.Status, resp.IncompleteDetails, blocks))
|
||||
|
||||
if resp.Usage != nil {
|
||||
out.Usage = anthropicUsageFromResponsesUsage(resp.Usage)
|
||||
@@ -177,7 +182,10 @@ type ResponsesEventToAnthropicState struct {
|
||||
CurrentToolName string
|
||||
CurrentToolArgs string
|
||||
CurrentToolHadDelta bool
|
||||
HasToolCall bool
|
||||
// PendingThinkingSignature is filled from reasoning.encrypted_content and
|
||||
// emitted as signature_delta before the thinking block is closed.
|
||||
PendingThinkingSignature string
|
||||
HasToolCall bool
|
||||
|
||||
// OutputIndexToBlockIdx maps Responses output_index → Anthropic content block index.
|
||||
OutputIndexToBlockIdx map[int]int
|
||||
@@ -228,7 +236,10 @@ func ResponsesEventToAnthropicEvents(
|
||||
"response.reasoning_text.delta":
|
||||
return resToAnthHandleReasoningDelta(evt, state)
|
||||
case "response.reasoning_summary_text.done":
|
||||
return resToAnthHandleBlockDone(state)
|
||||
// Keep the thinking block open until response.output_item.done.
|
||||
// Grok/Codex attach encrypted_content on the finished reasoning item;
|
||||
// closing early would drop signature_delta and break multi-turn cache.
|
||||
return nil
|
||||
// response.done 是 Realtime/WS 与项目透传路径使用的终止别名;
|
||||
// 普通 Responses HTTP SSE 的公开终止事件仍以 response.completed 为主。
|
||||
case "response.completed", "response.done", "response.incomplete", "response.failed":
|
||||
@@ -297,14 +308,19 @@ func resToAnthHandleCreated(evt *ResponsesStreamEvent, state *ResponsesEventToAn
|
||||
}
|
||||
state.MessageStartSent = true
|
||||
|
||||
// Official Anthropic message_start uses stop_reason: null and usage with
|
||||
// input_tokens when known. We leave StopReason nil (JSON null) and usage
|
||||
// zeros until response.completed; never emit stop_reason:"" which breaks
|
||||
// strict clients' turn-finalization / session usage accounting.
|
||||
return []AnthropicStreamEvent{{
|
||||
Type: "message_start",
|
||||
Message: &AnthropicResponse{
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
ID: state.ResponseID,
|
||||
Type: "message",
|
||||
Role: "assistant",
|
||||
Content: []AnthropicContentBlock{},
|
||||
Model: state.Model,
|
||||
StopReason: nil,
|
||||
Usage: AnthropicUsage{
|
||||
InputTokens: 0,
|
||||
OutputTokens: 0,
|
||||
@@ -354,6 +370,7 @@ func resToAnthHandleOutputItemAdded(evt *ResponsesStreamEvent, state *ResponsesE
|
||||
state.OutputIndexToBlockIdx[evt.OutputIndex] = idx
|
||||
state.ContentBlockOpen = true
|
||||
state.CurrentBlockType = "thinking"
|
||||
state.PendingThinkingSignature = strings.TrimSpace(evt.Item.EncryptedContent)
|
||||
|
||||
events = append(events, AnthropicStreamEvent{
|
||||
Type: "content_block_start",
|
||||
@@ -537,6 +554,13 @@ func resToAnthHandleOutputItemDone(evt *ResponsesStreamEvent, state *ResponsesEv
|
||||
return resToAnthHandleWebSearchDone(evt, state)
|
||||
}
|
||||
|
||||
// Capture encrypted_content on reasoning item done (often only present here).
|
||||
if evt.Item.Type == "reasoning" {
|
||||
if sig := strings.TrimSpace(evt.Item.EncryptedContent); sig != "" {
|
||||
state.PendingThinkingSignature = sig
|
||||
}
|
||||
}
|
||||
|
||||
if state.ContentBlockOpen {
|
||||
return closeCurrentBlock(state)
|
||||
}
|
||||
@@ -658,13 +682,30 @@ func closeCurrentBlock(state *ResponsesEventToAnthropicState) []AnthropicStreamE
|
||||
return nil
|
||||
}
|
||||
idx := state.ContentBlockIndex
|
||||
var events []AnthropicStreamEvent
|
||||
// Emit signature_delta before stop so Claude clients retain encrypted
|
||||
// reasoning for the next turn (required for Grok multi-turn cache).
|
||||
if state.CurrentBlockType == "thinking" {
|
||||
if sig := strings.TrimSpace(state.PendingThinkingSignature); sig != "" {
|
||||
events = append(events, AnthropicStreamEvent{
|
||||
Type: "content_block_delta",
|
||||
Index: &idx,
|
||||
Delta: &AnthropicDelta{
|
||||
Type: "signature_delta",
|
||||
Signature: sig,
|
||||
},
|
||||
})
|
||||
}
|
||||
state.PendingThinkingSignature = ""
|
||||
}
|
||||
state.ContentBlockOpen = false
|
||||
state.ContentBlockIndex++
|
||||
state.CurrentToolName = ""
|
||||
state.CurrentToolArgs = ""
|
||||
state.CurrentToolHadDelta = false
|
||||
return []AnthropicStreamEvent{{
|
||||
events = append(events, AnthropicStreamEvent{
|
||||
Type: "content_block_stop",
|
||||
Index: &idx,
|
||||
}}
|
||||
})
|
||||
return events
|
||||
}
|
||||
|
||||
@@ -60,6 +60,9 @@ type AnthropicContentBlock struct {
|
||||
|
||||
// type=thinking
|
||||
Thinking string `json:"thinking,omitempty"`
|
||||
// Signature carries provider encrypted reasoning (e.g. xAI encrypted_content)
|
||||
// so multi-turn Claude clients can round-trip it back on subsequent turns.
|
||||
Signature string `json:"signature,omitempty"`
|
||||
|
||||
// type=image
|
||||
Source *AnthropicImageSource `json:"source,omitempty"`
|
||||
@@ -121,17 +124,34 @@ type AnthropicCacheControl struct {
|
||||
}
|
||||
|
||||
// AnthropicResponse is the non-streaming response from POST /v1/messages.
|
||||
//
|
||||
// StopReason is a pointer so streaming message_start can emit JSON null
|
||||
// (official Anthropic wire format). A plain string zero-value would marshal as
|
||||
// "" which strict clients treat as invalid mid-stream state.
|
||||
type AnthropicResponse struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"` // "message"
|
||||
Role string `json:"role"` // "assistant"
|
||||
Content []AnthropicContentBlock `json:"content"`
|
||||
Model string `json:"model"`
|
||||
StopReason string `json:"stop_reason"`
|
||||
StopReason *string `json:"stop_reason"`
|
||||
StopSequence *string `json:"stop_sequence,omitempty"`
|
||||
Usage AnthropicUsage `json:"usage"`
|
||||
}
|
||||
|
||||
// AnthropicStopReasonPtr returns a non-nil pointer to s for final stop reasons.
|
||||
func AnthropicStopReasonPtr(s string) *string {
|
||||
return &s
|
||||
}
|
||||
|
||||
// AnthropicStopReasonString returns the stop reason value, or "" when unset/null.
|
||||
func AnthropicStopReasonString(p *string) string {
|
||||
if p == nil {
|
||||
return ""
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
// AnthropicUsage holds token counts in Anthropic format.
|
||||
type AnthropicUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
@@ -230,6 +250,9 @@ type ResponsesInputItem struct {
|
||||
Role string `json:"role,omitempty"`
|
||||
Content json.RawMessage `json:"content,omitempty"` // string or []ResponsesContentPart
|
||||
|
||||
// type=reasoning (multi-turn replay of encrypted reasoning)
|
||||
EncryptedContent string `json:"encrypted_content,omitempty"`
|
||||
|
||||
// type=function_call
|
||||
CallID string `json:"call_id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
|
||||
+140
-39
@@ -8,40 +8,138 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetClientIP 从 Gin Context 中提取客户端真实 IP 地址。
|
||||
// 按以下优先级检查 Header:
|
||||
// 1. CF-Connecting-IP (Cloudflare)
|
||||
// 2. X-Real-IP (Nginx)
|
||||
// 3. X-Forwarded-For (取第一个非私有 IP)
|
||||
// 4. c.ClientIP() (Gin 内置方法)
|
||||
const forwardedIPSettingsKey = "sub2api.forwarded_ip_settings"
|
||||
|
||||
type forwardedIPSettings struct {
|
||||
trustForwarded bool
|
||||
headers []string
|
||||
}
|
||||
|
||||
// SetForwardedIPSettings snapshots the forwarded-IP mode and custom header list
|
||||
// for this request.
|
||||
func SetForwardedIPSettings(c *gin.Context, enabled bool, headers []string) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
c.Set(forwardedIPSettingsKey, forwardedIPSettings{
|
||||
trustForwarded: enabled,
|
||||
headers: append([]string(nil), headers...),
|
||||
})
|
||||
}
|
||||
|
||||
// SetLegacyForwardedIPTrust records whether raw forwarding headers override
|
||||
// Gin's server.trusted_proxies chain for this request.
|
||||
func SetLegacyForwardedIPTrust(c *gin.Context, enabled bool) {
|
||||
SetForwardedIPSettings(c, enabled, nil)
|
||||
}
|
||||
|
||||
func requestForwardedIPSettings(c *gin.Context) (forwardedIPSettings, bool) {
|
||||
if c == nil {
|
||||
return forwardedIPSettings{}, false
|
||||
}
|
||||
value, ok := c.Get(forwardedIPSettingsKey)
|
||||
if !ok {
|
||||
return forwardedIPSettings{}, false
|
||||
}
|
||||
settings, ok := value.(forwardedIPSettings)
|
||||
return settings, ok
|
||||
}
|
||||
|
||||
func requestUsesLegacyForwardedIPTrust(c *gin.Context) bool {
|
||||
settings, ok := requestForwardedIPSettings(c)
|
||||
return !ok || settings.trustForwarded
|
||||
}
|
||||
|
||||
// GetClientIP resolves the client address using the legacy forwarding-header
|
||||
// precedence used before the trusted-proxy hardening. It remains the
|
||||
// compatibility path for request metadata and usage/error logs; security-
|
||||
// sensitive callers must use GetTrustedClientIP or GetSecurityClientIP.
|
||||
func GetClientIP(c *gin.Context) string {
|
||||
// 1. Cloudflare
|
||||
if ip := c.GetHeader("CF-Connecting-IP"); ip != "" {
|
||||
return normalizeIP(ip)
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
if !requestUsesLegacyForwardedIPTrust(c) {
|
||||
return GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
// 2. Nginx X-Real-IP
|
||||
if ip := c.GetHeader("X-Real-IP"); ip != "" {
|
||||
return normalizeIP(ip)
|
||||
settings, _ := requestForwardedIPSettings(c)
|
||||
customIP, customFallback := resolveCustomForwardedClientIP(c, settings.headers)
|
||||
if customIP != "" {
|
||||
return customIP
|
||||
}
|
||||
|
||||
// 3. X-Forwarded-For (多个 IP 时取第一个公网 IP)
|
||||
if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
|
||||
ips := strings.Split(xff, ",")
|
||||
for _, ip := range ips {
|
||||
ip = strings.TrimSpace(ip)
|
||||
if ip != "" && !isPrivateIP(ip) {
|
||||
return normalizeIP(ip)
|
||||
// Preserve the historical precedence used by existing reverse-proxy
|
||||
// deployments, while skipping an internal proxy address when a public XFF
|
||||
// value is available. This covers Docker/Nginx setups that accidentally
|
||||
// write the bridge address into X-Real-IP.
|
||||
legacyIP, legacyFallback := resolveLegacyForwardedHeaderIP(c)
|
||||
if legacyIP != "" {
|
||||
return legacyIP
|
||||
}
|
||||
if customFallback != "" {
|
||||
return customFallback
|
||||
}
|
||||
if legacyFallback != "" {
|
||||
return legacyFallback
|
||||
}
|
||||
return normalizeIP(c.ClientIP())
|
||||
}
|
||||
|
||||
func resolveCustomForwardedClientIP(c *gin.Context, headers []string) (string, string) {
|
||||
if c == nil {
|
||||
return "", ""
|
||||
}
|
||||
var fallback string
|
||||
for _, header := range headers {
|
||||
for _, value := range c.Request.Header.Values(header) {
|
||||
for _, candidate := range strings.Split(value, ",") {
|
||||
parsed := net.ParseIP(strings.TrimSpace(candidate))
|
||||
if parsed == nil {
|
||||
continue
|
||||
}
|
||||
normalized := parsed.String()
|
||||
if isPrivateIP(normalized) {
|
||||
if fallback == "" {
|
||||
fallback = normalized
|
||||
}
|
||||
continue
|
||||
}
|
||||
return normalized, fallback
|
||||
}
|
||||
}
|
||||
// 如果都是私有 IP,返回第一个
|
||||
if len(ips) > 0 {
|
||||
return normalizeIP(strings.TrimSpace(ips[0]))
|
||||
}
|
||||
return "", fallback
|
||||
}
|
||||
|
||||
func resolveLegacyForwardedHeaderIP(c *gin.Context) (string, string) {
|
||||
var fallback string
|
||||
if forwarded := normalizeIP(c.GetHeader("CF-Connecting-IP")); forwarded != "" {
|
||||
fallback = forwarded
|
||||
if !isPrivateIP(forwarded) {
|
||||
return forwarded, fallback
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Gin 内置方法
|
||||
return normalizeIP(c.ClientIP())
|
||||
if realIP := normalizeIP(c.GetHeader("X-Real-IP")); realIP != "" {
|
||||
if fallback == "" {
|
||||
fallback = realIP
|
||||
}
|
||||
if !isPrivateIP(realIP) {
|
||||
return realIP, fallback
|
||||
}
|
||||
}
|
||||
if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
|
||||
ips := strings.Split(xff, ",")
|
||||
for _, candidate := range ips {
|
||||
candidate = strings.TrimSpace(candidate)
|
||||
if candidate != "" && !isPrivateIP(candidate) {
|
||||
return normalizeIP(candidate), fallback
|
||||
}
|
||||
}
|
||||
if fallback == "" && len(ips) > 0 {
|
||||
fallback = normalizeIP(strings.TrimSpace(ips[0]))
|
||||
}
|
||||
}
|
||||
return "", fallback
|
||||
}
|
||||
|
||||
// GetTrustedClientIP 从 Gin 的可信代理解析链提取客户端 IP。
|
||||
@@ -54,11 +152,14 @@ func GetTrustedClientIP(c *gin.Context) string {
|
||||
return normalizeIP(c.ClientIP())
|
||||
}
|
||||
|
||||
// GetSecurityClientIP 返回安全敏感场景(API Key IP 限制、审计日志、会话 IP/UA 绑定)
|
||||
// 使用的客户端 IP。trustForwarded 对应系统设置「信任反代传递的客户端 IP」:
|
||||
// 开启时信任反代转发头(CF-Connecting-IP / X-Real-IP / X-Forwarded-For),
|
||||
// 关闭时走 Gin trusted_proxies 解析链。
|
||||
// GetSecurityClientIP returns the address used by security-sensitive paths.
|
||||
// When legacy forwarded-IP trust is enabled, raw forwarding headers take over
|
||||
// client-IP resolution. When disabled, Gin's server.trusted_proxies chain is
|
||||
// authoritative.
|
||||
func GetSecurityClientIP(c *gin.Context, trustForwarded bool) string {
|
||||
if requestSettings, ok := requestForwardedIPSettings(c); ok {
|
||||
trustForwarded = requestSettings.trustForwarded
|
||||
}
|
||||
if trustForwarded {
|
||||
return GetClientIP(c)
|
||||
}
|
||||
@@ -75,17 +176,10 @@ func normalizeIP(ip string) string {
|
||||
return ip
|
||||
}
|
||||
|
||||
// privateNets 预编译私有 IP CIDR 块,避免每次调用 isPrivateIP 时重复解析
|
||||
// privateNets contains the private/loopback ranges skipped while selecting a
|
||||
// public address from a legacy X-Forwarded-For chain.
|
||||
var privateNets []*net.IPNet
|
||||
|
||||
// CompiledIPRules 表示预编译的 IP 匹配规则。
|
||||
// PatternCount 记录原始规则数量,用于保留“规则存在但全无效”时的行为语义。
|
||||
type CompiledIPRules struct {
|
||||
CIDRs []*net.IPNet
|
||||
IPs []net.IP
|
||||
PatternCount int
|
||||
}
|
||||
|
||||
func init() {
|
||||
for _, cidr := range []string{
|
||||
"10.0.0.0/8",
|
||||
@@ -103,6 +197,14 @@ func init() {
|
||||
}
|
||||
}
|
||||
|
||||
// CompiledIPRules 表示预编译的 IP 匹配规则。
|
||||
// PatternCount 记录原始规则数量,用于保留“规则存在但全无效”时的行为语义。
|
||||
type CompiledIPRules struct {
|
||||
CIDRs []*net.IPNet
|
||||
IPs []net.IP
|
||||
PatternCount int
|
||||
}
|
||||
|
||||
// CompileIPRules 将 IP/CIDR 字符串规则预编译为可复用结构。
|
||||
// 非法规则会被忽略,但 PatternCount 会保留原始规则条数。
|
||||
func CompileIPRules(patterns []string) *CompiledIPRules {
|
||||
@@ -150,7 +252,6 @@ func matchesCompiledRules(parsedIP net.IP, rules *CompiledIPRules) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// isPrivateIP 检查 IP 是否为私有地址。
|
||||
func isPrivateIP(ipStr string) bool {
|
||||
ip := net.ParseIP(ipStr)
|
||||
if ip == nil {
|
||||
|
||||
@@ -10,48 +10,6 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsPrivateIP(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ip string
|
||||
expected bool
|
||||
}{
|
||||
// 私有 IPv4
|
||||
{"10.x 私有地址", "10.0.0.1", true},
|
||||
{"10.x 私有地址段末", "10.255.255.255", true},
|
||||
{"172.16.x 私有地址", "172.16.0.1", true},
|
||||
{"172.31.x 私有地址", "172.31.255.255", true},
|
||||
{"192.168.x 私有地址", "192.168.1.1", true},
|
||||
{"127.0.0.1 本地回环", "127.0.0.1", true},
|
||||
{"127.x 回环段", "127.255.255.255", true},
|
||||
|
||||
// 公网 IPv4
|
||||
{"8.8.8.8 公网 DNS", "8.8.8.8", false},
|
||||
{"1.1.1.1 公网", "1.1.1.1", false},
|
||||
{"172.15.255.255 非私有", "172.15.255.255", false},
|
||||
{"172.32.0.0 非私有", "172.32.0.0", false},
|
||||
{"11.0.0.1 公网", "11.0.0.1", false},
|
||||
|
||||
// IPv6
|
||||
{"::1 IPv6 回环", "::1", true},
|
||||
{"fc00:: IPv6 私有", "fc00::1", true},
|
||||
{"fd00:: IPv6 私有", "fd00::1", true},
|
||||
{"2001:db8::1 IPv6 公网", "2001:db8::1", false},
|
||||
|
||||
// 无效输入
|
||||
{"空字符串", "", false},
|
||||
{"非法字符串", "not-an-ip", false},
|
||||
{"不完整 IP", "192.168", false},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := isPrivateIP(tc.ip)
|
||||
require.Equal(t, tc.expected, got, "isPrivateIP(%q)", tc.ip)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -74,6 +32,26 @@ func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) {
|
||||
require.Equal(t, "9.9.9.9", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetClientIPPreservesLegacyDockerForwardedHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetClientIP(c))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "192.168.32.1:12345"
|
||||
req.Header.Set("X-Forwarded-For", "10.0.0.2, 203.0.113.42")
|
||||
req.Header.Set("X-Real-IP", "192.168.32.1")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, "203.0.113.42", w.Body.String())
|
||||
}
|
||||
|
||||
func TestCheckIPRestrictionWithCompiledRules(t *testing.T) {
|
||||
whitelist := CompileIPRules([]string{"10.0.0.0/8", "192.168.1.2"})
|
||||
blacklist := CompileIPRules([]string{"10.1.1.1"})
|
||||
@@ -95,22 +73,195 @@ func TestCheckIPRestrictionWithCompiledRules_InvalidWhitelistStillDenies(t *test
|
||||
require.Equal(t, "access denied", reason)
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPHonorsTrustToggle(t *testing.T) {
|
||||
func TestGetSecurityClientIPSwitchEnabledUsesLegacyHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, tc := range []struct {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetSecurityClientIP(c, true))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPCustomHeaderPrecedenceAndFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
trustForwarded bool
|
||||
trustForward bool
|
||||
headers []string
|
||||
requestHeaders map[string]string
|
||||
want string
|
||||
}{
|
||||
{name: "trust disabled uses trusted proxy chain", trustForwarded: false, want: "9.9.9.9"},
|
||||
{name: "trust enabled uses forwarded header", trustForwarded: true, want: "1.2.3.4"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
{
|
||||
name: "configured order precedes built-ins",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-First", "X-CDN-Second"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-First": "198.51.100.10",
|
||||
"X-CDN-Second": "203.0.113.20",
|
||||
"CF-Connecting-IP": "8.8.8.8",
|
||||
},
|
||||
want: "198.51.100.10",
|
||||
},
|
||||
{
|
||||
name: "comma candidates skip invalid and private values",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-First", "X-CDN-Second"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-First": "not-an-ip, 10.0.0.8",
|
||||
"X-CDN-Second": "also-bad, 203.0.113.9",
|
||||
},
|
||||
want: "203.0.113.9",
|
||||
},
|
||||
{
|
||||
name: "legacy public header wins over custom private fallback",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "10.0.0.8",
|
||||
"X-Real-IP": "1.2.3.4",
|
||||
},
|
||||
want: "1.2.3.4",
|
||||
},
|
||||
{
|
||||
name: "custom private fallback retains configured precedence",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "10.0.0.8",
|
||||
"X-Real-IP": "192.168.1.4",
|
||||
},
|
||||
want: "10.0.0.8",
|
||||
},
|
||||
{
|
||||
name: "invalid custom value continues to built-ins",
|
||||
trustForward: true,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "1.2.3.4:443",
|
||||
"CF-Connecting-IP": "4.4.4.4",
|
||||
},
|
||||
want: "4.4.4.4",
|
||||
},
|
||||
{
|
||||
name: "disabled mode ignores custom and legacy headers",
|
||||
trustForward: false,
|
||||
headers: []string{"X-CDN-IP"},
|
||||
requestHeaders: map[string]string{
|
||||
"X-CDN-IP": "1.2.3.4",
|
||||
"X-Real-IP": "4.4.4.4",
|
||||
},
|
||||
want: "9.9.9.9",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetSecurityClientIP(c, tc.trustForwarded))
|
||||
SetForwardedIPSettings(c, test.trustForward, test.headers)
|
||||
c.String(200, GetSecurityClientIP(c, !test.trustForward))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
for name, value := range test.requestHeaders {
|
||||
req.Header.Set(name, value)
|
||||
}
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, test.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPSwitchDisabledUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies([]string{"9.9.9.9"}))
|
||||
r.GET("/t", func(c *gin.Context) { c.String(200, GetSecurityClientIP(c, false)) })
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetClientIPSwitchDisabledUsesTrustedProxyChain(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
SetLegacyForwardedIPTrust(c, false)
|
||||
c.String(200, GetClientIP(c))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "9.9.9.9", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPRequestSnapshotCopiesCustomHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
headers := []string{"X-Original-IP"}
|
||||
SetForwardedIPSettings(c, true, headers)
|
||||
headers[0] = "X-Mutated-IP"
|
||||
c.String(200, GetSecurityClientIP(c, false))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Original-IP", "1.2.3.4")
|
||||
req.Header.Set("X-Mutated-IP", "4.4.4.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPRequestSnapshotOverridesLiveFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
requestTrust bool
|
||||
fallbackTrust bool
|
||||
want string
|
||||
}{
|
||||
{name: "captured secure mode wins", requestTrust: false, fallbackTrust: true, want: "9.9.9.9"},
|
||||
{name: "captured compatibility mode wins", requestTrust: true, fallbackTrust: false, want: "1.2.3.4"},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
SetLegacyForwardedIPTrust(c, test.requestTrust)
|
||||
c.String(200, GetSecurityClientIP(c, test.fallbackTrust))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
@@ -119,8 +270,7 @@ func TestGetSecurityClientIPHonorsTrustToggle(t *testing.T) {
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, tc.want, w.Body.String())
|
||||
require.Equal(t, test.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,10 @@ const (
|
||||
LevelWarn = zapcore.WarnLevel
|
||||
LevelError = zapcore.ErrorLevel
|
||||
LevelFatal = zapcore.FatalLevel
|
||||
|
||||
// OpsSystemLogSkipField keeps an event in the standard logger while
|
||||
// preventing the database-backed Ops system-log sink from indexing it.
|
||||
OpsSystemLogSkipField = "ops_system_log_skip"
|
||||
)
|
||||
|
||||
type Sink interface {
|
||||
|
||||
@@ -7,6 +7,18 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
const GrokFreeRolling24hTokenLimit int64 = 1_000_000
|
||||
|
||||
var grokFreeRolling24hTokenLimits = map[int64]struct{}{
|
||||
GrokFreeRolling24hTokenLimit: {},
|
||||
2_000_000: {}, // Legacy Free limit observed before July 2026.
|
||||
}
|
||||
|
||||
func IsGrokFreeRolling24hTokenLimit(limit int64) bool {
|
||||
_, ok := grokFreeRolling24hTokenLimits[limit]
|
||||
return ok
|
||||
}
|
||||
|
||||
type QuotaWindow struct {
|
||||
Limit *int64 `json:"limit,omitempty"`
|
||||
Remaining *int64 `json:"remaining,omitempty"`
|
||||
|
||||
@@ -61,3 +61,11 @@ func TestObserveQuotaHeadersRecordsNoHeaderProbe(t *testing.T) {
|
||||
require.Nil(t, snapshot.Requests)
|
||||
require.Nil(t, snapshot.Tokens)
|
||||
}
|
||||
|
||||
func TestIsGrokFreeRolling24hTokenLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.True(t, IsGrokFreeRolling24hTokenLimit(GrokFreeRolling24hTokenLimit))
|
||||
require.True(t, IsGrokFreeRolling24hTokenLimit(2_000_000), "legacy snapshots remain classifiable")
|
||||
require.False(t, IsGrokFreeRolling24hTokenLimit(3_000_000))
|
||||
}
|
||||
|
||||
@@ -893,6 +893,20 @@ func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platfor
|
||||
func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
|
||||
sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
|
||||
sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc)
|
||||
if sortBy == "upstream_billing_rate" {
|
||||
direction := "ASC"
|
||||
tieOrder := entsql.Asc
|
||||
if sortOrder == pagination.SortOrderDesc {
|
||||
direction = "DESC"
|
||||
tieOrder = entsql.Desc
|
||||
}
|
||||
return []func(*entsql.Selector){func(s *entsql.Selector) {
|
||||
extra := s.C(dbaccount.FieldExtra)
|
||||
expression := upstreamBillingRateSortExpression(extra)
|
||||
s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST"))
|
||||
s.OrderBy(tieOrder(s.C(dbaccount.FieldID)))
|
||||
}}
|
||||
}
|
||||
|
||||
field := dbaccount.FieldName
|
||||
defaultOrder := true
|
||||
@@ -934,6 +948,40 @@ func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selecto
|
||||
return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)}
|
||||
}
|
||||
|
||||
func upstreamBillingRateSortExpression(extra string) string {
|
||||
status := extra + " #>> '{upstream_billing_probe,status}'"
|
||||
effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'"
|
||||
effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'"
|
||||
resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
|
||||
resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
|
||||
peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'"
|
||||
peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'"
|
||||
peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'"
|
||||
peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'"
|
||||
peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'"
|
||||
peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'"
|
||||
peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)"
|
||||
billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'"
|
||||
timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'"
|
||||
validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'"
|
||||
startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)"
|
||||
endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)"
|
||||
localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))"
|
||||
validPeakWindow := peakStart + " ~ " + validClock + " AND " +
|
||||
peakEnd + " ~ " + validClock + " AND " +
|
||||
startMinute + " < " + endMinute
|
||||
validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " +
|
||||
"EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")"
|
||||
dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig +
|
||||
" THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute +
|
||||
" THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END"
|
||||
legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL"
|
||||
|
||||
return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" +
|
||||
resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot +
|
||||
" AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END"
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||||
accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
|
||||
status: service.StatusActive,
|
||||
@@ -1877,6 +1925,46 @@ func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Con
|
||||
})
|
||||
}
|
||||
|
||||
// ListModelAvailabilityCandidates returns the persistently configured account
|
||||
// pool used to decide whether a model is supported. Unlike scheduling queries,
|
||||
// it intentionally ignores transient runtime state (rate limits, overload,
|
||||
// temporary unschedulability, and expiry windows).
|
||||
func (r *accountRepository) ListModelAvailabilityCandidates(
|
||||
ctx context.Context,
|
||||
groupID *int64,
|
||||
platforms []string,
|
||||
includeGrouped bool,
|
||||
) ([]service.Account, error) {
|
||||
if len(platforms) == 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
if groupID != nil {
|
||||
return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{
|
||||
status: service.StatusActive,
|
||||
schedulable: true,
|
||||
ignoreTransientState: true,
|
||||
platforms: platforms,
|
||||
})
|
||||
}
|
||||
|
||||
preds := []dbpredicate.Account{
|
||||
dbaccount.StatusEQ(service.StatusActive),
|
||||
dbaccount.SchedulableEQ(true),
|
||||
dbaccount.PlatformIn(platforms...),
|
||||
}
|
||||
if !includeGrouped {
|
||||
preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups()))
|
||||
}
|
||||
accounts, err := r.client.Account.Query().
|
||||
Where(preds...).
|
||||
Order(dbent.Asc(dbaccount.FieldPriority)).
|
||||
All(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.accountsToService(ctx, accounts)
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
now := time.Now()
|
||||
_, err := r.client.Account.Update().
|
||||
@@ -2575,6 +2663,12 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
args = append(args, *updates.Schedulable)
|
||||
idx++
|
||||
}
|
||||
if updates.ProbeEnabled != nil {
|
||||
if updates.Extra == nil {
|
||||
updates.Extra = make(map[string]any)
|
||||
}
|
||||
updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled
|
||||
}
|
||||
// JSONB 需要合并而非覆盖,使用 raw SQL 保持旧行为。
|
||||
if len(updates.Credentials) > 0 {
|
||||
payload, err := json.Marshal(updates.Credentials)
|
||||
@@ -2605,8 +2699,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
|
||||
setClauses = append(setClauses, "updated_at = NOW()")
|
||||
|
||||
query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
|
||||
whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
|
||||
args = append(args, pq.Array(ids))
|
||||
idx++
|
||||
if updates.ProbeEnabled != nil {
|
||||
whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1)
|
||||
args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey)
|
||||
}
|
||||
query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause
|
||||
|
||||
baseCtx := ctx
|
||||
contextTx := dbent.TxFromContext(ctx)
|
||||
@@ -2635,6 +2735,20 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if updates.ProbeEnabled != nil {
|
||||
expectedRows := int64(0)
|
||||
seenIDs := make(map[int64]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if _, seen := seenIDs[id]; seen {
|
||||
continue
|
||||
}
|
||||
seenIDs[id] = struct{}{}
|
||||
expectedRows++
|
||||
}
|
||||
if rows != expectedRows {
|
||||
return 0, service.ErrUpstreamBillingProbeAccountInvalid
|
||||
}
|
||||
}
|
||||
if rows > 0 {
|
||||
payload := map[string]any{"account_ids": ids}
|
||||
if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
|
||||
@@ -2662,9 +2776,10 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
}
|
||||
|
||||
type accountGroupQueryOptions struct {
|
||||
status string
|
||||
schedulable bool
|
||||
platforms []string // 允许的多个平台,空切片表示不进行平台过滤
|
||||
status string
|
||||
schedulable bool
|
||||
ignoreTransientState bool
|
||||
platforms []string // 允许的多个平台,空切片表示不进行平台过滤
|
||||
}
|
||||
|
||||
func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) {
|
||||
@@ -2681,14 +2796,16 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in
|
||||
preds = append(preds, dbaccount.PlatformIn(opts.platforms...))
|
||||
}
|
||||
if opts.schedulable {
|
||||
now := time.Now()
|
||||
preds = append(preds,
|
||||
dbaccount.SchedulableEQ(true),
|
||||
tempUnschedulablePredicate(),
|
||||
notExpiredPredicate(now),
|
||||
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
|
||||
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
|
||||
)
|
||||
preds = append(preds, dbaccount.SchedulableEQ(true))
|
||||
if !opts.ignoreTransientState {
|
||||
now := time.Now()
|
||||
preds = append(preds,
|
||||
tempUnschedulablePredicate(),
|
||||
notExpiredPredicate(now),
|
||||
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
|
||||
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if len(preds) > 0 {
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
_ "github.com/Wei-Shaw/sub2api/ent/runtime"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
)
|
||||
|
||||
func TestListModelAvailabilityCandidates_GroupQueryIgnoresTransientState(t *testing.T) {
|
||||
var capturedSQL string
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(captureEntQueryMatcher{actual: &capturedSQL}))
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
driver := entsql.OpenDB(dialect.Postgres, db)
|
||||
client := dbent.NewClient(dbent.Driver(driver))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
|
||||
mock.ExpectQuery("model availability candidates").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
groupID := int64(42)
|
||||
accounts, err := repo.ListModelAvailabilityCandidates(
|
||||
context.Background(),
|
||||
&groupID,
|
||||
[]string{service.PlatformAnthropic},
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, accounts)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
|
||||
normalized := normalizeSQLWhitespace(capturedSQL)
|
||||
_, whereClause, found := strings.Cut(normalized, " WHERE ")
|
||||
require.True(t, found, "expected WHERE clause in query: %s", normalized)
|
||||
whereClause, _, _ = strings.Cut(whereClause, " ORDER BY ")
|
||||
for _, configuredPredicate := range []string{"group_id", "status", "schedulable", "platform"} {
|
||||
require.Contains(t, whereClause, configuredPredicate)
|
||||
}
|
||||
for _, transientPredicate := range []string{
|
||||
"rate_limit_reset_at",
|
||||
"overload_until",
|
||||
"temp_unschedulable_until",
|
||||
"expires_at",
|
||||
"auto_pause_on_expired",
|
||||
} {
|
||||
require.NotContains(t, whereClause, transientPredicate, "configured-state diagnosis must not filter transient predicate %q", transientPredicate)
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,9 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
@@ -33,3 +36,105 @@ func (s *AccountRepoSuite) TestListWithFilters_SortByPriorityDesc() {
|
||||
s.Require().Equal("high-priority", accounts[0].Name)
|
||||
s.Require().Equal("low-priority", accounts[1].Name)
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestListWithFilters_SortByUpstreamBillingRateWithMissingLast() {
|
||||
makeAccount := func(name, status string, rate any) {
|
||||
extra := map[string]any{}
|
||||
if rate != nil {
|
||||
extra[service.UpstreamBillingProbeExtraKey] = map[string]any{
|
||||
"status": status,
|
||||
"data": map[string]any{"effective_rate_multiplier": rate},
|
||||
}
|
||||
}
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{Name: name, Extra: extra})
|
||||
}
|
||||
makeAccount("high-rate", service.UpstreamBillingProbeStatusOK, 0.8)
|
||||
makeAccount("low-rate", service.UpstreamBillingProbeStatusOK, 0.03)
|
||||
makeAccount("missing-rate", "", nil)
|
||||
makeAccount("unsupported-with-retained-rate", service.UpstreamBillingProbeStatusUnsupported, 0.01)
|
||||
|
||||
for _, tc := range []struct {
|
||||
order string
|
||||
want []string
|
||||
}{
|
||||
{order: "asc", want: []string{"low-rate", "high-rate", "missing-rate", "unsupported-with-retained-rate"}},
|
||||
{order: "desc", want: []string{"high-rate", "low-rate", "unsupported-with-retained-rate", "missing-rate"}},
|
||||
} {
|
||||
accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{
|
||||
Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order,
|
||||
}, "", "", "", "", 0, "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 4)
|
||||
for i, name := range tc.want {
|
||||
s.Require().Equal(name, accounts[i].Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestListWithFilters_SortByCurrentUpstreamBillingRateDuringPeak() {
|
||||
now := time.Now()
|
||||
locations := []string{"UTC", "Asia/Shanghai", "America/New_York", "Europe/London"}
|
||||
var timezone string
|
||||
var minute int
|
||||
for _, name := range locations {
|
||||
location, err := time.LoadLocation(name)
|
||||
s.Require().NoError(err)
|
||||
local := now.In(location)
|
||||
candidate := local.Hour()*60 + local.Minute()
|
||||
if candidate >= 2 && candidate <= 1436 {
|
||||
timezone = name
|
||||
minute = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
s.Require().NotEmpty(timezone)
|
||||
|
||||
peakStart := fmt.Sprintf("%02d:%02d", (minute-2)/60, (minute-2)%60)
|
||||
peakEnd := fmt.Sprintf("%02d:%02d", (minute+3)/60, (minute+3)%60)
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{
|
||||
Name: "current-peak-rate",
|
||||
Extra: map[string]any{
|
||||
service.UpstreamBillingProbeExtraKey: map[string]any{
|
||||
"status": service.UpstreamBillingProbeStatusOK,
|
||||
"data": map[string]any{
|
||||
"billing_scope": "token",
|
||||
"resolved_rate_multiplier": 1.0,
|
||||
"effective_rate_multiplier": 1.0,
|
||||
"peak_rate_enabled": true,
|
||||
"peak_start": peakStart,
|
||||
"peak_end": peakEnd,
|
||||
"peak_rate_multiplier": 10.0,
|
||||
"timezone": timezone,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{
|
||||
Name: "current-off-peak-rate",
|
||||
Extra: map[string]any{
|
||||
service.UpstreamBillingProbeExtraKey: map[string]any{
|
||||
"status": service.UpstreamBillingProbeStatusOK,
|
||||
"data": map[string]any{
|
||||
"effective_rate_multiplier": 5.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
for _, tc := range []struct {
|
||||
order string
|
||||
want []string
|
||||
}{
|
||||
{order: "asc", want: []string{"current-off-peak-rate", "current-peak-rate"}},
|
||||
{order: "desc", want: []string{"current-peak-rate", "current-off-peak-rate"}},
|
||||
} {
|
||||
accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{
|
||||
Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order,
|
||||
}, "", "", "", "", 0, "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 2)
|
||||
for i, name := range tc.want {
|
||||
s.Require().Equal(name, accounts[i].Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -163,6 +163,46 @@ func TestBulkUpdateNilProbeRemovesKeyInsteadOfWritingJSONNull(t *testing.T) {
|
||||
require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'")
|
||||
}
|
||||
|
||||
func TestBulkUpdateDisablingProbeRemovesSnapshot(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{27}, service.AccountBulkUpdate{
|
||||
Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'")
|
||||
payload, ok := exec.execArgs[0][0].([]byte)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, `{"upstream_billing_probe_enabled":false}`, string(payload))
|
||||
}
|
||||
|
||||
func TestBulkUpdateProbeEligibilityMismatchRollsBack(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
enabled := true
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectRollback()
|
||||
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{
|
||||
ProbeEnabled: &enabled,
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeAccountInvalid)
|
||||
require.Zero(t, rows)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -110,28 +110,25 @@ func (c *apiKeyCache) SubscribeAuthCacheInvalidation(ctx context.Context, handle
|
||||
return fmt.Errorf("subscribe to auth cache invalidation: %w", err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
if err := pubsub.Close(); err != nil {
|
||||
log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if msg != nil {
|
||||
handler(msg.Payload)
|
||||
}
|
||||
}
|
||||
defer func() {
|
||||
if err := pubsub.Close(); err != nil {
|
||||
log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err)
|
||||
}
|
||||
}()
|
||||
service.NotifyAuthCacheSubscriptionReady(ctx)
|
||||
|
||||
return nil
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return errors.New("auth cache invalidation pubsub channel closed")
|
||||
}
|
||||
if msg != nil {
|
||||
handler(msg.Payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyCacheSubscriber_BlocksUntilContextCancellation(t *testing.T) {
|
||||
server := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
||||
defer func() { _ = client.Close() }()
|
||||
cache := NewAPIKeyCache(client)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
received := make(chan string, 1)
|
||||
returned := make(chan error, 1)
|
||||
go func() {
|
||||
returned <- cache.SubscribeAuthCacheInvalidation(ctx, func(value string) { received <- value })
|
||||
}()
|
||||
|
||||
var value string
|
||||
require.Eventually(t, func() bool {
|
||||
require.NoError(t, client.Publish(context.Background(), authCacheInvalidateChannel, "hash").Err())
|
||||
select {
|
||||
case value = <-received:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
require.Equal(t, "hash", value)
|
||||
select {
|
||||
case err := <-returned:
|
||||
t.Fatalf("subscriber returned while connection was active: %v", err)
|
||||
default:
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case err := <-returned:
|
||||
require.True(t, errors.Is(err, context.Canceled))
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("subscriber did not stop after context cancellation")
|
||||
}
|
||||
}
|
||||
@@ -326,16 +326,14 @@ func (r *apiKeyRepository) Delete(ctx context.Context, id int64) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// DeleteWithAudit 在同一事务内:
|
||||
// 1. 把(明文 key、所有者、key 名称)写入 deleted_api_key_audits;
|
||||
// 2. 软删除该 key(tombstone 覆盖 key 列以释放唯一约束)。
|
||||
//
|
||||
// 保证"被删除的 key 一定能反查到所有者"。事务模式与 group_repo.DeleteCascade 一致。
|
||||
// DeleteWithAudit keeps the legacy method name for rolling-upgrade compatibility.
|
||||
// It atomically tombstones and soft-deletes the key without retaining credential
|
||||
// material. Tombstoning releases the unique key value for safe reuse.
|
||||
func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error {
|
||||
tombstoneKey := fmt.Sprintf("__deleted__%d__%d", id, time.Now().UnixNano())
|
||||
|
||||
if existingTx := dbent.TxFromContext(ctx); existingTx != nil {
|
||||
return r.deleteWithAudit(ctx, existingTx.Client(), id, tombstoneKey)
|
||||
return r.deleteWithTombstone(ctx, existingTx.Client(), id, tombstoneKey)
|
||||
}
|
||||
|
||||
tx, err := r.client.Tx(ctx)
|
||||
@@ -348,7 +346,7 @@ func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error
|
||||
exec = tx.Client()
|
||||
}
|
||||
|
||||
if err := r.deleteWithAudit(ctx, exec, id, tombstoneKey); err != nil {
|
||||
if err := r.deleteWithTombstone(ctx, exec, id, tombstoneKey); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -358,17 +356,7 @@ func (r *apiKeyRepository) DeleteWithAudit(ctx context.Context, id int64) error
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *apiKeyRepository) deleteWithAudit(ctx context.Context, exec *dbent.Client, id int64, tombstoneKey string) error {
|
||||
// 1. 审计:数据源即 api_keys 当前行;WHERE deleted_at IS NULL 保证只对未删除行写一次。
|
||||
if _, err := exec.ExecContext(ctx, `
|
||||
INSERT INTO deleted_api_key_audits (key, api_key_id, user_id, key_name, deleted_at)
|
||||
SELECT key, id, user_id, name, NOW()
|
||||
FROM api_keys
|
||||
WHERE id = $1 AND deleted_at IS NULL`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 2. 软删除(tombstone 覆盖 key)。
|
||||
func (r *apiKeyRepository) deleteWithTombstone(ctx context.Context, exec *dbent.Client, id int64, tombstoneKey string) error {
|
||||
res, err := exec.ExecContext(ctx, `
|
||||
UPDATE api_keys
|
||||
SET key = $1, deleted_at = NOW(), updated_at = NOW()
|
||||
|
||||
@@ -556,7 +556,7 @@ func TestIncrementQuotaUsed_Concurrent(t *testing.T) {
|
||||
"并发递增后总和应为 %v,实际为 %v", float64(goroutines)*increment, got.QuotaUsed)
|
||||
}
|
||||
|
||||
func (s *APIKeyRepoSuite) TestDeleteWithAudit_WritesAuditAndSoftDeletes() {
|
||||
func (s *APIKeyRepoSuite) TestDeleteWithAudit_TombstonesWithoutRetainingCredential() {
|
||||
user := s.mustCreateUser("delwithaudit@test.com")
|
||||
key := &service.APIKey{
|
||||
UserID: user.ID,
|
||||
@@ -571,18 +571,24 @@ func (s *APIKeyRepoSuite) TestDeleteWithAudit_WritesAuditAndSoftDeletes() {
|
||||
_, err := s.repo.GetByID(s.ctx, key.ID)
|
||||
s.Require().Error(err)
|
||||
|
||||
rows, qErr := s.client.QueryContext(s.ctx,
|
||||
`SELECT key, key_name, user_id, api_key_id FROM deleted_api_key_audits WHERE api_key_id = $1`, key.ID)
|
||||
s.Require().NoError(qErr)
|
||||
defer rows.Close()
|
||||
s.Require().True(rows.Next(), "expected one audit row")
|
||||
var auditKey, auditName string
|
||||
var auditUserID, auditAPIKeyID int64
|
||||
s.Require().NoError(rows.Scan(&auditKey, &auditName, &auditUserID, &auditAPIKeyID))
|
||||
s.Require().Equal("sk-del-audit-1", auditKey)
|
||||
s.Require().Equal("Audit Me", auditName)
|
||||
s.Require().Equal(user.ID, auditUserID)
|
||||
s.Require().Equal(key.ID, auditAPIKeyID)
|
||||
var tombstone string
|
||||
var deletedAt time.Time
|
||||
rows, err := s.repo.sql.QueryContext(s.ctx, `SELECT key, deleted_at FROM api_keys WHERE id = $1`, key.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().True(rows.Next())
|
||||
s.Require().NoError(rows.Scan(&tombstone, &deletedAt))
|
||||
s.Require().NoError(rows.Close())
|
||||
s.Require().NotEqual("sk-del-audit-1", tombstone)
|
||||
s.Require().Contains(tombstone, "__deleted__")
|
||||
|
||||
var auditCount int
|
||||
auditRows, err := s.repo.sql.QueryContext(s.ctx,
|
||||
`SELECT COUNT(*) FROM deleted_api_key_audits WHERE api_key_id = $1`, key.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().True(auditRows.Next())
|
||||
s.Require().NoError(auditRows.Scan(&auditCount))
|
||||
s.Require().NoError(auditRows.Close())
|
||||
s.Require().Zero(auditCount, "deleted credentials must not be retained")
|
||||
}
|
||||
|
||||
func (s *APIKeyRepoSuite) TestDeleteWithAudit_RepeatIsIdempotent() {
|
||||
|
||||
@@ -0,0 +1,120 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAuthCacheInvalidationTriggers_CoverSecurityMutationsOnly(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
suffix := time.Now().UnixNano()
|
||||
group := mustCreateGroup(t, integrationEntClient, &service.Group{
|
||||
Name: fmt.Sprintf("auth-outbox-group-%d", suffix), RateMultiplier: 1, IsExclusive: true,
|
||||
})
|
||||
user := mustCreateUser(t, integrationEntClient, &service.User{
|
||||
Email: fmt.Sprintf("auth-outbox-%d@example.com", suffix), Concurrency: 5,
|
||||
})
|
||||
groupID := group.ID
|
||||
keyValue := fmt.Sprintf("sk-auth-outbox-%d", suffix)
|
||||
apiKeyRepo := NewAPIKeyRepository(integrationEntClient, integrationDB)
|
||||
key := &service.APIKey{UserID: user.ID, GroupID: &groupID, Key: keyValue, Name: "outbox", Status: service.StatusActive}
|
||||
require.NoError(t, apiKeyRepo.Create(ctx, key))
|
||||
|
||||
sum := sha256.Sum256([]byte(keyValue))
|
||||
cacheKey := hex.EncodeToString(sum[:])
|
||||
clear := func() {
|
||||
_, err := integrationDB.ExecContext(ctx, "DELETE FROM auth_cache_invalidation_outbox WHERE cache_key = $1", cacheKey)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
count := func() int {
|
||||
var value int
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx,
|
||||
"SELECT COUNT(*) FROM auth_cache_invalidation_outbox WHERE cache_key = $1", cacheKey).Scan(&value))
|
||||
return value
|
||||
}
|
||||
clear()
|
||||
t.Cleanup(clear)
|
||||
t.Cleanup(func() {
|
||||
// Keep the shared integration database isolated for suites that assert
|
||||
// platform-wide group counts. The final clear cleanup runs after this one
|
||||
// and removes invalidations emitted by these hard deletes.
|
||||
_, err := integrationDB.ExecContext(ctx, "DELETE FROM user_allowed_groups WHERE user_id = $1 OR group_id = $2", user.ID, group.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM api_keys WHERE id = $1", key.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM users WHERE id = $1", user.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id = $1", group.ID)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
_, err := integrationDB.ExecContext(ctx, `
|
||||
UPDATE api_keys
|
||||
SET quota_used = quota_used + 1,
|
||||
usage_5h = usage_5h + 1,
|
||||
last_used_at = NOW()
|
||||
WHERE id = $1`, key.ID)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count(), "usage-only key updates must not enqueue")
|
||||
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE api_keys SET status = 'disabled' WHERE id = $1", key.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count(), "key disable must enqueue")
|
||||
clear()
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE api_keys SET status = 'active' WHERE id = $1", key.ID)
|
||||
require.NoError(t, err)
|
||||
clear()
|
||||
|
||||
userRepo := NewUserRepository(integrationEntClient, integrationDB)
|
||||
loadedUser, err := userRepo.GetByID(ctx, user.ID)
|
||||
require.NoError(t, err)
|
||||
loadedUser.Balance += 10
|
||||
require.NoError(t, userRepo.Update(ctx, loadedUser))
|
||||
require.Zero(t, count(), "balance update with unchanged allowed groups must not enqueue")
|
||||
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE users SET status = 'disabled' WHERE id = $1", user.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count(), "user disable must enqueue all active keys")
|
||||
clear()
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE users SET status = 'active' WHERE id = $1", user.ID)
|
||||
require.NoError(t, err)
|
||||
clear()
|
||||
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE groups SET name = name || '-cosmetic' WHERE id = $1", group.ID)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, count(), "cosmetic group update must not enqueue")
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE groups SET status = 'disabled' WHERE id = $1", group.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count(), "group disable must enqueue bound keys")
|
||||
clear()
|
||||
_, err = integrationDB.ExecContext(ctx, "UPDATE groups SET status = 'active' WHERE id = $1", group.ID)
|
||||
require.NoError(t, err)
|
||||
clear()
|
||||
|
||||
_, err = integrationDB.ExecContext(ctx,
|
||||
"INSERT INTO user_allowed_groups (user_id, group_id) VALUES ($1, $2)", user.ID, group.ID)
|
||||
require.NoError(t, err)
|
||||
clear()
|
||||
_, err = integrationDB.ExecContext(ctx,
|
||||
"DELETE FROM user_allowed_groups WHERE user_id = $1 AND group_id = $2", user.ID, group.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, count(), "exclusive-group revocation must enqueue")
|
||||
clear()
|
||||
|
||||
require.NoError(t, apiKeyRepo.DeleteWithAudit(ctx, key.ID))
|
||||
require.Equal(t, 1, count(), "tombstone delete must hash OLD.key exactly once")
|
||||
var stored string
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx,
|
||||
"SELECT cache_key FROM auth_cache_invalidation_outbox WHERE cache_key = $1 LIMIT 1", cacheKey).Scan(&stored))
|
||||
require.Equal(t, cacheKey, stored)
|
||||
require.NotContains(t, stored, keyValue)
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
type authCacheInvalidationOutboxRepository struct {
|
||||
db *sql.DB
|
||||
}
|
||||
|
||||
func NewAuthCacheInvalidationOutboxRepository(db *sql.DB) service.AuthCacheInvalidationOutboxRepository {
|
||||
return &authCacheInvalidationOutboxRepository{db: db}
|
||||
}
|
||||
|
||||
func (r *authCacheInvalidationOutboxRepository) Claim(ctx context.Context, workerID string, limit int, lease time.Duration) ([]service.AuthCacheInvalidationEvent, error) {
|
||||
if r == nil || r.db == nil {
|
||||
return nil, errors.New("nil auth cache invalidation outbox database")
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 100
|
||||
}
|
||||
leaseSeconds := int64(lease / time.Second)
|
||||
if leaseSeconds < 1 {
|
||||
leaseSeconds = 30
|
||||
}
|
||||
rows, err := r.db.QueryContext(ctx, `
|
||||
WITH candidates AS (
|
||||
SELECT id
|
||||
FROM auth_cache_invalidation_outbox
|
||||
WHERE available_at <= NOW()
|
||||
AND (claimed_at IS NULL OR claimed_at < NOW() - ($3 * INTERVAL '1 second'))
|
||||
ORDER BY id ASC
|
||||
LIMIT $2
|
||||
FOR UPDATE SKIP LOCKED
|
||||
)
|
||||
UPDATE auth_cache_invalidation_outbox AS o
|
||||
SET claimed_at = NOW(), claimed_by = $1
|
||||
FROM candidates AS c
|
||||
WHERE o.id = c.id
|
||||
RETURNING o.id, o.cache_key, o.attempts, o.delivery_stage, o.created_at
|
||||
`, workerID, limit, leaseSeconds)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
|
||||
events := make([]service.AuthCacheInvalidationEvent, 0, limit)
|
||||
for rows.Next() {
|
||||
var event service.AuthCacheInvalidationEvent
|
||||
if err := rows.Scan(&event.ID, &event.CacheKey, &event.Attempts, &event.Stage, &event.CreatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
event.CacheKey = strings.TrimSpace(event.CacheKey)
|
||||
events = append(events, event)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return events, nil
|
||||
}
|
||||
|
||||
func (r *authCacheInvalidationOutboxRepository) ScheduleSecondPass(ctx context.Context, id int64, workerID string, availableAt time.Time) error {
|
||||
result, err := r.db.ExecContext(ctx, `
|
||||
UPDATE auth_cache_invalidation_outbox
|
||||
SET delivery_stage = 1,
|
||||
available_at = $3,
|
||||
last_error = NULL,
|
||||
claimed_at = NULL,
|
||||
claimed_by = NULL
|
||||
WHERE id = $1 AND claimed_by = $2 AND delivery_stage = 0
|
||||
`, id, workerID, availableAt)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected != 1 {
|
||||
return fmt.Errorf("auth cache invalidation claim %d cannot schedule second pass", id)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authCacheInvalidationOutboxRepository) DeleteClaimed(ctx context.Context, id int64, workerID string) error {
|
||||
result, err := r.db.ExecContext(ctx, `
|
||||
DELETE FROM auth_cache_invalidation_outbox
|
||||
WHERE id = $1 AND claimed_by = $2
|
||||
`, id, workerID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected != 1 {
|
||||
return fmt.Errorf("auth cache invalidation claim %d is no longer owned by %s", id, workerID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authCacheInvalidationOutboxRepository) RetryClaimed(ctx context.Context, id int64, workerID string, availableAt time.Time, lastError string) error {
|
||||
result, err := r.db.ExecContext(ctx, `
|
||||
UPDATE auth_cache_invalidation_outbox
|
||||
SET attempts = attempts + 1,
|
||||
available_at = $3,
|
||||
last_error = $4,
|
||||
claimed_at = NULL,
|
||||
claimed_by = NULL
|
||||
WHERE id = $1 AND claimed_by = $2
|
||||
`, id, workerID, availableAt, lastError)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if affected != 1 {
|
||||
return fmt.Errorf("auth cache invalidation claim %d is no longer owned by %s", id, workerID)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *authCacheInvalidationOutboxRepository) Stats(ctx context.Context) (service.AuthCacheInvalidationOutboxStats, error) {
|
||||
var (
|
||||
stats service.AuthCacheInvalidationOutboxStats
|
||||
oldest sql.NullTime
|
||||
lastError sql.NullString
|
||||
)
|
||||
err := r.db.QueryRowContext(ctx, `
|
||||
SELECT COUNT(*), MIN(created_at), COALESCE(MAX(attempts), 0),
|
||||
(SELECT last_error
|
||||
FROM auth_cache_invalidation_outbox
|
||||
WHERE last_error IS NOT NULL
|
||||
ORDER BY available_at DESC, id DESC
|
||||
LIMIT 1)
|
||||
FROM auth_cache_invalidation_outbox
|
||||
`).Scan(&stats.Pending, &oldest, &stats.MaxAttempts, &lastError)
|
||||
if err != nil {
|
||||
return stats, err
|
||||
}
|
||||
if oldest.Valid {
|
||||
value := oldest.Time
|
||||
stats.OldestCreatedAt = &value
|
||||
}
|
||||
if lastError.Valid {
|
||||
stats.LastError = lastError.String
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
sqlmock "github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/migrations"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAuthCacheInvalidationOutboxRepository_ClaimUsesLeaseAndSkipLocked(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
created := time.Now().UTC()
|
||||
mock.ExpectQuery("(?s)claimed_at < NOW\\(\\) - .*FOR UPDATE SKIP LOCKED.*RETURNING").
|
||||
WithArgs("worker-a", 100, int64(30)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "cache_key", "attempts", "delivery_stage", "created_at"}).
|
||||
AddRow(int64(4), strings.Repeat("a", 64), 2, 1, created))
|
||||
|
||||
repo := NewAuthCacheInvalidationOutboxRepository(db)
|
||||
events, err := repo.Claim(context.Background(), "worker-a", 100, 30*time.Second)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, events, 1)
|
||||
require.Equal(t, int64(4), events[0].ID)
|
||||
require.Equal(t, 1, events[0].Stage)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestAuthCacheInvalidationOutboxRepository_ClaimIsBoundedByDefault(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
mock.ExpectQuery("(?s)FROM auth_cache_invalidation_outbox.*LIMIT \\$2.*SKIP LOCKED").
|
||||
WithArgs("worker", 100, int64(30)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id", "cache_key", "attempts", "delivery_stage", "created_at"}))
|
||||
repo := NewAuthCacheInvalidationOutboxRepository(db)
|
||||
_, err = repo.Claim(context.Background(), "worker", 0, 0)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestAuthCacheInvalidationOutboxRepository_ClaimOwnershipTransitions(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
repo := NewAuthCacheInvalidationOutboxRepository(db)
|
||||
|
||||
next := time.Now().UTC().Add(time.Minute)
|
||||
mock.ExpectExec("UPDATE auth_cache_invalidation_outbox").
|
||||
WithArgs(int64(1), "worker", next).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
require.NoError(t, repo.ScheduleSecondPass(context.Background(), 1, "worker", next))
|
||||
|
||||
retryAt := next.Add(time.Minute)
|
||||
mock.ExpectExec("UPDATE auth_cache_invalidation_outbox").
|
||||
WithArgs(int64(2), "worker", retryAt, "publish failed").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
require.NoError(t, repo.RetryClaimed(context.Background(), 2, "worker", retryAt, "publish failed"))
|
||||
|
||||
mock.ExpectExec("DELETE FROM auth_cache_invalidation_outbox").
|
||||
WithArgs(int64(3), "worker").
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
require.NoError(t, repo.DeleteClaimed(context.Background(), 3, "worker"))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestAuthCacheInvalidationOutboxRepository_RejectsLostClaim(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
mock.ExpectExec("DELETE FROM auth_cache_invalidation_outbox").
|
||||
WithArgs(int64(3), "old-worker").
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
repo := NewAuthCacheInvalidationOutboxRepository(db)
|
||||
err = repo.DeleteClaimed(context.Background(), 3, "old-worker")
|
||||
require.ErrorContains(t, err, "no longer owned")
|
||||
}
|
||||
|
||||
func TestAuthCacheInvalidationOutboxRepository_StatsExposeDurableLagAndFailures(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
oldest := time.Now().UTC().Add(-time.Minute)
|
||||
mock.ExpectQuery("(?s)SELECT COUNT\\(\\*\\), MIN\\(created_at\\), COALESCE\\(MAX\\(attempts\\), 0\\)").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"count", "min", "max", "last_error"}).AddRow(5, oldest, 7, "redis down"))
|
||||
repo := NewAuthCacheInvalidationOutboxRepository(db)
|
||||
stats, err := repo.Stats(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(5), stats.Pending)
|
||||
require.Equal(t, 7, stats.MaxAttempts)
|
||||
require.Equal(t, "redis down", stats.LastError)
|
||||
require.NotNil(t, stats.OldestCreatedAt)
|
||||
}
|
||||
|
||||
func TestAuthCacheInvalidationMigration_SecurityCoverageAndNoPlaintextPayload(t *testing.T) {
|
||||
content, err := migrations.FS.ReadFile("184_auth_cache_invalidation_outbox.sql")
|
||||
require.NoError(t, err)
|
||||
sqlText := string(content)
|
||||
for _, required := range []string{
|
||||
"encode(sha256(convert_to(raw_key, 'UTF8')), 'hex')",
|
||||
"OLD.key", "OLD.status", "OLD.deleted_at", "OLD.user_id", "OLD.group_id",
|
||||
"OLD.ip_whitelist", "OLD.ip_blacklist", "OLD.expires_at",
|
||||
"trg_users_auth_cache_invalidation", "trg_groups_auth_cache_invalidation",
|
||||
"trg_user_allowed_groups_auth_cache_invalidation", "FOR EACH ROW",
|
||||
"delivery_stage", "claimed_at", "available_at",
|
||||
} {
|
||||
require.Contains(t, sqlText, required)
|
||||
}
|
||||
require.NotContains(t, sqlText, "quota_used IS DISTINCT")
|
||||
require.NotContains(t, sqlText, "last_used_at IS DISTINCT")
|
||||
|
||||
plaintext := "sk-plaintext-must-not-be-stored"
|
||||
sum := sha256.Sum256([]byte(plaintext))
|
||||
require.Len(t, hex.EncodeToString(sum[:]), 64)
|
||||
require.NotContains(t, sqlText, plaintext)
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -18,6 +19,7 @@ import (
|
||||
type githubReleaseClient struct {
|
||||
httpClient *http.Client
|
||||
downloadHTTPClient *http.Client
|
||||
updateGitHubToken string
|
||||
}
|
||||
|
||||
type githubReleaseClientError struct {
|
||||
@@ -43,6 +45,8 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi
|
||||
}
|
||||
sharedClient = &http.Client{Timeout: 30 * time.Second}
|
||||
}
|
||||
apiClient := cloneHTTPClient(sharedClient)
|
||||
apiClient.CheckRedirect = githubAPICheckRedirect(apiClient.CheckRedirect)
|
||||
|
||||
// 下载客户端需要更长的超时时间
|
||||
downloadClient, err := httpclient.GetClient(httpclient.Options{
|
||||
@@ -56,13 +60,50 @@ func NewGitHubReleaseClient(proxyURL string, allowDirectOnProxyError bool) servi
|
||||
}
|
||||
downloadClient = &http.Client{Timeout: 10 * time.Minute}
|
||||
}
|
||||
downloadClient = cloneHTTPClient(downloadClient)
|
||||
|
||||
return &githubReleaseClient{
|
||||
httpClient: sharedClient,
|
||||
httpClient: apiClient,
|
||||
downloadHTTPClient: downloadClient,
|
||||
updateGitHubToken: os.Getenv("UPDATE_GITHUB_TOKEN"),
|
||||
}
|
||||
}
|
||||
|
||||
func cloneHTTPClient(client *http.Client) *http.Client {
|
||||
cloned := *client
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func isGitHubAPIURL(url *url.URL) bool {
|
||||
return url != nil && strings.EqualFold(url.Scheme, "https") && url.User == nil &&
|
||||
strings.EqualFold(url.Host, "api.github.com")
|
||||
}
|
||||
|
||||
func githubAPICheckRedirect(previous func(*http.Request, []*http.Request) error) func(*http.Request, []*http.Request) error {
|
||||
return func(req *http.Request, via []*http.Request) error {
|
||||
if !isGitHubAPIURL(req.URL) {
|
||||
req.Header.Del("Authorization")
|
||||
}
|
||||
if previous != nil {
|
||||
return previous(req, via)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (c *githubReleaseClient) newAPIRequest(ctx context.Context, url string) (*http.Request, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
if c.updateGitHubToken != "" && isGitHubAPIURL(req.URL) {
|
||||
req.Header.Set("Authorization", "Bearer "+c.updateGitHubToken)
|
||||
}
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func (c *githubReleaseClientError) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
|
||||
return nil, c.err
|
||||
}
|
||||
@@ -82,12 +123,10 @@ func (c *githubReleaseClientError) FetchChecksumFile(ctx context.Context, url st
|
||||
func (c *githubReleaseClient) FetchLatestRelease(ctx context.Context, repo string) (*service.GitHubRelease, error) {
|
||||
url := fmt.Sprintf("https://api.github.com/repos/%s/releases/latest", repo)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
req, err := c.newAPIRequest(ctx, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
@@ -116,12 +155,10 @@ func (c *githubReleaseClient) FetchRecentReleases(ctx context.Context, repo stri
|
||||
}
|
||||
url := fmt.Sprintf("https://api.github.com/repos/%s/releases?per_page=%d", repo, perPage)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
req, err := c.newAPIRequest(ctx, url)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Accept", "application/vnd.github.v3+json")
|
||||
req.Header.Set("User-Agent", "Sub2API-Updater")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
|
||||
@@ -3,6 +3,7 @@ package repository
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
@@ -47,6 +48,97 @@ func newTestGitHubReleaseClient() *githubReleaseClient {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientAPIRequestAuthorization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
wantAuth string
|
||||
}{
|
||||
{name: "exact HTTPS authority", url: "https://api.github.com/repos/test/repo", wantAuth: "Bearer update-secret"},
|
||||
{name: "HTTP", url: "http://api.github.com/repos/test/repo"},
|
||||
{name: "subdomain", url: "https://sub.api.github.com/repos/test/repo"},
|
||||
{name: "userinfo", url: "https://user@api.github.com/repos/test/repo"},
|
||||
{name: "explicit default port", url: "https://api.github.com:443/repos/test/repo"},
|
||||
{name: "custom port", url: "https://api.github.com:8443/repos/test/repo"},
|
||||
{name: "different host", url: "https://github.com/test/repo"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := newTestGitHubReleaseClient()
|
||||
client.updateGitHubToken = "update-secret"
|
||||
req, err := client.newAPIRequest(context.Background(), tt.url)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.wantAuth, req.Header.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
|
||||
client := newTestGitHubReleaseClient()
|
||||
req, err := client.newAPIRequest(context.Background(), "https://api.github.com/repos/test/repo")
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, req.Header.Get("Authorization"))
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientRedirectAuthorization(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
url string
|
||||
wantAuth string
|
||||
}{
|
||||
{name: "same HTTPS authority", url: "https://api.github.com/redirected", wantAuth: "Bearer update-secret"},
|
||||
{name: "HTTP", url: "http://api.github.com/redirected"},
|
||||
{name: "subdomain", url: "https://sub.api.github.com/redirected"},
|
||||
{name: "userinfo", url: "https://user@api.github.com/redirected"},
|
||||
{name: "custom port", url: "https://api.github.com:8443/redirected"},
|
||||
{name: "different host", url: "https://example.com/redirected"},
|
||||
}
|
||||
|
||||
checkRedirect := githubAPICheckRedirect(nil)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodGet, tt.url, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer update-secret")
|
||||
|
||||
require.NoError(t, checkRedirect(req, nil))
|
||||
require.Equal(t, tt.wantAuth, req.Header.Get("Authorization"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGitHubReleaseClientDoesNotAuthorizeDownloads(t *testing.T) {
|
||||
client := newTestGitHubReleaseClient()
|
||||
client.updateGitHubToken = "update-secret"
|
||||
|
||||
var headers []http.Header
|
||||
transport := githubReleaseRoundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
headers = append(headers, req.Header.Clone())
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader("checksum")),
|
||||
Request: req,
|
||||
}, nil
|
||||
})
|
||||
client.httpClient.Transport = transport
|
||||
client.downloadHTTPClient.Transport = transport
|
||||
|
||||
dest := filepath.Join(t.TempDir(), "asset")
|
||||
require.NoError(t, client.DownloadFile(context.Background(), "https://objects.githubusercontent.com/asset", dest, 100))
|
||||
_, err := client.FetchChecksumFile(context.Background(), "https://github.com/test/repo/releases/download/v1/checksums.txt")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, headers, 2)
|
||||
for _, header := range headers {
|
||||
require.Empty(t, header.Get("Authorization"))
|
||||
}
|
||||
}
|
||||
|
||||
type githubReleaseRoundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f githubReleaseRoundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
|
||||
func (s *GitHubReleaseServiceSuite) SetupTest() {
|
||||
s.tempDir = s.T().TempDir()
|
||||
}
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -323,7 +324,7 @@ func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http.
|
||||
}
|
||||
|
||||
body, ok := bufferSmallResponseBody(resp, grokFallbackBodyLimit)
|
||||
if !ok || !bytes.Contains(bytes.ToLower(body), []byte("access denied")) {
|
||||
if !ok || !isGrokCLICompatibilityAccessDenied(body) {
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
@@ -350,6 +351,22 @@ func (t *grokAccessDeniedFallbackTransport) RoundTrip(req *http.Request) (*http.
|
||||
return fallbackResp, nil
|
||||
}
|
||||
|
||||
func isGrokCLICompatibilityAccessDenied(body []byte) bool {
|
||||
lower := bytes.ToLower(body)
|
||||
if bytes.Contains(lower, []byte("access denied")) {
|
||||
return true
|
||||
}
|
||||
var payload struct {
|
||||
Code string `json:"code"`
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &payload); err != nil || !strings.EqualFold(strings.TrimSpace(payload.Code), "permission_denied") {
|
||||
return false
|
||||
}
|
||||
const chatEndpointDeniedPrefix = "access to the chat endpoint is denied. please ensure you're using the correct credentials. if you believe this is a mistake, please"
|
||||
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(payload.Error)), chatEndpointDeniedPrefix)
|
||||
}
|
||||
|
||||
func isGrokCLIAccessDeniedFallbackCandidate(req *http.Request, resp *http.Response) bool {
|
||||
return req != nil && req.URL != nil && req.GetBody != nil && resp != nil &&
|
||||
resp.StatusCode == http.StatusForbidden &&
|
||||
|
||||
@@ -311,6 +311,119 @@ func TestHTTPUpstreamDoFallsBackToOfficialGrokAPIOnCLIAccessDenied(t *testing.T)
|
||||
require.Empty(t, fallbackHeaders.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestGrokAccessDeniedFallbackRecognizesChatEndpointPermissionDenied(t *testing.T) {
|
||||
var hosts []string
|
||||
transport := &grokAccessDeniedFallbackTransport{
|
||||
base: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
hosts = append(hosts, req.URL.Hostname())
|
||||
if req.URL.Hostname() == grokCLIProxyHost {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusForbidden,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`{"code":"permission_denied","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`,
|
||||
)),
|
||||
Request: req,
|
||||
}, nil
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(`{"id":"response-ok"}`)),
|
||||
Request: req,
|
||||
}, nil
|
||||
}),
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", strings.NewReader(`{"model":"grok-4.5"}`))
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer oauth-token")
|
||||
req.Header.Set("X-XAI-Token-Auth", "xai-grok-cli")
|
||||
|
||||
resp, err := transport.RoundTrip(req)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, http.StatusOK, resp.StatusCode)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
require.Equal(t, []string{grokCLIProxyHost, grokOfficialAPIHost}, hosts)
|
||||
}
|
||||
|
||||
func TestIsGrokCLICompatibilityAccessDenied(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
want bool
|
||||
}{
|
||||
{name: "legacy compatibility wording", body: `{"error":"Access denied"}`, want: true},
|
||||
{
|
||||
name: "observed chat endpoint permission denial",
|
||||
body: `{"code":"permission_denied","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`,
|
||||
want: true,
|
||||
},
|
||||
{
|
||||
name: "entitlement denial using the same broad terms",
|
||||
body: `{"code":"permission_denied","error":"Access to the chat endpoint is denied because a subscription is required"}`,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "different permission denied endpoint",
|
||||
body: `{"code":"permission_denied","error":"Access to the billing endpoint is denied."}`,
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "wrong structured error code",
|
||||
body: `{"code":"subscription_required","error":"Access to the chat endpoint is denied. Please ensure you're using the correct credentials. If you believe this is a mistake, please contact support."}`,
|
||||
want: false,
|
||||
},
|
||||
{name: "malformed response", body: `permission_denied: chat endpoint denied`, want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, isGrokCLICompatibilityAccessDenied([]byte(tt.body)))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsGrokCLIAccessDeniedFallbackCandidateRequiresAuthenticatedReplayableCLI403(t *testing.T) {
|
||||
newRequest := func() *http.Request {
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", strings.NewReader(`{"model":"grok-4.5"}`))
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Authorization", "Bearer oauth-token")
|
||||
req.Header.Set("X-XAI-Token-Auth", "xai-grok-cli")
|
||||
return req
|
||||
}
|
||||
newResponse := func() *http.Response { return &http.Response{StatusCode: http.StatusForbidden} }
|
||||
|
||||
t.Run("valid candidate", func(t *testing.T) {
|
||||
require.True(t, isGrokCLIAccessDeniedFallbackCandidate(newRequest(), newResponse()))
|
||||
})
|
||||
t.Run("non CLI host", func(t *testing.T) {
|
||||
req := newRequest()
|
||||
req.URL.Host = "api.x.ai"
|
||||
require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse()))
|
||||
})
|
||||
t.Run("missing CLI identity", func(t *testing.T) {
|
||||
req := newRequest()
|
||||
req.Header.Del("X-XAI-Token-Auth")
|
||||
require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse()))
|
||||
})
|
||||
t.Run("missing bearer authentication", func(t *testing.T) {
|
||||
req := newRequest()
|
||||
req.Header.Del("Authorization")
|
||||
require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse()))
|
||||
})
|
||||
t.Run("non forbidden response", func(t *testing.T) {
|
||||
resp := newResponse()
|
||||
resp.StatusCode = http.StatusUnauthorized
|
||||
require.False(t, isGrokCLIAccessDeniedFallbackCandidate(newRequest(), resp))
|
||||
})
|
||||
t.Run("non replayable request", func(t *testing.T) {
|
||||
req := newRequest()
|
||||
req.GetBody = nil
|
||||
require.False(t, isGrokCLIAccessDeniedFallbackCandidate(req, newResponse()))
|
||||
})
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamDoDoesNotFallbackForGrokEntitlementDenial(t *testing.T) {
|
||||
transport := &grokAccessDeniedFallbackTransport{
|
||||
base: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
|
||||
@@ -128,23 +128,31 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
|
||||
// 获取分布式锁,确保多实例部署时只有一个实例执行迁移。
|
||||
// 这是 PostgreSQL 特有的 Advisory Lock 机制。
|
||||
if err := pgAdvisoryLock(ctx, db); err != nil {
|
||||
lockConn, err := db.Conn(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("acquire migrations lock connection: %w", err)
|
||||
}
|
||||
defer func() { _ = lockConn.Close() }()
|
||||
if err := pgAdvisoryLock(ctx, lockConn); err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() {
|
||||
// 无论迁移是否成功,都要释放锁。
|
||||
// 使用 context.Background() 确保即使原 ctx 已取消也能释放锁。
|
||||
_ = pgAdvisoryUnlock(context.Background(), db)
|
||||
// 独立超时确保原 ctx 取消后仍会尝试释放,但数据库链路异常不会
|
||||
// 无限阻塞进程退出。
|
||||
unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = pgAdvisoryUnlock(unlockCtx, lockConn)
|
||||
}()
|
||||
|
||||
// 创建迁移记录表(如果不存在)。
|
||||
// 该表记录所有已应用的迁移及其校验和。
|
||||
if _, err := db.ExecContext(ctx, schemaMigrationsTableDDL); err != nil {
|
||||
if _, err := lockConn.ExecContext(ctx, schemaMigrationsTableDDL); err != nil {
|
||||
return fmt.Errorf("create schema_migrations: %w", err)
|
||||
}
|
||||
|
||||
// 自动对齐 Atlas 基线(如果检测到 legacy schema_migrations 且缺失 atlas_schema_revisions)。
|
||||
if err := ensureAtlasBaselineAligned(ctx, db, fsys); err != nil {
|
||||
if err := ensureAtlasBaselineAligned(ctx, lockConn, fsys); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -175,7 +183,7 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
|
||||
// 检查该迁移是否已经应用
|
||||
var existing string
|
||||
rowErr := db.QueryRowContext(ctx, "SELECT checksum FROM schema_migrations WHERE filename = $1", name).Scan(&existing)
|
||||
rowErr := lockConn.QueryRowContext(ctx, "SELECT checksum FROM schema_migrations WHERE filename = $1", name).Scan(&existing)
|
||||
if rowErr == nil {
|
||||
// 迁移已应用,验证校验和是否匹配
|
||||
if existing != checksum {
|
||||
@@ -207,7 +215,7 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
}
|
||||
|
||||
if nonTx {
|
||||
if err := prepareNonTransactionalMigration(ctx, db, name); err != nil {
|
||||
if err := prepareNonTransactionalMigration(ctx, lockConn, name); err != nil {
|
||||
return fmt.Errorf("prepare migration %s: %w", name, err)
|
||||
}
|
||||
|
||||
@@ -222,18 +230,18 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
if stripSQLLineComment(trimmed) == "" {
|
||||
continue
|
||||
}
|
||||
if _, err := db.ExecContext(ctx, trimmed); err != nil {
|
||||
if _, err := lockConn.ExecContext(ctx, trimmed); err != nil {
|
||||
return fmt.Errorf("apply migration %s (non-tx statement %d): %w", name, i+1, err)
|
||||
}
|
||||
}
|
||||
if _, err := db.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil {
|
||||
if _, err := lockConn.ExecContext(ctx, "INSERT INTO schema_migrations (filename, checksum) VALUES ($1, $2)", name, checksum); err != nil {
|
||||
return fmt.Errorf("record migration %s (non-tx): %w", name, err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// 默认迁移在事务中执行,确保原子性:要么完全成功,要么完全回滚。
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
tx, err := lockConn.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin migration %s: %w", name, err)
|
||||
}
|
||||
@@ -260,7 +268,14 @@ func applyMigrationsFS(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name string) error {
|
||||
type migrationConnection interface {
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error)
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error)
|
||||
}
|
||||
|
||||
func prepareNonTransactionalMigration(ctx context.Context, db migrationConnection, name string) error {
|
||||
switch name {
|
||||
case paymentOrdersOutTradeNoUniqueMigration:
|
||||
return preparePaymentOrdersOutTradeNoUniqueMigration(ctx, db)
|
||||
@@ -273,7 +288,7 @@ func prepareNonTransactionalMigration(ctx context.Context, db *sql.DB, name stri
|
||||
}
|
||||
}
|
||||
|
||||
func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql.DB) error {
|
||||
func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db migrationConnection) error {
|
||||
duplicates, err := findDuplicatePaymentOrderOutTradeNos(ctx, db)
|
||||
if err != nil {
|
||||
return fmt.Errorf("precheck duplicate out_trade_no: %w", err)
|
||||
@@ -289,7 +304,7 @@ func preparePaymentOrdersOutTradeNoUniqueMigration(ctx context.Context, db *sql.
|
||||
return dropInvalidIndexIfPresent(ctx, db, paymentOrdersOutTradeNoUniqueIndex)
|
||||
}
|
||||
|
||||
func dropInvalidIndexIfPresent(ctx context.Context, db *sql.DB, indexName string) error {
|
||||
func dropInvalidIndexIfPresent(ctx context.Context, db migrationConnection, indexName string) error {
|
||||
invalid, err := indexIsInvalid(ctx, db, indexName)
|
||||
if err != nil {
|
||||
return fmt.Errorf("check invalid index %s: %w", indexName, err)
|
||||
@@ -304,7 +319,7 @@ func dropInvalidIndexIfPresent(ctx context.Context, db *sql.DB, indexName string
|
||||
return nil
|
||||
}
|
||||
|
||||
func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db *sql.DB) ([]string, error) {
|
||||
func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db migrationConnection) ([]string, error) {
|
||||
rows, err := db.QueryContext(ctx, `
|
||||
SELECT out_trade_no, COUNT(*) AS duplicate_count
|
||||
FROM payment_orders
|
||||
@@ -336,7 +351,7 @@ func findDuplicatePaymentOrderOutTradeNos(ctx context.Context, db *sql.DB) ([]st
|
||||
return duplicates, nil
|
||||
}
|
||||
|
||||
func indexIsInvalid(ctx context.Context, db *sql.DB, indexName string) (bool, error) {
|
||||
func indexIsInvalid(ctx context.Context, db migrationConnection, indexName string) (bool, error) {
|
||||
var invalid bool
|
||||
err := db.QueryRowContext(ctx, `
|
||||
SELECT EXISTS (
|
||||
@@ -352,7 +367,7 @@ func indexIsInvalid(ctx context.Context, db *sql.DB, indexName string) (bool, er
|
||||
return invalid, err
|
||||
}
|
||||
|
||||
func ensureAtlasBaselineAligned(ctx context.Context, db *sql.DB, fsys fs.FS) error {
|
||||
func ensureAtlasBaselineAligned(ctx context.Context, db migrationConnection, fsys fs.FS) error {
|
||||
hasLegacy, err := tableExists(ctx, db, "schema_migrations")
|
||||
if err != nil {
|
||||
return fmt.Errorf("check schema_migrations: %w", err)
|
||||
@@ -393,7 +408,7 @@ func ensureAtlasBaselineAligned(ctx context.Context, db *sql.DB, fsys fs.FS) err
|
||||
return nil
|
||||
}
|
||||
|
||||
func tableExists(ctx context.Context, db *sql.DB, tableName string) (bool, error) {
|
||||
func tableExists(ctx context.Context, db migrationConnection, tableName string) (bool, error) {
|
||||
var exists bool
|
||||
err := db.QueryRowContext(ctx, `
|
||||
SELECT EXISTS (
|
||||
@@ -524,7 +539,12 @@ func stripSQLLineComment(s string) string {
|
||||
// pgAdvisoryLock 获取 PostgreSQL Advisory Lock。
|
||||
// Advisory Lock 是一种轻量级的锁机制,不与任何特定的数据库对象关联。
|
||||
// 它非常适合用于应用层面的分布式锁场景,如迁移序列化。
|
||||
func pgAdvisoryLock(ctx context.Context, db *sql.DB) error {
|
||||
type advisoryLockConnection interface {
|
||||
QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row
|
||||
ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error)
|
||||
}
|
||||
|
||||
func pgAdvisoryLock(ctx context.Context, db advisoryLockConnection) error {
|
||||
ticker := time.NewTicker(migrationsLockRetryInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
@@ -546,7 +566,7 @@ func pgAdvisoryLock(ctx context.Context, db *sql.DB) error {
|
||||
|
||||
// pgAdvisoryUnlock 释放 PostgreSQL Advisory Lock。
|
||||
// 必须在获取锁后确保释放,否则会阻塞其他实例的迁移操作。
|
||||
func pgAdvisoryUnlock(ctx context.Context, db *sql.DB) error {
|
||||
func pgAdvisoryUnlock(ctx context.Context, db advisoryLockConnection) error {
|
||||
_, err := db.ExecContext(ctx, "SELECT pg_advisory_unlock($1)", migrationsAdvisoryLockID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("release migrations lock: %w", err)
|
||||
|
||||
@@ -275,6 +275,9 @@ func TestApplyMigrationsFS_TransactionalMigration(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
// The advisory lock and all migration work must share one session. This also
|
||||
// proves startup cannot self-deadlock when deployments cap the pool at one.
|
||||
db.SetMaxOpenConns(1)
|
||||
|
||||
prepareMigrationsBootstrapExpectations(mock)
|
||||
mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1").
|
||||
|
||||
@@ -5,11 +5,32 @@ package repository
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMigrationsRunner_ConcurrentInstancesSerializeOnSessionLock(t *testing.T) {
|
||||
const instances = 2
|
||||
errorsByInstance := make([]error, instances)
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < instances; i++ {
|
||||
wg.Add(1)
|
||||
go func(index int) {
|
||||
defer wg.Done()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
errorsByInstance[index] = ApplyMigrations(ctx, integrationDB)
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
for i, err := range errorsByInstance {
|
||||
require.NoErrorf(t, err, "migration instance %d", i)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
|
||||
tx := testTx(t)
|
||||
|
||||
@@ -111,6 +132,13 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
|
||||
requireColumn(t, tx, "ops_system_logs", "api_key_id", "bigint", 0, true)
|
||||
requireIndex(t, tx, "ops_system_logs", "idx_ops_system_logs_api_key_id_created_at")
|
||||
|
||||
// Bounded ingress rejection security aggregates.
|
||||
requireColumn(t, tx, "ops_ingress_reject_aggregates", "bucket_start", "timestamp with time zone", 0, false)
|
||||
requireColumn(t, tx, "ops_ingress_reject_aggregates", "client_ip", "inet", 0, false)
|
||||
requireColumn(t, tx, "ops_ingress_reject_aggregates", "request_count", "bigint", 0, false)
|
||||
requireIndex(t, tx, "ops_ingress_reject_aggregates", "idx_ops_ingress_reject_aggregates_bucket")
|
||||
requireIndex(t, tx, "ops_ingress_reject_aggregates", "idx_ops_ingress_reject_aggregates_ip_bucket")
|
||||
|
||||
// user_allowed_groups table should exist
|
||||
var uagRegclass sql.NullString
|
||||
require.NoError(t, tx.QueryRowContext(context.Background(), "SELECT to_regclass('public.user_allowed_groups')").Scan(&uagRegclass))
|
||||
|
||||
@@ -137,26 +137,17 @@ func TestBuildOpsErrorLogsWhere_CyberPolicyStatusExemption(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildOpsErrorLogsWhere_MatchDeletedKeyOwner(t *testing.T) {
|
||||
func TestBuildOpsErrorLogsWhere_UserOwnershipIsDirectOnly(t *testing.T) {
|
||||
uid := int64(42)
|
||||
|
||||
// 开关开启 → 归属放宽为 OR(user_id 或 deleted_key_owner_user_id),且共用同一占位符
|
||||
on := &service.OpsErrorLogFilter{UserID: &uid, MatchDeletedKeyOwner: true}
|
||||
whereOn, argsOn := buildOpsErrorLogsWhere(on)
|
||||
if !strings.Contains(whereOn, "(e.user_id = $1 OR e.deleted_key_owner_user_id = $1)") {
|
||||
t.Fatalf("MatchDeletedKeyOwner=true should widen to OR, got: %s", whereOn)
|
||||
filter := &service.OpsErrorLogFilter{UserID: &uid}
|
||||
where, args := buildOpsErrorLogsWhere(filter)
|
||||
if !strings.Contains(where, "e.user_id = $1") {
|
||||
t.Fatalf("user scope should match user_id exactly, got: %s", where)
|
||||
}
|
||||
if len(argsOn) != 1 || argsOn[0] != uid {
|
||||
t.Fatalf("expected single reused arg %d, got %v", uid, argsOn)
|
||||
if len(args) != 1 || args[0] != uid {
|
||||
t.Fatalf("expected user id arg %d, got %v", uid, args)
|
||||
}
|
||||
|
||||
// 开关关闭(默认)→ 仅精确 user_id,绝不出现 deleted_key_owner_user_id(admin 回归)
|
||||
off := &service.OpsErrorLogFilter{UserID: &uid}
|
||||
whereOff, _ := buildOpsErrorLogsWhere(off)
|
||||
if !strings.Contains(whereOff, "e.user_id = $1") {
|
||||
t.Fatalf("default should match user_id exactly, got: %s", whereOff)
|
||||
}
|
||||
if strings.Contains(whereOff, "deleted_key_owner_user_id") {
|
||||
t.Fatalf("default must NOT include deleted_key_owner_user_id, got: %s", whereOff)
|
||||
if strings.Contains(where, "deleted_key_owner_user_id") {
|
||||
t.Fatalf("user ownership must not depend on deleted-key attribution: %s", where)
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user