From 4f01b82b81f72b99897261a7e52009bcf9bc455a Mon Sep 17 00:00:00 2001
From: InfinityPacer <160988576+InfinityPacer@users.noreply.github.com>
Date: Tue, 8 Oct 2024 14:32:29 +0800
Subject: [PATCH] feat(security): unify token validation for message endpoints

---
 app/api/endpoints/message.py     | 16 ++++++++--------
 app/modules/slack/__init__.py    |  4 ----
 app/modules/telegram/__init__.py |  4 ----
 app/modules/vocechat/__init__.py |  5 -----
 4 files changed, 8 insertions(+), 21 deletions(-)

diff --git a/app/api/endpoints/message.py b/app/api/endpoints/message.py
index 239b052f..48121d3f 100644
--- a/app/api/endpoints/message.py
+++ b/app/api/endpoints/message.py
@@ -9,7 +9,7 @@ from starlette.responses import PlainTextResponse
 from app import schemas
 from app.chain.message import MessageChain
 from app.core.config import settings, global_vars
-from app.core.security import verify_token
+from app.core.security import verify_token, verify_apitoken
 from app.db import get_db
 from app.db.models import User
 from app.db.models.message import Message
@@ -30,7 +30,8 @@ def start_message_chain(body: Any, form: Any, args: Any):
 
 
 @router.post("/", summary="接收用户消息", response_model=schemas.Response)
-async def user_message(background_tasks: BackgroundTasks, request: Request):
+async def user_message(background_tasks: BackgroundTasks, request: Request,
+                       _: schemas.TokenPayload = Depends(verify_apitoken)):
     """
     用户消息响应，配置请求中需要添加参数：token=API_TOKEN&source=消息配置名
     """
@@ -102,18 +103,17 @@ def wechat_verify(echostr: str, msg_signature: str, timestamp: Union[str, int],
     return "未找到对应的消息配置"
 
 
-def vocechat_verify(token: str) -> Any:
+def vocechat_verify() -> Any:
     """
     VoceChat验证响应
     """
-    if token == settings.API_TOKEN:
-        return {"status": "OK"}
-    return {"status": "API_TOKEN ERROR"}
+    return {"status": "OK"}
 
 
 @router.get("/", summary="回调请求验证")
 def incoming_verify(token: str = None, echostr: str = None, msg_signature: str = None,
-                    timestamp: Union[str, int] = None, nonce: str = None, source: str = None) -> Any:
+                    timestamp: Union[str, int] = None, nonce: str = None, source: str = None,
+                    _: schemas.TokenPayload = Depends(verify_apitoken)) -> Any:
     """
     微信/VoceChat等验证响应
     """
@@ -121,7 +121,7 @@ def incoming_verify(token: str = None, echostr: str = None, msg_signature: str =
                 f"msg_signature={msg_signature}, timestamp={timestamp}, nonce={nonce}")
     if echostr and msg_signature and timestamp and nonce:
         return wechat_verify(echostr, msg_signature, timestamp, nonce, source)
-    return vocechat_verify(token)
+    return vocechat_verify()
 
 
 @router.post("/webpush/subscribe", summary="客户端webpush通知订阅", response_model=schemas.Response)
diff --git a/app/modules/slack/__init__.py b/app/modules/slack/__init__.py
index f86d08ff..34476f85 100644
--- a/app/modules/slack/__init__.py
+++ b/app/modules/slack/__init__.py
@@ -171,10 +171,6 @@ class SlackModule(_ModuleBase, _MessageBase[Slack]):
         client_config = self.get_config(source)
         if not client_config:
             return None
-        # 校验token
-        token = args.get("token")
-        if not token or token != settings.API_TOKEN:
-            return None
         try:
             msg_json: dict = json.loads(body)
         except Exception as err:
diff --git a/app/modules/telegram/__init__.py b/app/modules/telegram/__init__.py
index af444c15..ad19647c 100644
--- a/app/modules/telegram/__init__.py
+++ b/app/modules/telegram/__init__.py
@@ -85,10 +85,6 @@ class TelegramModule(_ModuleBase, _MessageBase[Telegram]):
         if not client_config:
             return None
         client: Telegram = self.get_instance(source)
-        # 校验token
-        token = args.get("token")
-        if not token or token != settings.API_TOKEN:
-            return None
         try:
             message: dict = json.loads(body)
         except Exception as err:
diff --git a/app/modules/vocechat/__init__.py b/app/modules/vocechat/__init__.py
index 661bb580..479cd5f1 100644
--- a/app/modules/vocechat/__init__.py
+++ b/app/modules/vocechat/__init__.py
@@ -81,11 +81,6 @@ class VoceChatModule(_ModuleBase, _MessageBase[VoceChat]):
                 # 非新消息
                 return None
             logger.debug(f"收到VoceChat请求：{msg_body}")
-            # token校验
-            token = args.get("token")
-            if not token or token != settings.API_TOKEN:
-                logger.warn(f"VoceChat请求token校验失败：{token}")
-                return None
             # 文本内容
             content = msg_body.get("detail", {}).get("content")
             # 用户ID