From 148e8be28739711aee8631b77170fb2e39b1e2ec Mon Sep 17 00:00:00 2001
From: tbphp <admin@tbphp.net>
Date: Sun, 29 Jun 2025 09:50:54 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=94=AF=E6=8C=81gemini=E8=AE=A4?=
 =?UTF-8?q?=E8=AF=81key?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/middleware/middleware.go | 49 +++++++++++++++++++------------
 1 file changed, 31 insertions(+), 18 deletions(-)

diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go
index cda0e7e..9201d58 100644
--- a/internal/middleware/middleware.go
+++ b/internal/middleware/middleware.go
@@ -137,31 +137,19 @@ func Auth(config types.AuthConfig) gin.HandlerFunc {
 			return
 		}
 
-		// Get authorization header
-		authHeader := c.GetHeader("Authorization")
-		if authHeader == "" {
+		// Extract key from multiple sources
+		key := extractKey(c)
+		if key == "" {
 			c.JSON(401, gin.H{
-				"error": "Authorization header required",
+				"error": "Authorization required",
 				"code":  errors.ErrAuthMissing,
 			})
 			c.Abort()
 			return
 		}
 
-		// Check Bearer token format
-		const bearerPrefix = "Bearer "
-		if !strings.HasPrefix(authHeader, bearerPrefix) {
-			c.JSON(401, gin.H{
-				"error": "Invalid authorization format, expected 'Bearer <token>'",
-				"code":  errors.ErrAuthInvalid,
-			})
-			c.Abort()
-			return
-		}
-
-		// Extract and validate token
-		token := authHeader[len(bearerPrefix):]
-		if token != config.Key {
+		// Validate key
+		if key != config.Key {
 			c.JSON(401, gin.H{
 				"error": "Invalid authentication token",
 				"code":  errors.ErrAuthInvalid,
@@ -252,3 +240,28 @@ func isMonitoringEndpoint(path string) bool {
 	}
 	return false
 }
+
+// extractKey extracts the API key from the request, checking the Authorization header,
+// the X-Goog-Api-Key header, and the "key" query parameter.
+func extractKey(c *gin.Context) string {
+	// 1. Check Authorization header
+	authHeader := c.GetHeader("Authorization")
+	if authHeader != "" {
+		const bearerPrefix = "Bearer "
+		if strings.HasPrefix(authHeader, bearerPrefix) {
+			return authHeader[len(bearerPrefix):]
+		}
+	}
+
+	// 2. Check X-Goog-Api-Key header
+	if key := c.GetHeader("X-Goog-Api-Key"); key != "" {
+		return key
+	}
+
+	// 3. Check "key" query parameter
+	if key := c.Query("key"); key != "" {
+		return key
+	}
+
+	return ""
+}