feat: publish v1.0.0-rc.36-custom.1
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 9m20s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 55m7s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 9m20s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 55m7s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s
This commit is contained in:
commit
bb7968b612
2381 files changed
+551494
No files matched your search
@@ -0,0 +1,306 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// auditResponseWriter 包装 gin.ResponseWriter,捕获响应状态码并将响应体复制一份到
|
||||
// 有限大小的缓冲区,用于判断业务是否成功(解析响应 JSON 的 success 字段)。
|
||||
// 缓冲区有上限,避免大响应(如密钥导出)占用过多内存;超出上限则不再缓存,
|
||||
// 此时仅依据 HTTP 状态码判断成败。
|
||||
type auditResponseWriter struct {
|
||||
gin.ResponseWriter
|
||||
body *bytes.Buffer
|
||||
maxSize int
|
||||
}
|
||||
|
||||
func (w *auditResponseWriter) Write(b []byte) (int, error) {
|
||||
if w.body.Len() < w.maxSize {
|
||||
remain := w.maxSize - w.body.Len()
|
||||
if remain >= len(b) {
|
||||
w.body.Write(b)
|
||||
} else {
|
||||
w.body.Write(b[:remain])
|
||||
}
|
||||
}
|
||||
return w.ResponseWriter.Write(b)
|
||||
}
|
||||
|
||||
func (w *auditResponseWriter) WriteString(s string) (int, error) {
|
||||
return w.Write([]byte(s))
|
||||
}
|
||||
|
||||
// auditRouteActions 将「METHOD + 路由模板」映射为语言无关的操作标识 action。
|
||||
// 这些是未被 handler 手动埋点的写操作,由中间件兜底记录;前端依据 action 用 i18n 本地化展示。
|
||||
// 未命中的写操作回退为 action="generic",前端展示 "METHOD route"。
|
||||
var auditRouteActions = map[string]string{
|
||||
// 用户管理
|
||||
"POST /api/user/topup/complete": "user.topup_complete",
|
||||
"DELETE /api/user/:id/reset_passkey": "user.reset_passkey",
|
||||
"DELETE /api/user/:id/oauth/bindings/:provider_id": "user.oauth_unbind",
|
||||
|
||||
// 系统设置(root)
|
||||
"POST /api/option/payment_compliance": "option.payment_compliance",
|
||||
"POST /api/option/rest_model_ratio": "option.reset_ratio",
|
||||
"DELETE /api/option/channel_affinity_cache": "option.clear_affinity_cache",
|
||||
|
||||
// 自定义 OAuth(root)
|
||||
"POST /api/custom-oauth-provider/": "custom_oauth.create",
|
||||
"PUT /api/custom-oauth-provider/:id": "custom_oauth.update",
|
||||
"DELETE /api/custom-oauth-provider/:id": "custom_oauth.delete",
|
||||
|
||||
// 性能/缓存(root)
|
||||
"DELETE /api/performance/disk_cache": "performance.clear_disk_cache",
|
||||
"POST /api/performance/gc": "performance.gc",
|
||||
"DELETE /api/performance/logs": "performance.clear_logs",
|
||||
|
||||
// 兑换码
|
||||
"PUT /api/redemption/": "redemption.update",
|
||||
"POST /api/redemption/batch": "redemption.delete_batch",
|
||||
"DELETE /api/redemption/:id": "redemption.delete",
|
||||
"DELETE /api/redemption/invalid": "redemption.delete_invalid",
|
||||
|
||||
// 预填组
|
||||
"POST /api/prefill_group/": "prefill_group.create",
|
||||
"PUT /api/prefill_group/": "prefill_group.update",
|
||||
"DELETE /api/prefill_group/:id": "prefill_group.delete",
|
||||
|
||||
// 供应商
|
||||
"POST /api/vendors/": "vendor.create",
|
||||
"PUT /api/vendors/": "vendor.update",
|
||||
"DELETE /api/vendors/:id": "vendor.delete",
|
||||
|
||||
// 模型元数据
|
||||
"POST /api/models/": "model.create",
|
||||
"PUT /api/models/": "model.update",
|
||||
"DELETE /api/models/:id": "model.delete",
|
||||
"POST /api/models/sync_upstream": "model.sync_upstream",
|
||||
|
||||
// 部署
|
||||
"POST /api/deployments/": "deployment.create",
|
||||
"PUT /api/deployments/:id": "deployment.update",
|
||||
"DELETE /api/deployments/:id": "deployment.delete",
|
||||
|
||||
// 订阅(管理员)
|
||||
"POST /api/subscription/admin/plans": "subscription.plan_create",
|
||||
"PUT /api/subscription/admin/plans/:id": "subscription.plan_update",
|
||||
"POST /api/subscription/admin/bind": "subscription.bind",
|
||||
|
||||
// 日志
|
||||
"POST /api/system-task/log-cleanup": "log.cleanup_start",
|
||||
}
|
||||
|
||||
// beginAdminAudit 在管理/root 写操作进入 handler 前包装 ResponseWriter,
|
||||
// 以便事后解析响应判断业务是否成功。仅对写方法(POST/PUT/PATCH/DELETE)生效;
|
||||
// 只读请求返回 nil,调用方据此跳过事后兜底记录。
|
||||
//
|
||||
// 该函数由 authHelper 在鉴权通过、c.Next() 之前调用:因为任何管理/root 接口都
|
||||
// 必然经过 AdminAuth/RootAuth,将审计兜底内聚到鉴权链路即可保证「新增接口自动留痕」,
|
||||
// 无需在路由上再单独挂一层审计中间件(避免漏挂)。
|
||||
func beginAdminAudit(c *gin.Context) *auditResponseWriter {
|
||||
method := c.Request.Method
|
||||
if method != "POST" && method != "PUT" && method != "PATCH" && method != "DELETE" {
|
||||
return nil
|
||||
}
|
||||
writer := &auditResponseWriter{
|
||||
ResponseWriter: c.Writer,
|
||||
body: bytes.NewBuffer(nil),
|
||||
maxSize: 64 * 1024,
|
||||
}
|
||||
c.Writer = writer
|
||||
return writer
|
||||
}
|
||||
|
||||
// finishAdminAudit 在 c.Next() 之后对管理/高危写操作做兜底审计记录。
|
||||
// 若 handler 内已手动埋点(设置 ContextKeyAuditLogged),则跳过,避免重复。
|
||||
func finishAdminAudit(c *gin.Context, writer *auditResponseWriter) {
|
||||
if writer == nil {
|
||||
return
|
||||
}
|
||||
method := c.Request.Method
|
||||
|
||||
// handler 已手动记录更精细的审计日志,跳过兜底。
|
||||
if common.GetContextKeyBool(c, constant.ContextKeyAuditLogged) {
|
||||
return
|
||||
}
|
||||
|
||||
operatorId := c.GetInt("id")
|
||||
operatorName := c.GetString("username")
|
||||
operatorRole := c.GetInt("role")
|
||||
ip := c.ClientIP()
|
||||
status := writer.Status()
|
||||
success := auditResponseSuccess(status, writer.body.Bytes())
|
||||
|
||||
route := c.FullPath()
|
||||
action := auditRouteActions[method+" "+route]
|
||||
if action == "" {
|
||||
action = "generic"
|
||||
}
|
||||
|
||||
routeParams := map[string]string{}
|
||||
for _, p := range c.Params {
|
||||
routeParams[p.Key] = p.Value
|
||||
}
|
||||
|
||||
// op.params 为语言无关参数,供前端 i18n 渲染;generic 时携带 method/route。
|
||||
opParams := map[string]any{}
|
||||
if action == "generic" {
|
||||
opParams["method"] = method
|
||||
opParams["route"] = route
|
||||
}
|
||||
|
||||
// content 为英文兜底文本(供导出等非本地化消费者使用)。
|
||||
content := method + " " + route
|
||||
|
||||
adminInfo := &model.AuditAdminInfo{
|
||||
AdminID: operatorId,
|
||||
AdminUsername: operatorName,
|
||||
AdminRole: operatorRole,
|
||||
AuthMethod: auditAuthMethod(c),
|
||||
}
|
||||
auditInfo := &model.AuditRequestInfo{
|
||||
Method: method,
|
||||
Route: route,
|
||||
Path: route,
|
||||
Status: status,
|
||||
Success: success,
|
||||
Params: routeParams,
|
||||
}
|
||||
|
||||
model.RecordOperationAuditLog(operatorId, operatorRole, content, ip, action, opParams, adminInfo, auditInfo, c)
|
||||
}
|
||||
|
||||
func auditAuthMethod(c *gin.Context) string {
|
||||
if c.GetBool("use_access_token") {
|
||||
return "access_token"
|
||||
}
|
||||
return "session"
|
||||
}
|
||||
|
||||
// auditResponseSuccess 依据 HTTP 状态码与响应体推断操作是否成功。
|
||||
// 优先解析响应 JSON 中的 success 字段;无法解析时退回到状态码判断。
|
||||
func auditResponseSuccess(status int, body []byte) bool {
|
||||
if status >= 400 {
|
||||
return false
|
||||
}
|
||||
trimmed := bytes.TrimSpace(body)
|
||||
if len(trimmed) > 0 && trimmed[0] == '{' {
|
||||
var resp struct {
|
||||
Success *bool `json:"success"`
|
||||
}
|
||||
if err := common.Unmarshal(trimmed, &resp); err == nil && resp.Success != nil {
|
||||
return *resp.Success
|
||||
}
|
||||
}
|
||||
return status < 400
|
||||
}
|
||||
|
||||
const accessTokenAuditContextKey = "access_token_request_audit"
|
||||
|
||||
// TokenOperationAudit runs after UserAuth and before endpoint rate limits.
|
||||
// Handlers add only allowlisted metadata; neither bodies nor raw errors are persisted.
|
||||
func TokenOperationAudit() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
var action, content string
|
||||
switch c.Request.Method + " " + c.FullPath() {
|
||||
case "POST /api/token/":
|
||||
action, content = "token.create", "API token creation"
|
||||
case "PUT /api/token/":
|
||||
action, content = "token.update", "API token configuration update"
|
||||
if c.Query("status_only") != "" {
|
||||
action, content = "token.status_update", "API token status update"
|
||||
}
|
||||
case "DELETE /api/token/:id":
|
||||
action, content = "token.delete", "API token deletion"
|
||||
case "POST /api/token/batch":
|
||||
action, content = "token.delete_batch", "API token batch deletion"
|
||||
case "POST /api/token/:id/key":
|
||||
action, content = "token.key_view", "API token key access"
|
||||
case "POST /api/token/batch/keys":
|
||||
action, content = "token.key_view_batch", "API token batch key access"
|
||||
default:
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
params := model.AuditFields{}
|
||||
if id, err := strconv.Atoi(c.Param("id")); err == nil && id > 0 {
|
||||
params["id"] = id
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyTokenAuditParams, params)
|
||||
entry := model.AuditLog{
|
||||
UserId: c.GetInt("id"), Username: c.GetString("username"), ActorRole: c.GetInt("role"),
|
||||
Category: model.AuditCategorySecurity, Action: action, Content: content,
|
||||
Other: model.AuditOther{Op: &model.AuditOperation{Action: action, Params: params}},
|
||||
}
|
||||
writer := &auditResponseWriter{ResponseWriter: c.Writer, body: bytes.NewBuffer(nil), maxSize: 64 * 1024}
|
||||
c.Writer = writer
|
||||
c.Next()
|
||||
entry.Status = writer.Status()
|
||||
entry.Success = auditResponseSuccess(entry.Status, writer.body.Bytes())
|
||||
if writer.body.Len() == writer.maxSize {
|
||||
// JSON may be truncated before its success field. A completed handler
|
||||
// supplies the result without retaining an unbounded response body.
|
||||
entry.Success = entry.Status < 400 && common.GetContextKeyBool(c, constant.ContextKeyTokenAuditSucceeded)
|
||||
}
|
||||
model.RecordAuditLog(c, entry)
|
||||
}
|
||||
}
|
||||
|
||||
type accessTokenRequestAudit struct {
|
||||
entry model.AuditLog
|
||||
writer *auditResponseWriter
|
||||
}
|
||||
|
||||
// AccessTokenAudit also captures public reads and rejections before route-level
|
||||
// authentication (for example, rate limiting). It does not grant authentication.
|
||||
func AccessTokenAudit() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
raw, present := authorizationToken(c.GetHeader("Authorization"))
|
||||
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)
|
||||
defer finishAccessTokenAudit(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func beginAccessTokenAudit(c *gin.Context, user *model.User, token 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()},
|
||||
writer: writer,
|
||||
})
|
||||
}
|
||||
|
||||
func finishAccessTokenAudit(c *gin.Context) {
|
||||
value, exists := c.Get(accessTokenAuditContextKey)
|
||||
if !exists {
|
||||
return
|
||||
}
|
||||
audit, ok := value.(*accessTokenRequestAudit)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
audit.entry.Status = audit.writer.Status()
|
||||
audit.entry.Success = auditResponseSuccess(audit.entry.Status, audit.writer.body.Bytes())
|
||||
model.RecordAuditLog(c, audit.entry)
|
||||
}
|
||||
@@ -0,0 +1,545 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"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"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/service/authz"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
const authIdentityContextKey = "auth_identity"
|
||||
|
||||
type dashboardCredentialKind int
|
||||
|
||||
const (
|
||||
dashboardCredentialUnmatched dashboardCredentialKind = iota
|
||||
dashboardCredentialInternal
|
||||
dashboardCredentialPAT
|
||||
)
|
||||
|
||||
func validUserInfo(username string, role int) bool {
|
||||
// check username is empty
|
||||
if strings.TrimSpace(username) == "" {
|
||||
return false
|
||||
}
|
||||
if !common.IsValidateRole(role) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func authHelper(c *gin.Context, minRole int) {
|
||||
if _, started := c.Get(accessTokenAuditContextKey); !started {
|
||||
defer finishAccessTokenAudit(c)
|
||||
}
|
||||
user, identity, useAccessToken, err := authenticateDashboardRequest(c)
|
||||
if err != nil {
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if user.Status != common.UserStatusEnabled {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_USER_DISABLED", "message": common.TranslateMessage(c, i18n.MsgAuthUserBanned)})
|
||||
return
|
||||
}
|
||||
if user.Role < minRole {
|
||||
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"success": false, "code": "AUTH_INSUFFICIENT_PRIVILEGE", "message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege)})
|
||||
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)})
|
||||
return
|
||||
}
|
||||
setDashboardAuthContext(c, user, identity, useAccessToken)
|
||||
|
||||
// 管理/root 写操作审计兜底:内聚在鉴权链路里,保证任何经过 AdminAuth/RootAuth
|
||||
// 的写接口都会自动留痕(无需在路由上单独挂审计中间件,避免漏挂)。
|
||||
// handler 内手动埋点者会设置 ContextKeyAuditLogged,finishAdminAudit 据此跳过。
|
||||
var auditWriter *auditResponseWriter
|
||||
if minRole >= common.RoleAdminUser {
|
||||
auditWriter = beginAdminAudit(c)
|
||||
}
|
||||
|
||||
c.Next()
|
||||
|
||||
finishAdminAudit(c, auditWriter)
|
||||
}
|
||||
|
||||
func TryUserAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
if _, started := c.Get(accessTokenAuditContextKey); !started {
|
||||
defer finishAccessTokenAudit(c)
|
||||
}
|
||||
user, identity, credentialKind, err := classifyDashboardCredential(c)
|
||||
if err != nil {
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if credentialKind != dashboardCredentialUnmatched {
|
||||
setDashboardAuthContext(c, user, identity, credentialKind == dashboardCredentialPAT)
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func UserAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
authHelper(c, common.RoleCommonUser)
|
||||
}
|
||||
}
|
||||
|
||||
func AdminAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
authHelper(c, common.RoleAdminUser)
|
||||
}
|
||||
}
|
||||
|
||||
func RootAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
authHelper(c, common.RoleRootUser)
|
||||
}
|
||||
}
|
||||
|
||||
// GetAuthIdentity returns a dashboard session identity. PAT-authenticated
|
||||
// requests intentionally have no SessionID and cannot manage browser sessions.
|
||||
func GetAuthIdentity(c *gin.Context) (service.AuthIdentity, bool) {
|
||||
value, ok := c.Get(authIdentityContextKey)
|
||||
if !ok {
|
||||
return service.AuthIdentity{}, false
|
||||
}
|
||||
identity, ok := value.(service.AuthIdentity)
|
||||
return identity, ok
|
||||
}
|
||||
|
||||
// GetSessionAuthIdentity returns only identities backed by a live dashboard
|
||||
// session. PAT-authenticated requests intentionally fail this check.
|
||||
func GetSessionAuthIdentity(c *gin.Context) (service.AuthIdentity, bool) {
|
||||
identity, ok := GetAuthIdentity(c)
|
||||
if !ok {
|
||||
identity = service.AuthIdentity{
|
||||
UserID: c.GetInt("id"),
|
||||
SessionID: c.GetString("session_id"),
|
||||
UserAuthVersion: c.GetInt64("auth_version"),
|
||||
SessionVersion: c.GetInt64("session_version"),
|
||||
}
|
||||
}
|
||||
if identity.UserID <= 0 || identity.SessionID == "" || identity.UserAuthVersion <= 0 || identity.SessionVersion <= 0 {
|
||||
return service.AuthIdentity{}, false
|
||||
}
|
||||
return identity, true
|
||||
}
|
||||
|
||||
func authenticateDashboardRequest(c *gin.Context) (*model.UserBase, service.AuthIdentity, bool, error) {
|
||||
user, identity, credentialKind, err := classifyDashboardCredential(c)
|
||||
if err != nil {
|
||||
return nil, service.AuthIdentity{}, credentialKind == dashboardCredentialPAT, err
|
||||
}
|
||||
if credentialKind == dashboardCredentialUnmatched {
|
||||
return nil, service.AuthIdentity{}, false, service.ErrAuthTokenInvalid
|
||||
}
|
||||
return user, identity, credentialKind == dashboardCredentialPAT, nil
|
||||
}
|
||||
|
||||
func classifyDashboardCredential(c *gin.Context) (*model.UserBase, service.AuthIdentity, dashboardCredentialKind, error) {
|
||||
raw, ok := authorizationToken(c.GetHeader("Authorization"))
|
||||
if !ok {
|
||||
return nil, service.AuthIdentity{}, dashboardCredentialUnmatched, nil
|
||||
}
|
||||
identity, internal, err := service.ParseDashboardAccessToken(raw)
|
||||
if internal {
|
||||
if err != nil {
|
||||
return nil, service.AuthIdentity{}, dashboardCredentialInternal, err
|
||||
}
|
||||
_, user, err := service.ValidateLoginSession(identity)
|
||||
if err != nil {
|
||||
return nil, service.AuthIdentity{}, dashboardCredentialInternal, err
|
||||
}
|
||||
return user, identity, dashboardCredentialInternal, nil
|
||||
}
|
||||
patUser, err := model.ValidateAccessToken(raw)
|
||||
if err != nil {
|
||||
return nil, service.AuthIdentity{}, dashboardCredentialPAT, err
|
||||
}
|
||||
if patUser == nil || patUser.Id <= 0 {
|
||||
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 user, service.AuthIdentity{UserID: user.Id, UserAuthVersion: user.AuthVersion}, dashboardCredentialPAT, nil
|
||||
}
|
||||
|
||||
func authorizationToken(header string) (string, bool) {
|
||||
header = strings.TrimSpace(header)
|
||||
if header == "" {
|
||||
return "", false
|
||||
}
|
||||
parts := strings.Fields(header)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
header = parts[1]
|
||||
} else if len(parts) != 1 {
|
||||
return "", false
|
||||
}
|
||||
return header, header != ""
|
||||
}
|
||||
|
||||
func setDashboardAuthContext(c *gin.Context, user *model.UserBase, identity service.AuthIdentity, useAccessToken bool) {
|
||||
c.Header("Auth-Version", "864b7076dbcd0a3c01b5520316720ebf")
|
||||
c.Set("username", user.Username)
|
||||
c.Set("role", user.Role)
|
||||
c.Set("id", user.Id)
|
||||
c.Set("group", user.Group)
|
||||
c.Set("user_group", user.Group)
|
||||
c.Set("use_access_token", useAccessToken)
|
||||
c.Set("session_id", identity.SessionID)
|
||||
c.Set("auth_version", identity.UserAuthVersion)
|
||||
c.Set("session_version", identity.SessionVersion)
|
||||
c.Set(authIdentityContextKey, identity)
|
||||
user.WriteContext(c)
|
||||
}
|
||||
|
||||
func writeDashboardAuthError(c *gin.Context, err error) {
|
||||
if errors.Is(err, service.ErrAuthTokenExpired) {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_TOKEN_EXPIRED", "message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn)})
|
||||
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)})
|
||||
return
|
||||
}
|
||||
if errors.Is(err, service.ErrAuthTokenInvalid) {
|
||||
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"success": false, "code": "AUTH_UNAUTHORIZED", "message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid)})
|
||||
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)})
|
||||
}
|
||||
|
||||
func RequirePermission(permission authz.Permission) func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
role := c.GetInt("role")
|
||||
userID := c.GetInt("id")
|
||||
if authz.Can(userID, role, permission) {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege),
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
func WssAuth(c *gin.Context) {
|
||||
|
||||
}
|
||||
|
||||
// TokenOrUserAuth allows either session-based user auth or API token auth.
|
||||
// Used for endpoints that need to be accessible from both the dashboard and API clients.
|
||||
func TokenOrUserAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
raw, ok := authorizationToken(c.GetHeader("Authorization"))
|
||||
if ok {
|
||||
identity, internal, err := service.ParseDashboardAccessToken(raw)
|
||||
if !internal {
|
||||
TokenAuth()(c)
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
_, user, err := service.ValidateLoginSession(identity)
|
||||
if err != nil {
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
setDashboardAuthContext(c, user, identity, false)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
// Opaque credentials are relay API keys here, never dashboard PATs.
|
||||
TokenAuth()(c)
|
||||
}
|
||||
}
|
||||
|
||||
// TokenAuthReadOnly 宽松版本的令牌认证中间件,用于只读查询接口。
|
||||
// 只验证令牌 key 是否存在,不检查令牌状态、过期时间和额度。
|
||||
// 即使令牌已过期、已耗尽或已禁用,也允许访问。
|
||||
// 仍然检查用户是否被封禁。
|
||||
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),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
|
||||
key = strings.TrimSpace(key[7:])
|
||||
}
|
||||
key = strings.TrimPrefix(key, "sk-")
|
||||
parts := strings.Split(key, "-")
|
||||
key = parts[0]
|
||||
|
||||
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),
|
||||
})
|
||||
} else {
|
||||
common.SysLog("TokenAuthReadOnly GetTokenByKey database error: " + err.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
|
||||
})
|
||||
}
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 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),
|
||||
})
|
||||
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),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if userCache.Status != common.UserStatusEnabled {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserBanned),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Set("id", token.UserId)
|
||||
c.Set("token_id", token.Id)
|
||||
c.Set("token_key", token.Key)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func TokenAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
// 先检测是否为ws
|
||||
if c.Request.Header.Get("Sec-WebSocket-Protocol") != "" {
|
||||
// Sec-WebSocket-Protocol: realtime, openai-insecure-api-key.sk-xxx, openai-beta.realtime-v1
|
||||
// read sk from Sec-WebSocket-Protocol
|
||||
key := c.Request.Header.Get("Sec-WebSocket-Protocol")
|
||||
parts := strings.SplitSeq(key, ",")
|
||||
for part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if strings.HasPrefix(part, "openai-insecure-api-key") {
|
||||
key = strings.TrimPrefix(part, "openai-insecure-api-key.")
|
||||
break
|
||||
}
|
||||
}
|
||||
c.Request.Header.Set("Authorization", "Bearer "+key)
|
||||
}
|
||||
// 检查path包含/v1/messages 或 /v1/models
|
||||
if strings.Contains(c.Request.URL.Path, "/v1/messages") || strings.Contains(c.Request.URL.Path, "/v1/models") {
|
||||
anthropicKey := c.Request.Header.Get("x-api-key")
|
||||
if anthropicKey != "" {
|
||||
c.Request.Header.Set("Authorization", "Bearer "+anthropicKey)
|
||||
}
|
||||
}
|
||||
// gemini api 从query中获取key
|
||||
if c.Request.URL.Path == "/v1/models" ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1beta/models") ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1beta/openai/models") ||
|
||||
strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
|
||||
skKey := c.Query("key")
|
||||
if skKey != "" {
|
||||
c.Request.Header.Set("Authorization", "Bearer "+skKey)
|
||||
}
|
||||
// 从x-goog-api-key header中获取key
|
||||
xGoogKey := c.Request.Header.Get("x-goog-api-key")
|
||||
if xGoogKey != "" {
|
||||
c.Request.Header.Set("Authorization", "Bearer "+xGoogKey)
|
||||
}
|
||||
}
|
||||
key := c.Request.Header.Get("Authorization")
|
||||
parts := make([]string, 0)
|
||||
if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
|
||||
key = strings.TrimSpace(key[7:])
|
||||
}
|
||||
if key == "" || key == "midjourney-proxy" {
|
||||
key = c.Request.Header.Get("mj-api-secret")
|
||||
if strings.HasPrefix(key, "Bearer ") || strings.HasPrefix(key, "bearer ") {
|
||||
key = strings.TrimSpace(key[7:])
|
||||
}
|
||||
key = strings.TrimPrefix(key, "sk-")
|
||||
parts = strings.Split(key, "-")
|
||||
key = parts[0]
|
||||
} else {
|
||||
key = strings.TrimPrefix(key, "sk-")
|
||||
parts = strings.Split(key, "-")
|
||||
key = parts[0]
|
||||
}
|
||||
token, err := model.ValidateUserToken(key)
|
||||
if token != nil {
|
||||
id := c.GetInt("id")
|
||||
if id == 0 {
|
||||
c.Set("id", token.UserId)
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, model.ErrDatabase) {
|
||||
common.SysLog("TokenAuth ValidateUserToken database error: " + err.Error())
|
||||
abortWithOpenAiMessage(c, http.StatusInternalServerError,
|
||||
common.TranslateMessage(c, i18n.MsgDatabaseError))
|
||||
} else {
|
||||
abortWithOpenAiMessage(c, http.StatusUnauthorized,
|
||||
common.TranslateMessage(c, i18n.MsgTokenInvalid))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
allowIps := token.GetIpLimits()
|
||||
if len(allowIps) > 0 {
|
||||
clientIp := c.ClientIP()
|
||||
logger.LogDebug(c, "Token has IP restrictions, checking client IP %s", clientIp)
|
||||
ip := net.ParseIP(clientIp)
|
||||
if ip == nil {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, "无法解析客户端 IP 地址")
|
||||
return
|
||||
}
|
||||
if common.IsIpInCIDRList(ip, allowIps) == false {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, "您的 IP 不在令牌允许访问的列表中", types.ErrorCodeAccessDenied)
|
||||
return
|
||||
}
|
||||
logger.LogDebug(c, "Client IP %s passed the token IP restrictions check", clientIp)
|
||||
}
|
||||
|
||||
userCache, err := model.GetUserCache(token.UserId)
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("TokenAuth GetUserCache error for user %d: %v", token.UserId, err))
|
||||
abortWithOpenAiMessage(c, http.StatusInternalServerError,
|
||||
common.TranslateMessage(c, i18n.MsgDatabaseError))
|
||||
return
|
||||
}
|
||||
userEnabled := userCache.Status == common.UserStatusEnabled
|
||||
if !userEnabled {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, common.TranslateMessage(c, i18n.MsgAuthUserBanned))
|
||||
return
|
||||
}
|
||||
|
||||
userCache.WriteContext(c)
|
||||
|
||||
userGroup := userCache.Group
|
||||
tokenGroup := token.Group
|
||||
if tokenGroup != "" {
|
||||
// check common.UserUsableGroups[userGroup]
|
||||
if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("无权访问 %s 分组", tokenGroup))
|
||||
return
|
||||
}
|
||||
// check group in common.GroupRatio
|
||||
if !ratio_setting.ContainsGroupRatio(tokenGroup) {
|
||||
if tokenGroup != "auto" {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("分组 %s 已被弃用", tokenGroup))
|
||||
return
|
||||
}
|
||||
}
|
||||
userGroup = tokenGroup
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyUsingGroup, userGroup)
|
||||
|
||||
err = SetupContextForToken(c, token, parts...)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) error {
|
||||
if token == nil {
|
||||
return fmt.Errorf("token is nil")
|
||||
}
|
||||
c.Set("id", token.UserId)
|
||||
c.Set("token_id", token.Id)
|
||||
c.Set("token_key", token.Key)
|
||||
c.Set("token_name", token.Name)
|
||||
c.Set("token_unlimited_quota", token.UnlimitedQuota)
|
||||
if !token.UnlimitedQuota {
|
||||
c.Set("token_quota", token.RemainQuota)
|
||||
}
|
||||
if token.ModelLimitsEnabled {
|
||||
c.Set("token_model_limit_enabled", true)
|
||||
c.Set("token_model_limit", token.GetModelLimitsMap())
|
||||
} else {
|
||||
c.Set("token_model_limit_enabled", false)
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group)
|
||||
common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry)
|
||||
if token.AutoGroups != "" {
|
||||
autoGroups, err := token.GetAutoGroups()
|
||||
if err != nil {
|
||||
common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err))
|
||||
autoGroups = []string{}
|
||||
common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups)
|
||||
} else if len(autoGroups) > 0 {
|
||||
common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups)
|
||||
}
|
||||
}
|
||||
if len(parts) > 1 {
|
||||
if model.IsAdmin(token.UserId) {
|
||||
id, err := strconv.Atoi(parts[1])
|
||||
if err != nil {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
|
||||
return fmt.Errorf("invalid specific channel id")
|
||||
}
|
||||
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
||||
ChannelId: id,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
} else {
|
||||
c.Header("specific_channel_version", "701e3ae1dc3f7975556d354e0675168d004891c8")
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, "普通用户不支持指定渠道")
|
||||
return fmt.Errorf("普通用户不支持指定渠道")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"crypto/subtle"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// SessionCookieOriginGuard protects cookie-authenticated refresh/logout
|
||||
// endpoints when secure cookie mode is enabled. In insecure local development
|
||||
// mode it preserves the legacy behavior and intentionally performs no Origin
|
||||
// validation. It never adds CORS response headers and must not be installed on
|
||||
// relay routes.
|
||||
func SessionCookieOriginGuard() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if !common.SessionCookieSecure {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
origin, ok := requestBrowserOrigin(c.Request)
|
||||
if !ok || !isAllowedSessionOrigin(c.Request, origin) {
|
||||
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"code": "AUTH_ORIGIN_FORBIDDEN",
|
||||
"message": "request origin is not allowed",
|
||||
})
|
||||
return
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func requestBrowserOrigin(request *http.Request) (string, bool) {
|
||||
originValues := request.Header.Values("Origin")
|
||||
if len(originValues) > 1 {
|
||||
return "", false
|
||||
}
|
||||
if len(originValues) == 1 {
|
||||
if strings.Contains(originValues[0], ",") {
|
||||
return "", false
|
||||
}
|
||||
origin, err := common.NormalizeOrigin(originValues[0])
|
||||
return origin, err == nil
|
||||
}
|
||||
refererValues := request.Header.Values("Referer")
|
||||
if len(refererValues) != 1 {
|
||||
return "", false
|
||||
}
|
||||
referer, err := url.Parse(strings.TrimSpace(refererValues[0]))
|
||||
if err != nil || referer.Scheme == "" || referer.Host == "" || referer.User != nil {
|
||||
return "", false
|
||||
}
|
||||
origin, err := common.NormalizeOrigin(referer.Scheme + "://" + referer.Host)
|
||||
return origin, err == nil
|
||||
}
|
||||
|
||||
func isAllowedSessionOrigin(request *http.Request, origin string) bool {
|
||||
requestScheme := "http"
|
||||
if request.TLS != nil {
|
||||
requestScheme = "https"
|
||||
}
|
||||
requestOrigin, err := common.NormalizeOrigin(requestScheme + "://" + request.Host)
|
||||
if err == nil && subtle.ConstantTimeCompare([]byte(origin), []byte(requestOrigin)) == 1 {
|
||||
return true
|
||||
}
|
||||
for _, trustedOrigin := range common.SessionCookieTrustedURLs {
|
||||
if subtle.ConstantTimeCompare([]byte(origin), []byte(trustedOrigin)) == 1 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func runOriginGuardRequest(t *testing.T, origin, referer string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.POST("/api/user/auth/refresh", SessionCookieOriginGuard(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "https://panel.example.com/api/user/auth/refresh", nil)
|
||||
request.Host = "panel.example.com"
|
||||
request.Header.Set("Origin", origin)
|
||||
if origin == "" {
|
||||
request.Header.Del("Origin")
|
||||
}
|
||||
if referer != "" {
|
||||
request.Header.Set("Referer", referer)
|
||||
}
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
return response
|
||||
}
|
||||
|
||||
func TestSessionCookieOriginGuard(t *testing.T) {
|
||||
previousSecure := common.SessionCookieSecure
|
||||
previousTrustedURLs := common.SessionCookieTrustedURLs
|
||||
common.SessionCookieSecure = true
|
||||
common.SessionCookieTrustedURLs = []string{"https://trusted.example.com"}
|
||||
t.Cleanup(func() {
|
||||
common.SessionCookieSecure = previousSecure
|
||||
common.SessionCookieTrustedURLs = previousTrustedURLs
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
origin string
|
||||
referer string
|
||||
expected int
|
||||
}{
|
||||
{name: "same origin", origin: "https://panel.example.com", expected: http.StatusNoContent},
|
||||
{name: "trusted exact origin", origin: "https://trusted.example.com", expected: http.StatusNoContent},
|
||||
{name: "referer fallback", referer: "https://panel.example.com/profile", expected: http.StatusNoContent},
|
||||
{name: "missing both", expected: http.StatusForbidden},
|
||||
{name: "null origin", origin: "null", expected: http.StatusForbidden},
|
||||
{name: "suffix attack", origin: "https://trusted.example.com.evil.test", expected: http.StatusForbidden},
|
||||
{name: "scheme mismatch", origin: "http://panel.example.com", expected: http.StatusForbidden},
|
||||
{name: "path in origin", origin: "https://panel.example.com/profile", expected: http.StatusForbidden},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
response := runOriginGuardRequest(t, test.origin, test.referer)
|
||||
assert.Equal(t, test.expected, response.Code)
|
||||
assert.Empty(t, response.Header().Get("Access-Control-Allow-Origin"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCookieOriginGuardDevelopmentCompatibility(t *testing.T) {
|
||||
previousSecure := common.SessionCookieSecure
|
||||
previousTrustedURLs := common.SessionCookieTrustedURLs
|
||||
t.Cleanup(func() {
|
||||
common.SessionCookieSecure = previousSecure
|
||||
common.SessionCookieTrustedURLs = previousTrustedURLs
|
||||
})
|
||||
common.SessionCookieTrustedURLs = nil
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
secure bool
|
||||
origin string
|
||||
expected int
|
||||
}{
|
||||
{name: "insecure mode allows mismatched development origins", origin: "http://localhost:3001", expected: http.StatusNoContent},
|
||||
{name: "insecure mode allows missing origin", expected: http.StatusNoContent},
|
||||
{name: "secure mode rejects mismatched development origins", secure: true, origin: "http://localhost:3001", expected: http.StatusForbidden},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
common.SessionCookieSecure = test.secure
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.POST("/api/user/auth/refresh", SessionCookieOriginGuard(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "http://localhost:3000/api/user/auth/refresh", nil)
|
||||
request.Host = "localhost:3000"
|
||||
if test.origin != "" {
|
||||
request.Header.Set("Origin", test.origin)
|
||||
}
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
assert.Equal(t, test.expected, response.Code)
|
||||
assert.Empty(t, response.Header().Get("Access-Control-Allow-Origin"))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSessionCookieOriginGuardDoesNotTrustForwardedProtoFromClient(t *testing.T) {
|
||||
previousSecure := common.SessionCookieSecure
|
||||
previousTrustedURLs := common.SessionCookieTrustedURLs
|
||||
common.SessionCookieSecure = true
|
||||
common.SessionCookieTrustedURLs = nil
|
||||
t.Cleanup(func() {
|
||||
common.SessionCookieSecure = previousSecure
|
||||
common.SessionCookieTrustedURLs = previousTrustedURLs
|
||||
})
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.POST("/api/user/auth/refresh", SessionCookieOriginGuard(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "http://panel.example.com/api/user/auth/refresh", nil)
|
||||
request.Host = "panel.example.com"
|
||||
request.Header.Set("Origin", "https://panel.example.com")
|
||||
request.Header.Set("X-Forwarded-Proto", "https")
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
assert.Equal(t, http.StatusForbidden, response.Code)
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"fmt"
|
||||
"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/golang-jwt/jwt/v5"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func setupDashboardAuthMiddlewareTest(t *testing.T) {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousLogDB := model.LOG_DB
|
||||
previousType := common.MainDatabaseType()
|
||||
previousRedis := common.RedisEnabled
|
||||
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{}))
|
||||
model.DB = db
|
||||
model.LOG_DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.RedisEnabled = false
|
||||
common.SessionSecret = "middleware-auth-test-secret"
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLogDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
common.RedisEnabled = previousRedis
|
||||
common.SessionSecret = previousSecret
|
||||
})
|
||||
}
|
||||
|
||||
func issueExpiredDashboardAccessToken(t *testing.T, identity service.AuthIdentity) string {
|
||||
t.Helper()
|
||||
claims := jwt.MapClaims{
|
||||
"iss": "new-api",
|
||||
"aud": []string{"new-api-dashboard"},
|
||||
"sub": fmt.Sprintf("%d", identity.UserID),
|
||||
"token_use": "access",
|
||||
"sid": identity.SessionID,
|
||||
"uv": identity.UserAuthVersion,
|
||||
"sv": identity.SessionVersion,
|
||||
"exp": time.Now().Add(-time.Minute).Unix(),
|
||||
"nbf": time.Now().Add(-2 * time.Minute).Unix(),
|
||||
"iat": time.Now().Add(-2 * time.Minute).Unix(),
|
||||
}
|
||||
mac := hmac.New(sha256.New, []byte(common.SessionSecret))
|
||||
_, err := mac.Write([]byte("new-api/auth/access/v1"))
|
||||
require.NoError(t, err)
|
||||
token, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString(mac.Sum(nil))
|
||||
require.NoError(t, err)
|
||||
return token
|
||||
}
|
||||
|
||||
func tamperDashboardToken(token string) string {
|
||||
tamperAt := len(token) - 2
|
||||
replacement := "x"
|
||||
if token[tamperAt] == 'x' {
|
||||
replacement = "y"
|
||||
}
|
||||
return token[:tamperAt] + replacement + token[tamperAt+1:]
|
||||
}
|
||||
|
||||
func createMiddlewarePATUser(t *testing.T, username, token string) *model.User {
|
||||
t.Helper()
|
||||
user := &model.User{
|
||||
Username: username, Password: "password-placeholder", Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled, Group: "default", AccessToken: &token, AuthVersion: 1,
|
||||
AffCode: "middleware-aff-" + username,
|
||||
}
|
||||
require.NoError(t, model.DB.Create(user).Error)
|
||||
return user
|
||||
}
|
||||
|
||||
func TestUserAuthAllowsOpaqueDottedPAT(t *testing.T) {
|
||||
setupDashboardAuthMiddlewareTest(t)
|
||||
user := createMiddlewarePATUser(t, "dotted-pat-user", "opaque.key.with-dots")
|
||||
router := gin.New()
|
||||
router.GET("/protected", UserAuth(), func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"id": c.GetInt("id")})
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
request.Header.Set("Authorization", "Bearer opaque.key.with-dots")
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
assert.Equal(t, http.StatusOK, response.Code)
|
||||
var body struct {
|
||||
ID int `json:"id"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &body))
|
||||
assert.Equal(t, user.Id, body.ID)
|
||||
}
|
||||
|
||||
func TestUserAuthNeverFallsBackForRecognizedInvalidInternalJWT(t *testing.T) {
|
||||
setupDashboardAuthMiddlewareTest(t)
|
||||
identity := service.AuthIdentity{UserID: 42, SessionID: "session-42", UserAuthVersion: 1, SessionVersion: 1}
|
||||
token, _, err := service.IssueAccessToken(identity)
|
||||
require.NoError(t, err)
|
||||
tampered := tamperDashboardToken(token)
|
||||
createMiddlewarePATUser(t, "jwt-fallback-user", tampered)
|
||||
router := gin.New()
|
||||
router.GET("/protected", UserAuth(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodGet, "/protected", nil)
|
||||
request.Header.Set("Authorization", "Bearer "+tampered)
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
assert.Equal(t, http.StatusUnauthorized, response.Code)
|
||||
assert.Contains(t, response.Body.String(), "AUTH_UNAUTHORIZED")
|
||||
}
|
||||
|
||||
func TestTryUserAuthCredentialClassification(t *testing.T) {
|
||||
setupDashboardAuthMiddlewareTest(t)
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
patUser := createMiddlewarePATUser(t, "optional-pat-user", "optional.pat.with-dots")
|
||||
internalUser := createMiddlewarePATUser(t, "optional-session-user", "unrelated-pat")
|
||||
now := time.Now().Unix()
|
||||
session := &model.UserSession{
|
||||
SID: "optional-auth-session",
|
||||
UserID: internalUser.Id,
|
||||
Version: 1,
|
||||
UserAuthVersion: internalUser.AuthVersion,
|
||||
Status: model.UserSessionStatusActive,
|
||||
RefreshHash: "refresh-hash",
|
||||
LoginMethod: "password",
|
||||
LastActiveAt: now,
|
||||
ExpiresAt: now + 3600,
|
||||
}
|
||||
require.NoError(t, model.CreateUserSession(session))
|
||||
identity := service.AuthIdentity{
|
||||
UserID: internalUser.Id,
|
||||
SessionID: session.SID,
|
||||
UserAuthVersion: session.UserAuthVersion,
|
||||
SessionVersion: session.Version,
|
||||
}
|
||||
accessToken, _, err := service.IssueAccessToken(identity)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.AutoMigrate(&model.AuthFlow{}))
|
||||
binding, err := service.BindVerificationOperation(service.VerificationOperation{Scope: "channel.key.read", Context: []byte(`{"channel_id":123}`)})
|
||||
require.NoError(t, err)
|
||||
securityProof, _, err := service.IssueSecurityProof(identity, "2fa", binding)
|
||||
require.NoError(t, err)
|
||||
externalToken, err := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
|
||||
"iss": "external-issuer",
|
||||
"aud": "external-audience",
|
||||
"exp": time.Now().Add(time.Minute).Unix(),
|
||||
}).SignedString([]byte("external-secret"))
|
||||
require.NoError(t, err)
|
||||
|
||||
router := gin.New()
|
||||
router.GET("/optional", TryUserAuth(), func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"id": c.GetInt("id"),
|
||||
"use_access_token": c.GetBool("use_access_token"),
|
||||
})
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
token string
|
||||
wantStatus int
|
||||
wantUserID int
|
||||
wantPAT bool
|
||||
wantErrorCode string
|
||||
}{
|
||||
{name: "no authorization header", wantStatus: http.StatusOK},
|
||||
{name: "opaque unmatched credential", token: "opaque-relay-key", wantStatus: http.StatusOK},
|
||||
{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: "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"},
|
||||
{name: "security proof used as access", token: securityProof, wantStatus: http.StatusUnauthorized, wantErrorCode: "AUTH_UNAUTHORIZED"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodGet, "/optional", nil)
|
||||
if test.token != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+test.token)
|
||||
}
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
assert.Equal(t, test.wantStatus, response.Code)
|
||||
if test.wantErrorCode != "" {
|
||||
assert.Contains(t, response.Body.String(), test.wantErrorCode)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
ID int `json:"id"`
|
||||
UseAccessToken bool `json:"use_access_token"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &body))
|
||||
assert.Equal(t, test.wantUserID, body.ID)
|
||||
assert.Equal(t, test.wantPAT, body.UseAccessToken)
|
||||
})
|
||||
}
|
||||
|
||||
requiredRouter := gin.New()
|
||||
requiredRouter.GET("/required", UserAuth(), func(c *gin.Context) { c.Status(http.StatusNoContent) })
|
||||
requiredRequest := httptest.NewRequest(http.MethodGet, "/required", nil)
|
||||
requiredRequest.Header.Set("Authorization", "Bearer ordinary-unmatched-key")
|
||||
requiredResponse := httptest.NewRecorder()
|
||||
requiredRouter.ServeHTTP(requiredResponse, requiredRequest)
|
||||
assert.Equal(t, http.StatusUnauthorized, requiredResponse.Code, "required dashboard authentication must not adopt optional-auth fallback semantics")
|
||||
|
||||
var patUserQueries int
|
||||
forcedCacheError := errors.New("forced PAT user cache lookup failure")
|
||||
const callbackName = "test:optional-auth-pat-user-cache-failure"
|
||||
require.NoError(t, model.DB.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) {
|
||||
if tx.Statement.Table != "users" {
|
||||
return
|
||||
}
|
||||
patUserQueries++
|
||||
if patUserQueries == 2 {
|
||||
tx.AddError(forcedCacheError)
|
||||
}
|
||||
}))
|
||||
cacheFailureRequest := httptest.NewRequest(http.MethodGet, "/optional", nil)
|
||||
cacheFailureRequest.Header.Set("Authorization", "Bearer optional.pat.with-dots")
|
||||
cacheFailureResponse := httptest.NewRecorder()
|
||||
router.ServeHTTP(cacheFailureResponse, cacheFailureRequest)
|
||||
model.DB.Callback().Query().Remove(callbackName)
|
||||
assert.Equal(t, http.StatusInternalServerError, cacheFailureResponse.Code)
|
||||
assert.Contains(t, cacheFailureResponse.Body.String(), "AUTH_INTERNAL_ERROR")
|
||||
|
||||
sqlDB, err := model.DB.DB()
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, sqlDB.Close())
|
||||
databaseFailureRequest := httptest.NewRequest(http.MethodGet, "/optional", nil)
|
||||
databaseFailureRequest.Header.Set("Authorization", "Bearer database-failure-key")
|
||||
databaseFailureResponse := httptest.NewRecorder()
|
||||
router.ServeHTTP(databaseFailureResponse, databaseFailureRequest)
|
||||
assert.Equal(t, http.StatusInternalServerError, databaseFailureResponse.Code)
|
||||
assert.Contains(t, databaseFailureResponse.Body.String(), "AUTH_INTERNAL_ERROR")
|
||||
}
|
||||
@@ -0,0 +1,22 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// BodyStorageCleanup 请求体存储清理中间件
|
||||
// 在请求处理完成后自动清理磁盘/内存缓存
|
||||
func BodyStorageCleanup() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
defer func() {
|
||||
// 请求结束后清理存储
|
||||
common.CleanupBodyStorage(c)
|
||||
|
||||
// 清理文件缓存(URL 下载的文件等)
|
||||
service.CleanupFileSources(c)
|
||||
}()
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func Cache() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
if c.Request.RequestURI == "/" {
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
} else {
|
||||
c.Header("Cache-Control", "max-age=604800") // one week
|
||||
}
|
||||
c.Header("Cache-Version", "b688f2fb5be447c25e5aa3bd063087a83db32a288bf6a4f35f2d8db310e40b14")
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-contrib/cors"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func CORS() gin.HandlerFunc {
|
||||
config := cors.DefaultConfig()
|
||||
config.AllowAllOrigins = true
|
||||
config.AllowCredentials = true
|
||||
config.AllowMethods = []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"}
|
||||
config.AllowHeaders = []string{"*"}
|
||||
return cors.New(config)
|
||||
}
|
||||
|
||||
func Version() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("X-New-Api-Version", common.Version)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
package middleware
|
||||
|
||||
import "github.com/gin-gonic/gin"
|
||||
|
||||
func DisableCache() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("Cache-Control", "no-store, no-cache, must-revalidate, private, max-age=0")
|
||||
c.Header("Pragma", "no-cache")
|
||||
c.Header("Expires", "0")
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,735 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
taskdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type ModelRequest struct {
|
||||
Model string `json:"model"`
|
||||
Group string `json:"group,omitempty"`
|
||||
}
|
||||
|
||||
func Distribute() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
var channel *model.Channel
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddFilter(taskdto.ChannelFilter{
|
||||
Kind: taskdto.FilterRequestPath,
|
||||
RequestPath: c.Request.URL.Path,
|
||||
})
|
||||
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()}))
|
||||
return
|
||||
}
|
||||
if pin, found, overridden := constraints.ResolvedPin(); found {
|
||||
for _, lost := range overridden {
|
||||
logger.LogWarn(c, fmt.Sprintf(
|
||||
"channel pin overridden: winning_source=%s winning_channel_id=%d overridden_source=%s overridden_channel_id=%d",
|
||||
pin.Source, pin.ChannelId, lost.Source, lost.ChannelId,
|
||||
))
|
||||
}
|
||||
channel, err = model.CacheGetChannel(pin.ChannelId)
|
||||
if err != nil {
|
||||
if pin.Source == taskdto.PinSourceOriginTask {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, "origin_task_channel_disabled", types.ErrorCode("origin_task_channel_disabled"))
|
||||
} else {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
|
||||
}
|
||||
return
|
||||
}
|
||||
if channel.Status != common.ChannelStatusEnabled {
|
||||
if pin.Source == taskdto.PinSourceOriginTask {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, "origin_task_channel_disabled", types.ErrorCode("origin_task_channel_disabled"))
|
||||
} else {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
|
||||
}
|
||||
return
|
||||
}
|
||||
if ok, kind := model.ChannelSatisfiesFilters(channel, modelRequest.Model, constraints.Filters); !ok {
|
||||
if kind == taskdto.FilterTaskPluginIdentity {
|
||||
logTaskPluginChannelDecision(c, channel, modelRequest.Model, "channel_rejected", "identity_mismatch")
|
||||
}
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelRequest.Model}), types.ErrorCode(kind))
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// Select a channel for the user
|
||||
// check token model mapping
|
||||
modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled)
|
||||
if modelLimitEnable {
|
||||
s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit)
|
||||
if !ok {
|
||||
// token model limit is empty, all models are not allowed
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess))
|
||||
return
|
||||
}
|
||||
var tokenModelLimit map[string]bool
|
||||
tokenModelLimit, ok = s.(map[string]bool)
|
||||
if !ok {
|
||||
tokenModelLimit = map[string]bool{}
|
||||
}
|
||||
if !tokenModelLimitAllows(tokenModelLimit, modelRequest.Model) {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model}))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if shouldSelectChannel {
|
||||
if modelRequest.Model == "" {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired))
|
||||
return
|
||||
}
|
||||
var selectGroup string
|
||||
usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
|
||||
// check path is /pg/chat/completions
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
|
||||
playgroundRequest := &dto.PlayGroundRequest{}
|
||||
err = common.UnmarshalBodyReusable(c, playgroundRequest)
|
||||
if err != nil {
|
||||
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()}))
|
||||
return
|
||||
}
|
||||
if playgroundRequest.Group != "" {
|
||||
if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup {
|
||||
abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied))
|
||||
return
|
||||
}
|
||||
usingGroup = playgroundRequest.Group
|
||||
common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup)
|
||||
}
|
||||
}
|
||||
|
||||
if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
|
||||
affinityUsable := false
|
||||
preferred, err := model.CacheGetChannel(preferredChannelID)
|
||||
affinitySatisfied := false
|
||||
if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
|
||||
affinitySatisfied, _ = model.ChannelSatisfiesFilters(preferred, modelRequest.Model, constraints.Filters)
|
||||
}
|
||||
if affinitySatisfied {
|
||||
if usingGroup == "auto" {
|
||||
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
|
||||
autoGroups := service.GetRequestAutoGroups(c, userGroup)
|
||||
for _, g := range autoGroups {
|
||||
if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) {
|
||||
selectGroup = g
|
||||
common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
|
||||
channel = preferred
|
||||
affinityUsable = true
|
||||
service.MarkChannelAffinityUsed(c, g, preferred.Id)
|
||||
break
|
||||
}
|
||||
}
|
||||
} else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) {
|
||||
channel = preferred
|
||||
selectGroup = usingGroup
|
||||
affinityUsable = true
|
||||
service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
|
||||
}
|
||||
}
|
||||
if !affinityUsable && !service.ShouldKeepChannelAffinityOnChannelDisabled() {
|
||||
service.ClearCurrentChannelAffinityCache(c)
|
||||
}
|
||||
}
|
||||
|
||||
if channel == nil {
|
||||
channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{
|
||||
Ctx: c,
|
||||
ModelName: modelRequest.Model,
|
||||
TokenGroup: usingGroup,
|
||||
RequestPath: c.Request.URL.Path,
|
||||
Retry: common.GetPointer(0),
|
||||
})
|
||||
if err != nil {
|
||||
showGroup := usingGroup
|
||||
if usingGroup == "auto" {
|
||||
showGroup = fmt.Sprintf("auto(%s)", selectGroup)
|
||||
}
|
||||
message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()})
|
||||
// 如果错误,但是渠道不为空,说明是数据库一致性问题
|
||||
//if channel != nil {
|
||||
// common.SysError(fmt.Sprintf("渠道不存在:%d", channel.Id))
|
||||
// message = "数据库一致性已被破坏,请联系管理员"
|
||||
//}
|
||||
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound)
|
||||
return
|
||||
}
|
||||
if channel == nil {
|
||||
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, noAvailableChannelMessage(c, usingGroup, modelRequest.Model), types.ErrorCodeModelNotFound)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if channel != nil {
|
||||
if ok, kind := model.ChannelSatisfiesFilters(channel, modelRequest.Model, constraints.Filters); !ok {
|
||||
if kind == taskdto.FilterTaskPluginIdentity {
|
||||
logTaskPluginChannelDecision(c, channel, modelRequest.Model, "channel_rejected", "identity_mismatch")
|
||||
}
|
||||
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, noAvailableChannelMessage(c, common.GetContextKeyString(c, constant.ContextKeyUsingGroup), modelRequest.Model), types.ErrorCodeModelNotFound)
|
||||
return
|
||||
}
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
|
||||
SetupContextForSelectedChannel(c, channel, modelRequest.Model)
|
||||
c.Next()
|
||||
if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest {
|
||||
service.RecordChannelAffinity(c, channel.Id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// noAvailableChannelMessage explains a 503 for a task-plugin-claimed model.
|
||||
// A model claimed by a plugin is served only by that plugin's channels, so the
|
||||
// generic "no channel" text hides the real cause: the claiming plugin has no
|
||||
// enabled channel, and the operator must disable or override that plugin for
|
||||
// any other plugin or channel to take the model. Non-plugin requests keep the
|
||||
// generic message.
|
||||
func noAvailableChannelMessage(c *gin.Context, group, modelName string) string {
|
||||
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
|
||||
pinned, ok := value.(jsplugin.PinnedPlugin)
|
||||
if exists && ok && pinned.Plugin != nil {
|
||||
return i18n.T(c, i18n.MsgDistributorNoAvailableChannelTaskPlugin, map[string]any{"Group": group, "Model": modelName, "Plugin": pinned.Plugin.Meta.Key})
|
||||
}
|
||||
return i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": group, "Model": modelName})
|
||||
}
|
||||
|
||||
func channelMatchesExpectedTaskPlugin(c *gin.Context, channel *model.Channel, expected string) bool {
|
||||
if channel == nil {
|
||||
return false
|
||||
}
|
||||
if c != nil {
|
||||
if _, matched := pinnedEndpointCandidateForChannel(c, channel, expected); matched {
|
||||
return true
|
||||
}
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeTaskPlugin {
|
||||
return expected != "" && channel.GetSetting().TaskPluginKey == expected
|
||||
}
|
||||
if expected == "" {
|
||||
return true
|
||||
}
|
||||
|
||||
if c == nil {
|
||||
return false
|
||||
}
|
||||
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
|
||||
pinned, ok := value.(jsplugin.PinnedPlugin)
|
||||
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil || pinned.Plugin.Meta.Key != expected {
|
||||
return false
|
||||
}
|
||||
plugin, ok := pinned.Generation.GetByChannelType(channel.Type)
|
||||
return ok && plugin == pinned.Plugin
|
||||
}
|
||||
|
||||
func pinnedEndpointCandidateForChannel(c *gin.Context, channel *model.Channel, expected string) (jsplugin.ProtocolBinding, bool) {
|
||||
if c == nil || channel == nil || expected == "" {
|
||||
return jsplugin.ProtocolBinding{}, false
|
||||
}
|
||||
value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
|
||||
pinned, ok := value.(jsplugin.PinnedEndpoint)
|
||||
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil {
|
||||
return jsplugin.ProtocolBinding{}, false
|
||||
}
|
||||
candidates := pinned.Candidates
|
||||
if len(candidates) == 0 {
|
||||
candidates = []jsplugin.ProtocolBinding{{Plugin: pinned.Plugin, Protocol: pinned.Protocol, Operation: pinned.Operation, Model: pinned.Model}}
|
||||
}
|
||||
expectedOwned := false
|
||||
selected := jsplugin.ProtocolBinding{}
|
||||
for _, candidate := range candidates {
|
||||
if candidate.Plugin == nil {
|
||||
continue
|
||||
}
|
||||
if candidate.Plugin.Meta.Key == expected {
|
||||
expectedOwned = true
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeTaskPlugin {
|
||||
if channel.GetSetting().TaskPluginKey == candidate.Plugin.Meta.Key {
|
||||
selected = candidate
|
||||
}
|
||||
continue
|
||||
}
|
||||
plugin, indexed := pinned.Generation.GetByChannelType(channel.Type)
|
||||
if indexed && plugin == candidate.Plugin {
|
||||
selected = candidate
|
||||
}
|
||||
}
|
||||
return selected, expectedOwned && selected.Plugin != nil
|
||||
}
|
||||
|
||||
// getModelFromRequest 从请求中读取模型信息
|
||||
// 根据 Content-Type 自动处理:
|
||||
// - application/json
|
||||
// - application/x-www-form-urlencoded
|
||||
// - multipart/form-data
|
||||
func getModelFromRequest(c *gin.Context) (*ModelRequest, error) {
|
||||
if cached, exists := c.Get(contextKeyTaskPluginEndpointModel); exists {
|
||||
if modelRequest, ok := cached.(ModelRequest); ok {
|
||||
cachedRequest := modelRequest
|
||||
return &cachedRequest, nil
|
||||
}
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
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 &modelRequest, nil
|
||||
}
|
||||
|
||||
func getModelFromJSONBody(c *gin.Context) (*ModelRequest, error) {
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
requestBody, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !gjson.ValidBytes(requestBody) {
|
||||
return nil, errors.New("invalid JSON request body")
|
||||
}
|
||||
if countTopLevelJSONKey(requestBody, "model") > 1 {
|
||||
return nil, errors.New("model must be provided once")
|
||||
}
|
||||
|
||||
values := gjson.GetManyBytes(requestBody, "model", "group")
|
||||
model, err := getJSONStringValue(values[0], "model")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
group, err := getJSONStringValue(values[1], "group")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if _, seekErr := storage.Seek(0, io.SeekStart); seekErr != nil {
|
||||
return nil, seekErr
|
||||
}
|
||||
c.Request.Body = io.NopCloser(storage)
|
||||
|
||||
return &ModelRequest{
|
||||
Model: model,
|
||||
Group: group,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func countTopLevelJSONKey(data []byte, target string) int {
|
||||
depth := 0
|
||||
inString := false
|
||||
escaped := false
|
||||
stringStart := 0
|
||||
expectingKey := false
|
||||
count := 0
|
||||
for index, current := range data {
|
||||
if inString {
|
||||
if escaped {
|
||||
escaped = false
|
||||
continue
|
||||
}
|
||||
if current == '\\' {
|
||||
escaped = true
|
||||
continue
|
||||
}
|
||||
if current != '"' {
|
||||
continue
|
||||
}
|
||||
inString = false
|
||||
if depth == 1 && expectingKey {
|
||||
key := string(data[stringStart:index])
|
||||
var decodedKey string
|
||||
if common.Unmarshal(data[stringStart-1:index+1], &decodedKey) == nil {
|
||||
key = decodedKey
|
||||
}
|
||||
cursor := index + 1
|
||||
for cursor < len(data) && (data[cursor] == ' ' || data[cursor] == '\t' || data[cursor] == '\r' || data[cursor] == '\n') {
|
||||
cursor++
|
||||
}
|
||||
if cursor < len(data) && data[cursor] == ':' && key == target {
|
||||
count++
|
||||
}
|
||||
expectingKey = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch current {
|
||||
case '"':
|
||||
inString = true
|
||||
stringStart = index + 1
|
||||
case '{':
|
||||
depth++
|
||||
if depth == 1 {
|
||||
expectingKey = true
|
||||
}
|
||||
case '}':
|
||||
depth--
|
||||
case ',':
|
||||
if depth == 1 {
|
||||
expectingKey = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func getJSONStringValue(result gjson.Result, field string) (string, error) {
|
||||
if !result.Exists() || result.Type == gjson.Null {
|
||||
return "", nil
|
||||
}
|
||||
if result.Type != gjson.String {
|
||||
return "", fmt.Errorf("field %s must be a string", field)
|
||||
}
|
||||
return result.String(), nil
|
||||
}
|
||||
|
||||
func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
|
||||
var modelRequest ModelRequest
|
||||
shouldSelectChannel := true
|
||||
var err error
|
||||
if modelName := c.GetString("resolved_task_model"); modelName != "" {
|
||||
modelRequest.Model = modelName
|
||||
} else if strings.Contains(c.Request.URL.Path, "/mj/") {
|
||||
relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path)
|
||||
if relayMode == relayconstant.RelayModeMidjourneyTaskFetch ||
|
||||
relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition ||
|
||||
relayMode == relayconstant.RelayModeMidjourneyNotify ||
|
||||
relayMode == relayconstant.RelayModeMidjourneyTaskImageSeed {
|
||||
shouldSelectChannel = false
|
||||
} else {
|
||||
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()}))
|
||||
}
|
||||
midjourneyModel, mjErr, success := service.GetMjRequestModel(relayMode, &midjourneyRequest)
|
||||
if mjErr != nil {
|
||||
return nil, false, fmt.Errorf("%s", mjErr.Description)
|
||||
}
|
||||
if midjourneyModel == "" {
|
||||
if !success {
|
||||
return nil, false, fmt.Errorf("%s", i18n.T(c, i18n.MsgDistributorInvalidParseModel))
|
||||
} else {
|
||||
// task fetch, task fetch by condition, notify
|
||||
shouldSelectChannel = false
|
||||
}
|
||||
}
|
||||
modelRequest.Model = midjourneyModel
|
||||
}
|
||||
c.Set("relay_mode", relayMode)
|
||||
} else if strings.Contains(c.Request.URL.Path, "/v1/videos/") && strings.HasSuffix(c.Request.URL.Path, "/remix") {
|
||||
relayMode := relayconstant.RelayModeVideoSubmit
|
||||
c.Set("relay_mode", relayMode)
|
||||
shouldSelectChannel = false
|
||||
} else if strings.Contains(c.Request.URL.Path, "/v1/videos") {
|
||||
//curl https://api.openai.com/v1/videos \
|
||||
// -H "Authorization: Bearer $OPENAI_API_KEY" \
|
||||
// -F "model=sora-2" \
|
||||
// -F "prompt=A calico cat playing a piano on stage"
|
||||
// -F input_reference="@image.jpg"
|
||||
relayMode := relayconstant.RelayModeUnknown
|
||||
if c.Request.Method == http.MethodPost {
|
||||
relayMode = relayconstant.RelayModeVideoSubmit
|
||||
req, err := getModelFromRequest(c)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
if req != nil {
|
||||
modelRequest.Model = req.Model
|
||||
}
|
||||
} else if c.Request.Method == http.MethodGet {
|
||||
relayMode = relayconstant.RelayModeVideoFetchByID
|
||||
shouldSelectChannel = false
|
||||
modelRequest.Model = getTaskOriginModelName(c)
|
||||
}
|
||||
c.Set("relay_mode", relayMode)
|
||||
} else if strings.Contains(c.Request.URL.Path, "/v1/video/generations") {
|
||||
relayMode := relayconstant.RelayModeUnknown
|
||||
if c.Request.Method == http.MethodPost {
|
||||
req, err := getModelFromRequest(c)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
modelRequest.Model = req.Model
|
||||
relayMode = relayconstant.RelayModeVideoSubmit
|
||||
} else if c.Request.Method == http.MethodGet {
|
||||
relayMode = relayconstant.RelayModeVideoFetchByID
|
||||
shouldSelectChannel = false
|
||||
modelRequest.Model = getTaskOriginModelName(c)
|
||||
}
|
||||
if _, ok := c.Get("relay_mode"); !ok {
|
||||
c.Set("relay_mode", relayMode)
|
||||
}
|
||||
} else if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models/") || strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
|
||||
// Gemini API 路径处理: /v1beta/models/gemini-2.0-flash:generateContent
|
||||
relayMode := relayconstant.RelayModeGemini
|
||||
modelName := extractModelNameFromGeminiPath(c.Request.URL.Path)
|
||||
if modelName != "" {
|
||||
modelRequest.Model = modelName
|
||||
}
|
||||
c.Set("relay_mode", relayMode)
|
||||
} else if !strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") && !strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
|
||||
req, err := getModelFromRequest(c)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
modelRequest.Model = req.Model
|
||||
}
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") {
|
||||
//wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01
|
||||
modelRequest.Model = c.Query("model")
|
||||
}
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
|
||||
if modelRequest.Model == "" {
|
||||
modelRequest.Model = "text-moderation-stable"
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(c.Request.URL.Path, "embeddings") {
|
||||
if modelRequest.Model == "" {
|
||||
modelRequest.Model = c.Param("model")
|
||||
}
|
||||
}
|
||||
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") {
|
||||
//modelRequest.Model = common.GetStringIfEmpty(c.PostForm("model"), "gpt-image-1")
|
||||
contentType := c.ContentType()
|
||||
if slices.Contains([]string{gin.MIMEPOSTForm, gin.MIMEMultipartPOSTForm}, contentType) {
|
||||
req, err := getModelFromRequest(c)
|
||||
if err == nil && req.Model != "" {
|
||||
modelRequest.Model = req.Model
|
||||
}
|
||||
}
|
||||
}
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
|
||||
relayMode := relayconstant.RelayModeAudioSpeech
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
|
||||
|
||||
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "tts-1")
|
||||
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
|
||||
// 先尝试从请求读取
|
||||
if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
|
||||
modelRequest.Model = req.Model
|
||||
}
|
||||
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
|
||||
relayMode = relayconstant.RelayModeAudioTranslation
|
||||
} else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") {
|
||||
// 先尝试从请求读取
|
||||
if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
|
||||
modelRequest.Model = req.Model
|
||||
}
|
||||
modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
|
||||
relayMode = relayconstant.RelayModeAudioTranscription
|
||||
}
|
||||
c.Set("relay_mode", relayMode)
|
||||
}
|
||||
if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
|
||||
// playground chat completions
|
||||
req, err := getModelFromRequest(c)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
modelRequest.Model = req.Model
|
||||
modelRequest.Group = req.Group
|
||||
common.SetContextKey(c, constant.ContextKeyTokenGroup, modelRequest.Group)
|
||||
}
|
||||
|
||||
return &modelRequest, shouldSelectChannel, nil
|
||||
}
|
||||
|
||||
// tokenModelLimitAllows reports whether a token model-limit map authorizes
|
||||
// model. Exact name, wildcard-normalized name, and routing-normalized name
|
||||
// (modifiers and legacy aliases stripped) are all accepted.
|
||||
func tokenModelLimitAllows(limit map[string]bool, model string) bool {
|
||||
if limit[model] {
|
||||
return true
|
||||
}
|
||||
if formatted := ratio_setting.FormatMatchingModelName(model); limit[formatted] {
|
||||
return true
|
||||
}
|
||||
return limit[ratio_setting.RoutingMatchModelName(model)]
|
||||
}
|
||||
|
||||
// 修复 #4834: GET /v1/video/generations/:task_id && /v1/video/:task_id 此前不解析 model,
|
||||
// 当 token 启用「可用模型限制」时,下游 modelLimitEnable 校验会因
|
||||
// modelRequest.Model 为空而误报 "This token has no access to model"。
|
||||
// 从已存储的任务记录中回填 OriginModelName 即可让校验走在正确的模型上。
|
||||
func getTaskOriginModelName(c *gin.Context) string {
|
||||
if !common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) {
|
||||
return ""
|
||||
}
|
||||
|
||||
taskId := c.Param("task_id")
|
||||
if taskId == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
userId := c.GetInt("id")
|
||||
if task, exist, err := model.GetByTaskId(userId, taskId); err == nil && exist && task != nil {
|
||||
return task.Properties.OriginModelName
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
|
||||
c.Set("original_model", modelName) // for retry
|
||||
expectedPlugin := c.GetString("expected_task_plugin_key")
|
||||
if channel == nil {
|
||||
logTaskPluginChannelDecision(c, nil, modelName, "channel_rejected", "nil_channel")
|
||||
return types.NewError(errors.New("channel is nil"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
if expectedPlugin != "" && !channelMatchesExpectedTaskPlugin(c, channel, expectedPlugin) {
|
||||
logTaskPluginChannelDecision(c, channel, modelName, "channel_rejected", "identity_mismatch")
|
||||
return types.NewError(
|
||||
errors.New("selected channel does not match the pinned task plugin"),
|
||||
types.ErrorCodeGetChannelFailed,
|
||||
types.ErrOptionWithSkipRetry(),
|
||||
)
|
||||
}
|
||||
if candidate, matched := pinnedEndpointCandidateForChannel(c, channel, expectedPlugin); matched {
|
||||
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
|
||||
if pinned, ok := value.(jsplugin.PinnedEndpoint); ok && candidate.Plugin != nil && candidate.Plugin != pinned.Plugin {
|
||||
previousPlugin := pinned.Plugin.Meta.Key
|
||||
pinned.Plugin = candidate.Plugin
|
||||
pinned.Protocol = candidate.Protocol
|
||||
pinned.Operation = candidate.Operation
|
||||
c.Set(jsplugin.ContextKeyPinnedEndpoint, pinned)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: pinned.Generation, Plugin: candidate.Plugin})
|
||||
c.Set("expected_task_plugin_key", candidate.Plugin.Meta.Key)
|
||||
c.Set("task_plugin_key", candidate.Plugin.Meta.Key)
|
||||
c.Set("platform", candidate.Plugin.Meta.Key)
|
||||
logger.LogDebug(
|
||||
c,
|
||||
"task_plugin subsystem=endpoint event=provider_selected generation=%d previous_plugin=%q plugin=%q model=%q channel_id=%d channel_type=%d",
|
||||
pinned.Generation.Number,
|
||||
previousPlugin,
|
||||
candidate.Plugin.Meta.Key,
|
||||
modelName,
|
||||
channel.Id,
|
||||
channel.Type,
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
|
||||
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 {
|
||||
c.Set("task_plugin_key", channel.GetSetting().TaskPluginKey)
|
||||
}
|
||||
logTaskPluginChannelDecision(c, channel, modelName, "channel_selected", "")
|
||||
paramOverride := channel.GetParamOverride()
|
||||
headerOverride := channel.GetHeaderOverride()
|
||||
if mergedParam, applied := service.ApplyChannelAffinityOverrideTemplate(c, paramOverride); applied {
|
||||
paramOverride = mergedParam
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, paramOverride)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, headerOverride)
|
||||
if nil != channel.OpenAIOrganization && *channel.OpenAIOrganization != "" {
|
||||
common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization)
|
||||
}
|
||||
common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan())
|
||||
common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
|
||||
common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
|
||||
|
||||
key, index, newAPIError := channel.GetNextEnabledKey()
|
||||
if newAPIError != nil {
|
||||
return newAPIError
|
||||
}
|
||||
if channel.ChannelInfo.IsMultiKey {
|
||||
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index)
|
||||
} else {
|
||||
// 必须设置为 false,否则在重试到单个 key 的时候会导致日志显示错误
|
||||
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, false)
|
||||
}
|
||||
// c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
|
||||
common.SetContextKey(c, constant.ContextKeyChannelKey, key)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL())
|
||||
|
||||
common.SetContextKey(c, constant.ContextKeySystemPromptOverride, false)
|
||||
|
||||
// TODO: api_version统一
|
||||
switch channel.Type {
|
||||
case constant.ChannelTypeAzure:
|
||||
c.Set("api_version", channel.Other)
|
||||
case constant.ChannelTypeVertexAi:
|
||||
c.Set("region", channel.Other)
|
||||
case constant.ChannelTypeXunfei:
|
||||
c.Set("api_version", channel.Other)
|
||||
case constant.ChannelTypeGemini:
|
||||
c.Set("api_version", channel.Other)
|
||||
case constant.ChannelTypeAli:
|
||||
c.Set("plugin", channel.Other)
|
||||
case constant.ChannelCloudflare:
|
||||
c.Set("api_version", channel.Other)
|
||||
case constant.ChannelTypeMokaAI:
|
||||
c.Set("api_version", channel.Other)
|
||||
case constant.ChannelTypeCoze:
|
||||
c.Set("bot_id", channel.Other)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// extractModelNameFromGeminiPath 从 Gemini API URL 路径中提取模型名
|
||||
// 输入格式: /v1beta/models/gemini-2.0-flash:generateContent
|
||||
// 输出: gemini-2.0-flash
|
||||
func extractModelNameFromGeminiPath(path string) string {
|
||||
// 查找 "/models/" 的位置
|
||||
modelsPrefix := "/models/"
|
||||
modelsIndex := strings.Index(path, modelsPrefix)
|
||||
if modelsIndex == -1 {
|
||||
return ""
|
||||
}
|
||||
|
||||
// 从 "/models/" 之后开始提取
|
||||
startIndex := modelsIndex + len(modelsPrefix)
|
||||
if startIndex >= len(path) {
|
||||
return ""
|
||||
}
|
||||
|
||||
// 查找 ":" 的位置,模型名在 ":" 之前
|
||||
colonIndex := strings.Index(path[startIndex:], ":")
|
||||
if colonIndex == -1 {
|
||||
// 如果没有找到 ":",返回从 "/models/" 到路径结尾的部分
|
||||
return path[startIndex:]
|
||||
}
|
||||
|
||||
// 返回模型名部分
|
||||
return path[startIndex : startIndex+colonIndex]
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/setting/model_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestChannelMatchesExpectedTaskPluginUsesGenericChannelSetting(t *testing.T) {
|
||||
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin}
|
||||
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "generic-alpha"})
|
||||
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-alpha"))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-beta"))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, ""))
|
||||
}
|
||||
|
||||
func TestChannelMatchesExpectedTaskPluginUsesPinnedLegacyIndex(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
alpha, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
pinnedGeneration := registry.Generation()
|
||||
|
||||
require.NoError(t, registry.Unregister("legacy-alpha"))
|
||||
_, err = registry.Register(distributorTaskPluginSource("legacy-beta", constant.ChannelTypeKling), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: pinnedGeneration,
|
||||
Plugin: alpha,
|
||||
})
|
||||
channel := &model.Channel{Type: constant.ChannelTypeKling}
|
||||
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-alpha"))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-beta"))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
|
||||
}
|
||||
|
||||
func TestChannelMatchesExpectedTaskPluginRejectsUnindexedLegacyChannel(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
plugin, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: plugin,
|
||||
})
|
||||
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: 0}, "legacy-alpha"))
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, ""))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(nil, &model.Channel{Type: constant.ChannelTypeKling}, "legacy-alpha"))
|
||||
|
||||
c.Set("expected_task_plugin_key", "legacy-alpha")
|
||||
setupErr := SetupContextForSelectedChannel(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "task-model")
|
||||
require.NotNil(t, setupErr)
|
||||
assert.Contains(t, setupErr.Error(), "does not match")
|
||||
}
|
||||
|
||||
func TestSharedEndpointRebindsToSelectedLegacyProvider(t *testing.T) {
|
||||
registry := jsplugin.NewRegistry()
|
||||
_, err := registry.Register(distributorEndpointPluginSource("gemini-shared", constant.ChannelTypeGemini), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
_, err = registry.Register(distributorEndpointPluginSource("vertex-shared", constant.ChannelTypeVertexAi), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
candidates := registry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "task-model")
|
||||
require.Len(t, candidates, 2)
|
||||
|
||||
c, _ := gin.CreateTestContext(nil)
|
||||
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: registry.Generation(), Plugin: candidates[0].Plugin})
|
||||
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
|
||||
Generation: registry.Generation(),
|
||||
Plugin: candidates[0].Plugin,
|
||||
Protocol: candidates[0].Protocol,
|
||||
Operation: candidates[0].Operation,
|
||||
Model: "task-model",
|
||||
Candidates: candidates,
|
||||
})
|
||||
c.Set("expected_task_plugin_key", candidates[0].Plugin.Meta.Key)
|
||||
|
||||
geminiChannel := &model.Channel{Id: 1, Type: constant.ChannelTypeGemini}
|
||||
vertexChannel := &model.Channel{Id: 2, Type: constant.ChannelTypeVertexAi}
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, candidates[0].Plugin.Meta.Key))
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(c, vertexChannel, candidates[0].Plugin.Meta.Key))
|
||||
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeKling}, candidates[0].Plugin.Meta.Key))
|
||||
|
||||
require.Nil(t, SetupContextForSelectedChannel(c, vertexChannel, "task-model"))
|
||||
pinnedValue, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
|
||||
require.True(t, exists)
|
||||
pinned, ok := pinnedValue.(jsplugin.PinnedEndpoint)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, "vertex-shared", pinned.Plugin.Meta.Key)
|
||||
assert.Equal(t, "vertex-shared", c.GetString("expected_task_plugin_key"))
|
||||
assert.Equal(t, "vertex-shared", c.GetString("task_plugin_key"))
|
||||
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, "vertex-shared"), "a retry may select another declared provider")
|
||||
}
|
||||
|
||||
func distributorTaskPluginSource(key string, channelType int) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: %q,
|
||||
name: %q,
|
||||
version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
channelTypes: [%d],
|
||||
models: ["task-model"],
|
||||
fetchMode: "per_task",
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
`, key, key, channelType)
|
||||
}
|
||||
|
||||
func distributorEndpointPluginSource(key string, channelType int) string {
|
||||
return fmt.Sprintf(`
|
||||
export const meta = {
|
||||
apiVersion: 1,
|
||||
key: %q,
|
||||
name: %q,
|
||||
version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
channelTypes: [%d],
|
||||
models: ["task-model"],
|
||||
fetchMode: "per_task",
|
||||
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {taskId: "task"}; }
|
||||
export function buildQueryRequest() { return {}; }
|
||||
export function parseTaskResult() { return {status: "SUCCESS"}; }
|
||||
export const protocols = {openai_responses: {
|
||||
decodeRequest: function(ctx) { return {kind: "submit", model: "task-model", requestBody: ctx.body.value}; },
|
||||
renderEvents: function() { return {events: [], state: null, done: false}; },
|
||||
renderFinal: function() { return {output: []}; },
|
||||
}};
|
||||
`, key, key, channelType)
|
||||
}
|
||||
|
||||
func TestTokenModelLimitAllowsLegacyAliasAndModifierVariant(t *testing.T) {
|
||||
aliasOnly := map[string]bool{"claude-3-7-sonnet-thinking": true}
|
||||
assert.True(t, tokenModelLimitAllows(aliasOnly, "claude-3-7-sonnet-thinking"))
|
||||
assert.False(t, tokenModelLimitAllows(aliasOnly, "claude-3-7-sonnet"))
|
||||
|
||||
baseOnly := map[string]bool{"claude-3-7-sonnet": true}
|
||||
assert.True(t, tokenModelLimitAllows(baseOnly, "claude-3-7-sonnet@thinking:on"))
|
||||
assert.True(t, tokenModelLimitAllows(baseOnly, "claude-3-7-sonnet-thinking"))
|
||||
|
||||
wildcard := map[string]bool{"gemini-2.5-flash-thinking-*": true}
|
||||
assert.True(t, tokenModelLimitAllows(wildcard, "gemini-2.5-flash-thinking-8192"))
|
||||
}
|
||||
|
||||
func TestTokenModelLimitAllowsExemptAtNameByFullName(t *testing.T) {
|
||||
settings := model_setting.GetGlobalSettings()
|
||||
original := append([]string(nil), settings.ThinkingModelBlacklist...)
|
||||
t.Cleanup(func() { settings.ThinkingModelBlacklist = original })
|
||||
settings.ThinkingModelBlacklist = append(original, "re:.*@sha256:.*")
|
||||
|
||||
fullOnly := map[string]bool{"opaque@sha256:deadbeef": true}
|
||||
assert.True(t, tokenModelLimitAllows(fullOnly, "opaque@sha256:deadbeef"))
|
||||
|
||||
baseOnly := map[string]bool{"opaque": true}
|
||||
assert.False(t, tokenModelLimitAllows(baseOnly, "opaque@sha256:deadbeef"))
|
||||
}
|
||||
|
||||
func TestNoAvailableChannelMessageNamesClaimingTaskPlugin(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
registry := jsplugin.NewRegistry()
|
||||
plugin, err := registry.Register(distributorTaskPluginSource("claimer", constant.ChannelTypeKling), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
|
||||
pinned, _ := gin.CreateTestContext(nil)
|
||||
pinned.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
||||
pinned.Request.Header.Set("Accept-Language", "en")
|
||||
pinned.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
||||
message := noAvailableChannelMessage(pinned, "default", "kling-v1")
|
||||
assert.Contains(t, message, `"claimer"`)
|
||||
assert.Contains(t, message, "disable or override")
|
||||
assert.Contains(t, message, "kling-v1")
|
||||
|
||||
plain, _ := gin.CreateTestContext(nil)
|
||||
plain.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
plain.Request.Header.Set("Accept-Language", "en")
|
||||
generic := noAvailableChannelMessage(plain, "default", "gpt-4o")
|
||||
assert.NotContains(t, generic, "task plugin")
|
||||
assert.Contains(t, generic, "gpt-4o")
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
EmailVerificationRateLimitMark = "EV"
|
||||
EmailVerificationMaxRequests = 2 // 30秒内最多2次
|
||||
EmailVerificationDuration = 30 // 30秒时间窗口
|
||||
)
|
||||
|
||||
func redisEmailVerificationRateLimiter(c *gin.Context) {
|
||||
allowed, _, ttlSeconds, err := redisFixedWindowTake(
|
||||
c.Request.Context(),
|
||||
redisIPRateLimitKey(EmailVerificationRateLimitMark, c.ClientIP()),
|
||||
EmailVerificationMaxRequests,
|
||||
EmailVerificationDuration,
|
||||
)
|
||||
if err != nil {
|
||||
memoryEmailVerificationRateLimiter(c)
|
||||
return
|
||||
}
|
||||
if allowed {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
waitSeconds := int64(EmailVerificationDuration)
|
||||
if ttlSeconds > 0 {
|
||||
waitSeconds = ttlSeconds
|
||||
}
|
||||
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{
|
||||
"success": false,
|
||||
"message": fmt.Sprintf("发送过于频繁,请等待 %d 秒后再试", waitSeconds),
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
func memoryEmailVerificationRateLimiter(c *gin.Context) {
|
||||
key := EmailVerificationRateLimitMark + ":" + c.ClientIP()
|
||||
|
||||
if !inMemoryRateLimiter.Request(key, EmailVerificationMaxRequests, EmailVerificationDuration) {
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{
|
||||
"success": false,
|
||||
"message": "发送过于频繁,请稍后再试",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
|
||||
func EmailVerificationRateLimit() gin.HandlerFunc {
|
||||
// Keep the fallback ready before requests arrive so a concurrent Redis
|
||||
// outage cannot race the in-memory limiter's first initialization.
|
||||
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
|
||||
return func(c *gin.Context) {
|
||||
if common.RedisEnabled {
|
||||
redisEmailVerificationRateLimiter(c)
|
||||
} else {
|
||||
memoryEmailVerificationRateLimiter(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"compress/gzip"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/andybalholm/brotli"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
type readCloser struct {
|
||||
io.Reader
|
||||
closeFn func() error
|
||||
}
|
||||
|
||||
func (rc *readCloser) Close() error {
|
||||
if rc.closeFn != nil {
|
||||
return rc.closeFn()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func DecompressRequestMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if c.Request.Body == nil || c.Request.Method == http.MethodGet {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
maxMB := constant.MaxRequestBodyMB
|
||||
if maxMB <= 0 {
|
||||
maxMB = 32
|
||||
}
|
||||
maxBytes := int64(maxMB) << 20
|
||||
|
||||
origBody := c.Request.Body
|
||||
wrapMaxBytes := func(body io.ReadCloser) io.ReadCloser {
|
||||
return http.MaxBytesReader(c.Writer, body, maxBytes)
|
||||
}
|
||||
decompressed := false
|
||||
|
||||
switch c.GetHeader("Content-Encoding") {
|
||||
case "gzip":
|
||||
gzipReader, err := gzip.NewReader(origBody)
|
||||
if err != nil {
|
||||
_ = origBody.Close()
|
||||
c.AbortWithStatus(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
// Replace the request body with the decompressed data, and enforce a max size (post-decompression).
|
||||
c.Request.Body = wrapMaxBytes(&readCloser{
|
||||
Reader: gzipReader,
|
||||
closeFn: func() error {
|
||||
_ = gzipReader.Close()
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
decompressed = true
|
||||
case "br":
|
||||
reader := brotli.NewReader(origBody)
|
||||
c.Request.Body = wrapMaxBytes(&readCloser{
|
||||
Reader: reader,
|
||||
closeFn: func() error {
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
decompressed = true
|
||||
case "zstd":
|
||||
reader, err := zstd.NewReader(origBody)
|
||||
if err != nil {
|
||||
_ = origBody.Close()
|
||||
c.AbortWithStatus(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
c.Request.Body = wrapMaxBytes(&readCloser{
|
||||
Reader: reader,
|
||||
closeFn: func() error {
|
||||
reader.Close()
|
||||
return origBody.Close()
|
||||
},
|
||||
})
|
||||
decompressed = true
|
||||
default:
|
||||
// Even for uncompressed bodies, enforce a max size to avoid huge request allocations.
|
||||
c.Request.Body = wrapMaxBytes(origBody)
|
||||
}
|
||||
|
||||
if decompressed {
|
||||
c.Request.Header.Del("Content-Encoding")
|
||||
}
|
||||
|
||||
// Continue processing the request
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,135 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type headerNavAccess struct {
|
||||
Enabled bool
|
||||
RequireAuth bool
|
||||
}
|
||||
|
||||
func getHeaderNavAccess(module string) headerNavAccess {
|
||||
fallback := headerNavAccess{
|
||||
Enabled: true,
|
||||
RequireAuth: false,
|
||||
}
|
||||
|
||||
common.OptionMapRWMutex.RLock()
|
||||
raw := common.OptionMap["HeaderNavModules"]
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
|
||||
if strings.TrimSpace(raw) == "" {
|
||||
return fallback
|
||||
}
|
||||
|
||||
var parsed map[string]any
|
||||
if err := common.Unmarshal([]byte(raw), &parsed); err != nil {
|
||||
return fallback
|
||||
}
|
||||
|
||||
return parseHeaderNavAccess(parsed[module], fallback)
|
||||
}
|
||||
|
||||
func parseHeaderNavAccess(raw any, fallback headerNavAccess) headerNavAccess {
|
||||
switch value := raw.(type) {
|
||||
case bool:
|
||||
return headerNavAccess{
|
||||
Enabled: value,
|
||||
RequireAuth: fallback.RequireAuth,
|
||||
}
|
||||
case string:
|
||||
return headerNavAccess{
|
||||
Enabled: parseHeaderNavBool(value, fallback.Enabled),
|
||||
RequireAuth: fallback.RequireAuth,
|
||||
}
|
||||
case float64:
|
||||
return headerNavAccess{
|
||||
Enabled: parseHeaderNavBool(value, fallback.Enabled),
|
||||
RequireAuth: fallback.RequireAuth,
|
||||
}
|
||||
case map[string]any:
|
||||
access := fallback
|
||||
if enabled, ok := value["enabled"]; ok {
|
||||
access.Enabled = parseHeaderNavBool(enabled, fallback.Enabled)
|
||||
}
|
||||
if requireAuth, ok := value["requireAuth"]; ok {
|
||||
access.RequireAuth = parseHeaderNavBool(requireAuth, fallback.RequireAuth)
|
||||
}
|
||||
return access
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
func parseHeaderNavBool(value any, fallback bool) bool {
|
||||
switch v := value.(type) {
|
||||
case bool:
|
||||
return v
|
||||
case string:
|
||||
switch strings.ToLower(strings.TrimSpace(v)) {
|
||||
case "true", "1":
|
||||
return true
|
||||
case "false", "0":
|
||||
return false
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
case float64:
|
||||
if v == 1 {
|
||||
return true
|
||||
}
|
||||
if v == 0 {
|
||||
return false
|
||||
}
|
||||
return fallback
|
||||
case int:
|
||||
if v == 1 {
|
||||
return true
|
||||
}
|
||||
if v == 0 {
|
||||
return false
|
||||
}
|
||||
return fallback
|
||||
default:
|
||||
return fallback
|
||||
}
|
||||
}
|
||||
|
||||
func HeaderNavModuleAuth(module string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
access := getHeaderNavAccess(module)
|
||||
if !access.Enabled {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": fmt.Sprintf("%s is disabled", module),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
if access.RequireAuth {
|
||||
UserAuth()(c)
|
||||
return
|
||||
}
|
||||
|
||||
TryUserAuth()(c)
|
||||
}
|
||||
}
|
||||
|
||||
func HeaderNavModulePublicOrUserAuth(module string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
access := getHeaderNavAccess(module)
|
||||
if !access.Enabled || access.RequireAuth {
|
||||
UserAuth()(c)
|
||||
return
|
||||
}
|
||||
|
||||
TryUserAuth()(c)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"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/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func withHeaderNavModules(t *testing.T, raw string) {
|
||||
t.Helper()
|
||||
|
||||
common.OptionMapRWMutex.Lock()
|
||||
if common.OptionMap == nil {
|
||||
common.OptionMap = map[string]string{}
|
||||
}
|
||||
previous, hadPrevious := common.OptionMap["HeaderNavModules"]
|
||||
common.OptionMap["HeaderNavModules"] = raw
|
||||
common.OptionMapRWMutex.Unlock()
|
||||
|
||||
t.Cleanup(func() {
|
||||
common.OptionMapRWMutex.Lock()
|
||||
defer common.OptionMapRWMutex.Unlock()
|
||||
if hadPrevious {
|
||||
common.OptionMap["HeaderNavModules"] = previous
|
||||
return
|
||||
}
|
||||
delete(common.OptionMap, "HeaderNavModules")
|
||||
})
|
||||
}
|
||||
|
||||
func performHeaderNavRequest(t *testing.T, handler gin.HandlerFunc, authenticated bool) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.GET("/api/test", handler, func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
})
|
||||
|
||||
var accessToken string
|
||||
if authenticated {
|
||||
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{}))
|
||||
model.DB = db
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.RedisEnabled = previousRedis
|
||||
})
|
||||
accessToken = "header-nav-pat"
|
||||
user := model.User{
|
||||
Username: "tester",
|
||||
Password: "unused-password-hash",
|
||||
Role: common.RoleCommonUser,
|
||||
Status: common.UserStatusEnabled,
|
||||
Group: "default",
|
||||
AuthVersion: 1,
|
||||
}
|
||||
user.SetAccessToken(accessToken)
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
}
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/test", nil)
|
||||
if authenticated {
|
||||
request.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
}
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestHeaderNavModuleAuthAllowsDefaultPublicAccess(t *testing.T) {
|
||||
withHeaderNavModules(t, "")
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModuleAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModuleAuthRejectsDisabledPricing(t *testing.T) {
|
||||
raw := `{"pricing":{"enabled":false,"requireAuth":false}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModuleAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModuleAuthRequiresLoginForPricing(t *testing.T) {
|
||||
raw := `{"pricing":{"enabled":true,"requireAuth":true}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModuleAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModuleAuthRequiresLoginForRankings(t *testing.T) {
|
||||
raw := `{"rankings":{"enabled":true,"requireAuth":true}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModuleAuth("rankings"), false)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModuleAuthRejectsLegacyDisabledModule(t *testing.T) {
|
||||
raw := `{"rankings":false}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModuleAuth("rankings"), false)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModulePublicOrUserAuthAllowsDefaultPublicAccess(t *testing.T) {
|
||||
withHeaderNavModules(t, "")
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModulePublicOrUserAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModulePublicOrUserAuthRequiresLoginWhenDisabled(t *testing.T) {
|
||||
raw := `{"pricing":{"enabled":false,"requireAuth":false}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModulePublicOrUserAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModulePublicOrUserAuthAllowsLoggedInWhenDisabled(t *testing.T) {
|
||||
raw := `{"pricing":{"enabled":false,"requireAuth":false}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModulePublicOrUserAuth("pricing"), true)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModulePublicOrUserAuthRequiresLoginWhenRequireAuth(t *testing.T) {
|
||||
raw := `{"pricing":{"enabled":true,"requireAuth":true}}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModulePublicOrUserAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavModulePublicOrUserAuthRequiresLoginForLegacyDisabledModule(t *testing.T) {
|
||||
raw := `{"pricing":false}`
|
||||
withHeaderNavModules(t, raw)
|
||||
|
||||
recorder := performHeaderNavRequest(t, HeaderNavModulePublicOrUserAuth("pricing"), false)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, recorder.Code)
|
||||
}
|
||||
|
||||
func TestHeaderNavPublicRouteRejectsExpiredInternalAccessToken(t *testing.T) {
|
||||
setupDashboardAuthMiddlewareTest(t)
|
||||
withHeaderNavModules(t, "")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
router := gin.New()
|
||||
router.GET("/api/test", HeaderNavModuleAuth("pricing"), func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/test", nil)
|
||||
request.Header.Set("Authorization", "Bearer "+issueExpiredDashboardAccessToken(t, service.AuthIdentity{
|
||||
UserID: 1, SessionID: "expired-header-nav-session", UserAuthVersion: 1, SessionVersion: 1,
|
||||
}))
|
||||
response := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(response, request)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, response.Code)
|
||||
require.Contains(t, response.Body.String(), "AUTH_TOKEN_EXPIRED")
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
)
|
||||
|
||||
// I18n middleware detects and sets the language preference for the request
|
||||
func I18n() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
lang := detectLanguage(c)
|
||||
c.Set(string(constant.ContextKeyLanguage), lang)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// detectLanguage determines the language preference for the request
|
||||
// Priority: 1. User setting (if logged in) -> 2. Accept-Language header -> 3. Default language
|
||||
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 {
|
||||
if userSetting.Language != "" && i18n.IsSupported(userSetting.Language) {
|
||||
return userSetting.Language
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const RouteTagKey = "route_tag"
|
||||
|
||||
func RouteTag(tag string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Set(RouteTagKey, tag)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func SetUpLogger(server *gin.Engine) {
|
||||
server.Use(redactTaskArtifactAccessQuery())
|
||||
server.Use(gin.LoggerWithFormatter(func(param gin.LogFormatterParams) string {
|
||||
var requestID string
|
||||
if param.Keys != nil {
|
||||
requestID, _ = param.Keys[common.RequestIdKey].(string)
|
||||
}
|
||||
tag, _ := param.Keys[RouteTagKey].(string)
|
||||
if tag == "" {
|
||||
tag = "web"
|
||||
}
|
||||
path := param.Path
|
||||
// OAuth callbacks carry one-time codes and state in the query string.
|
||||
// Redact the log value only; the handler still needs the original query.
|
||||
if strings.HasPrefix(path, "/api/oauth/") || strings.HasPrefix(path, "/oauth/") {
|
||||
path, _, _ = strings.Cut(path, "?")
|
||||
}
|
||||
return fmt.Sprintf("[GIN] %s | %s | %s | %3d | %13v | %15s | %7s %s\n",
|
||||
param.TimeStamp.Format("2006/01/02 - 15:04:05"),
|
||||
tag,
|
||||
requestID,
|
||||
param.StatusCode,
|
||||
param.Latency,
|
||||
param.ClientIP,
|
||||
param.Method,
|
||||
path,
|
||||
)
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/common/limiter"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-redis/redis/v8"
|
||||
)
|
||||
|
||||
const (
|
||||
ModelRequestRateLimitCountMark = "MRRL"
|
||||
ModelRequestRateLimitSuccessCountMark = "MRRLS"
|
||||
modelRateLimitTimeFormat = "2006-01-02T15:04:05.000Z"
|
||||
)
|
||||
|
||||
// 检查Redis中的请求限制
|
||||
func checkRedisRateLimit(ctx context.Context, rdb *redis.Client, key string, maxCount int, duration int64) (bool, error) {
|
||||
// 如果maxCount为0,表示不限制
|
||||
if maxCount == 0 {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// 获取当前计数
|
||||
length, err := rdb.LLen(ctx, key).Result()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
// 如果未达到限制,允许请求
|
||||
if length < int64(maxCount) {
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// 检查时间窗口
|
||||
oldTimeStr, _ := rdb.LIndex(ctx, key, -1).Result()
|
||||
oldTime, err := time.Parse(modelRateLimitTimeFormat, oldTimeStr)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
nowTimeStr := time.Now().UTC().Format(modelRateLimitTimeFormat)
|
||||
nowTime, err := time.Parse(modelRateLimitTimeFormat, nowTimeStr)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
// 如果在时间窗口内已达到限制,拒绝请求
|
||||
subTime := nowTime.Sub(oldTime).Seconds()
|
||||
if int64(subTime) < duration {
|
||||
rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute)
|
||||
return false, nil
|
||||
}
|
||||
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// 记录Redis请求
|
||||
func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxCount int) {
|
||||
// 如果maxCount为0,不记录请求
|
||||
if maxCount == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now().UTC().Format(modelRateLimitTimeFormat)
|
||||
rdb.LPush(ctx, key, now)
|
||||
rdb.LTrim(ctx, key, 0, int64(maxCount-1))
|
||||
rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute)
|
||||
}
|
||||
|
||||
// Redis限流处理器
|
||||
func redisRateLimitHandler(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)
|
||||
allowed, err := checkRedisRateLimit(ctx, rdb, successKey, successMaxCount, duration)
|
||||
if err != nil {
|
||||
fmt.Println("检查成功请求数限制失败:", err.Error())
|
||||
abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed")
|
||||
return
|
||||
}
|
||||
if !allowed {
|
||||
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount))
|
||||
return
|
||||
}
|
||||
|
||||
//2.检查总请求数限制并记录总请求(当totalMaxCount为0时会自动跳过,使用令牌桶限流器
|
||||
if totalMaxCount > 0 {
|
||||
totalKey := fmt.Sprintf("rateLimit:%s", userId)
|
||||
// 初始化
|
||||
tb := limiter.New(ctx, rdb)
|
||||
allowed, err = tb.Allow(
|
||||
ctx,
|
||||
totalKey,
|
||||
limiter.WithCapacity(rateLimitCapacity(totalMaxCount, duration)),
|
||||
limiter.WithRate(int64(totalMaxCount)),
|
||||
limiter.WithRequested(duration),
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
fmt.Println("检查总请求数限制失败:", err.Error())
|
||||
abortWithOpenAiMessage(c, http.StatusInternalServerError, "rate_limit_check_failed")
|
||||
return
|
||||
}
|
||||
|
||||
if !allowed {
|
||||
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到总请求数限制:%d分钟内最多请求%d次,包括失败次数,请检查您的请求是否正确", setting.ModelRequestRateLimitDurationMinutes, totalMaxCount))
|
||||
}
|
||||
}
|
||||
|
||||
// 4. 处理请求
|
||||
c.Next()
|
||||
|
||||
// 5. 如果请求成功,记录成功请求
|
||||
if c.Writer.Status() < 400 {
|
||||
recordRedisRequest(ctx, rdb, successKey, successMaxCount)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 内存限流处理器
|
||||
func memoryRateLimitHandler(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
|
||||
|
||||
// 1. 检查总请求数限制(当totalMaxCount为0时跳过)
|
||||
if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) {
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 2. 检查成功请求数限制
|
||||
// 使用一个临时key来检查限制,这样可以避免实际记录
|
||||
checkKey := successKey + "_check"
|
||||
if !inMemoryRateLimiter.Request(checkKey, successMaxCount, duration) {
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 3. 处理请求
|
||||
c.Next()
|
||||
|
||||
// 4. 如果请求成功,记录到实际的成功请求计数中
|
||||
if c.Writer.Status() < 400 {
|
||||
inMemoryRateLimiter.Request(successKey, successMaxCount, duration)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ModelRequestRateLimit 模型请求限流中间件
|
||||
func ModelRequestRateLimit() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
// 在每个请求时检查是否启用限流
|
||||
if !setting.ModelRequestRateLimitEnabled {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 计算限流参数
|
||||
duration := rateLimitDurationSeconds(setting.ModelRequestRateLimitDurationMinutes)
|
||||
totalMaxCount := setting.ModelRequestRateLimitCount
|
||||
successMaxCount := setting.ModelRequestRateLimitSuccessCount
|
||||
|
||||
// 获取分组
|
||||
group := common.GetContextKeyString(c, constant.ContextKeyTokenGroup)
|
||||
if group == "" {
|
||||
group = common.GetContextKeyString(c, constant.ContextKeyUserGroup)
|
||||
}
|
||||
|
||||
//获取分组的限流配置
|
||||
groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group)
|
||||
if found {
|
||||
totalMaxCount = groupTotalCount
|
||||
successMaxCount = groupSuccessCount
|
||||
}
|
||||
|
||||
// 根据存储类型选择并执行限流处理器
|
||||
if common.RedisEnabled {
|
||||
redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
|
||||
} else {
|
||||
memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func rateLimitDurationSeconds(durationMinutes int) int64 {
|
||||
if durationMinutes <= 0 {
|
||||
return 0
|
||||
}
|
||||
minutes := int64(durationMinutes)
|
||||
if minutes > math.MaxInt64/60 {
|
||||
return math.MaxInt64
|
||||
}
|
||||
return minutes * 60
|
||||
}
|
||||
|
||||
func rateLimitCapacity(count int, durationSeconds int64) int64 {
|
||||
if count <= 0 || durationSeconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
c := int64(count)
|
||||
if c > math.MaxInt64/durationSeconds {
|
||||
return math.MaxInt64
|
||||
}
|
||||
return c * durationSeconds
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestModelRedisRateLimitUsesUTCRegardlessOfLocalTimezone(t *testing.T) {
|
||||
redisServer, redisClient := useRateLimitMiniRedis(t)
|
||||
previousLocation := time.Local
|
||||
time.Local = time.FixedZone("test-utc-plus-eight", 8*60*60)
|
||||
t.Cleanup(func() { time.Local = previousLocation })
|
||||
|
||||
ctx := context.Background()
|
||||
recordKey := "rateLimit:model-utc-record"
|
||||
recordRedisRequest(ctx, redisClient, recordKey, 2)
|
||||
recorded, err := redisClient.LIndex(ctx, recordKey, 0).Result()
|
||||
require.NoError(t, err)
|
||||
recordedAt, err := time.Parse(modelRateLimitTimeFormat, recorded)
|
||||
require.NoError(t, err)
|
||||
assert.WithinDuration(t, time.Now().UTC(), recordedAt, 2*time.Second)
|
||||
|
||||
checkKey := "rateLimit:model-utc-check"
|
||||
withinWindow := time.Now().UTC().Add(-30 * time.Second).Format(modelRateLimitTimeFormat)
|
||||
_, err = redisServer.Push(checkKey, withinWindow, withinWindow)
|
||||
require.NoError(t, err)
|
||||
allowed, err := checkRedisRateLimit(ctx, redisClient, checkKey, 2, 60)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, allowed, "an existing UTC timestamp inside the window must remain limited on a non-UTC host")
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// SystemPerformanceCheck 检查系统性能中间件
|
||||
func SystemPerformanceCheck() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 仅检查 Relay 接口 (/v1, /v1beta 等)
|
||||
// 这里简单判断路径前缀,可以根据实际路由调整
|
||||
path := c.Request.URL.Path
|
||||
if strings.HasPrefix(path, "/v1/messages") {
|
||||
if err := checkSystemPerformance(); err != nil {
|
||||
c.JSON(err.StatusCode, gin.H{
|
||||
"error": err.ToClaudeError(),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if err := checkSystemPerformance(); err != nil {
|
||||
c.JSON(err.StatusCode, gin.H{
|
||||
"error": err.ToOpenAIError(),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// checkSystemPerformance 检查系统性能是否超过阈值
|
||||
func checkSystemPerformance() *types.NewAPIError {
|
||||
config := common.GetPerformanceMonitorConfig()
|
||||
if !config.Enabled {
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const redisRateLimitNamespace = "rateLimit:v2"
|
||||
|
||||
// Redis rate limiting intentionally uses a fixed window. The single Lua script
|
||||
// makes increment, expiry, and the limit decision atomic, while retaining the
|
||||
// simple fixed-window behavior: traffic at a window boundary can burst up to
|
||||
// twice the configured limit. Do not replace this with a sliding-window ZSET
|
||||
// unless that externally visible behavior is intentionally changed.
|
||||
const redisFixedWindowScript = `
|
||||
local count = redis.call('INCR', KEYS[1])
|
||||
if count == 1 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
end
|
||||
local ttl = redis.call('TTL', KEYS[1])
|
||||
if ttl < 0 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[2])
|
||||
ttl = redis.call('TTL', KEYS[1])
|
||||
end
|
||||
if count > tonumber(ARGV[1]) then
|
||||
return {0, count, ttl}
|
||||
end
|
||||
return {1, count, ttl}
|
||||
`
|
||||
|
||||
var inMemoryRateLimiter common.InMemoryRateLimiter
|
||||
|
||||
var defNext = func(c *gin.Context) {
|
||||
c.Next()
|
||||
}
|
||||
|
||||
func redisIPRateLimitKey(mark string, clientIP string) string {
|
||||
return fmt.Sprintf("%s:ip:%s:%s", redisRateLimitNamespace, mark, clientIP)
|
||||
}
|
||||
|
||||
func redisUserRateLimitKey(mark string, userID int) string {
|
||||
return fmt.Sprintf("%s:user:%s:%d", redisRateLimitNamespace, mark, userID)
|
||||
}
|
||||
|
||||
func redisReplyInteger(value any) (int64, error) {
|
||||
switch typed := value.(type) {
|
||||
case int64:
|
||||
return typed, nil
|
||||
case string:
|
||||
return strconv.ParseInt(typed, 10, 64)
|
||||
case []byte:
|
||||
return strconv.ParseInt(string(typed), 10, 64)
|
||||
default:
|
||||
return 0, fmt.Errorf("unexpected Redis integer reply type %T", value)
|
||||
}
|
||||
}
|
||||
|
||||
func redisFixedWindowTake(ctx context.Context, key string, maxRequestNum int, duration int64) (bool, int64, int64, error) {
|
||||
if common.RDB == nil {
|
||||
return false, 0, 0, errors.New("Redis client is not initialized")
|
||||
}
|
||||
if key == "" {
|
||||
return false, 0, 0, errors.New("rate limit key is empty")
|
||||
}
|
||||
if maxRequestNum <= 0 {
|
||||
return false, 0, 0, errors.New("rate limit maximum must be positive")
|
||||
}
|
||||
if duration <= 0 {
|
||||
return false, 0, 0, errors.New("rate limit duration must be positive")
|
||||
}
|
||||
|
||||
values, err := common.RDB.Eval(
|
||||
ctx,
|
||||
redisFixedWindowScript,
|
||||
[]string{key},
|
||||
maxRequestNum,
|
||||
duration,
|
||||
).Slice()
|
||||
if err != nil {
|
||||
return false, 0, 0, err
|
||||
}
|
||||
if len(values) != 3 {
|
||||
return false, 0, 0, fmt.Errorf("unexpected Redis rate limit reply length %d", len(values))
|
||||
}
|
||||
|
||||
allowedValue, err := redisReplyInteger(values[0])
|
||||
if err != nil {
|
||||
return false, 0, 0, err
|
||||
}
|
||||
count, err := redisReplyInteger(values[1])
|
||||
if err != nil {
|
||||
return false, 0, 0, err
|
||||
}
|
||||
ttlSeconds, err := redisReplyInteger(values[2])
|
||||
if err != nil {
|
||||
return false, 0, 0, err
|
||||
}
|
||||
|
||||
return allowedValue == 1, count, ttlSeconds, nil
|
||||
}
|
||||
|
||||
func redisRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) {
|
||||
allowed, _, ttlSeconds, err := redisFixedWindowTake(
|
||||
c.Request.Context(),
|
||||
redisIPRateLimitKey(mark, c.ClientIP()),
|
||||
maxRequestNum,
|
||||
duration,
|
||||
)
|
||||
if err != nil {
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("rate limit check failed (mark=%s): %v", mark, err))
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if !allowed {
|
||||
writeRateLimited(c, ttlSeconds)
|
||||
}
|
||||
}
|
||||
|
||||
func memoryRateLimiter(c *gin.Context, maxRequestNum int, duration int64, mark string) {
|
||||
key := mark + c.ClientIP()
|
||||
if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) {
|
||||
writeRateLimited(c, duration)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// writeRateLimited rejects the request with 429 and a Retry-After hint so
|
||||
// clients can back off instead of treating the rejection as a fatal error.
|
||||
// The in-memory limiter cannot report the remaining window, so callers
|
||||
// without a TTL pass the full window duration as a conservative upper bound.
|
||||
func writeRateLimited(c *gin.Context, retryAfterSeconds int64) {
|
||||
if retryAfterSeconds > 0 {
|
||||
c.Header("Retry-After", strconv.FormatInt(retryAfterSeconds, 10))
|
||||
}
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gin.Context) {
|
||||
if common.RedisEnabled {
|
||||
return func(c *gin.Context) {
|
||||
redisRateLimiter(c, maxRequestNum, duration, mark)
|
||||
}
|
||||
}
|
||||
// It's safe to call multi times.
|
||||
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
|
||||
return func(c *gin.Context) {
|
||||
memoryRateLimiter(c, maxRequestNum, duration, mark)
|
||||
}
|
||||
}
|
||||
|
||||
func GlobalWebRateLimit() func(c *gin.Context) {
|
||||
if common.GlobalWebRateLimitEnable {
|
||||
return rateLimitFactory(common.GlobalWebRateLimitNum, common.GlobalWebRateLimitDuration, "GW")
|
||||
}
|
||||
return defNext
|
||||
}
|
||||
|
||||
func GlobalAPIRateLimit() func(c *gin.Context) {
|
||||
if common.GlobalApiRateLimitEnable {
|
||||
return rateLimitFactory(common.GlobalApiRateLimitNum, common.GlobalApiRateLimitDuration, "GA")
|
||||
}
|
||||
return defNext
|
||||
}
|
||||
|
||||
func CriticalRateLimit() func(c *gin.Context) {
|
||||
if common.CriticalRateLimitEnable {
|
||||
return rateLimitFactory(common.CriticalRateLimitNum, common.CriticalRateLimitDuration, "CT")
|
||||
}
|
||||
return defNext
|
||||
}
|
||||
|
||||
func UserCriticalRateLimit(scope string) func(c *gin.Context) {
|
||||
if !common.CriticalRateLimitEnable {
|
||||
return defNext
|
||||
}
|
||||
return userRateLimitFactory(
|
||||
common.CriticalRateLimitNum,
|
||||
common.CriticalRateLimitDuration,
|
||||
"UC:"+scope,
|
||||
)
|
||||
}
|
||||
|
||||
func DownloadRateLimit() func(c *gin.Context) {
|
||||
return rateLimitFactory(common.DownloadRateLimitNum, common.DownloadRateLimitDuration, "DW")
|
||||
}
|
||||
|
||||
func UploadRateLimit() func(c *gin.Context) {
|
||||
return rateLimitFactory(common.UploadRateLimitNum, common.UploadRateLimitDuration, "UP")
|
||||
}
|
||||
|
||||
// userRateLimitFactory creates a rate limiter keyed by authenticated user ID
|
||||
// instead of client IP, making it resistant to proxy rotation attacks.
|
||||
// Must be used AFTER authentication middleware (UserAuth).
|
||||
func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gin.Context) {
|
||||
if common.RedisEnabled {
|
||||
return func(c *gin.Context) {
|
||||
userID := c.GetInt("id")
|
||||
if userID == 0 {
|
||||
c.Status(http.StatusUnauthorized)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
userRedisRateLimiter(c, maxRequestNum, duration, redisUserRateLimitKey(mark, userID))
|
||||
}
|
||||
}
|
||||
// It's safe to call multi times.
|
||||
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
|
||||
return func(c *gin.Context) {
|
||||
userID := c.GetInt("id")
|
||||
if userID == 0 {
|
||||
c.Status(http.StatusUnauthorized)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
key := fmt.Sprintf("%s:user:%d", mark, userID)
|
||||
if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) {
|
||||
writeRateLimited(c, duration)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// userRedisRateLimiter is like redisRateLimiter but accepts a pre-built key
|
||||
// (to support user-ID-based keys).
|
||||
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))
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if !allowed {
|
||||
writeRateLimited(c, ttlSeconds)
|
||||
}
|
||||
}
|
||||
|
||||
// SearchRateLimit returns a per-user rate limiter for search endpoints.
|
||||
// Configurable via SEARCH_RATE_LIMIT_ENABLE / SEARCH_RATE_LIMIT / SEARCH_RATE_LIMIT_DURATION.
|
||||
func SearchRateLimit() func(c *gin.Context) {
|
||||
if !common.SearchRateLimitEnable {
|
||||
return defNext
|
||||
}
|
||||
return userRateLimitFactory(common.SearchRateLimitNum, common.SearchRateLimitDuration, "SR")
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-redis/redis/v8"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func useRateLimitMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) {
|
||||
t.Helper()
|
||||
|
||||
previousRedisEnabled := common.RedisEnabled
|
||||
previousRedisClient := common.RDB
|
||||
redisServer := miniredis.RunT(t)
|
||||
redisClient := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
|
||||
require.NoError(t, redisClient.Ping(context.Background()).Err())
|
||||
|
||||
common.RedisEnabled = true
|
||||
common.RDB = redisClient
|
||||
t.Cleanup(func() {
|
||||
_ = redisClient.Close()
|
||||
common.RedisEnabled = previousRedisEnabled
|
||||
common.RDB = previousRedisClient
|
||||
})
|
||||
|
||||
return redisServer, redisClient
|
||||
}
|
||||
|
||||
func performRateLimitRequest(router http.Handler, path string, remoteAddr string) *httptest.ResponseRecorder {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
request.RemoteAddr = remoteAddr
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestRedisIPRateLimiterThresholdTTLAndNamespace(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies(nil))
|
||||
router.GET("/limited", rateLimitFactory(2, 37, "TEST"), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
remoteAddr := "192.0.2.10:12345"
|
||||
legacyKey := "rateLimit:TEST192.0.2.10"
|
||||
_, err := redisServer.Push(legacyKey, "legacy-list-entry")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code)
|
||||
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", remoteAddr).Code)
|
||||
limitedResponse := performRateLimitRequest(router, "/limited", remoteAddr)
|
||||
assert.Equal(t, http.StatusTooManyRequests, limitedResponse.Code)
|
||||
assert.Equal(t, "37", limitedResponse.Header().Get("Retry-After"))
|
||||
|
||||
key := redisIPRateLimitKey("TEST", "192.0.2.10")
|
||||
count, err := redisServer.Get(key)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "3", count)
|
||||
assert.Equal(t, 37*time.Second, redisServer.TTL(key))
|
||||
assert.True(t, redisServer.Exists(legacyKey), "the v2 counter must not touch an old list key")
|
||||
}
|
||||
|
||||
func TestRedisUserRateLimiterUsesSharedFixedWindow(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
|
||||
router := gin.New()
|
||||
router.GET(
|
||||
"/limited",
|
||||
func(c *gin.Context) { c.Set("id", 42) },
|
||||
userRateLimitFactory(1, 23, "USER"),
|
||||
func(c *gin.Context) { c.Status(http.StatusNoContent) },
|
||||
)
|
||||
|
||||
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/limited", "192.0.2.20:12345").Code)
|
||||
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/limited", "198.51.100.20:12345").Code)
|
||||
|
||||
key := redisUserRateLimitKey("USER", 42)
|
||||
assert.True(t, redisServer.Exists(key))
|
||||
assert.Equal(t, 23*time.Second, redisServer.TTL(key))
|
||||
}
|
||||
|
||||
func TestRedisEmailVerificationRateLimiterPreservesResponseAndTTL(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies(nil))
|
||||
router.GET("/verify", EmailVerificationRateLimit(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
remoteAddr := "192.0.2.30:12345"
|
||||
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/verify", remoteAddr).Code)
|
||||
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())
|
||||
|
||||
key := redisIPRateLimitKey(EmailVerificationRateLimitMark, "192.0.2.30")
|
||||
assert.True(t, redisServer.Exists(key))
|
||||
assert.Equal(t, time.Duration(EmailVerificationDuration)*time.Second, redisServer.TTL(key))
|
||||
}
|
||||
|
||||
func TestRedisFixedWindowIsAtomicUnderConcurrency(t *testing.T) {
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
const (
|
||||
requestCount = 20
|
||||
maximumCount = 7
|
||||
duration = int64(41)
|
||||
)
|
||||
key := redisIPRateLimitKey("CONCURRENT", "192.0.2.40")
|
||||
|
||||
var allowedCount atomic.Int64
|
||||
errorsFound := make(chan error, requestCount)
|
||||
var waitGroup sync.WaitGroup
|
||||
waitGroup.Add(requestCount)
|
||||
for range requestCount {
|
||||
go func() {
|
||||
defer waitGroup.Done()
|
||||
allowed, _, _, err := redisFixedWindowTake(context.Background(), key, maximumCount, duration)
|
||||
if err != nil {
|
||||
errorsFound <- err
|
||||
return
|
||||
}
|
||||
if allowed {
|
||||
allowedCount.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
waitGroup.Wait()
|
||||
close(errorsFound)
|
||||
for err := range errorsFound {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
assert.Equal(t, int64(maximumCount), allowedCount.Load())
|
||||
count, err := redisServer.Get(key)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "20", count)
|
||||
assert.Equal(t, time.Duration(duration)*time.Second, redisServer.TTL(key))
|
||||
}
|
||||
|
||||
func TestRedisFixedWindowResetsAtBoundary(t *testing.T) {
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
const duration = int64(10)
|
||||
key := redisIPRateLimitKey("BOUNDARY", "192.0.2.50")
|
||||
|
||||
for range 2 {
|
||||
allowed, _, _, err := redisFixedWindowTake(context.Background(), key, 2, duration)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, allowed)
|
||||
}
|
||||
allowed, _, _, err := redisFixedWindowTake(context.Background(), key, 2, duration)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, allowed)
|
||||
|
||||
// This reset is intentional fixed-window behavior. A client can consume one
|
||||
// full allowance immediately before and another immediately after a boundary.
|
||||
redisServer.FastForward(time.Duration(duration) * time.Second)
|
||||
for range 2 {
|
||||
allowed, _, _, err = redisFixedWindowTake(context.Background(), key, 2, duration)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, allowed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRedisFixedWindowRepairsCounterWithoutTTL(t *testing.T) {
|
||||
redisServer, _ := useRateLimitMiniRedis(t)
|
||||
const duration = int64(29)
|
||||
key := redisIPRateLimitKey("MISSING-TTL", "192.0.2.51")
|
||||
redisServer.Set(key, "5")
|
||||
|
||||
allowed, count, ttl, err := redisFixedWindowTake(context.Background(), key, 3, duration)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, allowed)
|
||||
assert.Equal(t, int64(6), count)
|
||||
assert.Equal(t, duration, ttl)
|
||||
assert.Equal(t, time.Duration(duration)*time.Second, redisServer.TTL(key))
|
||||
|
||||
redisServer.FastForward(time.Duration(duration) * time.Second)
|
||||
assert.False(t, redisServer.Exists(key), "a recovered counter must not remain permanently rate-limited")
|
||||
}
|
||||
|
||||
func TestRedisFailurePolicies(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
_, redisClient := useRateLimitMiniRedis(t)
|
||||
require.NoError(t, redisClient.Close())
|
||||
|
||||
router := gin.New()
|
||||
require.NoError(t, router.SetTrustedProxies(nil))
|
||||
router.GET("/ip", rateLimitFactory(1, 30, "FAIL-IP"), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
router.GET(
|
||||
"/user",
|
||||
func(c *gin.Context) { c.Set("id", 7) },
|
||||
userRateLimitFactory(1, 30, "FAIL-USER"),
|
||||
func(c *gin.Context) { c.Status(http.StatusNoContent) },
|
||||
)
|
||||
router.GET("/email", EmailVerificationRateLimit(), func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
ipResponse := performRateLimitRequest(router, "/ip", "192.0.2.60:12345")
|
||||
assert.Equal(t, http.StatusInternalServerError, ipResponse.Code)
|
||||
assert.Empty(t, ipResponse.Body.String())
|
||||
userResponse := performRateLimitRequest(router, "/user", "192.0.2.61:12345")
|
||||
assert.Equal(t, http.StatusInternalServerError, userResponse.Code)
|
||||
assert.Empty(t, userResponse.Body.String())
|
||||
assert.Equal(t, http.StatusNoContent, performRateLimitRequest(router, "/email", "192.0.2.62:12345").Code)
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"runtime/debug"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
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())))
|
||||
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),
|
||||
"type": "new_api_panic",
|
||||
},
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
}()
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func RequestId() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
id := common.NewRequestId()
|
||||
c.Set(common.RequestIdKey, id)
|
||||
ctx := context.WithValue(c.Request.Context(), common.RequestIdKey, id)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
c.Header(common.RequestIdKey, id)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func AnonymousRequestBodyLimit() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
maxBytes := common.GetAnonymousRequestBodyLimitBytes()
|
||||
if maxBytes <= 0 || c.Request.Body == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
originalBody := c.Request.Body
|
||||
limitedBody, err := readAnonymousRequestBody(originalBody, maxBytes)
|
||||
_ = originalBody.Close()
|
||||
if err != nil {
|
||||
if common.IsRequestBodyTooLargeError(err) {
|
||||
c.AbortWithStatus(http.StatusRequestEntityTooLarge)
|
||||
return
|
||||
}
|
||||
c.AbortWithStatus(http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(limitedBody))
|
||||
c.Request.ContentLength = int64(len(limitedBody))
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func readAnonymousRequestBody(body io.Reader, maxBytes int64) ([]byte, error) {
|
||||
data, err := io.ReadAll(io.LimitReader(body, maxBytes+1))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if int64(len(data)) > maxBytes {
|
||||
return nil, common.ErrRequestBodyTooLarge
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// SecureVerificationRequired protects channel key disclosure. Other sensitive
|
||||
// operations validate their narrower proof scopes in their controller.
|
||||
func SecureVerificationRequired() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
channelID, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil || channelID <= 0 {
|
||||
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"success": false, "code": "SECURITY_CONTEXT_INVALID", "message": service.ErrVerificationContextInvalid.Error()})
|
||||
return
|
||||
}
|
||||
context, err := common.Marshal(service.ChannelKeyReadContext{ChannelID: channelID})
|
||||
if err != nil {
|
||||
c.AbortWithStatus(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if RequireSecurityProof(c, service.VerificationOperation{Scope: service.VerificationScopeChannelKeyRead, Context: context}) == nil {
|
||||
return
|
||||
}
|
||||
c.Set("secure_verified", true)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// RequireSecurityProof validates a proof against the authenticated dashboard
|
||||
// session and writes the shared proof error contract on failure.
|
||||
func RequireSecurityProof(c *gin.Context, operation service.VerificationOperation) *model.AuthFlowAuthorization {
|
||||
identity, ok := GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
securityProofError(c, "SECURITY_PROOF_INVALID", "安全验证状态无效")
|
||||
return nil
|
||||
}
|
||||
raw := strings.TrimSpace(c.GetHeader("X-Security-Proof"))
|
||||
if raw == "" {
|
||||
securityProofError(c, "SECURITY_PROOF_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", "安全验证已过期")
|
||||
case errors.Is(err, service.ErrProofScope):
|
||||
securityProofError(c, "SECURITY_PROOF_SCOPE_MISMATCH", "安全验证范围不匹配")
|
||||
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", "安全验证方式不匹配")
|
||||
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):
|
||||
securityProofError(c, "SECURITY_PROOF_CONTEXT_MISMATCH", "Verification does not match this action's details. Please verify again.")
|
||||
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", "安全验证状态无效")
|
||||
default:
|
||||
_ = c.Error(err)
|
||||
c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"success": false, "code": "AUTH_INTERNAL_ERROR", "message": "Please try again later."})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return authorization
|
||||
}
|
||||
|
||||
func securityProofError(c *gin.Context, code, message string) {
|
||||
c.Set("security_error_code", code)
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": message,
|
||||
"code": code,
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// HTTPStats 存储HTTP统计信息
|
||||
type HTTPStats struct {
|
||||
activeConnections int64
|
||||
}
|
||||
|
||||
var globalStats = &HTTPStats{}
|
||||
|
||||
// StatsMiddleware 统计中间件
|
||||
func StatsMiddleware() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
// 增加活跃连接数
|
||||
atomic.AddInt64(&globalStats.activeConnections, 1)
|
||||
|
||||
// 确保在请求结束时减少连接数
|
||||
defer func() {
|
||||
atomic.AddInt64(&globalStats.activeConnections, -1)
|
||||
}()
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// StatsInfo 统计信息结构
|
||||
type StatsInfo struct {
|
||||
ActiveConnections int64 `json:"active_connections"`
|
||||
}
|
||||
|
||||
// GetStats 获取统计信息
|
||||
func GetStats() StatsInfo {
|
||||
return StatsInfo{
|
||||
ActiveConnections: atomic.LoadInt64(&globalStats.activeConnections),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const TaskArtifactAccessContextKey = "task_artifact_access"
|
||||
|
||||
const (
|
||||
taskArtifactAccessRawContextKey = "task_artifact_access_raw"
|
||||
taskArtifactAccessPresentContextKey = "task_artifact_access_present"
|
||||
taskArtifactAccessInvalidContextKey = "task_artifact_access_invalid"
|
||||
taskArtifactAccessRateWindow = time.Minute
|
||||
taskArtifactAccessCleanupInterval = time.Minute
|
||||
maxEncodedTaskArtifactAccessQuerySize = 128
|
||||
)
|
||||
|
||||
type taskArtifactRateEntry struct {
|
||||
windowStart time.Time
|
||||
count int
|
||||
}
|
||||
|
||||
type taskArtifactAccessLimiter struct {
|
||||
mutex sync.Mutex
|
||||
global int
|
||||
byIP map[string]int
|
||||
byObject map[string]int
|
||||
rates map[string]taskArtifactRateEntry
|
||||
nextCleanup time.Time
|
||||
limits system_setting.TaskArtifactAccessLimits
|
||||
}
|
||||
|
||||
var taskArtifactAnonymousLimiter = newTaskArtifactAccessLimiter(
|
||||
system_setting.LoadTaskArtifactAccessLimits(),
|
||||
)
|
||||
|
||||
func newTaskArtifactAccessLimiter(limits system_setting.TaskArtifactAccessLimits) *taskArtifactAccessLimiter {
|
||||
return &taskArtifactAccessLimiter{
|
||||
byIP: make(map[string]int),
|
||||
byObject: make(map[string]int),
|
||||
rates: make(map[string]taskArtifactRateEntry),
|
||||
limits: limits,
|
||||
}
|
||||
}
|
||||
|
||||
func (l *taskArtifactAccessLimiter) invalidAttempt(now time.Time, ip string) bool {
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
|
||||
if l.nextCleanup.IsZero() || !now.Before(l.nextCleanup) {
|
||||
for key, entry := range l.rates {
|
||||
if now.Sub(entry.windowStart) >= taskArtifactAccessRateWindow {
|
||||
delete(l.rates, key)
|
||||
}
|
||||
}
|
||||
l.nextCleanup = now.Add(taskArtifactAccessCleanupInterval)
|
||||
}
|
||||
|
||||
rate := l.rates[ip]
|
||||
if rate.windowStart.IsZero() || now.Sub(rate.windowStart) >= taskArtifactAccessRateWindow {
|
||||
rate = taskArtifactRateEntry{windowStart: now}
|
||||
}
|
||||
if rate.count >= l.limits.InvalidRatePerMinute {
|
||||
return false
|
||||
}
|
||||
rate.count++
|
||||
l.rates[ip] = rate
|
||||
return true
|
||||
}
|
||||
|
||||
func (l *taskArtifactAccessLimiter) acquire(ip, taskID, artifactKey string) (func(), bool) {
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
objectKey := taskID + "\x00" + artifactKey
|
||||
if l.global >= l.limits.GlobalConcurrency ||
|
||||
l.byIP[ip] >= l.limits.IPConcurrency ||
|
||||
l.byObject[objectKey] >= l.limits.ObjectConcurrency {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
l.global++
|
||||
l.byIP[ip]++
|
||||
l.byObject[objectKey]++
|
||||
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
l.global--
|
||||
l.byIP[ip]--
|
||||
l.byObject[objectKey]--
|
||||
if l.byIP[ip] == 0 {
|
||||
delete(l.byIP, ip)
|
||||
}
|
||||
if l.byObject[objectKey] == 0 {
|
||||
delete(l.byObject, objectKey)
|
||||
}
|
||||
})
|
||||
}, true
|
||||
}
|
||||
|
||||
func redactTaskArtifactAccessQuery() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
path := c.Request.URL.Path
|
||||
isArtifactContent := strings.HasPrefix(path, "/v1/tasks/") &&
|
||||
strings.Contains(path, "/artifacts/") &&
|
||||
strings.HasSuffix(path, "/content")
|
||||
isLegacyVideoContent := strings.HasPrefix(path, "/v1/videos/") &&
|
||||
strings.HasSuffix(path, "/content")
|
||||
if !isArtifactContent && !isLegacyVideoContent {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
rawAccess, present, invalid := popTaskArtifactAccessQuery(c.Request)
|
||||
if present {
|
||||
c.Set(taskArtifactAccessRawContextKey, rawAccess)
|
||||
c.Set(taskArtifactAccessPresentContextKey, true)
|
||||
c.Set(taskArtifactAccessInvalidContextKey, invalid)
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func popTaskArtifactAccessQuery(request *http.Request) (string, bool, bool) {
|
||||
if request == nil || request.URL == nil {
|
||||
return "", false, false
|
||||
}
|
||||
rawAccess := ""
|
||||
count := 0
|
||||
invalid := false
|
||||
kept := make([]string, 0)
|
||||
for part := range strings.SplitSeq(request.URL.RawQuery, "&") {
|
||||
rawKey, rawValue, _ := strings.Cut(part, "=")
|
||||
key, err := url.QueryUnescape(rawKey)
|
||||
if err != nil || key != service.TaskArtifactAccessQueryParameter {
|
||||
kept = append(kept, part)
|
||||
continue
|
||||
}
|
||||
count++
|
||||
if len(rawValue) > maxEncodedTaskArtifactAccessQuerySize {
|
||||
invalid = true
|
||||
continue
|
||||
}
|
||||
if count == 1 {
|
||||
value, decodeErr := url.QueryUnescape(rawValue)
|
||||
if decodeErr != nil {
|
||||
invalid = true
|
||||
} else {
|
||||
rawAccess = value
|
||||
}
|
||||
}
|
||||
}
|
||||
if count == 0 {
|
||||
return "", false, false
|
||||
}
|
||||
invalid = invalid || count != 1
|
||||
request.URL.RawQuery = strings.Join(kept, "&")
|
||||
request.RequestURI = request.URL.RequestURI()
|
||||
return rawAccess, true, invalid
|
||||
}
|
||||
|
||||
// TokenOrTaskArtifactAccessAuth accepts the normal relay API Bearer token or a
|
||||
// route-bound capability. Capabilities are verified before any database read.
|
||||
func TokenOrTaskArtifactAccessAuth(taskParam, artifactParam string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
c.Header("Cache-Control", "private, no-store")
|
||||
|
||||
rawAccess := c.GetString(taskArtifactAccessRawContextKey)
|
||||
present := c.GetBool(taskArtifactAccessPresentContextKey)
|
||||
invalid := c.GetBool(taskArtifactAccessInvalidContextKey)
|
||||
if queryAccess, queryPresent, queryInvalid := popTaskArtifactAccessQuery(c.Request); queryPresent {
|
||||
present = true
|
||||
if rawAccess == "" {
|
||||
rawAccess = queryAccess
|
||||
}
|
||||
invalid = invalid || queryInvalid
|
||||
}
|
||||
if !present {
|
||||
TokenAuth()(c)
|
||||
return
|
||||
}
|
||||
|
||||
taskID := c.Param(taskParam)
|
||||
artifactKey := c.Param(artifactParam)
|
||||
ip := c.ClientIP()
|
||||
if ip == "" {
|
||||
ip = "unknown"
|
||||
}
|
||||
if invalid || !service.VerifyTaskArtifactAccess(rawAccess, taskID, artifactKey) {
|
||||
if !taskArtifactAnonymousLimiter.invalidAttempt(time.Now(), ip) {
|
||||
writeTaskArtifactAccessLimited(c)
|
||||
return
|
||||
}
|
||||
writeTaskArtifactAccessNotFound(c)
|
||||
return
|
||||
}
|
||||
|
||||
release, ok := taskArtifactAnonymousLimiter.acquire(ip, taskID, artifactKey)
|
||||
if !ok {
|
||||
writeTaskArtifactAccessLimited(c)
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
|
||||
c.Set(TaskArtifactAccessContextKey, true)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func IsTaskArtifactAccess(c *gin.Context) bool {
|
||||
return c != nil && c.GetBool(TaskArtifactAccessContextKey)
|
||||
}
|
||||
|
||||
func writeTaskArtifactAccessNotFound(c *gin.Context) {
|
||||
c.Header("Cache-Control", "private, no-store")
|
||||
c.AbortWithStatusJSON(http.StatusNotFound, gin.H{
|
||||
"error": gin.H{
|
||||
"message": "Task or artifact not found",
|
||||
"type": "artifact_not_found",
|
||||
"code": "artifact_not_found",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func writeTaskArtifactAccessLimited(c *gin.Context) {
|
||||
c.Header("Cache-Control", "private, no-store")
|
||||
c.Header("Retry-After", "60")
|
||||
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
|
||||
"error": gin.H{
|
||||
"message": "Artifact access limit exceeded",
|
||||
"type": "rate_limit_error",
|
||||
"code": "artifact_access_limited",
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestTaskArtifactAccessIsRedactedAndVerifiedBeforeHandler(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
previousSecret := common.CryptoSecret
|
||||
common.CryptoSecret = "task-artifact-middleware-secret"
|
||||
t.Cleanup(func() { common.CryptoSecret = previousSecret })
|
||||
|
||||
access, err := service.IssueTaskArtifactAccess("task-1", "video-main")
|
||||
require.NoError(t, err)
|
||||
|
||||
router := gin.New()
|
||||
router.Use(redactTaskArtifactAccessQuery())
|
||||
router.GET(
|
||||
"/v1/tasks/:key/artifacts/:artifact_key/content",
|
||||
TokenOrTaskArtifactAccessAuth("key", "artifact_key"),
|
||||
func(c *gin.Context) {
|
||||
assert.True(t, IsTaskArtifactAccess(c))
|
||||
assert.NotContains(t, c.Request.URL.RawQuery, service.TaskArtifactAccessQueryParameter)
|
||||
assert.Equal(t, "kept", c.Query("keep"))
|
||||
c.Status(http.StatusNoContent)
|
||||
},
|
||||
)
|
||||
request := httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/task-1/artifacts/video-main/content?access="+urlQueryEscape(access)+"&keep=kept",
|
||||
nil,
|
||||
)
|
||||
request.RemoteAddr = "192.0.2.1:1234"
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
|
||||
func TestTaskArtifactAccessRejectsTamperedAndEmptyCapabilitiesAsNotFound(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
previousSecret := common.CryptoSecret
|
||||
common.CryptoSecret = "task-artifact-middleware-reject-secret"
|
||||
t.Cleanup(func() { common.CryptoSecret = previousSecret })
|
||||
|
||||
router := gin.New()
|
||||
router.Use(redactTaskArtifactAccessQuery())
|
||||
router.GET(
|
||||
"/v1/tasks/:key/artifacts/:artifact_key/content",
|
||||
TokenOrTaskArtifactAccessAuth("key", "artifact_key"),
|
||||
func(c *gin.Context) { c.Status(http.StatusNoContent) },
|
||||
)
|
||||
|
||||
for _, query := range []string{
|
||||
"?access=",
|
||||
"?access=invalid",
|
||||
"?access=first&access=second",
|
||||
"?access=" + strings.Repeat("x", 1024),
|
||||
"?access=%20" + strings.Repeat("A", 43) + "%20",
|
||||
} {
|
||||
request := httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/task-1/artifacts/video-main/content"+query,
|
||||
nil,
|
||||
)
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.Equal(t, http.StatusNotFound, recorder.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTaskArtifactAccessLimiterDefaults(t *testing.T) {
|
||||
limits := system_setting.TaskArtifactAccessLimits{
|
||||
InvalidRatePerMinute: system_setting.DefaultTaskArtifactInvalidRateLimitPerMinute,
|
||||
GlobalConcurrency: system_setting.DefaultTaskArtifactGlobalConcurrency,
|
||||
IPConcurrency: system_setting.DefaultTaskArtifactIPConcurrency,
|
||||
ObjectConcurrency: system_setting.DefaultTaskArtifactObjectConcurrency,
|
||||
}
|
||||
limiter := newTaskArtifactAccessLimiter(limits)
|
||||
now := time.Unix(1000, 0)
|
||||
releases := make([]func(), 0, limits.ObjectConcurrency)
|
||||
for i := 0; i < limits.ObjectConcurrency; i++ {
|
||||
release, ok := limiter.acquire("192.0.2.1", "task-1", "video")
|
||||
require.True(t, ok)
|
||||
releases = append(releases, release)
|
||||
}
|
||||
_, ok := limiter.acquire("192.0.2.2", "task-1", "video")
|
||||
assert.False(t, ok, "task+key concurrency is shared across IPs")
|
||||
for _, release := range releases {
|
||||
release()
|
||||
}
|
||||
|
||||
rateLimiter := newTaskArtifactAccessLimiter(limits)
|
||||
for i := 0; i < limits.InvalidRatePerMinute; i++ {
|
||||
assert.True(t, rateLimiter.invalidAttempt(now, "192.0.2.10"))
|
||||
}
|
||||
assert.False(t, rateLimiter.invalidAttempt(now, "192.0.2.10"))
|
||||
assert.True(t, rateLimiter.invalidAttempt(now.Add(time.Minute), "192.0.2.10"))
|
||||
}
|
||||
|
||||
func TestRedactTaskArtifactAccessAlsoCoversLegacyVideoRoute(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(redactTaskArtifactAccessQuery())
|
||||
router.GET("/v1/videos/:task_id/content", func(c *gin.Context) {
|
||||
assert.NotContains(t, c.Request.URL.RawQuery, "access")
|
||||
assert.NotContains(t, c.Request.RequestURI, "secret-capability")
|
||||
assert.Equal(t, "ok", c.Query("keep"))
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/videos/task-1/content?access=secret-capability&keep=ok",
|
||||
nil,
|
||||
)
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
|
||||
func TestSetUpLoggerNeverWritesTaskArtifactAccess(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
previousWriter := gin.DefaultWriter
|
||||
var output bytes.Buffer
|
||||
gin.DefaultWriter = &output
|
||||
t.Cleanup(func() { gin.DefaultWriter = previousWriter })
|
||||
|
||||
router := gin.New()
|
||||
SetUpLogger(router)
|
||||
router.GET("/v1/tasks/:key/artifacts/:artifact_key/content", func(c *gin.Context) {
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(
|
||||
http.MethodGet,
|
||||
"/v1/tasks/task-1/artifacts/video/content?access=never-log-this&keep=ok",
|
||||
nil,
|
||||
)
|
||||
router.ServeHTTP(httptest.NewRecorder(), request)
|
||||
|
||||
assert.False(t, strings.Contains(output.String(), "never-log-this"))
|
||||
}
|
||||
|
||||
func urlQueryEscape(value string) string {
|
||||
replacer := strings.NewReplacer("+", "%2B", "=", "%3D")
|
||||
return replacer.Replace(value)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,45 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"io"
|
||||
"mime"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
// rewriteTaskPluginJSONModel patches the top-level JSON "model" field and
|
||||
// replaces BodyStorage. Non-JSON bodies are left untouched.
|
||||
func rewriteTaskPluginJSONModel(c *gin.Context, spelling string) error {
|
||||
mediaType, _, err := mime.ParseMediaType(c.GetHeader("Content-Type"))
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
if mediaType != "application/json" && !strings.HasSuffix(mediaType, "+json") {
|
||||
return nil
|
||||
}
|
||||
storage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
raw, err := storage.Bytes()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
patched, err := sjson.SetBytes(raw, "model", spelling)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
newStorage, err := common.CreateBodyStorage(patched)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_ = storage.Close()
|
||||
c.Set(common.KeyBodyStorage, newStorage)
|
||||
c.Set(common.KeyRequestBody, nil)
|
||||
c.Request.Body = io.NopCloser(newStorage)
|
||||
c.Request.ContentLength = int64(len(patched))
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,500 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
appI18n "github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"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"
|
||||
)
|
||||
|
||||
func setupOriginTaskDB(t *testing.T) {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousType := common.MainDatabaseType()
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}))
|
||||
model.DB = database
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
})
|
||||
}
|
||||
|
||||
func insertOriginTaskChannel(t *testing.T, status int) *model.Channel {
|
||||
t.Helper()
|
||||
channel := &model.Channel{
|
||||
Name: "origin-channel",
|
||||
Key: "sk-origin",
|
||||
Status: status,
|
||||
Type: constant.ChannelTypeDoubaoVideo,
|
||||
}
|
||||
require.NoError(t, model.DB.Create(channel).Error)
|
||||
return channel
|
||||
}
|
||||
|
||||
func insertOriginOwnedTask(t *testing.T, taskID string, userID, channelID int, platform constant.TaskPlatform) *model.Task {
|
||||
t.Helper()
|
||||
task := &model.Task{
|
||||
TaskID: taskID,
|
||||
UserId: userID,
|
||||
ChannelId: channelID,
|
||||
Platform: platform,
|
||||
Action: "text_to_video",
|
||||
Status: model.TaskStatusSuccess,
|
||||
PrivateData: model.TaskPrivateData{
|
||||
UpstreamTaskID: "upstream-" + taskID,
|
||||
},
|
||||
}
|
||||
data, err := common.Marshal(map[string]any{"id": "upstream-" + taskID})
|
||||
require.NoError(t, err)
|
||||
task.Data = data
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
return task
|
||||
}
|
||||
|
||||
func originTaskTestContext(userID int) *gin.Context {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", nil)
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, userID)
|
||||
return c
|
||||
}
|
||||
|
||||
func resolvedOriginPin(c *gin.Context) (int, bool) {
|
||||
pin, found, _ := service.GetChannelConstraints(c).ResolvedPin()
|
||||
if !found {
|
||||
return 0, false
|
||||
}
|
||||
return pin.ChannelId, true
|
||||
}
|
||||
|
||||
func TestApplyOriginTaskIntent(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
enabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
otherEnabled := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
disabled := insertOriginTaskChannel(t, common.ChannelStatusManuallyDisabled)
|
||||
insertOriginOwnedTask(t, "task-own", 7, enabled.Id, "origin-plugin")
|
||||
insertOriginOwnedTask(t, "task-own-b", 7, enabled.Id, "origin-plugin")
|
||||
insertOriginOwnedTask(t, "task-other-channel", 7, otherEnabled.Id, "origin-plugin")
|
||||
insertOriginOwnedTask(t, "task-foreign", 8, enabled.Id, "origin-plugin")
|
||||
insertOriginOwnedTask(t, "task-wrong-platform", 7, enabled.Id, "other-plugin")
|
||||
insertOriginOwnedTask(t, "task-legacy", 7, enabled.Id, constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeDoubaoVideo)))
|
||||
insertOriginOwnedTask(t, "task-disabled", 7, disabled.Id, "origin-plugin")
|
||||
|
||||
tooMany := make([]any, 17)
|
||||
for i := range tooMany {
|
||||
tooMany[i] = fmt.Sprintf("task-%d", i)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
userID int
|
||||
intent map[string]any
|
||||
channelType int
|
||||
wantCode string
|
||||
wantPinned int
|
||||
wantIDs []string
|
||||
}{
|
||||
{
|
||||
name: "valid single origin id pins channel",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-own"}},
|
||||
wantPinned: enabled.Id,
|
||||
wantIDs: []string{"task-own"},
|
||||
},
|
||||
{
|
||||
name: "dedupes preserving first order",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-own", " task-own ", "task-own-b"}},
|
||||
wantPinned: enabled.Id,
|
||||
wantIDs: []string{"task-own", "task-own-b"},
|
||||
},
|
||||
{
|
||||
name: "legacy platform matches channelType",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-legacy"}},
|
||||
channelType: constant.ChannelTypeDoubaoVideo,
|
||||
wantPinned: enabled.Id,
|
||||
wantIDs: []string{"task-legacy"},
|
||||
},
|
||||
{
|
||||
name: "unknown id",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-missing"}},
|
||||
wantCode: "origin_task_not_found",
|
||||
},
|
||||
{
|
||||
name: "other user's task is not found",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-foreign"}},
|
||||
wantCode: "origin_task_not_found",
|
||||
},
|
||||
{
|
||||
name: "platform mismatch",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-wrong-platform"}},
|
||||
wantCode: "origin_task_platform_mismatch",
|
||||
},
|
||||
{
|
||||
name: "two ids on different channels",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-own", "task-other-channel"}},
|
||||
wantCode: "origin_task_channel_conflict",
|
||||
},
|
||||
{
|
||||
name: "disabled channel",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-disabled"}},
|
||||
wantCode: "origin_task_channel_disabled",
|
||||
},
|
||||
{
|
||||
name: "more than 16 ids",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": tooMany},
|
||||
wantCode: "invalid_origin_task_ids",
|
||||
},
|
||||
{
|
||||
name: "non-array originTaskIds",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": "task-own"},
|
||||
wantCode: "invalid_origin_task_ids",
|
||||
},
|
||||
{
|
||||
name: "empty string entry",
|
||||
userID: 7,
|
||||
intent: map[string]any{"originTaskIds": []any{"task-own", " "}},
|
||||
wantCode: "invalid_origin_task_ids",
|
||||
},
|
||||
}
|
||||
|
||||
for _, testCase := range tests {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
c := originTaskTestContext(testCase.userID)
|
||||
intentErr := applyOriginTaskIntent(c, testCase.intent, jsplugin.Meta{Key: "origin-plugin", ChannelTypes: []int{testCase.channelType}})
|
||||
if testCase.wantCode != "" {
|
||||
require.NotNil(t, intentErr)
|
||||
assert.Equal(t, testCase.wantCode, intentErr.Code)
|
||||
assert.Equal(t, http.StatusBadRequest, intentErr.StatusCode)
|
||||
_, pinned := resolvedOriginPin(c)
|
||||
assert.False(t, pinned)
|
||||
return
|
||||
}
|
||||
require.Nil(t, intentErr)
|
||||
pinnedID, ok := resolvedOriginPin(c)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, testCase.wantPinned, pinnedID)
|
||||
tasks, ok := common.GetContextKeyType[[]*model.Task](c, constant.ContextKeyOriginTasks)
|
||||
require.True(t, ok)
|
||||
require.Len(t, tasks, len(testCase.wantIDs))
|
||||
for i, wantID := range testCase.wantIDs {
|
||||
assert.Equal(t, wantID, tasks[i].TaskID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestApplyOriginTaskIntentAbsentAndEmptyAreNoop(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
c := originTaskTestContext(7)
|
||||
require.Nil(t, applyOriginTaskIntent(c, map[string]any{}, jsplugin.Meta{Key: "origin-plugin"}))
|
||||
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{}}, jsplugin.Meta{Key: "origin-plugin"}))
|
||||
_, pinned := resolvedOriginPin(c)
|
||||
assert.False(t, pinned)
|
||||
}
|
||||
|
||||
func TestApplyOriginTaskAffinitySetsLockedChannel(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
insertOriginOwnedTask(t, "task-lock", 7, channel.Id, "origin-plugin")
|
||||
|
||||
c := originTaskTestContext(7)
|
||||
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{"task-lock"}}, jsplugin.Meta{Key: "origin-plugin"}))
|
||||
|
||||
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
taskErr := relay.ApplyOriginTaskAffinity(c, info)
|
||||
require.Nil(t, taskErr)
|
||||
locked, ok := info.LockedChannel.(*model.Channel)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, locked)
|
||||
assert.Equal(t, channel.Id, locked.Id)
|
||||
require.Len(t, info.OriginTasks, 1)
|
||||
assert.Equal(t, "task-lock", info.OriginTasks[0].TaskID)
|
||||
assert.Equal(t, "upstream-task-lock", info.OriginTasks[0].UpstreamTaskID)
|
||||
assert.Equal(t, "text_to_video", info.OriginTasks[0].Action)
|
||||
assert.Equal(t, string(model.TaskStatusSuccess), info.OriginTasks[0].Status)
|
||||
}
|
||||
|
||||
func TestPrepareTaskPluginRoutePinsOriginTaskChannel(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
insertOriginOwnedTask(t, "task-route", 7, channel.Id, "origin-route")
|
||||
plugin := compileTaskRoutePlugin(t, `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "origin-route", name: "Origin", version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
models: ["resolved-model"], fetchMode: "per_task",
|
||||
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
|
||||
};
|
||||
export const native = {
|
||||
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["task-route"], requestBody: {prompt: "ok"}}; },
|
||||
jobCreated: 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"}; }
|
||||
`)
|
||||
reached := false
|
||||
router := gin.New()
|
||||
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
||||
c.Next()
|
||||
}, PrepareTaskPluginRoute(), func(c *gin.Context) {
|
||||
reached = true
|
||||
pinnedID, ok := resolvedOriginPin(c)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, channel.Id, pinnedID)
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.True(t, reached)
|
||||
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
|
||||
func TestPrepareTaskPluginRouteRejectsUnknownOriginTask(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
plugin := compileTaskRoutePlugin(t, `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "origin-route-missing", name: "Origin", version: "1.0.0",
|
||||
author: {name: "Test"},
|
||||
models: ["resolved-model"], fetchMode: "per_task",
|
||||
routes: [{method: "POST", path: "/vendor/jobs", type: "submit", decode: "decodeJob", render: "jobCreated"}],
|
||||
};
|
||||
export const native = {
|
||||
decodeJob: function() { return {kind: "submit", model: "resolved-model", originTaskIds: ["missing"], requestBody: {}}; },
|
||||
jobCreated: 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"}; }
|
||||
`)
|
||||
reached := false
|
||||
router := gin.New()
|
||||
router.POST("/vendor/jobs", pinTaskPluginRoute(plugin, 0), func(c *gin.Context) {
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
||||
c.Next()
|
||||
}, PrepareTaskPluginRoute(), func(c *gin.Context) { reached = true })
|
||||
request := httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"resolved-model"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.False(t, reached)
|
||||
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
}
|
||||
|
||||
func TestPrepareTaskPluginEndpointPinsOriginTaskChannel(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
insertOriginOwnedTask(t, "task-endpoint", 7, channel.Id, "origin-endpoint")
|
||||
const key = "origin-endpoint"
|
||||
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
|
||||
key,
|
||||
"1.0.0",
|
||||
`["claimed-model"]`,
|
||||
"/v1/responses",
|
||||
`return {model: ctx.model, originTaskIds: ["task-endpoint"], requestBody: {prompt: "ok"}};`,
|
||||
), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
|
||||
|
||||
reached := false
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
||||
c.Next()
|
||||
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) {
|
||||
reached = true
|
||||
pinnedID, ok := resolvedOriginPin(c)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, channel.Id, pinnedID)
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.True(t, reached)
|
||||
assert.Equal(t, http.StatusNoContent, recorder.Code)
|
||||
}
|
||||
|
||||
func TestPrepareTaskPluginEndpointRejectsUnknownOriginTask(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
const key = "origin-endpoint-missing"
|
||||
_, err := jsplugin.DefaultRegistry.Register(taskProtocolPluginSource(
|
||||
key,
|
||||
"1.0.0",
|
||||
`["claimed-model"]`,
|
||||
"/v1/responses",
|
||||
`return {model: ctx.model, originTaskIds: ["missing"], requestBody: {prompt: "ok"}};`,
|
||||
), jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, jsplugin.DefaultRegistry.Unregister(key)) })
|
||||
|
||||
reached := false
|
||||
router := gin.New()
|
||||
router.POST("/v1/responses", func(c *gin.Context) {
|
||||
common.SetContextKey(c, constant.ContextKeyUserId, 7)
|
||||
c.Next()
|
||||
}, PinTaskPluginEndpoint(), PrepareTaskPluginEndpoint(), func(c *gin.Context) { reached = true })
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"claimed-model","input":"hello"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
assert.False(t, reached)
|
||||
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
assert.Contains(t, recorder.Body.String(), "origin_task_not_found")
|
||||
}
|
||||
|
||||
func TestDistributeHonorsOriginTaskChannelPin(t *testing.T) {
|
||||
require.NoError(t, appI18n.Init())
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set("resolved_task_model", "resolved-model")
|
||||
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
||||
ChannelId: channel.Id,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
|
||||
nextCalled := false
|
||||
handler := Distribute()
|
||||
handler(c)
|
||||
if !c.IsAborted() {
|
||||
nextCalled = true
|
||||
}
|
||||
assert.True(t, nextCalled)
|
||||
assert.Equal(t, channel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
|
||||
}
|
||||
|
||||
func TestDistributeTokenPinBeatsOriginPin(t *testing.T) {
|
||||
require.NoError(t, appI18n.Init())
|
||||
setupOriginTaskDB(t)
|
||||
tokenChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
originChannel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
|
||||
var warnBuf bytes.Buffer
|
||||
previousWriter := gin.DefaultErrorWriter
|
||||
gin.DefaultErrorWriter = &warnBuf
|
||||
t.Cleanup(func() { gin.DefaultErrorWriter = previousWriter })
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set("resolved_task_model", "resolved-model")
|
||||
constraints := service.GetChannelConstraints(c)
|
||||
constraints.AddPin(dto.ChannelPin{
|
||||
ChannelId: tokenChannel.Id,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
constraints.AddPin(dto.ChannelPin{
|
||||
ChannelId: originChannel.Id,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
nextCalled := false
|
||||
Distribute()(c)
|
||||
if !c.IsAborted() {
|
||||
nextCalled = true
|
||||
}
|
||||
assert.True(t, nextCalled)
|
||||
assert.Equal(t, tokenChannel.Id, common.GetContextKeyInt(c, constant.ContextKeyChannelId))
|
||||
warn := warnBuf.String()
|
||||
assert.Contains(t, warn, "winning_source=token")
|
||||
assert.Contains(t, warn, fmt.Sprintf("winning_channel_id=%d", tokenChannel.Id))
|
||||
assert.Contains(t, warn, "overridden_source=origin_task")
|
||||
assert.Contains(t, warn, fmt.Sprintf("overridden_channel_id=%d", originChannel.Id))
|
||||
}
|
||||
|
||||
func TestDistributePinViolatingIdentityFilterErrors(t *testing.T) {
|
||||
require.NoError(t, appI18n.Init())
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{}`))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Set("resolved_task_model", "resolved-model")
|
||||
c.Set("expected_task_plugin_key", "alpha")
|
||||
service.GetChannelConstraints(c).AddPin(dto.ChannelPin{
|
||||
ChannelId: channel.Id,
|
||||
Source: dto.PinSourceOriginTask,
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
|
||||
Distribute()(c)
|
||||
assert.True(t, c.IsAborted())
|
||||
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
assert.Contains(t, recorder.Body.String(), string(dto.FilterTaskPluginIdentity))
|
||||
}
|
||||
|
||||
func TestApplyChannelPinLocksOnlySameChannelRetry(t *testing.T) {
|
||||
setupOriginTaskDB(t)
|
||||
channel := insertOriginTaskChannel(t, common.ChannelStatusEnabled)
|
||||
insertOriginOwnedTask(t, "task-lock-mode", 7, channel.Id, "origin-plugin")
|
||||
|
||||
c := originTaskTestContext(7)
|
||||
require.Nil(t, applyOriginTaskIntent(c, map[string]any{"originTaskIds": []any{"task-lock-mode"}}, jsplugin.Meta{Key: "origin-plugin"}))
|
||||
|
||||
info := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
require.Nil(t, relay.ApplyChannelPin(c, info))
|
||||
locked, ok := info.LockedChannel.(*model.Channel)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, channel.Id, locked.Id)
|
||||
|
||||
tokenOnly := originTaskTestContext(7)
|
||||
service.GetChannelConstraints(tokenOnly).AddPin(dto.ChannelPin{
|
||||
ChannelId: channel.Id,
|
||||
Source: dto.PinSourceToken,
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
tokenInfo := &relaycommon.RelayInfo{TaskRelayInfo: &relaycommon.TaskRelayInfo{}}
|
||||
require.Nil(t, relay.ApplyChannelPin(tokenOnly, tokenInfo))
|
||||
assert.Nil(t, tokenInfo.LockedChannel)
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,48 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newTokenAutoGroupsContext() *gin.Context {
|
||||
gin.SetMode(gin.TestMode)
|
||||
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
return ctx
|
||||
}
|
||||
|
||||
func TestSetupContextForTokenPreservesCustomAutoGroupsOrder(t *testing.T) {
|
||||
ctx := newTokenAutoGroupsContext()
|
||||
token := &model.Token{Id: 1, UserId: 2, AutoGroups: `["vip","default"]`}
|
||||
|
||||
require.NoError(t, SetupContextForToken(ctx, token))
|
||||
value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, []string{"vip", "default"}, value)
|
||||
}
|
||||
|
||||
func TestSetupContextForTokenTreatsStoredEmptyArrayAsInheritance(t *testing.T) {
|
||||
ctx := newTokenAutoGroupsContext()
|
||||
token := &model.Token{Id: 1, UserId: 2, AutoGroups: `[]`}
|
||||
|
||||
require.NoError(t, SetupContextForToken(ctx, token))
|
||||
_, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups)
|
||||
assert.False(t, ok)
|
||||
}
|
||||
|
||||
func TestSetupContextForTokenMalformedAutoGroupsFailsClosed(t *testing.T) {
|
||||
ctx := newTokenAutoGroupsContext()
|
||||
token := &model.Token{Id: 1, UserId: 2, AutoGroups: `not-json`}
|
||||
|
||||
require.NoError(t, SetupContextForToken(ctx, token))
|
||||
value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups)
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, []string{}, value)
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func ConfigureTrustedProxies(engine *gin.Engine) error {
|
||||
trustedProxies, usedDefaults, err := common.ResolveTrustedProxies(os.Getenv("TRUSTED_PROXIES"))
|
||||
if err != nil {
|
||||
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.")
|
||||
}
|
||||
return common.ConfigureTrustedProxies(engine, trustedProxies)
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func requestClientIP(router http.Handler, remoteAddr string, forwardedFor string) string {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, "/client-ip", nil)
|
||||
request.RemoteAddr = remoteAddr
|
||||
if forwardedFor != "" {
|
||||
request.Header.Set("X-Forwarded-For", forwardedFor)
|
||||
}
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder.Body.String()
|
||||
}
|
||||
|
||||
func newClientIPRouter() *gin.Engine {
|
||||
router := gin.New()
|
||||
router.GET("/client-ip", func(c *gin.Context) {
|
||||
c.String(http.StatusOK, c.ClientIP())
|
||||
})
|
||||
return router
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesDefaultsToLoopbackAndPrivateNetworks(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Setenv("TRUSTED_PROXIES", "")
|
||||
router := newClientIPRouter()
|
||||
require.NoError(t, ConfigureTrustedProxies(router))
|
||||
|
||||
testCases := []struct {
|
||||
name string
|
||||
remoteAddr string
|
||||
}{
|
||||
{name: "IPv4 loopback", remoteAddr: "127.0.0.1:12345"},
|
||||
{name: "IPv6 loopback", remoteAddr: "[::1]:12345"},
|
||||
{name: "10 private network", remoteAddr: "10.20.30.40:12345"},
|
||||
{name: "172 private network", remoteAddr: "172.20.0.2:12345"},
|
||||
{name: "192 private network", remoteAddr: "192.168.10.2:12345"},
|
||||
{name: "IPv6 unique local network", remoteAddr: "[fd12:3456::2]:12345"},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
clientIP := requestClientIP(router, testCase.remoteAddr, "203.0.113.10")
|
||||
assert.Equal(t, "203.0.113.10", clientIP)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesDefaultRejectsPublicPeerHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Setenv("TRUSTED_PROXIES", " \t ")
|
||||
router := newClientIPRouter()
|
||||
require.NoError(t, ConfigureTrustedProxies(router))
|
||||
|
||||
clientIP := requestClientIP(router, "198.51.100.10:12345", "203.0.113.10")
|
||||
assert.Equal(t, "198.51.100.10", clientIP, "a public peer must not make a spoofed X-Forwarded-For authoritative")
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesDefaultStopsAtPublicClientInForwardedChain(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Setenv("TRUSTED_PROXIES", "")
|
||||
router := newClientIPRouter()
|
||||
require.NoError(t, ConfigureTrustedProxies(router))
|
||||
|
||||
clientIP := requestClientIP(router, "172.20.0.2:12345", "192.0.2.99, 203.0.113.10")
|
||||
assert.Equal(t, "203.0.113.10", clientIP, "the first public hop from the trusted proxy must win over a client-supplied prefix")
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesNoneDisablesForwardedHeaders(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Setenv("TRUSTED_PROXIES", " NoNe ")
|
||||
router := newClientIPRouter()
|
||||
require.NoError(t, ConfigureTrustedProxies(router))
|
||||
|
||||
clientIP := requestClientIP(router, "127.0.0.1:12345", "203.0.113.10")
|
||||
assert.Equal(t, "127.0.0.1", clientIP)
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesAcceptsTrimmedIPsAndCIDRs(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
t.Setenv("TRUSTED_PROXIES", " 192.0.2.0/24, 198.51.100.30 ")
|
||||
router := newClientIPRouter()
|
||||
require.NoError(t, ConfigureTrustedProxies(router))
|
||||
|
||||
trustedClientIP := requestClientIP(router, "192.0.2.10:12345", "203.0.113.20")
|
||||
assert.Equal(t, "203.0.113.20", trustedClientIP)
|
||||
|
||||
trustedExactIP := requestClientIP(router, "198.51.100.30:12345", "203.0.113.21")
|
||||
assert.Equal(t, "203.0.113.21", trustedExactIP)
|
||||
|
||||
untrustedClientIP := requestClientIP(router, "198.51.100.20:12345", "203.0.113.22")
|
||||
assert.Equal(t, "198.51.100.20", untrustedClientIP)
|
||||
|
||||
defaultProxyIP := requestClientIP(router, "127.0.0.1:12345", "203.0.113.23")
|
||||
assert.Equal(t, "127.0.0.1", defaultProxyIP, "an explicit list must replace, not extend, the compatibility defaults")
|
||||
}
|
||||
|
||||
func TestConfigureTrustedProxiesRejectsInvalidConfiguration(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
testCases := []struct {
|
||||
name string
|
||||
value string
|
||||
}{
|
||||
{name: "no entries", value: ", ,"},
|
||||
{name: "invalid entry", value: "not-an-ip"},
|
||||
{name: "mixed valid and invalid entries", value: "127.0.0.1, not-an-ip"},
|
||||
{name: "none mixed with valid entry", value: "none,127.0.0.1"},
|
||||
{name: "valid entry mixed with none", value: "127.0.0.1,NONE"},
|
||||
}
|
||||
|
||||
for _, testCase := range testCases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
t.Setenv("TRUSTED_PROXIES", testCase.value)
|
||||
router := newClientIPRouter()
|
||||
assert.Error(t, ConfigureTrustedProxies(router))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type turnstileCheckResponse struct {
|
||||
Success bool `json:"success"`
|
||||
}
|
||||
|
||||
func TurnstileCheck() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if common.TurnstileCheckEnabled {
|
||||
response := c.Query("turnstile")
|
||||
if response == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "Turnstile token 为空",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
rawRes, err := http.PostForm("https://challenges.cloudflare.com/turnstile/v0/siteverify", url.Values{
|
||||
"secret": {common.TurnstileSecretKey},
|
||||
"response": {response},
|
||||
"remoteip": {c.ClientIP()},
|
||||
})
|
||||
if err != nil {
|
||||
common.SysLog(err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
defer rawRes.Body.Close()
|
||||
var res turnstileCheckResponse
|
||||
err = common.DecodeJson(rawRes.Body, &res)
|
||||
if err != nil {
|
||||
common.SysLog(err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": err.Error(),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if !res.Success {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "Turnstile 校验失败,请刷新重试!",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func abortWithOpenAiMessage(c *gin.Context, statusCode int, message string, code ...types.ErrorCode) {
|
||||
codeStr := ""
|
||||
if len(code) > 0 {
|
||||
codeStr = string(code[0])
|
||||
}
|
||||
userId := c.GetInt("id")
|
||||
_, preparedPluginRoute := c.Get(pluginruntime.ContextKeyRouteRequest)
|
||||
if !preparedPluginRoute || !RespondTaskPluginError(c, &dto.TaskError{
|
||||
Code: codeStr,
|
||||
Message: message,
|
||||
StatusCode: statusCode,
|
||||
}) {
|
||||
c.JSON(statusCode, gin.H{
|
||||
"error": gin.H{
|
||||
"message": common.MessageWithRequestId(message, c.GetString(common.RequestIdKey)),
|
||||
"type": "new_api_error",
|
||||
"code": codeStr,
|
||||
},
|
||||
})
|
||||
}
|
||||
c.Abort()
|
||||
logger.LogError(c.Request.Context(), fmt.Sprintf("user %d | %s", userId, message))
|
||||
}
|
||||
|
||||
func abortWithMidjourneyMessage(c *gin.Context, statusCode int, code int, description string) {
|
||||
c.JSON(statusCode, gin.H{
|
||||
"description": description,
|
||||
"type": "new_api_error",
|
||||
"code": code,
|
||||
})
|
||||
c.Abort()
|
||||
logger.LogError(c.Request.Context(), description)
|
||||
}
|
||||
Reference in new issue
Block a user