From a100feb295892660757a42e95efd724449f5b310 Mon Sep 17 00:00:00 2001
From: "raoha.rh" <raohai.rh@antgroup.com>
Date: Sun, 18 Aug 2024 20:03:15 +0800
Subject: [PATCH] fix: user auth using petercat user token

---
 server/auth/get_user_info.py | 10 ++++++++++
 server/routers/chat.py       |  4 ++--
 server/verify/rate_limit.py  |  8 +++++---
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/server/auth/get_user_info.py b/server/auth/get_user_info.py
index 2e20d837..92aa3aad 100644
--- a/server/auth/get_user_info.py
+++ b/server/auth/get_user_info.py
@@ -63,6 +63,16 @@ async def getAnonymousUserInfoByToken(token: str):
     rows = supabase.table("profiles").select("*").eq("id", token).execute()
     return rows.data[0] if (len(rows.data) > 0) else None
 
+async def get_user_id(petercat_user_token: Annotated[str | None, Cookie()] = None):
+    try:
+        if petercat_user_token is None:
+            return None
+        user_info = await getUserInfoByToken(petercat_user_token)
+        return user_info['id']
+
+    except Exception as e:
+        return None
+
 async def get_user_access_token(petercat_user_token: Annotated[str | None, Cookie()] = None):
     try:
         if petercat_user_token is None:
diff --git a/server/routers/chat.py b/server/routers/chat.py
index 87036817..16ccb957 100644
--- a/server/routers/chat.py
+++ b/server/routers/chat.py
@@ -2,7 +2,7 @@
 from typing import Annotated, Optional
 from fastapi import APIRouter, Cookie, Depends
 from fastapi.responses import StreamingResponse
-from auth.get_user_info import get_user_access_token
+from auth.get_user_info import get_user_access_token, get_user_id
 from petercat_utils.data_class import ChatData
 from agent import qa_chat, bot_builder
 from verify.rate_limit import verify_rate_limit
@@ -32,7 +32,7 @@ async def run_issue_helper(input_data: ChatData):
 
 
 @router.post("/stream_builder", response_class=StreamingResponse, dependencies=[Depends(verify_rate_limit)])
-def run_bot_builder(input_data: ChatData, user_id: str = Cookie(None), bot_id: Optional[str] = None):
+def run_bot_builder(input_data: ChatData, bot_id: Optional[str] = None, user_id: Annotated[str | None, Depends(get_user_id)] = None):
     if not user_id:
         return StreamingResponse(generate_auth_failed_stream(), media_type="text/event-stream")
     result = bot_builder.agent_stream_chat(input_data, user_id, bot_id)
diff --git a/server/verify/rate_limit.py b/server/verify/rate_limit.py
index 4d0c7319..2a49e826 100644
--- a/server/verify/rate_limit.py
+++ b/server/verify/rate_limit.py
@@ -1,3 +1,4 @@
+from typing import Annotated
 from fastapi import Cookie, HTTPException
 from datetime import datetime, timedelta
 
@@ -8,13 +9,14 @@
 RATE_LIMIT_REQUESTS = get_env_variable("RATE_LIMIT_REQUESTS") or 100
 RATE_LIMIT_DURATION = timedelta(minutes=int(get_env_variable("RATE_LIMIT_DURATION") or 1))
 
-async def verify_rate_limit(petercat: str = Cookie(None)):
+async def verify_rate_limit(petercat_user_token: Annotated[str | None, Cookie()] = None):
     if not RATE_LIMIT_ENABLED:
         return
 
-    if not petercat:
+    if not petercat_user_token:
         raise HTTPException(status_code=403, detail="Must Login")
-    user = await getUserInfoByToken(petercat)
+    user = await getUserInfoByToken(petercat_user_token)
+
     if user is None:
         raise HTTPException(
             status_code=429,