From a6e7c4ffb2381076c9f21ba594383f16de92565e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=99=86=E6=B2=89?= <surgesoft@gmail.com>
Date: Fri, 6 Sep 2024 00:00:48 +0800
Subject: [PATCH] feat: support 3rd auth (#350)

* feat: support 3rd auth

* feat: support 3rd auth

* feat: support 3rd auth
---
 assistant/package.json                        |  2 +-
 assistant/src/Chat/index.tsx                  |  6 ++--
 assistant/src/services/ChatController.ts      | 11 ++++++--
 .../20240905155017_remote_schema.sql          |  7 +++++
 server/auth/middleware.py                     | 28 +++++++++++++++++++
 server/core/models/bot.py                     |  3 +-
 6 files changed, 49 insertions(+), 8 deletions(-)
 create mode 100644 migrations/supabase/migrations/20240905155017_remote_schema.sql

diff --git a/assistant/package.json b/assistant/package.json
index f6237c26..78624d64 100644
--- a/assistant/package.json
+++ b/assistant/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@petercatai/assistant",
-  "version": "1.0.3",
+  "version": "1.0.5",
   "description": "Peter Cat Assistant Application",
   "module": "dist/esm/index.js",
   "types": "dist/esm/index.d.ts",
diff --git a/assistant/src/Chat/index.tsx b/assistant/src/Chat/index.tsx
index 9a862128..a31af2c3 100644
--- a/assistant/src/Chat/index.tsx
+++ b/assistant/src/Chat/index.tsx
@@ -74,9 +74,9 @@ const Chat: FC<ChatProps> = memo(
     const proChatRef = useRef<ProChatInstance>();
     const [chats, setChats] = useState<ChatMessage<Record<string, any>>[]>();
     const [complete, setComplete] = useState(false);
-    const { data: detail } = useSWR<BotInfo>(
-      token ? [`${apiDomain}/api/bot/detail?id=${token}`] : null,
-      fetcher,
+    const { data: detail } = useSWR(
+      token ? [`${apiDomain}/api/bot/detail?id=${token}`, token] : null,
+      fetcher<BotInfo>,
     );
 
     const [botInfo, setBotInfo] = useState<BotInfo>({
diff --git a/assistant/src/services/ChatController.ts b/assistant/src/services/ChatController.ts
index 8a3f90e8..4d1016f6 100644
--- a/assistant/src/services/ChatController.ts
+++ b/assistant/src/services/ChatController.ts
@@ -19,6 +19,7 @@ export async function streamChat(
       'Content-Type': 'application/json',
       connection: 'keep-alive',
       'keep-alive': 'timeout=5',
+      'Authorization': `Bearer ${token}`
     },
     body: JSON.stringify({
       messages: messages,
@@ -28,9 +29,13 @@ export async function streamChat(
   });
 }
 
-export async function fetcher(url: string) {
-  const response = await axios.get(url);
-  return response.data.data;
+export async function fetcher<T>([url, token]: [string, string | undefined]) {
+  const response = await axios.get(url, {
+    headers: {
+      'Authorization': `Bearer ${token}`,
+    },
+  });
+  return response.data.data as T;
 }
 
 export const uploadImage = async (
diff --git a/migrations/supabase/migrations/20240905155017_remote_schema.sql b/migrations/supabase/migrations/20240905155017_remote_schema.sql
new file mode 100644
index 00000000..5e684837
--- /dev/null
+++ b/migrations/supabase/migrations/20240905155017_remote_schema.sql
@@ -0,0 +1,7 @@
+alter table "public"."bots" add column "domain_whitelist" text[] default '{}'::text[];
+
+alter table "public"."bots" add column "repo_name" text;
+
+alter table "public"."git_issue_tasks" add column "page_index" numeric;
+
+
diff --git a/server/auth/middleware.py b/server/auth/middleware.py
index df617324..6a4d0503 100644
--- a/server/auth/middleware.py
+++ b/server/auth/middleware.py
@@ -1,8 +1,14 @@
 from typing import Awaitable, Callable
 from fastapi import HTTPException, Request, status
 from fastapi.responses import JSONResponse
+from petercat_utils import get_env_variable
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import Response
+from fastapi.security import OAuth2PasswordBearer
+
+from core.dao.botDAO import BotDAO
+
+WEB_URL = get_env_variable("WEB_URL")
 
 ALLOW_LIST = [
   "/favicon.ico",
@@ -18,7 +24,25 @@
   "/api/chat/stream_qa",
 ]
 
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/api/auth/token")
 class AuthMiddleWare(BaseHTTPMiddleware):
+
+  async def oauth(self, request: Request):
+    referer = request.headers.get('referer')
+    if referer.startswith(WEB_URL):
+      return True
+    
+    token = await oauth2_scheme(request=request)
+
+    if token:
+      bot_dao = BotDAO()
+      bot = bot_dao.get_bot(bot_id=token)
+      return bot and (
+        "*" in bot.domain_whitelist
+        or
+        referer in bot.domain_whitelist
+      )
+        
   async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
     try:
       # Auth 相关的直接放过
@@ -28,6 +52,9 @@ async def dispatch(self, request: Request, call_next: Callable[[Request], Awaita
       if request.url.path in ALLOW_LIST:
         return await call_next(request)
       
+      if await self.oauth(request=request):
+        return await call_next(request)
+    
       # 获取 session 中的用户信息
       user = request.session.get("user") 
       if not user:
@@ -42,6 +69,7 @@ async def dispatch(self, request: Request, call_next: Callable[[Request], Awaita
     
       return await call_next(request)
     except HTTPException as e:
+      
         # 处理 HTTP 异常
         return JSONResponse(status_code=e.status_code, content={"detail": e.detail})
     except Exception as e:
diff --git a/server/core/models/bot.py b/server/core/models/bot.py
index 6feb509c..47776b76 100644
--- a/server/core/models/bot.py
+++ b/server/core/models/bot.py
@@ -10,4 +10,5 @@ class BotModel(BaseModel):
     prompt: Optional[str] = ""
     name: str
     llm: Optional[str] = "openai"
-    created_at: datetime = datetime.now()
\ No newline at end of file
+    created_at: datetime = datetime.now()
+    domain_whitelist: Optional[list[str]] = []
\ No newline at end of file