jwt handling improvements

Author: taylorwilsdon

auth/auth_info_middleware.py

@@ -2,7 +2,6 @@
 Authentication middleware to populate context state with user information
 """
 
-import jwt
 import logging
 import time
 from fastmcp.server.middleware import Middleware, MiddlewareContext
@@ -187,90 +186,11 @@ async def _process_request_for_auth(self, context: MiddlewareContext):
                                 )
 
                         else:
-                            # Decode JWT to get user info
-                            logger.info("Processing JWT token")
-                            try:
-                                token_payload = jwt.decode(
-                                    token_str, options={"verify_signature": False}
-                                )
-                                logger.info(
-                                    f"JWT payload decoded: {list(token_payload.keys())}"
-                                )
-
-                                # Create an AccessToken-like object
-                                access_token = WorkspaceAccessToken(
-                                    token=token_str,
-                                    client_id=token_payload.get("client_id", "unknown"),
-                                    scopes=token_payload.get("scope", "").split()
-                                    if token_payload.get("scope")
-                                    else [],
-                                    session_id=token_payload.get(
-                                        "sid",
-                                        token_payload.get(
-                                            "jti",
-                                            token_payload.get("session_id", "unknown"),
-                                        ),
-                                    ),
-                                    expires_at=token_payload.get("exp", 0),
-                                    claims=token_payload,
-                                    sub=token_payload.get("sub"),
-                                    email=token_payload.get("email"),
-                                )
-
-                                # Store in context state
-                                context.fastmcp_context.set_state(
-                                    "access_token", access_token
-                                )
-
-                                # Store additional user info
-                                context.fastmcp_context.set_state(
-                                    "user_id", token_payload.get("sub")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "username",
-                                    token_payload.get(
-                                        "username", token_payload.get("email")
-                                    ),
-                                )
-                                context.fastmcp_context.set_state(
-                                    "name", token_payload.get("name")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "auth_time", token_payload.get("auth_time")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "issuer", token_payload.get("iss")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "audience", token_payload.get("aud")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "jti", token_payload.get("jti")
-                                )
-                                context.fastmcp_context.set_state(
-                                    "auth_provider_type", self.auth_provider_type
-                                )
-
-                                # Set the definitive authentication state for JWT tokens
-                                user_email = token_payload.get(
-                                    "email", token_payload.get("username")
-                                )
-                                if user_email:
-                                    context.fastmcp_context.set_state(
-                                        "authenticated_user_email", user_email
-                                    )
-                                    context.fastmcp_context.set_state(
-                                        "authenticated_via", "jwt_token"
-                                    )
-                                    authenticated_user = user_email
-                                    auth_via = "jwt_token"
-
-                            except jwt.DecodeError:
-                                logger.error("Failed to decode JWT token")
-                            except Exception as e:
-                                logger.error(
-                                    f"Error processing JWT: {type(e).__name__}"
-                                )
+                            # Non-Google JWT tokens require verification
+                            # SECURITY: Never set authenticated_user_email from unverified tokens
+                            logger.debug(
+                                "Unverified JWT token rejected - only verified tokens accepted"
+                            )
                     else:
                         logger.debug("No Bearer token in Authorization header")
                 else:

auth/mcp_session_middleware.py

@@ -62,24 +62,8 @@ async def dispatch(self, request: Request, call_next: Callable) -> Any:
                 mcp_session_id = request.state.session_id
                 logger.debug(f"Found FastMCP session ID: {mcp_session_id}")
 
-            # Also check Authorization header for bearer tokens
-            auth_header = headers.get("authorization")
-            if (
-                auth_header
-                and auth_header.lower().startswith("bearer ")
-                and not user_email
-            ):
-                try:
-                    import jwt
-
-                    token = auth_header[7:]  # Remove "Bearer " prefix
-                    # Decode without verification to extract email
-                    claims = jwt.decode(token, options={"verify_signature": False})
-                    user_email = claims.get("email")
-                    if user_email:
-                        logger.debug(f"Extracted user email from JWT: {user_email}")
-                except Exception:
-                    pass
+            # SECURITY: Do not decode JWT without verification
+            # User email must come from verified sources only (FastMCP auth context)
 
             # Build session context
             if session_id or auth_context or user_email or mcp_session_id:

auth/oauth21_session_store.py

@@ -153,22 +153,22 @@ def extract_session_from_headers(headers: Dict[str, str]) -> Optional[str]:
     # Try Authorization header for Bearer token
     auth_header = headers.get("authorization") or headers.get("Authorization")
     if auth_header and auth_header.lower().startswith("bearer "):
-        # Extract bearer token and try to find associated session
         token = auth_header[7:]  # Remove "Bearer " prefix
+        # Intentionally ignore empty tokens - "Bearer " with no token should not
+        # create a session context (avoids hash collisions on empty string)
         if token:
-            # Look for a session that has this access token
-            # This requires scanning sessions, but bearer tokens should be unique
+            # Use thread-safe lookup to find session by access token
             store = get_oauth21_session_store()
-            for user_email, session_info in store._sessions.items():
-                if session_info.get("access_token") == token:
-                    return session_info.get("session_id") or f"bearer_{user_email}"
+            session_id = store.find_session_id_for_access_token(token)
+            if session_id:
+                return session_id
 
-        # If no session found, create a temporary session ID from token hash
-        # This allows header-based authentication to work with session context
-        import hashlib
+            # If no session found, create a temporary session ID from token hash
+            # This allows header-based authentication to work with session context
+            import hashlib
 
-        token_hash = hashlib.sha256(token.encode()).hexdigest()[:8]
-        return f"bearer_token_{token_hash}"
+            token_hash = hashlib.sha256(token.encode()).hexdigest()[:8]
+            return f"bearer_token_{token_hash}"
 
     return None
 
@@ -325,6 +325,32 @@ def store_session(
         """
         with self._lock:
             normalized_expiry = _normalize_expiry_to_naive_utc(expiry)
