refactor: enhance JWT token handling and session management

- Integrated PyJWKClient for secure JWT verification, improving token validation by using signing keys from JWKS.
- Updated UserSession initialization to decode tokens with verification, enhancing security and error handling.
- Added a caching mechanism for the JWKS client to optimize performance.
- Cleaned up token expiration checks to ensure accurate validation and error reporting.

Author: atyrode

src/backend/config.py

@@ -4,6 +4,7 @@
 import httpx
 import redis
 import jwt
+from jwt.jwks_client import PyJWKClient
 from typing import Optional, Dict, Any, Tuple
 from dotenv import load_dotenv
 
@@ -47,6 +48,9 @@
 CODER_DEFAULT_ORGANIZATION = os.getenv("CODER_DEFAULT_ORGANIZATION")
 CODER_WORKSPACE_NAME = os.getenv("CODER_WORKSPACE_NAME", "ubuntu")
 
+# Cache for JWKS client
+_jwks_client = None
+
 # Session management functions
 def get_session(session_id: str) -> Optional[Dict[str, Any]]:
     """Get session data from Redis"""
@@ -67,8 +71,6 @@ def delete_session(session_id: str) -> None:
     """Delete session data from Redis"""
     redis_client.delete(f"session:{session_id}")
 
-provisioning_times = {}
-
 def get_auth_url() -> str:
     """Generate the authentication URL for Keycloak login"""
     auth_url = f"{OIDC_SERVER_URL}/realms/{OIDC_REALM}/protocol/openid-connect/auth"
@@ -85,29 +87,28 @@ def get_token_url() -> str:
     return f"{OIDC_SERVER_URL}/realms/{OIDC_REALM}/protocol/openid-connect/token"
 
 def is_token_expired(token_data: Dict[str, Any], buffer_seconds: int = 30) -> bool:
-    """
-    Check if the access token is expired or about to expire
-    
-    Args:
-        token_data: The token data containing the access token
-        buffer_seconds: Buffer time in seconds to refresh token before it actually expires
-        
-    Returns:
-        bool: True if token is expired or about to expire, False otherwise
-    """
     if not token_data or 'access_token' not in token_data:
         return True
 
     try:
-        # Decode the JWT token without verification to get expiration time
-        decoded = jwt.decode(token_data['access_token'], options={"verify_signature": False})
+        # Get the signing key
+        jwks_client = get_jwks_client()
+        signing_key = jwks_client.get_signing_key_from_jwt(token_data['access_token'])
 
-        # Get expiration time from token
-        exp_time = decoded.get('exp', 0)
+        # Decode with verification
+        decoded = jwt.decode(
+            token_data['access_token'],
+            signing_key.key,
+            algorithms=["RS256"],  # Common algorithm for OIDC
+            audience=OIDC_CLIENT_ID,
+        )
 
-        # Check if token is expired or about to expire (with buffer)
+        # Check expiration
+        exp_time = decoded.get('exp', 0)
         current_time = time.time()
         return current_time + buffer_seconds >= exp_time
+    except jwt.ExpiredSignatureError:
+        return True
     except Exception as e:
         print(f"Error checking token expiration: {str(e)}")
         return True
@@ -153,3 +154,11 @@ async def refresh_token(session_id: str, token_data: Dict[str, Any]) -> Tuple[bo
     except Exception as e:
         print(f"Error refreshing token: {str(e)}")
         return False, token_data
+
+def get_jwks_client():
+    """Get or create a PyJWKClient for token verification"""
+    global _jwks_client
+    if _jwks_client is None:
+        jwks_url = f"{OIDC_SERVER_URL}/realms/{OIDC_REALM}/protocol/openid-connect/certs"
+        _jwks_client = PyJWKClient(jwks_url)
+    return _jwks_client

src/backend/dependencies.py

@@ -15,8 +15,25 @@ class UserSession:
     """
     def __init__(self, access_token: str, token_data: dict, user_id: UUID = None):
         self.access_token = access_token
-        self.token_data = jwt.decode(access_token, options={"verify_signature": False})
         self._user_data = None
+
+        # Get the signing key and decode with verification
+        from config import get_jwks_client, OIDC_CLIENT_ID
+        try:
+            jwks_client = get_jwks_client()
+            signing_key = jwks_client.get_signing_key_from_jwt(access_token)
+            
+            self.token_data = jwt.decode(
+                access_token,
+                signing_key.key,
+                algorithms=["RS256"],
+                audience=OIDC_CLIENT_ID
+            )
+
+        except jwt.InvalidTokenError as e:
+            # Log the error and raise an appropriate exception
+            print(f"Invalid token: {str(e)}")
+            raise ValueError(f"Invalid authentication token: {str(e)}")
 
     @property
     def is_authenticated(self) -> bool:

src/backend/requirements.txt

@@ -10,4 +10,5 @@ sqlalchemy
 posthog
 redis
 psycopg2-binary
-python-multipart
+python-multipart
+cryptography  # Required for JWT key handling