From eacf558e70a4b92afa98b26f9ca5eba33b4556d5 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 30 Aug 2023 17:29:07 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(chat.py):=20close=20websocke?=
 =?UTF-8?q?t=20connection=20with=20status=20code=201008=20and=20reason=20"?=
 =?UTF-8?q?Unauthorized"=20if=20user=20is=20not=20authenticated=20or=20ina?=
 =?UTF-8?q?ctive=20=F0=9F=90=9B=20fix(auth/utils.py):=20raise=20credential?=
 =?UTF-8?q?s=5Fexception=20if=20token=20has=20expired?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/chat.py         | 17 ++++++++++++++---
 src/backend/langflow/services/auth/utils.py |  4 ++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py
index 83d85b719..18c8a8a74 100644
--- a/src/backend/langflow/api/v1/chat.py
+++ b/src/backend/langflow/api/v1/chat.py
@@ -18,6 +18,10 @@ from langflow.services.utils import get_session
 from langflow.utils.logger import logger
 from cachetools import LRUCache
 from sqlmodel import Session
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from langflow.services.chat.manager import ChatManager
 
 router = APIRouter(tags=["Chat"])
 
@@ -33,16 +37,23 @@ async def chat(
 ):
     """Websocket endpoint for chat."""
     try:
+        await websocket.accept()
         user = await get_current_user(token, db)
+        if not user:
+            await websocket.close(
+                code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
+            )
         if not user.is_active:
-            raise HTTPException(status_code=401, detail="Invalid token")
-        chat_manager = service_manager.get(ServiceType.CHAT_MANAGER)
+            await websocket.close(
+                code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized"
+            )
+
+        chat_manager: "ChatManager" = service_manager.get(ServiceType.CHAT_MANAGER)
         if client_id in chat_manager.in_memory_cache:
             await chat_manager.handle_websocket(client_id, websocket)
         else:
             # We accept the connection but close it immediately
             # if the flow is not built yet
-            await websocket.accept()
             message = "Please, build the flow before sending messages"
             await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=message)
     except WebSocketException as exc:
diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py
index a8e3e1790..a434fefcb 100644
--- a/src/backend/langflow/services/auth/utils.py
+++ b/src/backend/langflow/services/auth/utils.py
@@ -88,6 +88,10 @@ async def get_current_user(
         )
         user_id: UUID = payload.get("sub")  # type: ignore
         token_type: str = payload.get("type")  # type: ignore
+        if expires := payload.get("exp", None):
+            expires_datetime = datetime.fromtimestamp(expires)
+            if datetime.now(timezone.utc) > expires_datetime:
+                raise credentials_exception
 
         if user_id is None or token_type:
             raise credentials_exception