From a5b30c3d85f2b78d1c88f2038f886b83faa091e4 Mon Sep 17 00:00:00 2001
From: prabinoid <38830224+prabinoid@users.noreply.github.com>
Date: Tue, 19 Nov 2024 18:04:03 +0545
Subject: [PATCH] fix: Invalid token exception in TokenAuthBackend

---
 backend/api/annotations/resources.py          | 13 +++----
 backend/main.py                               | 34 +++++++++----------
 .../services/users/authentication_service.py  | 22 ++++++------
 3 files changed, 35 insertions(+), 34 deletions(-)

diff --git a/backend/api/annotations/resources.py b/backend/api/annotations/resources.py
index 591421bb4a..ba23a6ef5c 100644
--- a/backend/api/annotations/resources.py
+++ b/backend/api/annotations/resources.py
@@ -1,14 +1,15 @@
+from databases import Database
+from fastapi import APIRouter, Depends, Request
+from loguru import logger
+from starlette.authentication import requires
+
+from backend.db import get_db
+from backend.models.dtos.user_dto import AuthUserDTO
 from backend.models.postgis.task import Task
 from backend.models.postgis.task_annotation import TaskAnnotation
 from backend.services.project_service import ProjectService
 from backend.services.task_annotations_service import TaskAnnotationsService
-from fastapi import APIRouter, Depends, Request
-from starlette.authentication import requires
-from loguru import logger
-from backend.db import get_db
-from databases import Database
 from backend.services.users.authentication_service import login_required
-from backend.models.dtos.user_dto import AuthUserDTO
 
 router = APIRouter(
     prefix="/projects",
diff --git a/backend/main.py b/backend/main.py
index 6cd4e21bfd..931f2f1952 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -42,6 +42,23 @@ async def lifespan(app):
     # Set custom logger
     # _app.logger = get_logger()
 
+    # Custom exception handler for 401 errors
+    @_app.exception_handler(HTTPException)
+    async def custom_http_exception_handler(request: Request, exc: HTTPException):
+        if exc.status_code == 401 and "InvalidToken" in exc.detail.get("SubCode", ""):
+            return JSONResponse(
+                content={
+                    "Error": exc.detail["Error"],
+                    "SubCode": exc.detail["SubCode"],
+                },
+                status_code=exc.status_code,
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return JSONResponse(
+            status_code=exc.status_code,
+            content={"detail": exc.detail},
+        )
+
     PROFILING = True  # Set this from a settings model
 
     if PROFILING:
@@ -71,23 +88,6 @@ async def pyinstrument_middleware(request, call_next):
         AuthenticationMiddleware, backend=TokenAuthBackend(), on_error=None
     )
 
-    # Custom exception handler for 401 errors
-    @_app.exception_handler(HTTPException)
-    async def custom_http_exception_handler(request: Request, exc: HTTPException):
-        if exc.status_code == 401 and "InvalidToken" in exc.detail.get("SubCode", ""):
-            return JSONResponse(
-                content={
-                    "Error": exc.detail["Error"],
-                    "SubCode": exc.detail["SubCode"],
-                },
-                status_code=exc.status_code,
-                headers={"WWW-Authenticate": "Bearer"},
-            )
-        return JSONResponse(
-            status_code=exc.status_code,
-            content={"detail": exc.detail},
-        )
-
     add_api_end_points(_app)
     return _app
 
diff --git a/backend/services/users/authentication_service.py b/backend/services/users/authentication_service.py
index d89dc82109..1bb77772d7 100644
--- a/backend/services/users/authentication_service.py
+++ b/backend/services/users/authentication_service.py
@@ -87,17 +87,17 @@ async def authenticate(self, conn):
             decoded_token, 604800
         )
         if not valid_token:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail={
-                    "Error": "Token is expired or invalid",
-                    "SubCode": "InvalidToken",
-                },
-                headers={"WWW-Authenticate": "Bearer"},
-            )
-        tm.authenticated_user_id = (
-            user_id  # Set the user ID on the decorator as a convenience
-        )
+            logger.debug("Token not valid...")
+            return
+            # raise HTTPException(
+            #     status_code=401,
+            #     detail={
+            #         "Error": "Token is expired or invalid",
+            #         "SubCode": "InvalidToken",
+            #     },
+            #     headers={"WWW-Authenticate": "Bearer"},
+            # )
+        tm.authenticated_user_id = user_id
         return AuthCredentials(["authenticated"]), SimpleUser(user_id)