From 4616b508b152930dd244643286ae3bf2ff920f15 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 5 Nov 2024 21:14:02 -0800
Subject: [PATCH] refac: token handling

---
 backend/open_webui/main.py        |  9 ++++++++-
 backend/open_webui/utils/utils.py | 10 +++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index 599a7b51e..45ea6867f 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -2224,7 +2224,14 @@ async def get_app_config(request: Request):
     user = None
     if "token" in request.cookies:
         token = request.cookies.get("token")
-        data = decode_token(token)
+        try:
+            data = decode_token(token)
+        except Exception as e:
+            log.debug(e)
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Invalid token",
+            )
         if data is not None and "id" in data:
             user = Users.get_user_by_id(data["id"])
 
diff --git a/backend/open_webui/utils/utils.py b/backend/open_webui/utils/utils.py
index 79faa1831..31fe227ed 100644
--- a/backend/open_webui/utils/utils.py
+++ b/backend/open_webui/utils/utils.py
@@ -91,7 +91,15 @@ def get_current_user(
         return get_current_user_by_api_key(token)
 
     # auth by jwt token
-    data = decode_token(token)
+
+    try:
+        data = decode_token(token)
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid token",
+        )
+
     if data is not None and "id" in data:
         user = Users.get_user_by_id(data["id"])
         if user is None: