refactor(auth): replace dashboard sessions with stateless tokens and session control (#6329)
* refactor(auth): replace dashboard sessions with stateless tokens * feat(auth): harden session issuance and distributed enforcement * fix(proxy): preserve trusted proxy compatibility defaults * refactor: address dashboard auth review feedback * refactor: remove classic frontend and flatten web app
This commit is contained in:
+2
-2
@@ -93,7 +93,7 @@ var auditRouteActions = map[string]string{
|
||||
"POST /api/subscription/admin/bind": "subscription.bind",
|
||||
|
||||
// 日志
|
||||
"DELETE /api/log/": "log.clear",
|
||||
"POST /api/system-task/log-cleanup": "log.cleanup_start",
|
||||
}
|
||||
|
||||
// beginAdminAudit 在管理/root 写操作进入 handler 前包装 ResponseWriter,
|
||||
@@ -155,7 +155,7 @@ func finishAdminAudit(c *gin.Context, writer *auditResponseWriter) {
|
||||
opParams["route"] = route
|
||||
}
|
||||
|
||||
// content 为英文兜底文本(导出/经典前端用)。
|
||||
// content 为英文兜底文本(供导出等非本地化消费者使用)。
|
||||
content := method + " " + route
|
||||
|
||||
adminInfo := map[string]interface{}{
|
||||
|
||||
+160
-122
@@ -5,7 +5,6 @@ import (
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
@@ -18,11 +17,20 @@ import (
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"github.com/QuantumNous/new-api/types"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"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) == "" {
|
||||
@@ -35,124 +43,24 @@ func validUserInfo(username string, role int) bool {
|
||||
}
|
||||
|
||||
func authHelper(c *gin.Context, minRole int) {
|
||||
session := sessions.Default(c)
|
||||
username := session.Get("username")
|
||||
role := session.Get("role")
|
||||
id := session.Get("id")
|
||||
status := session.Get("status")
|
||||
useAccessToken := false
|
||||
if username == nil {
|
||||
// Check access token
|
||||
accessToken := c.Request.Header.Get("Authorization")
|
||||
if accessToken == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthNotLoggedIn),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
user, authErr := model.ValidateAccessToken(accessToken)
|
||||
if authErr != nil {
|
||||
if errors.Is(authErr, model.ErrDatabase) {
|
||||
common.SysLog("ValidateAccessToken database error: " + authErr.Error())
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgDatabaseError),
|
||||
})
|
||||
} else {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid),
|
||||
})
|
||||
}
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if user != nil && user.Username != "" {
|
||||
if !validUserInfo(user.Username, user.Role) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
// Token is valid
|
||||
username = user.Username
|
||||
role = user.Role
|
||||
id = user.Id
|
||||
status = user.Status
|
||||
useAccessToken = true
|
||||
} else {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthAccessTokenInvalid),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
}
|
||||
// get header New-Api-User
|
||||
apiUserIdStr := c.Request.Header.Get("New-Api-User")
|
||||
if apiUserIdStr == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserIdNotProvided),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
apiUserId, err := strconv.Atoi(apiUserIdStr)
|
||||
user, identity, useAccessToken, err := authenticateDashboardRequest(c)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserIdFormatError),
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
|
||||
}
|
||||
if id != apiUserId {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserIdMismatch),
|
||||
})
|
||||
c.Abort()
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if status.(int) == common.UserStatusDisabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserBanned),
|
||||
})
|
||||
c.Abort()
|
||||
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 role.(int) < minRole {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthInsufficientPrivilege),
|
||||
})
|
||||
c.Abort()
|
||||
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(username.(string), role.(int)) {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": common.TranslateMessage(c, i18n.MsgAuthUserInfoInvalid),
|
||||
})
|
||||
c.Abort()
|
||||
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
|
||||
}
|
||||
// 防止不同newapi版本冲突,导致数据不通用
|
||||
c.Header("Auth-Version", "864b7076dbcd0a3c01b5520316720ebf")
|
||||
c.Set("username", username)
|
||||
c.Set("role", role)
|
||||
c.Set("id", id)
|
||||
c.Set("group", session.Get("group"))
|
||||
c.Set("user_group", session.Get("group"))
|
||||
c.Set("use_access_token", useAccessToken)
|
||||
setDashboardAuthContext(c, user, identity, useAccessToken)
|
||||
|
||||
// 管理/root 写操作审计兜底:内聚在鉴权链路里,保证任何经过 AdminAuth/RootAuth
|
||||
// 的写接口都会自动留痕(无需在路由上单独挂审计中间件,避免漏挂)。
|
||||
@@ -169,10 +77,13 @@ func authHelper(c *gin.Context, minRole int) {
|
||||
|
||||
func TryUserAuth() func(c *gin.Context) {
|
||||
return func(c *gin.Context) {
|
||||
session := sessions.Default(c)
|
||||
id := session.Get("id")
|
||||
if id != nil {
|
||||
c.Set("id", id)
|
||||
user, identity, credentialKind, err := classifyDashboardCredential(c)
|
||||
if err != nil {
|
||||
writeDashboardAuthError(c, err)
|
||||
return
|
||||
}
|
||||
if credentialKind != dashboardCredentialUnmatched {
|
||||
setDashboardAuthContext(c, user, identity, credentialKind == dashboardCredentialPAT)
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
@@ -196,6 +107,122 @@ func RootAuth() func(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
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")
|
||||
@@ -220,16 +247,27 @@ func WssAuth(c *gin.Context) {
|
||||
// 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) {
|
||||
// Try session auth first (dashboard users)
|
||||
session := sessions.Default(c)
|
||||
if id := session.Get("id"); id != nil {
|
||||
if status, ok := session.Get("status").(int); ok && status == common.UserStatusEnabled {
|
||||
c.Set("id", id)
|
||||
c.Next()
|
||||
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
|
||||
}
|
||||
// Fall back to token auth (API clients)
|
||||
// Opaque credentials are relay API keys here, never dashboard PATs.
|
||||
TokenAuth()(c)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,251 @@
|
||||
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
|
||||
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.DB = db
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.RedisEnabled = false
|
||||
common.SessionSecret = "middleware-auth-test-secret"
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
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)
|
||||
securityProof, _, err := service.IssueSecurityProof(identity, "2fa", []string{"channel.key.read"})
|
||||
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")
|
||||
}
|
||||
@@ -1,10 +1,8 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
|
||||
@@ -18,33 +16,24 @@ const (
|
||||
)
|
||||
|
||||
func redisEmailVerificationRateLimiter(c *gin.Context) {
|
||||
ctx := context.Background()
|
||||
rdb := common.RDB
|
||||
key := "emailVerification:" + EmailVerificationRateLimitMark + ":" + c.ClientIP()
|
||||
|
||||
count, err := rdb.Incr(ctx, key).Result()
|
||||
allowed, _, ttlSeconds, err := redisFixedWindowTake(
|
||||
c.Request.Context(),
|
||||
redisIPRateLimitKey(EmailVerificationRateLimitMark, c.ClientIP()),
|
||||
EmailVerificationMaxRequests,
|
||||
EmailVerificationDuration,
|
||||
)
|
||||
if err != nil {
|
||||
// fallback
|
||||
memoryEmailVerificationRateLimiter(c)
|
||||
return
|
||||
}
|
||||
|
||||
// 第一次设置键时设置过期时间
|
||||
if count == 1 {
|
||||
_ = rdb.Expire(ctx, key, time.Duration(EmailVerificationDuration)*time.Second).Err()
|
||||
}
|
||||
|
||||
// 检查是否超出限制
|
||||
if count <= int64(EmailVerificationMaxRequests) {
|
||||
if allowed {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
// 获取剩余等待时间
|
||||
ttl, err := rdb.TTL(ctx, key).Result()
|
||||
waitSeconds := int64(EmailVerificationDuration)
|
||||
if err == nil && ttl > 0 {
|
||||
waitSeconds = int64(ttl.Seconds())
|
||||
if ttlSeconds > 0 {
|
||||
waitSeconds = ttlSeconds
|
||||
}
|
||||
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{
|
||||
@@ -70,11 +59,13 @@ func memoryEmailVerificationRateLimiter(c *gin.Context) {
|
||||
}
|
||||
|
||||
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 {
|
||||
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
|
||||
memoryEmailVerificationRateLimiter(c)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,10 +6,12 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-contrib/sessions/cookie"
|
||||
"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) {
|
||||
@@ -39,40 +41,39 @@ func performHeaderNavRequest(t *testing.T, handler gin.HandlerFunc, authenticate
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(sessions.Sessions("session", cookie.NewStore([]byte("header-nav-test"))))
|
||||
router.GET("/login", func(c *gin.Context) {
|
||||
session := sessions.Default(c)
|
||||
session.Set("username", "tester")
|
||||
session.Set("role", common.RoleCommonUser)
|
||||
session.Set("id", 1)
|
||||
session.Set("status", common.UserStatusEnabled)
|
||||
session.Set("group", "default")
|
||||
if err := session.Save(); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"success": false})
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
router.GET("/api/test", handler, func(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
})
|
||||
|
||||
var cookies []*http.Cookie
|
||||
var accessToken string
|
||||
if authenticated {
|
||||
loginRecorder := httptest.NewRecorder()
|
||||
loginRequest := httptest.NewRequest(http.MethodGet, "/login", nil)
|
||||
router.ServeHTTP(loginRecorder, loginRequest)
|
||||
require.Equal(t, http.StatusNoContent, loginRecorder.Code)
|
||||
cookies = loginRecorder.Result().Cookies()
|
||||
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("New-Api-User", "1")
|
||||
for _, cookie := range cookies {
|
||||
request.AddCookie(cookie)
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
}
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
@@ -165,3 +166,24 @@ func TestHeaderNavModulePublicOrUserAuthRequiresLoginForLegacyDisabledModule(t *
|
||||
|
||||
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")
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
const (
|
||||
ModelRequestRateLimitCountMark = "MRRL"
|
||||
ModelRequestRateLimitSuccessCountMark = "MRRLS"
|
||||
modelRateLimitTimeFormat = "2006-01-02T15:04:05.000Z"
|
||||
)
|
||||
|
||||
// 检查Redis中的请求限制
|
||||
@@ -41,13 +42,13 @@ func checkRedisRateLimit(ctx context.Context, rdb *redis.Client, key string, max
|
||||
|
||||
// 检查时间窗口
|
||||
oldTimeStr, _ := rdb.LIndex(ctx, key, -1).Result()
|
||||
oldTime, err := time.Parse(timeFormat, oldTimeStr)
|
||||
oldTime, err := time.Parse(modelRateLimitTimeFormat, oldTimeStr)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
|
||||
nowTimeStr := time.Now().Format(timeFormat)
|
||||
nowTime, err := time.Parse(timeFormat, nowTimeStr)
|
||||
nowTimeStr := time.Now().UTC().Format(modelRateLimitTimeFormat)
|
||||
nowTime, err := time.Parse(modelRateLimitTimeFormat, nowTimeStr)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
@@ -68,7 +69,7 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now().Format(timeFormat)
|
||||
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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
+113
-84
@@ -2,15 +2,37 @@ package middleware
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
"strconv"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var timeFormat = "2006-01-02T15:04:05.000Z"
|
||||
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
|
||||
|
||||
@@ -18,49 +40,87 @@ 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 interface{}) (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) {
|
||||
ctx := context.Background()
|
||||
rdb := common.RDB
|
||||
key := "rateLimit:" + mark + c.ClientIP()
|
||||
listLength, err := rdb.LLen(ctx, key).Result()
|
||||
allowed, _, _, err := redisFixedWindowTake(
|
||||
c.Request.Context(),
|
||||
redisIPRateLimitKey(mark, c.ClientIP()),
|
||||
maxRequestNum,
|
||||
duration,
|
||||
)
|
||||
if err != nil {
|
||||
fmt.Println(err.Error())
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if listLength < int64(maxRequestNum) {
|
||||
rdb.LPush(ctx, key, time.Now().Format(timeFormat))
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
} else {
|
||||
oldTimeStr, _ := rdb.LIndex(ctx, key, -1).Result()
|
||||
oldTime, err := time.Parse(timeFormat, oldTimeStr)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
nowTimeStr := time.Now().Format(timeFormat)
|
||||
nowTime, err := time.Parse(timeFormat, nowTimeStr)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
// time.Since will return negative number!
|
||||
// See: https://stackoverflow.com/questions/50970900/why-is-time-since-returning-negative-durations-on-windows
|
||||
if int64(nowTime.Sub(oldTime).Seconds()) < duration {
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
return
|
||||
} else {
|
||||
rdb.LPush(ctx, key, time.Now().Format(timeFormat))
|
||||
rdb.LTrim(ctx, key, 0, int64(maxRequestNum-1))
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
}
|
||||
if !allowed {
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,12 +138,11 @@ func rateLimitFactory(maxRequestNum int, duration int64, mark string) func(c *gi
|
||||
return func(c *gin.Context) {
|
||||
redisRateLimiter(c, maxRequestNum, duration, mark)
|
||||
}
|
||||
} else {
|
||||
// It's safe to call multi times.
|
||||
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration)
|
||||
return func(c *gin.Context) {
|
||||
memoryRateLimiter(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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -122,26 +181,25 @@ func UploadRateLimit() func(c *gin.Context) {
|
||||
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 {
|
||||
userID := c.GetInt("id")
|
||||
if userID == 0 {
|
||||
c.Status(http.StatusUnauthorized)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
key := fmt.Sprintf("rateLimit:%s:user:%d", mark, userId)
|
||||
userRedisRateLimiter(c, maxRequestNum, duration, key)
|
||||
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 {
|
||||
userID := c.GetInt("id")
|
||||
if userID == 0 {
|
||||
c.Status(http.StatusUnauthorized)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
key := fmt.Sprintf("%s:user:%d", mark, userId)
|
||||
key := fmt.Sprintf("%s:user:%d", mark, userID)
|
||||
if !inMemoryRateLimiter.Request(key, maxRequestNum, duration) {
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
@@ -153,45 +211,16 @@ func userRateLimitFactory(maxRequestNum int, duration int64, mark string) func(c
|
||||
// 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) {
|
||||
ctx := context.Background()
|
||||
rdb := common.RDB
|
||||
listLength, err := rdb.LLen(ctx, key).Result()
|
||||
allowed, _, _, err := redisFixedWindowTake(c.Request.Context(), key, maxRequestNum, duration)
|
||||
if err != nil {
|
||||
fmt.Println(err.Error())
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if listLength < int64(maxRequestNum) {
|
||||
rdb.LPush(ctx, key, time.Now().Format(timeFormat))
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
} else {
|
||||
oldTimeStr, _ := rdb.LIndex(ctx, key, -1).Result()
|
||||
oldTime, err := time.Parse(timeFormat, oldTimeStr)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
nowTimeStr := time.Now().Format(timeFormat)
|
||||
nowTime, err := time.Parse(timeFormat, nowTimeStr)
|
||||
if err != nil {
|
||||
fmt.Println(err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
if int64(nowTime.Sub(oldTime).Seconds()) < duration {
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
return
|
||||
} else {
|
||||
rdb.LPush(ctx, key, time.Now().Format(timeFormat))
|
||||
rdb.LTrim(ctx, key, 0, int64(maxRequestNum-1))
|
||||
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration)
|
||||
}
|
||||
if !allowed {
|
||||
c.Status(http.StatusTooManyRequests)
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
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)
|
||||
assert.Equal(t, http.StatusTooManyRequests, performRateLimitRequest(router, "/limited", remoteAddr).Code)
|
||||
|
||||
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)
|
||||
}
|
||||
@@ -1,133 +1,60 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"time"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
const (
|
||||
// SecureVerificationSessionKey 安全验证的 session key(与 controller 保持一致)
|
||||
SecureVerificationSessionKey = "secure_verified_at"
|
||||
secureVerificationMethodSessionKey = "secure_verified_method"
|
||||
// SecureVerificationTimeout 验证有效期(秒)
|
||||
SecureVerificationTimeout = 300 // 5分钟
|
||||
)
|
||||
|
||||
// SecureVerificationRequired 安全验证中间件
|
||||
// 检查用户是否在有效时间内通过了安全验证
|
||||
// 如果未验证或验证已过期,返回 401 错误
|
||||
// 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) {
|
||||
// 检查用户是否已登录
|
||||
userId := c.GetInt("id")
|
||||
if userId == 0 {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"success": false,
|
||||
"message": "未登录",
|
||||
})
|
||||
c.Abort()
|
||||
if !RequireSecurityProof(c, "channel.key.read", []string{"2fa", "passkey"}) {
|
||||
return
|
||||
}
|
||||
|
||||
// 检查 session 中的验证时间戳
|
||||
session := sessions.Default(c)
|
||||
verifiedAtRaw := session.Get(SecureVerificationSessionKey)
|
||||
|
||||
if verifiedAtRaw == nil {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": "需要安全验证",
|
||||
"code": "VERIFICATION_REQUIRED",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
verifiedAt, ok := verifiedAtRaw.(int64)
|
||||
if !ok {
|
||||
// session 数据格式错误
|
||||
clearSecureVerificationSession(session)
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": "验证状态异常,请重新验证",
|
||||
"code": "VERIFICATION_INVALID",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
// 检查验证是否过期
|
||||
elapsed := time.Now().Unix() - verifiedAt
|
||||
if elapsed >= SecureVerificationTimeout {
|
||||
// 验证已过期,清除 session
|
||||
clearSecureVerificationSession(session)
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": "验证已过期,请重新验证",
|
||||
"code": "VERIFICATION_EXPIRED",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func clearSecureVerificationSession(session sessions.Session) {
|
||||
session.Delete(SecureVerificationSessionKey)
|
||||
session.Delete(secureVerificationMethodSessionKey)
|
||||
_ = session.Save()
|
||||
}
|
||||
|
||||
// OptionalSecureVerification 可选的安全验证中间件
|
||||
// 如果用户已验证,则在 context 中设置标记,但不阻止请求继续
|
||||
// 用于某些需要区分是否已验证的场景
|
||||
func OptionalSecureVerification() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
userId := c.GetInt("id")
|
||||
if userId == 0 {
|
||||
c.Set("secure_verified", false)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
session := sessions.Default(c)
|
||||
verifiedAtRaw := session.Get(SecureVerificationSessionKey)
|
||||
|
||||
if verifiedAtRaw == nil {
|
||||
c.Set("secure_verified", false)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
verifiedAt, ok := verifiedAtRaw.(int64)
|
||||
if !ok {
|
||||
c.Set("secure_verified", false)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
elapsed := time.Now().Unix() - verifiedAt
|
||||
if elapsed >= SecureVerificationTimeout {
|
||||
clearSecureVerificationSession(session)
|
||||
c.Set("secure_verified", false)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
c.Set("secure_verified", true)
|
||||
c.Set("secure_verified_at", verifiedAt)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// ClearSecureVerification 清除安全验证状态
|
||||
// 用于用户登出或需要强制重新验证的场景
|
||||
func ClearSecureVerification(c *gin.Context) {
|
||||
session := sessions.Default(c)
|
||||
clearSecureVerificationSession(session)
|
||||
// RequireSecurityProof validates a proof against the authenticated dashboard
|
||||
// session and writes the shared proof error contract on failure.
|
||||
func RequireSecurityProof(c *gin.Context, requiredScope string, allowedMethods []string) bool {
|
||||
identity, ok := GetSessionAuthIdentity(c)
|
||||
if !ok {
|
||||
securityProofError(c, "SECURITY_PROOF_INVALID", "安全验证状态无效")
|
||||
return false
|
||||
}
|
||||
raw := strings.TrimSpace(c.GetHeader("X-Security-Proof"))
|
||||
if raw == "" {
|
||||
securityProofError(c, "SECURITY_PROOF_REQUIRED", "需要安全验证")
|
||||
return false
|
||||
}
|
||||
if _, err := service.VerifySecurityProof(raw, identity, requiredScope, allowedMethods); 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.ErrProofMethod):
|
||||
securityProofError(c, "SECURITY_PROOF_METHOD_MISMATCH", "安全验证方式不匹配")
|
||||
default:
|
||||
securityProofError(c, "SECURITY_PROOF_INVALID", "安全验证状态无效")
|
||||
}
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func securityProofError(c *gin.Context, code, message string) {
|
||||
c.JSON(http.StatusForbidden, gin.H{
|
||||
"success": false,
|
||||
"message": message,
|
||||
"code": code,
|
||||
})
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
@@ -1,12 +1,10 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/gin-contrib/sessions"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
@@ -17,12 +15,6 @@ type turnstileCheckResponse struct {
|
||||
func TurnstileCheck() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
if common.TurnstileCheckEnabled {
|
||||
session := sessions.Default(c)
|
||||
turnstileChecked := session.Get("turnstile")
|
||||
if turnstileChecked != nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
response := c.Query("turnstile")
|
||||
if response == "" {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
@@ -48,7 +40,7 @@ func TurnstileCheck() gin.HandlerFunc {
|
||||
}
|
||||
defer rawRes.Body.Close()
|
||||
var res turnstileCheckResponse
|
||||
err = json.NewDecoder(rawRes.Body).Decode(&res)
|
||||
err = common.DecodeJson(rawRes.Body, &res)
|
||||
if err != nil {
|
||||
common.SysLog(err.Error())
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
@@ -66,15 +58,6 @@ func TurnstileCheck() gin.HandlerFunc {
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
session.Set("turnstile", true)
|
||||
err = session.Save()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"message": "无法保存会话信息,请重试",
|
||||
"success": false,
|
||||
})
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Next()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user