feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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
@@ -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"))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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) {
|
||||
|
||||
Reference in new issue
Block a user