tokrn version

core/models.py

@@ -54,6 +54,10 @@ class User(Base):
     # Legacy field (kept for migration, nullable now)
     secret_key_hash = Column(String(255), nullable=True)
     
+    # Token versioning for JWT invalidation
+    # Incrementing this invalidates all existing tokens for this user
+    token_version = Column(Integer, default=1, nullable=False)
+    
     # Credits and status
     credits = Column(Integer, default=100)
     created_at = Column(DateTime(timezone=True), server_default=func.now())

dependencies.py

@@ -77,6 +77,8 @@ async def get_current_user(
     Extract and verify JWT from Authorization header.
     Returns the authenticated user.
     
+    Also validates token_version to support instant logout/invalidation.
+    
     Usage:
         @router.get("/protected")
         async def protected_route(user: User = Depends(get_current_user)):
@@ -135,6 +137,15 @@ async def get_current_user(
             detail="User not found or inactive"
         )
     
+    # Validate token version - if user's version is higher, token is invalidated
+    if payload.token_version < user.token_version:
+        logger.info(f"Token invalidated for user {user.user_id}: token_version {payload.token_version} < {user.token_version}")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Token has been invalidated. Please sign in again.",
+            headers={"WWW-Authenticate": "Bearer"}
+        )
+    
     return user
 
 

routers/auth.py

@@ -167,8 +167,8 @@ async def google_auth(
     db.add(audit_log)
     await db.commit()
     
-    # Create our JWT access token
-    access_token = create_access_token(user.user_id, user.email)
+    # Create our JWT access token with current token_version
+    access_token = create_access_token(user.user_id, user.email, user.token_version)
     
     # Sync DB to Drive (Async)
     background_tasks.add_task(drive_service.upload_db)
@@ -214,6 +214,8 @@ async def refresh_token(
     Use this when the current token is about to expire
     (or has recently expired) to get a new one without
     requiring the user to sign in again.
+    
+    Validates that the token_version is still valid before refreshing.
     """
     ip = req.client.host
     
@@ -226,7 +228,42 @@ async def refresh_token(
     
     try:
         jwt_service = get_jwt_service()
-        new_token = jwt_service.refresh_token(request.token)
+        
+        # Decode the token (without verifying expiry) to get user info
+        import jwt as pyjwt
+        payload = pyjwt.decode(
+            request.token,
+            jwt_service.secret_key,
+            algorithms=[jwt_service.algorithm],
+            options={"verify_exp": False}
+        )
+        
+        user_id = payload.get("sub")
+        token_version = payload.get("tv", 1)
+        
+        if not user_id:
+            raise JWTInvalidTokenError("Token missing required claims")
+        
+        # Check if user exists and token version is still valid
+        query = select(User).where(User.user_id == user_id, User.is_active == True)
+        result = await db.execute(query)
+        user = result.scalar_one_or_none()
+        
+        if not user:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="User not found or inactive"
+            )
+        
+        # Validate token version
+        if token_version < user.token_version:
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Token has been invalidated. Please sign in again."
+            )
+        
+        # Create new token with current token_version
+        new_token = create_access_token(user.user_id, user.email, user.token_version)
         
         return TokenRefreshResponse(
             success=True,
@@ -249,14 +286,15 @@ async def logout(
     """
     Logout current user.
     
-    Note: JWT tokens are stateless, so this endpoint mainly
-    serves to log the action. Frontend should discard the token.
-    
-    For full session invalidation, consider implementing
-    a token blacklist or reducing token expiry times.
+    Increments the user's token_version which invalidates ALL existing
+    tokens for this user. This provides instant logout across all devices.
     """
     ip = req.client.host
     
+    # Increment token version to invalidate all existing tokens
+    user.token_version += 1
+    logger.info(f"User {user.user_id} logged out. Token version incremented to {user.token_version}")
+    
     # Log logout
     audit_log = AuditLog(
         user_id=user.user_id,
@@ -270,4 +308,4 @@ async def logout(
     # Sync DB to Drive (Async)
     background_tasks.add_task(drive_service.upload_db)
     
-    return {"success": True, "message": "Logged out successfully"}
+    return {"success": True, "message": "Logged out successfully. All sessions invalidated."}

services/jwt_service.py

@@ -52,12 +52,14 @@ class TokenPayload:
         email: The user's email address
         issued_at: When the token was issued
         expires_at: When the token expires
+        token_version: Version number for token invalidation
         extra: Any additional claims in the token
     """
     user_id: str
     email: str
     issued_at: datetime
     expires_at: datetime
+    token_version: int = 1
     extra: Dict[str, Any] = None
     
     def __post_init__(self):
@@ -169,6 +171,7 @@ class JWTService:
         self,
         user_id: str,
         email: str,
+        token_version: int = 1,
         extra_claims: Optional[Dict[str, Any]] = None,
         expiry_hours: Optional[int] = None
     ) -> str:
@@ -178,6 +181,7 @@ class JWTService:
         Args:
             user_id: The user's unique identifier.
             email: The user's email address.
+            token_version: User's current token version for invalidation.
             extra_claims: Additional claims to include in the token.
             expiry_hours: Custom expiry for this token (overrides default).
         
@@ -190,6 +194,7 @@ class JWTService:
         payload = {
             "sub": user_id,
             "email": email,
+            "tv": token_version,  # Token version for invalidation
             "iat": now,
             "exp": now + timedelta(hours=expiry),
         }
@@ -199,7 +204,7 @@ class JWTService:
         
         token = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
         
-        logger.debug(f"Created token for user_id={user_id}")
+        logger.debug(f"Created token for user_id={user_id} (version={token_version})")
         return token
     
     def verify_token(self, token: str) -> TokenPayload:
@@ -229,6 +234,7 @@ class JWTService:
             # Extract standard claims
             user_id = payload.get("sub")
             email = payload.get("email")
+            token_version = payload.get("tv", 1)  # Default to 1 for backward compatibility
             iat = payload.get("iat")
             exp = payload.get("exp")
             
@@ -240,7 +246,7 @@ class JWTService:
             expires_at = datetime.utcfromtimestamp(exp) if isinstance(exp, (int, float)) else exp
             
             # Extract extra claims
-            standard_claims = {"sub", "email", "iat", "exp"}
+            standard_claims = {"sub", "email", "tv", "iat", "exp"}
             extra = {k: v for k, v in payload.items() if k not in standard_claims}
             
             return TokenPayload(
@@ -248,6 +254,7 @@ class JWTService:
                 email=email,
                 issued_at=issued_at,
                 expires_at=expires_at,
+                token_version=token_version,
                 extra=extra
             )
             
@@ -346,19 +353,20 @@ def get_jwt_service() -> JWTService:
     return _default_service
 
 
-def create_access_token(user_id: str, email: str, **kwargs) -> str:
+def create_access_token(user_id: str, email: str, token_version: int = 1, **kwargs) -> str:
     """
     Convenience function to create a token using the default service.
     
     Args:
         user_id: The user's unique identifier.
         email: The user's email address.
+        token_version: User's current token version for invalidation.
         **kwargs: Additional arguments passed to create_token.
     
     Returns:
         str: The encoded JWT token.
     """
-    return get_jwt_service().create_token(user_id, email, **kwargs)
+    return get_jwt_service().create_token(user_id, email, token_version, **kwargs)
 
 
 def verify_access_token(token: str) -> TokenPayload: