diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3dda14..3b4eeb2e5d21 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -41,6 +41,7 @@ const ( ContextKeyAutoGroup ContextKey = "auto_group" ContextKeyAutoGroupIndex ContextKey = "auto_group_index" ContextKeyAutoGroupRetryIndex ContextKey = "auto_group_retry_index" + ContextKeyPaidOptimizedRoute ContextKey = "paid_optimized_route" /* user related keys */ ContextKeyUserId ContextKey = "id" diff --git a/controller/model.go b/controller/model.go index b32eebd7daac..7c79713a98de 100644 --- a/controller/model.go +++ b/controller/model.go @@ -190,7 +190,7 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) { return modelListGroups{ userGroup: userGroup, tokenGroup: tokenGroup, - ownerGroups: service.GetUserAutoGroup(userGroup), + ownerGroups: service.GetUserAutoGroupForRequest(c, userGroup), }, nil } diff --git a/docs/paid-optimized-api-route.md b/docs/paid-optimized-api-route.md new file mode 100644 index 000000000000..2e6fa28977aa --- /dev/null +++ b/docs/paid-optimized-api-route.md @@ -0,0 +1,224 @@ +# 付费优化 API 线路分组限制实现方案 + +## 背景 + +New API 的 Dashboard 可以展示多个 API 地址。新增一台网络优化机器作为独立 API 入口，为所有外部模型 API 提供更好的网络线路，但不允许零倍率免费流量占用该机器。 + +站点当前通过令牌分组区分免费和付费流量： + +- 分组倍率为 `0`：免费流量； +- 分组倍率大于 `0`：需要正常扣减额度的付费流量； +- 用户账号组目前通常为 `default`，创建令牌时选择的分组决定渠道路由和计费倍率。 + +因此，本功能不以用户账号组或购买记录判断权限，而是以当前令牌的实际计费分组倍率判断该请求能否使用优化线路。 + +## 目标 + +1. 只有实际计费分组倍率大于 `0` 的令牌可以调用付费优化线路。 +2. 倍率为 `0` 的免费分组令牌在优化线路上返回 `403`。 +3. 覆盖 OpenAI、Responses、Anthropic、Gemini、图片、音频、视频、Midjourney、Suno、Kling 和即梦等现有外部 API。 +4. 普通 API 地址、现有渠道路由、计费、限流和管理后台保持不变。 +5. 不增加数据库字段、第二套分组配置或额外的数据库查询。 + +## 非目标 + +- 不判断用户是否充值过或购买过订阅； +- 不修改用户账号组； +- 不改变令牌分组的现有创建和可用范围规则； +- 不开放优化域名下的 Dashboard 页面和 `/api/*` 管理接口； +- 不通过隐藏优化域名代替服务端权限校验。 + +## 识别方式 + +优化反代统一覆盖写入请求头： + +```http +X-NewAPI-Route: paid-optimized +``` + +该请求头只是“启用更严格限制”的线路标记，不是授权凭证。客户端伪造该请求头只能让自己的请求受到额外限制，不能获得额外权限。正式优化反代必须覆盖客户端传入的同名请求头，不能按客户端原值透传。 + +New API 的 `TokenAuth()` 已经完成以下工作： + +1. 验证 API Key； +2. 从令牌记录读取令牌分组； +3. 校验该用户是否有权使用所选令牌分组； +4. 将原始令牌分组写入 `ContextKeyTokenGroup`； +5. 将实际使用分组写入 `ContextKeyUsingGroup`； +6. 将用户账号组写入 `ContextKeyUserGroup`。 + +鉴权完成后调用 `service.GetUserGroupRatio(userGroup, usingGroup)` 获取与现有计费逻辑一致的实际分组倍率。该函数同时兼容用户组对计费分组的特殊倍率覆盖。检查集中在 `TokenAuth()` 中，因此当前和后续所有使用 API Key 鉴权的接口都会自动受到保护，不需要为每个协议重复挂载中间件。 + +## 请求流程 + +```text +客户端 + -> 付费优化反代（覆盖写入 X-NewAPI-Route） + -> New API TokenAuth / TokenOrUserAuth + -> 验证 API Key 或 Dashboard 会话 + -> 非优化线路：直接继续 + -> auto 令牌：只保留实际倍率 > 0 的候选组 + -> auto 没有付费候选组：403 + -> 实际分组倍率 <= 0：403 + -> 实际分组倍率 > 0：继续 + -> ModelRequestRateLimit / Distribute + -> 上游服务 +``` + +检查发生在现有令牌、用户和分组验证之后，并早于模型限流、渠道选择和上游连接。免费请求不会进入模型解析和渠道选择。 + +`/mj/image/:id` 在普通入口原本允许公开读取。为了不改变普通入口行为，该接口只在检测到付费优化线路标记时额外要求 API Key，然后复用相同的分组倍率检查。 + +## 覆盖接口 + +所有经过 `TokenAuth()` 的外部 API 都受限制，包括以下路由前缀及其 GET、POST、DELETE、WebSocket 和 SSE 请求： + +- `/v1/*`：OpenAI Chat Completions、Responses、Anthropic Messages、Realtime、模型、图片、音频、嵌入、重排和视频； +- `/v1beta/*`：Gemini 和 Gemini OpenAI 兼容接口； +- `/kling/v1/*`：Kling 视频任务； +- `/jimeng/*`：即梦任务； +- `/mj/*`、`/:mode/mj/*`：Midjourney 任务和结果； +- `/suno/*`：Suno 任务和结果。 + +Dashboard 的 `/api/*`、网页、登录、用户、渠道和管理员接口不属于外部模型 API，应由优化反代直接返回 `404`，不得代理到源站。`/pg/*` 是 Dashboard 游乐场接口，也不应在优化域名开放。 + +## 判定规则 + +| 场景 | 结果 | +| --- | --- | +| 普通 API 地址，任意令牌分组 | 保持原有行为 | +| 优化地址，分组倍率为 `0` | `403 access_denied` | +| 优化地址，分组倍率大于 `0` | 放行并按原逻辑计费 | +| 优化地址，令牌未指定分组 | 回退到现有实际使用分组并检查其倍率 | +| 优化地址，视频内容使用 Dashboard 会话 | 回退到用户账号组倍率 | +| 优化地址，用户特殊倍率为 `0` | `403 access_denied` | +| 优化地址，`auto` 令牌存在付费候选组 | 过滤免费组后继续选择和重试 | +| 优化地址，`auto` 令牌没有付费候选组 | `403 access_denied` | + +免费账号如果创建了非零倍率分组的令牌，该令牌可以进入优化线路，但请求仍会按对应倍率扣减现有额度；额度不足时继续由 New API 原有额度检查拒绝。这项功能限制的是零倍率免费流量，不是用户的历史付费身份。 + +## `auto` 分组处理 + +`auto` 是令牌的一种分组选择方式，不是另一种 Key。管理员在系统设置中维护 Auto 分组顺序；令牌选择 `auto` 后，请求会按顺序查找用户有权使用、支持当前模型且存在可用渠道的分组。启用跨组重试后，当前分组的渠道失败时还可以继续尝试后面的分组，最终按实际选中的分组倍率计费。 + +优化线路完整支持同一个 `auto` Key，但使用请求级候选组过滤： + +1. 普通入口继续使用完整 Auto 候选组列表，原有行为不变； +2. 优化入口根据用户组特殊倍率和分组倍率，删除实际倍率小于等于 `0` 的候选组； +3. 初次渠道选择、渠道亲和复用、跨组重试和模型列表共用过滤后的列表； +4. 即使全局 Auto 顺序中免费组排在前面，优化入口也会跳过它，只选择付费组； +5. 如果用户没有任何可用的付费 Auto 候选组，返回 `403 access_denied`。 + +该方案不是只检查第一次选中的分组，因此跨组重试不会重新进入免费组。普通入口仍可按完整 Auto 顺序使用免费组。 + +## OpenResty 反代配置 + +优化域名只允许模型 API 路径，不代理主站和 Dashboard： + +```nginx +# 明确禁止 Dashboard 后端和旧版 Dashboard API。必须放在模型 API 正则之前。 +location = /api { return 404; } +location ^~ /api/ { return 404; } +location = /pg { return 404; } +location ^~ /pg/ { return 404; } +location = /dashboard { return 404; } +location ^~ /dashboard/ { return 404; } + +location ~ ^/(?:v1(?:/|$)|v1beta(?:/|$)|kling/v1(?:/|$)|mj(?:/|$)|suno(?:/|$)|jimeng(?:/|$)|[^/]+/mj(?:/|$)) { + proxy_pass http://new-api:3000; + proxy_http_version 1.1; + + proxy_set_header Host $host; + proxy_set_header X-Real-IP $remote_addr; + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-NewAPI-Route "paid-optimized"; + + proxy_set_header Upgrade $http_upgrade; + proxy_set_header Connection $http_connection; + + proxy_buffering off; + proxy_request_buffering off; + proxy_cache off; + proxy_connect_timeout 30s; + proxy_send_timeout 3600s; + proxy_read_timeout 3600s; +} + +location / { + default_type application/json; + return 404 '{"error":{"message":"API endpoint only","type":"not_found"}}'; +} +``` + +部署时将 `proxy_pass` 替换为实际 New API 源站。必须删除原来的 `location ^~ /` 全站代理，否则主站仍会暴露，且 `^~` 可能使 API 白名单正则不生效。显式拒绝规则不能省略：动态 Midjourney 路由 `/:mode/mj/*` 的第一段理论上可以匹配 `api`、`pg` 或 `dashboard`，前置的 `^~` 拒绝规则可以确保这些 Dashboard 路径不会被正则重新放行。优化入口必须始终经过该反代配置；如果同一优化域名还有其他代理层，需要确认自定义请求头最终能到达 New API。 + +## 安全边界 + +- 分组来自服务端验证后的令牌记录，不读取客户端自报的分组字段； +- 倍率来自 New API 当前内存配置，不读取客户端参数； +- 普通入口没有线路标记，不受新增限制； +- 直接访问普通入口不等于绕过，因为普通入口本来就是保留的非优化线路； +- 若优化反代漏写线路标记，限制不会生效，因此部署后必须执行免费令牌验证； +- 请求头不需要设计成秘密，因为它只开启限制，不授予权限； +- 新增外部 API 只要复用 `TokenAuth()` 就会自动受到相同限制；若新增公开无鉴权接口，需要像 Midjourney 图片接口一样补充优化线路条件鉴权。 + +## 性能影响 + +每个标记请求只增加： + +1. 一次请求头字符串比较； +2. 数次 Gin Context 内存读取； +3. 一次内存倍率 Map 查询； +4. 一次浮点数大小判断。 + +不增加数据库、Redis 或外部 HTTP 请求。免费请求还会在渠道分发前提前结束。实际网络性能主要取决于优化反代和 New API 源站之间的链路，而不是该权限判断。 + +## 修改文件 + +- `middleware/auth.go`：在 API Key 和视频内容会话鉴权完成后统一执行线路权限检查； +- `middleware/paid_optimized_route.go`：线路权限判断，以及 Midjourney 公开图片接口的条件鉴权； +- `middleware/paid_optimized_route_test.go`：普通入口、免费、付费、GET、图片、会话回退、特殊倍率和 auto 回归测试； +- `router/relay-router.go`：在优化入口为 Midjourney 公开图片接口补充条件鉴权； +- `service/group.go`：按请求线路生成 Auto 候选组； +- `service/channel_select.go`、`middleware/distributor.go`：初选、亲和和跨组重试使用同一候选组； +- `controller/model.go`：优化入口的 Auto 模型列表只包含付费候选组模型； +- `i18n/keys.go`、`i18n/locales/*.yaml`：标准错误消息。 + +## 验证步骤 + +代码测试： + +```bash +go test ./middleware ./service ./controller ./router +go build ./... +``` + +部署后至少验证以下请求： + +```bash +# 优化入口：免费 Key 应返回 403 +curl -i https://optimized-api.example.com/v1/models \ + -H 'Authorization: Bearer sk-free-example' + +# 优化入口：付费 Key 应进入原有逻辑 +curl -i https://optimized-api.example.com/v1/models \ + -H 'Authorization: Bearer sk-paid-example' + +# 优化入口：Dashboard 应返回 404 +curl -i https://optimized-api.example.com/api/status + +# 普通入口：免费 Key 保持原有行为 +curl -i https://api.example.com/v1/models \ + -H 'Authorization: Bearer sk-free-example' +``` + +还应分别抽查 Chat Completions、Responses、Gemini、图片生成、视频创建和 `/content`、Midjourney、Suno、Kling 与即梦接口。 + +## 回滚方式 + +紧急回滚可以先在优化反代删除 `X-NewAPI-Route` 请求头，立即恢复所有请求的原有行为。代码回滚则移除 `TokenAuth()` 和 `TokenOrUserAuth()` 中的 `rejectUnpaidOptimizedRoute()` 调用，并恢复 Midjourney 图片路由。 + +## 后续工作 + +1. Dashboard 的 API 信息目前来自公开 `/api/status`，而且登录面板不知道用户准备使用哪一个 API Key。后续可先向所有用户展示优化地址并注明“仅非零倍率分组令牌可用”；若必须隐藏，需要另行设计账号级权限或令牌选择交互。 diff --git a/i18n/keys.go b/i18n/keys.go index 8e9a4b5694d4..3953d647da16 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -321,6 +321,8 @@ const ( MsgDistributorNoAvailableChannel = "distributor.no_available_channel" MsgDistributorInvalidMidjourney = "distributor.invalid_midjourney_request" MsgDistributorInvalidParseModel = "distributor.invalid_request_parse_model" + MsgPaidOptimizedRouteAutoGroupDenied = "paid_optimized_route.auto_group_denied" + MsgPaidOptimizedRouteFreeGroupDenied = "paid_optimized_route.free_group_denied" ) // Custom OAuth provider related messages diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 3f1fd03cb090..936b35ede26f 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -271,6 +271,8 @@ distributor.get_channel_failed: "Failed to get available channel for model {{.Mo distributor.no_available_channel: "No available channel for model {{.Model}} under group {{.Group}} (distributor)" distributor.invalid_midjourney_request: "Invalid Midjourney request: {{.Error}}" distributor.invalid_request_parse_model: "Invalid request, unable to parse model" +paid_optimized_route.auto_group_denied: "No paid auto group is available for the paid optimized route" +paid_optimized_route.free_group_denied: "The current token group cannot use the paid optimized route" # Custom OAuth provider messages custom_oauth.not_found: "Custom OAuth provider not found" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index fe982e59a0f3..6954c8ed303c 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -272,6 +272,8 @@ distributor.get_channel_failed: "获取分组 {{.Group}} 下模型 {{.Model}} distributor.no_available_channel: "分组 {{.Group}} 下模型 {{.Model}} 无可用渠道（distributor）" distributor.invalid_midjourney_request: "无效的midjourney请求，{{.Error}}" distributor.invalid_request_parse_model: "无效的请求，无法解析模型" +paid_optimized_route.auto_group_denied: "当前 auto 令牌没有可用于付费优化线路的付费分组" +paid_optimized_route.free_group_denied: "当前令牌分组不能使用付费优化线路" # Custom OAuth provider messages custom_oauth.not_found: "自定义 OAuth 提供商不存在" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 27759d07f378..7ba7ad66ffef 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -272,6 +272,8 @@ distributor.get_channel_failed: "獲取分組 {{.Group}} 下模型 {{.Model}} distributor.no_available_channel: "分組 {{.Group}} 下模型 {{.Model}} 無可用管道（distributor）" distributor.invalid_midjourney_request: "無效的midjourney請求，{{.Error}}" distributor.invalid_request_parse_model: "無效的請求，無法解析模型" +paid_optimized_route.auto_group_denied: "目前 auto 權杖沒有可用於付費最佳化線路的付費分組" +paid_optimized_route.free_group_denied: "目前令牌分組不能使用付費最佳化線路" # Custom OAuth provider messages custom_oauth.not_found: "自訂 OAuth 供應者不存在" diff --git a/middleware/auth.go b/middleware/auth.go index 2ad09a7a8d2d..6322434c9a7c 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -264,6 +264,9 @@ func TokenOrUserAuth() func(c *gin.Context) { return } setDashboardAuthContext(c, user, identity, false) + if rejectUnpaidOptimizedRoute(c) { + return + } c.Next() return } @@ -478,6 +481,9 @@ func TokenAuth() func(c *gin.Context) { if err != nil { return } + if rejectUnpaidOptimizedRoute(c) { + return + } c.Next() } } diff --git a/middleware/distributor.go b/middleware/distributor.go index bde639ddee5b..94089e2ecf88 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -109,7 +109,7 @@ func Distribute() func(c *gin.Context) { channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) + autoGroups := service.GetUserAutoGroupForRequest(c, userGroup) for _, g := range autoGroups { if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { selectGroup = g diff --git a/middleware/paid_optimized_route.go b/middleware/paid_optimized_route.go new file mode 100644 index 000000000000..60f284656037 --- /dev/null +++ b/middleware/paid_optimized_route.go @@ -0,0 +1,59 @@ +package middleware + +import ( + "net/http" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" +) + +const ( + paidOptimizedRouteHeader = "X-NewAPI-Route" + paidOptimizedRouteValue = "paid-optimized" +) + +// rejectUnpaidOptimizedRoute rejects zero-ratio groups on marked paid relay +// nodes after token or dashboard authentication has populated the group context. +func rejectUnpaidOptimizedRoute(c *gin.Context) bool { + if c.GetHeader(paidOptimizedRouteHeader) != paidOptimizedRouteValue { + return false + } + + common.SetContextKey(c, constant.ContextKeyPaidOptimizedRoute, true) + userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) + if common.GetContextKeyString(c, constant.ContextKeyTokenGroup) == "auto" { + if len(service.GetUserAutoGroupForRequest(c, userGroup)) == 0 { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgPaidOptimizedRouteAutoGroupDenied), types.ErrorCodeAccessDenied) + return true + } + return false + } + + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + if usingGroup == "" { + usingGroup = userGroup + } + if service.GetUserGroupRatio(userGroup, usingGroup) <= 0 { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgPaidOptimizedRouteFreeGroupDenied), types.ErrorCodeAccessDenied) + return true + } + + return false +} + +// TokenAuthOnPaidOptimizedRoute requires token authentication only when the +// paid optimized marker is present. It keeps public relay endpoints unchanged +// on ordinary nodes while making them group-aware on the optimized node. +func TokenAuthOnPaidOptimizedRoute() func(c *gin.Context) { + return func(c *gin.Context) { + if c.GetHeader(paidOptimizedRouteHeader) == paidOptimizedRouteValue { + TokenAuth()(c) + return + } + c.Next() + } +} diff --git a/middleware/paid_optimized_route_test.go b/middleware/paid_optimized_route_test.go new file mode 100644 index 000000000000..b5854906bdae --- /dev/null +++ b/middleware/paid_optimized_route_test.go @@ -0,0 +1,137 @@ +package middleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/i18n" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func setupPaidOptimizedRouteRatios(t *testing.T) { + t.Helper() + groupRatios := ratio_setting.GetGroupRatioSetting().GroupRatio + groupSpecialRatios := ratio_setting.GetGroupRatioSetting().GroupGroupRatio + originalGroupRatios := groupRatios.ReadAll() + originalGroupSpecialRatios := groupSpecialRatios.ReadAll() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + t.Cleanup(func() { + groupRatios.Clear() + groupRatios.AddAll(originalGroupRatios) + groupSpecialRatios.Clear() + groupSpecialRatios.AddAll(originalGroupSpecialRatios) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + }) + + groupRatios.Clear() + groupRatios.AddAll(map[string]float64{ + "default": 1, + "free": 0, + "paid": 1, + }) + groupSpecialRatios.Clear() + groupSpecialRatios.Set("sponsored", map[string]float64{"paid": 0}) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["free","paid"]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto":"Auto","default":"Default","free":"Free","paid":"Paid"}`)) +} + +func TestPaidOptimizedRoute(t *testing.T) { + require.NoError(t, i18n.Init()) + setupPaidOptimizedRouteRatios(t) + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + tests := []struct { + name string + method string + path string + marked bool + userGroup string + tokenGroup string + usingGroup string + expected int + handlerRuns bool + }{ + {name: "ordinary node keeps free group behavior", method: http.MethodPost, path: "/v1/chat/completions", userGroup: "default", tokenGroup: "free", usingGroup: "free", expected: http.StatusNoContent, handlerRuns: true}, + {name: "marked node rejects free chat group", method: http.MethodPost, path: "/v1/chat/completions", marked: true, userGroup: "default", tokenGroup: "free", usingGroup: "free", expected: http.StatusForbidden}, + {name: "marked node rejects free image group", method: http.MethodPost, path: "/v1/images/generations", marked: true, userGroup: "default", tokenGroup: "free", usingGroup: "free", expected: http.StatusForbidden}, + {name: "marked node rejects free read request", method: http.MethodGet, path: "/v1/models", marked: true, userGroup: "default", tokenGroup: "free", usingGroup: "free", expected: http.StatusForbidden}, + {name: "marked node allows paid group", method: http.MethodPost, path: "/v1/responses", marked: true, userGroup: "default", tokenGroup: "paid", usingGroup: "paid", expected: http.StatusNoContent, handlerRuns: true}, + {name: "marked node uses token fallback group ratio", method: http.MethodPost, path: "/v1/messages", marked: true, userGroup: "default", expected: http.StatusNoContent, handlerRuns: true}, + {name: "marked node uses session fallback group ratio", method: http.MethodGet, path: "/v1/videos/task/content", marked: true, userGroup: "free", expected: http.StatusForbidden}, + {name: "marked node allows auto with paid candidate", method: http.MethodPost, path: "/v1/responses/compact", marked: true, userGroup: "default", tokenGroup: "auto", usingGroup: "auto", expected: http.StatusNoContent, handlerRuns: true}, + {name: "marked node rejects auto without paid candidate", method: http.MethodPost, path: "/v1/responses", marked: true, userGroup: "sponsored", tokenGroup: "auto", usingGroup: "auto", expected: http.StatusForbidden}, + {name: "marked node honors special zero ratio", method: http.MethodPost, path: "/v1/completions", marked: true, userGroup: "sponsored", tokenGroup: "paid", usingGroup: "paid", expected: http.StatusForbidden}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + handlerRan := false + router := gin.New() + router.Handle(test.method, test.path, func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyUserGroup, test.userGroup) + common.SetContextKey(c, constant.ContextKeyTokenGroup, test.tokenGroup) + common.SetContextKey(c, constant.ContextKeyUsingGroup, test.usingGroup) + c.Next() + }, func(c *gin.Context) { + if rejectUnpaidOptimizedRoute(c) { + return + } + c.Next() + }, func(c *gin.Context) { + handlerRan = true + c.Status(http.StatusNoContent) + }) + + request := httptest.NewRequest(test.method, test.path, nil) + if test.marked { + request.Header.Set(paidOptimizedRouteHeader, paidOptimizedRouteValue) + } + response := httptest.NewRecorder() + router.ServeHTTP(response, request) + + assert.Equal(t, test.expected, response.Code) + assert.Equal(t, test.handlerRuns, handlerRan) + if test.expected == http.StatusForbidden { + assert.Contains(t, response.Body.String(), `"code":"access_denied"`) + } + }) + } +} + +func TestTokenAuthOnPaidOptimizedRoute(t *testing.T) { + require.NoError(t, i18n.Init()) + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + handlerRan := false + router := gin.New() + router.GET("/mj/image/test", TokenAuthOnPaidOptimizedRoute(), func(c *gin.Context) { + handlerRan = true + c.Status(http.StatusNoContent) + }) + + ordinaryRequest := httptest.NewRequest(http.MethodGet, "/mj/image/test", nil) + ordinaryResponse := httptest.NewRecorder() + router.ServeHTTP(ordinaryResponse, ordinaryRequest) + assert.Equal(t, http.StatusNoContent, ordinaryResponse.Code) + assert.True(t, handlerRan) + + handlerRan = false + optimizedRequest := httptest.NewRequest(http.MethodGet, "/mj/image/test", nil) + optimizedRequest.Header.Set(paidOptimizedRouteHeader, paidOptimizedRouteValue) + optimizedResponse := httptest.NewRecorder() + router.ServeHTTP(optimizedResponse, optimizedRequest) + assert.Equal(t, http.StatusUnauthorized, optimizedResponse.Code) + assert.False(t, handlerRan) +} diff --git a/router/relay-router.go b/router/relay-router.go index e08ecb14bc17..6d645a5d2eb3 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -206,7 +206,7 @@ func SetRelayRouter(router *gin.Engine) { } func registerMjRouterGroup(relayMjRouter *gin.RouterGroup) { - relayMjRouter.GET("/image/:id", relay.RelayMidjourneyImage) + relayMjRouter.GET("/image/:id", middleware.TokenAuthOnPaidOptimizedRoute(), relay.RelayMidjourneyImage) relayMjRouter.Use(middleware.TokenAuth(), middleware.Distribute()) { relayMjRouter.POST("/submit/action", controller.RelayMidjourney) diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252bfb3..0b937e8038f3 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -91,7 +91,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, if len(setting.GetAutoGroups()) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } - autoGroups := GetUserAutoGroup(userGroup) + autoGroups := GetUserAutoGroupForRequest(param.Ctx, userGroup) // startGroupIndex: the group index to start searching from // startGroupIndex: 开始搜索的分组索引 diff --git a/service/group.go b/service/group.go index 8cb359bcff63..7066a2ba9f2f 100644 --- a/service/group.go +++ b/service/group.go @@ -3,9 +3,12 @@ package service import ( "strings" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" ) func GetUserUsableGroups(userGroup string) map[string]string { @@ -54,6 +57,24 @@ func GetUserAutoGroup(userGroup string) []string { return autoGroups } +// GetUserAutoGroupForRequest returns the auto candidates available to the +// current request. Paid optimized routes exclude zero-ratio groups so the same +// auto token can safely be used on both ordinary and optimized API addresses. +func GetUserAutoGroupForRequest(c *gin.Context, userGroup string) []string { + autoGroups := GetUserAutoGroup(userGroup) + if !common.GetContextKeyBool(c, constant.ContextKeyPaidOptimizedRoute) { + return autoGroups + } + + paidGroups := make([]string, 0, len(autoGroups)) + for _, group := range autoGroups { + if GetUserGroupRatio(userGroup, group) > 0 { + paidGroups = append(paidGroups, group) + } + } + return paidGroups +} + // GetGroupsEnabledModels 按 groups 顺序获取各分组启用的模型并去重 func GetGroupsEnabledModels(groups []string) []string { seen := make(map[string]struct{}) diff --git a/service/group_test.go b/service/group_test.go new file mode 100644 index 000000000000..f4e2a8567386 --- /dev/null +++ b/service/group_test.go @@ -0,0 +1,50 @@ +package service + +import ( + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetUserAutoGroupForRequestFiltersFreeGroupsOnOptimizedRoute(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + groupRatios := ratio_setting.GetGroupRatioSetting().GroupRatio + groupSpecialRatios := ratio_setting.GetGroupRatioSetting().GroupGroupRatio + originalGroupRatios := groupRatios.ReadAll() + originalGroupSpecialRatios := groupSpecialRatios.ReadAll() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + t.Cleanup(func() { + groupRatios.Clear() + groupRatios.AddAll(originalGroupRatios) + groupSpecialRatios.Clear() + groupSpecialRatios.AddAll(originalGroupSpecialRatios) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + }) + + groupRatios.Clear() + groupRatios.AddAll(map[string]float64{"free": 0, "paid": 1, "premium": 2}) + groupSpecialRatios.Clear() + groupSpecialRatios.Set("sponsored", map[string]float64{"paid": 0, "premium": 0.5}) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["free","paid","premium"]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"free":"Free","paid":"Paid","premium":"Premium"}`)) + + ordinaryContext, _ := gin.CreateTestContext(httptest.NewRecorder()) + assert.Equal(t, []string{"free", "paid", "premium"}, GetUserAutoGroupForRequest(ordinaryContext, "default")) + + optimizedContext, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(optimizedContext, constant.ContextKeyPaidOptimizedRoute, true) + assert.Equal(t, []string{"paid", "premium"}, GetUserAutoGroupForRequest(optimizedContext, "default")) + assert.Equal(t, []string{"premium"}, GetUserAutoGroupForRequest(optimizedContext, "sponsored")) +}