Fix Services & Update the middleware && others
This commit is contained in:
@@ -1,23 +1,151 @@
|
||||
// Filename: internal/middleware/web.go
|
||||
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"gemini-balancer/internal/service"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
const (
|
||||
AdminSessionCookie = "gemini_admin_session"
|
||||
SessionMaxAge = 3600 * 24 * 7 // 7天
|
||||
CacheTTL = 5 * time.Minute
|
||||
CleanupInterval = 10 * time.Minute // 降低清理频率
|
||||
SessionRefreshTime = 30 * time.Minute
|
||||
)
|
||||
|
||||
// ==================== 缓存层 ====================
|
||||
|
||||
type authCacheEntry struct {
|
||||
Token interface{}
|
||||
ExpiresAt time.Time
|
||||
}
|
||||
|
||||
type authCache struct {
|
||||
mu sync.RWMutex
|
||||
cache map[string]*authCacheEntry
|
||||
ttl time.Duration
|
||||
}
|
||||
|
||||
var webAuthCache = newAuthCache(CacheTTL)
|
||||
|
||||
func newAuthCache(ttl time.Duration) *authCache {
|
||||
c := &authCache{
|
||||
cache: make(map[string]*authCacheEntry),
|
||||
ttl: ttl,
|
||||
}
|
||||
go c.cleanupLoop()
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *authCache) get(key string) (interface{}, bool) {
|
||||
c.mu.RLock()
|
||||
defer c.mu.RUnlock()
|
||||
|
||||
entry, exists := c.cache[key]
|
||||
if !exists || time.Now().After(entry.ExpiresAt) {
|
||||
return nil, false
|
||||
}
|
||||
return entry.Token, true
|
||||
}
|
||||
|
||||
func (c *authCache) set(key string, token interface{}) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
c.cache[key] = &authCacheEntry{
|
||||
Token: token,
|
||||
ExpiresAt: time.Now().Add(c.ttl),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *authCache) delete(key string) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
delete(c.cache, key)
|
||||
}
|
||||
|
||||
func (c *authCache) cleanupLoop() {
|
||||
ticker := time.NewTicker(CleanupInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
c.cleanup()
|
||||
}
|
||||
}
|
||||
|
||||
func (c *authCache) cleanup() {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
count := 0
|
||||
for key, entry := range c.cache {
|
||||
if now.After(entry.ExpiresAt) {
|
||||
delete(c.cache, key)
|
||||
count++
|
||||
}
|
||||
}
|
||||
|
||||
if count > 0 {
|
||||
logrus.Debugf("[AuthCache] Cleaned up %d expired entries", count)
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 会话刷新缓存 ====================
|
||||
|
||||
var sessionRefreshCache = struct {
|
||||
sync.RWMutex
|
||||
timestamps map[string]time.Time
|
||||
}{
|
||||
timestamps: make(map[string]time.Time),
|
||||
}
|
||||
|
||||
// 定期清理刷新时间戳
|
||||
func init() {
|
||||
go func() {
|
||||
ticker := time.NewTicker(1 * time.Hour)
|
||||
defer ticker.Stop()
|
||||
|
||||
for range ticker.C {
|
||||
sessionRefreshCache.Lock()
|
||||
now := time.Now()
|
||||
for key, ts := range sessionRefreshCache.timestamps {
|
||||
if now.Sub(ts) > 2*time.Hour {
|
||||
delete(sessionRefreshCache.timestamps, key)
|
||||
}
|
||||
}
|
||||
sessionRefreshCache.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// ==================== Cookie 操作 ====================
|
||||
|
||||
func SetAdminSessionCookie(c *gin.Context, adminToken string) {
|
||||
c.SetCookie(AdminSessionCookie, adminToken, 3600*24*7, "/", "", false, true)
|
||||
secure := c.Request.TLS != nil || c.GetHeader("X-Forwarded-Proto") == "https"
|
||||
c.SetSameSite(http.SameSiteStrictMode)
|
||||
c.SetCookie(AdminSessionCookie, adminToken, SessionMaxAge, "/", "", secure, true)
|
||||
}
|
||||
|
||||
func SetAdminSessionCookieWithAge(c *gin.Context, adminToken string, maxAge int) {
|
||||
secure := c.Request.TLS != nil || c.GetHeader("X-Forwarded-Proto") == "https"
|
||||
c.SetSameSite(http.SameSiteStrictMode)
|
||||
c.SetCookie(AdminSessionCookie, adminToken, maxAge, "/", "", secure, true)
|
||||
}
|
||||
|
||||
func ClearAdminSessionCookie(c *gin.Context) {
|
||||
c.SetSameSite(http.SameSiteStrictMode)
|
||||
c.SetCookie(AdminSessionCookie, "", -1, "/", "", false, true)
|
||||
}
|
||||
|
||||
@@ -29,26 +157,258 @@ func ExtractTokenFromCookie(c *gin.Context) string {
|
||||
return cookie
|
||||
}
|
||||
|
||||
// ==================== 认证中间件 ====================
|
||||
|
||||
func WebAdminAuthMiddleware(authService *service.SecurityService) gin.HandlerFunc {
|
||||
logger := logrus.New()
|
||||
logger.SetLevel(getLogLevel())
|
||||
|
||||
return func(c *gin.Context) {
|
||||
cookie := ExtractTokenFromCookie(c)
|
||||
log.Printf("[WebAuth_Guard] Intercepting request for: %s", c.Request.URL.Path)
|
||||
log.Printf("[WebAuth_Guard] Found session cookie value: '%s'", cookie)
|
||||
authToken, err := authService.AuthenticateToken(cookie)
|
||||
if err != nil {
|
||||
log.Printf("[WebAuth_Guard] FATAL: AuthenticateToken FAILED. Error: %v. Redirecting to /login.", err)
|
||||
} else if !authToken.IsAdmin {
|
||||
log.Printf("[WebAuth_Guard] FATAL: Token validated, but IsAdmin is FALSE. Redirecting to /login.")
|
||||
} else {
|
||||
log.Printf("[WebAuth_Guard] SUCCESS: Token validated and IsAdmin is TRUE. Allowing access.")
|
||||
}
|
||||
if err != nil || !authToken.IsAdmin {
|
||||
if cookie == "" {
|
||||
logger.Debug("[WebAuth] No session cookie found")
|
||||
ClearAdminSessionCookie(c)
|
||||
c.Redirect(http.StatusFound, "/login")
|
||||
c.Abort()
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
cacheKey := hashToken(cookie)
|
||||
|
||||
if cachedToken, found := webAuthCache.get(cacheKey); found {
|
||||
logger.Debug("[WebAuth] Using cached token")
|
||||
c.Set("adminUser", cachedToken)
|
||||
refreshSessionIfNeeded(c, cookie)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
logger.Debug("[WebAuth] Cache miss, authenticating...")
|
||||
authToken, err := authService.AuthenticateToken(cookie)
|
||||
|
||||
if err != nil {
|
||||
logger.WithError(err).Warn("[WebAuth] Authentication failed")
|
||||
ClearAdminSessionCookie(c)
|
||||
webAuthCache.delete(cacheKey)
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
if !authToken.IsAdmin {
|
||||
logger.Warn("[WebAuth] User is not admin")
|
||||
ClearAdminSessionCookie(c)
|
||||
webAuthCache.delete(cacheKey)
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
webAuthCache.set(cacheKey, authToken)
|
||||
logger.Debug("[WebAuth] Authentication success, token cached")
|
||||
|
||||
c.Set("adminUser", authToken)
|
||||
refreshSessionIfNeeded(c, cookie)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func WebAdminAuthMiddlewareWithLogger(authService *service.SecurityService, logger *logrus.Logger) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
cookie := ExtractTokenFromCookie(c)
|
||||
|
||||
if cookie == "" {
|
||||
logger.Debug("No session cookie found")
|
||||
ClearAdminSessionCookie(c)
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
cacheKey := hashToken(cookie)
|
||||
if cachedToken, found := webAuthCache.get(cacheKey); found {
|
||||
c.Set("adminUser", cachedToken)
|
||||
refreshSessionIfNeeded(c, cookie)
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
authToken, err := authService.AuthenticateToken(cookie)
|
||||
|
||||
if err != nil {
|
||||
logger.WithError(err).Warn("Token authentication failed")
|
||||
ClearAdminSessionCookie(c)
|
||||
webAuthCache.delete(cacheKey)
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
if !authToken.IsAdmin {
|
||||
logger.Warn("Token valid but user is not admin")
|
||||
ClearAdminSessionCookie(c)
|
||||
webAuthCache.delete(cacheKey)
|
||||
redirectToLogin(c)
|
||||
return
|
||||
}
|
||||
|
||||
webAuthCache.set(cacheKey, authToken)
|
||||
c.Set("adminUser", authToken)
|
||||
refreshSessionIfNeeded(c, cookie)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 辅助函数 ====================
|
||||
|
||||
func hashToken(token string) string {
|
||||
h := sha256.Sum256([]byte(token))
|
||||
return hex.EncodeToString(h[:])
|
||||
}
|
||||
|
||||
func redirectToLogin(c *gin.Context) {
|
||||
if isAjaxRequest(c) {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{
|
||||
"error": "Session expired",
|
||||
"code": "AUTH_REQUIRED",
|
||||
})
|
||||
c.Abort()
|
||||
return
|
||||
}
|
||||
|
||||
originalPath := c.Request.URL.Path
|
||||
if originalPath != "/" && originalPath != "/login" {
|
||||
c.Redirect(http.StatusFound, "/login?redirect="+originalPath)
|
||||
} else {
|
||||
c.Redirect(http.StatusFound, "/login")
|
||||
}
|
||||
c.Abort()
|
||||
}
|
||||
|
||||
func isAjaxRequest(c *gin.Context) bool {
|
||||
// 检查 Content-Type
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
if strings.Contains(contentType, "application/json") {
|
||||
return true
|
||||
}
|
||||
|
||||
// 检查 Accept（优先检查 JSON）
|
||||
accept := c.GetHeader("Accept")
|
||||
if strings.Contains(accept, "application/json") &&
|
||||
!strings.Contains(accept, "text/html") {
|
||||
return true
|
||||
}
|
||||
|
||||
// 兼容旧版 XMLHttpRequest
|
||||
return c.GetHeader("X-Requested-With") == "XMLHttpRequest"
|
||||
}
|
||||
|
||||
func refreshSessionIfNeeded(c *gin.Context, token string) {
|
||||
tokenHash := hashToken(token)
|
||||
|
||||
sessionRefreshCache.RLock()
|
||||
lastRefresh, exists := sessionRefreshCache.timestamps[tokenHash]
|
||||
sessionRefreshCache.RUnlock()
|
||||
|
||||
if !exists || time.Since(lastRefresh) > SessionRefreshTime {
|
||||
SetAdminSessionCookie(c, token)
|
||||
|
||||
sessionRefreshCache.Lock()
|
||||
sessionRefreshCache.timestamps[tokenHash] = time.Now()
|
||||
sessionRefreshCache.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func getLogLevel() logrus.Level {
|
||||
level := os.Getenv("LOG_LEVEL")
|
||||
switch strings.ToLower(level) {
|
||||
case "debug":
|
||||
return logrus.DebugLevel
|
||||
case "warn":
|
||||
return logrus.WarnLevel
|
||||
case "error":
|
||||
return logrus.ErrorLevel
|
||||
default:
|
||||
return logrus.InfoLevel
|
||||
}
|
||||
}
|
||||
|
||||
// ==================== 工具函数 ====================
|
||||
|
||||
func GetAdminUserFromContext(c *gin.Context) (interface{}, bool) {
|
||||
return c.Get("adminUser")
|
||||
}
|
||||
|
||||
func InvalidateTokenCache(token string) {
|
||||
tokenHash := hashToken(token)
|
||||
webAuthCache.delete(tokenHash)
|
||||
|
||||
// 同时清理刷新时间戳
|
||||
sessionRefreshCache.Lock()
|
||||
delete(sessionRefreshCache.timestamps, tokenHash)
|
||||
sessionRefreshCache.Unlock()
|
||||
}
|
||||
|
||||
func ClearAllAuthCache() {
|
||||
webAuthCache.mu.Lock()
|
||||
webAuthCache.cache = make(map[string]*authCacheEntry)
|
||||
webAuthCache.mu.Unlock()
|
||||
|
||||
sessionRefreshCache.Lock()
|
||||
sessionRefreshCache.timestamps = make(map[string]time.Time)
|
||||
sessionRefreshCache.Unlock()
|
||||
}
|
||||
|
||||
// ==================== 调试工具 ====================
|
||||
|
||||
type SessionInfo struct {
|
||||
HasCookie bool `json:"has_cookie"`
|
||||
IsValid bool `json:"is_valid"`
|
||||
IsAdmin bool `json:"is_admin"`
|
||||
IsCached bool `json:"is_cached"`
|
||||
LastActivity string `json:"last_activity"`
|
||||
}
|
||||
|
||||
func GetSessionInfo(c *gin.Context, authService *service.SecurityService) SessionInfo {
|
||||
info := SessionInfo{
|
||||
HasCookie: false,
|
||||
IsValid: false,
|
||||
IsAdmin: false,
|
||||
IsCached: false,
|
||||
LastActivity: time.Now().Format(time.RFC3339),
|
||||
}
|
||||
|
||||
cookie := ExtractTokenFromCookie(c)
|
||||
if cookie == "" {
|
||||
return info
|
||||
}
|
||||
|
||||
info.HasCookie = true
|
||||
|
||||
cacheKey := hashToken(cookie)
|
||||
if _, found := webAuthCache.get(cacheKey); found {
|
||||
info.IsCached = true
|
||||
}
|
||||
|
||||
authToken, err := authService.AuthenticateToken(cookie)
|
||||
if err != nil {
|
||||
return info
|
||||
}
|
||||
|
||||
info.IsValid = true
|
||||
info.IsAdmin = authToken.IsAdmin
|
||||
|
||||
return info
|
||||
}
|
||||
|
||||
func GetCacheStats() map[string]interface{} {
|
||||
webAuthCache.mu.RLock()
|
||||
cacheSize := len(webAuthCache.cache)
|
||||
webAuthCache.mu.RUnlock()
|
||||
|
||||
sessionRefreshCache.RLock()
|
||||
refreshSize := len(sessionRefreshCache.timestamps)
|
||||
sessionRefreshCache.RUnlock()
|
||||
return map[string]interface{}{
|
||||
"auth_cache_entries": cacheSize,
|
||||
"refresh_cache_entries": refreshSize,
|
||||
"ttl_seconds": int(webAuthCache.ttl.Seconds()),
|
||||
"cleanup_interval": int(CleanupInterval.Seconds()),
|
||||
"session_refresh_time": int(SessionRefreshTime.Seconds()),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user