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:
Calcium-Ion
2026-07-20 16:48:43 +08:00
committed by GitHub
parent 5a6c53d496
commit 31d70fca39
1605 changed files with 17511 additions and 147913 deletions
+2 -2
View File
@@ -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
View File
@@ -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)
}
}
+76
View File
@@ -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
}
+135
View File
@@ -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)
}
+251
View File
@@ -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")
}
+12 -21
View File
@@ -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)
}
}
+48 -26
View File
@@ -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")
}
+5 -4
View File
@@ -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)
+34
View File
@@ -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
View File
@@ -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()
}
}
+223
View File
@@ -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)
}
+42 -115
View File
@@ -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 -18
View File
@@ -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()
}