Initial commit

internal/middleware/auth.go

@@ -0,0 +1,82 @@
+// Filename: internal/middleware/auth.go
+package middleware
+
+import (
+	"gemini-balancer/internal/service"
+	"net/http"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+)
+
+// === API Admin 认证管道 (/admin/* API路由) ===
+
+func APIAdminAuthMiddleware(securityService *service.SecurityService) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		tokenValue := extractBearerToken(c)
+		if tokenValue == "" {
+			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Authorization token is missing"})
+			return
+		}
+		authToken, err := securityService.AuthenticateToken(tokenValue)
+		if err != nil || !authToken.IsAdmin {
+			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid or non-admin token"})
+			return
+		}
+		c.Set("adminUser", authToken)
+		c.Next()
+	}
+}
+
+// === /v1 Proxy 认证 ===
+
+func ProxyAuthMiddleware(securityService *service.SecurityService) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		tokenValue := extractProxyToken(c)
+		if tokenValue == "" {
+			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "API key is missing from request"})
+			return
+		}
+		authToken, err := securityService.AuthenticateToken(tokenValue)
+		if err != nil {
+			// 通用信息，避免泄露过多信息
+			c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{"error": "Invalid or inactive token provided"})
+			return
+		}
+		c.Set("authToken", authToken)
+		c.Next()
+	}
+}
+
+func extractProxyToken(c *gin.Context) string {
+	if key := c.Query("key"); key != "" {
+		return key
+	}
+	authHeader := c.GetHeader("Authorization")
+	if authHeader != "" {
+		if strings.HasPrefix(authHeader, "Bearer ") {
+			return strings.TrimPrefix(authHeader, "Bearer ")
+		}
+	}
+	if key := c.GetHeader("X-Api-Key"); key != "" {
+		return key
+	}
+	if key := c.GetHeader("X-Goog-Api-Key"); key != "" {
+		return key
+	}
+	return ""
+}
+
+// === 辅助函数 ===
+
+func extractBearerToken(c *gin.Context) string {
+	authHeader := c.GetHeader("Authorization")
+	if authHeader == "" {
+		return ""
+	}
+	parts := strings.Split(authHeader, " ")
+	if len(parts) == 2 && parts[0] == "Bearer" {
+		return parts[1]
+	}
+	return ""
+}

internal/middleware/log.go

@@ -0,0 +1,84 @@
+// Filename: internal/middleware/log_redaction.go
+package middleware
+
+import (
+	"bytes"
+	"io"
+	"regexp"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/sirupsen/logrus"
+)
+
+const RedactedBodyKey = "redactedBody"
+const RedactedAuthHeaderKey = "redactedAuthHeader"
+const RedactedValue = `"[REDACTED]"`
+
+func RedactionMiddleware() gin.HandlerFunc {
+	// Pre-compile regex for efficiency
+	jsonKeyPattern := regexp.MustCompile(`("api_key"|"keys")\s*:\s*"[^"]*"`)
+	bearerTokenPattern := regexp.MustCompile(`^(Bearer\s+)\S+$`)
+	return func(c *gin.Context) {
+		// --- 1. Redact Request Body ---
+		if c.Request.Method == "POST" || c.Request.Method == "PUT" || c.Request.Method == "DELETE" {
+			if bodyBytes, err := io.ReadAll(c.Request.Body); err == nil {
+
+				c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
+
+				bodyString := string(bodyBytes)
+
+				redactedBody := jsonKeyPattern.ReplaceAllString(bodyString, `$1:`+RedactedValue)
+
+				c.Set(RedactedBodyKey, redactedBody)
+			}
+		}
+		// --- 2. Redact Authorization Header ---
+		authHeader := c.GetHeader("Authorization")
+		if authHeader != "" {
+			if bearerTokenPattern.MatchString(authHeader) {
+				redactedHeader := bearerTokenPattern.ReplaceAllString(authHeader, `${1}[REDACTED]`)
+				c.Set(RedactedAuthHeaderKey, redactedHeader)
+			}
+		}
+		c.Next()
+	}
+}
+
+// LogrusLogger is a Gin middleware that logs requests using a Logrus logger.
+// It consumes redacted data prepared by the RedactionMiddleware.
+func LogrusLogger(logger *logrus.Logger) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		start := time.Now()
+		path := c.Request.URL.Path
+
+		// Process request
+		c.Next()
+
+		// After request, gather data and log
+		latency := time.Since(start)
+		statusCode := c.Writer.Status()
+
+		entry := logger.WithFields(logrus.Fields{
+			"status_code": statusCode,
+			"latency_ms":  latency.Milliseconds(),
+			"client_ip":   c.ClientIP(),
+			"method":      c.Request.Method,
+			"path":        path,
+		})
+
+		if redactedBody, exists := c.Get(RedactedBodyKey); exists {
+			entry = entry.WithField("body", redactedBody)
+		}
+
+		if redactedAuth, exists := c.Get(RedactedAuthHeaderKey); exists {
+			entry = entry.WithField("authorization", redactedAuth)
+		}
+
+		if len(c.Errors) > 0 {
+			entry.Error(c.Errors.String())
+		} else {
+			entry.Info("request handled")
+		}
+	}
+}

internal/middleware/security.go

@@ -0,0 +1,31 @@
+// Filename: internal/middleware/security.go
+
+package middleware
+
+import (
+	"gemini-balancer/internal/service"
+	"gemini-balancer/internal/settings"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+)
+
+func IPBanMiddleware(securityService *service.SecurityService, settingsManager *settings.SettingsManager) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		if !settingsManager.IsIPBanEnabled() {
+			c.Next()
+			return
+		}
+		ip := c.ClientIP()
+		isBanned, err := securityService.IsIPBanned(c.Request.Context(), ip)
+		if err != nil {
+			c.Next()
+			return
+		}
+		if isBanned {
+			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "您的IP已被暂时封禁，请稍后再试"})
+			return
+		}
+		c.Next()
+	}
+}

internal/middleware/web.go

@@ -0,0 +1,54 @@
+// Filename: internal/middleware/web.go
+package middleware
+
+import (
+	"gemini-balancer/internal/service"
+	"log"
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+)
+
+const (
+	AdminSessionCookie = "gemini_admin_session"
+)
+
+func SetAdminSessionCookie(c *gin.Context, adminToken string) {
+	c.SetCookie(AdminSessionCookie, adminToken, 3600*24*7, "/", "", false, true)
+}
+
+func ClearAdminSessionCookie(c *gin.Context) {
+	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 {
+	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 {
+			ClearAdminSessionCookie(c)
+			c.Redirect(http.StatusFound, "/login")
+			c.Abort()
+			return
+		}
+		c.Set("adminUser", authToken)
+		c.Next()
+	}
+}