fix(auth): 🐛 Fix auth mechanism after testing and integration

.python-version

@@ -1 +1 @@
-3.12
+3.13

protos/google/protobuf/empty.proto

@@ -0,0 +1,21 @@
+syntax = "proto3";
+
+package google.protobuf;
+
+option go_package = "google.golang.org/protobuf/types/known/emptypb";
+option java_package = "com.google.protobuf";
+option java_outer_classname = "EmptyProto";
+option java_multiple_files = true;
+option objc_class_prefix = "GPB";
+option csharp_namespace = "Google.Protobuf.WellKnownTypes";
+option cc_enable_arenas = true;
+
+// A generic empty message that you can re-use to avoid defining duplicated
+// empty messages in your APIs. A typical example is to use it as the request
+// or the response type of an API method. For instance:
+//
+// service Foo {
+//   rpc Bar(google.protobuf.Empty) returns (google.protobuf.Empty);
+// }
+//
+message Empty {}

pyproject.toml

@@ -1,15 +1,17 @@
 [project]
 name = "tp-auth-serverside"
-version = "0.1.2"
+version = "0.1.3"
 description = "A server side authentication utility which stores session tokens in memory db."
 readme = "README.md"
-requires-python = ">=3.12"
+requires-python = ">=3.13"
 dependencies = [
     "cryptography>=45.0.5",
     "fastapi>=0.116.1",
     "grpc-interceptor>=0.15.4",
     "grpcio>=1.75.0",
+    "grpcio-tools>=1.75.0",
     "httpx>=0.28.1",
+    "jetpack>=0.2.0",
     "mem-db-utils>=0.2.0",
     "orjson>=3.11.1",
     "pydantic>=2.11.7",
@@ -26,7 +28,6 @@ build-backend = "uv_build"
 [dependency-groups]
 dev = [
     "pre-commit>=4.2.0",
-    "grpcio-tools>=1.75.0",
 ]
 
 [[tool.uv.index]]

src/tp_auth_serverside/auth/auth_validator.py

@@ -29,8 +29,8 @@ async def __call__(
         self,
         security_scopes: SecurityScopes,
         token: Annotated[str, Depends(oauth2_scheme)],
-        user_id: Annotated[str, Cookie],
-        refresh: Annotated[bool | True, Header] = True,
+        user_id: Annotated[str, Cookie()],
+        refresh: Annotated[bool, Header()] = True,
     ) -> UserInfoSchema:
         if security_scopes.scopes:
             authenticate_value = f"Bearer scope={security_scopes.scope_str}"

src/tp_auth_serverside/auth/custom_auth_scheme.py

@@ -0,0 +1,19 @@
+from typing import Optional
+
+from fastapi import HTTPException, Request, status
+from fastapi.security import OAuth2PasswordBearer
+
+
+class CustomOAuth2PasswordBearer(OAuth2PasswordBearer):
+    async def __call__(self, request: Request) -> Optional[str]:
+        token = request.cookies.get("access_token")
+        if not token:
+            if self.auto_error:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Not authenticated",
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+            else:
+                return None
+        return token

src/tp_auth_serverside/config.py

@@ -1,11 +1,12 @@
 from enum import StrEnum
 from typing import Any, Optional
 
-from fastapi.security import OAuth2PasswordBearer
 from pydantic import BeforeValidator, Field, model_validator
 from pydantic_settings import BaseSettings
 from typing_extensions import Annotated
 
+from tp_auth_serverside.auth.custom_auth_scheme import CustomOAuth2PasswordBearer
+
 
 def options_decoder(v):
     if isinstance(v, str):
@@ -46,14 +47,15 @@ class _Secrets(BaseSettings):
     leeway: Optional[int] = Field(10, env="LEEWAY")
     expiry: Optional[int] = Field(1440, env="EXPIRY")
     authorization_server: Optional[bool] = Field(False, env="AUTHORIZATION_SERVER")
-    scopes: Optional[dict] = Field(None, env="AUTH_SCOPES")
+    scopes: Optional[dict] = Field(None, alias="AUTH_SCOPES")
     token_url: Optional[str] = Field("/token", env="TOKEN_URL")
     refresh_url: Optional[str] = Field("/refresh", env="REFRESH_URL")
     refresh_restrict_minutes: Optional[int] = Field(2, env="REFRESH_RESTRICT_MINUTES")
 
     @model_validator(mode="before")
     def check_secrets(cls, values) -> dict:
-        if values["algorithm"] == SupportedAlgorithms.RS256:
+        algorithm = values.get("algorithm", SupportedAlgorithms.HS256)
+        if algorithm == SupportedAlgorithms.RS256:
             import base64
 
             if not values.get("public_key"):
@@ -66,7 +68,7 @@ def check_secrets(cls, values) -> dict:
             if private_bytes:
                 private_bytes = private_bytes.encode("utf-8")
                 values["private_key"] = base64.b64decode(private_bytes).decode("utf-8")
-        elif values["algorithm"] == SupportedAlgorithms.HS256:
+        elif algorithm == SupportedAlgorithms.HS256:
             if not values.get("secret_key"):
                 raise ValueError("Secret key must be provided for HS256 algorithm")
         return values
@@ -75,6 +77,6 @@ def check_secrets(cls, values) -> dict:
 Secrets = _Secrets()
 Service = _Service()
 Database = _Database()
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl=Secrets.token_url, scopes=Secrets.scopes)
+oauth2_scheme = CustomOAuth2PasswordBearer(tokenUrl=Secrets.token_url, scopes=Secrets.scopes)
 
 __all__ = ["Secrets", "SupportedAlgorithms", "Database", "Service", "oauth2_scheme"]