+
+            # Clean up previous session mappings for this user before storing new one
+            old_session = self._sessions.get(user_email)
+            if old_session:
+                old_mcp_session_id = old_session.get("mcp_session_id")
+                old_session_id = old_session.get("session_id")
+                # Remove old MCP session mapping if it differs from new one
+                if old_mcp_session_id and old_mcp_session_id != mcp_session_id:
+                    if old_mcp_session_id in self._mcp_session_mapping:
+                        del self._mcp_session_mapping[old_mcp_session_id]
+                        logger.debug(
+                            f"Removed stale MCP session mapping: {old_mcp_session_id}"
+                        )
+                    if old_mcp_session_id in self._session_auth_binding:
+                        del self._session_auth_binding[old_mcp_session_id]
+                        logger.debug(
+                            f"Removed stale auth binding: {old_mcp_session_id}"
+                        )
+                # Remove old OAuth session binding if it differs from new one
+                if old_session_id and old_session_id != session_id:
+                    if old_session_id in self._session_auth_binding:
+                        del self._session_auth_binding[old_session_id]
+                        logger.debug(
+                            f"Removed stale OAuth session binding: {old_session_id}"
+                        )
+
             session_info = {
                 "access_token": access_token,
                 "refresh_token": refresh_token,
@@ -570,6 +596,9 @@ def remove_session(self, user_email: str):
                 if not mcp_session_id:
                     logger.info(f"Removed OAuth 2.1 session for {user_email}")
 
+            # Clean up any orphaned mappings that may have accumulated
+            self._cleanup_orphaned_mappings_locked()
+
     def has_session(self, user_email: str) -> bool:
         """Check if a user has an active session."""
         with self._lock:
@@ -597,6 +626,71 @@ def get_stats(self) -> Dict[str, Any]:
                 "mcp_sessions": list(self._mcp_session_mapping.keys()),
             }
 
