gemini-banlancer/internal/middleware/web.go

// Filename: internal/middleware/web.go

package middleware

import (
	"crypto/sha256"
	"encoding/hex"
	"gemini-balancer/internal/service"
	"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) {
	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)
}

func ExtractTokenFromCookie(c *gin.Context) string {
	cookie, err := c.Cookie(AdminSessionCookie)
	if err != nil {
		return ""
	}
	return cookie
}

// ==================== 认证中间件 ====================

func WebAdminAuthMiddleware(authService *service.SecurityService) gin.HandlerFunc {
	logger := logrus.New()
	logger.SetLevel(getLogLevel())

	return func(c *gin.Context) {
		cookie := ExtractTokenFromCookie(c)
		if cookie == "" {
			logger.Debug("[WebAuth] No session cookie found")
			ClearAdminSessionCookie(c)
			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()),
	}
}