fix: 🐛 authentication fixes for login and logout functionality

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ repos:
     -   id: trailing-whitespace
     -   id: requirements-txt-fixer
 -   repo: https://github.com/charliermarsh/ruff-pre-commit
-    rev: v0.13.0
+    rev: v0.14.4
     hooks:
     -   id: ruff
         args:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tp-auth-serverside"
-version = "0.1.4"
+version = "0.1.5"
 description = "A server side authentication utility which stores session tokens in memory db."
 readme = "README.md"
 requires-python = ">=3.13"

src/tp_auth_serverside/config.py

@@ -77,6 +77,6 @@ def check_secrets(cls, values) -> dict:
 Secrets = _Secrets()
 Service = _Service()
 Database = _Database()
-oauth2_scheme = CustomOAuth2PasswordBearer(tokenUrl=Secrets.token_url, scopes=Secrets.scopes)
+oauth2_scheme = CustomOAuth2PasswordBearer(tokenUrl=Secrets.token_url, scopes=Secrets.scopes, auto_error=False)
 
 __all__ = ["Secrets", "SupportedAlgorithms", "Database", "Service", "oauth2_scheme"]

src/tp_auth_serverside/core/fastapi_configurer.py

@@ -4,7 +4,7 @@
 from typing import Callable, Optional, Tuple
 
 from asgi_correlation_id import CorrelationIdMiddleware
-from fastapi import APIRouter, Depends, FastAPI, Request, Response
+from fastapi import APIRouter, Cookie, Depends, FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
 from fastapi.responses import ORJSONResponse
@@ -14,6 +14,7 @@
 
 from tp_auth_serverside.auth.auth_validator import AuthValidatorInstance
 from tp_auth_serverside.auth.schemas import Token
+from tp_auth_serverside.auth.user_specs import UserInfoSchema
 from tp_auth_serverside.config import Secrets, Service
 from tp_auth_serverside.core.handler.authentication_handler import AuthenticationHandler
 
@@ -120,6 +121,25 @@ async def token(
     return app
 
 
+def add_logout_route(app: FastAPI, handler: Callable = None, asynced: bool = False) -> FastAPI:
+    @app.post("/logout", response_model=StatusResponse, tags=["Authentication"])
+    async def logout(
+        request: Request,
+        response: Response,
+        access_token: Annotated[str, Cookie()],
+        user: Annotated[UserInfoSchema, Depends(AuthValidatorInstance)],
+    ) -> StatusResponse:
+        if handler is not None:
+            if asynced:
+                await handler(request, response, user.user_id, user)
+            else:
+                handler(request, response, user.user_id, user)
+        await AuthenticationHandler().revoke_authentication(response, user.user_id, access_token)
+        return StatusResponse()
+
+    return app
+
+
 async def start_refresh_service():
     """
     Start the gRPC refresh service as a background task.
@@ -144,6 +164,7 @@ def generate_fastapi_app(
     routers: list[APIRouter],
     disable_operation_default: bool = False,
     token_route_handler: Optional[Callable | Tuple[Callable, bool]] = None,
+    logout_route_handler: Optional[Callable | Tuple[Callable, bool]] = None,
     health_check_routine: Optional[Callable | Tuple[Callable, bool]] = None,
 ) -> FastAPI:
     # Create lifespan context manager for gRPC server lifecycle
@@ -195,6 +216,13 @@ async def lifespan_with_grpc(app: FastAPI):
             app = add_token_route(app, *token_route_handler)
         else:
             app = add_token_route(app, token_route_handler)
+    if logout_route_handler:
+        if isinstance(logout_route_handler, tuple):
+            app = add_logout_route(app, *logout_route_handler)
+        else:
+            app = add_logout_route(app, logout_route_handler)
+    else:
+        app = add_logout_route(app)
     app.add_middleware(CorrelationIdMiddleware)
     return app
 

src/tp_auth_serverside/core/handler/authentication_handler.py

@@ -1,7 +1,7 @@
 from fastapi import Response
 
 from tp_auth_serverside.auth.user_specs import UserInfoSchema
-from tp_auth_serverside.db.memorydb.login import set_token
+from tp_auth_serverside.db.memorydb.login import revoke_token, set_token
 from tp_auth_serverside.db.memorydb.refresh import set_restrict_refresh
 from tp_auth_serverside.utilities.jwt_util import JWTUtil
 
@@ -15,3 +15,8 @@ async def authenticate(self, response: Response, user_id: str, payload: UserInfo
         response.set_cookie(key="user_id", value=user_id, httponly=True, secure=True, samesite="strict")
         response.set_cookie(key="access_token", value=token, httponly=True, secure=True, samesite="strict")
         return token
+
+    async def revoke_authentication(self, response: Response, user_id: str, token: str):
+        await revoke_token(user_id, token)
+        response.delete_cookie(key="user_id")
+        response.delete_cookie(key="access_token")

src/tp_auth_serverside/db/memorydb/login.py

@@ -6,7 +6,7 @@
 
 
 async def set_token(user_id: str, token: str, expire_minutes: int = Secrets.expiry, short_token: str = None) -> str:
-    short_token = short_token or shortuuid.uuid(name=user_id)
+    short_token = short_token or shortuuid.uuid()
     await login_db.hset(user_id, short_token, orjson.dumps({"token": token, "expire": expire_minutes}))
     await login_db.hexpire(user_id, expire_minutes * 60, short_token)
     return short_token
@@ -17,3 +17,10 @@ async def get_token(user_id: str, short_token: str) -> str | None:
     if token_data:
         return orjson.loads(token_data).get("token")
     return None
+
+
+async def revoke_token(user_id: str, short_token: str = None):
+    if short_token:
+        await login_db.hdel(user_id, short_token)
+    else:
+        await login_db.delete(user_id)