src/tp_auth_serverside/core/fastapi_configurer.py

@@ -1,7 +1,9 @@
 import logging
-from asyncio import futures
+from concurrent import futures
+from contextlib import asynccontextmanager
 from typing import Callable, Optional, Tuple
 
+from asgi_correlation_id import CorrelationIdMiddleware
 from fastapi import APIRouter, Depends, FastAPI, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.openapi.utils import get_openapi
@@ -100,22 +102,29 @@ def add_cors(app: FastAPI) -> FastAPI:
     return app
 
 
-def add_token_route(app: FastAPI, handler: Callable, asynced: bool = False) -> FastAPI:
-    @app.post("/token", response_model=Token)
+def add_token_route(app: FastAPI, handler: Callable, asynced: bool = False, dependency=None) -> FastAPI:
+    @app.post("/token", response_model=Token, tags=["Authentication"])
     async def token(
-        form_data: Annotated[OAuth2PasswordRequestForm, Depends()], request: Request, response: Response
+        form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
+        request: Request,
+        response: Response,
+        dependency=Depends(dependency),
     ) -> Token:
         if asynced:
-            user_id, payload = await handler(form_data, request, response)
+            user_id, payload = await handler(form_data, request, response, dependency)
         else:
-            user_id, payload = handler(form_data, request, response)
+            user_id, payload = handler(form_data, request, response, dependency)
         token = await AuthenticationHandler().authenticate(response, user_id, payload)
         return Token(user_id=user_id, token=token)
 
     return app
 
 
-def start_refresh_service():
+async def start_refresh_service():
+    """
+    Start the gRPC refresh service as a background task.
+    Returns the server instance for lifecycle management.
+    """
     import grpc
 
     from tp_auth_serverside.core.handler.refresh_handler import RefreshHandler
@@ -125,9 +134,9 @@ def start_refresh_service():
     refresh_pb2_grpc.add_RefreshServiceServicer_to_server(RefreshHandler(), server)
     server.add_insecure_port(Secrets.refresh_url)
     logging.info(f"Starting refresh service on {Secrets.refresh_url}")
-    server.start()
+    await server.start()
     logging.info("Refresh service started")
