Refactor auth service to improve refresh token revocation handling and logging

app/api/v1/auth.py

@@ -8,12 +8,9 @@ from app.core.config import settings
 from app.core.dependencies import get_auth_service, get_current_active_user
 from app.core.logging import get_logger
 from app.models.user import User
-from app.schemas.auth import (
-    UserLoginRequest,
-    UserRegisterRequest,
-    UserResponse,
-)
+from app.schemas.auth import UserLoginRequest, UserRegisterRequest, UserResponse
 from app.services.auth import AuthService
+from app.utils.auth import JWTUtils
 
 router = APIRouter()
 logger = get_logger(__name__)
@@ -48,7 +45,10 @@ async def register(
         response.set_cookie(
             key="refresh_token",
             value=refresh_token,
-            max_age=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS * 24 * 60 * 60,  # Convert days to seconds
+            max_age=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS
+            * 24
+            * 60
+            * 60,  # Convert days to seconds
             httponly=True,
             secure=settings.COOKIE_SECURE,
             samesite=settings.COOKIE_SAMESITE,
@@ -92,7 +92,10 @@ async def login(
         response.set_cookie(
             key="refresh_token",
             value=refresh_token,
-            max_age=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS * 24 * 60 * 60,  # Convert days to seconds
+            max_age=settings.JWT_REFRESH_TOKEN_EXPIRE_DAYS
+            * 24
+            * 60
+            * 60,  # Convert days to seconds
             httponly=True,
             secure=settings.COOKIE_SECURE,
             samesite=settings.COOKIE_SAMESITE,
@@ -167,32 +170,56 @@ async def refresh_token(
 @router.post("/logout")
 async def logout(
     response: Response,
-    current_user: Annotated[User, Depends(get_current_active_user)],
+    access_token: Annotated[str | None, Cookie()],
+    refresh_token: Annotated[str | None, Cookie()],
     auth_service: Annotated[AuthService, Depends(get_auth_service)],
 ) -> dict[str, str]:
     """Logout endpoint - clears cookies and revokes refresh token."""
-    try:
-        # Revoke refresh token from database
-        await auth_service.revoke_refresh_token(current_user)
+    user = None
     
-        # Clear both cookies
-        response.delete_cookie(
-            key="access_token",
-            httponly=True,
-            secure=settings.COOKIE_SECURE,
-            samesite=settings.COOKIE_SAMESITE,
-        )
-        response.delete_cookie(
-            key="refresh_token",
-            httponly=True,
-            secure=settings.COOKIE_SECURE,
-            samesite=settings.COOKIE_SAMESITE,
-        )
+    # Try to get user from access token first
+    if access_token:
+        try:
+            payload = JWTUtils.decode_access_token(access_token)
+            user_id_str = payload.get("sub")
+            if user_id_str:
+                user_id = int(user_id_str)
+                user = await auth_service.get_current_user(user_id)
+                logger.info("Found user from access token: %s", user.email)
+        except (HTTPException, Exception) as e:
+            logger.info("Access token validation failed: %s", str(e))
     
-        return {"message": "Successfully logged out"}
-    except Exception as e:
-        logger.exception("Logout failed")
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail="Logout failed",
-        ) from e
+    # If no user found, try refresh token
+    if not user and refresh_token:
+        try:
+            payload = JWTUtils.decode_refresh_token(refresh_token)
+            user_id_str = payload.get("sub")
+            if user_id_str:
+                user_id = int(user_id_str)
+                user = await auth_service.get_current_user(user_id)
+                logger.info("Found user from refresh token: %s", user.email)
+        except (HTTPException, Exception) as e:
+            logger.info("Refresh token validation failed: %s", str(e))
+    
+    # If we found a user, revoke their refresh token
+    if user:
+        await auth_service.revoke_refresh_token(user)
+        logger.info("Successfully revoked refresh token for user: %s", user.email)
+    else:
+        logger.info("No user found, skipping token revocation")
+    
+    # Always clear both cookies regardless of token validity
+    response.delete_cookie(
+        key="access_token",
+        httponly=True,
+        secure=settings.COOKIE_SECURE,
+        samesite=settings.COOKIE_SAMESITE,
+    )
+    response.delete_cookie(
+        key="refresh_token",
+        httponly=True,
+        secure=settings.COOKIE_SECURE,
+        samesite=settings.COOKIE_SAMESITE,
+    )
+    
+    return {"message": "Successfully logged out"}

app/services/auth.py

@@ -231,11 +231,17 @@ class AuthService:
 
     async def revoke_refresh_token(self, user: User) -> None:
         """Revoke a user's refresh token."""
-        user.refresh_token_hash = None
-        user.refresh_token_expires_at = None
-        self.session.add(user)
-        await self.session.commit()
-        logger.info("Refresh token revoked for user: %s", user.email)
+        try:
+            # Use the repository to update the user to ensure proper session handling
+            update_data = {
+                "refresh_token_hash": None,
+                "refresh_token_expires_at": None,
+            }
+            await self.user_repo.update(user, update_data)
+            logger.info("Refresh token revoked for user: %s", user.email)
+        except Exception:
+            logger.exception("Failed to revoke refresh token for user: %s", user.email)
+            raise
 
     async def create_user_response(self, user: User) -> UserResponse:
         """Create a user response from a user model."""