From 0b1f3dfc047927befab2be1732982b003702e292 Mon Sep 17 00:00:00 2001
From: snaily <snaily@linux.do>
Date: Fri, 28 Mar 2025 19:27:42 +0800
Subject: [PATCH] =?UTF-8?q?feat(auth):=20=E6=94=AF=E6=8C=81x-goog-api-key?=
 =?UTF-8?q?=E8=AF=B7=E6=B1=82=E5=A4=B4=E8=AE=A4=E8=AF=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 添加verify_key_or_goog_api_key方法，支持同时验证URL参数中的key和请求头中的x-goog-api-key
- 更新models接口使用新的认证方法，提高与Google API客户端的兼容性
---
 app/core/security.py        | 19 +++++++++++++++++++
 app/router/gemini_routes.py |  2 +-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/app/core/security.py b/app/core/security.py
index 060928a..a759823 100644
--- a/app/core/security.py
+++ b/app/core/security.py
@@ -72,3 +72,22 @@ class SecurityService:
             raise HTTPException(status_code=401, detail="Invalid auth_token")
 
         return token
+
+    async def verify_key_or_goog_api_key(
+        self, key: Optional[str] = None , x_goog_api_key: Optional[str] = Header(None)
+    ) -> str:
+        """验证URL中的key或请求头中的x-goog-api-key"""
+        # 如果URL中的key有效，直接返回
+        if key in self.allowed_tokens or key == self.auth_token:
+            return key
+        
+        # 否则检查请求头中的x-goog-api-key
+        if not x_goog_api_key:
+            logger.error("Invalid key and missing x-goog-api-key header")
+            raise HTTPException(status_code=401, detail="Invalid key and missing x-goog-api-key header")
+        
+        if x_goog_api_key not in self.allowed_tokens and x_goog_api_key != self.auth_token:
+            logger.error("Invalid key and invalid x-goog-api-key")
+            raise HTTPException(status_code=401, detail="Invalid key and invalid x-goog-api-key")
+        
+        return x_goog_api_key
\ No newline at end of file
diff --git a/app/router/gemini_routes.py b/app/router/gemini_routes.py
index 06d6700..50a9efc 100644
--- a/app/router/gemini_routes.py
+++ b/app/router/gemini_routes.py
@@ -34,7 +34,7 @@ async def get_next_working_key(key_manager: KeyManager = Depends(get_key_manager
 @router.get("/models")
 @router_v1beta.get("/models")
 async def list_models(
-    _=Depends(security_service.verify_key),
+    _=Depends(security_service.verify_key_or_goog_api_key),
     key_manager: KeyManager = Depends(get_key_manager)
 ):
     """获取可用的Gemini模型列表"""