-    server.wait_for_termination()
+    return server
 
 
 def generate_fastapi_app(
@@ -137,6 +146,31 @@ def generate_fastapi_app(
     token_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
+    @asynccontextmanager
+    async def lifespan_with_grpc(app: FastAPI):
+        grpc_server = None
+        # Startup: Start the gRPC refresh service
+        if Secrets.authorization_server:
+            logging.info("Initializing gRPC refresh service...")
+            grpc_server = await start_refresh_service()
+
+        # Call user-provided lifespan if exists
+        if app_config.lifespan:
+            async with app_config.lifespan(app):
+                yield
+        else:
+            yield
+
+        # Shutdown: Stop the gRPC server gracefully
+        if grpc_server:
+            logging.info("Shutting down gRPC refresh service...")
+            await grpc_server.stop(grace=5)
+            logging.info("gRPC refresh service stopped")
+
+    # Use the combined lifespan if authorization_server is enabled
+    final_lifespan = lifespan_with_grpc if Secrets.authorization_server else app_config.lifespan
+
     app = FastAPI(
         title=app_config.title,
         version=app_config.version,
@@ -145,7 +179,7 @@ def generate_fastapi_app(
         openapi_url=app_config.openapi_url,
         docs_url=app_config.docs_url,
         redoc_url=app_config.redoc_url,
-        lifespan=app_config.lifespan,
+        lifespan=final_lifespan,
         exception_handlers=app_config.exception_handlers,
         default_response_class=ORJSONResponse,
     )
@@ -158,11 +192,10 @@ def generate_fastapi_app(
     app = add_cors(app)
     if token_route_handler:
         if isinstance(token_route_handler, tuple):
-            app = add_token_route(app, token_route_handler[0], token_route_handler[1])
+            app = add_token_route(app, *token_route_handler)
         else:
             app = add_token_route(app, token_route_handler)
-    if Secrets.authorization_server:
-        start_refresh_service()
+    app.add_middleware(CorrelationIdMiddleware)
     return app
 
 

src/tp_auth_serverside/core/handler/refresh_handler.py

@@ -1,4 +1,6 @@
-import asyncio
+import logging
+
+from google.protobuf import empty_pb2
 
 from tp_auth_serverside.db.memorydb.login import get_token, set_token
 from tp_auth_serverside.db.memorydb.refresh import is_refresh_restricted, set_restrict_refresh
@@ -7,21 +9,22 @@
 
 
 class RefreshHandler(RefreshServiceServicer):
-    @staticmethod
-    def RefreshToken(request, context):
+    async def RefreshToken(self, request, context):
         user_id = request.user_id
         token = request.token
-        if asyncio.run(is_refresh_restricted(user_id, token)):
-            return None  # Since the response type is google.protobuf.empty_pb2.Empty
-        jwt_token = asyncio.run(get_token(user_id, token))
+        if await is_refresh_restricted(user_id, token):
+            logging.warning(f"Refresh token is restricted for user_id: {user_id}, token: {token}")
+            return empty_pb2.Empty()
+        jwt_token = await get_token(user_id, token)
         if not jwt_token:
-            return None  # Since the response type is google.protobuf.empty_pb2.Empty
+            return empty_pb2.Empty()
         jwt_util = JWTUtil()
         try:
+            logging.info(f"Refreshing token for user_id: {user_id}, token: {token}")
             payload = jwt_util.decode(jwt_token)
             jwt_token = jwt_util.encode(payload=payload)
-            asyncio.run(set_token(user_id, jwt_token, short_token=token))
-            asyncio.run(set_restrict_refresh(user_id, token))
-        except Exception:
-            pass
-        return None  # Since the response type is google.protobuf.empty_pb2.Empty
+            await set_token(user_id, jwt_token, short_token=token)
+            await set_restrict_refresh(user_id, token)
+        except Exception as e:
+            logging.error(f"Error refreshing token for user_id: {user_id}, token: {token}, error: {e}")
+        return empty_pb2.Empty()

src/tp_auth_serverside/db/memorydb/__init__.py

@@ -1,9 +1,11 @@
+import asyncio
+
 from mem_db_utils.asyncio import MemDBConnector
 from redis import Redis
 
 from tp_auth_serverside.config import Database
 
-mem_connector = MemDBConnector(db=Database.login_redis_db)
+mem_connector = MemDBConnector()
 
-login_db: Redis = mem_connector.connect(db=Database.login_redis_db, decode_response=True)
-refresh_restrict_db: Redis = mem_connector.connect(db=Database.refresh_restrict_db, decode_response=True)
+login_db: Redis = asyncio.run(mem_connector.connect(db=Database.login_redis_db, decode_response=True))
+refresh_restrict_db: Redis = asyncio.run(mem_connector.connect(db=Database.refresh_restrict_db, decode_response=True))

src/tp_auth_serverside/db/memorydb/login.py

@@ -1,3 +1,4 @@
+import orjson
 import shortuuid
 
 from tp_auth_serverside.config import Secrets
@@ -6,13 +7,13 @@
 
 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)
-    await login_db.hset(user_id, short_token, {"token": token, "expire": expire_minutes})
-    await login_db.hexpire(user_id, expire_minutes * 60, fields=[short_token])
+    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
 
 
 async def get_token(user_id: str, short_token: str) -> str | None:
     token_data = await login_db.hget(user_id, short_token)
     if token_data:
-        return token_data.get("token")
+        return orjson.loads(token_data).get("token")
     return None

src/tp_auth_serverside/db/memorydb/refresh.py

@@ -1,10 +1,15 @@
+import logging
+
 from tp_auth_serverside.config import Secrets
 from tp_auth_serverside.db.memorydb import refresh_restrict_db
 
 
 async def set_restrict_refresh(user_id: str, token: str) -> None:
+    logging.info(f"Setting refresh restrict for user_id: {user_id}, token: {token}")
     await refresh_restrict_db.set(f"{user_id}__{token}", "restricted", ex=Secrets.refresh_restrict_minutes * 60)
 
 
 async def is_refresh_restricted(user_id: str, token: str) -> bool:
-    return await refresh_restrict_db.exists(f"{user_id}__{token}") == 1
+    result = await refresh_restrict_db.exists(f"{user_id}__{token}")
+    logging.info(f"Checking refresh restrict for user_id: {user_id}, token: {token}, exists: {result}")
+    return result == 1

src/tp_auth_serverside/pb/refresh_pb2_grpc.py

@@ -6,7 +6,7 @@
 
 import tp_auth_serverside.pb.refresh_pb2 as refresh__pb2
 
-GRPC_GENERATED_VERSION = "1.75.0"
+GRPC_GENERATED_VERSION = "1.75.1"
 GRPC_VERSION = grpc.__version__
 _version_not_supported = False