+    def find_session_id_for_access_token(self, token: str) -> Optional[str]:
+        """
+        Thread-safe lookup of session ID by access token.
+
+        Args:
+            token: The access token to search for
+
+        Returns:
+            Session ID if found, None otherwise
+        """
+        with self._lock:
+            for user_email, session_info in self._sessions.items():
+                if session_info.get("access_token") == token:
+                    return session_info.get("session_id") or f"bearer_{user_email}"
+            return None
+
+    def _cleanup_orphaned_mappings_locked(self) -> int:
+        """Remove orphaned mappings. Caller must hold lock."""
+        # Collect valid session IDs and mcp_session_ids from active sessions
+        valid_session_ids = set()
+        valid_mcp_session_ids = set()
+        for session_info in self._sessions.values():
+            if session_info.get("session_id"):
+                valid_session_ids.add(session_info["session_id"])
+            if session_info.get("mcp_session_id"):
+                valid_mcp_session_ids.add(session_info["mcp_session_id"])
+
+        removed = 0
+
+        # Remove orphaned MCP session mappings
+        orphaned_mcp = [
+            sid for sid in self._mcp_session_mapping
+            if sid not in valid_mcp_session_ids
+        ]
+        for sid in orphaned_mcp:
+            del self._mcp_session_mapping[sid]
+            removed += 1
+            logger.debug(f"Removed orphaned MCP session mapping: {sid}")
+
+        # Remove orphaned auth bindings
+        valid_bindings = valid_session_ids | valid_mcp_session_ids
+        orphaned_bindings = [
+            sid for sid in self._session_auth_binding
+            if sid not in valid_bindings
+        ]
+        for sid in orphaned_bindings:
+            del self._session_auth_binding[sid]
+            removed += 1
+            logger.debug(f"Removed orphaned auth binding: {sid}")
+
+        if removed > 0:
+            logger.info(f"Cleaned up {removed} orphaned session mappings/bindings")
+
+        return removed
+
+    def cleanup_orphaned_mappings(self) -> int:
+        """
+        Remove orphaned entries from mcp_session_mapping and session_auth_binding.
+
+        Returns:
+            Number of orphaned entries removed
+        """
+        with self._lock:
+            return self._cleanup_orphaned_mappings_locked()
+
 
 # Global instance
 _global_store = OAuth21SessionStore()

auth/oauth_config.py

@@ -9,6 +9,7 @@
 """
 
 import os
+from threading import RLock
 from urllib.parse import urlparse
 from typing import List, Optional, Dict, Any
 
@@ -356,35 +357,40 @@ def get_authorization_server_metadata(
         return metadata
 
 
-# Global configuration instance
+# Global configuration instance with thread-safe access
 _oauth_config = None
+_oauth_config_lock = RLock()
 
 
 def get_oauth_config() -> OAuthConfig:
     """
     Get the global OAuth configuration instance.
 
+    Thread-safe singleton accessor.
+
     Returns:
         The singleton OAuth configuration instance
     """
     global _oauth_config
-    if _oauth_config is None:
-        _oauth_config = OAuthConfig()
-    return _oauth_config
+    with _oauth_config_lock:
+        if _oauth_config is None:
+            _oauth_config = OAuthConfig()
+        return _oauth_config
 
 
 def reload_oauth_config() -> OAuthConfig:
     """
     Reload the OAuth configuration from environment variables.
 
-    This is useful for testing or when environment variables change.
+    Thread-safe reload that prevents races with concurrent access.
 
     Returns:
         The reloaded OAuth configuration instance
     """
     global _oauth_config
-    _oauth_config = OAuthConfig()
-    return _oauth_config
+    with _oauth_config_lock:
+        _oauth_config = OAuthConfig()
+        return _oauth_config
 
 
 # Convenience functions for backward compatibility