feat: publish sanitized custom.97 source

This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207

No files matched your search

+370
View File
@@ -0,0 +1,370 @@
package middleware
import (
"fmt"
"maps"
"slices"
"sync"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/service/authz"
)
type accessTokenRuleKind int
const (
// accessTokenRuleScope requires the token to carry scope.
accessTokenRuleScope accessTokenRuleKind = iota + 1
// accessTokenRuleAny accepts every valid scoped token; the handler applies
// its own checks.
accessTokenRuleAny
// accessTokenRuleSession rejects every access token, legacy ones included.
accessTokenRuleSession
)
type accessTokenRouteRule struct {
kind accessTokenRuleKind
scope string
}
var (
accessTokenAnyRule = accessTokenRouteRule{kind: accessTokenRuleAny}
accessTokenSessionRule = accessTokenRouteRule{kind: accessTokenRuleSession}
)
func accessTokenScopeRule(scope string) accessTokenRouteRule {
return accessTokenRouteRule{kind: accessTokenRuleScope, scope: scope}
}
// accessTokenRouteRules declares the token rule of every dashboard route that
// is not guarded by RequirePermission, keyed by method and gin full path.
// Casbin-guarded routes declare themselves through
// DeclareAccessTokenPermissionRoute when they are registered.
var accessTokenRouteRules = map[string]accessTokenRouteRule{
// router/api-router.go: top level
"GET /api/models": accessTokenScopeRule("profile:read"),
"GET /api/status/test": accessTokenScopeRule("log:read"),
"GET /api/pricing": accessTokenAnyRule,
"GET /api/perf-metrics/summary": accessTokenAnyRule,
"GET /api/perf-metrics": accessTokenAnyRule,
"GET /api/rankings": accessTokenAnyRule,
"POST /api/oauth/state": accessTokenSessionRule,
"POST /api/oauth/email/bind/start": accessTokenScopeRule("account_security:write"),
"POST /api/oauth/email/bind/resend": accessTokenScopeRule("account_security:write"),
"POST /api/oauth/email/bind": accessTokenScopeRule("account_security:write"),
"POST /api/oauth/wechat/bind": accessTokenScopeRule("account_security:write"),
"POST /api/oauth/telegram/bind/start": accessTokenSessionRule,
"GET /api/oauth/:provider": accessTokenSessionRule,
"GET /api/verify/methods": accessTokenAnyRule,
"POST /api/verify": accessTokenAnyRule,
// router/api-router.go: /api/user (self)
"GET /api/user/sessions": accessTokenSessionRule,
"DELETE /api/user/sessions/:sid": accessTokenSessionRule,
"POST /api/user/sessions/revoke-others": accessTokenSessionRule,
"GET /api/user/self/groups": accessTokenScopeRule("profile:read"),
"GET /api/user/self": accessTokenScopeRule("profile:read"),
"GET /api/user/models": accessTokenScopeRule("profile:read"),
"PUT /api/user/self": accessTokenScopeRule("profile:write"),
"DELETE /api/user/self": accessTokenScopeRule("account_security:write"),
"GET /api/user/access_tokens": accessTokenSessionRule,
"GET /api/user/access_tokens/catalog": accessTokenSessionRule,
"GET /api/user/access_tokens/scopes": accessTokenSessionRule,
"POST /api/user/access_tokens": accessTokenSessionRule,
"PATCH /api/user/access_tokens/:id": accessTokenSessionRule,
"DELETE /api/user/access_tokens/:id": accessTokenSessionRule,
"DELETE /api/user/access_tokens/legacy": accessTokenSessionRule,
"GET /api/user/passkey": accessTokenScopeRule("account_security:read"),
"POST /api/user/passkey/register/begin": accessTokenScopeRule("account_security:write"),
"POST /api/user/passkey/register/finish": accessTokenScopeRule("account_security:write"),
"POST /api/user/passkey/verify/begin": accessTokenAnyRule,
"POST /api/user/passkey/verify/finish": accessTokenAnyRule,
"DELETE /api/user/passkey": accessTokenScopeRule("account_security:write"),
"GET /api/user/aff": accessTokenScopeRule("wallet:read"),
"GET /api/user/topup/info": accessTokenScopeRule("wallet:read"),
"GET /api/user/topup/self": accessTokenScopeRule("wallet:read"),
"POST /api/user/topup": accessTokenScopeRule("wallet:write"),
"POST /api/user/pay": accessTokenScopeRule("wallet:write"),
"POST /api/user/amount": accessTokenScopeRule("wallet:read"),
"POST /api/user/stripe/pay": accessTokenScopeRule("wallet:write"),
"POST /api/user/stripe/amount": accessTokenScopeRule("wallet:read"),
"POST /api/user/creem/pay": accessTokenScopeRule("wallet:write"),
"POST /api/user/waffo/amount": accessTokenScopeRule("wallet:read"),
"POST /api/user/waffo/pay": accessTokenScopeRule("wallet:write"),
"POST /api/user/waffo-pancake/amount": accessTokenScopeRule("wallet:read"),
"POST /api/user/waffo-pancake/pay": accessTokenScopeRule("wallet:write"),
"POST /api/user/aff_transfer": accessTokenScopeRule("wallet:write"),
"PUT /api/user/setting": accessTokenScopeRule("profile:write"),
"GET /api/user/2fa/status": accessTokenScopeRule("account_security:read"),
"POST /api/user/2fa/setup": accessTokenScopeRule("account_security:write"),
"POST /api/user/2fa/enable": accessTokenScopeRule("account_security:write"),
"POST /api/user/2fa/disable": accessTokenScopeRule("account_security:write"),
"POST /api/user/2fa/backup_codes": accessTokenScopeRule("account_security:write"),
"GET /api/user/checkin": accessTokenScopeRule("wallet:read"),
"POST /api/user/checkin": accessTokenScopeRule("wallet:write"),
"GET /api/user/oauth/bindings": accessTokenScopeRule("account_security:read"),
"DELETE /api/user/oauth/bindings/:provider_id": accessTokenScopeRule("account_security:write"),
// router/api-router.go: /api/user (admin)
"GET /api/user/": accessTokenScopeRule("user:read"),
"GET /api/user/topup": accessTokenScopeRule("billing:read"),
"POST /api/user/topup/complete": accessTokenScopeRule("billing:write"),
"GET /api/user/search": accessTokenScopeRule("user:read"),
"GET /api/user/:id/oauth/bindings": accessTokenScopeRule("user:read"),
"DELETE /api/user/:id/oauth/bindings/:provider_id": accessTokenScopeRule("user:write"),
"DELETE /api/user/:id/bindings/:binding_type": accessTokenScopeRule("user:write"),
"GET /api/user/:id": accessTokenScopeRule("user:read"),
"POST /api/user/": accessTokenScopeRule("user:write"),
"POST /api/user/manage": accessTokenScopeRule("user:write"),
"PUT /api/user/": accessTokenScopeRule("user:write"),
"DELETE /api/user/:id": accessTokenScopeRule("user:write"),
"DELETE /api/user/:id/reset_passkey": accessTokenScopeRule("user:write"),
"GET /api/user/2fa/stats": accessTokenScopeRule("user:read"),
"DELETE /api/user/:id/2fa": accessTokenScopeRule("user:write"),
// router/api-router.go: /api/subscription
"GET /api/subscription/plans": accessTokenScopeRule("wallet:read"),
"GET /api/subscription/self": accessTokenScopeRule("wallet:read"),
"PUT /api/subscription/self/preference": accessTokenScopeRule("wallet:write"),
"POST /api/subscription/balance/pay": accessTokenScopeRule("wallet:write"),
"POST /api/subscription/epay/pay": accessTokenScopeRule("wallet:write"),
"POST /api/subscription/stripe/pay": accessTokenScopeRule("wallet:write"),
"POST /api/subscription/creem/pay": accessTokenScopeRule("wallet:write"),
"POST /api/subscription/waffo-pancake/pay": accessTokenScopeRule("wallet:write"),
"GET /api/subscription/admin/plans": accessTokenScopeRule("billing:read"),
"POST /api/subscription/admin/plans": accessTokenScopeRule("billing:write"),
"PUT /api/subscription/admin/plans/:id": accessTokenScopeRule("billing:write"),
"PATCH /api/subscription/admin/plans/:id": accessTokenScopeRule("billing:write"),
"POST /api/subscription/admin/bind": accessTokenScopeRule("billing:write"),
"POST /api/subscription/admin/plans/:id/subscriptions/reset": accessTokenScopeRule("billing:write"),
"GET /api/subscription/admin/users/:id/subscriptions": accessTokenScopeRule("billing:read"),
"POST /api/subscription/admin/users/:id/subscriptions": accessTokenScopeRule("billing:write"),
"POST /api/subscription/admin/users/:id/subscriptions/reset": accessTokenScopeRule("billing:write"),
"POST /api/subscription/admin/user_subscriptions/:id/invalidate": accessTokenScopeRule("billing:write"),
"DELETE /api/subscription/admin/user_subscriptions/:id": accessTokenScopeRule("billing:write"),
// router/api-router.go: /api/option
"GET /api/option/": accessTokenScopeRule("option:read"),
"GET /api/option/request_policy": accessTokenScopeRule("option:read"),
"PATCH /api/option/request_policy": accessTokenScopeRule("option:write"),
"PUT /api/option/": accessTokenScopeRule("option:write"),
"PUT /api/option/passkey/domains": accessTokenScopeRule("option:write"),
"GET /api/option/model_pricing": accessTokenScopeRule("option:read"),
"PATCH /api/option/model_pricing": accessTokenScopeRule("option:write"),
"POST /api/option/model_pricing/convert": accessTokenScopeRule("option:read"),
"POST /api/option/model_pricing/preview": accessTokenScopeRule("option:read"),
"POST /api/option/payment_compliance": accessTokenSessionRule,
"GET /api/option/channel_affinity_cache": accessTokenScopeRule("option:read"),
"DELETE /api/option/channel_affinity_cache": accessTokenScopeRule("option:write"),
"POST /api/option/rest_model_ratio": accessTokenScopeRule("option:write"),
"GET /api/option/waffo-pancake/catalog": accessTokenScopeRule("option:read"),
"POST /api/option/waffo-pancake/pair": accessTokenScopeRule("option:write"),
"POST /api/option/waffo-pancake/save": accessTokenScopeRule("option:write"),
"POST /api/option/waffo-pancake/subscription-product": accessTokenScopeRule("option:write"),
"GET /api/option/waffo-pancake/subscription-product-options": accessTokenScopeRule("option:read"),
// router/api-router.go: /api/custom-oauth-provider, /api/ratio_sync
"POST /api/custom-oauth-provider/discovery": accessTokenScopeRule("option:write"),
"GET /api/custom-oauth-provider/": accessTokenScopeRule("option:read"),
"GET /api/custom-oauth-provider/:id": accessTokenScopeRule("option:read"),
"POST /api/custom-oauth-provider/": accessTokenScopeRule("option:write"),
"PUT /api/custom-oauth-provider/:id": accessTokenScopeRule("option:write"),
"DELETE /api/custom-oauth-provider/:id": accessTokenScopeRule("option:write"),
"GET /api/ratio_sync/channels": accessTokenScopeRule("option:read"),
"POST /api/ratio_sync/fetch": accessTokenScopeRule("option:write"),
// router/api-router.go: /api/performance, /api/system-task, /api/system-info
"GET /api/performance/stats": accessTokenScopeRule("ops:read"),
"DELETE /api/performance/disk_cache": accessTokenScopeRule("ops:write"),
"POST /api/performance/reset_stats": accessTokenScopeRule("ops:write"),
"POST /api/performance/gc": accessTokenScopeRule("ops:write"),
"GET /api/performance/logs": accessTokenScopeRule("ops:read"),
"DELETE /api/performance/logs": accessTokenScopeRule("ops:write"),
"POST /api/system-task/log-cleanup": accessTokenScopeRule("ops:write"),
"GET /api/system-task/list": accessTokenScopeRule("ops:read"),
"DELETE /api/system-task/history": accessTokenScopeRule("ops:write"),
"GET /api/system-task/current": accessTokenScopeRule("ops:read"),
"GET /api/system-task/:task_id": accessTokenScopeRule("ops:read"),
"GET /api/system-info/instances": accessTokenScopeRule("ops:read"),
"DELETE /api/system-info/stale-instances": accessTokenScopeRule("ops:write"),
"DELETE /api/system-info/instances/:node_name": accessTokenScopeRule("ops:write"),
// router/api-router.go: /api/plugin/task
"GET /api/plugin/task": accessTokenScopeRule("plugin:read"),
"POST /api/plugin/task": accessTokenScopeRule("plugin:write"),
"PUT /api/plugin/task": accessTokenScopeRule("plugin:write"),
"GET /api/plugin/task/runtime/status": accessTokenScopeRule("plugin:read"),
"GET /api/plugin/task/marketplace/sources": accessTokenScopeRule("plugin:read"),
"PUT /api/plugin/task/marketplace/sources": accessTokenScopeRule("plugin:write"),
"GET /api/plugin/task/:key": accessTokenScopeRule("plugin:read"),
"GET /api/plugin/task/:key/icon": accessTokenScopeRule("plugin:read"),
"GET /api/plugin/task/:key/versions": accessTokenScopeRule("plugin:read"),
"POST /api/plugin/task/:key/activate": accessTokenScopeRule("plugin:write"),
"POST /api/plugin/task/:key/status": accessTokenScopeRule("plugin:write"),
"POST /api/plugin/task/:key/dryrun": accessTokenScopeRule("plugin:write"),
"DELETE /api/plugin/task/:key/versions/:version": accessTokenScopeRule("plugin:write"),
// router/channel-router.go: the key route is RootAuth, not Casbin-guarded.
"POST /api/channel/:id/key": accessTokenScopeRule(service.AccessTokenScopeOf(authz.ChannelSecretView)),
// router/authz-router.go
"GET /api/authz/catalog": accessTokenScopeRule("user:read"),
// router/api-router.go: /api/token
"GET /api/token/": accessTokenScopeRule("api_key:read"),
"GET /api/token/search": accessTokenScopeRule("api_key:read"),
"GET /api/token/auto-groups": accessTokenScopeRule("api_key:read"),
"GET /api/token/:id": accessTokenScopeRule("api_key:read"),
"POST /api/token/:id/key": accessTokenScopeRule("api_key:reveal"),
"POST /api/token/": accessTokenScopeRule("api_key:write"),
"PUT /api/token/": accessTokenScopeRule("api_key:write"),
"DELETE /api/token/:id": accessTokenScopeRule("api_key:write"),
"POST /api/token/batch": accessTokenScopeRule("api_key:write"),
"POST /api/token/batch/keys": accessTokenScopeRule("api_key:reveal"),
// router/api-router.go: /api/redemption
"GET /api/redemption/": accessTokenScopeRule("redemption:read"),
"GET /api/redemption/search": accessTokenScopeRule("redemption:read"),
"GET /api/redemption/:id": accessTokenScopeRule("redemption:read"),
"POST /api/redemption/": accessTokenScopeRule("redemption:write"),
"POST /api/redemption/batch": accessTokenScopeRule("redemption:write"),
"PUT /api/redemption/": accessTokenScopeRule("redemption:write"),
"DELETE /api/redemption/invalid": accessTokenScopeRule("redemption:write"),
"DELETE /api/redemption/:id": accessTokenScopeRule("redemption:write"),
// router/api-router.go: /api/audit, /api/log, /api/data
"GET /api/audit/self": accessTokenScopeRule("usage:read"),
"GET /api/log/": accessTokenScopeRule("log:read"),
"GET /api/log/stat": accessTokenScopeRule("log:read"),
"GET /api/log/self/stat": accessTokenScopeRule("usage:read"),
"GET /api/log/channel_affinity_usage_cache": accessTokenScopeRule("log:read"),
"GET /api/log/search": accessTokenScopeRule("log:read"),
"GET /api/log/self": accessTokenScopeRule("usage:read"),
"GET /api/log/self/search": accessTokenScopeRule("usage:read"),
"GET /api/log/export": accessTokenScopeRule("log:read"),
"GET /api/log/self/export": accessTokenScopeRule("usage:read"),
"GET /api/data/": accessTokenScopeRule("log:read"),
"GET /api/data/users": accessTokenScopeRule("log:read"),
"GET /api/data/self": accessTokenScopeRule("usage:read"),
"GET /api/data/flow": accessTokenScopeRule("log:read"),
"GET /api/data/flow/self": accessTokenScopeRule("usage:read"),
// router/api-router.go: /api/group, /api/prefill_group
"GET /api/group/": accessTokenScopeRule("group:read"),
"GET /api/prefill_group/": accessTokenScopeRule("group:read"),
"POST /api/prefill_group/": accessTokenScopeRule("group:write"),
"PUT /api/prefill_group/": accessTokenScopeRule("group:write"),
"DELETE /api/prefill_group/:id": accessTokenScopeRule("group:write"),
// router/api-router.go: /api/mj, /api/task
"GET /api/mj/self": accessTokenScopeRule("usage:read"),
"GET /api/mj/": accessTokenScopeRule("log:read"),
"GET /api/task/self": accessTokenScopeRule("usage:read"),
"GET /api/task": accessTokenScopeRule("log:read"),
"GET /api/task/:task_id/artifacts": accessTokenScopeRule("usage:read"),
// router/api-router.go: /api/vendors, /api/models (admin)
"POST /api/vendors/operations/preview": accessTokenScopeRule("model:read"),
"POST /api/vendors/operations": accessTokenScopeRule("model:write"),
"GET /api/vendors/": accessTokenScopeRule("model:read"),
"GET /api/vendors/search": accessTokenScopeRule("model:read"),
"GET /api/vendors/:id": accessTokenScopeRule("model:read"),
"POST /api/vendors/": accessTokenScopeRule("model:write"),
"PUT /api/vendors/": accessTokenScopeRule("model:write"),
"DELETE /api/vendors/:id": accessTokenScopeRule("model:write"),
"GET /api/models/sync_upstream/preview": accessTokenScopeRule("model:read"),
"POST /api/models/sync_upstream": accessTokenScopeRule("model:write"),
"POST /api/models/delete": accessTokenScopeRule("model:write"),
"GET /api/models/missing": accessTokenScopeRule("model:read"),
"GET /api/models/": accessTokenScopeRule("model:read"),
"GET /api/models/search": accessTokenScopeRule("model:read"),
"GET /api/models/:id": accessTokenScopeRule("model:read"),
"POST /api/models/": accessTokenScopeRule("model:write"),
"PUT /api/models/": accessTokenScopeRule("model:write"),
"DELETE /api/models/:id": accessTokenScopeRule("model:write"),
// router/api-router.go: /api/deployments
"GET /api/deployments/settings": accessTokenScopeRule("deployment:read"),
"POST /api/deployments/settings/test-connection": accessTokenScopeRule("deployment:write"),
"GET /api/deployments/": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/search": accessTokenScopeRule("deployment:read"),
"POST /api/deployments/test-connection": accessTokenScopeRule("deployment:write"),
"GET /api/deployments/hardware-types": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/locations": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/available-replicas": accessTokenScopeRule("deployment:read"),
"POST /api/deployments/price-estimation": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/check-name": accessTokenScopeRule("deployment:read"),
"POST /api/deployments/": accessTokenScopeRule("deployment:write"),
"GET /api/deployments/:id": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/:id/logs": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/:id/containers": accessTokenScopeRule("deployment:read"),
"GET /api/deployments/:id/containers/:container_id": accessTokenScopeRule("deployment:read"),
"PUT /api/deployments/:id": accessTokenScopeRule("deployment:write"),
"PUT /api/deployments/:id/name": accessTokenScopeRule("deployment:write"),
"POST /api/deployments/:id/extend": accessTokenScopeRule("deployment:write"),
"DELETE /api/deployments/:id": accessTokenScopeRule("deployment:write"),
// router/relay-router.go
"POST /pg/chat/completions": accessTokenSessionRule,
}
var (
permissionRouteRulesMu sync.RWMutex
permissionRouteRules = map[string]accessTokenRouteRule{}
)
// DeclareAccessTokenPermissionRoute records that the Casbin permission guarding
// a route is also the token scope it requires. It runs while routes are
// registered and panics on conflicting declarations.
func DeclareAccessTokenPermissionRoute(method, fullPath string, permission authz.Permission) {
key := method + " " + fullPath
rule := accessTokenScopeRule(service.AccessTokenScopeOf(permission))
if _, ok := accessTokenRouteRules[key]; ok {
panic(fmt.Sprintf("access token route %s is declared twice", key))
}
permissionRouteRulesMu.Lock()
defer permissionRouteRulesMu.Unlock()
if existing, ok := permissionRouteRules[key]; ok && existing != rule {
panic(fmt.Sprintf("access token route %s is declared with %s and %s", key, existing.scope, rule.scope))
}
permissionRouteRules[key] = rule
}
// AccessTokenRouteRule returns the rule for "METHOD /full/path".
func AccessTokenRouteRule(key string) (accessTokenRouteRule, bool) {
if rule, ok := accessTokenRouteRules[key]; ok {
return rule, true
}
permissionRouteRulesMu.RLock()
defer permissionRouteRulesMu.RUnlock()
rule, ok := permissionRouteRules[key]
return rule, ok
}
// AccessTokenRouteRuleKeys lists every declared route key, sorted.
func AccessTokenRouteRuleKeys() []string {
permissionRouteRulesMu.RLock()
keys := slices.Collect(maps.Keys(permissionRouteRules))
permissionRouteRulesMu.RUnlock()
keys = slices.AppendSeq(keys, maps.Keys(accessTokenRouteRules))
slices.Sort(keys)
return keys
}
// Scope returns the token scope the rule requires, or "" for any/session rules.
func (rule accessTokenRouteRule) Scope() string {
return rule.scope
}
// Kind names the rule category: "scope", "any" or "session".
func (rule accessTokenRouteRule) Kind() string {
switch rule.kind {
case accessTokenRuleScope:
return "scope"
case accessTokenRuleAny:
return "any"
case accessTokenRuleSession:
return "session"
}
return ""
}
+18 -5
View File
@@ -268,9 +268,8 @@ func AccessTokenAudit() gin.HandlerFunc {
if present {
_, internal, _ := service.ParseDashboardAccessToken(raw)
if !internal {
user, err := model.ValidateAccessToken(raw)
if err == nil && user != nil && user.Id > 0 {
beginAccessTokenAudit(c, user, raw)
if lookup := lookupAccessToken(c, raw); lookup.user != nil {
beginAccessTokenAudit(c, lookup.user, lookup.ref)
defer finishAccessTokenAudit(c)
}
}
@@ -279,18 +278,32 @@ func AccessTokenAudit() gin.HandlerFunc {
}
}
func beginAccessTokenAudit(c *gin.Context, user *model.User, token string) {
func beginAccessTokenAudit(c *gin.Context, user *model.UserBase, tokenRef string) {
if _, exists := c.Get(accessTokenAuditContextKey); exists {
return
}
writer := &auditResponseWriter{ResponseWriter: c.Writer, body: bytes.NewBuffer(nil), maxSize: 64 * 1024}
c.Writer = writer
c.Set(accessTokenAuditContextKey, &accessTokenRequestAudit{
entry: model.AuditLog{UserId: user.Id, Username: user.Username, ActorRole: user.Role, Category: model.AuditCategoryAccessToken, AuthMethod: "access_token", TokenRef: model.AccessTokenFingerprint(token), CreatedAt: common.GetTimestamp(), EventId: common.NewRequestId()},
entry: model.AuditLog{UserId: user.Id, Username: user.Username, ActorRole: user.Role, Category: model.AuditCategoryAccessToken, AuthMethod: "access_token", TokenRef: tokenRef, CreatedAt: common.GetTimestamp(), EventId: common.NewRequestId()},
writer: writer,
})
}
// setAccessTokenAuditParam adds non-secret context, such as why a token was
// rejected, to the current request's access audit entry.
func setAccessTokenAuditParam(c *gin.Context, key string, value any) {
stored, _ := c.Get(accessTokenAuditContextKey)
audit, ok := stored.(*accessTokenRequestAudit)
if !ok {
return
}
if audit.entry.Other.Op == nil {
audit.entry.Other.Op = &model.AuditOperation{Params: model.AuditFields{}}
}
audit.entry.Other.Op.Params[key] = value
}
func finishAccessTokenAudit(c *gin.Context) {
value, exists := c.Get(accessTokenAuditContextKey)
if !exists {
+228 -61
View File
@@ -23,7 +23,12 @@ import (
"gorm.io/gorm"
)
const authIdentityContextKey = "auth_identity"
const (
authIdentityContextKey = "auth_identity"
accessTokenLookupContextKey = "access_token_lookup"
accessTokenIdentityContextKey = "access_token_identity"
accessTokenSessionRequiredMsg = "a dashboard login session is required"
)
type dashboardCredentialKind int
@@ -54,18 +59,28 @@ func authHelper(c *gin.Context, minRole int) {
return
}
if user.Status != common.UserStatusEnabled {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_USER_DISABLED", "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned)})
abortWithWebMessage(c, http.StatusUnauthorized, "AUTH_USER_DISABLED", common.NewMessage("User has been banned"))
return
}
if user.Role < minRole {
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"success": false, "code": "AUTH_INSUFFICIENT_PRIVILEGE", "message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege)})
abortWithWebMessage(c, http.StatusForbidden, "AUTH_INSUFFICIENT_PRIVILEGE", common.NewMessage("Unauthorized, insufficient privileges"))
return
}
var lookup *accessTokenLookup
if useAccessToken {
lookup = requestAccessTokenLookup(c)
if !enforceAccessTokenRoute(c, lookup, true) {
return
}
}
if !validUserInfo(user.Username, user.Role) {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_USER_INVALID", "message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid)})
abortWithWebMessage(c, http.StatusUnauthorized, "AUTH_USER_INVALID", common.NewMessage("Unauthorized, invalid user info"))
return
}
setDashboardAuthContext(c, user, identity, useAccessToken)
if lookup != nil {
setAccessTokenContext(c, lookup)
}
// 管理/root 写操作审计兜底:内聚在鉴权链路里,保证任何经过 AdminAuth/RootAuth
// 的写接口都会自动留痕(无需在路由上单独挂审计中间件,避免漏挂)。
@@ -86,12 +101,29 @@ func TryUserAuth() func(c *gin.Context) {
defer finishAccessTokenAudit(c)
}
user, identity, credentialKind, err := classifyDashboardCredential(c)
if errors.Is(err, model.ErrLegacyAccessTokenRetired) {
// After the deadline no opaque value is looked up, so a retired
// token is indistinguishable from any unknown value: anonymous.
credentialKind, err = dashboardCredentialUnmatched, nil
}
if err != nil {
writeDashboardAuthError(c, err)
return
}
// Optional authentication never elevates a token beyond its grant: a
// token the route would reject leaves the request anonymous.
var lookup *accessTokenLookup
if credentialKind == dashboardCredentialPAT {
lookup = requestAccessTokenLookup(c)
if user.Status != common.UserStatusEnabled || !enforceAccessTokenRoute(c, lookup, false) {
credentialKind = dashboardCredentialUnmatched
}
}
if credentialKind != dashboardCredentialUnmatched {
setDashboardAuthContext(c, user, identity, credentialKind == dashboardCredentialPAT)
if credentialKind == dashboardCredentialPAT {
setAccessTokenContext(c, lookup)
}
}
c.Next()
}
@@ -144,6 +176,21 @@ func GetSessionAuthIdentity(c *gin.Context) (service.AuthIdentity, bool) {
return identity, true
}
// GetStepUpIdentity returns the identity that may request and consume
// security verification proofs: a browser session, or a scoped access token
// bound to its own pat:<id> session. Legacy access tokens have neither.
func GetStepUpIdentity(c *gin.Context) (service.AuthIdentity, bool) {
if identity, ok := GetSessionAuthIdentity(c); ok {
return identity, true
}
value, ok := c.Get(accessTokenIdentityContextKey)
if !ok {
return service.AuthIdentity{}, false
}
identity, ok := value.(service.AuthIdentity)
return identity, ok && identity.UserID > 0 && identity.SessionID != ""
}
func authenticateDashboardRequest(c *gin.Context) (*model.UserBase, service.AuthIdentity, bool, error) {
user, identity, credentialKind, err := classifyDashboardCredential(c)
if err != nil {
@@ -171,19 +218,128 @@ func classifyDashboardCredential(c *gin.Context) (*model.UserBase, service.AuthI
}
return user, identity, dashboardCredentialInternal, nil
}
patUser, err := model.ValidateAccessToken(raw)
if err != nil {
return nil, service.AuthIdentity{}, dashboardCredentialPAT, err
lookup := lookupAccessToken(c, raw)
if lookup.user != nil {
beginAccessTokenAudit(c, lookup.user, lookup.ref)
}
if patUser == nil || patUser.Id <= 0 {
if lookup.err != nil {
if errors.Is(lookup.err, model.ErrAccessTokenExpired) {
setAccessTokenAuditParam(c, "failure_reason", "expired")
}
return nil, service.AuthIdentity{}, dashboardCredentialPAT, lookup.err
}
if lookup.user == nil {
return nil, service.AuthIdentity{}, dashboardCredentialUnmatched, nil
}
beginAccessTokenAudit(c, patUser, raw)
user, err := model.GetUserCache(patUser.Id)
if err != nil {
return nil, service.AuthIdentity{}, dashboardCredentialPAT, err
return lookup.user, service.AuthIdentity{UserID: lookup.user.Id, UserAuthVersion: lookup.user.AuthVersion}, dashboardCredentialPAT, nil
}
// accessTokenLookup is the per-request result of resolving an opaque dashboard
// credential. A nil user with a nil err means the value is not a PAT.
type accessTokenLookup struct {
ref string
legacy bool
token *model.UserAccessToken
user *model.UserBase
err error
}
// lookupAccessToken resolves raw once per request, so the access audit and
// authentication share one database lookup.
func lookupAccessToken(c *gin.Context, raw string) *accessTokenLookup {
ref := model.AccessTokenFingerprint(raw)
if cached := requestAccessTokenLookup(c); cached != nil && cached.ref == ref {
return cached
}
lookup := &accessTokenLookup{ref: ref}
c.Set(accessTokenLookupContextKey, lookup)
if !strings.HasPrefix(raw, model.AccessTokenPrefix) {
lookup.legacy = true
legacyUser, err := model.ValidateAccessToken(raw)
if err != nil || legacyUser == nil || legacyUser.Id <= 0 {
lookup.err = err
return lookup
}
lookup.user, lookup.err = model.GetUserCache(legacyUser.Id)
return lookup
}
token, err := model.FindUserAccessTokenByHash(ref)
if err != nil || token == nil {
lookup.err = err
return lookup
}
lookup.token = token
lookup.user, lookup.err = model.GetUserCache(token.UserId)
if lookup.err == nil && token.Expired(common.GetTimestamp()) {
// Keep the owner so the rejected attempt is still audited.
lookup.err = model.ErrAccessTokenExpired
}
return lookup
}
func requestAccessTokenLookup(c *gin.Context) *accessTokenLookup {
value, _ := c.Get(accessTokenLookupContextKey)
lookup, _ := value.(*accessTokenLookup)
return lookup
}
// enforceAccessTokenRoute applies the route's access token rule (fail closed:
// scoped tokens cannot call undeclared routes). With abort set it writes the
// 403 response and records the reason in the access audit.
func enforceAccessTokenRoute(c *gin.Context, lookup *accessTokenLookup, abort bool) bool {
key := c.Request.Method + " " + c.FullPath()
rule, declared := AccessTokenRouteRule(key)
code, reason := "", ""
var message *common.Message
switch {
case declared && rule.kind == accessTokenRuleSession:
code, reason, message = "AUTH_SESSION_REQUIRED", "session_required", common.NewMessage(accessTokenSessionRequiredMsg)
case lookup.legacy:
return true
case !declared:
common.SysError(common.LogText("access token route is not declared: %s", key))
code, reason, message = "ACCESS_TOKEN_ROUTE_UNDECLARED", "route_undeclared", common.NewMessage("Unauthorized, insufficient privileges")
case rule.kind == accessTokenRuleAny, service.AccessTokenScopeGranted(lookup.token.GetScopes(), rule.scope):
return true
default:
code, reason, message = "ACCESS_TOKEN_SCOPE_DENIED", "scope_denied", common.NewMessage("The access token lacks the required permission: {{scope}}", map[string]any{"scope": rule.scope})
if abort {
setAccessTokenAuditParam(c, "required_scope", rule.scope)
}
}
if abort {
setAccessTokenAuditParam(c, "failure_reason", reason)
abortWithWebMessage(c, http.StatusForbidden, code, message)
}
return false
}
// setAccessTokenContext exposes the token grant to handlers. Scoped tokens
// also get a step-up identity bound to the token instead of a browser session.
func setAccessTokenContext(c *gin.Context, lookup *accessTokenLookup) {
_, alreadySet := c.Get(accessTokenIdentityContextKey)
c.Set("access_token_legacy", lookup.legacy)
// Security audit params identify the token the same way the access audit does.
c.Set("access_token_ref", lookup.ref)
if lookup.legacy {
c.Set("access_token_id", 0)
c.Set("access_token_scopes", []string{})
return
}
c.Set("access_token_id", lookup.token.Id)
c.Set("access_token_scopes", lookup.token.GetScopes())
c.Set(accessTokenIdentityContextKey, service.AuthIdentity{
UserID: lookup.user.Id,
SessionID: model.AccessTokenSessionID(lookup.token.Id),
UserAuthVersion: lookup.user.AuthVersion,
SessionVersion: model.AccessTokenSessionVersion,
})
if alreadySet {
return
}
if err := model.TouchUserAccessToken(lookup.token.Id, c.ClientIP(), common.GetTimestamp()); err != nil {
logger.LogError(c.Request.Context(), common.LogText("access token last-use update failed (token_id=%d): %v", lookup.token.Id, err))
}
return user, service.AuthIdentity{UserID: user.Id, UserAuthVersion: user.AuthVersion}, dashboardCredentialPAT, nil
}
func authorizationToken(header string) (string, bool) {
@@ -215,21 +371,40 @@ func setDashboardAuthContext(c *gin.Context, user *model.UserBase, identity serv
user.WriteContext(c)
}
// abortWithWebMessage answers a web console request; the web console
// translates the message, and code, when set, stays a stable error code.
func abortWithWebMessage(c *gin.Context, status int, code string, message *common.Message) {
var fields []gin.H
if code != "" {
fields = append(fields, gin.H{"code": code})
}
common.ApiErrorStatus(c, status, message, fields...)
c.Abort()
}
func writeDashboardAuthError(c *gin.Context, err error) {
if errors.Is(err, model.ErrAccessTokenExpired) {
abortWithWebMessage(c, http.StatusUnauthorized, "ACCESS_TOKEN_EXPIRED", common.NewMessage("Unauthorized, the access token has expired"))
return
}
if errors.Is(err, model.ErrLegacyAccessTokenRetired) {
abortWithWebMessage(c, http.StatusUnauthorized, "ACCESS_TOKEN_LEGACY_RETIRED", common.NewMessage("Legacy access tokens have been retired. Create a new access token in security settings"))
return
}
if errors.Is(err, service.ErrAuthTokenExpired) {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_TOKEN_EXPIRED", "message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn)})
abortWithWebMessage(c, http.StatusUnauthorized, "AUTH_TOKEN_EXPIRED", common.NewMessage("Unauthorized, not logged in and no access token provided"))
return
}
if errors.Is(err, service.ErrLoginSessionRevoked) || errors.Is(err, gorm.ErrRecordNotFound) {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_SESSION_REVOKED", "message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn)})
abortWithWebMessage(c, http.StatusUnauthorized, "AUTH_SESSION_REVOKED", common.NewMessage("Unauthorized, not logged in and no access token provided"))
return
}
if errors.Is(err, service.ErrAuthTokenInvalid) {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_UNAUTHORIZED", "message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid)})
abortWithWebMessage(c, http.StatusUnauthorized, "AUTH_UNAUTHORIZED", common.NewMessage("Unauthorized, invalid access token"))
return
}
common.SysLog("dashboard authentication error: " + err.Error())
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"success": false, "code": "AUTH_INTERNAL_ERROR", "message": common.TranslateMessage(c, i18n.MsgDatabaseError)})
common.SysLog(common.LogText("dashboard authentication error: %s", err.Error()))
abortWithWebMessage(c, http.StatusInternalServerError, "AUTH_INTERNAL_ERROR", common.NewMessage("Database error, please contact the administrator"))
}
func RequirePermission(permission authz.Permission) func(c *gin.Context) {
@@ -240,11 +415,7 @@ func RequirePermission(permission authz.Permission) func(c *gin.Context) {
c.Next()
return
}
c.JSON(http.StatusForbidden, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege),
})
c.Abort()
abortWithWebMessage(c, http.StatusForbidden, "", common.NewMessage("Unauthorized, insufficient privileges"))
}
}
@@ -282,17 +453,14 @@ func TokenOrUserAuth() func(c *gin.Context) {
}
// TokenAuthReadOnly 宽松版本的令牌认证中间件,用于只读查询接口。
// 只验证令牌 key 是否存在,不检查令牌状态、过期时间和额度。
// 即使令牌已过期、已耗尽或已禁用,也允许访问。
// 仍然检查用户是否被封禁。
// 只验证令牌 key 是否存在,不检查过期时间和额度。
// 即使令牌已过期或已耗尽,也允许访问;已禁用的令牌会被拒绝。
// 仍然检查令牌 IP 白名单和用户是否被封禁。
func TokenAuthReadOnly() func(c *gin.Context) {
return func(c *gin.Context) {
key := c.Request.Header.Get("Authorization")
if key == "" {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgTokenNotProvided),
})
common.ApiErrorStatus(c, http.StatusUnauthorized, errors.New(common.TranslateMessage(c, i18n.MsgTokenNotProvided)))
c.Abort()
return
}
@@ -306,16 +474,10 @@ func TokenAuthReadOnly() func(c *gin.Context) {
token, err := model.GetTokenByKey(key, false)
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgTokenInvalid),
})
common.ApiErrorStatus(c, http.StatusUnauthorized, errors.New(common.TranslateMessage(c, i18n.MsgTokenInvalid)))
} else {
common.SysLog("TokenAuthReadOnly GetTokenByKey database error: " + err.Error())
c.JSON(http.StatusInternalServerError, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
})
common.SysLog(common.LogText("TokenAuthReadOnly GetTokenByKey database error: %s", err.Error()))
common.ApiErrorStatus(c, http.StatusInternalServerError, errors.New(common.TranslateMessage(c, i18n.MsgDatabaseError)))
}
c.Abort()
return
@@ -324,29 +486,34 @@ func TokenAuthReadOnly() func(c *gin.Context) {
// TokenAuthReadOnly must keep allowing other token states to query read-only
// data, such as token usage logs; only explicitly disabled tokens are denied.
if token.Status == common.TokenStatusDisabled {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgTokenStatusUnavailable),
})
common.ApiErrorStatus(c, http.StatusUnauthorized, errors.New(common.TranslateMessage(c, i18n.MsgTokenStatusUnavailable)))
c.Abort()
return
}
if allowIps := token.GetIpLimits(); len(allowIps) > 0 {
ip := net.ParseIP(c.ClientIP())
if ip == nil {
common.ApiErrorStatus(c, http.StatusForbidden, errors.New(common.TranslateMessage(c, i18n.MsgAuthClientIPUnresolved)))
c.Abort()
return
}
if !common.IsIpInCIDRList(ip, allowIps) {
common.ApiErrorStatus(c, http.StatusForbidden, errors.New(common.TranslateMessage(c, i18n.MsgAuthTokenIPNotAllowed)))
c.Abort()
return
}
}
userCache, err := model.GetUserCache(token.UserId)
if err != nil {
common.SysLog(fmt.Sprintf("TokenAuthReadOnly GetUserCache error for user %d: %v", token.UserId, err))
c.JSON(http.StatusInternalServerError, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
})
common.SysLog(common.LogText("TokenAuthReadOnly GetUserCache error for user %d: %v", token.UserId, err))
common.ApiErrorStatus(c, http.StatusInternalServerError, errors.New(common.TranslateMessage(c, i18n.MsgDatabaseError)))
c.Abort()
return
}
if userCache.Status != common.UserStatusEnabled {
c.JSON(http.StatusForbidden, gin.H{
"success": false,
"message": common.TranslateMessage(c, i18n.MsgAuthUserBanned),
})
common.ApiErrorStatus(c, http.StatusForbidden, errors.New(common.TranslateMessage(c, i18n.MsgAuthUserBanned)))
c.Abort()
return
}
@@ -411,7 +578,7 @@ func TokenAuth() func(c *gin.Context) {
}
if err != nil {
if errors.Is(err, model.ErrDatabase) {
common.SysLog("TokenAuth ValidateUserToken database error: " + err.Error())
common.SysLog(common.LogText("TokenAuth ValidateUserToken database error: %s", err.Error()))
abortWithOpenAiMessage(c, http.StatusInternalServerError,
common.TranslateMessage(c, i18n.MsgDatabaseError))
} else {
@@ -427,11 +594,11 @@ func TokenAuth() func(c *gin.Context) {
logger.LogDebug(c, "Token has IP restrictions, checking client IP %s", clientIp)
ip := net.ParseIP(clientIp)
if ip == nil {
abortWithOpenAiMessage(c, http.StatusForbidden, "无法解析客户端 IP 地址")
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgAuthClientIPUnresolved))
return
}
if common.IsIpInCIDRList(ip, allowIps) == false {
abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgAuthTokenIPNotAllowed), types.ErrorCodeAccessDenied)
return
}
logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)
@@ -439,7 +606,7 @@ func TokenAuth() func(c *gin.Context) {
userCache, err := model.GetUserCache(token.UserId)
if err != nil {
common.SysLog(fmt.Sprintf("TokenAuth GetUserCache error for user %d: %v", token.UserId, err))
common.SysLog(common.LogText("TokenAuth GetUserCache error for user %d: %v", token.UserId, err))
abortWithOpenAiMessage(c, http.StatusInternalServerError,
common.TranslateMessage(c, i18n.MsgDatabaseError))
return
@@ -462,13 +629,13 @@ func TokenAuth() func(c *gin.Context) {
if tokenGroup != "" && token.RoutingPolicy == "" {
// check common.UserUsableGroups[userGroup]
if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("无权访问 %s 分组", tokenGroup))
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgAuthTokenGroupForbidden, map[string]any{"Group": tokenGroup}))
return
}
// check group in common.GroupRatio
if !ratio_setting.ContainsGroupRatio(tokenGroup) {
if tokenGroup != "auto" {
abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgAuthTokenGroupDeprecated, map[string]any{"Group": tokenGroup}))
return
}
}
@@ -542,7 +709,7 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e
if token.AutoGroups != "" {
groupPool, err := token.GetGroupPool()
if err != nil {
common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err))
common.SysError(common.LogText("failed to parse auto groups for token %d: %v", token.Id, err))
common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, []string{})
} else if len(groupPool) > 0 {
autoGroups := make([]string, 0, len(groupPool))
@@ -568,8 +735,8 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e
})
} else {
c.Header("specific_channel_version", "701e3ae1dc3f7975556d354e0675168d004891c8")
abortWithOpenAiMessage(c, http.StatusForbidden, "普通用户不支持指定渠道")
return fmt.Errorf("普通用户不支持指定渠道")
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgAuthChannelSelectForbidden))
return errors.New("regular users cannot specify a channel")
}
}
return nil
+201 -1
View File
@@ -7,12 +7,15 @@ import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/service/authz"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/golang-jwt/jwt/v5"
@@ -30,12 +33,13 @@ func setupDashboardAuthMiddlewareTest(t *testing.T) {
previousSecret := common.SessionSecret
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.AuditLog{}))
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.AuditLog{}, &model.UserAccessToken{}, &model.Option{}))
model.DB = db
model.LOG_DB = db
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
common.RedisEnabled = false
common.SessionSecret = "middleware-auth-test-secret"
require.NoError(t, model.EnsureLegacyAccessTokenRetireAt(time.Now().Unix()))
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLogDB
@@ -87,6 +91,29 @@ func createMiddlewarePATUser(t *testing.T, username, token string) *model.User {
return user
}
// Read-only API key routes (/api/usage/token, /api/log/token) are called by
// scripts, not the web console, so their errors are translated by the backend
// in the language the request asks for.
func TestTokenAuthReadOnlyTranslatesErrors(t *testing.T) {
require.NoError(t, i18n.Init())
for _, tc := range []struct {
acceptLanguage string
want string
}{
{"zh-CN", "未提供令牌"},
{"en-US", "Token not provided"},
} {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/log/token", nil)
c.Request.Header.Set("Accept-Language", tc.acceptLanguage)
TokenAuthReadOnly()(c)
assert.True(t, c.IsAborted(), tc.acceptLanguage)
assert.Equal(t, http.StatusUnauthorized, recorder.Code, tc.acceptLanguage)
assert.JSONEq(t, `{"success":false,"message":"`+tc.want+`"}`, recorder.Body.String(), tc.acceptLanguage)
}
}
func TestUserAuthAllowsOpaqueDottedPAT(t *testing.T) {
setupDashboardAuthMiddlewareTest(t)
user := createMiddlewarePATUser(t, "dotted-pat-user", "opaque.key.with-dots")
@@ -135,6 +162,7 @@ func TestTryUserAuthCredentialClassification(t *testing.T) {
patUser := createMiddlewarePATUser(t, "optional-pat-user", "optional.pat.with-dots")
internalUser := createMiddlewarePATUser(t, "optional-session-user", "unrelated-pat")
scopedPAT, _ := createMiddlewareScopedToken(t, patUser.Id, 0, "profile:read")
now := time.Now().Unix()
session := &model.UserSession{
SID: "optional-auth-session",
@@ -189,6 +217,7 @@ func TestTryUserAuthCredentialClassification(t *testing.T) {
{name: "dotted unmatched credential", token: "ordinary.key.with-dots", wantStatus: http.StatusOK},
{name: "third party jwt", token: externalToken, wantStatus: http.StatusOK},
{name: "valid pat", token: "optional.pat.with-dots", wantStatus: http.StatusOK, wantUserID: patUser.Id, wantPAT: true},
{name: "scoped pat outside its grant stays anonymous", token: scopedPAT, wantStatus: http.StatusOK},
{name: "valid internal access jwt", token: accessToken, wantStatus: http.StatusOK, wantUserID: internalUser.Id},
{name: "expired internal access jwt", token: issueExpiredDashboardAccessToken(t, identity), wantStatus: http.StatusUnauthorized, wantErrorCode: "AUTH_TOKEN_EXPIRED"},
{name: "tampered internal access jwt", token: tamperDashboardToken(accessToken), wantStatus: http.StatusUnauthorized, wantErrorCode: "AUTH_UNAUTHORIZED"},
@@ -256,6 +285,177 @@ func TestTryUserAuthCredentialClassification(t *testing.T) {
assert.Contains(t, databaseFailureResponse.Body.String(), "AUTH_INTERNAL_ERROR")
}
func createMiddlewareScopedToken(t *testing.T, userID int, expiresAt int64, scopes ...string) (string, *model.UserAccessToken) {
t.Helper()
suffix, err := common.GenerateRandomCharsKey(43)
require.NoError(t, err)
raw := model.AccessTokenPrefix + suffix
token := &model.UserAccessToken{Name: "middleware token", TokenHash: model.AccessTokenFingerprint(raw), TokenHint: model.AccessTokenHint(raw), ExpiresAt: expiresAt}
require.NoError(t, token.SetScopes(scopes))
require.NoError(t, model.CreateUserAccessToken(userID, token, service.AccessTokenMaxPerUser))
return raw, token
}
func middlewareBearerRequest(router *gin.Engine, path, token string) *httptest.ResponseRecorder {
request := httptest.NewRequest(http.MethodGet, path, nil)
request.Header.Set("Authorization", "Bearer "+token)
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
return response
}
func TestUserAuthAppliesAccessTokenRouteRules(t *testing.T) {
setupDashboardAuthMiddlewareTest(t)
gin.SetMode(gin.TestMode)
wasMaster := common.IsMasterNode
common.IsMasterNode = true
t.Cleanup(func() { common.IsMasterNode = wasMaster })
require.NoError(t, model.DB.AutoMigrate(&model.CasbinRule{}, &model.AuthzRole{}))
require.NoError(t, authz.Init(model.DB))
// The channel router declares this rule when it registers the route.
DeclareAccessTokenPermissionRoute(http.MethodGet, "/api/channel/", authz.ChannelRead)
legacy := "middleware-legacy-token"
admin := createMiddlewarePATUser(t, "route-rule-admin", legacy)
require.NoError(t, model.DB.Model(admin).Update("role", common.RoleAdminUser).Error)
profile, _ := createMiddlewareScopedToken(t, admin.Id, 0, "profile:read")
channel, _ := createMiddlewareScopedToken(t, admin.Id, time.Now().Unix()+3600, "channel:read")
expired, _ := createMiddlewareScopedToken(t, admin.Id, time.Now().Unix()-1, "profile:read")
ok := func(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"id": c.GetInt("id")}) }
router := gin.New()
router.GET("/api/user/self", UserAuth(), ok)
router.GET("/api/user/access_tokens", UserAuth(), ok)
router.GET("/api/undeclared", UserAuth(), ok)
router.GET("/api/channel/", AdminAuth(), RequirePermission(authz.ChannelRead), ok)
router.GET("/api/pricing", TryUserAuth(), ok)
for _, test := range []struct {
name, path, token, code, reason string
status int
}{
{name: "granted scope on a never expiring token", path: "/api/user/self", token: profile, status: http.StatusOK},
{name: "missing scope", path: "/api/user/self", token: channel, status: http.StatusForbidden, code: "ACCESS_TOKEN_SCOPE_DENIED", reason: "scope_denied"},
{name: "undeclared route", path: "/api/undeclared", token: profile, status: http.StatusForbidden, code: "ACCESS_TOKEN_ROUTE_UNDECLARED", reason: "route_undeclared"},
{name: "session route", path: "/api/user/access_tokens", token: profile, status: http.StatusForbidden, code: "AUTH_SESSION_REQUIRED", reason: "session_required"},
{name: "permission route with its scope", path: "/api/channel/", token: channel, status: http.StatusOK},
{name: "permission route without its scope", path: "/api/channel/", token: profile, status: http.StatusForbidden, code: "ACCESS_TOKEN_SCOPE_DENIED", reason: "scope_denied"},
{name: "expired token", path: "/api/user/self", token: expired, status: http.StatusUnauthorized, code: "ACCESS_TOKEN_EXPIRED", reason: "expired"},
{name: "legacy token before the deadline", path: "/api/user/self", token: legacy, status: http.StatusOK},
{name: "legacy token on a session route", path: "/api/user/access_tokens", token: legacy, status: http.StatusForbidden, code: "AUTH_SESSION_REQUIRED", reason: "session_required"},
} {
t.Run(test.name, func(t *testing.T) {
response := middlewareBearerRequest(router, test.path, test.token)
assert.Equal(t, test.status, response.Code, response.Body.String())
if test.code != "" {
assert.Contains(t, response.Body.String(), `"code":"`+test.code+`"`)
}
var audit model.AuditLog
require.NoError(t, model.LOG_DB.Where("category = ?", model.AuditCategoryAccessToken).Order("id desc").First(&audit).Error)
assert.Equal(t, model.AccessTokenFingerprint(test.token), audit.TokenRef)
assert.Equal(t, test.status, audit.Status)
if test.reason == "" {
return
}
require.NotNil(t, audit.Other.Op)
params, err := common.Marshal(audit.Other.Op.Params)
require.NoError(t, err)
assert.Contains(t, string(params), `"failure_reason":"`+test.reason+`"`)
if test.code == "ACCESS_TOKEN_SCOPE_DENIED" {
assert.Contains(t, string(params), `"required_scope":"`)
}
})
}
// The token grant never widens the account: revoking the Casbin permission
// rejects a token that still carries the scope.
require.NoError(t, authz.SetUserPermissions(admin.Id, authz.PermissionsMap{authz.ResourceChannel: {authz.ActionRead: false}}))
response := middlewareBearerRequest(router, "/api/channel/", channel)
assert.Equal(t, http.StatusForbidden, response.Code)
assert.NotContains(t, response.Body.String(), "ACCESS_TOKEN_SCOPE_DENIED")
require.NoError(t, model.DB.Model(admin).Update("status", common.UserStatusDisabled).Error)
response = middlewareBearerRequest(router, "/api/user/self", profile)
assert.Equal(t, http.StatusUnauthorized, response.Code)
assert.Contains(t, response.Body.String(), `"code":"AUTH_USER_DISABLED"`)
require.NoError(t, model.DB.Model(admin).Update("status", common.UserStatusEnabled).Error)
assert.Equal(t, http.StatusOK, middlewareBearerRequest(router, "/api/user/self", profile).Code, "disabling a user keeps the token")
require.NoError(t, model.DB.Model(&model.Option{}).Where(&model.Option{Key: "LegacyAccessTokenRetireAt"}).Update("value", fmt.Sprint(time.Now().Unix()-1)).Error)
require.NoError(t, model.EnsureLegacyAccessTokenRetireAt(time.Now().Unix()))
response = middlewareBearerRequest(router, "/api/user/self", legacy)
assert.Equal(t, http.StatusUnauthorized, response.Code)
assert.Contains(t, response.Body.String(), `"code":"ACCESS_TOKEN_LEGACY_RETIRED"`)
response = middlewareBearerRequest(router, "/api/pricing", legacy)
assert.Equal(t, http.StatusOK, response.Code)
assert.JSONEq(t, `{"id":0}`, response.Body.String(), "optional authentication treats a retired token as anonymous")
assert.Equal(t, http.StatusOK, middlewareBearerRequest(router, "/api/user/self", profile).Code, "the deadline only retires legacy tokens")
}
func TestAccessTokenIdentityAndSingleLookup(t *testing.T) {
setupDashboardAuthMiddlewareTest(t)
gin.SetMode(gin.TestMode)
legacy := "middleware-identity-legacy"
user := createMiddlewarePATUser(t, "identity-user", legacy)
scoped, token := createMiddlewareScopedToken(t, user.Id, 0, "profile:read")
router := gin.New()
router.Use(AccessTokenAudit())
router.GET("/api/user/self", UserAuth(), func(c *gin.Context) {
_, session := GetSessionAuthIdentity(c)
identity, stepUp := GetStepUpIdentity(c)
c.JSON(http.StatusOK, gin.H{"session": session, "step_up": stepUp, "identity": identity})
})
var queries []string
const callbackName = "test:access-token-lookup-count"
require.NoError(t, model.DB.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) {
queries = append(queries, tx.Statement.Table+" "+tx.Statement.SQL.String())
}))
t.Cleanup(func() { model.DB.Callback().Query().Remove(callbackName) })
countQueries := func(table, fragment string) int {
count := 0
for _, query := range queries {
if strings.HasPrefix(query, table+" ") && strings.Contains(query, fragment) {
count++
}
}
return count
}
for _, test := range []struct {
name, token string
stepUp bool
tokenTable int
legacyQuery int
}{
{name: "scoped token", token: scoped, stepUp: true, tokenTable: 1},
{name: "legacy token", token: legacy, legacyQuery: 1},
} {
t.Run(test.name, func(t *testing.T) {
queries = nil
response := middlewareBearerRequest(router, "/api/user/self", test.token)
require.Equal(t, http.StatusOK, response.Code, response.Body.String())
var body struct {
Session bool `json:"session"`
StepUp bool `json:"step_up"`
Identity service.AuthIdentity `json:"identity"`
}
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &body))
assert.False(t, body.Session, "an access token is never a browser session")
assert.Equal(t, test.stepUp, body.StepUp)
if test.stepUp {
assert.Equal(t, service.AuthIdentity{UserID: user.Id, SessionID: model.AccessTokenSessionID(token.Id), UserAuthVersion: user.AuthVersion, SessionVersion: model.AccessTokenSessionVersion}, body.Identity)
}
// The access audit and authentication share one lookup: the
// credential query plus the owner's user row.
assert.Len(t, queries, 2)
assert.Equal(t, test.tokenTable, countQueries("user_access_tokens", "token_hash"))
assert.Equal(t, test.legacyQuery, countQueries("users", "access_token = "))
})
}
}
func TestAPIKeyFromWebSocketSubprotocol(t *testing.T) {
tests := []struct {
name string
+39 -24
View File
@@ -57,7 +57,7 @@ func Distribute() func(c *gin.Context) {
service.AppendTaskPluginIdentityFilter(c, c.GetString("expected_task_plugin_key"))
modelRequest, shouldSelectChannel, err := getModelRequest(c)
if err != nil {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
abortWithOpenAiMessage(c, invalidRequestStatus(err), i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return
}
requestedStream = modelRequest.Stream || strings.Contains(c.Request.URL.Path, "streamGenerateContent") || strings.EqualFold(c.GetHeader("Upgrade"), "websocket")
@@ -130,14 +130,7 @@ func Distribute() func(c *gin.Context) {
message = i18n.T(c, selectErr.MessageID, selectErr.Params)
}
if service.RequestRoutingPolicy(c) != nil && selectErr.NoAvailableChannel {
if c.GetBool("pool_degradation_unavailable") {
message = "Group quality checks are temporarily unavailable. Please retry later."
} else if wait := c.GetInt("pool_retry_after"); wait > 0 {
message = fmt.Sprintf("No eligible group is currently available; failed groups are cooling down or being checked. Retry after %d seconds.", wait)
c.Header("Retry-After", fmt.Sprint(wait))
} else {
message = "No eligible group is available for this model under the current pool and quality settings."
}
message = noAvailableGroupPoolMessage(c)
}
rejectionMessage = message
abortWithOpenAiMessage(c, selectErr.StatusCode, message, selectErr.Code)
@@ -157,6 +150,13 @@ func Distribute() func(c *gin.Context) {
}
}
func noAvailableGroupPoolMessage(c *gin.Context) string {
if wait := c.GetInt("pool_retry_after"); wait > 0 {
c.Header("Retry-After", fmt.Sprint(wait))
}
return "分组池中无条件内可用分组"
}
// noAvailableChannelMessage explains a 503 for a task-plugin-claimed model.
// The response tells the caller the model is plugin-claimed without naming the
// plugin; the candidate plugin keys go to the server log under the request id.
@@ -175,7 +175,7 @@ func noAvailableChannelMessage(c *gin.Context, group, modelName string) string {
}
}
}
logger.LogWarn(c, "task_plugin subsystem=distribution event=no_available_channel group=%q model=%q plugins=%q reason=no_eligible_channel", group, modelName, strings.Join(keys, ","))
logger.LogWarn(c, common.LogText("task_plugin subsystem=distribution event=no_available_channel group=%q model=%q plugins=%q reason=no_eligible_channel", group, modelName, strings.Join(keys, ",")))
return i18n.T(c, i18n.MsgDistributorNoAvailableChannelTaskPlugin, map[string]any{"Group": group, "Model": modelName})
}
return i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": group, "Model": modelName})
@@ -190,11 +190,11 @@ func channelMatchesExpectedTaskPlugin(c *gin.Context, channel *model.Channel, ex
return true
}
}
if channel.Type == constant.ChannelTypeTaskPlugin {
return expected != "" && channel.GetSetting().TaskPluginKey == expected
}
if expected == "" {
return true
return channel.Type != constant.ChannelTypeTaskPlugin
}
if channel.Type == constant.ChannelTypeTaskPlugin || channel.Type == constant.ChannelTypeNewAPI {
return channel.GetSetting().BindsTaskPlugin(expected)
}
if c == nil {
@@ -224,6 +224,7 @@ func pinnedEndpointCandidateForChannel(c *gin.Context, channel *model.Channel, e
}
expectedOwned := false
selected := jsplugin.ProtocolBinding{}
setting := channel.GetSetting()
for _, candidate := range candidates {
if candidate.Plugin == nil {
continue
@@ -231,8 +232,12 @@ func pinnedEndpointCandidateForChannel(c *gin.Context, channel *model.Channel, e
if candidate.Plugin.Meta.Key == expected {
expectedOwned = true
}
if channel.Type == constant.ChannelTypeTaskPlugin {
if channel.GetSetting().TaskPluginKey == candidate.Plugin.Meta.Key {
if channel.Type == constant.ChannelTypeTaskPlugin || channel.Type == constant.ChannelTypeNewAPI {
// A New API channel may bind several candidates. The first bound
// candidate in generation order executes, so the billing provider
// depends only on the channel and the request, never on which
// channels an earlier retry attempt happened to try.
if selected.Plugin == nil && setting.BindsTaskPlugin(candidate.Plugin.Meta.Key) {
selected = candidate
}
continue
@@ -258,17 +263,13 @@ func getModelFromRequest(c *gin.Context) (*ModelRequest, error) {
}
}
if strings.HasPrefix(c.Request.Header.Get("Content-Type"), "application/json") {
modelRequest, err := getModelFromJSONBody(c)
if err != nil {
return nil, errors.New(i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
}
return modelRequest, nil
return getModelFromJSONBody(c)
}
var modelRequest ModelRequest
err := common.UnmarshalBodyReusable(c, &modelRequest)
if err != nil {
return nil, errors.New(i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return nil, err
}
return &modelRequest, nil
}
@@ -402,7 +403,7 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
midjourneyRequest := taskdto.MidjourneyRequest{}
err = common.UnmarshalBodyReusable(c, &midjourneyRequest)
if err != nil {
return nil, false, errors.New(i18n.T(c, i18n.MsgDistributorInvalidMidjourney, map[string]any{"Error": err.Error()}))
return nil, false, err
}
midjourneyModel, mjErr, success := service.GetMjRequestModel(relayMode, &midjourneyRequest)
if mjErr != nil {
@@ -491,6 +492,13 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
modelRequest.Model = c.Param("model")
}
}
// The routing name is also the billing name, so a model named in the body
// must not replace it: DeepSeek ignores that field.
if strings.HasPrefix(c.Request.URL.Path, "/v1/tools/search") {
modelRequest.Model = constant.DeepSeekSearchModel
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/tools/fetch") {
modelRequest.Model = constant.DeepSeekFetchModel
}
if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "dall-e")
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/edits") {
@@ -620,8 +628,15 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime)
common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
if channel.Type == constant.ChannelTypeTaskPlugin {
switch channel.Type {
case constant.ChannelTypeTaskPlugin:
c.Set("task_plugin_key", channel.GetSetting().TaskPluginKey)
case constant.ChannelTypeNewAPI:
// The bound plugin verified above executes; ordinary requests through
// the same gateway channel carry no pinned plugin.
if executing := c.GetString("expected_task_plugin_key"); executing != "" {
c.Set("task_plugin_key", executing)
}
}
logTaskPluginChannelDecision(c, channel, modelName, "channel_selected", "")
paramOverride := channel.GetParamOverride()
+80
View File
@@ -168,6 +168,27 @@ func TestTokenModelLimitAllowsLegacyAliasAndModifierVariant(t *testing.T) {
assert.True(t, TokenModelLimitAllows(wildcard, "gemini-2.5-flash-thinking-8192"))
}
// /v1/tools/search and /v1/tools/fetch always route and bill by the DeepSeek
// routing names. A model named in the body is ignored: a chat model there
// would price the call at zero tokens.
func TestGetModelRequestAssignsDeepSeekToolRoutingNames(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct{ path, body, want string }{
{"/v1/tools/search", `{"query":"deepseek","count":3}`, constant.DeepSeekSearchModel},
{"/v1/tools/fetch", `{"url":"https://example.com/"}`, constant.DeepSeekFetchModel},
{"/v1/tools/search", `{"query":"deepseek","model":"deepseek-flash"}`, constant.DeepSeekSearchModel},
{"/v1/tools/fetch", `{"url":"https://example.com/","model":"deepseek-v4-pro"}`, constant.DeepSeekFetchModel},
} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
modelRequest, shouldSelectChannel, err := getModelRequest(c)
require.NoError(t, err, tc.body)
assert.True(t, shouldSelectChannel, tc.body)
assert.Equal(t, tc.want, modelRequest.Model, tc.body)
}
}
func TestTokenModelLimitAllowsExemptAtNameByFullName(t *testing.T) {
settings := model_setting.GetGlobalSettings()
original := append([]string(nil), settings.ThinkingModelBlacklist...)
@@ -267,6 +288,14 @@ func TestNoAvailableChannelMessageWithoutPlugin(t *testing.T) {
assert.Contains(t, generic, "gpt-4o")
}
func TestNoAvailableGroupPoolMessage(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Set("pool_retry_after", 9)
assert.Equal(t, "分组池中无条件内可用分组", noAvailableGroupPoolMessage(c))
assert.Equal(t, "9", c.Writer.Header().Get("Retry-After"))
}
func TestSharedEndpointRebindsToSelectedType61Plugin(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"alpha", "beta"} {
@@ -290,3 +319,54 @@ func TestSharedEndpointRebindsToSelectedType61Plugin(t *testing.T) {
assert.Equal(t, "beta", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "beta", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
}
func TestChannelMatchesExpectedTaskPluginUsesNewAPIExtensions(t *testing.T) {
channel := &model.Channel{Type: constant.ChannelTypeNewAPI}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "single", TaskExtendPluginKeys: []string{"ext-alpha", "ext-beta"}})
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "ext-alpha"))
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "single"), "the single key stays valid on a New API channel")
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, "ext-gamma"))
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, ""), "requests without a pinned plugin still use the gateway")
assert.False(t, channelMatchesExpectedTaskPlugin(nil, &model.Channel{Type: constant.ChannelTypeNewAPI}, "ext-alpha"))
}
func TestSharedEndpointRebindsToBoundNewAPIExtension(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"alpha", "beta"} {
source := strings.Replace(distributorEndpointPluginSource(key, 0), "channelTypes: [0],", "", 1)
_, err := registry.Register(source, jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
pin := func(expected string) *gin.Context {
pinned := candidates[0]
if candidates[1].Plugin.Meta.Key == expected {
pinned = candidates[1]
}
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: pinned.Plugin})
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Generation: generation, Plugin: pinned.Plugin, Protocol: pinned.Protocol, Operation: pinned.Operation, Model: "task-model", Candidates: candidates})
c.Set("expected_task_plugin_key", expected)
return c
}
channel := &model.Channel{Id: 3, Type: constant.ChannelTypeNewAPI}
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"unrelated"}})
assert.False(t, channelMatchesExpectedTaskPlugin(pin("alpha"), channel, "alpha"))
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"beta"}})
c := pin("alpha")
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
assert.Equal(t, "beta", c.GetString("task_plugin_key"), "the only bound candidate executes")
assert.Equal(t, "beta", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "beta", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"alpha", "beta"}})
c = pin("beta")
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
assert.Equal(t, "alpha", c.GetString("task_plugin_key"), "the first bound candidate executes regardless of the earlier pin")
assert.Equal(t, "alpha", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
}
+2 -9
View File
@@ -1,7 +1,6 @@
package middleware
import (
"fmt"
"net/http"
"github.com/QuantumNous/new-api/common"
@@ -36,10 +35,7 @@ func redisEmailVerificationRateLimiter(c *gin.Context) {
waitSeconds = ttlSeconds
}
c.JSON(http.StatusTooManyRequests, gin.H{
"success": false,
"message": fmt.Sprintf("发送过于频繁,请等待 %d 秒后再试", waitSeconds),
})
common.ApiErrorStatus(c, http.StatusTooManyRequests, common.NewMessage("Please wait {{seconds}} seconds before requesting another verification code.", map[string]any{"seconds": waitSeconds}))
c.Abort()
}
@@ -47,10 +43,7 @@ func memoryEmailVerificationRateLimiter(c *gin.Context) {
key := EmailVerificationRateLimitMark + ":" + c.ClientIP()
if !inMemoryRateLimiter.Request(key, EmailVerificationMaxRequests, EmailVerificationDuration) {
c.JSON(http.StatusTooManyRequests, gin.H{
"success": false,
"message": "发送过于频繁,请稍后再试",
})
common.ApiErrorStatus(c, http.StatusTooManyRequests, common.NewMessage("Please wait before requesting another verification code."))
c.Abort()
return
}
+109
View File
@@ -0,0 +1,109 @@
package middleware
import (
"bytes"
"compress/gzip"
"errors"
"io"
"mime"
"net/http"
"path"
"strconv"
"strings"
"sync"
ginGzip "github.com/gin-contrib/gzip"
"github.com/gin-contrib/static"
"github.com/gin-gonic/gin"
)
// precompressedFile is the gzip encoding of one frontend build file.
type precompressedFile struct {
once sync.Once
body []byte
contentType string
err error
}
// ServeFrontendFiles serves files from the embedded frontend build. The build
// never changes while the process runs, so each file is gzip-compressed once
// and the bytes are reused, instead of recompressing multi-megabyte chunks on
// every request. The cache holds at most one entry per build file. Clients
// without gzip, the image types the gzip middleware skips, directories, and
// the index.html redirect keep using http.FileServer.
func ServeFrontendFiles(frontendFS static.ServeFileSystem) gin.HandlerFunc {
fileServer := http.FileServer(frontendFS)
var cache sync.Map // URL path -> *precompressedFile
return func(c *gin.Context) {
urlPath := c.Request.URL.Path
if !frontendFS.Exists("/", urlPath) {
return
}
c.Abort()
if !strings.Contains(c.GetHeader("Accept-Encoding"), "gzip") ||
ginGzip.DefaultExcludedExtentions.Contains(path.Ext(urlPath)) ||
strings.HasSuffix(urlPath, "/index.html") {
fileServer.ServeHTTP(c.Writer, c.Request)
return
}
value, ok := cache.Load(urlPath)
if !ok {
value, _ = cache.LoadOrStore(urlPath, &precompressedFile{})
}
file := value.(*precompressedFile)
file.once.Do(func() {
f, err := frontendFS.Open(urlPath)
if err != nil {
file.err = err
return
}
defer f.Close()
info, err := f.Stat()
if err != nil {
file.err = err
return
}
if info.IsDir() {
file.err = errors.New("directory")
return
}
raw, err := io.ReadAll(f)
if err != nil {
file.err = err
return
}
var compressed bytes.Buffer
writer := gzip.NewWriter(&compressed)
if _, err := writer.Write(raw); err != nil {
file.err = err
return
}
if err := writer.Close(); err != nil {
file.err = err
return
}
// Match http.FileServer: extension first, then content sniffing.
file.contentType = mime.TypeByExtension(path.Ext(urlPath))
if file.contentType == "" {
file.contentType = http.DetectContentType(raw)
}
file.body = bytes.Clone(compressed.Bytes())
})
if file.err != nil {
fileServer.ServeHTTP(c.Writer, c.Request)
return
}
// Range requests get the whole body; a server may ignore Range.
header := c.Writer.Header()
header.Set("Content-Encoding", "gzip")
header.Set("Vary", "Accept-Encoding")
header.Set("Content-Type", file.contentType)
header.Set("Content-Length", strconv.Itoa(len(file.body)))
c.Writer.WriteHeader(http.StatusOK)
if c.Request.Method != http.MethodHead {
_, _ = c.Writer.Write(file.body)
}
}
}
+106
View File
@@ -0,0 +1,106 @@
package middleware
import (
"bytes"
"compress/gzip"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"github.com/gin-contrib/static"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type countingFrontendFS struct {
static.ServeFileSystem
opens atomic.Int32
}
func (f *countingFrontendFS) Open(name string) (http.File, error) {
f.opens.Add(1)
return f.ServeFileSystem.Open(name)
}
func performFrontendRequest(router http.Handler, method, target, acceptEncoding string) *httptest.ResponseRecorder {
recorder := httptest.NewRecorder()
request := httptest.NewRequest(method, target, nil)
if acceptEncoding != "" {
request.Header.Set("Accept-Encoding", acceptEncoding)
}
router.ServeHTTP(recorder, request)
return recorder
}
func TestServeFrontendFilesCompressesEachFileOnce(t *testing.T) {
gin.SetMode(gin.TestMode)
frontendDir := t.TempDir()
script := []byte(strings.Repeat("export const answer = 42;\n", 400))
image := []byte("\x89PNG\r\n\x1a\nnot-really-an-image")
require.NoError(t, os.MkdirAll(filepath.Join(frontendDir, "static", "js"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(frontendDir, "static", "js", "app.js"), script, 0o644))
require.NoError(t, os.WriteFile(filepath.Join(frontendDir, "logo.png"), image, 0o644))
require.NoError(t, os.WriteFile(filepath.Join(frontendDir, "index.html"), []byte("<html></html>"), 0o644))
frontendFS := &countingFrontendFS{ServeFileSystem: static.LocalFile(frontendDir, false)}
router := gin.New()
router.NoRoute(ServeFrontendFiles(frontendFS), func(c *gin.Context) {
c.String(http.StatusTeapot, "fallback")
})
plain := performFrontendRequest(router, http.MethodGet, "/static/js/app.js", "")
require.Equal(t, http.StatusOK, plain.Code)
assert.Empty(t, plain.Header().Get("Content-Encoding"))
assert.Equal(t, script, plain.Body.Bytes())
frontendFS.opens.Store(0)
responses := make([]*httptest.ResponseRecorder, 8)
var wg sync.WaitGroup
for i := range responses {
wg.Go(func() {
responses[i] = performFrontendRequest(router, http.MethodGet, "/static/js/app.js", "gzip, deflate, br")
})
}
wg.Wait()
for _, response := range responses {
require.Equal(t, http.StatusOK, response.Code)
assert.Equal(t, "gzip", response.Header().Get("Content-Encoding"))
assert.Equal(t, "Accept-Encoding", response.Header().Get("Vary"))
assert.Equal(t, plain.Header().Get("Content-Type"), response.Header().Get("Content-Type"))
assert.Equal(t, strconv.Itoa(response.Body.Len()), response.Header().Get("Content-Length"))
assert.Equal(t, responses[0].Body.Bytes(), response.Body.Bytes())
}
reader, err := gzip.NewReader(bytes.NewReader(responses[0].Body.Bytes()))
require.NoError(t, err)
decoded, err := io.ReadAll(reader)
require.NoError(t, err)
assert.Equal(t, script, decoded)
head := performFrontendRequest(router, http.MethodHead, "/static/js/app.js", "gzip")
assert.Equal(t, http.StatusOK, head.Code)
assert.Equal(t, "gzip", head.Header().Get("Content-Encoding"))
assert.Equal(t, strconv.Itoa(responses[0].Body.Len()), head.Header().Get("Content-Length"))
assert.Empty(t, head.Body.Bytes())
assert.Equal(t, int32(1), frontendFS.opens.Load(), "the file must be read and compressed once")
png := performFrontendRequest(router, http.MethodGet, "/logo.png", "gzip")
assert.Equal(t, http.StatusOK, png.Code)
assert.Empty(t, png.Header().Get("Content-Encoding"))
assert.Equal(t, image, png.Body.Bytes())
index := performFrontendRequest(router, http.MethodGet, "/index.html", "gzip")
assert.Equal(t, http.StatusMovedPermanently, index.Code)
assert.Equal(t, "./", index.Header().Get("Location"))
missing := performFrontendRequest(router, http.MethodGet, "/static/js/missing.js", "gzip")
assert.Equal(t, http.StatusTeapot, missing.Code)
assert.Equal(t, "fallback", missing.Body.String())
}
+5
View File
@@ -105,6 +105,11 @@ func HeaderNavModuleAuth(module string) gin.HandlerFunc {
return func(c *gin.Context) {
access := getHeaderNavAccess(module)
if !access.Enabled {
// Keep root pricing management available without reopening the public model square.
if module == "pricing" && c.GetHeader("Authorization") != "" {
RootAuth()(c)
return
}
c.JSON(http.StatusForbidden, gin.H{
"success": false,
"message": fmt.Sprintf("%s is disabled", module),
+77 -1
View File
@@ -4,12 +4,14 @@ import (
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
@@ -50,9 +52,10 @@ func performHeaderNavRequest(t *testing.T, handler gin.HandlerFunc, authenticate
previousDB, previousRedis := model.DB, common.RedisEnabled
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, db.AutoMigrate(&model.User{}))
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Option{}))
model.DB = db
common.RedisEnabled = false
require.NoError(t, model.EnsureLegacyAccessTokenRetireAt(time.Now().Unix()))
t.Cleanup(func() {
model.DB = previousDB
common.RedisEnabled = previousRedis
@@ -96,6 +99,79 @@ func TestHeaderNavModuleAuthRejectsDisabledPricing(t *testing.T) {
require.Equal(t, http.StatusForbidden, recorder.Code)
}
func TestDisabledPricingRequiresRootAuthentication(t *testing.T) {
for _, tc := range []struct {
name string
module string
options string
role int
credential string
wantStatus int
}{
{name: "root", role: common.RoleRootUser, wantStatus: http.StatusOK},
{name: "root with login required", options: `{"pricing":{"enabled":false,"requireAuth":true}}`, role: common.RoleRootUser, wantStatus: http.StatusOK},
{name: "root with legacy disabled setting", options: `{"pricing":false}`, role: common.RoleRootUser, wantStatus: http.StatusOK},
{name: "ordinary user", role: common.RoleCommonUser, wantStatus: http.StatusForbidden},
{name: "non-root administrator", role: common.RoleAdminUser, wantStatus: http.StatusForbidden},
{name: "expired root token", role: common.RoleRootUser, credential: "expired", wantStatus: http.StatusUnauthorized},
{name: "tampered root token", role: common.RoleRootUser, credential: "tampered", wantStatus: http.StatusUnauthorized},
{name: "revoked root session", role: common.RoleRootUser, credential: "revoked", wantStatus: http.StatusUnauthorized},
{name: "disabled root account", role: common.RoleRootUser, credential: "banned", wantStatus: http.StatusUnauthorized},
{name: "root cannot bypass disabled rankings", module: "rankings", options: `{"rankings":false}`, role: common.RoleRootUser, wantStatus: http.StatusForbidden},
} {
t.Run(tc.name, func(t *testing.T) {
setupDashboardAuthMiddlewareTest(t)
gin.SetMode(gin.TestMode)
options := tc.options
if options == "" {
options = `{"pricing":{"enabled":false,"requireAuth":false}}`
}
withHeaderNavModules(t, options)
user := model.User{
Username: "pricing-reader", Password: "unused-password-hash", Role: tc.role,
Status: common.UserStatusEnabled, Group: "default", AuthVersion: 1,
}
require.NoError(t, model.DB.Create(&user).Error)
bundle, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "header-nav-test")
require.NoError(t, err)
token := bundle.AccessToken
switch tc.credential {
case "expired":
identity, internal, err := service.ParseDashboardAccessToken(token)
require.NoError(t, err)
require.True(t, internal)
token = issueExpiredDashboardAccessToken(t, identity)
case "tampered":
token = tamperDashboardToken(token)
case "revoked":
revoked, err := model.RevokeUserSession(user.Id, bundle.Session.SID, "test")
require.NoError(t, err)
require.True(t, revoked)
case "banned":
require.NoError(t, model.DB.Model(&user).Update("status", common.UserStatusDisabled).Error)
}
module := tc.module
if module == "" {
module = "pricing"
}
handled := false
router := gin.New()
router.GET("/api/"+module, HeaderNavModuleAuth(module), func(c *gin.Context) {
handled = true
assert.Equal(t, user.Id, c.GetInt("id"))
assert.Equal(t, common.RoleRootUser, c.GetInt("role"))
c.JSON(http.StatusOK, gin.H{"success": true})
})
response := middlewareBearerRequest(router, "/api/"+module, token)
assert.Equal(t, tc.wantStatus, response.Code)
assert.Equal(t, tc.wantStatus == http.StatusOK, handled)
})
}
}
func TestHeaderNavModuleAuthRequiresLoginForPricing(t *testing.T) {
raw := `{"pricing":{"enabled":true,"requireAuth":true}}`
withHeaderNavModules(t, raw)
+3 -20
View File
@@ -18,8 +18,8 @@ func I18n() gin.HandlerFunc {
}
}
// detectLanguage determines the language preference for the request
// Priority: 1. User setting (if logged in) -> 2. Accept-Language header -> 3. Default language
// detectLanguage returns the language the request states, or "" when it states
// none. Priority: 1. User setting (if logged in) -> 2. Accept-Language header
func detectLanguage(c *gin.Context) string {
// 1. Try to get language from user setting (set by auth middleware)
if userSetting, ok := common.GetContextKeyType[dto.UserSetting](c, constant.ContextKeyUserSetting); ok {
@@ -29,22 +29,5 @@ func detectLanguage(c *gin.Context) string {
}
// 2. Parse Accept-Language header
acceptLang := c.GetHeader("Accept-Language")
if acceptLang != "" {
lang := i18n.ParseAcceptLanguage(acceptLang)
if i18n.IsSupported(lang) {
return lang
}
}
// 3. Return default language
return i18n.DefaultLang
}
// GetLanguage returns the current language from gin context
func GetLanguage(c *gin.Context) string {
if lang := c.GetString(string(constant.ContextKeyLanguage)); lang != "" {
return lang
}
return i18n.DefaultLang
return i18n.ParseAcceptLanguage(c.GetHeader("Accept-Language"))
}
+19 -18
View File
@@ -11,6 +11,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/common/limiter"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting"
@@ -78,28 +79,27 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC
}
// Redis限流处理器
func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc {
func redisRateLimitHandler(subject string, duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc {
return func(c *gin.Context) {
userId := strconv.Itoa(c.GetInt("id"))
ctx := context.Background()
rdb := common.RDB
// 1. 检查成功请求数限制
successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, userId)
successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, subject)
allowed, err := checkRedisRateLimit(ctx, rdb, successKey, successMaxCount, duration)
if err != nil {
fmt.Println("检查成功请求数限制失败:", err.Error())
fmt.Println(common.LogText("failed to check the successful request limit: %s", err.Error()))
abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed")
return
}
if !allowed {
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount))
abortWithOpenAiMessage(c, http.StatusTooManyRequests, i18n.T(c, i18n.MsgRateLimitReached, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": successMaxCount}))
return
}
//2.检查总请求数限制并记录总请求(当totalMaxCount为0时会自动跳过,使用令牌桶限流器
if totalMaxCount > 0 {
totalKey := fmt.Sprintf("rateLimit:%s", userId)
totalKey := fmt.Sprintf("rateLimit:%s", subject)
// 初始化
tb := limiter.New(ctx, rdb)
allowed, err = tb.Allow(
@@ -111,13 +111,13 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g
)
if err != nil {
fmt.Println("检查总请求数限制失败:", err.Error())
fmt.Println(common.LogText("failed to check the total request limit: %s", err.Error()))
abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed")
return
}
if !allowed {
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到总请求数限制:%d分钟内最多请求%d次,包括失败次数,请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount))
abortWithOpenAiMessage(c, http.StatusTooManyRequests, i18n.T(c, i18n.MsgRateLimitTotalReached, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": totalMaxCount}))
return
}
}
@@ -133,18 +133,16 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g
}
// 内存限流处理器
func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc {
func memoryRateLimitHandler(subject string, duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc {
inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute)
return func(c *gin.Context) {
userId := strconv.Itoa(c.GetInt("id"))
totalKey := ModelRequestRateLimitCountMark + userId
successKey := ModelRequestRateLimitSuccessCountMark + userId
totalKey := ModelRequestRateLimitCountMark + subject
successKey := ModelRequestRateLimitSuccessCountMark + subject
// 1. 检查总请求数限制(当totalMaxCount为0时跳过)
if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) {
c.Status(http.StatusTooManyRequests)
c.Abort()
abortWithOpenAiMessage(c, http.StatusTooManyRequests, i18n.T(c, i18n.MsgRateLimitTotalReached, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": totalMaxCount}))
return
}
@@ -152,7 +150,7 @@ func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int)
if successMaxCount > 0 {
reservation = inMemoryRateLimiter.Reserve(successKey, successMaxCount, duration)
if reservation == nil {
c.AbortWithStatus(http.StatusTooManyRequests)
abortWithOpenAiMessage(c, http.StatusTooManyRequests, i18n.T(c, i18n.MsgRateLimitReached, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": successMaxCount}))
return
}
defer reservation.Complete(false)
@@ -191,18 +189,21 @@ func ModelRequestRateLimit() func(c *gin.Context) {
group = common.GetContextKeyString(c, constant.ContextKeyUserGroup)
}
//获取分组的限流配置
// 获取分组的限流配置。配置了独立限流的分组按用户和分组单独计数,
// 其余分组共用该用户的全局计数
subject := strconv.Itoa(c.GetInt("id"))
groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group)
if found {
totalMaxCount = groupTotalCount
successMaxCount = groupSuccessCount
subject += ":" + group
}
// 根据存储类型选择并执行限流处理器
if common.RedisEnabled {
redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
redisRateLimitHandler(subject, duration, totalMaxCount, successMaxCount)(c)
} else {
memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
memoryRateLimitHandler(subject, duration, totalMaxCount, successMaxCount)(c)
}
}
}
+103 -6
View File
@@ -4,13 +4,16 @@ import (
"context"
"fmt"
"net/http"
"strconv"
"sync/atomic"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -43,17 +46,18 @@ func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) {
var modelRateLimitTestUsers atomic.Int64
func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) {
require.NoError(t, i18n.Init())
for _, backend := range []string{"memory", "redis"} {
for _, totalLimit := range []int{0, 2} {
t.Run(fmt.Sprintf("%s/total=%d", backend, totalLimit), func(t *testing.T) {
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
handler := memoryRateLimitHandler(60, totalLimit, 1)
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
handler := memoryRateLimitHandler(subject, 60, totalLimit, 1)
if backend == "redis" {
useRateLimitMiniRedis(t)
handler = redisRateLimitHandler(60, totalLimit, 1)
handler = redisRateLimitHandler(subject, 60, totalLimit, 1)
}
router := gin.New()
router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, handler, func(c *gin.Context) {
router.GET("/:outcome", handler, func(c *gin.Context) {
status := relaycommon.NewStreamStatus()
if c.Param("outcome") == "failed" {
status.MarkFailed("server_error", "", 0)
@@ -76,10 +80,11 @@ func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) {
}
func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) {
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
require.NoError(t, i18n.Init())
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{})
router := gin.New()
router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, memoryRateLimitHandler(60, 0, 1), func(c *gin.Context) {
router.GET("/:outcome", memoryRateLimitHandler(subject, 60, 0, 1), func(c *gin.Context) {
if c.Param("outcome") == "failed" {
close(entered)
<-release
@@ -100,3 +105,95 @@ func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) {
assert.Equal(t, http.StatusOK, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/completed", "127.0.0.1:1000").Code)
}
// The memory limiter answers like the Redis limiter: an OpenAI error body
// naming the limit that was reached, not an empty 429.
func TestModelMemoryRateLimitRejectionsReturnOpenAIError(t *testing.T) {
require.NoError(t, i18n.Init())
for _, tc := range []struct {
name string
totalLimit, successLimit int
messageKey string
}{
{name: "total", totalLimit: 1, messageKey: i18n.MsgRateLimitTotalReached},
{name: "success", successLimit: 1, messageKey: i18n.MsgRateLimitReached},
} {
t.Run(tc.name, func(t *testing.T) {
subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1)))
router := gin.New()
router.GET("/", func(c *gin.Context) {
c.Set(common.RequestIdKey, "req-limited")
}, memoryRateLimitHandler(subject, 60, tc.totalLimit, tc.successLimit), func(c *gin.Context) {
c.Status(http.StatusOK)
})
require.Equal(t, http.StatusOK, performRateLimitRequest(router, "/", "127.0.0.1:1000").Code)
resp := performRateLimitRequest(router, "/", "127.0.0.1:1000")
assert.Equal(t, http.StatusTooManyRequests, resp.Code)
message := i18n.Translate("", tc.messageKey, map[string]any{"Minutes": setting.ModelRequestRateLimitDurationMinutes, "Max": 1})
want, err := common.Marshal(gin.H{"error": gin.H{
"message": common.MessageWithRequestId(message, "req-limited"),
"type": "new_api_error",
"code": "",
}})
require.NoError(t, err)
assert.JSONEq(t, string(want), resp.Body.String())
})
}
}
// A group with its own limit counts each user's requests on its own; groups
// without one keep sharing the user's global limit.
func TestModelRateLimitCountsConfiguredGroupsSeparately(t *testing.T) {
require.NoError(t, i18n.Init())
previousEnabled, previousMinutes := setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes
previousTotal, previousSuccess := setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount
previousGroups, previousRedisEnabled := setting.ModelRequestRateLimitGroup, common.RedisEnabled
t.Cleanup(func() {
setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = previousEnabled, previousMinutes
setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = previousTotal, previousSuccess
setting.ModelRequestRateLimitGroup, common.RedisEnabled = previousGroups, previousRedisEnabled
})
setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = true, 1
common.RedisEnabled = false
// Groups a and b have their own limits; c and d have none and share the global one.
steps := []struct {
group string
want int
}{
{"a", http.StatusOK}, {"a", http.StatusTooManyRequests},
{"b", http.StatusOK}, {"b", http.StatusTooManyRequests},
{"c", http.StatusOK}, {"d", http.StatusTooManyRequests},
}
// The Redis token bucket binds to the first Redis client it sees and another
// test already drives it, so the total limit runs on the memory backend only.
for _, tc := range []struct {
backend, limit string
counts [2]int // [total, success] for groups a and b and for the global limit
}{
{"memory", "success", [2]int{0, 1}},
{"redis", "success", [2]int{0, 1}},
{"memory", "total", [2]int{1, 1000}},
} {
t.Run(tc.backend+"/"+tc.limit, func(t *testing.T) {
if tc.backend == "redis" {
useRateLimitMiniRedis(t)
}
setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = tc.counts[0], tc.counts[1]
setting.ModelRequestRateLimitGroup = map[string][2]int{"a": tc.counts, "b": tc.counts}
userID := 7200000 + int(modelRateLimitTestUsers.Add(1))
router := gin.New()
router.GET("/:group", func(c *gin.Context) {
c.Set("id", userID)
common.SetContextKey(c, constant.ContextKeyTokenGroup, c.Param("group"))
}, ModelRequestRateLimit(), func(c *gin.Context) {
c.Status(http.StatusOK)
})
for _, step := range steps {
assert.Equal(t, step.want, performRateLimitRequest(router, "/"+step.group, "127.0.0.1:1000").Code, "group %s", step.group)
}
})
}
}
+6 -28
View File
@@ -1,9 +1,9 @@
package middleware
import (
"fmt"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -39,33 +39,11 @@ func SystemPerformanceCheck() gin.HandlerFunc {
// checkSystemPerformance 检查系统性能是否超过阈值
func checkSystemPerformance() *types.NewAPIError {
config := common.GetPerformanceMonitorConfig()
if !config.Enabled {
// 过载标记在采样侧完成:CPU / 内存按容器 cgroup 配额计算,并要求连续多次
// 超阈值;判定时再按当前阈值复核,并忽略过期采样。
overload := common.CheckSystemOverload(common.GetSystemStatus(), common.GetPerformanceMonitorConfig(), time.Now())
if overload == nil {
return nil
}
status := common.GetSystemStatus()
// 检查 CPU
if config.CPUThreshold > 0 && int(status.CPUUsage) > config.CPUThreshold {
return types.NewErrorWithStatusCode(
fmt.Errorf("system cpu overloaded (current: %.1f%%, threshold: %d%%)", status.CPUUsage, config.CPUThreshold),
"system_cpu_overloaded", http.StatusServiceUnavailable)
}
// 检查内存
if config.MemoryThreshold > 0 && int(status.MemoryUsage) > config.MemoryThreshold {
return types.NewErrorWithStatusCode(
fmt.Errorf("system memory overloaded (current: %.1f%%, threshold: %d%%)", status.MemoryUsage, config.MemoryThreshold),
"system_memory_overloaded", http.StatusServiceUnavailable)
}
// 检查磁盘
if config.DiskThreshold > 0 && int(status.DiskUsage) > config.DiskThreshold {
return types.NewErrorWithStatusCode(
fmt.Errorf("system disk overloaded (current: %.1f%%, threshold: %d%%)", status.DiskUsage, config.DiskThreshold),
"system_disk_overloaded", http.StatusServiceUnavailable)
}
return nil
return types.NewErrorWithStatusCode(overload, types.ErrorCode(overload.Code), http.StatusServiceUnavailable)
}
+21 -5
View File
@@ -9,6 +9,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/gin-contrib/static"
"github.com/gin-gonic/gin"
)
@@ -114,7 +115,7 @@ func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark st
duration,
)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit check failed (mark=%s): %v", mark, err))
logger.LogError(c.Request.Context(), common.LogText("rate limit check failed (mark=%s): %v", mark, err))
c.Status(http.StatusInternalServerError)
c.Abort()
return
@@ -157,11 +158,26 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi
}
}
func GlobalWebRateLimit() func(c *gin.Context) {
// GlobalWebRateLimit limits dashboard pages and other fallback paths per IP.
// Files that exist in the frontend build count against a separate static
// limiter instead, so lazy-loaded chunks and icons do not use up the page
// budget. The static limiter is disabled by default.
func GlobalWebRateLimit(frontendFS static.ServeFileSystem) func(c *gin.Context) {
webLimit := defNext
if common.GlobalWebRateLimitEnable {
return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
webLimit = rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
}
staticLimit := defNext
if common.GlobalStaticRateLimitEnable {
staticLimit = rateLimitFactory(common.GlobalStaticRateLimitNum, common.GlobalStaticRateLimitDuration, "GS")
}
return func(c *gin.Context) {
if frontendFS.Exists("/", c.Request.URL.Path) {
staticLimit(c)
return
}
webLimit(c)
}
return defNext
}
func GlobalAPIRateLimit() func(c *gin.Context) {
@@ -234,7 +250,7 @@ func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c
func userRedisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, key string) {
allowed, _, ttlSeconds, err := redisFixedWindowTake(c.Request.Context(), key, maxRequestNum, duration)
if err != nil {
logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit check failed (key=%s): %v", key, err))
logger.LogError(c.Request.Context(), common.LogText("rate limit check failed (key=%s): %v", key, err))
c.Status(http.StatusInternalServerError)
c.Abort()
return
+68 -1
View File
@@ -4,6 +4,8 @@ import (
"context"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync"
"sync/atomic"
"testing"
@@ -11,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/alicebob/miniredis/v2"
"github.com/gin-contrib/static"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
"github.com/stretchr/testify/assert"
@@ -108,7 +111,7 @@ func TestRedisEmailVerificationRateLimiterPreservesResponseAndTTL(t *testing.T)
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify", remoteAddr).Code)
response := performRateLimitRequest(router, "/verify", remoteAddr)
assert.Equal(t, http.StatusTooManyRequests, response.Code)
assert.JSONEq(t, `{"success":false,"message":"发送过于频繁,请等待 30 秒后再试"}`, response.Body.String())
assert.JSONEq(t, `{"success":false,"message":"Please wait 30 seconds before requesting another verification code.","message_key":"Please wait {{seconds}} seconds before requesting another verification code.","message_params":{"seconds":30}}`, response.Body.String())
key := redisIPRateLimitKey(EmailVerificationRateLimitMark, "192.0.2.30")
assert.True(t, redisServer.Exists(key))
@@ -223,3 +226,67 @@ func TestRedisFailurePolicies(t *testing.T) {
assert.Empty(t, userResponse.Body.String())
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/email", "192.0.2.62:12345").Code)
}
func TestGlobalWebRateLimitCountsFrontendFilesInStaticBucket(t *testing.T) {
gin.SetMode(gin.TestMode)
previousWebEnable, previousWebNum, previousWebDuration := common.GlobalWebRateLimitEnable, common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration
previousStaticEnable, previousStaticNum, previousStaticDuration := common.GlobalStaticRateLimitEnable, common.GlobalStaticRateLimitNum, common.GlobalStaticRateLimitDuration
t.Cleanup(func() {
common.GlobalWebRateLimitEnable, common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration = previousWebEnable, previousWebNum, previousWebDuration
common.GlobalStaticRateLimitEnable, common.GlobalStaticRateLimitNum, common.GlobalStaticRateLimitDuration = previousStaticEnable, previousStaticNum, previousStaticDuration
})
common.GlobalWebRateLimitEnable, common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration = true, 2, 41
common.GlobalStaticRateLimitNum, common.GlobalStaticRateLimitDuration = 2, 43
frontendDir := t.TempDir()
require.NoError(t, os.MkdirAll(filepath.Join(frontendDir, "static", "js"), 0o755))
require.NoError(t, os.WriteFile(filepath.Join(frontendDir, "static", "js", "app.js"), []byte("app"), 0o644))
frontendFS := static.LocalFile(frontendDir, false)
const clientIP = "192.0.2.70"
const remoteAddr = clientIP + ":12345"
t.Run("static limiter disabled by default", func(t *testing.T) {
redisServer, _ := useRateLimitMiniRedis(t)
common.GlobalStaticRateLimitEnable = false
router := gin.New()
require.NoError(t, router.SetTrustedProxies(nil))
router.NoRoute(GlobalWebRateLimit(frontendFS), func(c *gin.Context) { c.Status(http.StatusNoContent) })
for range 5 {
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/static/js/app.js", remoteAddr).Code)
}
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/channels", remoteAddr).Code)
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/static/js/missing.js", remoteAddr).Code)
pageResponse := performRateLimitRequest(router, "/channels", remoteAddr)
assert.Equal(t, http.StatusTooManyRequests, pageResponse.Code)
assert.Equal(t, "41", pageResponse.Header().Get("Retry-After"))
webCount, err := redisServer.Get(redisIPRateLimitKey("GW", clientIP))
require.NoError(t, err)
assert.Equal(t, "3", webCount)
assert.False(t, redisServer.Exists(redisIPRateLimitKey("GS", clientIP)))
})
t.Run("static limiter enabled uses its own bucket", func(t *testing.T) {
redisServer, _ := useRateLimitMiniRedis(t)
common.GlobalStaticRateLimitEnable = true
router := gin.New()
require.NoError(t, router.SetTrustedProxies(nil))
router.NoRoute(GlobalWebRateLimit(frontendFS), func(c *gin.Context) { c.Status(http.StatusNoContent) })
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/static/js/app.js", remoteAddr).Code)
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/static/js/app.js", remoteAddr).Code)
staticResponse := performRateLimitRequest(router, "/static/js/app.js", remoteAddr)
assert.Equal(t, http.StatusTooManyRequests, staticResponse.Code)
assert.Equal(t, "43", staticResponse.Header().Get("Retry-After"))
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/channels", remoteAddr).Code)
staticCount, err := redisServer.Get(redisIPRateLimitKey("GS", clientIP))
require.NoError(t, err)
assert.Equal(t, "3", staticCount)
webCount, err := redisServer.Get(redisIPRateLimitKey("GW", clientIP))
require.NoError(t, err)
assert.Equal(t, "1", webCount)
})
}
+2 -2
View File
@@ -13,8 +13,8 @@ func RelayPanicRecover() gin.HandlerFunc {
return func(c *gin.Context) {
defer func() {
if err := recover(); err != nil {
common.SysLog(fmt.Sprintf("panic detected: %v", err))
common.SysLog(fmt.Sprintf("stacktrace from panic: %s", string(debug.Stack())))
common.SysLog(common.LogText("panic detected: %v", err))
common.SysLog(common.LogText("stacktrace from panic: %s", string(debug.Stack())))
c.JSON(http.StatusInternalServerError, gin.H{
"error": gin.H{
"message": fmt.Sprintf("Panic detected, error: %v. Please submit a issue here: https://github.com/Calcium-Ion/new-api", err),
+4
View File
@@ -2,6 +2,7 @@ package middleware
import (
"context"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/gin-gonic/gin"
@@ -9,6 +10,9 @@ import (
func RequestId() func(c *gin.Context) {
return func(c *gin.Context) {
if c.Request.URL.Path == "/v1/responses" {
c.Set("request_received_at", time.Now())
}
id := common.NewRequestId()
c.Set(common.RequestIdKey, id)
ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id)
+12 -8
View File
@@ -35,29 +35,33 @@ func SecureVerificationRequired() gin.HandlerFunc {
}
// RequireSecurityProof validates a proof against the authenticated dashboard
// session and writes the shared proof error contract on failure.
// session or scoped access token and writes the shared proof error contract on
// failure.
func RequireSecurityProof(c *gin.Context, operation service.VerificationOperation) *model.AuthFlowAuthorization {
identity, ok := GetSessionAuthIdentity(c)
identity, ok := GetStepUpIdentity(c)
if !ok {
securityProofError(c, "SECURITY_PROOF_INVALID", "安全验证状态无效")
securityProofError(c, "SECURITY_PROOF_INVALID", "Security verification is no longer valid. Please verify again.")
return nil
}
raw := strings.TrimSpace(c.GetHeader("X-Security-Proof"))
if raw == "" {
securityProofError(c, "SECURITY_PROOF_REQUIRED", "需要安全验证")
securityProofError(c, "SECURITY_PROOF_REQUIRED", "Additional verification required")
return nil
}
authorization, err := service.ConsumeOperationProof(raw, identity, operation)
if err != nil {
switch {
case errors.Is(err, service.ErrAuthTokenExpired):
securityProofError(c, "SECURITY_PROOF_EXPIRED", "安全验证已过期")
securityProofError(c, "SECURITY_PROOF_EXPIRED", "Security verification has expired. Please verify again.")
case errors.Is(err, service.ErrProofScope):
securityProofError(c, "SECURITY_PROOF_SCOPE_MISMATCH", "安全验证范围不匹配")
securityProofError(c, "SECURITY_PROOF_SCOPE_MISMATCH", "Verification does not match this action.")
case errors.Is(err, service.ErrVerificationContextInvalid):
c.Set("security_error_code", "SECURITY_CONTEXT_INVALID")
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"success": false, "code": "SECURITY_CONTEXT_INVALID", "message": service.ErrVerificationContextInvalid.Error()})
case errors.Is(err, service.ErrVerificationUnavailable):
securityProofError(c, "SECURITY_METHOD_UNAVAILABLE", service.ErrVerificationUnavailable.Error())
case errors.Is(err, service.ErrProofMethod):
securityProofError(c, "SECURITY_PROOF_METHOD_MISMATCH", "安全验证方式不匹配")
securityProofError(c, "SECURITY_PROOF_METHOD_MISMATCH", "This verification method is not allowed for this action.")
case errors.Is(err, service.ErrProofConsumed):
securityProofError(c, "SECURITY_PROOF_CONSUMED", "This verification has already been used. Please verify again.")
case errors.Is(err, service.ErrProofContext):
@@ -65,7 +69,7 @@ func RequireSecurityProof(c *gin.Context, operation service.VerificationOperatio
case errors.Is(err, service.ErrVerificationForbidden):
securityProofError(c, "SECURITY_ACTION_FORBIDDEN", service.ErrVerificationForbidden.Error())
case errors.Is(err, service.ErrAuthTokenInvalid), errors.Is(err, service.ErrLoginSessionInvalid), errors.Is(err, service.ErrLoginSessionRevoked), errors.Is(err, model.ErrUserSessionInactive):
securityProofError(c, "SECURITY_PROOF_INVALID", "安全验证状态无效")
securityProofError(c, "SECURITY_PROOF_INVALID", "Security verification is no longer valid. Please verify again.")
default:
_ = c.Error(err)
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"success": false, "code": "AUTH_INTERNAL_ERROR", "message": "Please try again later."})
+134 -105
View File
@@ -3,6 +3,7 @@ package middleware
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
@@ -19,6 +20,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
@@ -77,11 +79,12 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if err != nil {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_request",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_request err=%q",
generation,
pinned.Plugin.Meta.Key,
err.Error()),
)
status := http.StatusBadRequest
status := invalidRequestStatus(err)
if errors.Is(err, errTaskPluginUnsupportedMediaType) {
status = http.StatusUnsupportedMediaType
}
@@ -93,10 +96,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if pinned.Route.Type == pluginruntime.RouteTypeQuery && bodyKind != string(pluginruntime.BodyNone) || pinned.Route.Type != pluginruntime.RouteTypeQuery && bodyKind != string(pluginruntime.BodyJSON) {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=body_kind_mismatch body_kind=%q",
generation,
pinned.Plugin.Meta.Key,
bodyKind,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=body_kind_mismatch body_kind=%q",
generation,
pinned.Plugin.Meta.Key,
bodyKind),
)
detail := "this route requires a JSON body"
if pinned.Route.Type == pluginruntime.RouteTypeQuery {
@@ -124,10 +127,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if claimedModel == "" || !slices.Contains(pinned.Route.Models, claimedModel) {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=model_not_allowed model=%q",
generation,
pinned.Plugin.Meta.Key,
claimedModel,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=model_not_allowed model=%q",
generation,
pinned.Plugin.Meta.Key,
claimedModel),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, fmt.Sprintf("model %q is not allowed on this route", claimedModel))
return
@@ -135,15 +138,18 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
}
hookStarted := time.Now()
resolvedValue, err := pinned.Plugin.Engine.CallMember(c.Request.Context(), "native", pinned.Route.Decode, requestContext.JSValue())
resolvedValue, requestBodyText, err := pinned.Plugin.Engine.CallPathWithMemberJSON(
c.Request.Context(), 0, pinned.Plugin.Meta.JSONTextMember("requestBody"),
"native", []string{pinned.Route.Decode}, requestContext.JSValueFor(pinned.Plugin.Meta),
)
if err != nil {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=hook_failed err=%q elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
err.Error(),
time.Since(hookStarted).Milliseconds(),
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=hook_failed err=%q elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
err.Error(),
time.Since(hookStarted).Milliseconds()),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginHookDetail(err))
return
@@ -152,10 +158,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !ok {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=result_not_object elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
time.Since(hookStarted).Milliseconds(),
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=result_not_object elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
time.Since(hookStarted).Milliseconds()),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
@@ -164,16 +170,16 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !ok {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=missing_kind elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
time.Since(hookStarted).Milliseconds(),
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=missing_kind elapsed_ms=%d",
generation,
pinned.Plugin.Meta.Key,
time.Since(hookStarted).Milliseconds()),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
}
if _, forbidden := resolved["renderer"]; forbidden {
logger.LogWarn(c, "task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=forbidden_renderer", generation, pinned.Plugin.Meta.Key)
logger.LogWarn(c, common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=forbidden_renderer", generation, pinned.Plugin.Meta.Key))
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
}
@@ -184,9 +190,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !valid || strings.TrimSpace(modelName) == "" {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_model",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_model",
generation,
pinned.Plugin.Meta.Key),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, "decoded request is missing a model")
return
@@ -195,10 +201,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !owned {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=model_not_owned model=%q",
generation,
pinned.Plugin.Meta.Key,
modelName,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=model_not_owned model=%q",
generation,
pinned.Plugin.Meta.Key,
modelName),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, fmt.Sprintf("model %q is not served by this plugin", modelName))
return
@@ -206,10 +212,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if len(pinned.Route.Models) > 0 && !slices.Contains(pinned.Route.Models, modelName) {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=resolved_model_not_allowed model=%q",
generation,
pinned.Plugin.Meta.Key,
modelName,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=resolved_model_not_allowed model=%q",
generation,
pinned.Plugin.Meta.Key,
modelName),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, fmt.Sprintf("model %q is not allowed on this route", modelName))
return
@@ -220,9 +226,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !actionOK {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_action",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_action",
generation,
pinned.Plugin.Meta.Key),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
@@ -236,6 +242,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
c.Set(pluginruntime.ContextKeyRouteRequest, requestContext)
}
c.Set("task_request", requestContext.RequestBody)
if requestBodyText != nil {
c.Set(pluginruntime.ContextKeyRequestBodyText, requestBodyText)
}
c.Set("resolved_task_model", modelName)
c.Set("expected_task_plugin_key", pinned.Plugin.Meta.Key)
c.Set("task_plugin_key", pinned.Plugin.Meta.Key)
@@ -248,10 +257,10 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if intentErr := applyOriginTaskIntent(c, resolved, pinned.Plugin.Meta); intentErr != nil {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=origin_task reason=%s",
generation,
pinned.Plugin.Meta.Key,
intentErr.Code,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=origin_task reason=%s",
generation,
pinned.Plugin.Meta.Key,
intentErr.Code),
)
abortTaskPluginRouteErrorDetail(c, intentErr.StatusCode, intentErr.Message)
return
@@ -272,9 +281,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if pinned.Route.Type != pluginruntime.RouteTypeDynamic {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=query_from_non_dynamic_route",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=query_from_non_dynamic_route",
generation,
pinned.Plugin.Meta.Key),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
@@ -283,9 +292,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
if !valid {
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_query_result",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=invalid_query_result",
generation,
pinned.Plugin.Meta.Key),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
@@ -303,9 +312,9 @@ func PrepareTaskPluginRoute() gin.HandlerFunc {
default:
logger.LogWarn(
c,
"task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=unsupported_kind",
generation,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=route event=prepare_rejected generation=%d plugin=%q stage=resolve_request reason=unsupported_kind",
generation,
pinned.Plugin.Meta.Key),
)
abortTaskPluginRouteErrorDetail(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
}
@@ -326,7 +335,7 @@ func PinTaskPluginEndpoint() gin.HandlerFunc {
modelRequest, err := getModelFromRequest(c)
if err != nil {
if _, _, protocolPath := pluginruntime.LookupHostProtocolOperation(c.Request.Method, c.Request.URL.Path); protocolPath {
abortWithOpenAiMessage(c, http.StatusBadRequest, "Invalid task protocol request")
abortWithOpenAiMessage(c, invalidRequestStatus(err), i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return
}
c.Next()
@@ -525,9 +534,9 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if !pluginruntime.SupportsHostProtocol(pinned.Protocol) {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=protocol_check reason=unsupported_protocol",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=protocol_check reason=unsupported_protocol",
pinned.Generation.Number,
pinned.Plugin.Meta.Key),
)
abortWithOpenAiMessage(c, http.StatusNotImplemented, "Task protocol bridge is not available")
return
@@ -536,9 +545,9 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if err != nil {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_request",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_request",
pinned.Generation.Number,
pinned.Plugin.Meta.Key),
)
status := http.StatusBadRequest
if errors.Is(err, errTaskPluginUnsupportedMediaType) {
@@ -566,10 +575,10 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if !allowedBody {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=body_kind_mismatch body_kind=%q",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
bodyKind,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=body_kind_mismatch body_kind=%q",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
bodyKind),
)
detail := "unsupported request body for this operation"
if len(pinned.Operation.BodyKinds) == 1 && pinned.Operation.BodyKinds[0] == pluginruntime.BodyJSON {
@@ -586,9 +595,9 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if !ok {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_stream_flag",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=request_decode reason=invalid_stream_flag",
pinned.Generation.Number,
pinned.Plugin.Meta.Key),
)
abortWithOpenAiMessage(c, http.StatusBadRequest, "stream must be a boolean")
return
@@ -610,6 +619,7 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
}
accepted := make([]pluginruntime.ProtocolBinding, 0, len(candidates))
var resolved map[string]any
var resolvedBodyText json.RawMessage
var failures []string
rejectedPlugins := make(map[string][]string)
for _, candidate := range candidates {
@@ -618,9 +628,10 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
candidateContext.Operation = candidate.Operation.Name
// Parsing belongs to durable submission; disconnecting only stops
// the later Responses observation.
resolvedValue, callErr := candidate.Plugin.Engine.CallPathWithAdmissionTimeout(
resolvedValue, requestBodyText, callErr := candidate.Plugin.Engine.CallPathWithMemberJSON(
context.WithoutCancel(c.Request.Context()), pluginruntime.DefaultCallTimeout,
"protocols", []string{candidate.Protocol, "decodeRequest"}, candidateContext.JSValue(),
candidate.Plugin.Meta.JSONTextMember("requestBody"),
"protocols", []string{candidate.Protocol, "decodeRequest"}, candidateContext.JSValueFor(candidate.Plugin.Meta),
)
result, resultOK := resolvedValue.(map[string]any)
detail := ""
@@ -645,8 +656,8 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
detail = fmt.Sprintf("model %q is not served by this plugin", model)
}
if detail != "" {
logger.LogWarn(c, "task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=%s err=%q elapsed_ms=%d",
pinned.Generation.Number, candidate.Plugin.Meta.Key, reason, detail, time.Since(hookStarted).Milliseconds())
logger.LogWarn(c, common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=%s err=%q elapsed_ms=%d",
pinned.Generation.Number, candidate.Plugin.Meta.Key, reason, detail, time.Since(hookStarted).Milliseconds()))
if _, seen := rejectedPlugins[detail]; !seen {
failures = append(failures, detail)
}
@@ -656,6 +667,7 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
accepted = append(accepted, candidate)
if resolved == nil {
resolved = result
resolvedBodyText = requestBodyText
protocolContext = candidateContext
}
}
@@ -683,9 +695,9 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if !ok {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=invalid_action",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=parse_request reason=invalid_action",
pinned.Generation.Number,
pinned.Plugin.Meta.Key),
)
abortWithOpenAiMessage(c, http.StatusBadRequest, taskPluginInvalidRouteResult)
return
@@ -697,6 +709,9 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
}
c.Set(pluginruntime.ContextKeyRouteRequest, requestContext)
c.Set("task_request", requestContext.RequestBody)
if resolvedBodyText != nil {
c.Set(pluginruntime.ContextKeyRequestBodyText, resolvedBodyText)
}
c.Set("resolved_task_model", resolvedModel)
c.Set("expected_task_plugin_key", pinned.Plugin.Meta.Key)
c.Set("task_plugin_key", pinned.Plugin.Meta.Key)
@@ -709,10 +724,10 @@ func PrepareTaskPluginEndpoint() gin.HandlerFunc {
if intentErr := applyOriginTaskIntent(c, resolved, pinned.Plugin.Meta); intentErr != nil {
logger.LogWarn(
c,
"task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=origin_task reason=%s",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
intentErr.Code,
common.LogText("task_plugin subsystem=endpoint event=prepare_rejected generation=%d plugin=%q stage=origin_task reason=%s",
pinned.Generation.Number,
pinned.Plugin.Meta.Key,
intentErr.Code),
)
abortWithOpenAiMessage(c, intentErr.StatusCode, intentErr.Message, types.ErrorCode(intentErr.Code))
return
@@ -791,6 +806,9 @@ func buildTaskPluginRouteRequest(c *gin.Context) (pluginruntime.RouteRequestCont
return requestContext, err
}
requestContext.Body = map[string]any{"kind": string(pluginruntime.BodyJSON), "value": value}
// Body storage never changes its bytes in place (a rewrite stores a
// new copy), so the text needs no copy of its own.
requestContext.BodyText = raw
case mediaType == "application/x-www-form-urlencoded":
storage, storageErr := common.GetBodyStorage(c)
if storageErr != nil {
@@ -933,7 +951,7 @@ func buildTaskPluginRouteRequest(c *gin.Context) (pluginruntime.RouteRequestCont
if !utf8.ValidString(field) || len(field) > maxTaskPluginFieldNameBytes {
return requestContext, fmt.Errorf("invalid multipart file field name")
}
for _, header := range headers {
for index, header := range headers {
if !utf8.ValidString(header.Filename) || len(header.Filename) > maxTaskPluginFilenameBytes {
return requestContext, fmt.Errorf("invalid multipart filename")
}
@@ -947,7 +965,7 @@ func buildTaskPluginRouteRequest(c *gin.Context) (pluginruntime.RouteRequestCont
if header.Size < 0 || header.Size > int64(fileLimitMB)<<20 {
return requestContext, fmt.Errorf("multipart file exceeds %d MB", fileLimitMB)
}
ref := "request_file:" + field
ref := pluginruntime.FileReference(field, index)
files = append(files, map[string]any{"ref": ref, "field": field, "filename": header.Filename, "mimeType": header.Header.Get("Content-Type"), "size": header.Size})
}
}
@@ -1163,10 +1181,10 @@ func renderTaskPluginQuery(
for _, task := range tasks {
tasksByID[task.TaskID] = task
}
views := make([]map[string]any, 0, len(taskIDs))
views := make([]dto.TaskView, 0, len(taskIDs))
for _, taskID := range taskIDs {
task := tasksByID[taskID]
if task == nil {
if task == nil || !task.ResultRetrievable() {
logger.LogDebug(
c,
"task_plugin subsystem=query event=lookup_failed generation=%d plugin=%q reason=task_not_found requested=%d found=%d",
@@ -1183,25 +1201,20 @@ func renderTaskPluginQuery(
abortTaskPluginRouteError(c, http.StatusInternalServerError)
return
}
var viewValue map[string]any
encoded, marshalErr := common.Marshal(view)
if marshalErr != nil {
abortTaskPluginRouteError(c, http.StatusInternalServerError)
return
}
if unmarshalErr := common.Unmarshal(encoded, &viewValue); unmarshalErr != nil {
abortTaskPluginRouteError(c, http.StatusInternalServerError)
return
}
views = append(views, viewValue)
views = append(views, view)
}
var rendererInput any = views
if !multiple {
rendererInput = views[0]
}
encoded, err := common.Marshal(rendererInput)
if err != nil {
abortTaskPluginRouteError(c, http.StatusInternalServerError)
return
}
renderStarted := time.Now()
result, err := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{renderer}, requestContext.JSValue(), rendererInput)
result, err := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{renderer}, requestContext.JSValue(), pluginruntime.RawJSON(encoded))
if err != nil {
logger.LogDebug(
c,
@@ -1264,18 +1277,18 @@ func RespondTaskPluginError(c *gin.Context, taskErr *dto.TaskError) bool {
}
logger.LogWarn(
c,
"task_plugin subsystem=route event=error_renderer_failed plugin=%q reason=hook_failed status=%d err=%q",
pinned.Plugin.Meta.Key,
sanitized.HTTPStatus,
callErr.Error(),
common.LogText("task_plugin subsystem=route event=error_renderer_failed plugin=%q reason=hook_failed status=%d err=%q",
pinned.Plugin.Meta.Key,
sanitized.HTTPStatus,
callErr.Error()),
)
}
logger.LogWarn(
c,
"task_plugin subsystem=route event=error_rendered plugin=%q renderer=host_fallback status=%d code=%q",
pinned.Plugin.Meta.Key,
sanitized.HTTPStatus,
sanitized.Code,
common.LogText("task_plugin subsystem=route event=error_rendered plugin=%q renderer=host_fallback status=%d code=%q",
pinned.Plugin.Meta.Key,
sanitized.HTTPStatus,
sanitized.Code),
)
message := sanitized.Message
if requestID != "" {
@@ -1373,8 +1386,11 @@ func logTaskPluginChannelDecision(c *gin.Context, channel *model.Channel, modelN
return
}
identityMode := "legacy_channel_type"
if channel.Type == constant.ChannelTypeTaskPlugin {
switch channel.Type {
case constant.ChannelTypeTaskPlugin:
identityMode = "type59_setting"
case constant.ChannelTypeNewAPI:
identityMode = "type60_setting"
}
logger.LogDebug(
c,
@@ -1412,7 +1428,7 @@ func PrepareTaskPluginSubmit() gin.HandlerFunc {
)
var requestBody map[string]any
if err := common.UnmarshalBodyReusable(c, &requestBody); err != nil {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": err.Error(), "type": "invalid_request_error"}})
c.AbortWithStatusJSON(invalidRequestStatus(err), gin.H{"error": gin.H{"message": err.Error(), "type": "invalid_request_error"}})
return
}
modelName, _ := requestBody["model"].(string)
@@ -1442,6 +1458,19 @@ func PrepareTaskPluginSubmit() gin.HandlerFunc {
}
}
c.Set("task_request", requestBody)
if plugin.Meta.PreservesJSONOrder() {
// The body as stored after any model rewrite: the text hooks get.
storage, err := common.GetBodyStorage(c)
var raw []byte
if err == nil {
raw, err = storage.Bytes()
}
if err != nil {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": gin.H{"message": err.Error(), "type": "invalid_request_error"}})
return
}
c.Set(pluginruntime.ContextKeyRequestBodyText, json.RawMessage(raw))
}
c.Set("resolved_task_model", modelName)
c.Set("expected_task_plugin_key", pluginKey)
service.AppendTaskPluginIdentityFilter(c, pluginKey)
+251 -1
View File
@@ -3,6 +3,7 @@ package middleware
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"mime/multipart"
@@ -11,6 +12,7 @@ import (
"net/textproto"
"strings"
"testing"
"testing/iotest"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
@@ -97,6 +99,85 @@ export function parseTaskResult() { return {status: "SUCCESS"}; }
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
// Native routes and the generic submit entry hand a plugin declaring
// json-order@1 the client's member order and keep its decoded requestBody as
// ordered JSON text beside the Go value; other plugins see sorted members.
func TestTaskPluginEntriesKeepJSONOrderForDeclaringPlugins(t *testing.T) {
source := `
export const meta = {
apiVersion: 1, key: "route-order-test", name: "Order", version: "1.0.0", author: {name: "Test"},
models: ["m"], fetchMode: "per_task", requiredCapabilities: ["json-order@1"],
routes: [{method: "POST", path: "/vendor/eval", type: "submit", decode: "decodeEval", render: "evalDone"}],
};
export const native = {
decodeEval: function(ctx) {
const state = ctx.body.value.state;
return {kind: "submit", model: ctx.body.value.model, requestBody: {keys: Object.keys(state), state: state}};
},
evalDone: function(ctx, task) { return task; },
};
export function buildSubmitRequest() { return {url: "https://example.com"}; }
export function parseSubmitResponse() { return {taskId: "one"}; }
export function buildQueryRequest() { return {url: "https://example.com"}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`
cases := []struct {
name string
source string
wantText string
wantKeys []any
}{
{name: "declared", source: source, wantText: `{"keys":["zeta","alpha"],"state":{"zeta":1,"alpha":2}}`, wantKeys: []any{"zeta", "alpha"}},
{name: "undeclared", source: strings.Replace(source, `requiredCapabilities: ["json-order@1"],`, "", 1), wantKeys: []any{"alpha", "zeta"}},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
plugin := compileTaskRoutePlugin(t, tc.source)
router := gin.New()
reached := false
router.POST("/vendor/eval", pinTaskPluginRoute(plugin, 0), PrepareTaskPluginRoute(), func(c *gin.Context) {
reached = true
request, ok := c.MustGet("task_request").(map[string]any)
require.True(t, ok)
assert.Equal(t, tc.wantKeys, request["keys"])
text, exists := c.Get(jsplugin.ContextKeyRequestBodyText)
if tc.wantText == "" {
assert.False(t, exists)
} else {
assert.Equal(t, json.RawMessage(tc.wantText), text)
}
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(http.MethodPost, "/vendor/eval", strings.NewReader(`{"model":"m","state":{"zeta":1,"alpha":2}}`))
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.True(t, reached, recorder.Body.String())
})
}
t.Run("generic submit entry", func(t *testing.T) {
declaring := strings.Replace(genericTaskPluginSource, `fetchMode: "per_task"`, `fetchMode: "per_task", requiredCapabilities: ["json-order@1"]`, 1)
_, err := jsplugin.DefaultRegistry.Register(declaring, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("generic-entry-test") })
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Params = gin.Params{{Key: "key", Value: "generic-entry-test"}}
c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/generic-entry-test", strings.NewReader(`{"model":"doc","zeta":1,"alpha":2}`))
c.Request.Header.Set("Content-Type", "application/json")
PrepareTaskPluginSubmit()(c)
require.False(t, c.IsAborted(), recorder.Body.String())
text, exists := c.Get(jsplugin.ContextKeyRequestBodyText)
require.True(t, exists)
assert.Equal(t, json.RawMessage(`{"model":"doc","zeta":1,"alpha":2}`), text)
})
}
func TestPrepareTaskPluginNativeRouteRejectsMultipartBeforeDecoder(t *testing.T) {
plugin := compileTaskRoutePlugin(t, `
export const meta = {
@@ -704,6 +785,100 @@ func TestTaskPluginEndpointUsesOneCanonicalModelForDuplicateJSONKeys(t *testing.
}
}
func TestRequestBodyFailuresReportTheirCause(t *testing.T) {
require.NoError(t, appI18n.Init())
require.NotNil(t, jsplugin.DefaultRegistry.Generation(), "the built-in plugins keep a routing generation, as in production")
_, err := jsplugin.DefaultRegistry.Register(genericTaskPluginSource, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("generic-entry-test") })
routePlugin := compileTaskRoutePlugin(t, genericTaskPluginSource)
previous := constant.MaxRequestBodyMB
constant.MaxRequestBodyMB = 1
t.Cleanup(func() { constant.MaxRequestBodyMB = previous })
reached := func(c *gin.Context) { c.Status(http.StatusNoContent) }
router := gin.New()
router.POST("/v1/responses", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), Distribute(), reached)
router.POST("/v1/video/generations", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), Distribute(), reached)
router.POST("/v1/chat/completions", Distribute(), reached)
router.POST("/mj/submit/imagine", Distribute(), reached)
router.POST("/v1/tasks/:key", PrepareTaskPluginSubmit(), reached)
router.POST("/vendor/jobs", func(c *gin.Context) {
c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Plugin: routePlugin})
}, PrepareTaskPluginRoute(), reached)
oversized := `{"model":"gpt-x","prompt":"` + strings.Repeat("a", 1<<20) + `"}`
tests := []struct {
name string
path string
body io.Reader
status int
message string
}{
{
name: "oversized body on a plugin protocol path",
path: "/v1/responses",
body: strings.NewReader(oversized),
status: http.StatusRequestEntityTooLarge,
message: "Invalid request: request body exceeds 1 MB: request body too large",
},
{
name: "body cut off on a plugin protocol path",
path: "/v1/responses",
body: io.MultiReader(strings.NewReader(`{"model":"gpt-x","input":"`), iotest.ErrReader(io.ErrUnexpectedEOF)),
status: http.StatusBadRequest,
message: "Invalid request: unexpected EOF",
},
{
// The endpoint check lets this path through and the distributor reads the body again.
name: "oversized body on the legacy video path",
path: "/v1/video/generations",
body: strings.NewReader(oversized),
status: http.StatusRequestEntityTooLarge,
message: "Invalid request: request body exceeds 1 MB: request body too large",
},
{
name: "oversized Midjourney body",
path: "/mj/submit/imagine",
body: strings.NewReader(oversized),
status: http.StatusRequestEntityTooLarge,
message: "Invalid request: request body exceeds 1 MB: request body too large",
},
{
name: "oversized task plugin submission",
path: "/v1/tasks/generic-entry-test",
body: strings.NewReader(oversized),
status: http.StatusRequestEntityTooLarge,
message: "request body exceeds 1 MB: request body too large",
},
{
name: "oversized plugin route body",
path: "/vendor/jobs",
body: strings.NewReader(oversized),
status: http.StatusRequestEntityTooLarge,
message: "request body exceeds 1 MB: request body too large",
},
{
name: "invalid JSON on an ordinary path",
path: "/v1/chat/completions",
body: strings.NewReader(`{"model":"gpt-x",`),
status: http.StatusBadRequest,
message: "Invalid request: invalid JSON request body",
},
}
for _, testCase := range tests {
t.Run(testCase.name, func(t *testing.T) {
request := httptest.NewRequest(http.MethodPost, testCase.path, testCase.body)
request.Header.Set("Content-Type", "application/json")
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
assert.Equal(t, testCase.status, recorder.Code)
assert.Contains(t, recorder.Body.String(), `"message":"`+testCase.message)
})
}
}
func TestTaskPluginEndpointOnlyPreservesConditionalMiddlewareSemantics(t *testing.T) {
tests := []struct {
name string
@@ -857,6 +1032,70 @@ func TestPrepareTaskPluginEndpointAcceptsRegisteredVideoMultipartBody(t *testing
assert.Equal(t, http.StatusNoContent, recorder.Code)
}
// The OpenAI Images edits endpoint accepts multipart uploads. Every file of a
// repeated image[] field is exposed to the decoder with its own ref, and an
// unclaimed model on the shared endpoint still reaches the ordinary relay.
func TestPrepareTaskPluginEndpointExposesIndexedRefsForRepeatedImageFiles(t *testing.T) {
const key = "endpoint-image-edits-test"
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
key,
"1.0.0",
`["image-edit-model"]`,
"/v1/images/edits",
`if (ctx.body.kind !== "multipart" || ctx.body.fields.prompt[0] !== "watercolor") throw new Error("bad prompt");
return {model: ctx.model, action: "image_to_image", requestBody: {prompt: ctx.body.fields.prompt[0], refs: ctx.body.files.map(function(file) { return file.ref; })}};`,
), jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("model", "image-edit-model"))
require.NoError(t, writer.WriteField("prompt", "watercolor"))
for _, name := range []string{"first.png", "second.png"} {
file, createErr := writer.CreateFormFile("image[]", name)
require.NoError(t, createErr)
_, err = file.Write([]byte(name))
require.NoError(t, err)
}
require.NoError(t, writer.Close())
router := gin.New()
var pinnedProtocol string
var taskRequest any
router.POST("/v1/images/edits", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) {
if pinned, ok := c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint); ok {
pinnedProtocol = pinned.Protocol
}
taskRequest, _ = c.Get("task_request")
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes()))
request.Header.Set("Content-Type", writer.FormDataContentType())
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusNoContent, recorder.Code, recorder.Body.String())
assert.Equal(t, jsplugin.ProtocolOpenAIImage, pinnedProtocol)
decoded, ok := taskRequest.(map[string]any)
require.True(t, ok)
assert.Equal(t, []any{"request_file:image[]", "request_file:image[]#1"}, decoded["refs"])
unclaimed := gin.New()
reachedRelay := false
unclaimed.POST("/v1/images/generations", PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) {
_, pinned := c.Get(jsplugin.ContextKeyPinnedEndpoint)
reachedRelay = !pinned
c.Status(http.StatusNoContent)
})
plain := httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(`{"model":"dall-e-3","prompt":"a cat"}`))
plain.Header.Set("Content-Type", "application/json")
recorder = httptest.NewRecorder()
unclaimed.ServeHTTP(recorder, plain)
assert.Equal(t, http.StatusNoContent, recorder.Code)
assert.True(t, reachedRelay, "an unclaimed image model keeps the built-in relay")
}
func TestVideoGenerationsIsNotClaimedByOpenAIVideoProtocol(t *testing.T) {
const key = "endpoint-video-gen-test"
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
@@ -1065,6 +1304,12 @@ export const native = {status: function(ctx, task) { return {id: task.task_id};
insertTaskPluginRouteTask(t, &model.Task{
TaskID: "legacy-task", UserId: 7, Platform: constant.TaskPlatform("651"),
})
// Submitted through a route with retainResult: false; owned by the caller
// but not retrievable, so it must be indistinguishable from a missing task.
insertTaskPluginRouteTask(t, &model.Task{
TaskID: "discarded-task", UserId: 7, Platform: constant.TaskPlatform("route-static-query-test"),
Status: model.TaskStatusSuccess, PrivateData: model.TaskPrivateData{ResultDiscarded: true},
})
router := gin.New()
router.GET("/vendor/jobs/:id",
@@ -1082,7 +1327,7 @@ export const native = {status: function(ctx, task) { return {id: task.task_id};
assert.JSONEq(t, `{"id":"legacy-task"}`, legacyRecorder.Body.String())
var firstBody string
for _, taskID := range []string{"missing-task", "foreign-task", "wrong-platform", "wrong-legacy-platform"} {
for _, taskID := range []string{"missing-task", "foreign-task", "wrong-platform", "wrong-legacy-platform", "discarded-task"} {
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, httptest.NewRequest(http.MethodGet, "/vendor/jobs/"+taskID, nil))
assert.Equal(t, http.StatusNotFound, recorder.Code)
@@ -1637,6 +1882,11 @@ func taskProtocolPluginSource(key, version, models, endpoint, parseRequestBody s
protocolClaim = `"openai_video"`
presenters = `render: function() { return {}; },`
}
if endpoint == "/v1/images/edits" || endpoint == "/v1/images/generations" {
protocol = "openai_image"
protocolClaim = `"openai_image"`
presenters = `render: function() { return {data: []}; },`
}
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
+1 -1
View File
@@ -14,7 +14,7 @@ func ConfigureTrustedProxies(engine *gin.Engine) error {
return err
}
if usedDefaults {
log.Print("WARNING: TRUSTED_PROXIES is unset or blank; trusting loopback, RFC 1918, and IPv6 ULA proxy addresses for compatibility. Set TRUSTED_PROXIES=none to trust no proxies, or configure explicit proxy IPs/CIDRs to replace these defaults.")
log.Print(common.LogText("WARNING: TRUSTED_PROXIES is unset or blank; trusting loopback, RFC 1918, and IPv6 ULA proxy addresses for compatibility. Set TRUSTED_PROXIES=none to trust no proxies, or configure explicit proxy IPs/CIDRs to replace these defaults."))
}
return common.ConfigureTrustedProxies(engine, trustedProxies)
}
+2 -8
View File
@@ -17,10 +17,7 @@ func TurnstileCheck() gin.HandlerFunc {
if common.TurnstileCheckEnabled {
response := c.Query("turnstile")
if response == "" {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "Turnstile token 为空",
})
common.ApiErrorT(c, "Turnstile token is empty")
c.Abort()
return
}
@@ -51,10 +48,7 @@ func TurnstileCheck() gin.HandlerFunc {
return
}
if !res.Success {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "Turnstile 校验失败,请刷新重试!",
})
common.ApiErrorT(c, "Turnstile verification failed, please refresh and try again!")
c.Abort()
return
}
+11 -2
View File
@@ -1,7 +1,7 @@
package middleware
import (
"fmt"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/dto"
@@ -32,7 +32,16 @@ func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code
})
}
c.Abort()
logger.LogError(c.Request.Context(), fmt.Sprintf("user %d | %s", userId, message))
logger.LogError(c.Request.Context(), common.LogText("user %d | %s", userId, message))
}
// invalidRequestStatus is the status of a request whose body could not be
// read or parsed: 413 when the body exceeds MAX_REQUEST_BODY_MB, else 400.
func invalidRequestStatus(err error) int {
if common.IsRequestBodyTooLargeError(err) {
return http.StatusRequestEntityTooLarge
}
return http.StatusBadRequest
}
func abortWithMidjourneyMessage(c *gin.Context, statusCode int, code int, description string) {