mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 20:12:39 +08:00
Compare commits
69
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0f9d460dc6 | ||
|
|
6e315616ed | ||
|
|
d97ccc952f | ||
|
|
33d9e72f90 | ||
|
|
d1aac6b989 | ||
|
|
bbba01dae9 | ||
|
|
b8dece9000 | ||
|
|
3040209f20 | ||
|
|
f5aca74623 | ||
|
|
432a6a461f | ||
|
|
c2d5bbd93e | ||
|
|
ae501cc220 | ||
|
|
99efaa477f | ||
|
|
d9d4333795 | ||
|
|
8b9788eeb4 | ||
|
|
66a9e2e6bb | ||
|
|
458b92abd4 | ||
|
|
5106065716 | ||
|
|
46f5f5e8ba | ||
|
|
2529371bab | ||
|
|
43c4c882fd | ||
|
|
719013053a | ||
|
|
f3dc0835df | ||
|
|
5637f469fa | ||
|
|
34fe845f5a | ||
|
|
fa49647c6d | ||
|
|
9eb394e3e9 | ||
|
|
2d8a386811 | ||
|
|
bb97ab364a | ||
|
|
7225538358 | ||
|
|
0cecf3c367 | ||
|
|
1aa34d478d | ||
|
|
1b1039f4fd | ||
|
|
d6adebd22d | ||
|
|
53158587a6 | ||
|
|
ad05eda157 | ||
|
|
fdce5b8d01 | ||
|
|
00940743d2 | ||
|
|
2e3c2cf592 | ||
|
|
42bc7f6cff | ||
|
|
96f4c115c9 | ||
|
|
327c322184 | ||
|
|
48b7109588 | ||
|
|
4f0c079e27 | ||
|
|
b0cfff9826 | ||
|
|
017bbcb901 | ||
|
|
9ecb340822 | ||
|
|
a0f41f95a0 | ||
|
|
b5efbe3f4c | ||
|
|
e6d191a83f | ||
|
|
c91bb6124a | ||
|
|
9688571a83 | ||
|
|
514ad0aa8e | ||
|
|
dc2efc6b9a | ||
|
|
1872d44b58 | ||
|
|
196b47757a | ||
|
|
a60a29549f | ||
|
|
475c0839c6 | ||
|
|
b886fb8edb | ||
|
|
652a7d74f9 | ||
|
|
52c123e6f4 | ||
|
|
43f88c99aa | ||
|
|
b1ce3c6d59 | ||
|
|
d359196adc | ||
|
|
cf19f30d26 | ||
|
|
58d605083e | ||
|
|
0da18e4a5e | ||
|
|
b31b435091 | ||
|
|
d1ba57977d |
@@ -0,0 +1,107 @@
|
||||
# Security Policy
|
||||
|
||||
[English](#english) | [中文](#中文)
|
||||
|
||||
## English
|
||||
|
||||
### Reporting a Vulnerability
|
||||
|
||||
**Please do not report security vulnerabilities through public GitHub issues, pull requests, or discussions.**
|
||||
|
||||
Report them privately through GitHub's private vulnerability reporting:
|
||||
|
||||
**[Report a vulnerability](https://github.com/Wei-Shaw/sub2api/security/advisories/new)**
|
||||
|
||||
(Repository → **Security** tab → **Report a vulnerability**.) Only the maintainers can see the report, and the advisory becomes the place to discuss, fix, and eventually publish the issue.
|
||||
|
||||
If you cannot use GitHub private vulnerability reporting, email **dev@sub2api.org** instead, with `[Security]` in the subject line.
|
||||
|
||||
A useful report includes:
|
||||
|
||||
- The affected component and version (release tag or commit)
|
||||
- Impact: what an attacker can do and which privileges they need
|
||||
- Minimal reproduction steps or a proof of concept
|
||||
- Any non-default configuration required to trigger the issue
|
||||
- A suggested fix, if you have one
|
||||
|
||||
### Supported Versions
|
||||
|
||||
Sub2API releases frequently, and fixes ship in a new release rather than being backported.
|
||||
|
||||
| Version | Supported |
|
||||
| ------- | --------- |
|
||||
| Latest release | ✅ |
|
||||
| Older releases | ❌ |
|
||||
|
||||
Please confirm that the issue reproduces on the latest release or on `main` before reporting.
|
||||
|
||||
### Scope
|
||||
|
||||
In scope: the Sub2API backend, frontend, official Docker images, and deployment files in this repository.
|
||||
|
||||
Usually out of scope:
|
||||
|
||||
- Issues that require already-compromised admin credentials or server access
|
||||
- Risks arising from explicitly weakening documented security settings (for example, disabling `security.url_allowlist`)
|
||||
- Vulnerabilities in upstream AI providers or third-party dependencies with no Sub2API-specific impact (please report those upstream)
|
||||
- Missing best-practice headers or scanner output without a demonstrated impact
|
||||
|
||||
### Disclosure Process
|
||||
|
||||
1. We acknowledge the report and assess its validity and severity.
|
||||
2. We develop a fix privately in the advisory and may ask you to verify it.
|
||||
3. We publish a release containing the fix, then publish the GitHub Security Advisory (with a CVE when applicable), crediting you unless you prefer to stay anonymous.
|
||||
|
||||
Please keep details private until the advisory is published. This is a volunteer-maintained project, so response times are best effort.
|
||||
|
||||
## 中文
|
||||
|
||||
### 报告漏洞
|
||||
|
||||
**请不要通过公开的 Issue、Pull Request 或 Discussion 报告安全漏洞。**
|
||||
|
||||
请使用 GitHub 私密漏洞报告功能提交:
|
||||
|
||||
**[提交漏洞报告](https://github.com/Wei-Shaw/sub2api/security/advisories/new)**
|
||||
|
||||
(仓库 → **Security** 标签页 → **Report a vulnerability**。)报告仅维护者可见,后续的讨论、修复与公告发布都在该安全公告中进行。
|
||||
|
||||
如果无法使用 GitHub 私密漏洞报告,也可以发送邮件至 **dev@sub2api.org**,邮件标题请注明 `[Security]`。
|
||||
|
||||
报告中建议包含:
|
||||
|
||||
- 受影响的组件与版本(Release tag 或 commit)
|
||||
- 影响:攻击者能做什么、需要什么权限
|
||||
- 最小复现步骤或 PoC
|
||||
- 触发问题所需的非默认配置
|
||||
- 修复建议(如有)
|
||||
|
||||
### 支持的版本
|
||||
|
||||
Sub2API 发版频繁,安全修复随新版本发布,不向旧版本回移。
|
||||
|
||||
| 版本 | 是否支持 |
|
||||
| ---- | -------- |
|
||||
| 最新 Release | ✅ |
|
||||
| 更早的版本 | ❌ |
|
||||
|
||||
报告前请先确认问题在最新 Release 或 `main` 上仍可复现。
|
||||
|
||||
### 范围
|
||||
|
||||
范围内:本仓库中的 Sub2API 后端、前端、官方 Docker 镜像与部署文件。
|
||||
|
||||
通常不在范围内:
|
||||
|
||||
- 需要事先已掌握管理员凭据或服务器权限才能利用的问题
|
||||
- 因主动关闭文档中的安全配置(如关闭 `security.url_allowlist`)而产生的风险
|
||||
- 上游 AI 服务商或第三方依赖自身的漏洞,且对 Sub2API 无特定影响(请向上游报告)
|
||||
- 仅缺少最佳实践响应头或扫描器输出,但未证明实际影响
|
||||
|
||||
### 披露流程
|
||||
|
||||
1. 我们确认收到报告,并评估其有效性与严重程度。
|
||||
2. 我们在安全公告中私下开发修复,可能请你协助验证。
|
||||
3. 我们先发布包含修复的版本,再公开 GitHub 安全公告(适用时申请 CVE),并致谢报告者(如你希望匿名则不署名)。
|
||||
|
||||
在公告发布前请勿公开漏洞细节。本项目由志愿者维护,响应时间尽力而为。
|
||||
@@ -3,16 +3,16 @@ exceptions:
|
||||
- package: xlsx
|
||||
advisory: "GHSA-4r6h-8v6p-xvw6"
|
||||
severity: high
|
||||
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2023-30533)"
|
||||
mitigation: "Load only on export; restrict export permissions and data scope"
|
||||
expires_on: "2026-10-06"
|
||||
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2023-30533)"
|
||||
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
|
||||
expires_on: "2027-01-06"
|
||||
owner: "security@your-domain"
|
||||
- package: xlsx
|
||||
advisory: "GHSA-5pgg-2g8v-p4x9"
|
||||
severity: high
|
||||
reason: "Admin export only; switched to dynamic import to reduce exposure (CVE-2024-22363)"
|
||||
mitigation: "Load only on export; restrict export permissions and data scope"
|
||||
expires_on: "2026-10-06"
|
||||
reason: "Write-only: admin UsageView lazily imports xlsx solely for XLSX.write export; the advisory requires parsing a crafted file and no XLSX.read path exists (CVE-2024-22363)"
|
||||
mitigation: "Never parse untrusted spreadsheets with xlsx; no patched version on npm, migrate to SheetJS CDN build or another writer before expiry"
|
||||
expires_on: "2027-01-06"
|
||||
owner: "security@your-domain"
|
||||
- package: lodash
|
||||
advisory: "GHSA-r5fr-rjxr-66jc"
|
||||
|
||||
@@ -354,8 +354,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# Optional: Admin account
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# Leave empty to auto-generate a random email (login username) and password, shown in logs on first startup.
|
||||
# Avoid guessable values such as admin@example.com: they are brute-force targets.
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# Optional: Custom port
|
||||
SERVER_PORT=8080
|
||||
@@ -404,9 +406,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
Open `http://YOUR_SERVER_IP:8080` in your browser.
|
||||
|
||||
If admin password was auto-generated, find it in logs:
|
||||
If the admin email (login username) or password was auto-generated, find them in logs:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### Upgrade
|
||||
@@ -827,6 +829,31 @@ Administrators can override automatic media eligibility through the account crea
|
||||
|
||||
---
|
||||
|
||||
## TypeSafe / Jev Support
|
||||
|
||||
Sub2API supports TypeSafe API-key accounts through Jev's native, non-streaming System One protocol.
|
||||
|
||||
- Platform: `typesafe`; account type: API Key
|
||||
- Default upstream: `https://api.typesafe.ai`
|
||||
- Public endpoint: `POST /v1/systemone`
|
||||
- Model: `jev-latest`, also returned by `/v1/models` for TypeSafe groups
|
||||
- Questions: `noul`, `choice`, and `score`
|
||||
|
||||
Requests and successful responses retain the native System One JSON structure. This endpoint is not compatible with Chat Completions, Responses, Anthropic Messages, or streaming clients.
|
||||
|
||||
Question validation follows the TypeSafe OpenAPI wire schema (also used by SDK v0.5.7). `instructions` may be omitted or `null` for all question types. Noul `criteria` may be omitted or `null`; its `true`/`false` descriptions and Choice descriptions accept strings, objects, arrays, or `null`. Score `criteria` must be a non-empty array of string, object, or array descriptions; a single level is valid. SDK integer-keyed Score maps are normalized to arrays by the SDK before sending.
|
||||
|
||||
```bash
|
||||
curl https://your-sub2api.example.com/v1/systemone \
|
||||
-H 'Authorization: Bearer sk-your-sub2api-key' \
|
||||
-H 'Content-Type: application/json' \
|
||||
--data '{"model":"jev-latest","state":"Text to evaluate","questions":{"safety":{"type":"noul","instructions":"Evaluate whether the text is unsafe"}}}'
|
||||
```
|
||||
|
||||
The built-in `jev-latest` price is `$0.042` per million input tokens and `$0` for output tokens. Channel pricing can override both values. Credential, billing, permission, rate-limit, overload, server, and network failures (`401`, `402`, `403`, `429`, `529`, `5xx`, transport errors) use the existing account error policy (including custom error codes and temporary-unschedulable rules) and fail over to another account; request errors (`400`, `413`, and `422`) are returned without retrying another account and never change account state. TypeSafe groups (and Composite requests routed to TypeSafe) reject Messages, Chat Completions, Responses, and count_tokens requests with `404`.
|
||||
|
||||
---
|
||||
|
||||
## Antigravity Support
|
||||
|
||||
Sub2API supports [Antigravity](https://antigravity.so/) accounts. After authorization, dedicated endpoints are available for Claude and Gemini models.
|
||||
|
||||
+31
-4
@@ -380,8 +380,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# 可选:管理员账号
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# 留空则首次启动时自动生成随机邮箱(登录用户名)和密码,并输出到日志。
|
||||
# 请勿使用 admin@example.com 这类可被猜中的值,它们是暴力破解的目标。
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# 可选:自定义端口
|
||||
SERVER_PORT=8080
|
||||
@@ -442,9 +444,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
在浏览器中打开 `http://你的服务器IP:8080`
|
||||
|
||||
如果管理员密码是自动生成的,在日志中查找:
|
||||
如果管理员邮箱(登录用户名)或密码是自动生成的,在日志中查找:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### 升级
|
||||
@@ -743,6 +745,31 @@ OAuth / Setup Token 图片请求使用 Responses 主控模型调用 `image_gener
|
||||
|
||||
---
|
||||
|
||||
## TypeSafe / Jev 使用说明
|
||||
|
||||
Sub2API 支持使用 TypeSafe API Key 账户,通过 Jev 原生、非流式的 System One 协议调用模型。
|
||||
|
||||
- 平台:`typesafe`;账号类型:API Key
|
||||
- 默认上游:`https://api.typesafe.ai`
|
||||
- 对外端点:`POST /v1/systemone`
|
||||
- 模型:`jev-latest`,TypeSafe 分组的 `/v1/models` 也会返回该模型
|
||||
- 问题类型:`noul`、`choice`、`score`
|
||||
|
||||
请求和成功响应保持 System One 原生 JSON 结构。该端点不兼容 Chat Completions、Responses、Anthropic Messages 或流式客户端。
|
||||
|
||||
问题校验遵循 TypeSafe OpenAPI 的线上协议 schema(SDK v0.5.7 也使用该 schema)。所有问题的 `instructions` 都可以省略或为 `null`。Noul 的 `criteria` 可以省略或为 `null`,其中 `true`/`false` 的描述和 Choice 描述支持字符串、对象、数组或 `null`。Score 的 `criteria` 必须是至少包含一档描述的数组,每档支持字符串、对象或数组;单档也合法。SDK 的整数键 Score 映射会由 SDK 在发送前转换为数组。
|
||||
|
||||
```bash
|
||||
curl https://your-sub2api.example.com/v1/systemone \
|
||||
-H 'Authorization: Bearer sk-your-sub2api-key' \
|
||||
-H 'Content-Type: application/json' \
|
||||
--data '{"model":"jev-latest","state":"待评估文本","questions":{"safety":{"type":"noul","instructions":"评估文本是否不安全"}}}'
|
||||
```
|
||||
|
||||
`jev-latest` 内置价格为输入 `$0.042/百万 tokens`、输出 `$0`,渠道定价可以覆盖。凭据、欠费、权限、限流、过载、服务端和网络错误(`401`、`402`、`403`、`429`、`529`、`5xx`、传输错误)沿用现有账号错误策略(含自定义错误码与临时不可调度规则)并切换账号;请求错误(`400`、`413`、`422`)不会切换账号重试,也不会改变账号状态。TypeSafe 分组(以及路由到 TypeSafe 的 Composite 请求)调用 Messages、Chat Completions、Responses、count_tokens 时返回 `404`。
|
||||
|
||||
---
|
||||
|
||||
## Antigravity 使用说明
|
||||
|
||||
Sub2API 支持 [Antigravity](https://antigravity.so/) 账户,授权后可通过专用端点访问 Claude 和 Gemini 模型。
|
||||
|
||||
+6
-4
@@ -355,8 +355,10 @@ JWT_SECRET=your_jwt_secret_here
|
||||
TOTP_ENCRYPTION_KEY=your_totp_key_here
|
||||
|
||||
# オプション: 管理者アカウント
|
||||
ADMIN_EMAIL=admin@example.com
|
||||
ADMIN_PASSWORD=your_admin_password
|
||||
# 空欄の場合、初回起動時にランダムなメールアドレス(ログインユーザー名)とパスワードが自動生成され、ログに出力されます。
|
||||
# admin@example.com のような推測されやすい値は総当たり攻撃の標的になるため避けてください。
|
||||
ADMIN_EMAIL=
|
||||
ADMIN_PASSWORD=
|
||||
|
||||
# オプション: カスタムポート
|
||||
SERVER_PORT=8080
|
||||
@@ -405,9 +407,9 @@ docker compose -f docker-compose.local.yml logs -f sub2api
|
||||
|
||||
ブラウザで `http://YOUR_SERVER_IP:8080` を開いてください。
|
||||
|
||||
管理者パスワードが自動生成された場合は、ログで確認できます:
|
||||
管理者メールアドレス(ログインユーザー名)またはパスワードが自動生成された場合は、ログで確認できます:
|
||||
```bash
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "admin password"
|
||||
docker compose -f docker-compose.local.yml logs sub2api | grep "Generated admin"
|
||||
```
|
||||
|
||||
#### アップグレード
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.9
|
||||
0.2.13
|
||||
|
||||
@@ -288,7 +288,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
openCodeGoUsageService := service.ProvideOpenCodeGoUsageService(accountRepository, httpUpstream, settingService, leaderLockCache, db)
|
||||
claudeResetCreditService := service.NewClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
claudeResetCreditService := service.ProvideClaudeResetCreditService(accountRepository, claudeTokenProvider, proxyRepository, settingService, idempotencyCoordinator, leaderLockCache)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, pluginHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService, openCodeGoUsageService, claudeResetCreditService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
@@ -322,7 +323,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageHandler := handler.ProvideBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService, openAIGatewayHandler)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
|
||||
openAIQuotaAutoResetService := service.ProvideOpenAIQuotaAutoResetService(accountRepository, openAIQuotaService, rateLimitService, idempotencyCoordinator, auditLogService, settingService, leaderLockCache)
|
||||
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, channelMonitorV2Handler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, passkeyHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, modelPlazaHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService, openAIQuotaAutoResetService)
|
||||
|
||||
@@ -1128,6 +1128,7 @@ var (
|
||||
{Name: "amount", Type: field.TypeFloat64, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "pay_amount", Type: field.TypeFloat64, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "fee_rate", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
|
||||
{Name: "bonus_amount", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,2)"}},
|
||||
{Name: "recharge_code", Type: field.TypeString, Size: 64},
|
||||
{Name: "out_trade_no", Type: field.TypeString, Size: 64, Default: ""},
|
||||
{Name: "payment_type", Type: field.TypeString, Size: 30},
|
||||
@@ -1170,7 +1171,7 @@ var (
|
||||
ForeignKeys: []*schema.ForeignKey{
|
||||
{
|
||||
Symbol: "payment_orders_users_payment_orders",
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[39]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[40]},
|
||||
RefColumns: []*schema.Column{UsersColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
@@ -1179,7 +1180,7 @@ var (
|
||||
{
|
||||
Name: "paymentorder_out_trade_no",
|
||||
Unique: true,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[8]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[9]},
|
||||
Annotation: &entsql.IndexAnnotation{
|
||||
Where: "out_trade_no <> ''",
|
||||
},
|
||||
@@ -1187,37 +1188,37 @@ var (
|
||||
{
|
||||
Name: "paymentorder_user_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[39]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[40]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_status",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[21]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[22]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_expires_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[29]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[30]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[37]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[38]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_paid_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[30]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[31]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_payment_type_paid_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[9], PaymentOrdersColumns[30]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[10], PaymentOrdersColumns[31]},
|
||||
},
|
||||
{
|
||||
Name: "paymentorder_order_type",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[14]},
|
||||
Columns: []*schema.Column{PaymentOrdersColumns[15]},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+88
-1
@@ -30113,6 +30113,8 @@ type PaymentOrderMutation struct {
|
||||
addpay_amount *float64
|
||||
fee_rate *float64
|
||||
addfee_rate *float64
|
||||
bonus_amount *float64
|
||||
addbonus_amount *float64
|
||||
recharge_code *string
|
||||
out_trade_no *string
|
||||
payment_type *string
|
||||
@@ -30580,6 +30582,62 @@ func (m *PaymentOrderMutation) ResetFeeRate() {
|
||||
m.addfee_rate = nil
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) SetBonusAmount(f float64) {
|
||||
m.bonus_amount = &f
|
||||
m.addbonus_amount = nil
|
||||
}
|
||||
|
||||
// BonusAmount returns the value of the "bonus_amount" field in the mutation.
|
||||
func (m *PaymentOrderMutation) BonusAmount() (r float64, exists bool) {
|
||||
v := m.bonus_amount
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldBonusAmount returns the old "bonus_amount" field's value of the PaymentOrder entity.
|
||||
// If the PaymentOrder object wasn't provided to the builder, the object is fetched from the database.
|
||||
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
|
||||
func (m *PaymentOrderMutation) OldBonusAmount(ctx context.Context) (v float64, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldBonusAmount is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldBonusAmount requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldBonusAmount: %w", err)
|
||||
}
|
||||
return oldValue.BonusAmount, nil
|
||||
}
|
||||
|
||||
// AddBonusAmount adds f to the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) AddBonusAmount(f float64) {
|
||||
if m.addbonus_amount != nil {
|
||||
*m.addbonus_amount += f
|
||||
} else {
|
||||
m.addbonus_amount = &f
|
||||
}
|
||||
}
|
||||
|
||||
// AddedBonusAmount returns the value that was added to the "bonus_amount" field in this mutation.
|
||||
func (m *PaymentOrderMutation) AddedBonusAmount() (r float64, exists bool) {
|
||||
v := m.addbonus_amount
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// ResetBonusAmount resets all changes to the "bonus_amount" field.
|
||||
func (m *PaymentOrderMutation) ResetBonusAmount() {
|
||||
m.bonus_amount = nil
|
||||
m.addbonus_amount = nil
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (m *PaymentOrderMutation) SetRechargeCode(s string) {
|
||||
m.recharge_code = &s
|
||||
@@ -32123,7 +32181,7 @@ func (m *PaymentOrderMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *PaymentOrderMutation) Fields() []string {
|
||||
fields := make([]string, 0, 39)
|
||||
fields := make([]string, 0, 40)
|
||||
if m.user != nil {
|
||||
fields = append(fields, paymentorder.FieldUserID)
|
||||
}
|
||||
@@ -32145,6 +32203,9 @@ func (m *PaymentOrderMutation) Fields() []string {
|
||||
if m.fee_rate != nil {
|
||||
fields = append(fields, paymentorder.FieldFeeRate)
|
||||
}
|
||||
if m.bonus_amount != nil {
|
||||
fields = append(fields, paymentorder.FieldBonusAmount)
|
||||
}
|
||||
if m.recharge_code != nil {
|
||||
fields = append(fields, paymentorder.FieldRechargeCode)
|
||||
}
|
||||
@@ -32263,6 +32324,8 @@ func (m *PaymentOrderMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.PayAmount()
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.FeeRate()
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.BonusAmount()
|
||||
case paymentorder.FieldRechargeCode:
|
||||
return m.RechargeCode()
|
||||
case paymentorder.FieldOutTradeNo:
|
||||
@@ -32350,6 +32413,8 @@ func (m *PaymentOrderMutation) OldField(ctx context.Context, name string) (ent.V
|
||||
return m.OldPayAmount(ctx)
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.OldFeeRate(ctx)
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.OldBonusAmount(ctx)
|
||||
case paymentorder.FieldRechargeCode:
|
||||
return m.OldRechargeCode(ctx)
|
||||
case paymentorder.FieldOutTradeNo:
|
||||
@@ -32472,6 +32537,13 @@ func (m *PaymentOrderMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetFeeRate(v)
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
v, ok := value.(float64)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetBonusAmount(v)
|
||||
return nil
|
||||
case paymentorder.FieldRechargeCode:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
@@ -32713,6 +32785,9 @@ func (m *PaymentOrderMutation) AddedFields() []string {
|
||||
if m.addfee_rate != nil {
|
||||
fields = append(fields, paymentorder.FieldFeeRate)
|
||||
}
|
||||
if m.addbonus_amount != nil {
|
||||
fields = append(fields, paymentorder.FieldBonusAmount)
|
||||
}
|
||||
if m.addplan_id != nil {
|
||||
fields = append(fields, paymentorder.FieldPlanID)
|
||||
}
|
||||
@@ -32739,6 +32814,8 @@ func (m *PaymentOrderMutation) AddedField(name string) (ent.Value, bool) {
|
||||
return m.AddedPayAmount()
|
||||
case paymentorder.FieldFeeRate:
|
||||
return m.AddedFeeRate()
|
||||
case paymentorder.FieldBonusAmount:
|
||||
return m.AddedBonusAmount()
|
||||
case paymentorder.FieldPlanID:
|
||||
return m.AddedPlanID()
|
||||
case paymentorder.FieldSubscriptionGroupID:
|
||||
@@ -32777,6 +32854,13 @@ func (m *PaymentOrderMutation) AddField(name string, value ent.Value) error {
|
||||
}
|
||||
m.AddFeeRate(v)
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
v, ok := value.(float64)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.AddBonusAmount(v)
|
||||
return nil
|
||||
case paymentorder.FieldPlanID:
|
||||
v, ok := value.(int64)
|
||||
if !ok {
|
||||
@@ -32976,6 +33060,9 @@ func (m *PaymentOrderMutation) ResetField(name string) error {
|
||||
case paymentorder.FieldFeeRate:
|
||||
m.ResetFeeRate()
|
||||
return nil
|
||||
case paymentorder.FieldBonusAmount:
|
||||
m.ResetBonusAmount()
|
||||
return nil
|
||||
case paymentorder.FieldRechargeCode:
|
||||
m.ResetRechargeCode()
|
||||
return nil
|
||||
|
||||
@@ -33,6 +33,8 @@ type PaymentOrder struct {
|
||||
PayAmount float64 `json:"pay_amount,omitempty"`
|
||||
// FeeRate holds the value of the "fee_rate" field.
|
||||
FeeRate float64 `json:"fee_rate,omitempty"`
|
||||
// BonusAmount holds the value of the "bonus_amount" field.
|
||||
BonusAmount float64 `json:"bonus_amount,omitempty"`
|
||||
// RechargeCode holds the value of the "recharge_code" field.
|
||||
RechargeCode string `json:"recharge_code,omitempty"`
|
||||
// OutTradeNo holds the value of the "out_trade_no" field.
|
||||
@@ -132,7 +134,7 @@ func (*PaymentOrder) scanValues(columns []string) ([]any, error) {
|
||||
values[i] = new([]byte)
|
||||
case paymentorder.FieldForceRefund:
|
||||
values[i] = new(sql.NullBool)
|
||||
case paymentorder.FieldAmount, paymentorder.FieldPayAmount, paymentorder.FieldFeeRate, paymentorder.FieldRefundAmount:
|
||||
case paymentorder.FieldAmount, paymentorder.FieldPayAmount, paymentorder.FieldFeeRate, paymentorder.FieldBonusAmount, paymentorder.FieldRefundAmount:
|
||||
values[i] = new(sql.NullFloat64)
|
||||
case paymentorder.FieldID, paymentorder.FieldUserID, paymentorder.FieldPlanID, paymentorder.FieldSubscriptionGroupID, paymentorder.FieldSubscriptionDays:
|
||||
values[i] = new(sql.NullInt64)
|
||||
@@ -204,6 +206,12 @@ func (_m *PaymentOrder) assignValues(columns []string, values []any) error {
|
||||
} else if value.Valid {
|
||||
_m.FeeRate = value.Float64
|
||||
}
|
||||
case paymentorder.FieldBonusAmount:
|
||||
if value, ok := values[i].(*sql.NullFloat64); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field bonus_amount", values[i])
|
||||
} else if value.Valid {
|
||||
_m.BonusAmount = value.Float64
|
||||
}
|
||||
case paymentorder.FieldRechargeCode:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field recharge_code", values[i])
|
||||
@@ -480,6 +488,9 @@ func (_m *PaymentOrder) String() string {
|
||||
builder.WriteString("fee_rate=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.FeeRate))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("bonus_amount=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.BonusAmount))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("recharge_code=")
|
||||
builder.WriteString(_m.RechargeCode)
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -28,6 +28,8 @@ const (
|
||||
FieldPayAmount = "pay_amount"
|
||||
// FieldFeeRate holds the string denoting the fee_rate field in the database.
|
||||
FieldFeeRate = "fee_rate"
|
||||
// FieldBonusAmount holds the string denoting the bonus_amount field in the database.
|
||||
FieldBonusAmount = "bonus_amount"
|
||||
// FieldRechargeCode holds the string denoting the recharge_code field in the database.
|
||||
FieldRechargeCode = "recharge_code"
|
||||
// FieldOutTradeNo holds the string denoting the out_trade_no field in the database.
|
||||
@@ -115,6 +117,7 @@ var Columns = []string{
|
||||
FieldAmount,
|
||||
FieldPayAmount,
|
||||
FieldFeeRate,
|
||||
FieldBonusAmount,
|
||||
FieldRechargeCode,
|
||||
FieldOutTradeNo,
|
||||
FieldPaymentType,
|
||||
@@ -166,6 +169,8 @@ var (
|
||||
UserNameValidator func(string) error
|
||||
// DefaultFeeRate holds the default value on creation for the "fee_rate" field.
|
||||
DefaultFeeRate float64
|
||||
// DefaultBonusAmount holds the default value on creation for the "bonus_amount" field.
|
||||
DefaultBonusAmount float64
|
||||
// RechargeCodeValidator is a validator for the "recharge_code" field. It is called by the builders before save.
|
||||
RechargeCodeValidator func(string) error
|
||||
// DefaultOutTradeNo holds the default value on creation for the "out_trade_no" field.
|
||||
@@ -249,6 +254,11 @@ func ByFeeRate(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldFeeRate, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByBonusAmount orders the results by the bonus_amount field.
|
||||
func ByBonusAmount(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldBonusAmount, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByRechargeCode orders the results by the recharge_code field.
|
||||
func ByRechargeCode(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldRechargeCode, opts...).ToFunc()
|
||||
|
||||
@@ -90,6 +90,11 @@ func FeeRate(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldFeeRate, v))
|
||||
}
|
||||
|
||||
// BonusAmount applies equality check predicate on the "bonus_amount" field. It's identical to BonusAmountEQ.
|
||||
func BonusAmount(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// RechargeCode applies equality check predicate on the "recharge_code" field. It's identical to RechargeCodeEQ.
|
||||
func RechargeCode(v string) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldRechargeCode, v))
|
||||
@@ -590,6 +595,46 @@ func FeeRateLTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLTE(FieldFeeRate, v))
|
||||
}
|
||||
|
||||
// BonusAmountEQ applies the EQ predicate on the "bonus_amount" field.
|
||||
func BonusAmountEQ(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountNEQ applies the NEQ predicate on the "bonus_amount" field.
|
||||
func BonusAmountNEQ(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldNEQ(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountIn applies the In predicate on the "bonus_amount" field.
|
||||
func BonusAmountIn(vs ...float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldIn(FieldBonusAmount, vs...))
|
||||
}
|
||||
|
||||
// BonusAmountNotIn applies the NotIn predicate on the "bonus_amount" field.
|
||||
func BonusAmountNotIn(vs ...float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldNotIn(FieldBonusAmount, vs...))
|
||||
}
|
||||
|
||||
// BonusAmountGT applies the GT predicate on the "bonus_amount" field.
|
||||
func BonusAmountGT(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldGT(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountGTE applies the GTE predicate on the "bonus_amount" field.
|
||||
func BonusAmountGTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldGTE(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountLT applies the LT predicate on the "bonus_amount" field.
|
||||
func BonusAmountLT(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLT(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// BonusAmountLTE applies the LTE predicate on the "bonus_amount" field.
|
||||
func BonusAmountLTE(v float64) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldLTE(FieldBonusAmount, v))
|
||||
}
|
||||
|
||||
// RechargeCodeEQ applies the EQ predicate on the "recharge_code" field.
|
||||
func RechargeCodeEQ(v string) predicate.PaymentOrder {
|
||||
return predicate.PaymentOrder(sql.FieldEQ(FieldRechargeCode, v))
|
||||
|
||||
@@ -81,6 +81,20 @@ func (_c *PaymentOrderCreate) SetNillableFeeRate(v *float64) *PaymentOrderCreate
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_c *PaymentOrderCreate) SetBonusAmount(v float64) *PaymentOrderCreate {
|
||||
_c.mutation.SetBonusAmount(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_c *PaymentOrderCreate) SetNillableBonusAmount(v *float64) *PaymentOrderCreate {
|
||||
if v != nil {
|
||||
_c.SetBonusAmount(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_c *PaymentOrderCreate) SetRechargeCode(v string) *PaymentOrderCreate {
|
||||
_c.mutation.SetRechargeCode(v)
|
||||
@@ -517,6 +531,10 @@ func (_c *PaymentOrderCreate) defaults() {
|
||||
v := paymentorder.DefaultFeeRate
|
||||
_c.mutation.SetFeeRate(v)
|
||||
}
|
||||
if _, ok := _c.mutation.BonusAmount(); !ok {
|
||||
v := paymentorder.DefaultBonusAmount
|
||||
_c.mutation.SetBonusAmount(v)
|
||||
}
|
||||
if _, ok := _c.mutation.OutTradeNo(); !ok {
|
||||
v := paymentorder.DefaultOutTradeNo
|
||||
_c.mutation.SetOutTradeNo(v)
|
||||
@@ -577,6 +595,9 @@ func (_c *PaymentOrderCreate) check() error {
|
||||
if _, ok := _c.mutation.FeeRate(); !ok {
|
||||
return &ValidationError{Name: "fee_rate", err: errors.New(`ent: missing required field "PaymentOrder.fee_rate"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.BonusAmount(); !ok {
|
||||
return &ValidationError{Name: "bonus_amount", err: errors.New(`ent: missing required field "PaymentOrder.bonus_amount"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.RechargeCode(); !ok {
|
||||
return &ValidationError{Name: "recharge_code", err: errors.New(`ent: missing required field "PaymentOrder.recharge_code"`)}
|
||||
}
|
||||
@@ -725,6 +746,10 @@ func (_c *PaymentOrderCreate) createSpec() (*PaymentOrder, *sqlgraph.CreateSpec)
|
||||
_spec.SetField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
_node.FeeRate = value
|
||||
}
|
||||
if value, ok := _c.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
_node.BonusAmount = value
|
||||
}
|
||||
if value, ok := _c.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
_node.RechargeCode = value
|
||||
@@ -1030,6 +1055,24 @@ func (u *PaymentOrderUpsert) AddFeeRate(v float64) *PaymentOrderUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsert) SetBonusAmount(v float64) *PaymentOrderUpsert {
|
||||
u.Set(paymentorder.FieldBonusAmount, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsert) UpdateBonusAmount() *PaymentOrderUpsert {
|
||||
u.SetExcluded(paymentorder.FieldBonusAmount)
|
||||
return u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsert) AddBonusAmount(v float64) *PaymentOrderUpsert {
|
||||
u.Add(paymentorder.FieldBonusAmount, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsert) SetRechargeCode(v string) *PaymentOrderUpsert {
|
||||
u.Set(paymentorder.FieldRechargeCode, v)
|
||||
@@ -1711,6 +1754,27 @@ func (u *PaymentOrderUpsertOne) UpdateFeeRate() *PaymentOrderUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertOne) SetBonusAmount(v float64) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.SetBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertOne) AddBonusAmount(v float64) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.AddBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsertOne) UpdateBonusAmount() *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.UpdateBonusAmount()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsertOne) SetRechargeCode(v string) *PaymentOrderUpsertOne {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
@@ -2643,6 +2707,27 @@ func (u *PaymentOrderUpsertBulk) UpdateFeeRate() *PaymentOrderUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertBulk) SetBonusAmount(v float64) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.SetBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// AddBonusAmount adds v to the "bonus_amount" field.
|
||||
func (u *PaymentOrderUpsertBulk) AddBonusAmount(v float64) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.AddBonusAmount(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateBonusAmount sets the "bonus_amount" field to the value that was provided on create.
|
||||
func (u *PaymentOrderUpsertBulk) UpdateBonusAmount() *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
s.UpdateBonusAmount()
|
||||
})
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (u *PaymentOrderUpsertBulk) SetRechargeCode(v string) *PaymentOrderUpsertBulk {
|
||||
return u.Update(func(s *PaymentOrderUpsert) {
|
||||
|
||||
@@ -154,6 +154,27 @@ func (_u *PaymentOrderUpdate) AddFeeRate(v float64) *PaymentOrderUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdate) SetBonusAmount(v float64) *PaymentOrderUpdate {
|
||||
_u.mutation.ResetBonusAmount()
|
||||
_u.mutation.SetBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_u *PaymentOrderUpdate) SetNillableBonusAmount(v *float64) *PaymentOrderUpdate {
|
||||
if v != nil {
|
||||
_u.SetBonusAmount(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds value to the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdate) AddBonusAmount(v float64) *PaymentOrderUpdate {
|
||||
_u.mutation.AddBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_u *PaymentOrderUpdate) SetRechargeCode(v string) *PaymentOrderUpdate {
|
||||
_u.mutation.SetRechargeCode(v)
|
||||
@@ -881,6 +902,12 @@ func (_u *PaymentOrderUpdate) sqlSave(ctx context.Context) (_node int, err error
|
||||
if value, ok := _u.mutation.AddedFeeRate(); ok {
|
||||
_spec.AddField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AddedBonusAmount(); ok {
|
||||
_spec.AddField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
}
|
||||
@@ -1217,6 +1244,27 @@ func (_u *PaymentOrderUpdateOne) AddFeeRate(v float64) *PaymentOrderUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetBonusAmount sets the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdateOne) SetBonusAmount(v float64) *PaymentOrderUpdateOne {
|
||||
_u.mutation.ResetBonusAmount()
|
||||
_u.mutation.SetBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableBonusAmount sets the "bonus_amount" field if the given value is not nil.
|
||||
func (_u *PaymentOrderUpdateOne) SetNillableBonusAmount(v *float64) *PaymentOrderUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetBonusAmount(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// AddBonusAmount adds value to the "bonus_amount" field.
|
||||
func (_u *PaymentOrderUpdateOne) AddBonusAmount(v float64) *PaymentOrderUpdateOne {
|
||||
_u.mutation.AddBonusAmount(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetRechargeCode sets the "recharge_code" field.
|
||||
func (_u *PaymentOrderUpdateOne) SetRechargeCode(v string) *PaymentOrderUpdateOne {
|
||||
_u.mutation.SetRechargeCode(v)
|
||||
@@ -1974,6 +2022,12 @@ func (_u *PaymentOrderUpdateOne) sqlSave(ctx context.Context) (_node *PaymentOrd
|
||||
if value, ok := _u.mutation.AddedFeeRate(); ok {
|
||||
_spec.AddField(paymentorder.FieldFeeRate, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.BonusAmount(); ok {
|
||||
_spec.SetField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.AddedBonusAmount(); ok {
|
||||
_spec.AddField(paymentorder.FieldBonusAmount, field.TypeFloat64, value)
|
||||
}
|
||||
if value, ok := _u.mutation.RechargeCode(); ok {
|
||||
_spec.SetField(paymentorder.FieldRechargeCode, field.TypeString, value)
|
||||
}
|
||||
|
||||
@@ -1323,70 +1323,74 @@ func init() {
|
||||
paymentorderDescFeeRate := paymentorderFields[6].Descriptor()
|
||||
// paymentorder.DefaultFeeRate holds the default value on creation for the fee_rate field.
|
||||
paymentorder.DefaultFeeRate = paymentorderDescFeeRate.Default.(float64)
|
||||
// paymentorderDescBonusAmount is the schema descriptor for bonus_amount field.
|
||||
paymentorderDescBonusAmount := paymentorderFields[7].Descriptor()
|
||||
// paymentorder.DefaultBonusAmount holds the default value on creation for the bonus_amount field.
|
||||
paymentorder.DefaultBonusAmount = paymentorderDescBonusAmount.Default.(float64)
|
||||
// paymentorderDescRechargeCode is the schema descriptor for recharge_code field.
|
||||
paymentorderDescRechargeCode := paymentorderFields[7].Descriptor()
|
||||
paymentorderDescRechargeCode := paymentorderFields[8].Descriptor()
|
||||
// paymentorder.RechargeCodeValidator is a validator for the "recharge_code" field. It is called by the builders before save.
|
||||
paymentorder.RechargeCodeValidator = paymentorderDescRechargeCode.Validators[0].(func(string) error)
|
||||
// paymentorderDescOutTradeNo is the schema descriptor for out_trade_no field.
|
||||
paymentorderDescOutTradeNo := paymentorderFields[8].Descriptor()
|
||||
paymentorderDescOutTradeNo := paymentorderFields[9].Descriptor()
|
||||
// paymentorder.DefaultOutTradeNo holds the default value on creation for the out_trade_no field.
|
||||
paymentorder.DefaultOutTradeNo = paymentorderDescOutTradeNo.Default.(string)
|
||||
// paymentorder.OutTradeNoValidator is a validator for the "out_trade_no" field. It is called by the builders before save.
|
||||
paymentorder.OutTradeNoValidator = paymentorderDescOutTradeNo.Validators[0].(func(string) error)
|
||||
// paymentorderDescPaymentType is the schema descriptor for payment_type field.
|
||||
paymentorderDescPaymentType := paymentorderFields[9].Descriptor()
|
||||
paymentorderDescPaymentType := paymentorderFields[10].Descriptor()
|
||||
// paymentorder.PaymentTypeValidator is a validator for the "payment_type" field. It is called by the builders before save.
|
||||
paymentorder.PaymentTypeValidator = paymentorderDescPaymentType.Validators[0].(func(string) error)
|
||||
// paymentorderDescPaymentTradeNo is the schema descriptor for payment_trade_no field.
|
||||
paymentorderDescPaymentTradeNo := paymentorderFields[10].Descriptor()
|
||||
paymentorderDescPaymentTradeNo := paymentorderFields[11].Descriptor()
|
||||
// paymentorder.PaymentTradeNoValidator is a validator for the "payment_trade_no" field. It is called by the builders before save.
|
||||
paymentorder.PaymentTradeNoValidator = paymentorderDescPaymentTradeNo.Validators[0].(func(string) error)
|
||||
// paymentorderDescOrderType is the schema descriptor for order_type field.
|
||||
paymentorderDescOrderType := paymentorderFields[14].Descriptor()
|
||||
paymentorderDescOrderType := paymentorderFields[15].Descriptor()
|
||||
// paymentorder.DefaultOrderType holds the default value on creation for the order_type field.
|
||||
paymentorder.DefaultOrderType = paymentorderDescOrderType.Default.(string)
|
||||
// paymentorder.OrderTypeValidator is a validator for the "order_type" field. It is called by the builders before save.
|
||||
paymentorder.OrderTypeValidator = paymentorderDescOrderType.Validators[0].(func(string) error)
|
||||
// paymentorderDescProviderInstanceID is the schema descriptor for provider_instance_id field.
|
||||
paymentorderDescProviderInstanceID := paymentorderFields[18].Descriptor()
|
||||
paymentorderDescProviderInstanceID := paymentorderFields[19].Descriptor()
|
||||
// paymentorder.ProviderInstanceIDValidator is a validator for the "provider_instance_id" field. It is called by the builders before save.
|
||||
paymentorder.ProviderInstanceIDValidator = paymentorderDescProviderInstanceID.Validators[0].(func(string) error)
|
||||
// paymentorderDescProviderKey is the schema descriptor for provider_key field.
|
||||
paymentorderDescProviderKey := paymentorderFields[19].Descriptor()
|
||||
paymentorderDescProviderKey := paymentorderFields[20].Descriptor()
|
||||
// paymentorder.ProviderKeyValidator is a validator for the "provider_key" field. It is called by the builders before save.
|
||||
paymentorder.ProviderKeyValidator = paymentorderDescProviderKey.Validators[0].(func(string) error)
|
||||
// paymentorderDescStatus is the schema descriptor for status field.
|
||||
paymentorderDescStatus := paymentorderFields[21].Descriptor()
|
||||
paymentorderDescStatus := paymentorderFields[22].Descriptor()
|
||||
// paymentorder.DefaultStatus holds the default value on creation for the status field.
|
||||
paymentorder.DefaultStatus = paymentorderDescStatus.Default.(string)
|
||||
// paymentorder.StatusValidator is a validator for the "status" field. It is called by the builders before save.
|
||||
paymentorder.StatusValidator = paymentorderDescStatus.Validators[0].(func(string) error)
|
||||
// paymentorderDescRefundAmount is the schema descriptor for refund_amount field.
|
||||
paymentorderDescRefundAmount := paymentorderFields[22].Descriptor()
|
||||
paymentorderDescRefundAmount := paymentorderFields[23].Descriptor()
|
||||
// paymentorder.DefaultRefundAmount holds the default value on creation for the refund_amount field.
|
||||
paymentorder.DefaultRefundAmount = paymentorderDescRefundAmount.Default.(float64)
|
||||
// paymentorderDescForceRefund is the schema descriptor for force_refund field.
|
||||
paymentorderDescForceRefund := paymentorderFields[25].Descriptor()
|
||||
paymentorderDescForceRefund := paymentorderFields[26].Descriptor()
|
||||
// paymentorder.DefaultForceRefund holds the default value on creation for the force_refund field.
|
||||
paymentorder.DefaultForceRefund = paymentorderDescForceRefund.Default.(bool)
|
||||
// paymentorderDescRefundRequestedBy is the schema descriptor for refund_requested_by field.
|
||||
paymentorderDescRefundRequestedBy := paymentorderFields[28].Descriptor()
|
||||
paymentorderDescRefundRequestedBy := paymentorderFields[29].Descriptor()
|
||||
// paymentorder.RefundRequestedByValidator is a validator for the "refund_requested_by" field. It is called by the builders before save.
|
||||
paymentorder.RefundRequestedByValidator = paymentorderDescRefundRequestedBy.Validators[0].(func(string) error)
|
||||
// paymentorderDescClientIP is the schema descriptor for client_ip field.
|
||||
paymentorderDescClientIP := paymentorderFields[34].Descriptor()
|
||||
paymentorderDescClientIP := paymentorderFields[35].Descriptor()
|
||||
// paymentorder.ClientIPValidator is a validator for the "client_ip" field. It is called by the builders before save.
|
||||
paymentorder.ClientIPValidator = paymentorderDescClientIP.Validators[0].(func(string) error)
|
||||
// paymentorderDescSrcHost is the schema descriptor for src_host field.
|
||||
paymentorderDescSrcHost := paymentorderFields[35].Descriptor()
|
||||
paymentorderDescSrcHost := paymentorderFields[36].Descriptor()
|
||||
// paymentorder.SrcHostValidator is a validator for the "src_host" field. It is called by the builders before save.
|
||||
paymentorder.SrcHostValidator = paymentorderDescSrcHost.Validators[0].(func(string) error)
|
||||
// paymentorderDescCreatedAt is the schema descriptor for created_at field.
|
||||
paymentorderDescCreatedAt := paymentorderFields[37].Descriptor()
|
||||
paymentorderDescCreatedAt := paymentorderFields[38].Descriptor()
|
||||
// paymentorder.DefaultCreatedAt holds the default value on creation for the created_at field.
|
||||
paymentorder.DefaultCreatedAt = paymentorderDescCreatedAt.Default.(func() time.Time)
|
||||
// paymentorderDescUpdatedAt is the schema descriptor for updated_at field.
|
||||
paymentorderDescUpdatedAt := paymentorderFields[38].Descriptor()
|
||||
paymentorderDescUpdatedAt := paymentorderFields[39].Descriptor()
|
||||
// paymentorder.DefaultUpdatedAt holds the default value on creation for the updated_at field.
|
||||
paymentorder.DefaultUpdatedAt = paymentorderDescUpdatedAt.Default.(func() time.Time)
|
||||
// paymentorder.UpdateDefaultUpdatedAt holds the default value on update for the updated_at field.
|
||||
|
||||
@@ -50,6 +50,10 @@ func (PaymentOrder) Fields() []ent.Field {
|
||||
field.Float("fee_rate").
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}).
|
||||
Default(0),
|
||||
// 充值赠送额度(USD)。已计入 amount;单独记录用于订单展示与推广返利基数剔除。
|
||||
field.Float("bonus_amount").
|
||||
SchemaType(map[string]string{dialect.Postgres: "decimal(20,2)"}).
|
||||
Default(0),
|
||||
field.String("recharge_code").
|
||||
MaxLen(64),
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ func (UserPlatformQuota) Fields() []ent.Field {
|
||||
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
|
||||
switch s {
|
||||
case "anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go":
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe":
|
||||
return nil
|
||||
default:
|
||||
return fmt.Errorf("platform %q is not allowed", s)
|
||||
|
||||
@@ -89,6 +89,7 @@ type Config struct {
|
||||
Pricing PricingConfig `mapstructure:"pricing"`
|
||||
Gateway GatewayConfig `mapstructure:"gateway"`
|
||||
APIKeyAuth APIKeyAuthCacheConfig `mapstructure:"api_key_auth_cache"`
|
||||
APIKeyCreate APIKeyCreateConfig `mapstructure:"api_key_create"`
|
||||
SubscriptionCache SubscriptionCacheConfig `mapstructure:"subscription_cache"`
|
||||
SubscriptionMaintenance SubscriptionMaintenanceConfig `mapstructure:"subscription_maintenance"`
|
||||
Dashboard DashboardCacheConfig `mapstructure:"dashboard_cache"`
|
||||
@@ -921,6 +922,29 @@ type BillingConfig struct {
|
||||
// UserPlatformQuotaSentinelTTLSeconds sentinel(无 limit 占位)entry 的 TTL,
|
||||
// 显著短于 quota cache 默认 86400s 以控 Redis 内存;默认 3600=1h。
|
||||
UserPlatformQuotaSentinelTTLSeconds int `mapstructure:"user_platform_quota_sentinel_ttl_seconds"`
|
||||
// InflightReservation 余额模式在途请求预留(Redis),防止并发请求在预检时看到同一份余额而集体透支。
|
||||
InflightReservation InflightReservationConfig `mapstructure:"inflight_reservation"`
|
||||
}
|
||||
|
||||
// InflightReservationConfig 余额模式在途预留配置。
|
||||
// 准入时按 输入估算 + 输出单价 × max_tokens 估算单请求费用,在 Redis 中原子地
|
||||
// 校验 缓存余额 - 在途预留合计 >= 估算 后登记预留,请求结束(任意路径)释放。
|
||||
// 估算失败或 Redis 不可用时 fail-open,退回旧的仅余额 > 阈值检查。
|
||||
type InflightReservationConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
// TTLSeconds 单条预留的最长存活时间;进程崩溃等泄漏的预留到期自动失效。
|
||||
TTLSeconds int `mapstructure:"ttl_seconds"`
|
||||
// DefaultMaxTokens 请求未携带 max_tokens 时用于估算的输出 token 数。
|
||||
DefaultMaxTokens int `mapstructure:"default_max_tokens"`
|
||||
// MaxOutputTokens 估算输出 token 的上限(max_tokens 超出时截断)。
|
||||
MaxOutputTokens int `mapstructure:"max_output_tokens"`
|
||||
// MaxInputTokens 输入 token 估算(请求体字节数 / 4)的上限。
|
||||
MaxInputTokens int `mapstructure:"max_input_tokens"`
|
||||
// MaxReservationUSD 单请求预留金额上限;0 表示不设上限。
|
||||
MaxReservationUSD float64 `mapstructure:"max_reservation_usd"`
|
||||
// FailClosedOnUnpriced 无法为请求估算费用(模型/分组/渠道均无定价)时是否拒绝请求。
|
||||
// 默认 false:放行且不预留(fail-open,节流告警日志)。
|
||||
FailClosedOnUnpriced bool `mapstructure:"fail_closed_on_unpriced"`
|
||||
}
|
||||
|
||||
type CircuitBreakerConfig struct {
|
||||
@@ -1696,6 +1720,14 @@ type APIKeyAuthCacheConfig struct {
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
// APIKeyCreateConfig 用户创建 API Key 的防滥用限制(0 表示不限制)
|
||||
type APIKeyCreateConfig struct {
|
||||
// MaxActivePerUser 单个用户同时存在(未删除)的 API Key 上限
|
||||
MaxActivePerUser int `mapstructure:"max_active_per_user"`
|
||||
// MaxPerUserPerHour 单个用户每小时可创建的 API Key 次数(删除不返还次数)
|
||||
MaxPerUserPerHour int `mapstructure:"max_per_user_per_hour"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
@@ -2093,6 +2125,13 @@ func setDefaults() {
|
||||
viper.SetDefault("billing.minimum_balance_reserve", 0.000001)
|
||||
viper.SetDefault("billing.user_platform_quota_cache_ttl_seconds", 86400)
|
||||
viper.SetDefault("billing.user_platform_quota_sentinel_ttl_seconds", 3600)
|
||||
viper.SetDefault("billing.inflight_reservation.enabled", true)
|
||||
viper.SetDefault("billing.inflight_reservation.ttl_seconds", 900)
|
||||
viper.SetDefault("billing.inflight_reservation.default_max_tokens", 8192)
|
||||
viper.SetDefault("billing.inflight_reservation.max_output_tokens", 128000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_input_tokens", 200000)
|
||||
viper.SetDefault("billing.inflight_reservation.max_reservation_usd", 0)
|
||||
viper.SetDefault("billing.inflight_reservation.fail_closed_on_unpriced", false)
|
||||
|
||||
// Turnstile
|
||||
viper.SetDefault("turnstile.required", false)
|
||||
@@ -2332,6 +2371,8 @@ func setDefaults() {
|
||||
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)
|
||||
viper.SetDefault("api_key_create.max_active_per_user", 200)
|
||||
viper.SetDefault("api_key_create.max_per_user_per_hour", 60)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2711,6 +2752,12 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyCreate.MaxActivePerUser < 0 {
|
||||
return fmt.Errorf("api_key_create.max_active_per_user must be non-negative")
|
||||
}
|
||||
if c.APIKeyCreate.MaxPerUserPerHour < 0 {
|
||||
return fmt.Errorf("api_key_create.max_per_user_per_hour must be non-negative")
|
||||
}
|
||||
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")
|
||||
@@ -3062,6 +3109,11 @@ func (c *Config) Validate() error {
|
||||
if c.Billing.MinimumBalanceReserve < 0 {
|
||||
return fmt.Errorf("billing.minimum_balance_reserve must be non-negative")
|
||||
}
|
||||
if c.Billing.InflightReservation.TTLSeconds < 0 || c.Billing.InflightReservation.DefaultMaxTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxOutputTokens < 0 || c.Billing.InflightReservation.MaxInputTokens < 0 ||
|
||||
c.Billing.InflightReservation.MaxReservationUSD < 0 {
|
||||
return fmt.Errorf("billing.inflight_reservation values must be non-negative")
|
||||
}
|
||||
if c.Database.MaxOpenConns <= 0 {
|
||||
return fmt.Errorf("database.max_open_conns must be positive")
|
||||
}
|
||||
|
||||
@@ -28,6 +28,7 @@ const (
|
||||
PlatformZhipu = "zhipu" // 智谱 GLM (bigmodel)
|
||||
PlatformDeepseek = "deepseek" // DeepSeek
|
||||
PlatformMiniMax = "minimax" // MiniMax (M 系列)
|
||||
PlatformTypeSafe = "typesafe" // TypeSafe AI System One (Jev)
|
||||
// PlatformOpenCodeGo 是 OpenCode 平台(账号类型 Zen 按量 / Go 订阅)。
|
||||
// 值保持 opencode_go 以兼容已落库的分组、配额与 Composite 路由 CHECK。
|
||||
PlatformOpenCodeGo = "opencode_go"
|
||||
|
||||
@@ -27,6 +27,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -2950,6 +2951,12 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// TypeSafe accounts serve only the native System One model.
|
||||
if account.IsTypeSafe() {
|
||||
response.Success(c, []claude.Model{{ID: typesafe.JevLatestModel, Type: "model", DisplayName: typesafe.JevLatestModel}})
|
||||
return
|
||||
}
|
||||
|
||||
// Handle Claude/Anthropic accounts
|
||||
// For OAuth and Setup-Token accounts: return default models
|
||||
if account.IsOAuth() {
|
||||
|
||||
@@ -526,3 +526,36 @@ func TestAccountHandlerSyncUpstreamModels_MetadataEnrichmentFailureReturnsWarnin
|
||||
require.Len(t, resp.Data.Warnings, 1)
|
||||
require.Equal(t, "upstream_model_metadata_incomplete", resp.Data.Warnings[0].Code)
|
||||
}
|
||||
|
||||
func TestAccountHandlerGetAvailableModels_TypeSafeOnlyReturnsJev(t *testing.T) {
|
||||
for _, credentials := range []map[string]any{
|
||||
{"api_key": "ts-secret"},
|
||||
{"api_key": "ts-secret", "model_mapping": map[string]any{"jev-latest": "jev-latest"}},
|
||||
} {
|
||||
svc := &availableModelsAdminService{
|
||||
stubAdminService: newStubAdminService(),
|
||||
account: service.Account{
|
||||
ID: 46,
|
||||
Name: "typesafe",
|
||||
Platform: service.PlatformTypeSafe,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Credentials: credentials,
|
||||
},
|
||||
}
|
||||
router := setupAvailableModelsRouter(svc)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/46/models", nil))
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Len(t, resp.Data, 1)
|
||||
require.Equal(t, "jev-latest", resp.Data[0].ID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -644,6 +644,7 @@ var platformToLiteLLMProvider = map[string]string{
|
||||
service.PlatformDeepseek: "deepseek",
|
||||
service.PlatformMiniMax: "minimax",
|
||||
service.PlatformOpenCodeGo: "opencode-go",
|
||||
service.PlatformTypeSafe: "typesafe",
|
||||
}
|
||||
|
||||
// SyncPricingModels 返回 LiteLLM 定价目录中指定平台的最新模型列表
|
||||
|
||||
@@ -565,7 +565,7 @@ func TestSyncPricingModels_ValidPlatform_EmptyService(t *testing.T) {
|
||||
svc := service.NewPricingService(nil, nil)
|
||||
router := setupSyncPricingModelsRouter(svc)
|
||||
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity", "grok", "kimi", "zhipu", "deepseek", "minimax"} {
|
||||
for _, platform := range []string{"anthropic", "openai", "gemini", "antigravity", "grok", "kimi", "zhipu", "deepseek", "minimax", "typesafe"} {
|
||||
req := httptest.NewRequest(http.MethodGet, "/channels/pricing/sync-models?platform="+platform, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
|
||||
type claudeResetReader interface {
|
||||
Query(context.Context, int64) (*service.ClaudeResetCredits, error)
|
||||
Redeem(context.Context, int64, string) (*service.ClaudeResetOutcome, error)
|
||||
}
|
||||
|
||||
func (h *AccountHandler) SetClaudeResetCreditService(s *service.ClaudeResetCreditService) {
|
||||
@@ -32,3 +33,24 @@ func (h *AccountHandler) ClaudeResetCredits(c *gin.Context) {
|
||||
}
|
||||
response.Success(c, status)
|
||||
}
|
||||
|
||||
// RedeemClaudeResetCredit consumes the upstream next reset credit. The request has
|
||||
// no body: the server picks the grant; the Idempotency-Key header identifies one
|
||||
// operator confirmation and replays its outcome.
|
||||
func (h *AccountHandler) RedeemClaudeResetCredit(c *gin.Context) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return
|
||||
}
|
||||
if h.claudeResetCredits == nil {
|
||||
response.Error(c, 503, "Claude reset service unavailable")
|
||||
return
|
||||
}
|
||||
result, err := h.claudeResetCredits.Redeem(c.Request.Context(), id, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type claudeResetHandlerStub struct {
|
||||
key string
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Query(context.Context, int64) (*service.ClaudeResetCredits, error) {
|
||||
return &service.ClaudeResetCredits{Credits: []service.ClaudeResetCredit{}}, nil
|
||||
}
|
||||
|
||||
func (s *claudeResetHandlerStub) Redeem(_ context.Context, _ int64, key string) (*service.ClaudeResetOutcome, error) {
|
||||
s.calls++
|
||||
s.key = key
|
||||
if key == "" {
|
||||
return nil, service.ErrIdempotencyKeyRequired
|
||||
}
|
||||
return &service.ClaudeResetOutcome{Outcome: "reset", Cleared: []string{"five_hour"}}, nil
|
||||
}
|
||||
|
||||
func TestClaudeResetHandlerRedeemContract(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
stub := &claudeResetHandlerStub{}
|
||||
h := &AccountHandler{claudeResetCredits: stub}
|
||||
r := gin.New()
|
||||
r.POST("/accounts/:id/claude/reset-credits/redeem", h.RedeemClaudeResetCredit)
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil)
|
||||
req.Header.Set("Idempotency-Key", "same-operation")
|
||||
w := httptest.NewRecorder()
|
||||
r.ServeHTTP(w, req)
|
||||
require.Equal(t, http.StatusOK, w.Code)
|
||||
require.Equal(t, "same-operation", stub.key)
|
||||
require.Contains(t, w.Body.String(), `"outcome":"reset"`)
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/1/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Contains(t, w.Body.String(), "IDEMPOTENCY_KEY_REQUIRED")
|
||||
|
||||
w = httptest.NewRecorder()
|
||||
r.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/accounts/abc/claude/reset-credits/redeem", nil))
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
require.Equal(t, 2, stub.calls)
|
||||
}
|
||||
@@ -184,7 +184,7 @@ func sanitizeUpdateGroupRequestForSimpleMode(req *UpdateGroupRequest) {
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
|
||||
RateMultiplier float64 `json:"rate_multiplier"`
|
||||
IsExclusive bool `json:"is_exclusive"`
|
||||
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
|
||||
@@ -258,7 +258,7 @@ type CreateGroupRequest struct {
|
||||
type UpdateGroupRequest struct {
|
||||
Name string `json:"name"`
|
||||
Description *string `json:"description"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go composite"`
|
||||
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe composite"`
|
||||
RateMultiplier *float64 `json:"rate_multiplier"`
|
||||
IsExclusive *bool `json:"is_exclusive"`
|
||||
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
|
||||
@@ -332,7 +332,7 @@ type UpdateGroupRequest struct {
|
||||
type CompositeRouteRequest struct {
|
||||
PublicModel string `json:"public_model" binding:"required"`
|
||||
MatchType string `json:"match_type" binding:"omitempty,oneof=exact prefix"`
|
||||
TargetPlatform string `json:"target_platform" binding:"required,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go"`
|
||||
TargetPlatform string `json:"target_platform" binding:"required,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek minimax opencode_go typesafe"`
|
||||
UpstreamModel string `json:"upstream_model"`
|
||||
Endpoint string `json:"endpoint" binding:"omitempty,oneof=any messages count_tokens responses chat_completions embeddings images gemini"`
|
||||
Priority int `json:"priority"`
|
||||
|
||||
@@ -28,6 +28,7 @@ func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
|
||||
allowed := []string{
|
||||
"anthropic", "openai", "gemini", "antigravity", "grok",
|
||||
"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "composite",
|
||||
"typesafe",
|
||||
}
|
||||
for _, platform := range allowed {
|
||||
t.Run("create_"+platform, func(t *testing.T) {
|
||||
@@ -71,8 +72,8 @@ func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompositeRouteTargetPlatform_AllowsCNProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go"} {
|
||||
func TestCompositeRouteTargetPlatform_AllowsConcreteProviders(t *testing.T) {
|
||||
for _, platform := range []string{"kimi", "zhipu", "deepseek", "minimax", "opencode_go", "typesafe"} {
|
||||
var req CompositeRouteRequest
|
||||
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
|
||||
require.NoError(t, bindGroupPlatformJSON(t, &req, body))
|
||||
|
||||
@@ -125,6 +125,7 @@ type AdminPaymentOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
RechargeCode string `json:"recharge_code,omitempty"`
|
||||
OutTradeNo string `json:"out_trade_no"`
|
||||
@@ -182,6 +183,7 @@ func sanitizeAdminPaymentOrderForResponse(order *dbent.PaymentOrder) *AdminPayme
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
RechargeCode: order.RechargeCode,
|
||||
OutTradeNo: order.OutTradeNo,
|
||||
|
||||
@@ -362,6 +362,9 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
PaymentBalanceRechargeMultiplier: paymentCfg.BalanceRechargeMultiplier,
|
||||
PaymentSubscriptionUSDToCNYRate: paymentCfg.SubscriptionUSDToCNYRate,
|
||||
PaymentRechargeFeeRate: paymentCfg.RechargeFeeRate,
|
||||
PaymentRechargeBonusTiers: rechargeBonusTiersToDTO(paymentCfg.RechargeBonusTiers),
|
||||
PaymentRechargeBonusMode: rechargeBonusModeToDTO(paymentCfg.RechargeBonusMode),
|
||||
PaymentRechargeBonusNotice: paymentCfg.RechargeBonusNotice,
|
||||
PaymentLoadBalanceStrat: paymentCfg.LoadBalanceStrategy,
|
||||
PaymentProductNamePrefix: paymentCfg.ProductNamePrefix,
|
||||
PaymentProductNameSuffix: paymentCfg.ProductNameSuffix,
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
// rechargeBonusTiersFromDTO 请求 nil 表示未携带该字段(保持现值);空数组表示清空阶梯。
|
||||
func rechargeBonusTiersFromDTO(items *[]dto.RechargeBonusTier) *[]service.RechargeBonusTier {
|
||||
if items == nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]service.RechargeBonusTier, 0, len(*items))
|
||||
for _, item := range *items {
|
||||
out = append(out, service.RechargeBonusTier{MinAmount: item.MinAmount, BonusPercent: item.BonusPercent})
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// rechargeBonusModeToDTO 输出已归一化的模式(空/非法按 bonus)。
|
||||
func rechargeBonusModeToDTO(mode string) string {
|
||||
normalized, _ := service.NormalizeRechargeBonusMode(mode)
|
||||
return normalized
|
||||
}
|
||||
|
||||
// rechargeBonusTiersToDTO 始终返回非 nil 切片,空配置输出 []。
|
||||
func rechargeBonusTiersToDTO(items []service.RechargeBonusTier) []dto.RechargeBonusTier {
|
||||
out := make([]dto.RechargeBonusTier, 0, len(items))
|
||||
for _, item := range items {
|
||||
out = append(out, dto.RechargeBonusTier{MinAmount: item.MinAmount, BonusPercent: item.BonusPercent})
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -312,11 +312,15 @@ type UpdateSettingsRequest struct {
|
||||
PaymentBalanceRechargeMultiplier *float64 `json:"payment_balance_recharge_multiplier"`
|
||||
PaymentSubscriptionUSDToCNYRate *float64 `json:"payment_subscription_usd_to_cny_rate"`
|
||||
PaymentRechargeFeeRate *float64 `json:"payment_recharge_fee_rate"`
|
||||
PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix *string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix *string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL *string `json:"payment_help_image_url"`
|
||||
PaymentHelpText *string `json:"payment_help_text"`
|
||||
// nil 表示不更新;空数组表示清空阶梯
|
||||
PaymentRechargeBonusTiers *[]dto.RechargeBonusTier `json:"payment_recharge_bonus_tiers"`
|
||||
PaymentRechargeBonusMode *string `json:"payment_recharge_bonus_mode"`
|
||||
PaymentRechargeBonusNotice *string `json:"payment_recharge_bonus_notice"`
|
||||
PaymentLoadBalanceStrat *string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix *string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix *string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL *string `json:"payment_help_image_url"`
|
||||
PaymentHelpText *string `json:"payment_help_text"`
|
||||
|
||||
// Cancel rate limit
|
||||
PaymentCancelRateLimitEnabled *bool `json:"payment_cancel_rate_limit_enabled"`
|
||||
@@ -2128,6 +2132,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
BalanceRechargeMultiplier: req.PaymentBalanceRechargeMultiplier,
|
||||
SubscriptionUSDToCNYRate: req.PaymentSubscriptionUSDToCNYRate,
|
||||
RechargeFeeRate: req.PaymentRechargeFeeRate,
|
||||
RechargeBonusTiers: rechargeBonusTiersFromDTO(req.PaymentRechargeBonusTiers),
|
||||
RechargeBonusMode: req.PaymentRechargeBonusMode,
|
||||
RechargeBonusNotice: req.PaymentRechargeBonusNotice,
|
||||
LoadBalanceStrategy: req.PaymentLoadBalanceStrat,
|
||||
ProductNamePrefix: req.PaymentProductNamePrefix,
|
||||
ProductNameSuffix: req.PaymentProductNameSuffix,
|
||||
@@ -2407,6 +2414,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
PaymentBalanceRechargeMultiplier: updatedPaymentCfg.BalanceRechargeMultiplier,
|
||||
PaymentSubscriptionUSDToCNYRate: updatedPaymentCfg.SubscriptionUSDToCNYRate,
|
||||
PaymentRechargeFeeRate: updatedPaymentCfg.RechargeFeeRate,
|
||||
PaymentRechargeBonusTiers: rechargeBonusTiersToDTO(updatedPaymentCfg.RechargeBonusTiers),
|
||||
PaymentRechargeBonusMode: rechargeBonusModeToDTO(updatedPaymentCfg.RechargeBonusMode),
|
||||
PaymentRechargeBonusNotice: updatedPaymentCfg.RechargeBonusNotice,
|
||||
PaymentLoadBalanceStrat: updatedPaymentCfg.LoadBalanceStrategy,
|
||||
PaymentProductNamePrefix: updatedPaymentCfg.ProductNamePrefix,
|
||||
PaymentProductNameSuffix: updatedPaymentCfg.ProductNameSuffix,
|
||||
@@ -2483,6 +2493,7 @@ func hasPaymentFields(req UpdateSettingsRequest) bool {
|
||||
req.PaymentEnabledTypes != nil || req.PaymentBalanceDisabled != nil ||
|
||||
req.PaymentBalanceRechargeMultiplier != nil || req.PaymentSubscriptionUSDToCNYRate != nil ||
|
||||
req.PaymentRechargeFeeRate != nil ||
|
||||
req.PaymentRechargeBonusTiers != nil || req.PaymentRechargeBonusMode != nil || req.PaymentRechargeBonusNotice != nil ||
|
||||
req.PaymentLoadBalanceStrat != nil || req.PaymentProductNamePrefix != nil ||
|
||||
req.PaymentProductNameSuffix != nil || req.PaymentHelpImageURL != nil ||
|
||||
req.PaymentHelpText != nil || req.PaymentCancelRateLimitEnabled != nil ||
|
||||
|
||||
@@ -3593,3 +3593,20 @@ func (oauthPendingFlowTotpEncryptorStub) Encrypt(plaintext string) (string, erro
|
||||
func (oauthPendingFlowTotpEncryptorStub) Decrypt(ciphertext string) (string, error) {
|
||||
return ciphertext, nil
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) IncrVerificationCodeAttempts(_ context.Context, email string) (int, error) {
|
||||
data := s.verificationCodes[email]
|
||||
if data == nil {
|
||||
return 0, errors.New("verification code not found")
|
||||
}
|
||||
data.Attempts++
|
||||
return data.Attempts, nil
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) IncrNotifyVerifyCodeAttempts(context.Context, string) (int, error) {
|
||||
return 0, errors.New("notify verification code not found")
|
||||
}
|
||||
|
||||
func (s *oauthPendingFlowEmailCacheStub) ConsumePasswordResetToken(context.Context, string, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
package dto
|
||||
|
||||
// RechargeBonusTier 充值赠送档位:余额充值支付金额 ≥ MinAmount 时,在到账基数上赠送 BonusPercent%。
|
||||
type RechargeBonusTier struct {
|
||||
MinAmount float64 `json:"min_amount"`
|
||||
BonusPercent float64 `json:"bonus_percent"`
|
||||
}
|
||||
@@ -279,11 +279,15 @@ type SystemSettings struct {
|
||||
PaymentBalanceRechargeMultiplier float64 `json:"payment_balance_recharge_multiplier"`
|
||||
PaymentSubscriptionUSDToCNYRate float64 `json:"payment_subscription_usd_to_cny_rate"`
|
||||
PaymentRechargeFeeRate float64 `json:"payment_recharge_fee_rate"`
|
||||
PaymentLoadBalanceStrat string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL string `json:"payment_help_image_url"`
|
||||
PaymentHelpText string `json:"payment_help_text"`
|
||||
// 充值赠送阶梯与活动文案
|
||||
PaymentRechargeBonusTiers []RechargeBonusTier `json:"payment_recharge_bonus_tiers"`
|
||||
PaymentRechargeBonusMode string `json:"payment_recharge_bonus_mode"`
|
||||
PaymentRechargeBonusNotice string `json:"payment_recharge_bonus_notice"`
|
||||
PaymentLoadBalanceStrat string `json:"payment_load_balance_strategy"`
|
||||
PaymentProductNamePrefix string `json:"payment_product_name_prefix"`
|
||||
PaymentProductNameSuffix string `json:"payment_product_name_suffix"`
|
||||
PaymentHelpImageURL string `json:"payment_help_image_url"`
|
||||
PaymentHelpText string `json:"payment_help_text"`
|
||||
|
||||
// Cancel rate limit
|
||||
PaymentCancelRateLimitEnabled bool `json:"payment_cancel_rate_limit_enabled"`
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
|
||||
const (
|
||||
EndpointMessages = "/v1/messages"
|
||||
EndpointSystemOne = "/v1/systemone"
|
||||
EndpointChatCompletions = "/v1/chat/completions"
|
||||
EndpointEmbeddings = "/v1/embeddings"
|
||||
EndpointAlphaSearch = "/v1/alpha/search"
|
||||
@@ -94,6 +95,8 @@ func NormalizeInboundEndpoint(path string) string {
|
||||
return EndpointChatCompletions
|
||||
case strings.Contains(path, EndpointMessages):
|
||||
return EndpointMessages
|
||||
case strings.Contains(path, EndpointSystemOne):
|
||||
return EndpointSystemOne
|
||||
case strings.Contains(path, EndpointImagesGenerations) || strings.Contains(path, "/images/generations"):
|
||||
return EndpointImagesGenerations
|
||||
case strings.Contains(path, EndpointImagesEdits) || strings.Contains(path, "/images/edits"):
|
||||
@@ -223,6 +226,9 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string {
|
||||
case service.PlatformAnthropic:
|
||||
return EndpointMessages
|
||||
|
||||
case service.PlatformTypeSafe:
|
||||
return EndpointSystemOne
|
||||
|
||||
case service.PlatformGemini:
|
||||
return EndpointGeminiModels
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ func TestNormalizeInboundEndpoint(t *testing.T) {
|
||||
}{
|
||||
// Direct canonical paths.
|
||||
{"/v1/messages", EndpointMessages},
|
||||
{"/v1/systemone", EndpointSystemOne},
|
||||
{"/v1/chat/completions", EndpointChatCompletions},
|
||||
{"/v1/embeddings", EndpointEmbeddings},
|
||||
{"/v1/alpha/search", EndpointAlphaSearch},
|
||||
@@ -97,6 +98,7 @@ func TestDeriveUpstreamEndpoint(t *testing.T) {
|
||||
}{
|
||||
// Anthropic.
|
||||
{"anthropic messages", EndpointMessages, "/v1/messages", service.PlatformAnthropic, EndpointMessages},
|
||||
{"typesafe system one", EndpointSystemOne, "/v1/systemone", service.PlatformTypeSafe, EndpointSystemOne},
|
||||
|
||||
// Gemini.
|
||||
{"gemini models", EndpointGeminiModels, "/v1beta/models/gemini:gen", service.PlatformGemini, EndpointGeminiModels},
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
@@ -222,6 +223,9 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.errorResponse) {
|
||||
return
|
||||
}
|
||||
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
@@ -263,6 +267,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 设置请求所属分组 ID(用于渠道级功能判断,如 WebSearch 模拟)
|
||||
parsedReq.GroupID = apiKey.GroupID
|
||||
|
||||
@@ -1149,7 +1166,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
}
|
||||
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, true)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
@@ -1193,6 +1210,10 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
writeGrokModelsList(c, xai.DefaultModelIDs())
|
||||
return
|
||||
}
|
||||
if platform == service.PlatformTypeSafe {
|
||||
writeModelsList(c, platform, []string{typesafe.JevLatestModel})
|
||||
return
|
||||
}
|
||||
|
||||
writeModelsListResponse(c, claude.DefaultModels)
|
||||
}
|
||||
@@ -1244,7 +1265,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
platform = group.Platform
|
||||
}
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID)
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID, false)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||
if group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
@@ -1270,14 +1291,20 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
return fallbackModels
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64) []string {
|
||||
// compositeAvailableModels lists the models the composite group can serve.
|
||||
// includeSystemOne adds TypeSafe models, which only work through /v1/systemone;
|
||||
// LLM client catalogs (Codex) must exclude them.
|
||||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, includeSystemOne bool) []string {
|
||||
if h == nil || h.gatewayService == nil {
|
||||
return nil
|
||||
}
|
||||
seen := make(map[string]struct{})
|
||||
models := make([]string, 0)
|
||||
schedulablePlatforms := h.gatewayService.GetSchedulablePlatforms(ctx, groupID)
|
||||
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo} {
|
||||
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo, service.PlatformTypeSafe} {
|
||||
if platform == service.PlatformTypeSafe && !includeSystemOne {
|
||||
continue
|
||||
}
|
||||
platformModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
|
||||
if len(platformModels) == 0 {
|
||||
// CN 供应商没有静态默认模型列表(defaultModelIDsForPlatform 的
|
||||
@@ -1461,9 +1488,14 @@ func defaultModelIDsForPlatform(platform string) []string {
|
||||
return xai.DefaultModelIDs()
|
||||
case service.PlatformOpenCodeGo:
|
||||
return service.DefaultOpenCodeGoModelIDs()
|
||||
case service.PlatformTypeSafe:
|
||||
return []string{"jev-latest"}
|
||||
case service.PlatformComposite:
|
||||
ids := make([]string, 0)
|
||||
seen := make(map[string]struct{})
|
||||
// TypeSafe is deliberately absent: jev-latest only works through
|
||||
// /v1/systemone, so the static fallback never advertises it to LLM
|
||||
// clients. compositeAvailableModels lists it when the group can serve it.
|
||||
for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformMiniMax, service.PlatformOpenCodeGo} {
|
||||
for _, id := range defaultModelIDsForPlatform(concretePlatform) {
|
||||
if _, ok := seen[id]; ok {
|
||||
@@ -2149,6 +2181,9 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.errorResponse) {
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, parsedReq.Model, parsedReq.Stream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsedReq.Stream, false)))
|
||||
@@ -2514,9 +2549,12 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -2544,7 +2582,7 @@ func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context,
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
|
||||
@@ -82,6 +82,9 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.chatCompletionsErrorResponse) {
|
||||
return
|
||||
}
|
||||
reqStream, ok := parseOpenAICompatibleStream(body)
|
||||
if !ok {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
|
||||
@@ -101,8 +104,11 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/chat/completions is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.chatCompletionsErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -144,6 +150,19 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.cc.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.chatCompletionsErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "chat_completions")
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
middleware "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"
|
||||
)
|
||||
|
||||
// groupScopedSchedulerCache 按桶的分组只返回该分组的成员账号,并记录被查询过的分组。
|
||||
type groupScopedSchedulerCache struct {
|
||||
*fakeSchedulerCache
|
||||
mu sync.Mutex
|
||||
groupIDs []int64
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) GetSnapshot(_ context.Context, bucket service.SchedulerBucket) ([]*service.Account, bool, error) {
|
||||
c.mu.Lock()
|
||||
c.groupIDs = append(c.groupIDs, bucket.GroupID)
|
||||
c.mu.Unlock()
|
||||
var members []*service.Account
|
||||
for _, account := range c.accounts {
|
||||
for _, ag := range account.AccountGroups {
|
||||
if ag.GroupID == bucket.GroupID {
|
||||
members = append(members, account)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
return members, true, nil
|
||||
}
|
||||
|
||||
func (c *groupScopedSchedulerCache) queriedGroupIDs() []int64 {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return append([]int64(nil), c.groupIDs...)
|
||||
}
|
||||
|
||||
type groupMapRepo struct {
|
||||
*fakeGroupRepo
|
||||
groups map[int64]*service.Group
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByID(_ context.Context, id int64) (*service.Group, error) {
|
||||
if group, ok := r.groups[id]; ok {
|
||||
return group, nil
|
||||
}
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
|
||||
func (r *groupMapRepo) GetByIDLite(ctx context.Context, id int64) (*service.Group, error) {
|
||||
return r.GetByID(ctx, id)
|
||||
}
|
||||
|
||||
func TestGatewayOpenAICompatibleHandlersClaudeCodeOnlyFallback(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
const (
|
||||
primaryGroupID = int64(9300)
|
||||
fallbackGroupID = int64(9301)
|
||||
primaryAccountID = int64(9310)
|
||||
fallbackAccountID = int64(9311)
|
||||
)
|
||||
|
||||
// 账号不带 api_key:转发在取令牌时失败,请求不会触达上游,断言只看选号结果。
|
||||
newAccount := func(id, groupID int64) *service.Account {
|
||||
return &service.Account{
|
||||
ID: id, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true, Concurrency: 1,
|
||||
AccountGroups: []service.AccountGroup{{AccountID: id, GroupID: groupID}},
|
||||
}
|
||||
}
|
||||
|
||||
endpoints := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
call func(*GatewayHandler, *gin.Context)
|
||||
}{
|
||||
{
|
||||
name: "responses", path: "/v1/responses",
|
||||
body: `{"model":"claude-sonnet-4-5","input":"hello","stream":false}`,
|
||||
call: (*GatewayHandler).Responses,
|
||||
},
|
||||
{
|
||||
name: "chat completions", path: "/v1/chat/completions",
|
||||
body: `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hello"}],"stream":false}`,
|
||||
call: (*GatewayHandler).ChatCompletions,
|
||||
},
|
||||
}
|
||||
|
||||
for _, ep := range endpoints {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
hasFallback bool
|
||||
}{
|
||||
{name: "with fallback group", hasFallback: true},
|
||||
{name: "without fallback group", hasFallback: false},
|
||||
} {
|
||||
t.Run(ep.name+"/"+tc.name, func(t *testing.T) {
|
||||
primary := &service.Group{
|
||||
ID: primaryGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive, ClaudeCodeOnly: true,
|
||||
}
|
||||
if tc.hasFallback {
|
||||
fallbackID := fallbackGroupID
|
||||
primary.FallbackGroupID = &fallbackID
|
||||
}
|
||||
fallback := &service.Group{
|
||||
ID: fallbackGroupID, Hydrated: true, Platform: service.PlatformAnthropic,
|
||||
Status: service.StatusActive,
|
||||
}
|
||||
|
||||
schedulerCache := &groupScopedSchedulerCache{fakeSchedulerCache: &fakeSchedulerCache{accounts: []*service.Account{
|
||||
newAccount(primaryAccountID, primaryGroupID),
|
||||
newAccount(fallbackAccountID, fallbackGroupID),
|
||||
}}}
|
||||
gatewayService := service.NewGatewayService(
|
||||
nil, &groupMapRepo{fakeGroupRepo: &fakeGroupRepo{}, groups: map[int64]*service.Group{
|
||||
primaryGroupID: primary,
|
||||
fallbackGroupID: fallback,
|
||||
}}, nil, nil, nil, nil, nil, nil, nil,
|
||||
service.NewSchedulerSnapshotService(schedulerCache, nil, nil, nil, nil),
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
billingCacheService := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCacheService.Stop)
|
||||
h := &GatewayHandler{
|
||||
gatewayService: gatewayService,
|
||||
billingCacheService: billingCacheService,
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(&fakeConcurrencyCache{}), SSEPingFormatClaude, 0),
|
||||
maxAccountSwitches: 1,
|
||||
cfg: cfg,
|
||||
}
|
||||
|
||||
primaryGroupIDRef := primaryGroupID
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9320, UserID: 9330, GroupID: &primaryGroupIDRef, Group: primary, Status: service.StatusActive,
|
||||
User: &service.User{ID: 9330, Concurrency: 10, Balance: 100},
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.Group, primary)
|
||||
req := httptest.NewRequest(http.MethodPost, ep.path, bytes.NewBufferString(ep.body)).WithContext(ctx)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.Request = req
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
|
||||
|
||||
ep.call(h, c)
|
||||
|
||||
selected, reachedSelection := c.Get(opsAccountIDKey)
|
||||
if !tc.hasFallback {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "This group is restricted to Claude Code clients")
|
||||
require.False(t, reachedSelection, "a Claude Code only group without fallback must be rejected before account selection")
|
||||
require.Empty(t, schedulerCache.queriedGroupIDs())
|
||||
return
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "restricted to Claude Code clients")
|
||||
require.True(t, reachedSelection, "a Claude Code only group with fallback must reach account selection")
|
||||
require.Equal(t, fallbackAccountID, selected)
|
||||
queried := schedulerCache.queriedGroupIDs()
|
||||
require.Contains(t, queried, fallbackGroupID)
|
||||
require.NotContains(t, queried, primaryGroupID)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -82,6 +82,9 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||||
return
|
||||
}
|
||||
if rejectSystemOneOnlyPlatform(c, apiKey, h.responsesErrorResponse) {
|
||||
return
|
||||
}
|
||||
reqStream, ok := parseOpenAICompatibleStream(body)
|
||||
if !ok {
|
||||
h.responsesErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
|
||||
@@ -105,13 +108,11 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
// 解析渠道级模型映射
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(requestCtx, apiKey.GroupID, reqModel)
|
||||
|
||||
// Claude Code only restriction:
|
||||
// /v1/responses is never a Claude Code endpoint.
|
||||
// When claude_code_only is enabled, this endpoint is rejected.
|
||||
// The existing service-layer checkClaudeCodeRestriction handles degradation
|
||||
// to fallback groups when the Forward path calls SelectAccountForModelWithExclusions.
|
||||
// Here we just reject at handler level since /v1/responses clients can't be Claude Code.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly {
|
||||
// Claude Code only restriction: /v1/responses is never a Claude Code
|
||||
// endpoint. With a fallback group the request continues and account selection
|
||||
// (checkClaudeCodeRestriction) schedules it in the fallback group; without one
|
||||
// it is rejected here.
|
||||
if apiKey.Group != nil && apiKey.Group.ClaudeCodeOnly && apiKey.Group.FallbackGroupID == nil {
|
||||
h.responsesErrorResponse(c, http.StatusForbidden, "permission_error",
|
||||
"This group is restricted to Claude Code clients (/v1/messages only)")
|
||||
return
|
||||
@@ -153,6 +154,19 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gateway.responses.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.responsesErrorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Parse request for session hash
|
||||
bodyRef := service.NewRequestBodyRef(body)
|
||||
parsedReq, _ := service.ParseGatewayRequest(bodyRef, "responses")
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// inflightReservationEstimator 估算单请求在途预留金额(USD);false 表示无法定价。
|
||||
type inflightReservationEstimator interface {
|
||||
EstimateInflightReservation(ctx context.Context, apiKey *service.APIKey, req service.InflightEstimateRequest) (float64, bool)
|
||||
}
|
||||
|
||||
// requestMaxOutputTokens 从请求体中提取输出 token 上限(兼容 Anthropic / OpenAI Chat / Responses / Gemini)。
|
||||
func requestMaxOutputTokens(body []byte) int {
|
||||
for _, path := range []string{"max_tokens", "max_completion_tokens", "max_output_tokens", "generationConfig.maxOutputTokens", "generation_config.max_output_tokens"} {
|
||||
if v := gjson.GetBytes(body, path); v.Exists() && v.Type == gjson.Number && v.Int() > 0 {
|
||||
return int(v.Int())
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// tokenInflightEstimate 文本类请求的估算输入。
|
||||
func tokenInflightEstimate(model string, body []byte) service.InflightEstimateRequest {
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
MaxTokens: requestMaxOutputTokens(body),
|
||||
Kind: service.InflightEstimateToken,
|
||||
}
|
||||
}
|
||||
|
||||
func inflightNoop() {}
|
||||
|
||||
// reserveInflightBalance 在 CheckBillingEligibility 之后为余额模式请求登记在途预留。
|
||||
//
|
||||
// 成功时把预留句柄挂到 c.Request 的 context 上:之后通过 submit*UsageRecordTask 提交的
|
||||
// 计费任务会接管一个引用,直到余额缓存实际扣减后才释放。返回的 done 必须 defer 调用
|
||||
// (停止续期并归还 handler 引用;没有提交计费任务时即刻释放)。
|
||||
// 开关关闭、订阅模式、Redis 故障时返回 no-op(fail-open);无法定价时默认 fail-open,
|
||||
// 配置 fail_closed_on_unpriced=true 时返回 ErrInsufficientBalance。
|
||||
func reserveInflightBalance(
|
||||
c *gin.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (func(), error) {
|
||||
if c == nil || c.Request == nil {
|
||||
return inflightNoop, nil
|
||||
}
|
||||
ctx, done, err := reserveInflightBalanceCtx(c.Request.Context(), billing, estimator, apiKey, subscription, req)
|
||||
if err != nil {
|
||||
return inflightNoop, err
|
||||
}
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
return done, nil
|
||||
}
|
||||
|
||||
// reserveInflightBalanceCtx 同 reserveInflightBalance,但返回携带预留句柄的新 context
|
||||
// (供 WebSocket 等自管 context 的路径)。
|
||||
func reserveInflightBalanceCtx(
|
||||
ctx context.Context,
|
||||
billing *service.BillingCacheService,
|
||||
estimator inflightReservationEstimator,
|
||||
apiKey *service.APIKey,
|
||||
subscription *service.UserSubscription,
|
||||
req service.InflightEstimateRequest,
|
||||
) (context.Context, func(), error) {
|
||||
if billing == nil || estimator == nil || apiKey == nil || apiKey.User == nil || !billing.InflightReservationEnabled() {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() && subscription != nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
estimate, priced := estimator.EstimateInflightReservation(ctx, apiKey, req)
|
||||
if !priced && billing.InflightReservationFailClosedOnUnpriced() {
|
||||
return ctx, inflightNoop, service.ErrInsufficientBalance
|
||||
}
|
||||
if estimate <= 0 {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
res, err := billing.ReserveInflight(ctx, apiKey.User, apiKey.Group, subscription, estimate)
|
||||
if err != nil {
|
||||
return ctx, inflightNoop, err
|
||||
}
|
||||
if res == nil {
|
||||
return ctx, inflightNoop, nil
|
||||
}
|
||||
return service.WithInflightReservation(ctx, res), res.HandlerDone, nil
|
||||
}
|
||||
|
||||
// grokMediaInflightEstimate 媒体生成请求的估算输入;状态/内容查询返回空模型(不预留:
|
||||
// 查询会为已生成的媒体计费,不能因余额预留而拦截用户取回已付费结果)。
|
||||
func grokMediaInflightEstimate(endpoint service.GrokMediaEndpoint, model string, info service.GrokMediaRequestInfo, body []byte) service.InflightEstimateRequest {
|
||||
if !endpoint.IsGenerationRequest() {
|
||||
return service.InflightEstimateRequest{}
|
||||
}
|
||||
switch endpoint {
|
||||
case service.GrokMediaEndpointImagesGenerations, service.GrokMediaEndpointImagesEdits:
|
||||
return service.InflightEstimateRequest{Model: model, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: info.N}
|
||||
default:
|
||||
return service.InflightEstimateRequest{
|
||||
Model: model,
|
||||
BodyBytes: len(body),
|
||||
Kind: service.InflightEstimateVideo,
|
||||
Units: 1,
|
||||
VideoResolution: info.Resolution,
|
||||
VideoDurationSeconds: info.DurationSeconds,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// grokVoiceSTTBytesPerSecond STT 时长粗估(~128kbps 压缩音频)。
|
||||
const grokVoiceSTTBytesPerSecond = 16000
|
||||
|
||||
// grokVoiceInflightEstimate 语音 HTTP 接口估算:TTS 按输入字符数(百万字符),STT 按音频字节粗估时长(小时)。
|
||||
// 其他接口(custom-voices)无音频计量,返回的估算为 0(不预留)。
|
||||
func grokVoiceInflightEstimate(endpoint string, body []byte) service.InflightEstimateRequest {
|
||||
req := service.InflightEstimateRequest{Model: endpoint, Kind: service.InflightEstimateAudio, AudioMode: endpoint}
|
||||
switch endpoint {
|
||||
case "tts":
|
||||
req.AudioUnits = float64(len([]rune(extractGrokTTSInputText(body)))) / 1e6
|
||||
case "stt":
|
||||
req.AudioUnits = float64(len(body)) / grokVoiceSTTBytesPerSecond / 3600
|
||||
}
|
||||
return req
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"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"
|
||||
)
|
||||
|
||||
func TestRequestMaxOutputTokens(t *testing.T) {
|
||||
require.Equal(t, 1024, requestMaxOutputTokens([]byte(`{"max_tokens":1024}`)))
|
||||
require.Equal(t, 2048, requestMaxOutputTokens([]byte(`{"max_completion_tokens":2048}`)))
|
||||
require.Equal(t, 4096, requestMaxOutputTokens([]byte(`{"max_output_tokens":4096}`)))
|
||||
require.Equal(t, 512, requestMaxOutputTokens([]byte(`{"generationConfig":{"maxOutputTokens":512}}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{"max_tokens":"x"}`)))
|
||||
require.Equal(t, 0, requestMaxOutputTokens([]byte(`{}`)))
|
||||
}
|
||||
|
||||
type countingEstimator struct {
|
||||
calls int
|
||||
cost float64
|
||||
priced bool
|
||||
}
|
||||
|
||||
func (e *countingEstimator) EstimateInflightReservation(context.Context, *service.APIKey, service.InflightEstimateRequest) (float64, bool) {
|
||||
e.calls++
|
||||
return e.cost, e.priced
|
||||
}
|
||||
|
||||
func newInflightTestGinContext() *gin.Context {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_SkipsWhenDisabledOrSubscription(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
est := &countingEstimator{cost: 1, priced: true}
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "disabled switch must not even estimate")
|
||||
|
||||
cfg.Billing.InflightReservation.Enabled = true
|
||||
apiKey.Group = &service.Group{SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
done, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, &service.UserSubscription{}, tokenInflightEstimate("m", []byte(`{}`)))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
require.Equal(t, 0, est.calls, "subscription mode must be unaffected")
|
||||
}
|
||||
|
||||
func TestReserveInflightBalance_UnpricedFailOpenByDefaultFailClosedOptIn(t *testing.T) {
|
||||
cache := newHandlerInflightCache(10)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 1}}
|
||||
est := &countingEstimator{priced: false}
|
||||
|
||||
done, err := reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.NoError(t, err)
|
||||
done()
|
||||
|
||||
cfg.Billing.InflightReservation.FailClosedOnUnpriced = true
|
||||
_, err = reserveInflightBalance(newInflightTestGinContext(), billing, est, apiKey, nil, tokenInflightEstimate("unknown", nil))
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
}
|
||||
|
||||
// handlerInflightCache 内存版余额缓存 + 在途预留(语义同 Redis Lua)。
|
||||
type handlerInflightCache struct {
|
||||
service.BillingCache
|
||||
mu sync.Mutex
|
||||
balance float64
|
||||
res map[string]float64
|
||||
}
|
||||
|
||||
func newHandlerInflightCache(balance float64) *handlerInflightCache {
|
||||
return &handlerInflightCache{balance: balance, res: map[string]float64{}}
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserBalance(context.Context, int64) (float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.balance, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) GetUserPlatformQuotaCache(context.Context, int64, string) (*service.UserPlatformQuotaCacheEntry, bool, error) {
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReserveInflightBalance(_ context.Context, _ int64, id string, amount, balance float64, _ time.Duration) (bool, float64, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
sum := 0.0
|
||||
for _, v := range m.res {
|
||||
sum += v
|
||||
}
|
||||
if len(m.res) > 0 && balance-sum < amount {
|
||||
return false, sum, nil
|
||||
}
|
||||
m.res[id] = amount
|
||||
return true, sum, nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) ReleaseInflightBalance(_ context.Context, _ int64, id string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
delete(m.res, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *handlerInflightCache) count() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return len(m.res)
|
||||
}
|
||||
|
||||
func TestWrapUsageRecordTaskContext_HandsReservationToBillingTask(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billing := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billing.Stop)
|
||||
apiKey := &service.APIKey{User: &service.User{ID: 5}}
|
||||
|
||||
c := newInflightTestGinContext()
|
||||
done, err := reserveInflightBalance(c, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, cache.count())
|
||||
|
||||
ran := false
|
||||
task, abandon := wrapUsageRecordTaskContext(c.Request.Context(), func(context.Context) { ran = true })
|
||||
done() // handler returns; billing still pending
|
||||
require.Equal(t, 1, cache.count(), "reservation held until the billing task finishes")
|
||||
task(context.Background())
|
||||
require.True(t, ran)
|
||||
require.Equal(t, 0, cache.count())
|
||||
abandon() // idempotent with the task's own done
|
||||
|
||||
// Dropped task: the submitter abandons it and the reservation is released.
|
||||
c2 := newInflightTestGinContext()
|
||||
done2, err := reserveInflightBalance(c2, billing, &countingEstimator{cost: 0.9, priced: true}, apiKey, nil, tokenInflightEstimate("m", nil))
|
||||
require.NoError(t, err)
|
||||
_, abandon2 := wrapUsageRecordTaskContext(c2.Request.Context(), func(context.Context) {})
|
||||
done2()
|
||||
require.Equal(t, 1, cache.count())
|
||||
abandon2()
|
||||
require.Equal(t, 0, cache.count())
|
||||
}
|
||||
|
||||
// 新接入的端点(独立 web_search):在途预留超过余额时拒绝,且不残留预留。
|
||||
func TestWebSearch_RejectsWhenInflightExceedsBalance(t *testing.T) {
|
||||
cache := newHandlerInflightCache(1.5)
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
billingCache := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(billingCache.Stop)
|
||||
billing := service.NewBillingService(cfg, nil)
|
||||
gw := service.NewGatewayService(
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil, nil, billing, nil, nil,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, service.NewModelPricingResolver(nil, billing), nil, nil, nil,
|
||||
)
|
||||
h := &GatewayHandler{gatewayService: gw, billingCacheService: billingCache}
|
||||
|
||||
groupID := int64(3)
|
||||
perK := 1000.0 // $1 per search
|
||||
apiKey := &service.APIKey{
|
||||
ID: 9, User: &service.User{ID: 42, Balance: 1.5}, GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, RateMultiplier: 1, SearchPricePer1k: &perK},
|
||||
}
|
||||
|
||||
// Another in-flight request of this user already holds $1.
|
||||
held, err := billingCache.ReserveInflight(context.Background(), apiKey.User, apiKey.Group, nil, 1.0)
|
||||
require.NoError(t, err)
|
||||
defer held.HandlerDone()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/web_search", bytes.NewBufferString(`{"query":"sub2api"}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), apiKey)
|
||||
|
||||
h.WebSearch(c)
|
||||
|
||||
require.NotEqual(t, http.StatusOK, w.Code)
|
||||
require.Contains(t, w.Body.String(), "balance")
|
||||
require.Equal(t, 1, cache.count(), "rejected request must not leave a reservation behind")
|
||||
}
|
||||
@@ -27,7 +27,7 @@ func requestModelForTest(h *GatewayHandler, group *service.Group, modelID, etag
|
||||
|
||||
func TestRetrieveModelMatchesVisibleCatalogue(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformComposite} {
|
||||
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformTypeSafe, service.PlatformComposite} {
|
||||
for _, mapped := range []bool{false, true} {
|
||||
name := platform + "/fallback"
|
||||
if mapped {
|
||||
|
||||
@@ -932,6 +932,10 @@ func TestDefaultCodexModelIDsForPlatform_DeepSeekUsesDeepSeekModels(t *testing.T
|
||||
require.Equal(t, defaultModelIDsForPlatform(service.PlatformAnthropic), defaultCodexModelIDsForPlatform(service.PlatformAnthropic))
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForPlatform_TypeSafeUsesJev(t *testing.T) {
|
||||
require.Equal(t, []string{"jev-latest"}, defaultModelIDsForPlatform(service.PlatformTypeSafe))
|
||||
}
|
||||
|
||||
func TestGatewayCodexModels_DeepSeekWithoutMappingUsesDeepSeekDefaults(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 130
|
||||
@@ -1499,9 +1503,9 @@ func TestGatewayModels_GPT6SolLunaDiscoveryRespectsGroupAndAccountRestrictions(t
|
||||
restricted bool
|
||||
want []string
|
||||
}{
|
||||
{"selected and ordered", []string{"gpt-6-luna", "gpt-6-sol"}, false, []string{"gpt-6-luna", "gpt-6-sol"}},
|
||||
{"selected and ordered", []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}, false, []string{"gpt-6.1-sol", "gpt-6-luna", "gpt-6-sol"}},
|
||||
{"group excludes new models", []string{"gpt-5.6-sol"}, false, []string{"gpt-5.6-sol"}},
|
||||
{"account restricts new models", []string{"gpt-6-sol", "gpt-6-luna", "gpt-5.6-sol"}, true, []string{"gpt-5.6-sol"}},
|
||||
{"account restricts new models", []string{"gpt-6.1-sol", "gpt-6-sol", "gpt-6-luna", "gpt-5.6-sol"}, true, []string{"gpt-5.6-sol"}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
groupID := int64(25)
|
||||
@@ -1522,3 +1526,47 @@ func TestGatewayModels_GPT6SolLunaDiscoveryRespectsGroupAndAccountRestrictions(t
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Scenario: jev-latest only works through /v1/systemone, so Composite groups list
|
||||
// it in /v1/models only when they can serve it, and never in the Codex manifest.
|
||||
func TestGatewayModels_CompositeTypeSafeListingScope(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(66)
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {{ID: 1, Platform: service.PlatformAnthropic}, {ID: 2, Platform: service.PlatformTypeSafe, Type: service.AccountTypeAPIKey}},
|
||||
},
|
||||
})
|
||||
newContext := func(path string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, path, nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
c, rec := newContext("/v1/models")
|
||||
h.Models(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var models gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &models))
|
||||
require.Contains(t, modelIDsForTest(models.Data), "jev-latest")
|
||||
require.Contains(t, modelIDsForTest(models.Data), "claude-opus-4-6")
|
||||
|
||||
c, rec = newContext("/models?client_version=0.147.0")
|
||||
h.CodexModels(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var manifest codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &manifest))
|
||||
slugs := codexModelSlugsForTest(manifest.Models)
|
||||
require.Contains(t, slugs, "claude-opus-4-6")
|
||||
require.NotContains(t, slugs, "jev-latest")
|
||||
}
|
||||
|
||||
func TestDefaultModelIDsForPlatform_CompositeFallbackExcludesTypeSafe(t *testing.T) {
|
||||
require.NotContains(t, defaultModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
require.NotContains(t, defaultCodexModelIDsForPlatform(service.PlatformComposite), "jev-latest")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,287 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/typesafe"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const systemOneOnlyPlatformMessage = "TypeSafe models are only available through POST /v1/systemone"
|
||||
|
||||
// rejectSystemOneOnlyPlatform stops TypeSafe traffic from entering a non System
|
||||
// One protocol chain. TypeSafe accounts only speak the native System One
|
||||
// protocol; letting them reach the Anthropic/OpenAI converters would send
|
||||
// foreign payloads (and the account key) to the wrong upstream path and feed
|
||||
// the resulting auth failures back into account state.
|
||||
func rejectSystemOneOnlyPlatform(c *gin.Context, apiKey *service.APIKey, writeError func(*gin.Context, int, string, string)) bool {
|
||||
if _, forced := middleware2.GetForcePlatformFromContext(c); forced {
|
||||
return false
|
||||
}
|
||||
if effectiveAPIKeyPlatform(c, apiKey) != service.PlatformTypeSafe {
|
||||
return false
|
||||
}
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
writeError(c, http.StatusNotFound, "not_found_error", systemOneOnlyPlatformMessage)
|
||||
return true
|
||||
}
|
||||
|
||||
// SystemOne proxies TypeSafe's native, non-streaming System One protocol.
|
||||
func (h *GatewayHandler) SystemOne(c *gin.Context) {
|
||||
requestStart := time.Now()
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil || apiKey.Group == nil {
|
||||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||||
return
|
||||
}
|
||||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||||
return
|
||||
}
|
||||
reqLog := requestLogger(c, "handler.gateway.systemone",
|
||||
zap.Int64("user_id", subject.UserID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Any("group_id", apiKey.GroupID),
|
||||
)
|
||||
|
||||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||||
if err != nil {
|
||||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||||
return
|
||||
}
|
||||
model, err := typesafe.ValidateSystemOneRequest(body)
|
||||
if err != nil {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, model)
|
||||
if apiKey.Group.Platform != service.PlatformTypeSafe &&
|
||||
(apiKey.Group.Platform != service.PlatformComposite || !compositeTargetPlatformAllowed(c, apiKey, model, service.PlatformTypeSafe)) {
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "System One is only available for TypeSafe and compatible Composite groups")
|
||||
return
|
||||
}
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolTypeSafeSystemOne, model, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
setOpsRequestContext(c, model, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeSync))
|
||||
service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds())
|
||||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||||
c.Request = c.Request.WithContext(pricingCtx)
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, model)
|
||||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||||
|
||||
streamStarted := false
|
||||
userRelease, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("systemone.user_slot_acquire_failed", zap.Error(err))
|
||||
h.handleConcurrencyError(c, err, "user", false)
|
||||
return
|
||||
}
|
||||
// 在请求结束或 Context 取消时确保释放槽位,避免客户端断开造成泄漏。
|
||||
userRelease = wrapReleaseOnDone(c.Request.Context(), userRelease)
|
||||
if userRelease != nil {
|
||||
defer userRelease()
|
||||
}
|
||||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||||
reqLog.Info("systemone.billing_eligibility_check_failed", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(model, body))
|
||||
if err != nil {
|
||||
reqLog.Info("systemone.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
fs := NewFailoverState(h.maxAccountSwitches, false)
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, "", model, fs.FailedAccountIDs, "", subject.UserID)
|
||||
if err == nil && (selection == nil || selection.Account == nil) {
|
||||
err = service.ErrNoAvailableAccounts
|
||||
}
|
||||
if err != nil {
|
||||
if failoverClientGone(c) {
|
||||
reqLog.Info("systemone.account_select_aborted_client_disconnected", zap.Error(err))
|
||||
return
|
||||
}
|
||||
if len(fs.FailedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, model, model, service.PlatformTypeSafe)
|
||||
cls = classifySelectionFailureError(err, cls)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
reqLog.Warn("systemone.select_account_no_available", zap.Bool("model_not_found", cls.ModelNotFound), zap.Error(err))
|
||||
h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
|
||||
return
|
||||
}
|
||||
switch fs.HandleSelectionExhausted(c.Request.Context()) {
|
||||
case FailoverContinue:
|
||||
continue
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
default:
|
||||
if fs.LastFailoverErr != nil {
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformTypeSafe, false)
|
||||
} else {
|
||||
h.handleFailoverExhaustedSimple(c, http.StatusBadGateway, false)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
account := selection.Account
|
||||
|
||||
accountRelease := selection.ReleaseFunc
|
||||
if !selection.Acquired {
|
||||
if selection.WaitPlan == nil {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts")
|
||||
return
|
||||
}
|
||||
accountRelease, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(c, account.ID, selection.WaitPlan.MaxConcurrency, selection.WaitPlan.Timeout, false, &streamStarted)
|
||||
if err != nil {
|
||||
reqLog.Warn("systemone.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||||
h.handleConcurrencyError(c, err, "account", false)
|
||||
return
|
||||
}
|
||||
}
|
||||
// 准入终检:与其他网关入口一致,利润控制否决的账号不得承接本次请求。
|
||||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||||
if vetoed {
|
||||
if accountRelease != nil {
|
||||
accountRelease()
|
||||
}
|
||||
reqLog.Debug("systemone.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||||
reqLog.Warn("systemone.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", profitVetoExhaustedMessage)
|
||||
return
|
||||
}
|
||||
continue
|
||||
}
|
||||
account = latest
|
||||
accountRelease = wrapReleaseOnDone(c.Request.Context(), accountRelease)
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
service.SetOpsUpstreamModel(c, model)
|
||||
|
||||
forwardStart := time.Now()
|
||||
result, forwardErr := h.gatewayService.ForwardSystemOne(c.Request.Context(), c, account, body)
|
||||
if accountRelease != nil {
|
||||
accountRelease()
|
||||
}
|
||||
service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds())
|
||||
|
||||
if forwardErr != nil {
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(forwardErr, &failoverErr) {
|
||||
switch fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr) {
|
||||
case FailoverContinue:
|
||||
reqLog.Warn("systemone.upstream_failover_switching",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.Int("upstream_status", failoverErr.StatusCode),
|
||||
zap.Int("switch_count", fs.SwitchCount),
|
||||
)
|
||||
continue
|
||||
case FailoverExhausted:
|
||||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformTypeSafe, false)
|
||||
return
|
||||
case FailoverCanceled:
|
||||
failoverClientGone(c)
|
||||
return
|
||||
}
|
||||
}
|
||||
if failoverClientGone(c) {
|
||||
return
|
||||
}
|
||||
var upstreamErr *service.SystemOneUpstreamError
|
||||
if errors.As(forwardErr, &upstreamErr) {
|
||||
status := upstreamErr.StatusCode
|
||||
if !service.IsSystemOneRequestErrorStatus(status) {
|
||||
status = http.StatusBadGateway
|
||||
}
|
||||
h.errorResponse(c, status, "upstream_error", "TypeSafe rejected the request")
|
||||
return
|
||||
}
|
||||
reqLog.Warn("systemone.forward_failed", zap.Int64("account_id", account.ID), zap.Error(forwardErr))
|
||||
if errors.Is(forwardErr, typesafe.ErrSystemOneResponseTooLarge) {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "TypeSafe response exceeds the gateway size limit")
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "TypeSafe upstream request failed")
|
||||
return
|
||||
}
|
||||
|
||||
c.Data(result.StatusCode, result.ContentType, result.Body)
|
||||
h.recordSystemOneUsage(c, apiKey, account, subscription, channelMapping, model, body, result, subject.UserID, pricingAt)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) recordSystemOneUsage(c *gin.Context, apiKey *service.APIKey, account *service.Account, subscription *service.UserSubscription, mapping service.ChannelMappingResult, model string, body []byte, result *service.SystemOneForwardResult, userID int64, pricingAt time.Time) {
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
clientIP := ip.GetClientIP(c)
|
||||
inboundEndpoint := GetInboundEndpoint(c)
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
sessionID := service.ExtractClientSessionID(c)
|
||||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||||
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: &result.ForwardResult,
|
||||
APIKey: apiKey,
|
||||
User: apiKey.User,
|
||||
Account: account,
|
||||
Subscription: subscription,
|
||||
PricingAt: pricingAt,
|
||||
InboundEndpoint: inboundEndpoint,
|
||||
UpstreamEndpoint: upstreamEndpoint,
|
||||
UserAgent: userAgent,
|
||||
IPAddress: clientIP,
|
||||
SessionID: sessionID,
|
||||
RequestPayloadHash: requestPayloadHash,
|
||||
APIKeyService: h.apiKeyService,
|
||||
QuotaPlatform: quotaPlatform,
|
||||
ChannelUsageFields: clientRequestedUsageFields(c, mapping, model, result.UpstreamModel),
|
||||
}); err != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gateway.systemone"),
|
||||
zap.Int64("user_id", userID),
|
||||
zap.Int64("api_key_id", apiKey.ID),
|
||||
zap.Int64("account_id", account.ID),
|
||||
).Error("systemone.record_usage_failed", zap.Error(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"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"
|
||||
)
|
||||
|
||||
const validSystemOneHandlerBody = `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","instructions":"Evaluate"}}}`
|
||||
|
||||
func newSystemOneHandlerContext(body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/systemone", strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func TestSystemOneRequiresAuthentication(t *testing.T) {
|
||||
c, recorder := newSystemOneHandlerContext(validSystemOneHandlerBody)
|
||||
(&GatewayHandler{}).SystemOne(c)
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "authentication_error")
|
||||
}
|
||||
|
||||
func TestSystemOneRejectsNonTypeSafeGroupBeforeScheduling(t *testing.T) {
|
||||
c, recorder := newSystemOneHandlerContext(validSystemOneHandlerBody)
|
||||
groupID := int64(3)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 4, UserID: 5, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI}})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 5, Concurrency: 1})
|
||||
|
||||
(&GatewayHandler{cfg: &config.Config{Gateway: config.GatewayConfig{MaxBodySize: 1 << 20}}}).SystemOne(c)
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "only available for TypeSafe")
|
||||
}
|
||||
|
||||
func newTypeSafeGroupContext(t *testing.T, path, body, groupPlatform string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
c, recorder := newSystemOneHandlerContext(body)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
groupID := int64(9)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 4, UserID: 5, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: groupPlatform}})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 5, Concurrency: 1})
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func TestRejectSystemOneOnlyPlatform(t *testing.T) {
|
||||
write := func(c *gin.Context, status int, errType, message string) {
|
||||
c.JSON(status, gin.H{"type": errType, "message": message})
|
||||
}
|
||||
|
||||
c, recorder := newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformTypeSafe)
|
||||
require.True(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), systemOneOnlyPlatformMessage)
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformComposite)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
c.Request = c.Request.WithContext(service.WithResolvedTargetPlatform(c.Request.Context(), service.PlatformTypeSafe))
|
||||
require.True(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/v1/messages", `{}`, service.PlatformAnthropic)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
|
||||
c, _ = newTypeSafeGroupContext(t, "/antigravity/v1/messages", `{}`, service.PlatformTypeSafe)
|
||||
c.Set(string(middleware2.ContextKeyForcePlatform), service.PlatformAntigravity)
|
||||
require.False(t, rejectSystemOneOnlyPlatform(c, mustAPIKey(t, c), write))
|
||||
}
|
||||
|
||||
func mustAPIKey(t *testing.T, c *gin.Context) *service.APIKey {
|
||||
t.Helper()
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
require.True(t, ok)
|
||||
return apiKey
|
||||
}
|
||||
|
||||
func TestTypeSafeGroupsRejectNonSystemOneProtocolsBeforeScheduling(t *testing.T) {
|
||||
h := &GatewayHandler{
|
||||
cfg: &config.Config{Gateway: config.GatewayConfig{MaxBodySize: 1 << 20}},
|
||||
gatewayService: &service.GatewayService{},
|
||||
}
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
handler func(*gin.Context)
|
||||
}{
|
||||
{"messages", "/v1/messages", `{"model":"jev-latest","max_tokens":1,"messages":[{"role":"user","content":"hi"}]}`, h.Messages},
|
||||
{"count_tokens", "/v1/messages/count_tokens", `{"model":"jev-latest","messages":[{"role":"user","content":"hi"}]}`, h.CountTokens},
|
||||
{"chat_completions", "/v1/chat/completions", `{"model":"jev-latest","messages":[{"role":"user","content":"hi"}]}`, h.ChatCompletions},
|
||||
{"responses", "/v1/responses", `{"model":"jev-latest","input":"hi"}`, h.Responses},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, recorder := newTypeSafeGroupContext(t, tc.path, tc.body, service.PlatformTypeSafe)
|
||||
tc.handler(c)
|
||||
require.Equal(t, http.StatusNotFound, recorder.Code, recorder.Body.String())
|
||||
require.Contains(t, recorder.Body.String(), systemOneOnlyPlatformMessage)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -88,6 +88,18 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"), Kind: service.InflightEstimatePerRequest, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{"type": code, "message": message}})
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
subject, _ := middleware2.GetAuthSubjectFromContext(c)
|
||||
reqLog := requestLogger(c, "handler.gateway.web_search")
|
||||
// Audit user search query before upstream Grok web_search traffic.
|
||||
|
||||
@@ -390,6 +390,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(modelName, body))
|
||||
if err != nil {
|
||||
reqLog.Info("gemini.inflight_reservation_rejected", zap.Error(err))
|
||||
status, _, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
googleError(c, status, message)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// 3) select account (sticky session based on request body)
|
||||
// 优先使用 Gemini CLI 的会话标识(privileged-user-id + tmp 目录哈希)
|
||||
sessionHash := extractGeminiCLISessionHash(c, body)
|
||||
|
||||
@@ -49,6 +49,18 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
if strings.TrimSpace(model) == "" {
|
||||
model = "grok-voice-latest"
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: model, Kind: service.InflightEstimateAudio, AudioMode: "realtime", AudioUnits: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
// Keep the HTTP response uncommitted while selecting and probing an account.
|
||||
// Realtime is not an HTTP streaming response; using reqStream=true here would
|
||||
// let the wait queue flush an SSE ping before the WebSocket handshake succeeds.
|
||||
@@ -212,6 +224,18 @@ func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokVoiceInflightEstimate(endpoint, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
if strings.TrimSpace(contentType) == "" {
|
||||
contentType = "application/json"
|
||||
|
||||
@@ -173,6 +173,18 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, grokMediaInflightEstimate(endpoint, routingModel, requestInfo, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionSeed := body
|
||||
if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" {
|
||||
sessionSeed = []byte(requestID)
|
||||
|
||||
@@ -109,6 +109,18 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: requestedModel, BodyBytes: len(body), MaxTokens: requestMaxOutputTokens(body), Kind: service.InflightEstimateToken, SearchCalls: 1})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
searchID := strings.TrimSpace(gjson.GetBytes(body, "id").String())
|
||||
sessionHash := h.gatewayService.GenerateSessionHashWithFallback(c, nil, searchID)
|
||||
profitVetoCount := 0
|
||||
|
||||
@@ -146,6 +146,19 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai_chat_completions.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
|
||||
|
||||
@@ -108,6 +108,18 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: reqModel, BodyBytes: len(body), MaxTokens: 1, Kind: service.InflightEstimateToken})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.errorResponse(c, status, code, message)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
profitVetoCount := 0
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
|
||||
@@ -258,13 +258,21 @@ func usageRecordContext(parent context.Context, base context.Context) context.Co
|
||||
return base
|
||||
}
|
||||
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) service.UsageRecordTask {
|
||||
// wrapUsageRecordTaskContext 包装计费任务:复制请求级 context 值,并接管请求的在途余额预留引用。
|
||||
// 返回的 abandon 在任务未被执行(被丢弃)时必须调用以归还预留引用;任务执行结束时(含 panic)
|
||||
// 自动归还,此时余额缓存已在计费路径中同步扣减。
|
||||
func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecordTask) (service.UsageRecordTask, func()) {
|
||||
if task == nil {
|
||||
return nil
|
||||
return nil, func() {}
|
||||
}
|
||||
done := func() {}
|
||||
if parent != nil {
|
||||
done = service.InflightReservationFromContext(parent).Acquire()
|
||||
}
|
||||
return func(ctx context.Context) {
|
||||
defer done()
|
||||
task(usageRecordContext(parent, ctx))
|
||||
}
|
||||
}, done
|
||||
}
|
||||
|
||||
func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string {
|
||||
@@ -607,6 +615,19 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留:防止并发请求在预检时看到同一份余额而集体透支。
|
||||
inflightRelease, err := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if err != nil {
|
||||
reqLog.Info("openai.inflight_reservation_rejected", zap.Error(err))
|
||||
status, code, message, retryAfter := billingErrorDetails(err)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightRelease()
|
||||
|
||||
// Generate session hash (header first; fallback to prompt_cache_key)
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, sessionHashBody)
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
|
||||
@@ -1243,6 +1264,18 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, body))
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.anthropicStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateSessionHash(c, body)
|
||||
promptCacheKey := h.gatewayService.ExtractSessionID(c, body)
|
||||
sessionHash, promptCacheKey = resolveOpenAIMessagesMetadataSession(c, sessionHash, promptCacheKey, reqModel, body)
|
||||
@@ -2541,6 +2574,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(会话级):按首帧估算一次,会话期间续期;每轮计费任务接管引用,
|
||||
// 会话结束且所有轮次扣减落地后释放。
|
||||
inflightCtx, inflightDone, inflightErr := reserveInflightBalanceCtx(ctx, h.billingCacheService, h.gatewayService, apiKey, subscription, tokenInflightEstimate(reqModel, firstMessage))
|
||||
if inflightErr != nil {
|
||||
reqLog.Info("openai.websocket_inflight_reservation_rejected", zap.Error(inflightErr))
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "billing check failed")
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
ctx = inflightCtx
|
||||
|
||||
// A WebSocket may outlive a key's remaining spending window. Recheck
|
||||
// after acquiring turn slots, including the first account-selection wait.
|
||||
// Restrict this extra check to the opt-in mode so standard-mode RPM checks
|
||||
@@ -3290,9 +3334,12 @@ func (h *OpenAIGatewayHandler) submitUsageRecordTask(parent context.Context, tas
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, abandon := wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||||
if mode.Dropped() {
|
||||
abandon()
|
||||
}
|
||||
return
|
||||
}
|
||||
// 池已停止(进程关停窗口):计费任务不能静默丢失,降级为内联同步执行。
|
||||
@@ -3329,7 +3376,7 @@ func (h *OpenAIGatewayHandler) submitMandatoryUsageRecordTask(parent context.Con
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
task, _ = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
|
||||
@@ -143,6 +143,18 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 余额模式在途预留(与计费同口径估算;计费任务扣减余额缓存后才释放)。
|
||||
inflightDone, inflightErr := reserveInflightBalance(c, h.billingCacheService, h.gatewayService, apiKey, subscription, service.InflightEstimateRequest{Model: routingModel, BodyBytes: len(body), Kind: service.InflightEstimateImage, Units: parsed.N})
|
||||
if inflightErr != nil {
|
||||
status, code, message, retryAfter := billingErrorDetails(inflightErr)
|
||||
if retryAfter > 0 {
|
||||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||||
}
|
||||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||||
return
|
||||
}
|
||||
defer inflightDone()
|
||||
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body)
|
||||
requestCtx := service.WithOpenAIImagesEndpoint(service.WithOpenAIImageGenerationIntent(c.Request.Context()))
|
||||
|
||||
|
||||
@@ -149,6 +149,9 @@ func (h *PaymentHandler) GetCheckoutInfo(c *gin.Context) {
|
||||
BalanceRechargeMultiplier: cfg.BalanceRechargeMultiplier,
|
||||
SubscriptionUSDToCNYRate: cfg.SubscriptionUSDToCNYRate,
|
||||
RechargeFeeRate: cfg.RechargeFeeRate,
|
||||
RechargeBonusTiers: cfg.RechargeBonusTiers,
|
||||
RechargeBonusMode: cfg.RechargeBonusMode,
|
||||
RechargeBonusNotice: cfg.RechargeBonusNotice,
|
||||
HelpText: cfg.HelpText,
|
||||
HelpImageURL: cfg.HelpImageURL,
|
||||
StripePublishableKey: cfg.StripePublishableKey,
|
||||
@@ -166,6 +169,9 @@ type checkoutInfoResponse struct {
|
||||
BalanceRechargeMultiplier float64 `json:"balance_recharge_multiplier"`
|
||||
SubscriptionUSDToCNYRate float64 `json:"subscription_usd_to_cny_rate"`
|
||||
RechargeFeeRate float64 `json:"recharge_fee_rate"`
|
||||
RechargeBonusTiers []service.RechargeBonusTier `json:"recharge_bonus_tiers"`
|
||||
RechargeBonusMode string `json:"recharge_bonus_mode"`
|
||||
RechargeBonusNotice string `json:"recharge_bonus_notice"`
|
||||
HelpText string `json:"help_text"`
|
||||
HelpImageURL string `json:"help_image_url"`
|
||||
StripePublishableKey string `json:"stripe_publishable_key"`
|
||||
@@ -482,6 +488,7 @@ type PublicOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
PaymentType string `json:"payment_type"`
|
||||
OrderType string `json:"order_type"`
|
||||
@@ -517,6 +524,7 @@ func buildPublicOrderResult(order *dbent.PaymentOrder) PublicOrderResult {
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
PaymentType: order.PaymentType,
|
||||
OrderType: order.OrderType,
|
||||
@@ -625,6 +633,7 @@ type PaymentOrderResult struct {
|
||||
Amount float64 `json:"amount"`
|
||||
PayAmount float64 `json:"pay_amount"`
|
||||
FeeRate float64 `json:"fee_rate"`
|
||||
BonusAmount float64 `json:"bonus_amount"`
|
||||
Currency string `json:"currency"`
|
||||
PaymentType string `json:"payment_type"`
|
||||
OutTradeNo string `json:"out_trade_no"`
|
||||
@@ -663,6 +672,7 @@ func sanitizePaymentOrderForResponse(order *dbent.PaymentOrder) *PaymentOrderRes
|
||||
Amount: order.Amount,
|
||||
PayAmount: order.PayAmount,
|
||||
FeeRate: order.FeeRate,
|
||||
BonusAmount: order.BonusAmount,
|
||||
Currency: service.PaymentOrderCurrency(order),
|
||||
PaymentType: order.PaymentType,
|
||||
OutTradeNo: order.OutTradeNo,
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
@@ -811,3 +812,19 @@ func TestUserHandlerStartIdentityBindingReturnsAuthorizeURL(t *testing.T) {
|
||||
require.Contains(t, resp.Data.AuthorizeURL, "intent=bind_current_user")
|
||||
require.Contains(t, resp.Data.AuthorizeURL, "redirect=%2Fsettings%2Fprofile")
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) IncrVerificationCodeAttempts(context.Context, string) (int, error) {
|
||||
if s.data == nil {
|
||||
return 0, errors.New("verification code not found")
|
||||
}
|
||||
s.data.Attempts++
|
||||
return s.data.Attempts, nil
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) IncrNotifyVerifyCodeAttempts(context.Context, string) (int, error) {
|
||||
return 0, errors.New("notify verification code not found")
|
||||
}
|
||||
|
||||
func (s *userHandlerEmailCacheStub) ConsumePasswordResetToken(context.Context, string, string) (bool, error) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
@@ -264,10 +264,11 @@ func loginTestUser(t *testing.T) string {
|
||||
t.Helper()
|
||||
|
||||
// 先尝试用管理员账户登录
|
||||
adminEmail := getEnv("ADMIN_EMAIL", "admin@sub2api.local")
|
||||
// 管理员邮箱不再有固定默认值(未设置时安装阶段随机生成),需同时提供邮箱与密码。
|
||||
adminEmail := getEnv("ADMIN_EMAIL", "")
|
||||
adminPassword := getEnv("ADMIN_PASSWORD", "")
|
||||
|
||||
if adminPassword == "" {
|
||||
if adminEmail == "" || adminPassword == "" {
|
||||
// 尝试用测试用户
|
||||
adminEmail = testUserEmail
|
||||
adminPassword = testUserPassword
|
||||
|
||||
@@ -46,6 +46,7 @@ const (
|
||||
PlatformDeepseek = domain.PlatformDeepseek
|
||||
PlatformMiniMax = domain.PlatformMiniMax
|
||||
PlatformOpenCodeGo = domain.PlatformOpenCodeGo
|
||||
PlatformTypeSafe = domain.PlatformTypeSafe
|
||||
)
|
||||
|
||||
// AllPlatforms 返回所有支持的平台列表
|
||||
@@ -61,6 +62,7 @@ func AllPlatforms() []string {
|
||||
PlatformDeepseek,
|
||||
PlatformMiniMax,
|
||||
PlatformOpenCodeGo,
|
||||
PlatformTypeSafe,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -18,5 +18,6 @@ func TestAllPlatformsIncludesEveryConcretePlatform(t *testing.T) {
|
||||
"deepseek",
|
||||
"minimax",
|
||||
"opencode_go",
|
||||
"typesafe",
|
||||
}, AllPlatforms())
|
||||
}
|
||||
|
||||
@@ -376,6 +376,9 @@ func (e *EasyPay) VerifyNotification(_ context.Context, rawBody string, _ map[st
|
||||
// url.ParseQuery already decodes values — no additional decode needed.
|
||||
params := make(map[string]string)
|
||||
for k := range values {
|
||||
if !easyPayNotifyAllowedParams[k] {
|
||||
return nil, fmt.Errorf("unexpected notify param: %s", k)
|
||||
}
|
||||
params[k] = values.Get(k)
|
||||
}
|
||||
sign := params["sign"]
|
||||
@@ -601,3 +604,26 @@ func easyPaySign(params map[string]string, pkey string) string {
|
||||
func easyPayVerifySign(params map[string]string, pkey string, sign string) bool {
|
||||
return hmac.Equal([]byte(easyPaySign(params, pkey)), []byte(sign))
|
||||
}
|
||||
|
||||
// easyPayNotifyAllowedParams is the exact parameter set a genuine EasyPay
|
||||
// (彩虹易支付-compatible) async notification carries. Order-creation-only
|
||||
// fields (notify_url, return_url, cid, device, clientip, ...) must never
|
||||
// appear in a callback: because the sign base string concatenates values
|
||||
// unescaped, a signed order URL whose return_url embeds e.g.
|
||||
// "trade_status=TRADE_SUCCESS" could otherwise be replayed as a forged
|
||||
// payment-success notification (issue #7881). Rejecting unknown keys closes
|
||||
// the whole smuggling class; genuinely paid orders rejected by an exotic
|
||||
// upstream variant are still recovered by the upstream QueryOrder reconcile
|
||||
// path.
|
||||
var easyPayNotifyAllowedParams = map[string]bool{
|
||||
"pid": true,
|
||||
"trade_no": true,
|
||||
"out_trade_no": true,
|
||||
"type": true,
|
||||
"name": true,
|
||||
"money": true,
|
||||
"trade_status": true,
|
||||
"param": true,
|
||||
"sign": true,
|
||||
"sign_type": true,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
package provider
|
||||
|
||||
// Regression tests for the EasyPay forged-callback vulnerability
|
||||
// (Wei-Shaw/sub2api issue #7875).
|
||||
//
|
||||
// Attack recap: the order-creation signature (visible to the payer in the
|
||||
// submit.php popup URL) signs return_url among other fields, and the sign
|
||||
// base string concatenates values unescaped. A client-supplied return_url
|
||||
// ending in "&trade_status=TRADE_SUCCESS" therefore produced a signature
|
||||
// that was byte-identical to a payment-success notification's signature
|
||||
// once the callback decoded return_url partially and promoted the smuggled
|
||||
// pair to a top-level parameter. Merchant-key rotation could not help
|
||||
// because the attacker never needed the key.
|
||||
//
|
||||
// Fixes covered here:
|
||||
// 1. service.CanonicalizeReturnURL drops client-supplied query parameters
|
||||
// (see payment_resume_service_test.go).
|
||||
// 2. VerifyNotification rejects any parameter outside the genuine async
|
||||
// notify set, so smuggled order-creation fields (return_url et al.)
|
||||
// fail closed even if a signed value smuggles a fake pair.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/url"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
)
|
||||
|
||||
// easyPayPoCProvider returns a provider with a fixed test credential set.
|
||||
func easyPayPoCProvider() *EasyPay {
|
||||
return &EasyPay{config: map[string]string{
|
||||
"pid": "1000",
|
||||
"pkey": "MERCHANT_SECRET_KEY",
|
||||
"apiBase": "https://pay.example.com",
|
||||
"notifyUrl": "https://site.example.com/api/v1/payment/webhook/easypay",
|
||||
"returnUrl": "https://site.example.com/payment/result",
|
||||
}}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsForgedSignReuseCallback is the exact PoC payload
|
||||
// that was accepted before the fix: the order's own submit.php signature
|
||||
// replayed with trade_status smuggled out of the return_url value.
|
||||
func TestEasyPayNotifyRejectsForgedSignReuseCallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
// Order creation (popup mode): return_url carries the smuggled pair as
|
||||
// its last sorted inner query key, exactly as buildPaymentReturnURL
|
||||
// would have produced before CanonicalizeReturnURL stripped user query.
|
||||
returnURL := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success&trade_status=TRADE_SUCCESS"
|
||||
createParams := map[string]string{
|
||||
"pid": "1000",
|
||||
"type": "alipay",
|
||||
"out_trade_no": "ORDER123",
|
||||
"notify_url": e.config["notifyUrl"],
|
||||
"return_url": returnURL,
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
}
|
||||
sign := easyPaySign(createParams, e.config["pkey"])
|
||||
|
||||
// Forged callback: return_url encoded only up to status=success, then a
|
||||
// raw &trade_status=TRADE_SUCCESS promotes it to a top-level parameter.
|
||||
prefix := "https://site.example.com/payment/result?order_id=99&out_trade_no=ORDER123&status=success"
|
||||
cb := url.Values{}
|
||||
cb.Set("pid", "1000")
|
||||
cb.Set("type", "alipay")
|
||||
cb.Set("out_trade_no", "ORDER123")
|
||||
cb.Set("notify_url", e.config["notifyUrl"])
|
||||
cb.Set("name", "balance recharge")
|
||||
cb.Set("money", "650.00")
|
||||
cb.Set("return_url", prefix)
|
||||
rawCallback := cb.Encode() + "&trade_status=TRADE_SUCCESS" + "&sign=" + sign + "&sign_type=MD5"
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), rawCallback, nil); err == nil {
|
||||
t.Fatal("forged sign-reuse callback must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsOrderURLReplay covers the milder variant where the
|
||||
// attacker replays the complete signed pay URL unchanged: return_url itself
|
||||
// is not a legitimate notify parameter and must also be rejected.
|
||||
func TestEasyPayNotifyRejectsOrderURLReplay(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
createParams := map[string]string{
|
||||
"pid": "1000",
|
||||
"type": "alipay",
|
||||
"out_trade_no": "ORDER123",
|
||||
"notify_url": e.config["notifyUrl"],
|
||||
"return_url": "https://site.example.com/payment/result",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
}
|
||||
createParams["sign"] = easyPaySign(createParams, e.config["pkey"])
|
||||
createParams["sign_type"] = signTypeMD5
|
||||
q := url.Values{}
|
||||
for k, v := range createParams {
|
||||
q.Set(k, v)
|
||||
}
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
|
||||
t.Fatal("replayed order URL must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyAcceptsGenuineCallback locks in the legitimate notify
|
||||
// contract: the canonical parameter set with a valid signature succeeds.
|
||||
func TestEasyPayNotifyAcceptsGenuineCallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
params := map[string]string{
|
||||
"pid": "1000",
|
||||
"trade_no": "2026100622001400000001",
|
||||
"out_trade_no": "ORDER123",
|
||||
"type": "alipay",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
"trade_status": tradeStatusSuccess,
|
||||
}
|
||||
sign := easyPaySign(params, e.config["pkey"])
|
||||
params["sign"] = sign
|
||||
params["sign_type"] = signTypeMD5
|
||||
q := url.Values{}
|
||||
for k, v := range params {
|
||||
q.Set(k, v)
|
||||
}
|
||||
|
||||
n, err := e.VerifyNotification(context.Background(), q.Encode(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("genuine callback rejected: %v", err)
|
||||
}
|
||||
if n.Status != payment.ProviderStatusSuccess {
|
||||
t.Fatalf("status = %v, want success", n.Status)
|
||||
}
|
||||
if n.OrderID != "ORDER123" || n.TradeNo != "2026100622001400000001" || n.Amount != 650.00 {
|
||||
t.Fatalf("unexpected notification: %+v", n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEasyPayNotifyRejectsUnknownParam ensures any parameter outside the
|
||||
// canonical notify set fails closed, including empty-valued ones that the
|
||||
// signer itself would skip.
|
||||
func TestEasyPayNotifyRejectsUnknownParam(t *testing.T) {
|
||||
t.Parallel()
|
||||
e := easyPayPoCProvider()
|
||||
|
||||
params := map[string]string{
|
||||
"pid": "1000",
|
||||
"trade_no": "T1",
|
||||
"out_trade_no": "ORDER123",
|
||||
"type": "alipay",
|
||||
"name": "balance recharge",
|
||||
"money": "650.00",
|
||||
"trade_status": tradeStatusSuccess,
|
||||
}
|
||||
sign := easyPaySign(params, e.config["pkey"])
|
||||
q := url.Values{}
|
||||
for k, v := range params {
|
||||
q.Set(k, v)
|
||||
}
|
||||
q.Set("sign", sign)
|
||||
q.Set("sign_type", signTypeMD5)
|
||||
q.Set("device", "") // empty value: invisible to the signer, still rejected
|
||||
|
||||
if _, err := e.VerifyNotification(context.Background(), q.Encode(), nil); err == nil {
|
||||
t.Fatal("callback with unknown param must be rejected")
|
||||
}
|
||||
}
|
||||
@@ -1979,3 +1979,16 @@ func TestMessageStartSSE_StopReasonIsJSONNull(t *testing.T) {
|
||||
require.Contains(t, sse, `"stop_reason":null`)
|
||||
require.NotContains(t, sse, `"stop_reason":""`)
|
||||
}
|
||||
|
||||
func TestGPT61SolCacheOptionsAndBreakpointsSurviveChatBridge(t *testing.T) {
|
||||
sampling := 0.7
|
||||
for _, effort := range []string{"low", "medium", "high", "xhigh", "max"} {
|
||||
out, err := ChatCompletionsToResponses(&ChatCompletionsRequest{Model: "gpt-6.1-sol", ReasoningEffort: effort, Temperature: &sampling, TopP: &sampling, PromptCacheOptions: json.RawMessage(`{"ttl":"30m","mode":"explicit"}`), Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`[{"type":"text","text":"prefix","prompt_cache_breakpoint":{"mode":"explicit"}}]`)}}})
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, out.Temperature)
|
||||
require.Nil(t, out.TopP)
|
||||
require.Equal(t, effort, out.Reasoning.Effort)
|
||||
require.JSONEq(t, `{"ttl":"30m","mode":"explicit"}`, string(out.PromptCacheOptions))
|
||||
require.Contains(t, string(out.Input), "prompt_cache_breakpoint")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,9 @@ import (
|
||||
// Chat Completions intermediary round-trip (e.g. thinking, cache_control,
|
||||
// structured system prompts).
|
||||
func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
if err := openai.ValidateGPT61SolReasoningEffort(req.Model, anthropicReasoningEffort(req)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
input, err := convertAnthropicToResponsesInput(req.System, req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -59,6 +62,9 @@ func AnthropicToResponses(req *AnthropicRequest) (*ResponsesRequest, error) {
|
||||
|
||||
// An explicit thinking disable takes precedence over output_config.effort.
|
||||
effort := anthropicReasoningEffort(req)
|
||||
if openai.IsGPT61SolModelSpelling(req.Model) && req.OutputConfig != nil && req.OutputConfig.Effort == "max" && effort != "none" {
|
||||
effort = "max"
|
||||
}
|
||||
out.Reasoning = &ResponsesReasoning{
|
||||
Effort: effort,
|
||||
}
|
||||
|
||||
@@ -18,6 +18,9 @@ type chatMessageContent struct {
|
||||
// true. store is always false and reasoning.encrypted_content is always
|
||||
// included so that the response translator has full context.
|
||||
func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest, error) {
|
||||
if err := openai.ValidateGPT61SolReasoningEffort(req.Model, req.ReasoningEffort); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
input, err := convertChatMessagesToResponsesInput(req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
File diff suppressed because one or more lines are too long
@@ -3,6 +3,8 @@ package openai
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
@@ -23,6 +25,7 @@ var DefaultModels = []Model{
|
||||
{ID: "gpt-5.6", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 (Sol)"},
|
||||
{ID: "gpt-5.6-terra", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Terra"},
|
||||
{ID: "gpt-5.6-luna", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Luna"},
|
||||
{ID: "gpt-6.1-sol", Object: "model", Created: 1790640000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6.1 Sol"},
|
||||
{ID: "gpt-6-sol", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Sol"},
|
||||
{ID: "gpt-6-luna", Object: "model", Created: 1790035200, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Luna"},
|
||||
{ID: "gpt-6-astra", Object: "model", Created: 1788480000, OwnedBy: "openai", Type: "model", DisplayName: "GPT-6 Astra"},
|
||||
@@ -78,6 +81,12 @@ var instructionsGPT55 string
|
||||
//go:embed instructions_gpt6_astra.txt
|
||||
var instructionsGPT6Astra string
|
||||
|
||||
// CodexGPT61SolMetadata is the complete official descriptor from openai/codex
|
||||
// b1e72963c3b71a9265a551e54beff078384efed9, codex-rs/models-manager/models.json.
|
||||
//
|
||||
//go:embed codex_gpt61_sol.json
|
||||
var CodexGPT61SolMetadata []byte
|
||||
|
||||
// latestCodexInstructions 返回当前已知最新版本的 Codex base instructions,
|
||||
// 当前为 GPT-5.5;若 5.5 prompt 意外为空则回退到 DefaultInstructions 保证非空。
|
||||
func latestCodexInstructions() string {
|
||||
@@ -140,6 +149,16 @@ func CanonicalizeOpenAIModelAliasSpelling(model string) string {
|
||||
func CodexBaseInstructionsForModel(model string) string {
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
switch {
|
||||
case IsGPT61SolModelSpelling(canonical):
|
||||
var metadata struct {
|
||||
ModelMessages struct {
|
||||
InstructionsTemplate string `json:"instructions_template"`
|
||||
} `json:"model_messages"`
|
||||
}
|
||||
if err := json.Unmarshal(CodexGPT61SolMetadata, &metadata); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return metadata.ModelMessages.InstructionsTemplate
|
||||
case canonical == "gpt-6" || canonical == "gpt-6-astra" || strings.HasPrefix(canonical, "gpt-6-astra-"):
|
||||
if v := strings.TrimSpace(instructionsGPT6Astra); v != "" {
|
||||
return instructionsGPT6Astra
|
||||
@@ -177,3 +196,34 @@ func IsGPT6SolOrLunaModelSpelling(model string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsGPT61SolModelSpelling recognizes the published model and local effort/compact
|
||||
// spellings. Invalid effort suffixes remain identifiable for request validation.
|
||||
func IsGPT61SolModelSpelling(model string) bool {
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
if canonical == "gpt-6.1-sol" {
|
||||
return true
|
||||
}
|
||||
suffix, ok := strings.CutPrefix(canonical, "gpt-6.1-sol-")
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
switch suffix {
|
||||
case "none", "minimal", "low", "medium", "high", "xhigh", "max", "openai-compact":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ValidateGPT61SolReasoningEffort rejects disabled reasoning instead of silently
|
||||
// increasing the client's requested effort on compatibility paths.
|
||||
func ValidateGPT61SolReasoningEffort(model, effort string) error {
|
||||
if IsGPT61SolModelSpelling(model) {
|
||||
switch strings.ToLower(strings.TrimSpace(effort)) {
|
||||
case "none", "minimal":
|
||||
return fmt.Errorf("gpt-6.1-sol does not support reasoning effort %q; use low, medium, high, xhigh or max", effort)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -42,3 +42,22 @@ func TestGPT6SolLunaModelIdentity(t *testing.T) {
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-solitude"))
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling("gpt-6-luna-preview"))
|
||||
}
|
||||
|
||||
func TestGPT61SolIdentityAndEffort(t *testing.T) {
|
||||
require.Contains(t, DefaultModelIDs(), "gpt-6.1-sol")
|
||||
for _, id := range []string{"gpt-6.1-sol", "openai/gpt-6.1-sol-max", "GPT_6.1_SOL", "gpt-6.1-sol-openai-compact"} {
|
||||
require.True(t, IsGPT61SolModelSpelling(id), id)
|
||||
require.False(t, IsGPT6SolOrLunaModelSpelling(id), id)
|
||||
}
|
||||
for _, id := range []string{"gpt-6.1", "gpt-6.1-solitude", "gpt-6.1-sol-preview", "gpt-6-sol"} {
|
||||
require.False(t, IsGPT61SolModelSpelling(id), id)
|
||||
}
|
||||
for _, effort := range []string{"", "low", "medium", "high", "xhigh", "max"} {
|
||||
require.NoError(t, ValidateGPT61SolReasoningEffort("gpt-6.1-sol", effort))
|
||||
}
|
||||
for _, effort := range []string{"none", "minimal"} {
|
||||
require.Error(t, ValidateGPT61SolReasoningEffort("gpt-6.1-sol", effort))
|
||||
require.NoError(t, ValidateGPT61SolReasoningEffort("gpt-6-sol", effort))
|
||||
}
|
||||
require.Contains(t, CodexBaseInstructionsForModel("gpt-6.1-sol"), "based on GPT-6")
|
||||
}
|
||||
|
||||
@@ -13,6 +13,12 @@ import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultBaseURL = "https://api.typesafe.ai"
|
||||
SystemOnePath = "/v1/systemone"
|
||||
JevLatestModel = "jev-latest"
|
||||
)
|
||||
|
||||
type Question struct {
|
||||
Type string `json:"type"`
|
||||
Instructions string `json:"instructions"`
|
||||
@@ -35,22 +41,103 @@ type Usage struct {
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
|
||||
type SystemOneResponse struct {
|
||||
Body []byte
|
||||
Model string
|
||||
Usage Usage
|
||||
}
|
||||
|
||||
func NewSystemOneRequest(ctx context.Context, baseURL, key string, body []byte) (*http.Request, error) {
|
||||
endpoint, err := url.JoinPath(strings.TrimRight(baseURL, "/"), SystemOnePath)
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid endpoint")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid request")
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
return req, nil
|
||||
}
|
||||
|
||||
// MaxSystemOneResponseBytes bounds a buffered System One response body.
|
||||
const MaxSystemOneResponseBytes = 4 << 20
|
||||
|
||||
var ErrSystemOneResponseTooLarge = errors.New("typesafe response exceeds size limit")
|
||||
|
||||
func DecodeSystemOneResponse(r io.Reader) (*SystemOneResponse, error) {
|
||||
body, err := io.ReadAll(io.LimitReader(r, MaxSystemOneResponseBytes+1))
|
||||
if err != nil {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
// A truncated body would otherwise surface as a misleading "invalid JSON".
|
||||
if len(body) > MaxSystemOneResponseBytes {
|
||||
return nil, ErrSystemOneResponseTooLarge
|
||||
}
|
||||
if !json.Valid(body) {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
var envelope map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &envelope); err != nil || envelope == nil {
|
||||
return nil, errors.New("typesafe invalid response")
|
||||
}
|
||||
// The upstream already answered (and charged); an unexpected model or usage
|
||||
// shape must not discard the answer, so both are decoded leniently.
|
||||
var model string
|
||||
_ = json.Unmarshal(envelope["model"], &model)
|
||||
var usage map[string]json.RawMessage
|
||||
_ = json.Unmarshal(envelope["usage"], &usage)
|
||||
return &SystemOneResponse{
|
||||
Body: body,
|
||||
Model: model,
|
||||
Usage: Usage{
|
||||
InputTokens: systemOneTokenCount(usage["input_tokens"]),
|
||||
OutputTokens: systemOneTokenCount(usage["output_tokens"]),
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// maxSystemOneTokenCount bounds a reported token count before int conversion.
|
||||
const maxSystemOneTokenCount = 1 << 40
|
||||
|
||||
// systemOneTokenCount accepts integer, float, or numeric-string token counts.
|
||||
func systemOneTokenCount(raw json.RawMessage) int {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 {
|
||||
return 0
|
||||
}
|
||||
if raw[0] == '"' {
|
||||
var text string
|
||||
if json.Unmarshal(raw, &text) != nil {
|
||||
return 0
|
||||
}
|
||||
raw = json.RawMessage(strings.TrimSpace(text))
|
||||
}
|
||||
var number json.Number
|
||||
if json.Unmarshal(raw, &number) != nil {
|
||||
return 0
|
||||
}
|
||||
value, err := number.Float64()
|
||||
if err != nil || math.IsNaN(value) || value <= 0 {
|
||||
return 0
|
||||
}
|
||||
if value > maxSystemOneTokenCount {
|
||||
value = maxSystemOneTokenCount
|
||||
}
|
||||
return int(math.Round(value))
|
||||
}
|
||||
|
||||
// Evaluate performs one attempt. The caller owns timeouts, retries and key rotation.
|
||||
func Evaluate(ctx context.Context, client *http.Client, baseURL, key string, input Request) (*Result, int, error) {
|
||||
endpoint, err := url.JoinPath(strings.TrimRight(baseURL, "/"), "/v1/systemone")
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid endpoint")
|
||||
}
|
||||
body, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(body))
|
||||
req, err := NewSystemOneRequest(ctx, baseURL, key, body)
|
||||
if err != nil {
|
||||
return nil, 0, errors.New("typesafe invalid request")
|
||||
return nil, 0, err
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var ErrStreamingUnsupported = errors.New("typesafe system one does not support streaming")
|
||||
|
||||
type systemOneEnvelope struct {
|
||||
Model json.RawMessage `json:"model"`
|
||||
State json.RawMessage `json:"state"`
|
||||
Questions json.RawMessage `json:"questions"`
|
||||
Stream json.RawMessage `json:"stream"`
|
||||
}
|
||||
|
||||
var (
|
||||
systemOneRequestFields = []string{"model", "state", "questions", "stream"}
|
||||
systemOneQuestionFields = []string{"type", "instructions", "criteria"}
|
||||
)
|
||||
|
||||
func ValidateSystemOneRequest(body []byte) (string, error) {
|
||||
var envelope systemOneEnvelope
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
return "", errors.New("invalid JSON request")
|
||||
}
|
||||
// encoding/json matches struct fields case-insensitively and keeps the last
|
||||
// duplicate, while the raw body is forwarded upstream unchanged. Reject any
|
||||
// spelling the upstream could read differently from this validator.
|
||||
if err := checkSystemOneObjectKeys(body, "request", systemOneRequestFields); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
model, err := requiredString(envelope.Model, "model")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if model != JevLatestModel {
|
||||
return "", fmt.Errorf("model must be %s", JevLatestModel)
|
||||
}
|
||||
if err := validateStringObjectOrArray(envelope.State, "state"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
questionsRaw := bytes.TrimSpace(envelope.Questions)
|
||||
var questions map[string]json.RawMessage
|
||||
if len(questionsRaw) == 0 || questionsRaw[0] != '{' || json.Unmarshal(questionsRaw, &questions) != nil || len(questions) == 0 {
|
||||
return "", errors.New("questions must be a non-empty object")
|
||||
}
|
||||
if err := checkSystemOneObjectKeys(questionsRaw, "questions", nil); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(envelope.Stream) > 0 && string(envelope.Stream) != "null" {
|
||||
var stream bool
|
||||
if err := json.Unmarshal(envelope.Stream, &stream); err != nil {
|
||||
return "", errors.New("stream must be a boolean")
|
||||
}
|
||||
if stream {
|
||||
return "", ErrStreamingUnsupported
|
||||
}
|
||||
}
|
||||
for id, raw := range questions {
|
||||
if err := validateQuestion(id, raw); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return model, nil
|
||||
}
|
||||
|
||||
func validateQuestion(id string, raw json.RawMessage) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
var question struct {
|
||||
Type json.RawMessage `json:"type"`
|
||||
Instructions json.RawMessage `json:"instructions"`
|
||||
Criteria json.RawMessage `json:"criteria"`
|
||||
}
|
||||
if len(raw) == 0 || raw[0] != '{' || json.Unmarshal(raw, &question) != nil {
|
||||
return fmt.Errorf("question %q must be an object", id)
|
||||
}
|
||||
if err := checkSystemOneObjectKeys(raw, fmt.Sprintf("question %q", id), systemOneQuestionFields); err != nil {
|
||||
return err
|
||||
}
|
||||
typ, err := requiredString(question.Type, "question type")
|
||||
if err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
// Mirror the wire schema generated from https://api.typesafe.ai/openapi.json
|
||||
// in TypeSafe SDK v0.5.7: instructions are optional and nullable for all types.
|
||||
if err := validateOptionalDescription(question.Instructions, "instructions"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
|
||||
switch typ {
|
||||
case "noul":
|
||||
criteria := bytes.TrimSpace(question.Criteria)
|
||||
if len(criteria) == 0 || bytes.Equal(criteria, []byte("null")) {
|
||||
return nil
|
||||
}
|
||||
var descriptions map[string]json.RawMessage
|
||||
if criteria[0] != '{' || json.Unmarshal(criteria, &descriptions) != nil {
|
||||
return fmt.Errorf("question %q: noul criteria must be an object", id)
|
||||
}
|
||||
for _, outcome := range []string{"true", "false"} {
|
||||
if err := validateOptionalDescription(descriptions[outcome], "noul criteria "+outcome); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
case "choice":
|
||||
var criteria map[string]json.RawMessage
|
||||
if json.Unmarshal(question.Criteria, &criteria) != nil || criteria == nil {
|
||||
return fmt.Errorf("question %q: choice criteria must be an object", id)
|
||||
}
|
||||
for _, value := range criteria {
|
||||
if err := validateOptionalDescription(value, "choice criteria value"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
case "score":
|
||||
var criteria []json.RawMessage
|
||||
if json.Unmarshal(question.Criteria, &criteria) != nil || len(criteria) == 0 {
|
||||
return fmt.Errorf("question %q: score criteria must contain at least one level", id)
|
||||
}
|
||||
for _, value := range criteria {
|
||||
if err := validateStringObjectOrArray(value, "score criteria value"); err != nil {
|
||||
return fmt.Errorf("question %q: %w", id, err)
|
||||
}
|
||||
}
|
||||
default:
|
||||
return fmt.Errorf("question %q: unsupported type %q", id, typ)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateOptionalDescription(raw json.RawMessage, name string) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 || bytes.Equal(raw, []byte("null")) {
|
||||
return nil
|
||||
}
|
||||
return validateStringObjectOrArray(raw, name)
|
||||
}
|
||||
|
||||
func requiredString(raw json.RawMessage, name string) (string, error) {
|
||||
var value string
|
||||
if len(raw) == 0 || json.Unmarshal(raw, &value) != nil || strings.TrimSpace(value) == "" {
|
||||
return "", fmt.Errorf("%s must be a non-empty string", name)
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func validateStringObjectOrArray(raw json.RawMessage, name string) error {
|
||||
raw = bytes.TrimSpace(raw)
|
||||
if len(raw) == 0 || string(raw) == "null" {
|
||||
return fmt.Errorf("%s is required", name)
|
||||
}
|
||||
switch raw[0] {
|
||||
case '{', '[':
|
||||
return nil
|
||||
case '"':
|
||||
if rawString(raw) {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("%s must be a string, object, or array", name)
|
||||
}
|
||||
|
||||
func rawString(raw json.RawMessage) bool {
|
||||
var value string
|
||||
return json.Unmarshal(raw, &value) == nil
|
||||
}
|
||||
|
||||
// checkSystemOneObjectKeys rejects duplicate keys and non-canonical spellings
|
||||
// (any case variant) of the known fields in one JSON object level.
|
||||
func checkSystemOneObjectKeys(raw []byte, scope string, canonical []string) error {
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
if token, err := decoder.Token(); err != nil || token != json.Delim('{') {
|
||||
return fmt.Errorf("%s must be an object", scope)
|
||||
}
|
||||
seen := make(map[string]struct{})
|
||||
for decoder.More() {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return errors.New("invalid JSON request")
|
||||
}
|
||||
key, _ := token.(string)
|
||||
if _, duplicate := seen[key]; duplicate {
|
||||
return fmt.Errorf("%s contains duplicate field %q", scope, key)
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
for _, field := range canonical {
|
||||
if key != field && strings.EqualFold(key, field) {
|
||||
return fmt.Errorf("%s field %q must be written as %q", scope, key, field)
|
||||
}
|
||||
}
|
||||
var value json.RawMessage
|
||||
if err := decoder.Decode(&value); err != nil {
|
||||
return errors.New("invalid JSON request")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package typesafe
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestValidateSystemOneRequestValidQuestionTypes(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{"noul string state", `{"model":"jev-latest","state":"sample","questions":{"safety":{"type":"noul","instructions":"Evaluate safety","criteria":{"safe":"No harm"}}}}`},
|
||||
{"choice object state", `{"model":"jev-latest","state":{"text":"sample"},"questions":{"label":{"type":"choice","instructions":{"task":"Classify"},"criteria":{"safe":"Allowed","unsafe":null}}},"stream":false}`},
|
||||
{"score array state", `{"model":"jev-latest","state":["sample"],"questions":{"quality":{"type":"score","instructions":["Rate quality"],"criteria":["poor","good"]}}}`},
|
||||
{"noul omitted instructions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul"}}}`},
|
||||
{"noul nullable instructions and criteria", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","instructions":null,"criteria":null}}}`},
|
||||
{"noul structured descriptions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","criteria":{"true":{"examples":["yes",null]},"false":["no",null],"extension":42}}}}`},
|
||||
{"choice structured descriptions", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"choice","instructions":null,"criteria":{"a":{"description":"A","extra":null},"b":["B",null],"c":null}}}}`},
|
||||
{"choice empty criteria", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"choice","criteria":{}}}}`},
|
||||
{"score one level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","criteria":["only"]}}}`},
|
||||
{"score object level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","instructions":null,"criteria":[{"description":"only","extra":null}]}}}`},
|
||||
{"score array level", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"score","criteria":[["only",null]]}}}`},
|
||||
{"native extensions preserved", `{"model":"jev-latest","state":"sample","questions":{"q":{"type":"noul","extension":{"kept":true}}},"extension":[1,null]}`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
model, err := ValidateSystemOneRequest([]byte(tc.body))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, JevLatestModel, model)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateSystemOneRequestRejectsInvalidRequests(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
want string
|
||||
}{
|
||||
{"invalid json", `{`, "invalid JSON"},
|
||||
{"missing model", `{"state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "model"},
|
||||
{"illegal model", `{"model":"jev-old","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "jev-latest"},
|
||||
{"model with whitespace", `{"model":" jev-latest ","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`, "jev-latest"},
|
||||
{"missing state", `{"model":"jev-latest","questions":{"q":{"type":"noul","instructions":"x"}}}`, "state"},
|
||||
{"scalar state", `{"model":"jev-latest","state":42,"questions":{"q":{"type":"noul","instructions":"x"}}}`, "state"},
|
||||
{"empty questions", `{"model":"jev-latest","state":"x","questions":{}}`, "questions"},
|
||||
{"array questions", `{"model":"jev-latest","state":"x","questions":[{"type":"noul"}]}`, "questions must be a non-empty object"},
|
||||
{"null questions", `{"model":"jev-latest","state":"x","questions":null}`, "questions must be a non-empty object"},
|
||||
{"unknown question type", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"boolean","instructions":"x"}}}`, "unsupported type"},
|
||||
{"question type with whitespace", `{"model":"jev-latest","state":"x","questions":{"q":{"type":" noul ","instructions":"x"}}}`, "unsupported type"},
|
||||
{"noul criteria array", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"x","criteria":[]}}}`, "noul criteria"},
|
||||
{"noul numeric description", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","criteria":{"true":1}}}}`, "noul criteria"},
|
||||
{"noul boolean description", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","criteria":{"false":false}}}}`, "noul criteria"},
|
||||
{"missing choice criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice"}}}`, "choice criteria"},
|
||||
{"null choice criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","criteria":null}}}`, "choice criteria"},
|
||||
{"choice criteria array", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","criteria":[]}}}`, "choice criteria"},
|
||||
{"choice numeric value", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","instructions":"x","criteria":{"one":1}}}}`, "choice criteria"},
|
||||
{"missing score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score"}}}`, "score criteria"},
|
||||
{"empty score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":[]}}}`, "score criteria"},
|
||||
{"null score criteria", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":null}}}`, "score criteria"},
|
||||
{"score map is not wire protocol", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":{"0":"low"}}}}`, "score criteria"},
|
||||
{"numeric score level", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":["low",2]}}}`, "score criteria"},
|
||||
{"null score level", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"score","criteria":[null]}}}`, "score criteria"},
|
||||
{"numeric instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":1}}}`, "instructions"},
|
||||
{"boolean instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"choice","instructions":true,"criteria":{}}}}`, "instructions"},
|
||||
{"stream true", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"x"}},"stream":true}`, "streaming"},
|
||||
{"case variant model smuggles upstream model", `{"model":"jev-pro","MODEL":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `must be written as "model"`},
|
||||
{"case variant stream hides streaming", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}},"stream":true,"Stream":false}`, `must be written as "stream"`},
|
||||
{"unicode fold variant state", `{"model":"jev-latest","state":"x","ſtate":"y","questions":{"q":{"type":"noul"}}}`, `must be written as "state"`},
|
||||
{"duplicate model", `{"model":"jev-pro","model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `duplicate field "model"`},
|
||||
{"escaped duplicate model", `{"\u006dodel":"jev-pro","model":"jev-latest","state":"x","questions":{"q":{"type":"noul"}}}`, `duplicate field "model"`},
|
||||
{"duplicate state", `{"model":"jev-latest","state":"benign","state":"payload","questions":{"q":{"type":"noul"}}}`, `duplicate field "state"`},
|
||||
{"duplicate question id", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul"},"q":{"type":"noul"}}}`, `duplicate field "q"`},
|
||||
{"duplicate question instructions", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","instructions":"a","instructions":"b"}}}`, `duplicate field "instructions"`},
|
||||
{"case variant question type", `{"model":"jev-latest","state":"x","questions":{"q":{"type":"noul","Type":"choice"}}}`, `must be written as "type"`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := ValidateSystemOneRequest([]byte(tc.body))
|
||||
require.Error(t, err)
|
||||
require.Contains(t, err.Error(), tc.want)
|
||||
if tc.name == "stream true" {
|
||||
require.True(t, errors.Is(err, ErrStreamingUnsupported))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeSystemOneResponseToleratesUsageShapes(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
model string
|
||||
input, output int
|
||||
}{
|
||||
{"integers", `{"model":"jev-1","usage":{"input_tokens":12,"output_tokens":3}}`, "jev-1", 12, 3},
|
||||
{"floats", `{"model":"jev-1","usage":{"input_tokens":12.0,"output_tokens":2.6}}`, "jev-1", 12, 3},
|
||||
{"numeric strings", `{"usage":{"input_tokens":"15","output_tokens":" 4 "}}`, "", 15, 4},
|
||||
{"non numeric usage", `{"model":7,"usage":{"input_tokens":"many","output_tokens":null}}`, "", 0, 0},
|
||||
{"negative usage", `{"usage":{"input_tokens":-5}}`, "", 0, 0},
|
||||
{"usage not object", `{"model":"jev-1","usage":"none","answers":{}}`, "jev-1", 0, 0},
|
||||
{"missing usage", `{"answers":{}}`, "", 0, 0},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
decoded, err := DecodeSystemOneResponse(strings.NewReader(tc.body))
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, []byte(tc.body), decoded.Body)
|
||||
require.Equal(t, tc.model, decoded.Model)
|
||||
require.Equal(t, tc.input, decoded.Usage.InputTokens)
|
||||
require.Equal(t, tc.output, decoded.Usage.OutputTokens)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDecodeSystemOneResponseRejectsNonObjects(t *testing.T) {
|
||||
for _, body := range []string{`[]`, `"text"`, `null`, `{`, `data: {}`} {
|
||||
_, err := DecodeSystemOneResponse(strings.NewReader(body))
|
||||
require.Error(t, err, body)
|
||||
}
|
||||
_, err := DecodeSystemOneResponse(strings.NewReader(`{"pad":"` + strings.Repeat("a", MaxSystemOneResponseBytes) + `"}`))
|
||||
require.ErrorIs(t, err, ErrSystemOneResponseTooLarge)
|
||||
}
|
||||
@@ -19,10 +19,7 @@ const (
|
||||
// repository and service layers build their own client identity from it, so
|
||||
// one bump here covers OAuth traffic and billing probes together.
|
||||
// Keep in sync with https://x.ai/cli/stable.
|
||||
CLIClientVersion = "0.2.120"
|
||||
// billingCLIUserAgent is the legacy pager/shell UA used by billing probes.
|
||||
// Distinct from CLIUserAgent() in cli_identity.go (workspace-style UA).
|
||||
billingCLIUserAgent = "grok-pager/" + CLIClientVersion + " grok-shell/" + CLIClientVersion + " (macos; aarch64)"
|
||||
CLIClientVersion = "1.0.46"
|
||||
|
||||
BillingWeeklyPath = "/billing?format=credits"
|
||||
BillingMonthlyPath = "/billing"
|
||||
@@ -148,7 +145,8 @@ func ApplyCLIBillingHeaders(req *http.Request, accessToken string) {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set(CLITokenAuthHeader, CLITokenAuthValue)
|
||||
req.Header.Set(CLIClientVersionHeader, CLIClientVersion)
|
||||
req.Header.Set("User-Agent", billingCLIUserAgent)
|
||||
req.Header.Set("User-Agent", CLIUserAgent(CLIClientVersion))
|
||||
req.Header.Set("x-grok-client-mode", CLIClientMode)
|
||||
}
|
||||
|
||||
// ParseBillingPayload unmarshals a billing API response body.
|
||||
|
||||
@@ -38,7 +38,8 @@ func TestApplyCLIBillingHeaders(t *testing.T) {
|
||||
require.Equal(t, "Bearer token", req.Header.Get("Authorization"))
|
||||
require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader))
|
||||
require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader))
|
||||
require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent())
|
||||
require.Equal(t, CLIUserAgent(CLIClientVersion), req.UserAgent())
|
||||
require.Equal(t, "interactive", req.Header.Get("x-grok-client-mode"))
|
||||
}
|
||||
|
||||
func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) {
|
||||
|
||||
@@ -3,6 +3,7 @@ package xai
|
||||
import (
|
||||
"net/http"
|
||||
"os"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/mod/semver"
|
||||
@@ -16,7 +17,7 @@ const (
|
||||
CLIProxyHost = "cli-chat-proxy.grok.com"
|
||||
|
||||
// CLIStableVersion is the known-good minimum client version accepted by cli-chat-proxy.
|
||||
CLIStableVersion = "0.2.93"
|
||||
CLIStableVersion = "1.0.13"
|
||||
|
||||
// CLIVersionEnv is the optional operator override for CLIStableVersion.
|
||||
CLIVersionEnv = "XAI_GROK_CLI_VERSION"
|
||||
@@ -24,11 +25,11 @@ const (
|
||||
// CLITokenAuth is required by cli-chat-proxy for Grok Build OAuth tokens.
|
||||
CLITokenAuth = "xai-grok-cli"
|
||||
|
||||
// CLIClientIdentifier is the x-grok-client-identifier value used by Grok shell/CLI.
|
||||
CLIClientIdentifier = "grok-shell"
|
||||
// CLIClientIdentifier 对齐官方交互式 CLI 主请求的客户端标识。
|
||||
CLIClientIdentifier = "grok-pager"
|
||||
|
||||
// CLIClientMode is used by billing / quota probes on the CLI surface.
|
||||
CLIClientMode = "cli"
|
||||
// CLIClientMode 对齐官方 CLI 正常交互模式。
|
||||
CLIClientMode = "interactive"
|
||||
)
|
||||
|
||||
// ResolveCLIVersion returns a supported CLI client version.
|
||||
@@ -54,12 +55,24 @@ func IsSupportedCLIVersion(version string) bool {
|
||||
semver.Compare(canonical, minimum) >= 0
|
||||
}
|
||||
|
||||
// CLIUserAgent builds the workspace-style User-Agent for a CLI client version.
|
||||
// CLIUserAgent 对齐官方交互式 CLI 的 UA,平台名称使用 Rust 的格式。
|
||||
func CLIUserAgent(version string) string {
|
||||
if strings.TrimSpace(version) == "" {
|
||||
version = CLIClientVersion
|
||||
}
|
||||
return "xai-grok-workspace/" + version
|
||||
platform, arch := runtime.GOOS, runtime.GOARCH
|
||||
if platform == "darwin" {
|
||||
platform = "macos"
|
||||
}
|
||||
switch arch {
|
||||
case "amd64":
|
||||
arch = "x86_64"
|
||||
case "arm64":
|
||||
arch = "aarch64"
|
||||
case "386":
|
||||
arch = "x86"
|
||||
}
|
||||
return "grok-pager/" + version + " grok-shell/" + version + " (" + platform + "; " + arch + ")"
|
||||
}
|
||||
|
||||
// ApplyCLIProxyHeaders stamps the fixed Grok CLI identity when the request
|
||||
@@ -75,5 +88,7 @@ func ApplyCLIProxyHeaders(req *http.Request) {
|
||||
req.Header.Set("X-XAI-Token-Auth", CLITokenAuth)
|
||||
req.Header.Set("x-grok-client-version", version)
|
||||
req.Header.Set("x-grok-client-identifier", CLIClientIdentifier)
|
||||
req.Header.Set("x-grok-client-mode", CLIClientMode)
|
||||
req.Header.Set("x-authenticateresponse", "authenticate-response")
|
||||
req.Header.Set("User-Agent", CLIUserAgent(version))
|
||||
}
|
||||
|
||||
@@ -2,11 +2,40 @@ package xai
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"runtime"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/mod/semver"
|
||||
)
|
||||
|
||||
func TestApplyCLIProxyHeadersMeetsUpstreamMinimumVersion(t *testing.T) {
|
||||
// 上游 426 明确要求至少 1.0.13;断言独立于生产版本常量,防止旧覆盖值绕过下限。
|
||||
for _, override := range []string{"", "0.2.93", "0.2.120", "1.0.12", "1.0.13-beta.1", "1.0.13", "1.0.14-alpha.1"} {
|
||||
t.Run("override="+override, func(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, override)
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
ApplyCLIProxyHeaders(req)
|
||||
|
||||
version := req.Header.Get("x-grok-client-version")
|
||||
require.True(t, semver.IsValid("v"+version))
|
||||
require.GreaterOrEqual(t, semver.Compare("v"+version, "v1.0.13"), 0,
|
||||
"Grok Responses 会以 426 拒绝版本 %s,最低要求为 1.0.13", version)
|
||||
require.Contains(t, req.Header.Get("User-Agent"), "grok-shell/"+version+" (")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCLIUserAgentMatchesOfficialInteractiveCapture(t *testing.T) {
|
||||
// 来源:官方 CLI 1.0.46 交互式界面在 Linux x86_64 的主 Responses 请求抓包。
|
||||
if runtime.GOOS != "linux" || runtime.GOARCH != "amd64" {
|
||||
t.Skip("该抓包样本来自 Linux x86_64")
|
||||
}
|
||||
require.Equal(t, "grok-pager/1.0.46 grok-shell/1.0.46 (linux; x86_64)", CLIUserAgent("1.0.46"))
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionDefaultsToPinnedClientVersion(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "")
|
||||
// Default advertise pin is CLIClientVersion; CLIStableVersion is only the floor.
|
||||
@@ -16,18 +45,18 @@ func TestResolveCLIVersionDefaultsToPinnedClientVersion(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionAcceptsValidOverride(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "0.2.95-alpha.1")
|
||||
require.Equal(t, "0.2.95-alpha.1", ResolveCLIVersion())
|
||||
t.Setenv(CLIVersionEnv, "1.0.14-alpha.1")
|
||||
require.Equal(t, "1.0.14-alpha.1", ResolveCLIVersion())
|
||||
}
|
||||
|
||||
func TestResolveCLIVersionRejectsUnsafeOrTooOld(t *testing.T) {
|
||||
for _, version := range []string{
|
||||
"0.2.92",
|
||||
"0.2.93-beta.1",
|
||||
"0.2.95\r\nX-Injected: true",
|
||||
"0.2.093",
|
||||
"0.3",
|
||||
"1",
|
||||
"1.0.12",
|
||||
"1.0.13-beta.1",
|
||||
"1.0.14\r\nX-Injected: true",
|
||||
"1.0.014",
|
||||
"1.1",
|
||||
"2",
|
||||
} {
|
||||
t.Run(version, func(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, version)
|
||||
@@ -48,11 +77,13 @@ func TestApplyCLIProxyHeaders(t *testing.T) {
|
||||
require.Equal(t, CLIClientVersion, req.Header.Get("x-grok-client-version"))
|
||||
require.Equal(t, CLIClientIdentifier, req.Header.Get("x-grok-client-identifier"))
|
||||
require.Equal(t, CLITokenAuth, req.Header.Get("X-XAI-Token-Auth"))
|
||||
require.Equal(t, "interactive", req.Header.Get("x-grok-client-mode"))
|
||||
require.Equal(t, "authenticate-response", req.Header.Get("x-authenticateresponse"))
|
||||
require.Equal(t, CLIUserAgent(CLIClientVersion), req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
func TestApplyCLIProxyHeadersLeavesAPIHostUnchanged(t *testing.T) {
|
||||
t.Setenv(CLIVersionEnv, "0.2.95")
|
||||
t.Setenv(CLIVersionEnv, "1.0.14")
|
||||
|
||||
req, err := http.NewRequest(http.MethodPost, "https://api.x.ai/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
@@ -63,5 +94,7 @@ func TestApplyCLIProxyHeadersLeavesAPIHostUnchanged(t *testing.T) {
|
||||
require.Empty(t, req.Header.Get("x-grok-client-version"))
|
||||
require.Empty(t, req.Header.Get("x-grok-client-identifier"))
|
||||
require.Empty(t, req.Header.Get("X-XAI-Token-Auth"))
|
||||
require.Empty(t, req.Header.Get("x-grok-client-mode"))
|
||||
require.Empty(t, req.Header.Get("x-authenticateresponse"))
|
||||
require.Equal(t, "direct-api-client/1.0", req.Header.Get("User-Agent"))
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
const (
|
||||
apiKeyRateLimitKeyPrefix = "apikey:ratelimit:"
|
||||
apiKeyRateLimitDuration = 24 * time.Hour
|
||||
apiKeyCreateCountKeyPrefix = "apikey:create_count:"
|
||||
apiKeyAuthCachePrefix = "apikey:auth:"
|
||||
authCacheInvalidateChannel = "auth:cache:invalidate"
|
||||
)
|
||||
@@ -24,6 +25,11 @@ func apiKeyRateLimitKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyRateLimitKeyPrefix, userID)
|
||||
}
|
||||
|
||||
// apiKeyCreateCountKey generates the Redis key for per-user API key creation counting.
|
||||
func apiKeyCreateCountKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", apiKeyCreateCountKeyPrefix, userID)
|
||||
}
|
||||
|
||||
func apiKeyAuthCacheKey(key string) string {
|
||||
return fmt.Sprintf("%s%s", apiKeyAuthCachePrefix, key)
|
||||
}
|
||||
@@ -54,9 +60,18 @@ func (c *apiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID in
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
key := apiKeyRateLimitKey(userID)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
// IncrementCreateCount 在固定窗口内累加创建次数并返回累加后的值。
|
||||
// ExpireNX 只在首次创建计数键时设置过期,后续创建不会延长窗口;
|
||||
// MULTI 保证 INCR 与 EXPIRE 同时生效,避免计数键丢失 TTL 后永久封禁。
|
||||
func (c *apiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
key := apiKeyCreateCountKey(userID)
|
||||
pipe := c.rdb.TxPipeline()
|
||||
incr := pipe.Incr(ctx, key)
|
||||
pipe.ExpireNX(ctx, key, window)
|
||||
if _, err := pipe.Exec(ctx); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return incr.Val(), nil
|
||||
}
|
||||
|
||||
func (c *apiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -51,19 +51,6 @@ func (s *ApiKeyCacheSuite) TestCreateAttemptCount() {
|
||||
s.AssertTTLWithin(ttl, 1*time.Second, apiKeyRateLimitDuration)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "delete_removes_key",
|
||||
fn: func(ctx context.Context, rdb *redis.Client, cache *apiKeyCache) {
|
||||
userID := int64(1)
|
||||
|
||||
require.NoError(s.T(), cache.IncrementCreateAttemptCount(ctx, userID))
|
||||
require.NoError(s.T(), cache.DeleteCreateAttemptCount(ctx, userID), "DeleteCreateAttemptCount")
|
||||
|
||||
count, err := cache.GetCreateAttemptCount(ctx, userID)
|
||||
require.NoError(s.T(), err, "expected nil error after delete")
|
||||
require.Equal(s.T(), 0, count, "expected zero count after delete")
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyCacheIncrementCreateCountUsesFixedWindow(t *testing.T) {
|
||||
server := miniredis.RunT(t)
|
||||
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
||||
defer func() { _ = client.Close() }()
|
||||
cache := NewAPIKeyCache(client)
|
||||
ctx := context.Background()
|
||||
key := apiKeyCreateCountKey(7)
|
||||
|
||||
count, err := cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
require.Equal(t, time.Hour, server.TTL(key))
|
||||
|
||||
// 后续创建不得延长窗口,否则持续创建会让计数永不过期。
|
||||
server.FastForward(40 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(2), count)
|
||||
require.Equal(t, 20*time.Minute, server.TTL(key))
|
||||
|
||||
// 其他用户独立计数。
|
||||
count, err = cache.IncrementCreateCount(ctx, 8, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
|
||||
// 窗口到期后重新计数。
|
||||
server.FastForward(21 * time.Minute)
|
||||
count, err = cache.IncrementCreateCount(ctx, 7, time.Hour)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), count)
|
||||
}
|
||||
@@ -652,6 +652,17 @@ func apiKeyListOrder(params pagination.PaginationParams) []func(*entsql.Selector
|
||||
sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
|
||||
sortOrder := params.NormalizedSortOrder(pagination.SortOrderDesc)
|
||||
|
||||
if sortBy == "group" {
|
||||
// Sort before pagination, keeping ungrouped keys last in either direction.
|
||||
opts := []entsql.OrderTermOption{entsql.OrderNullsLast()}
|
||||
tieOrder := dbent.Asc(apikey.FieldID)
|
||||
if sortOrder == pagination.SortOrderDesc {
|
||||
opts = append(opts, entsql.OrderDesc())
|
||||
tieOrder = dbent.Desc(apikey.FieldID)
|
||||
}
|
||||
return []func(*entsql.Selector){apikey.ByGroupField(group.FieldName, opts...), tieOrder}
|
||||
}
|
||||
|
||||
var field string
|
||||
switch sortBy {
|
||||
case "name":
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAPIKeyRepositoryListByUserIDSortByGroup(t *testing.T) {
|
||||
repo, client := newAPIKeyRepoSQLite(t)
|
||||
ctx := context.Background()
|
||||
user := mustCreateAPIKeyRepoUser(t, ctx, client, "group-sort@test.com")
|
||||
otherUser := mustCreateAPIKeyRepoUser(t, ctx, client, "other-group-sort@test.com")
|
||||
createGroup := func(name string) *dbent.Group {
|
||||
g, err := client.Group.Create().SetName(name).Save(ctx)
|
||||
require.NoError(t, err)
|
||||
return g
|
||||
}
|
||||
// Create groups in reverse name order so sorting by group ID cannot pass.
|
||||
zulu := createGroup("Zulu")
|
||||
alpha := createGroup("Alpha")
|
||||
createKey := func(userID int64, name string, groupID *int64, status string) int64 {
|
||||
key := &service.APIKey{
|
||||
UserID: userID, Key: "sk-" + name, Name: name, GroupID: groupID, Status: status,
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, key))
|
||||
return key.ID
|
||||
}
|
||||
zuluFirst := createKey(user.ID, "match-zulu-first", &zulu.ID, service.StatusActive)
|
||||
ungroupedFirst := createKey(user.ID, "match-ungrouped-first", nil, service.StatusActive)
|
||||
alphaFirst := createKey(user.ID, "match-alpha-first", &alpha.ID, service.StatusActive)
|
||||
zuluSecond := createKey(user.ID, "match-zulu-second", &zulu.ID, service.StatusActive)
|
||||
alphaSecond := createKey(user.ID, "match-alpha-second", &alpha.ID, service.StatusActive)
|
||||
ungroupedSecond := createKey(user.ID, "match-ungrouped-second", nil, service.StatusActive)
|
||||
createKey(otherUser.ID, "match-other-user", &alpha.ID, service.StatusActive)
|
||||
createKey(user.ID, "excluded-search", &alpha.ID, service.StatusActive)
|
||||
createKey(user.ID, "match-inactive", &alpha.ID, service.StatusDisabled)
|
||||
deleted := createKey(user.ID, "match-deleted", &alpha.ID, service.StatusActive)
|
||||
require.NoError(t, repo.Delete(ctx, deleted))
|
||||
|
||||
ungroupedID := int64(0)
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
order string
|
||||
groupID *int64
|
||||
want []int64
|
||||
}{
|
||||
{"ascending", "asc", nil, []int64{alphaFirst, alphaSecond, zuluFirst, zuluSecond, ungroupedFirst, ungroupedSecond}},
|
||||
{"descending", "desc", nil, []int64{zuluSecond, zuluFirst, alphaSecond, alphaFirst, ungroupedSecond, ungroupedFirst}},
|
||||
{"group filter", "asc", &alpha.ID, []int64{alphaFirst, alphaSecond}},
|
||||
{"ungrouped filter", "desc", &ungroupedID, []int64{ungroupedSecond, ungroupedFirst}},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
var got []int64
|
||||
const pageSize = 3
|
||||
pages := (len(tc.want) + pageSize - 1) / pageSize
|
||||
for page := 1; page <= pages; page++ {
|
||||
keys, result, err := repo.ListByUserID(ctx, user.ID, pagination.PaginationParams{
|
||||
Page: page, PageSize: pageSize, SortBy: "group", SortOrder: tc.order,
|
||||
}, service.APIKeyListFilters{Search: "match", Status: service.StatusActive, GroupID: tc.groupID})
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, len(tc.want), result.Total)
|
||||
require.Equal(t, pages, result.Pages)
|
||||
for _, key := range keys {
|
||||
got = append(got, key.ID)
|
||||
if key.GroupID != nil {
|
||||
require.NotNil(t, key.Group, "sorting must preserve group preloading")
|
||||
}
|
||||
}
|
||||
}
|
||||
require.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// 余额在途预留:
|
||||
// - billing:inflight:{uid} ZSET member=requestID score=过期时间(ms)
|
||||
// - billing:inflight_amt:{uid} HASH field=requestID value=预留金额(USD)
|
||||
//
|
||||
// 两个 key 使用同一 hash tag,保证 Redis Cluster 下 Lua 可同时访问。
|
||||
const (
|
||||
billingInflightKeyPrefix = "billing:inflight:"
|
||||
billingInflightAmtKeyPrefix = "billing:inflight_amt:"
|
||||
)
|
||||
|
||||
func billingInflightKeys(userID int64) (string, string) {
|
||||
return fmt.Sprintf("%s{%d}", billingInflightKeyPrefix, userID),
|
||||
fmt.Sprintf("%s{%d}", billingInflightAmtKeyPrefix, userID)
|
||||
}
|
||||
|
||||
var (
|
||||
// KEYS: [1]=zset [2]=hash
|
||||
// ARGV: [1]=now_ms [2]=expire_at_ms [3]=member [4]=amount [5]=balance [6]=key_ttl_ms
|
||||
// 返回 {allowed, inflight_sum(string), inflight_count}
|
||||
// 规则:先清理已过期成员;若无在途预留则直接放行(与旧行为一致,外层已校验余额 > 阈值);
|
||||
// 否则要求 balance - sum(在途) >= amount。放行时登记预留。
|
||||
reserveInflightBalanceScript = redis.NewScript(`
|
||||
local now = tonumber(ARGV[1])
|
||||
local expired = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
if #expired > 0 then
|
||||
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now)
|
||||
redis.call('HDEL', KEYS[2], unpack(expired))
|
||||
end
|
||||
local vals = redis.call('HVALS', KEYS[2])
|
||||
local sum = 0
|
||||
for i = 1, #vals do
|
||||
sum = sum + (tonumber(vals[i]) or 0)
|
||||
end
|
||||
local count = redis.call('ZCARD', KEYS[1])
|
||||
local amount = tonumber(ARGV[4])
|
||||
local balance = tonumber(ARGV[5])
|
||||
if count > 0 and (balance - sum) < amount then
|
||||
return {0, tostring(sum), count}
|
||||
end
|
||||
redis.call('ZADD', KEYS[1], tonumber(ARGV[2]), ARGV[3])
|
||||
redis.call('HSET', KEYS[2], ARGV[3], ARGV[4])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[6])
|
||||
redis.call('PEXPIRE', KEYS[2], ARGV[6])
|
||||
return {1, tostring(sum), count}
|
||||
`)
|
||||
|
||||
// KEYS: [1]=zset [2]=hash ARGV: [1]=member [2]=expire_at_ms [3]=key_ttl_ms [4]=now_ms
|
||||
// 仅当成员仍存在且尚未过期(score > now)时续期(XX),并把两个 key 的 TTL 至少延长到 key_ttl_ms。
|
||||
// 已过期但尚未被惰性清理的成员不得被复活:顺手清掉并返回 0。
|
||||
renewInflightBalanceScript = redis.NewScript(`
|
||||
local score = redis.call('ZSCORE', KEYS[1], ARGV[1])
|
||||
if not score then
|
||||
return 0
|
||||
end
|
||||
if tonumber(score) <= tonumber(ARGV[4]) then
|
||||
redis.call('ZREM', KEYS[1], ARGV[1])
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 0
|
||||
end
|
||||
redis.call('ZADD', KEYS[1], 'XX', tonumber(ARGV[2]), ARGV[1])
|
||||
local ttl = tonumber(ARGV[3])
|
||||
if redis.call('PTTL', KEYS[1]) < ttl then
|
||||
redis.call('PEXPIRE', KEYS[1], ttl)
|
||||
end
|
||||
if redis.call('PTTL', KEYS[2]) < ttl then
|
||||
redis.call('PEXPIRE', KEYS[2], ttl)
|
||||
end
|
||||
return 1
|
||||
`)
|
||||
|
||||
releaseInflightBalanceScript = redis.NewScript(`
|
||||
redis.call('ZREM', KEYS[1], ARGV[1])
|
||||
redis.call('HDEL', KEYS[2], ARGV[1])
|
||||
return 1
|
||||
`)
|
||||
)
|
||||
|
||||
// ReserveInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
func (c *billingCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
now := time.Now().UnixMilli()
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
res, err := reserveInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey},
|
||||
now,
|
||||
now+ttlMs,
|
||||
requestID,
|
||||
strconv.FormatFloat(amount, 'f', -1, 64),
|
||||
strconv.FormatFloat(balance, 'f', -1, 64),
|
||||
ttlMs,
|
||||
).Slice()
|
||||
if err != nil {
|
||||
return false, 0, err
|
||||
}
|
||||
if len(res) < 2 {
|
||||
return false, 0, fmt.Errorf("unexpected inflight reservation reply: %v", res)
|
||||
}
|
||||
allowed, _ := res[0].(int64)
|
||||
var sum float64
|
||||
if s, ok := res[1].(string); ok {
|
||||
sum, _ = strconv.ParseFloat(s, 64)
|
||||
}
|
||||
return allowed == 1, sum, nil
|
||||
}
|
||||
|
||||
// ReleaseInflightBalance 实现 service.InflightBalanceReservationCache。
|
||||
func (c *billingCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
return releaseInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID).Err()
|
||||
}
|
||||
|
||||
// RenewInflightBalance 实现 service.InflightBalanceReservationRenewer。
|
||||
func (c *billingCache) RenewInflightBalance(ctx context.Context, userID int64, requestID string, ttl time.Duration) (bool, error) {
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
ttlMs := ttl.Milliseconds()
|
||||
if ttlMs <= 0 {
|
||||
ttlMs = 1
|
||||
}
|
||||
now := time.Now().UnixMilli()
|
||||
n, err := renewInflightBalanceScript.Run(ctx, c.rdb, []string{zkey, hkey}, requestID, now+ttlMs, ttlMs, now).Int64()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
var (
|
||||
_ service.InflightBalanceReservationCache = (*billingCache)(nil)
|
||||
_ service.InflightBalanceReservationRenewer = (*billingCache)(nil)
|
||||
)
|
||||
@@ -0,0 +1,287 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newInflightTestEnv(t *testing.T, enabled bool, ttlSeconds int) (*miniredis.Miniredis, *billingCache, *service.BillingCacheService) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
cache := &billingCache{rdb: rdb}
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{
|
||||
Enabled: enabled,
|
||||
TTLSeconds: ttlSeconds,
|
||||
DefaultMaxTokens: 8192,
|
||||
}
|
||||
svc := service.NewBillingCacheService(cache, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
return mr, cache, svc
|
||||
}
|
||||
|
||||
func inflightCount(t *testing.T, cache *billingCache, userID int64) int64 {
|
||||
t.Helper()
|
||||
zkey, _ := billingInflightKeys(userID)
|
||||
n, err := cache.rdb.ZCard(context.Background(), zkey).Result()
|
||||
require.NoError(t, err)
|
||||
return n
|
||||
}
|
||||
|
||||
func TestInflightReservation_ConcurrentAdmitsOnlyWhatBalanceCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 42}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
const n = 20
|
||||
var admitted, rejected atomic.Int32
|
||||
releases := make(chan func(), n)
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
for i := 0; i < n; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3)
|
||||
if err != nil {
|
||||
if !errors.Is(err, service.ErrInsufficientBalance) {
|
||||
t.Errorf("unexpected error: %v", err)
|
||||
}
|
||||
rejected.Add(1)
|
||||
return
|
||||
}
|
||||
admitted.Add(1)
|
||||
releases <- release
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(releases)
|
||||
|
||||
// balance 1.0, estimate 0.3: first always admitted, then while 1.0 - inflight >= 0.3 → 3 total.
|
||||
require.Equal(t, int32(3), admitted.Load())
|
||||
require.Equal(t, int32(n-3), rejected.Load())
|
||||
require.Equal(t, int64(3), inflightCount(t, cache, user.ID))
|
||||
|
||||
for release := range releases {
|
||||
release()
|
||||
release() // idempotent
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
|
||||
// After release, capacity is available again.
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.3)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
|
||||
func TestInflightReservation_FirstRequestAlwaysAdmitted(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 7}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01))
|
||||
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 5)
|
||||
require.NoError(t, err, "a lone request keeps legacy behavior even if estimate exceeds balance")
|
||||
_, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
release()
|
||||
release2, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.001)
|
||||
require.NoError(t, err)
|
||||
release2()
|
||||
}
|
||||
|
||||
func TestInflightReservation_TTLExpiryFreesLeakedReservation(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 1)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 9}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8) // leaked: never released
|
||||
require.NoError(t, err)
|
||||
_, err = svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
|
||||
time.Sleep(1100 * time.Millisecond)
|
||||
release, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0.8)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), inflightCount(t, cache, user.ID))
|
||||
release()
|
||||
}
|
||||
|
||||
func TestInflightReservation_DisabledKeepsLegacyBehavior(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, false, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 11}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0.01))
|
||||
require.False(t, svc.InflightReservationEnabled())
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_SubscriptionUnaffected(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 12}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 0))
|
||||
group := &service.Group{ID: 1, SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
sub := &service.UserSubscription{ID: 1}
|
||||
for i := 0; i < 5; i++ {
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, group, sub, 10)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_ZeroEstimateSkips(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 13}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1))
|
||||
_, err := svc.ReserveInflightBalance(ctx, user, nil, nil, 0)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
// downReserveCache 余额读取正常,但预留走一个已关闭的 Redis。
|
||||
type downReserveCache struct {
|
||||
service.BillingCache
|
||||
down *billingCache
|
||||
}
|
||||
|
||||
func (d *downReserveCache) ReserveInflightBalance(ctx context.Context, userID int64, requestID string, amount, balance float64, ttl time.Duration) (bool, float64, error) {
|
||||
return d.down.ReserveInflightBalance(ctx, userID, requestID, amount, balance, ttl)
|
||||
}
|
||||
|
||||
func (d *downReserveCache) ReleaseInflightBalance(ctx context.Context, userID int64, requestID string) error {
|
||||
return d.down.ReleaseInflightBalance(ctx, userID, requestID)
|
||||
}
|
||||
|
||||
func TestInflightReservation_RedisDownFailsOpen(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
up := &billingCache{rdb: rdb}
|
||||
require.NoError(t, up.SetUserBalance(context.Background(), 5, 0.01))
|
||||
|
||||
downMR := miniredis.RunT(t)
|
||||
downRDB := redis.NewClient(&redis.Options{Addr: downMR.Addr(), MaxRetries: -1, DialTimeout: 100 * time.Millisecond})
|
||||
t.Cleanup(func() { _ = downRDB.Close() })
|
||||
downMR.Close()
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.Billing.InflightReservation = config.InflightReservationConfig{Enabled: true, TTLSeconds: 60}
|
||||
svc := service.NewBillingCacheService(&downReserveCache{BillingCache: up, down: &billingCache{rdb: downRDB}}, nil, nil, nil, nil, nil, cfg, nil)
|
||||
t.Cleanup(svc.Stop)
|
||||
|
||||
user := &service.User{ID: 5}
|
||||
for i := 0; i < 3; i++ {
|
||||
release, err := svc.ReserveInflightBalance(context.Background(), user, nil, nil, 100)
|
||||
require.NoError(t, err)
|
||||
release()
|
||||
}
|
||||
}
|
||||
|
||||
func TestInflightReservation_ScriptErrorSurfaces(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr(), MaxRetries: -1})
|
||||
cache := &billingCache{rdb: rdb}
|
||||
mr.Close()
|
||||
_, _, err := cache.ReserveInflightBalance(context.Background(), 1, "r", 1, 1, time.Second)
|
||||
require.Error(t, err)
|
||||
require.False(t, errors.Is(err, context.Canceled))
|
||||
_ = rdb.Close()
|
||||
}
|
||||
|
||||
// 回归:原实现在 handler 返回时即释放预留,而计费(RecordUsage → 余额缓存扣减)是异步的,
|
||||
// 窗口内的顺序请求看到「在途=0 且余额未扣」全部放行(余额 $1、5×$0.9 → -4.40)。
|
||||
// 现在预留由计费任务在余额缓存扣减之后才释放。
|
||||
func TestInflightReservation_SequentialInBillingWindowAdmitsOnlyWhatBalanceCovers(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 77}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
const cost = 0.9
|
||||
var pendingBilling []func()
|
||||
admitted := 0
|
||||
for i := 0; i < 5; i++ {
|
||||
res, err := svc.ReserveInflight(ctx, user, nil, nil, cost)
|
||||
if err != nil {
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance)
|
||||
continue
|
||||
}
|
||||
admitted++
|
||||
taskDone := res.Acquire() // handler submits the async billing task
|
||||
res.HandlerDone() // handler returns before billing lands
|
||||
pendingBilling = append(pendingBilling, func() {
|
||||
require.NoError(t, cache.DeductUserBalance(ctx, user.ID, cost))
|
||||
taskDone()
|
||||
})
|
||||
}
|
||||
require.Equal(t, 1, admitted, "requests arriving before billing lands must not be admitted")
|
||||
|
||||
for _, bill := range pendingBilling {
|
||||
bill()
|
||||
}
|
||||
bal, err := cache.GetUserBalance(ctx, user.ID)
|
||||
require.NoError(t, err)
|
||||
require.InDelta(t, 0.1, bal, 1e-9)
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
}
|
||||
|
||||
func TestInflightReservation_RenewalKeepsStreamingReservationAlive(t *testing.T) {
|
||||
_, cache, svc := newInflightTestEnv(t, true, 1)
|
||||
ctx := context.Background()
|
||||
user := &service.User{ID: 78}
|
||||
require.NoError(t, cache.SetUserBalance(ctx, user.ID, 1.0))
|
||||
|
||||
res, err := svc.ReserveInflight(ctx, user, nil, nil, 0.8)
|
||||
require.NoError(t, err)
|
||||
time.Sleep(2500 * time.Millisecond) // streaming well past the 1s TTL
|
||||
_, err = svc.ReserveInflight(ctx, user, nil, nil, 0.8)
|
||||
require.ErrorIs(t, err, service.ErrInsufficientBalance, "renewed reservation must still count")
|
||||
require.Equal(t, int64(1), inflightCount(t, cache, user.ID))
|
||||
|
||||
res.HandlerDone()
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, user.ID))
|
||||
|
||||
ok, err := cache.RenewInflightBalance(ctx, user.ID, "missing", time.Second)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok, "renew must not resurrect a released reservation")
|
||||
}
|
||||
|
||||
func TestInflightReservation_RenewDoesNotResurrectExpiredMember(t *testing.T) {
|
||||
_, cache, _ := newInflightTestEnv(t, true, 60)
|
||||
ctx := context.Background()
|
||||
userID := int64(79)
|
||||
zkey, hkey := billingInflightKeys(userID)
|
||||
// 已过期但尚未被惰性清理的成员(score 在过去)。
|
||||
require.NoError(t, cache.rdb.ZAdd(ctx, zkey, redis.Z{Score: float64(time.Now().UnixMilli() - 1000), Member: "stale"}).Err())
|
||||
require.NoError(t, cache.rdb.HSet(ctx, hkey, "stale", "0.5").Err())
|
||||
|
||||
ok, err := cache.RenewInflightBalance(ctx, userID, "stale", time.Minute)
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok, "renew must not resurrect an expired member")
|
||||
require.Equal(t, int64(0), inflightCount(t, cache, userID))
|
||||
exists, err := cache.rdb.HExists(ctx, hkey, "stale").Result()
|
||||
require.NoError(t, err)
|
||||
require.False(t, exists)
|
||||
}
|
||||
@@ -17,8 +17,42 @@ const (
|
||||
passwordResetKeyPrefix = "password_reset:"
|
||||
passwordResetSentAtKeyPrefix = "password_reset_sent:"
|
||||
notifyCodeUserRateKeyPrefix = "notify_code_user_rate:"
|
||||
|
||||
// attemptsKeySuffix stores the failed-attempt counter next to a verification code.
|
||||
// Kept in a separate key so it can be incremented atomically with INCR.
|
||||
attemptsKeySuffix = ":attempts"
|
||||
)
|
||||
|
||||
// incrAttemptsScript atomically increments the attempt counter for an existing
|
||||
// verification code and aligns the counter TTL with the code TTL.
|
||||
// KEYS[1] = code key, KEYS[2] = attempts key. Returns -1 when the code is missing.
|
||||
var incrAttemptsScript = redis.NewScript(`
|
||||
if redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
return -1
|
||||
end
|
||||
local n = redis.call('INCR', KEYS[2])
|
||||
local ttl = redis.call('PTTL', KEYS[1])
|
||||
if ttl > 0 then
|
||||
redis.call('PEXPIRE', KEYS[2], ttl)
|
||||
end
|
||||
return n
|
||||
`)
|
||||
|
||||
// consumeResetTokenScript atomically compares the stored token hash and deletes it.
|
||||
// KEYS[1] = reset key, ARGV[1] = expected token hash. Returns 1 on success, 0 otherwise.
|
||||
var consumeResetTokenScript = redis.NewScript(`
|
||||
local v = redis.call('GET', KEYS[1])
|
||||
if not v then
|
||||
return 0
|
||||
end
|
||||
local ok, d = pcall(cjson.decode, v)
|
||||
if not ok or type(d) ~= 'table' or d['Token'] ~= ARGV[1] then
|
||||
return 0
|
||||
end
|
||||
redis.call('DEL', KEYS[1])
|
||||
return 1
|
||||
`)
|
||||
|
||||
// verifyCodeKey generates the Redis key for email verification code.
|
||||
// Email is lowercased for case-insensitive consistency.
|
||||
func verifyCodeKey(email string) string {
|
||||
@@ -50,8 +84,7 @@ func NewEmailCache(rdb *redis.Client) service.EmailCache {
|
||||
return &emailCache{rdb: rdb}
|
||||
}
|
||||
|
||||
func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
key := verifyCodeKey(email)
|
||||
func (c *emailCache) getCode(ctx context.Context, key string) (*service.VerificationCodeData, error) {
|
||||
val, err := c.rdb.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -60,21 +93,56 @@ func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*se
|
||||
if err := json.Unmarshal([]byte(val), &data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if n, err := c.rdb.Get(ctx, key+attemptsKeySuffix).Int(); err == nil && n > data.Attempts {
|
||||
data.Attempts = n
|
||||
}
|
||||
return &data, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
key := verifyCodeKey(email)
|
||||
func (c *emailCache) setCode(ctx context.Context, key string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
val, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
pipe := c.rdb.TxPipeline()
|
||||
pipe.Set(ctx, key, val, ttl)
|
||||
pipe.Del(ctx, key+attemptsKeySuffix)
|
||||
if data.Attempts > 0 {
|
||||
pipe.Set(ctx, key+attemptsKeySuffix, data.Attempts, ttl)
|
||||
}
|
||||
_, err = pipe.Exec(ctx)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *emailCache) incrCodeAttempts(ctx context.Context, key string) (int, error) {
|
||||
n, err := incrAttemptsScript.Run(ctx, c.rdb, []string{key, key + attemptsKeySuffix}).Int()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if n < 0 {
|
||||
return 0, redis.Nil
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) deleteCode(ctx context.Context, key string) error {
|
||||
return c.rdb.Del(ctx, key, key+attemptsKeySuffix).Err()
|
||||
}
|
||||
|
||||
func (c *emailCache) GetVerificationCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
return c.getCode(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) SetVerificationCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
return c.setCode(ctx, verifyCodeKey(email), data, ttl)
|
||||
}
|
||||
|
||||
func (c *emailCache) IncrVerificationCodeAttempts(ctx context.Context, email string) (int, error) {
|
||||
return c.incrCodeAttempts(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) DeleteVerificationCode(ctx context.Context, email string) error {
|
||||
key := verifyCodeKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
return c.deleteCode(ctx, verifyCodeKey(email))
|
||||
}
|
||||
|
||||
// Password reset token methods
|
||||
@@ -101,6 +169,16 @@ func (c *emailCache) SetPasswordResetToken(ctx context.Context, email string, da
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
}
|
||||
|
||||
// ConsumePasswordResetToken atomically deletes the stored reset token when its
|
||||
// stored hash equals tokenHash. Returns true only for the single winning caller.
|
||||
func (c *emailCache) ConsumePasswordResetToken(ctx context.Context, email, tokenHash string) (bool, error) {
|
||||
n, err := consumeResetTokenScript.Run(ctx, c.rdb, []string{passwordResetKey(email)}, tokenHash).Int()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return n == 1, nil
|
||||
}
|
||||
|
||||
func (c *emailCache) DeletePasswordResetToken(ctx context.Context, email string) error {
|
||||
key := passwordResetKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
@@ -122,30 +200,19 @@ func (c *emailCache) SetPasswordResetEmailCooldown(ctx context.Context, email st
|
||||
// Notify email verification code methods
|
||||
|
||||
func (c *emailCache) GetNotifyVerifyCode(ctx context.Context, email string) (*service.VerificationCodeData, error) {
|
||||
key := notifyVerifyKey(email)
|
||||
val, err := c.rdb.Get(ctx, key).Result()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var data service.VerificationCodeData
|
||||
if err := json.Unmarshal([]byte(val), &data); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &data, nil
|
||||
return c.getCode(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) SetNotifyVerifyCode(ctx context.Context, email string, data *service.VerificationCodeData, ttl time.Duration) error {
|
||||
key := notifyVerifyKey(email)
|
||||
val, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.rdb.Set(ctx, key, val, ttl).Err()
|
||||
return c.setCode(ctx, notifyVerifyKey(email), data, ttl)
|
||||
}
|
||||
|
||||
func (c *emailCache) IncrNotifyVerifyCodeAttempts(ctx context.Context, email string) (int, error) {
|
||||
return c.incrCodeAttempts(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
func (c *emailCache) DeleteNotifyVerifyCode(ctx context.Context, email string) error {
|
||||
key := notifyVerifyKey(email)
|
||||
return c.rdb.Del(ctx, key).Err()
|
||||
return c.deleteCode(ctx, notifyVerifyKey(email))
|
||||
}
|
||||
|
||||
// User-level rate limiting for notify email verification codes
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newMiniredisEmailCache(t *testing.T) (service.EmailCache, *miniredis.Miniredis, *redis.Client) {
|
||||
t.Helper()
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
return NewEmailCache(rdb), mr, rdb
|
||||
}
|
||||
|
||||
func TestEmailCache_ConcurrentWrongCodesCannotExceedAttemptCap(t *testing.T) {
|
||||
cache, _, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "user@example.com"
|
||||
|
||||
svc := service.NewEmailService(nil, cache)
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{
|
||||
Code: "123456",
|
||||
CreatedAt: time.Now(),
|
||||
ExpiresAt: time.Now().Add(15 * time.Minute),
|
||||
}, 15*time.Minute))
|
||||
|
||||
const workers = 50
|
||||
var invalid, maxed atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
err := svc.VerifyCode(ctx, email, "000000")
|
||||
switch {
|
||||
case errors.Is(err, service.ErrInvalidVerifyCode):
|
||||
invalid.Add(1)
|
||||
case errors.Is(err, service.ErrVerifyCodeMaxAttempts):
|
||||
maxed.Add(1)
|
||||
default:
|
||||
t.Errorf("unexpected result: %v", err)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Only attempts 1..4 may return "invalid"; every other guess is rejected by the cap.
|
||||
require.LessOrEqual(t, int(invalid.Load()), 4)
|
||||
require.Equal(t, workers, int(invalid.Load()+maxed.Load()))
|
||||
|
||||
// Even the correct code is now rejected.
|
||||
require.ErrorIs(t, svc.VerifyCode(ctx, email, "123456"), service.ErrVerifyCodeMaxAttempts)
|
||||
|
||||
data, err := cache.GetVerificationCode(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.GreaterOrEqual(t, data.Attempts, 5)
|
||||
}
|
||||
|
||||
func TestEmailCache_AttemptsResetOnNewCodeAndTTLFollowsCode(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "User@Example.com"
|
||||
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "1"}, time.Minute))
|
||||
n, err := cache.IncrVerificationCodeAttempts(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, n)
|
||||
require.Greater(t, mr.TTL(verifyCodeKey(email)+attemptsKeySuffix), time.Duration(0))
|
||||
|
||||
require.NoError(t, cache.SetVerificationCode(ctx, email, &service.VerificationCodeData{Code: "2"}, time.Minute))
|
||||
data, err := cache.GetVerificationCode(ctx, email)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 0, data.Attempts)
|
||||
|
||||
require.NoError(t, cache.DeleteVerificationCode(ctx, email))
|
||||
_, err = cache.IncrVerificationCodeAttempts(ctx, email)
|
||||
require.Error(t, err)
|
||||
require.False(t, mr.Exists(verifyCodeKey(email)+attemptsKeySuffix))
|
||||
}
|
||||
|
||||
func TestEmailCache_PasswordResetTokenHashedAndSingleUse(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "reset@example.com"
|
||||
|
||||
svc := service.NewEmailService(nil, cache)
|
||||
|
||||
// Seed the token the same way SendPasswordResetEmail does (hash only).
|
||||
token, err := svc.GeneratePasswordResetToken()
|
||||
require.NoError(t, err)
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{
|
||||
Token: hex.EncodeToString(sum[:]), CreatedAt: time.Now(),
|
||||
}, 30*time.Minute))
|
||||
|
||||
raw, err := mr.Get(passwordResetKey(email))
|
||||
require.NoError(t, err)
|
||||
require.False(t, strings.Contains(raw, token), "plaintext token must not be stored")
|
||||
|
||||
require.NoError(t, svc.VerifyPasswordResetToken(ctx, email, token))
|
||||
require.ErrorIs(t, svc.ConsumePasswordResetToken(ctx, email, "wrong"), service.ErrInvalidResetToken)
|
||||
|
||||
const workers = 30
|
||||
var ok atomic.Int32
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
if svc.ConsumePasswordResetToken(ctx, email, token) == nil {
|
||||
ok.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
require.Equal(t, int32(1), ok.Load())
|
||||
require.False(t, mr.Exists(passwordResetKey(email)))
|
||||
}
|
||||
|
||||
func TestEmailCache_ConsumePasswordResetTokenMismatchKeepsToken(t *testing.T) {
|
||||
cache, mr, _ := newMiniredisEmailCache(t)
|
||||
ctx := context.Background()
|
||||
email := "keep@example.com"
|
||||
require.NoError(t, cache.SetPasswordResetToken(ctx, email, &service.PasswordResetTokenData{Token: "abc"}, time.Minute))
|
||||
|
||||
ok, err := cache.ConsumePasswordResetToken(ctx, email, "xyz")
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok)
|
||||
require.True(t, mr.Exists(passwordResetKey(email)))
|
||||
|
||||
ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
|
||||
require.NoError(t, err)
|
||||
require.True(t, ok)
|
||||
ok, err = cache.ConsumePasswordResetToken(ctx, email, "abc")
|
||||
require.NoError(t, err)
|
||||
require.False(t, ok)
|
||||
}
|
||||
@@ -466,6 +466,8 @@ func newGrokOfficialAPIFallbackRequest(req *http.Request) (*http.Request, error)
|
||||
for _, header := range []string{
|
||||
"X-XAI-Token-Auth",
|
||||
"X-Grok-Client-Version",
|
||||
"X-Grok-Client-Mode",
|
||||
"X-Authenticateresponse",
|
||||
"X-Grok-Client-Surface",
|
||||
"X-UserID",
|
||||
"X-Email",
|
||||
@@ -531,6 +533,8 @@ func applyGrokCLIProxyHeaders(req *http.Request) {
|
||||
req.Header.Set("X-XAI-Token-Auth", xai.CLITokenAuth)
|
||||
req.Header.Set("x-grok-client-version", version)
|
||||
req.Header.Set("x-grok-client-identifier", xai.CLIClientIdentifier)
|
||||
req.Header.Set("x-grok-client-mode", xai.CLIClientMode)
|
||||
req.Header.Set("x-authenticateresponse", "authenticate-response")
|
||||
req.Header.Set("User-Agent", xai.CLIUserAgent(version))
|
||||
}
|
||||
|
||||
|
||||
@@ -237,6 +237,10 @@ func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeOAuthRoundTrip(t *testing.T)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
|
||||
require.Equal(t, xai.CLIClientVersion, capturedHeaders.Get("x-grok-client-version"))
|
||||
require.Equal(t, "1.0.46", capturedHeaders.Get("x-grok-client-version"))
|
||||
require.Equal(t, "grok-pager", capturedHeaders.Get("x-grok-client-identifier"))
|
||||
require.Equal(t, "interactive", capturedHeaders.Get("x-grok-client-mode"))
|
||||
require.Equal(t, "authenticate-response", capturedHeaders.Get("x-authenticateresponse"))
|
||||
require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth"))
|
||||
require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), capturedHeaders.Get("User-Agent"))
|
||||
})
|
||||
@@ -309,6 +313,8 @@ func TestHTTPUpstreamDoFallsBackToOfficialGrokAPIOnCLIAccessDenied(t *testing.T)
|
||||
require.Equal(t, "Bearer oauth-token", fallbackHeaders.Get("Authorization"))
|
||||
require.Empty(t, fallbackHeaders.Get("X-XAI-Token-Auth"))
|
||||
require.Empty(t, fallbackHeaders.Get("x-grok-client-version"))
|
||||
require.Empty(t, fallbackHeaders.Get("x-grok-client-mode"))
|
||||
require.Empty(t, fallbackHeaders.Get("x-authenticateresponse"))
|
||||
require.Empty(t, fallbackHeaders.Get("User-Agent"))
|
||||
}
|
||||
|
||||
@@ -465,18 +471,18 @@ func TestApplyGrokCLIProxyHeaders(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("accepts a valid operator override", func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "0.2.121-alpha.1")
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "1.0.47-alpha.1")
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/chat/completions", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
applyGrokCLIProxyHeaders(req)
|
||||
|
||||
require.Equal(t, "0.2.121-alpha.1", req.Header.Get("x-grok-client-version"))
|
||||
require.Equal(t, xai.CLIUserAgent("0.2.121-alpha.1"), req.Header.Get("User-Agent"))
|
||||
require.Equal(t, "1.0.47-alpha.1", req.Header.Get("x-grok-client-version"))
|
||||
require.Equal(t, xai.CLIUserAgent("1.0.47-alpha.1"), req.Header.Get("User-Agent"))
|
||||
})
|
||||
|
||||
t.Run("rejects an unsafe override", func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "0.2.121\r\nX-Injected: true")
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "1.0.47\r\nX-Injected: true")
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -487,7 +493,7 @@ func TestApplyGrokCLIProxyHeaders(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("rejects an override below the supported minimum", func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "0.2.119")
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "1.0.45")
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -498,7 +504,7 @@ func TestApplyGrokCLIProxyHeaders(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("rejects a prerelease override at the minimum version", func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "0.2.120-beta.1")
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "1.0.46-beta.1")
|
||||
req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -511,11 +517,11 @@ func TestApplyGrokCLIProxyHeaders(t *testing.T) {
|
||||
// Every entry sits above the pinned minimum, so a rejection here can only be
|
||||
// caused by the malformed semver and never by the version being too old.
|
||||
for _, version := range []string{
|
||||
"0.2.0121",
|
||||
"0.2.121-alpha..1",
|
||||
"0.3",
|
||||
"1",
|
||||
"0.2.121+build.1",
|
||||
"1.0.047",
|
||||
"1.0.47-alpha..1",
|
||||
"1.1",
|
||||
"2",
|
||||
"1.0.47+build.1",
|
||||
} {
|
||||
t.Run("rejects invalid semver "+version, func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", version)
|
||||
@@ -530,7 +536,7 @@ func TestApplyGrokCLIProxyHeaders(t *testing.T) {
|
||||
}
|
||||
|
||||
t.Run("leaves direct xAI API requests unchanged", func(t *testing.T) {
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "0.2.95")
|
||||
t.Setenv("XAI_GROK_CLI_VERSION", "1.0.47")
|
||||
req, err := http.NewRequest(http.MethodPost, "https://api.x.ai/v1/responses", nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("User-Agent", "direct-api-client/1.0")
|
||||
|
||||
@@ -187,16 +187,17 @@ func (r *usageBillingRepository) applyUsageBillingEffects(ctx context.Context, t
|
||||
result.BalanceOverdrafted = !sufficient
|
||||
}
|
||||
|
||||
// Key 已不存在时跳过其自身的额度/限速计数,其余结算项不受影响。
|
||||
if cmd.APIKeyQuotaCost > 0 {
|
||||
exhausted, err := incrementUsageBillingAPIKeyQuota(ctx, tx, cmd.APIKeyID, cmd.APIKeyQuotaCost)
|
||||
if err != nil {
|
||||
if err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
return err
|
||||
}
|
||||
result.APIKeyQuotaExhausted = exhausted
|
||||
}
|
||||
|
||||
if cmd.APIKeyRateLimitCost > 0 {
|
||||
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil {
|
||||
if err := incrementUsageBillingAPIKeyRateLimit(ctx, tx, cmd.APIKeyID, cmd.APIKeyRateLimitCost); err != nil && !errors.Is(err, service.ErrAPIKeyNotFound) {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
@@ -162,6 +162,61 @@ func TestUsageBillingRepositoryApply_RequestFingerprintConflict(t *testing.T) {
|
||||
require.ErrorIs(t, err, service.ErrUsageBillingRequestConflict)
|
||||
}
|
||||
|
||||
func TestUsageBillingRepositoryApply_DeletedAPIKeyStillBillsBalance(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
repo := NewUsageBillingRepository(client, integrationDB)
|
||||
|
||||
user := mustCreateUser(t, client, &service.User{
|
||||
Email: fmt.Sprintf("usage-billing-deleted-key-user-%d@example.com", time.Now().UnixNano()),
|
||||
PasswordHash: "hash",
|
||||
Balance: 100,
|
||||
})
|
||||
apiKey := mustCreateApiKey(t, client, &service.APIKey{
|
||||
UserID: user.ID,
|
||||
Key: "sk-usage-billing-deleted-key-" + uuid.NewString(),
|
||||
Name: "billing-deleted-key",
|
||||
Quota: 50,
|
||||
RateLimit5h: 50,
|
||||
})
|
||||
account := mustCreateAccount(t, client, &service.Account{
|
||||
Name: "usage-billing-deleted-key-account-" + uuid.NewString(),
|
||||
Type: service.AccountTypeAPIKey,
|
||||
})
|
||||
|
||||
_, err := integrationDB.ExecContext(ctx, "UPDATE api_keys SET deleted_at = NOW() WHERE id = $1", apiKey.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
requestID := uuid.NewString()
|
||||
result, err := repo.Apply(ctx, &service.UsageBillingCommand{
|
||||
RequestID: requestID,
|
||||
APIKeyID: apiKey.ID,
|
||||
UserID: user.ID,
|
||||
AccountID: account.ID,
|
||||
AccountType: service.AccountTypeAPIKey,
|
||||
BalanceCost: 1.25,
|
||||
APIKeyQuotaCost: 1.25,
|
||||
APIKeyRateLimitCost: 1.25,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.True(t, result.Applied)
|
||||
require.False(t, result.APIKeyQuotaExhausted)
|
||||
|
||||
var balance float64
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT balance FROM users WHERE id = $1", user.ID).Scan(&balance))
|
||||
require.InDelta(t, 98.75, balance, 0.000001)
|
||||
|
||||
var quotaUsed, usage5h float64
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT quota_used, usage_5h FROM api_keys WHERE id = $1", apiKey.ID).Scan("aUsed, &usage5h))
|
||||
require.InDelta(t, 0, quotaUsed, 0.000001)
|
||||
require.InDelta(t, 0, usage5h, 0.000001)
|
||||
|
||||
var dedupCount int
|
||||
require.NoError(t, integrationDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM usage_billing_dedup WHERE request_id = $1 AND api_key_id = $2", requestID, apiKey.ID).Scan(&dedupCount))
|
||||
require.Equal(t, 1, dedupCount)
|
||||
}
|
||||
|
||||
func TestUsageBillingRepositoryApply_UpdatesAccountQuota(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
client := testEntClient(t)
|
||||
|
||||
@@ -20,6 +20,8 @@ const (
|
||||
captureBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance\s+\+ CASE WHEN \$1 > \$2 THEN \$1 - \$2 ELSE 0 END\s+- CASE WHEN \$2 > \$1 THEN \$2 - \$1 ELSE 0 END,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$3 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
||||
releaseBatchImageHoldSQL = `(?s)UPDATE users\s+SET balance = balance \+ \$1,\s+frozen_balance = COALESCE\(frozen_balance, 0\) - \$1,\s+updated_at = NOW\(\)\s+WHERE id = \$2 AND deleted_at IS NULL AND COALESCE\(frozen_balance, 0\) >= \$1\s+RETURNING balance, frozen_balance`
|
||||
userExistsForBillingSQL = `(?s)SELECT 1\s+FROM users\s+WHERE id = \$1 AND deleted_at IS NULL`
|
||||
apiKeyQuotaIncrementSQL = `(?s)UPDATE api_keys\s+SET quota_used = quota_used \+ \$1,.*WHERE id = \$2 AND deleted_at IS NULL\s+RETURNING`
|
||||
apiKeyRateLimitIncrementSQL = `(?s)UPDATE api_keys SET\s+usage_5h = .*WHERE id = \$2 AND deleted_at IS NULL`
|
||||
)
|
||||
|
||||
func TestDeductUsageBillingBalance_UsesSufficientBalanceGuard(t *testing.T) {
|
||||
@@ -99,6 +101,66 @@ func TestApplyUsageBillingEffects_FlagsBalanceOverdraft(t *testing.T) {
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyUsageBillingEffects_DeletedAPIKeyStillBillsBalance(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
mock.ExpectBegin()
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectQuery(conditionalBalanceDeductSQL).
|
||||
WithArgs(10.0, int64(42)).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"balance"}).AddRow(90.0))
|
||||
mock.ExpectQuery(apiKeyQuotaIncrementSQL).
|
||||
WithArgs(10.0, int64(7), service.StatusAPIKeyActive, service.StatusAPIKeyQuotaExhausted).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
mock.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||
WithArgs(10.0, int64(7)).
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectCommit()
|
||||
|
||||
result := &service.UsageBillingApplyResult{Applied: true}
|
||||
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||
UserID: 42,
|
||||
APIKeyID: 7,
|
||||
BalanceCost: 10,
|
||||
APIKeyQuotaCost: 10,
|
||||
APIKeyRateLimitCost: 10,
|
||||
}, result)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result.NewBalance)
|
||||
require.InDelta(t, 90.0, *result.NewBalance, 0.000001)
|
||||
require.False(t, result.APIKeyQuotaExhausted)
|
||||
require.NoError(t, tx.Commit())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyUsageBillingEffects_APIKeyCounterErrorStillFails(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
mock.ExpectBegin()
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(apiKeyRateLimitIncrementSQL).
|
||||
WithArgs(10.0, int64(7)).
|
||||
WillReturnError(sql.ErrConnDone)
|
||||
mock.ExpectRollback()
|
||||
|
||||
err = (&usageBillingRepository{}).applyUsageBillingEffects(ctx, tx, &service.UsageBillingCommand{
|
||||
UserID: 42,
|
||||
APIKeyID: 7,
|
||||
APIKeyRateLimitCost: 10,
|
||||
}, &service.UsageBillingApplyResult{Applied: true})
|
||||
require.ErrorIs(t, err, sql.ErrConnDone)
|
||||
require.NoError(t, tx.Rollback())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestDeductUsageBillingBalance_ReturnsUserNotFoundWhenNoUserUpdated(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
db, mock, err := sqlmock.New()
|
||||
|
||||
@@ -106,6 +106,8 @@ func extractProtocolSegments(protocol string, document any) []promptSegment {
|
||||
return append(extractInstructions(root["instructions"]), extractResponses(root["input"])...)
|
||||
case "openai_images", "grok_media", "media", "images":
|
||||
return userPromptSegments(extractMediaPrompts(root))
|
||||
case "typesafe_systemone":
|
||||
return extractSystemOneSegments(root)
|
||||
default:
|
||||
if segments := extractChatLikeSegments(root); len(segments) > 0 {
|
||||
return segments
|
||||
@@ -120,6 +122,79 @@ func extractProtocolSegments(protocol string, document any) []promptSegment {
|
||||
}
|
||||
}
|
||||
|
||||
// extractSystemOneSegments collects every client-controlled text of a TypeSafe
|
||||
// System One request: question IDs, every question field except the validated
|
||||
// type, unknown top-level extension fields, and the evaluated state. Object keys
|
||||
// are text too (Jev reads the whole JSON), so they are collected with values.
|
||||
// The state comes last so it is the prioritized segment, and keys are visited
|
||||
// in sorted order to keep the prompt hash stable across requests.
|
||||
func extractSystemOneSegments(root map[string]any) []promptSegment {
|
||||
if root == nil {
|
||||
return nil
|
||||
}
|
||||
texts := make([]string, 0, 4)
|
||||
questions, isObject := root["questions"].(map[string]any)
|
||||
if !isObject {
|
||||
texts = appendJSONStringLeaves(texts, root["questions"])
|
||||
}
|
||||
for _, id := range sortedJSONKeys(questions) {
|
||||
texts = appendJSONStringLeaves(texts, id)
|
||||
question, ok := questions[id].(map[string]any)
|
||||
if !ok {
|
||||
texts = appendJSONStringLeaves(texts, questions[id])
|
||||
continue
|
||||
}
|
||||
for _, field := range sortedJSONKeys(question) {
|
||||
switch field {
|
||||
case "type":
|
||||
continue
|
||||
case "instructions", "criteria":
|
||||
default:
|
||||
texts = appendJSONStringLeaves(texts, field)
|
||||
}
|
||||
texts = appendJSONStringLeaves(texts, question[field])
|
||||
}
|
||||
}
|
||||
for _, field := range sortedJSONKeys(root) {
|
||||
switch field {
|
||||
case "model", "stream", "state", "questions":
|
||||
continue
|
||||
}
|
||||
texts = appendJSONStringLeaves(texts, field)
|
||||
texts = appendJSONStringLeaves(texts, root[field])
|
||||
}
|
||||
texts = appendJSONStringLeaves(texts, root["state"])
|
||||
return userPromptSegments(texts)
|
||||
}
|
||||
|
||||
func appendJSONStringLeaves(texts []string, value any) []string {
|
||||
switch typed := value.(type) {
|
||||
case string:
|
||||
if text := strings.TrimSpace(typed); text != "" {
|
||||
texts = append(texts, text)
|
||||
}
|
||||
case []any:
|
||||
for _, item := range typed {
|
||||
texts = appendJSONStringLeaves(texts, item)
|
||||
}
|
||||
case map[string]any:
|
||||
for _, key := range sortedJSONKeys(typed) {
|
||||
texts = appendJSONStringLeaves(texts, key)
|
||||
texts = appendJSONStringLeaves(texts, typed[key])
|
||||
}
|
||||
}
|
||||
return texts
|
||||
}
|
||||
|
||||
func sortedJSONKeys(values map[string]any) []string {
|
||||
keys := make([]string, 0, len(values))
|
||||
for key := range values {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
|
||||
// clientInstructionRoles are roles a client may freely populate. Attackers can
|
||||
// place jailbreak/PII text in assistant/tool turns, so blocking audit must scan
|
||||
// them too—not only user/system/developer instructions.
|
||||
|
||||
@@ -412,3 +412,53 @@ func mustJSON(t *testing.T, value string) []byte {
|
||||
func metadataTextForTest(scanText string) string {
|
||||
return strings.Replace(scanText, promptAuditPrioritySeparator, "\n\n", 1)
|
||||
}
|
||||
|
||||
func TestPromptSnapshotTypeSafeSystemOneCollectsStateAndQuestions(t *testing.T) {
|
||||
body := `{"model":"jev-latest","state":{"STATE_KEY":"STATE_TITLE","items":["STATE_ITEM",{"text":"STATE_NESTED"},3]},` +
|
||||
`"questions":{"b":{"type":"choice","instructions":"CHOICE_INSTRUCTIONS","criteria":{"OPTION_LABEL":{"description":"OPTION_DESC"},"empty":null}},` +
|
||||
`"a":{"type":"noul","instructions":["NOUL_INSTRUCTIONS"],"criteria":{"true":"NOUL_TRUE","false":"NOUL_FALSE"},"EXTENSION_KEY":"EXTENSION_VALUE"},` +
|
||||
`"QUESTION_ID":{"type":"score","criteria":["SCORE_LOW",{"description":"SCORE_HIGH"}]}},"TOP_EXTENSION":{"x":"TOP_VALUE"}}`
|
||||
|
||||
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "typesafe_systemone", Body: []byte(body)})
|
||||
require.NoError(t, err)
|
||||
for _, text := range []string{"STATE_KEY", "STATE_TITLE", "STATE_ITEM", "STATE_NESTED", "CHOICE_INSTRUCTIONS", "OPTION_LABEL", "OPTION_DESC",
|
||||
"NOUL_INSTRUCTIONS", "NOUL_TRUE", "NOUL_FALSE", "EXTENSION_KEY", "EXTENSION_VALUE", "QUESTION_ID", "SCORE_LOW", "SCORE_HIGH",
|
||||
"TOP_EXTENSION", "TOP_VALUE"} {
|
||||
require.Contains(t, snapshot.ScanText, text)
|
||||
}
|
||||
// Canonical field names and validated enum values are not client text.
|
||||
for _, text := range []string{"jev-latest", "instructions", "criteria", "questions", "choice", "score"} {
|
||||
require.NotContains(t, snapshot.ScanText, text)
|
||||
}
|
||||
|
||||
// Map iteration order must not change the audited text or its hash.
|
||||
for range 20 {
|
||||
again, err := ExtractPromptSnapshot(Request{Protocol: "typesafe_systemone", Body: []byte(body)})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, snapshot.PromptHash, again.PromptHash)
|
||||
require.Equal(t, snapshot.ScanText, again.ScanText)
|
||||
}
|
||||
|
||||
blocking, err := ExtractBlockingPromptSnapshot(Request{Protocol: "typesafe_systemone", Body: []byte(body)}, true)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, blocking.ScanText, "STATE_TITLE")
|
||||
require.Contains(t, blocking.ScanText, "CHOICE_INSTRUCTIONS")
|
||||
require.Contains(t, blocking.ScanText, "QUESTION_ID")
|
||||
}
|
||||
|
||||
func TestPromptSnapshotTypeSafeSystemOneStringStateIsPrioritized(t *testing.T) {
|
||||
snapshot, err := ExtractPromptSnapshot(Request{Protocol: "typesafe_systemone", Body: []byte(`{"model":"jev-latest","state":"plain state","questions":{"q":{"type":"noul"}}}`)})
|
||||
require.NoError(t, err)
|
||||
require.True(t, strings.HasPrefix(snapshot.ScanText, "plain state"))
|
||||
require.Equal(t, 2, snapshot.MessageCount)
|
||||
}
|
||||
|
||||
func TestPromptSnapshotTypeSafeSystemOneAuditsKeyOnlyPayloads(t *testing.T) {
|
||||
body := `{"model":"jev-latest","state":{"HIDDEN_STATE_KEY":1},"questions":{"HIDDEN_QUESTION_ID":{"type":"noul"}}}`
|
||||
for _, latestTurnOnly := range []bool{false, true} {
|
||||
snapshot, err := ExtractBlockingPromptSnapshot(Request{Protocol: "typesafe_systemone", Body: []byte(body)}, latestTurnOnly)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, snapshot.ScanText, "HIDDEN_STATE_KEY")
|
||||
require.Contains(t, snapshot.ScanText, "HIDDEN_QUESTION_ID")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -863,7 +863,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"force_email_on_third_party_signup": false,
|
||||
"default_concurrency": 5,
|
||||
"default_balance": 1.25,
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": null,
|
||||
@@ -971,6 +971,9 @@ func TestAPIContracts(t *testing.T) {
|
||||
"payment_balance_recharge_multiplier": 0,
|
||||
"payment_subscription_usd_to_cny_rate": 0,
|
||||
"payment_recharge_fee_rate": 0,
|
||||
"payment_recharge_bonus_tiers": [],
|
||||
"payment_recharge_bonus_mode": "bonus",
|
||||
"payment_recharge_bonus_notice": "",
|
||||
"payment_load_balance_strategy": "",
|
||||
"payment_product_name_prefix": "",
|
||||
"payment_product_name_suffix": "",
|
||||
@@ -1186,7 +1189,7 @@ func TestAPIContracts(t *testing.T) {
|
||||
"purchase_subscription_url": "",
|
||||
"table_default_page_size": 20,
|
||||
"table_page_size_options": [10, 20, 50],
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"minimax":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"opencode_go":{"daily":null,"weekly":null,"monthly":null},"typesafe":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
|
||||
"auth_source_default_email_platform_quotas": null,
|
||||
"auth_source_default_github_platform_quotas": null,
|
||||
"auth_source_default_google_platform_quotas": null,
|
||||
@@ -1291,6 +1294,9 @@ func TestAPIContracts(t *testing.T) {
|
||||
"payment_balance_recharge_multiplier": 0,
|
||||
"payment_subscription_usd_to_cny_rate": 0,
|
||||
"payment_recharge_fee_rate": 0,
|
||||
"payment_recharge_bonus_tiers": [],
|
||||
"payment_recharge_bonus_mode": "bonus",
|
||||
"payment_recharge_bonus_notice": "",
|
||||
"payment_load_balance_strategy": "",
|
||||
"payment_product_name_prefix": "",
|
||||
"payment_product_name_suffix": "",
|
||||
@@ -1737,8 +1743,8 @@ func (stubApiKeyCache) IncrementCreateAttemptCount(ctx context.Context, userID i
|
||||
return nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) DeleteCreateAttemptCount(ctx context.Context, userID int64) error {
|
||||
return nil
|
||||
func (stubApiKeyCache) IncrementCreateCount(ctx context.Context, userID int64, window time.Duration) (int64, error) {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (stubApiKeyCache) IncrementDailyUsage(ctx context.Context, apiKey string) error {
|
||||
|
||||
@@ -129,7 +129,7 @@ func registerRoutes(
|
||||
routes.RegisterModelPlazaRoutes(v1, h, optionalJWTAuth, settingService, panelRateLimiter)
|
||||
routes.RegisterAdminRoutes(v1, h, adminAuth, auditLog, stepUpAuth, settingService, panelRateLimiter)
|
||||
routes.RegisterGatewayRoutes(r, h, apiKeyAuth, apiKeyService, subscriptionService, opsService, settingService, compositeResolver, cfg)
|
||||
routes.RegisterPaymentRoutes(v1, h.Payment, h.PaymentWebhook, h.Admin.Payment, jwtAuth, adminAuth, auditLog, settingService, panelRateLimiter)
|
||||
routes.RegisterPaymentRoutes(v1, h.Payment, h.PaymentWebhook, h.Admin.Payment, jwtAuth, adminAuth, auditLog, settingService, panelRateLimiter, redisClient)
|
||||
|
||||
handler.RegisterPageRoutes(v1, cfg.Pricing.DataDir, gin.HandlerFunc(jwtAuth), gin.HandlerFunc(adminAuth), settingService)
|
||||
}
|
||||
|
||||
@@ -368,6 +368,8 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAu
|
||||
accounts.PUT("/opencode-go-usage/settings", h.Admin.Account.UpdateOpenCodeGoUsageSettings)
|
||||
accounts.GET("/:id", h.Admin.Account.GetByID)
|
||||
accounts.GET("/:id/claude/reset-credits", h.Admin.Account.ClaudeResetCredits)
|
||||
// Same protection as the Codex reset-quota route (admin auth, audit, compliance guard).
|
||||
accounts.POST("/:id/claude/reset-credits/redeem", h.Admin.Account.RedeemClaudeResetCredit)
|
||||
accounts.POST("", h.Admin.Account.Create)
|
||||
accounts.POST("/:id/duplicate", h.Admin.Account.Duplicate)
|
||||
accounts.POST("/check-mixed-channel", h.Admin.Account.CheckMixedChannel)
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
adminhandler "github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// The Claude redeem route consumes an irreversible credit just like the Codex
|
||||
// reset-quota route, so it must sit behind exactly the same middleware chain.
|
||||
func TestClaudeResetRedeemRouteMatchesCodexResetProtection(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handlers := &handler.Handlers{Admin: &handler.AdminHandlers{Account: &adminhandler.AccountHandler{}, OpenAIOAuth: &adminhandler.OpenAIOAuthHandler{}}}
|
||||
chains := map[string][]string{}
|
||||
adminAuth := servermiddleware.AdminAuthMiddleware(func(c *gin.Context) {
|
||||
names := c.HandlerNames()
|
||||
chains[c.FullPath()] = names[:len(names)-1]
|
||||
if c.GetHeader("Authorization") == "" {
|
||||
servermiddleware.AbortWithError(c, http.StatusUnauthorized, "UNAUTHORIZED", "Authorization required")
|
||||
return
|
||||
}
|
||||
servermiddleware.AbortWithError(c, http.StatusForbidden, "FORBIDDEN", "Admin access required")
|
||||
})
|
||||
auditLog := servermiddleware.AuditLogMiddleware(func(c *gin.Context) { c.Next() })
|
||||
stepUp := servermiddleware.StepUpAuthMiddleware(func(c *gin.Context) { c.Next() })
|
||||
RegisterAdminRoutes(router.Group("/api/v1"), handlers, adminAuth, auditLog, stepUp, nil, nil)
|
||||
|
||||
codex := "/api/v1/admin/openai/accounts/1/reset-quota"
|
||||
claude := "/api/v1/admin/accounts/1/claude/reset-credits/redeem"
|
||||
for _, path := range []string{codex, claude} {
|
||||
for _, auth := range []string{"", "Bearer user-token"} {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, path, nil)
|
||||
if auth != "" {
|
||||
request.Header.Set("Authorization", auth)
|
||||
}
|
||||
router.ServeHTTP(recorder, request)
|
||||
if auth == "" {
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code, path)
|
||||
} else {
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code, path)
|
||||
}
|
||||
}
|
||||
}
|
||||
codexChain := chains["/api/v1/admin/openai/accounts/:id/reset-quota"]
|
||||
claudeChain := chains["/api/v1/admin/accounts/:id/claude/reset-credits/redeem"]
|
||||
require.NotEmpty(t, codexChain)
|
||||
require.Equal(t, strings.Join(codexChain, "\n"), strings.Join(claudeChain, "\n"))
|
||||
}
|
||||
@@ -202,6 +202,8 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.Messages(c)
|
||||
})
|
||||
// System One carries only JSON text, so it uses the text body limit.
|
||||
gateway.POST("/systemone", textBodyLimit, h.Gateway.SystemOne)
|
||||
// /v1/messages/count_tokens: OpenAI bridges upstream, Grok estimates
|
||||
// locally, and Anthropic-compatible platforms retain their existing path.
|
||||
gateway.POST("/messages/count_tokens", countTokensHandler)
|
||||
|
||||
@@ -152,6 +152,7 @@ func TestGatewayRoutesGroupModelAllowlistCoversRootAliasRoutes(t *testing.T) {
|
||||
{http.MethodGet, "/realtime?model=gpt-4.1", ""},
|
||||
{http.MethodPost, "/v1/responses", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/messages", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/systemone", `{"model":"gpt-4.1","state":"x","questions":{"q":{"type":"noul","instructions":"x"}}}`},
|
||||
{http.MethodPost, "/v1/messages/count_tokens", `{"model":"gpt-4.1","messages":[]}`},
|
||||
{http.MethodPost, "/v1/chat/completions", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/embeddings", `{"model":"gpt-4.1","input":"hi"}`},
|
||||
|
||||
@@ -1,12 +1,25 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
ratelimit "github.com/Wei-Shaw/sub2api/internal/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// publicOrderVerifyRateLimit caps anonymous legacy out_trade_no lookups per
|
||||
// client IP. The payment result page polls at most a handful of times per
|
||||
// order, so this leaves ample headroom for real users while making
|
||||
// out_trade_no enumeration impractical.
|
||||
const (
|
||||
publicOrderVerifyRateLimit = 20
|
||||
publicOrderVerifyRateLimitWindow = time.Minute
|
||||
)
|
||||
|
||||
// RegisterPaymentRoutes registers all payment-related routes:
|
||||
@@ -21,6 +34,7 @@ func RegisterPaymentRoutes(
|
||||
auditLog middleware.AuditLogMiddleware,
|
||||
settingService *service.SettingService,
|
||||
panelRateLimiter *middleware.PanelRateLimiter,
|
||||
redisClient *redis.Client,
|
||||
) {
|
||||
// --- User-facing payment endpoints (authenticated) ---
|
||||
authenticated := v1.Group("/payment")
|
||||
@@ -50,9 +64,15 @@ func RegisterPaymentRoutes(
|
||||
// Signed resume-token recovery is the preferred public lookup path.
|
||||
// The legacy anonymous out_trade_no verify endpoint remains available as a
|
||||
// persisted-state compatibility path for staggered upgrades.
|
||||
// The anonymous verify endpoint is IP rate-limited to prevent out_trade_no
|
||||
// enumeration. It fails open on Redis errors so an outage never blocks a
|
||||
// user who is mid-payment from seeing their result.
|
||||
publicRateLimiter := ratelimit.NewRateLimiter(redisClient)
|
||||
public := v1.Group("/payment/public")
|
||||
{
|
||||
public.POST("/orders/verify", paymentHandler.VerifyOrderPublic)
|
||||
public.POST("/orders/verify",
|
||||
publicRateLimiter.Limit("payment-public-order-verify", publicOrderVerifyRateLimit, publicOrderVerifyRateLimitWindow),
|
||||
paymentHandler.VerifyOrderPublic)
|
||||
public.POST("/orders/resolve", paymentHandler.ResolveOrderPublicByResumeToken)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newPaymentRoutesTestRouter(redisClient *redis.Client) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
v1 := router.Group("/api/v1")
|
||||
noop := func(c *gin.Context) { c.Next() }
|
||||
|
||||
RegisterPaymentRoutes(
|
||||
v1,
|
||||
&handler.PaymentHandler{},
|
||||
&handler.PaymentWebhookHandler{},
|
||||
&admin.PaymentHandler{},
|
||||
servermiddleware.JWTAuthMiddleware(noop),
|
||||
servermiddleware.AdminAuthMiddleware(noop),
|
||||
servermiddleware.AuditLogMiddleware(noop),
|
||||
nil,
|
||||
nil,
|
||||
redisClient,
|
||||
)
|
||||
return router
|
||||
}
|
||||
|
||||
func postPublicOrderVerify(router *gin.Engine, remoteAddr string) *httptest.ResponseRecorder {
|
||||
// Empty body fails binding in the handler, so no service is touched and a
|
||||
// non-429 response proves the request passed the limiter.
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/payment/public/orders/verify", strings.NewReader(`{}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.RemoteAddr = remoteAddr
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
return w
|
||||
}
|
||||
|
||||
func TestPublicOrderVerifyRateLimitedPerIP(t *testing.T) {
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
|
||||
router := newPaymentRoutesTestRouter(rdb)
|
||||
|
||||
for i := 1; i <= publicOrderVerifyRateLimit; i++ {
|
||||
w := postPublicOrderVerify(router, "198.51.100.20:1234")
|
||||
require.Equal(t, http.StatusBadRequest, w.Code, "request %d should reach the handler", i)
|
||||
}
|
||||
|
||||
w := postPublicOrderVerify(router, "198.51.100.20:1234")
|
||||
require.Equal(t, http.StatusTooManyRequests, w.Code)
|
||||
require.Contains(t, w.Body.String(), "rate limit exceeded")
|
||||
|
||||
// A different client IP has its own budget.
|
||||
w = postPublicOrderVerify(router, "198.51.100.21:1234")
|
||||
require.Equal(t, http.StatusBadRequest, w.Code)
|
||||
}
|
||||
|
||||
func TestPublicOrderVerifyRateLimitFailsOpenWhenRedisUnavailable(t *testing.T) {
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: "127.0.0.1:1",
|
||||
DialTimeout: 50 * time.Millisecond,
|
||||
ReadTimeout: 50 * time.Millisecond,
|
||||
WriteTimeout: 50 * time.Millisecond,
|
||||
})
|
||||
t.Cleanup(func() { _ = rdb.Close() })
|
||||
|
||||
router := newPaymentRoutesTestRouter(rdb)
|
||||
w := postPublicOrderVerify(router, "203.0.113.30:1234")
|
||||
require.Equal(t, http.StatusBadRequest, w.Code, "users mid-payment must not be blocked by a Redis outage")
|
||||
}
|
||||
@@ -29,6 +29,7 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
|
||||
|
||||
audited := map[string][]string{
|
||||
"/messages": {"gateway_handler.go", "openai_gateway_handler.go"},
|
||||
"/systemone": {"gateway_systemone.go"},
|
||||
"/responses": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
|
||||
"/responses/*subpath": {"gateway_handler_responses.go", "openai_gateway_handler.go"},
|
||||
"/chat/completions": {"gateway_handler_chat_completions.go", "openai_chat_completions.go"},
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user