apply ruff formatting
This commit is contained in:
Taylor Wilsdon
@@ -1,6 +1,7 @@
|
||||
"""
|
||||
Authentication middleware to populate context state with user information
|
||||
"""
|
||||
|
||||
import jwt
|
||||
import logging
|
||||
import os
|
||||
@@ -20,11 +21,11 @@ class AuthInfoMiddleware(Middleware):
|
||||
Middleware to extract authentication information from JWT tokens
|
||||
and populate the FastMCP context state for use in tools and prompts.
|
||||
"""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.auth_provider_type = "GoogleProvider"
|
||||
|
||||
|
||||
async def _process_request_for_auth(self, context: MiddlewareContext):
|
||||
"""Helper to extract, verify, and store auth info from a request."""
|
||||
if not context.fastmcp_context:
|
||||
@@ -42,69 +43,101 @@ class AuthInfoMiddleware(Middleware):
|
||||
headers = get_http_headers()
|
||||
if headers:
|
||||
logger.debug("Processing HTTP headers for authentication")
|
||||
|
||||
|
||||
# Get the Authorization header
|
||||
auth_header = headers.get("authorization", "")
|
||||
if auth_header.startswith("Bearer "):
|
||||
token_str = auth_header[7:] # Remove "Bearer " prefix
|
||||
logger.debug("Found Bearer token")
|
||||
|
||||
|
||||
# For Google OAuth tokens (ya29.*), we need to verify them differently
|
||||
if token_str.startswith("ya29."):
|
||||
logger.debug("Detected Google OAuth access token format")
|
||||
|
||||
|
||||
# Verify the token to get user info
|
||||
from core.server import get_auth_provider
|
||||
|
||||
auth_provider = get_auth_provider()
|
||||
|
||||
|
||||
if auth_provider:
|
||||
try:
|
||||
# Verify the token
|
||||
verified_auth = await auth_provider.verify_token(token_str)
|
||||
verified_auth = await auth_provider.verify_token(
|
||||
token_str
|
||||
)
|
||||
if verified_auth:
|
||||
# Extract user info from verified token
|
||||
user_email = None
|
||||
if hasattr(verified_auth, 'claims'):
|
||||
if hasattr(verified_auth, "claims"):
|
||||
user_email = verified_auth.claims.get("email")
|
||||
|
||||
|
||||
# Get expires_at, defaulting to 1 hour from now if not available
|
||||
if hasattr(verified_auth, 'expires_at'):
|
||||
if hasattr(verified_auth, "expires_at"):
|
||||
expires_at = verified_auth.expires_at
|
||||
else:
|
||||
expires_at = int(time.time()) + 3600 # Default to 1 hour
|
||||
|
||||
expires_at = (
|
||||
int(time.time()) + 3600
|
||||
) # Default to 1 hour
|
||||
|
||||
# Get client_id from verified auth or use default
|
||||
client_id = getattr(verified_auth, 'client_id', None) or "google"
|
||||
|
||||
client_id = (
|
||||
getattr(verified_auth, "client_id", None)
|
||||
or "google"
|
||||
)
|
||||
|
||||
access_token = SimpleNamespace(
|
||||
token=token_str,
|
||||
client_id=client_id,
|
||||
scopes=verified_auth.scopes if hasattr(verified_auth, 'scopes') else [],
|
||||
scopes=verified_auth.scopes
|
||||
if hasattr(verified_auth, "scopes")
|
||||
else [],
|
||||
session_id=f"google_oauth_{token_str[:8]}",
|
||||
expires_at=expires_at,
|
||||
# Add other fields that might be needed
|
||||
sub=verified_auth.sub if hasattr(verified_auth, 'sub') else user_email,
|
||||
email=user_email
|
||||
sub=verified_auth.sub
|
||||
if hasattr(verified_auth, "sub")
|
||||
else user_email,
|
||||
email=user_email,
|
||||
)
|
||||
|
||||
|
||||
# Store in context state - this is the authoritative authentication state
|
||||
context.fastmcp_context.set_state("access_token", access_token)
|
||||
mcp_session_id = getattr(context.fastmcp_context, "session_id", None)
|
||||
context.fastmcp_context.set_state(
|
||||
"access_token", access_token
|
||||
)
|
||||
mcp_session_id = getattr(
|
||||
context.fastmcp_context, "session_id", None
|
||||
)
|
||||
ensure_session_from_access_token(
|
||||
verified_auth,
|
||||
user_email,
|
||||
mcp_session_id,
|
||||
)
|
||||
context.fastmcp_context.set_state("access_token_obj", verified_auth)
|
||||
context.fastmcp_context.set_state("auth_provider_type", self.auth_provider_type)
|
||||
context.fastmcp_context.set_state("token_type", "google_oauth")
|
||||
context.fastmcp_context.set_state("user_email", user_email)
|
||||
context.fastmcp_context.set_state("username", user_email)
|
||||
context.fastmcp_context.set_state(
|
||||
"access_token_obj", verified_auth
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", self.auth_provider_type
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"token_type", "google_oauth"
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"user_email", user_email
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"username", user_email
|
||||
)
|
||||
# Set the definitive authentication state
|
||||
context.fastmcp_context.set_state("authenticated_user_email", user_email)
|
||||
context.fastmcp_context.set_state("authenticated_via", "bearer_token")
|
||||
|
||||
logger.info(f"Authenticated via Google OAuth: {user_email}")
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_user_email", user_email
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_via", "bearer_token"
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Authenticated via Google OAuth: {user_email}"
|
||||
)
|
||||
else:
|
||||
logger.error("Failed to verify Google OAuth token")
|
||||
# Don't set authenticated_user_email if verification failed
|
||||
@@ -113,18 +146,29 @@ class AuthInfoMiddleware(Middleware):
|
||||
# Still store the unverified token - service decorator will handle verification
|
||||
access_token = SimpleNamespace(
|
||||
token=token_str,
|
||||
client_id=os.getenv("GOOGLE_OAUTH_CLIENT_ID", "google"),
|
||||
client_id=os.getenv(
|
||||
"GOOGLE_OAUTH_CLIENT_ID", "google"
|
||||
),
|
||||
scopes=[],
|
||||
session_id=f"google_oauth_{token_str[:8]}",
|
||||
expires_at=int(time.time()) + 3600, # Default to 1 hour
|
||||
expires_at=int(time.time())
|
||||
+ 3600, # Default to 1 hour
|
||||
sub="unknown",
|
||||
email=""
|
||||
email="",
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"access_token", access_token
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", self.auth_provider_type
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"token_type", "google_oauth"
|
||||
)
|
||||
context.fastmcp_context.set_state("access_token", access_token)
|
||||
context.fastmcp_context.set_state("auth_provider_type", self.auth_provider_type)
|
||||
context.fastmcp_context.set_state("token_type", "google_oauth")
|
||||
else:
|
||||
logger.warning("No auth provider available to verify Google token")
|
||||
logger.warning(
|
||||
"No auth provider available to verify Google token"
|
||||
)
|
||||
# Store unverified token
|
||||
access_token = SimpleNamespace(
|
||||
token=token_str,
|
||||
@@ -133,51 +177,93 @@ class AuthInfoMiddleware(Middleware):
|
||||
session_id=f"google_oauth_{token_str[:8]}",
|
||||
expires_at=int(time.time()) + 3600, # Default to 1 hour
|
||||
sub="unknown",
|
||||
email=""
|
||||
email="",
|
||||
)
|
||||
context.fastmcp_context.set_state("access_token", access_token)
|
||||
context.fastmcp_context.set_state("auth_provider_type", self.auth_provider_type)
|
||||
context.fastmcp_context.set_state("token_type", "google_oauth")
|
||||
|
||||
context.fastmcp_context.set_state(
|
||||
"access_token", access_token
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", self.auth_provider_type
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"token_type", "google_oauth"
|
||||
)
|
||||
|
||||
else:
|
||||
# Decode JWT to get user info
|
||||
try:
|
||||
token_payload = jwt.decode(
|
||||
token_str,
|
||||
options={"verify_signature": False}
|
||||
token_str, options={"verify_signature": False}
|
||||
)
|
||||
logger.debug(f"JWT payload decoded: {list(token_payload.keys())}")
|
||||
|
||||
logger.debug(
|
||||
f"JWT payload decoded: {list(token_payload.keys())}"
|
||||
)
|
||||
|
||||
# Create an AccessToken-like object
|
||||
access_token = SimpleNamespace(
|
||||
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)
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
# Store in context state
|
||||
context.fastmcp_context.set_state("access_token", access_token)
|
||||
|
||||
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)
|
||||
|
||||
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"))
|
||||
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")
|
||||
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_user_email", user_email
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_via", "jwt_token"
|
||||
)
|
||||
|
||||
logger.debug("JWT token processed successfully")
|
||||
|
||||
|
||||
except jwt.DecodeError as e:
|
||||
logger.error(f"Failed to decode JWT: {e}")
|
||||
except Exception as e:
|
||||
@@ -185,44 +271,58 @@ class AuthInfoMiddleware(Middleware):
|
||||
else:
|
||||
logger.debug("No Bearer token in Authorization header")
|
||||
else:
|
||||
logger.debug("No HTTP headers available (might be using stdio transport)")
|
||||
logger.debug(
|
||||
"No HTTP headers available (might be using stdio transport)"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Could not get HTTP request: {e}")
|
||||
|
||||
|
||||
# After trying HTTP headers, check for other authentication methods
|
||||
# This consolidates all authentication logic in the middleware
|
||||
if not context.fastmcp_context.get_state("authenticated_user_email"):
|
||||
logger.debug("No authentication found via bearer token, checking other methods")
|
||||
|
||||
logger.debug(
|
||||
"No authentication found via bearer token, checking other methods"
|
||||
)
|
||||
|
||||
# Check transport mode
|
||||
from core.config import get_transport_mode
|
||||
|
||||
transport_mode = get_transport_mode()
|
||||
|
||||
|
||||
if transport_mode == "stdio":
|
||||
# In stdio mode, check if there's a session with credentials
|
||||
# This is ONLY safe in stdio mode because it's single-user
|
||||
logger.debug("Checking for stdio mode authentication")
|
||||
|
||||
|
||||
# Get the requested user from the context if available
|
||||
requested_user = None
|
||||
if hasattr(context, 'request') and hasattr(context.request, 'params'):
|
||||
requested_user = context.request.params.get('user_google_email')
|
||||
elif hasattr(context, 'arguments'):
|
||||
if hasattr(context, "request") and hasattr(context.request, "params"):
|
||||
requested_user = context.request.params.get("user_google_email")
|
||||
elif hasattr(context, "arguments"):
|
||||
# FastMCP may store arguments differently
|
||||
requested_user = context.arguments.get('user_google_email')
|
||||
|
||||
requested_user = context.arguments.get("user_google_email")
|
||||
|
||||
if requested_user:
|
||||
try:
|
||||
from auth.oauth21_session_store import get_oauth21_session_store
|
||||
|
||||
store = get_oauth21_session_store()
|
||||
|
||||
|
||||
# Check if user has a recent session
|
||||
if store.has_session(requested_user):
|
||||
logger.debug(f"Using recent stdio session for {requested_user}")
|
||||
logger.debug(
|
||||
f"Using recent stdio session for {requested_user}"
|
||||
)
|
||||
# In stdio mode, we can trust the user has authenticated recently
|
||||
context.fastmcp_context.set_state("authenticated_user_email", requested_user)
|
||||
context.fastmcp_context.set_state("authenticated_via", "stdio_session")
|
||||
context.fastmcp_context.set_state("auth_provider_type", "oauth21_stdio")
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_user_email", requested_user
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_via", "stdio_session"
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", "oauth21_stdio"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Error checking stdio session: {e}")
|
||||
|
||||
@@ -230,73 +330,95 @@ class AuthInfoMiddleware(Middleware):
|
||||
if not context.fastmcp_context.get_state("authenticated_user_email"):
|
||||
try:
|
||||
from auth.oauth21_session_store import get_oauth21_session_store
|
||||
|
||||
store = get_oauth21_session_store()
|
||||
single_user = store.get_single_user_email()
|
||||
if single_user:
|
||||
logger.debug(
|
||||
f"Defaulting to single stdio OAuth session for {single_user}"
|
||||
)
|
||||
context.fastmcp_context.set_state("authenticated_user_email", single_user)
|
||||
context.fastmcp_context.set_state("authenticated_via", "stdio_single_session")
|
||||
context.fastmcp_context.set_state("auth_provider_type", "oauth21_stdio")
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_user_email", single_user
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_via", "stdio_single_session"
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", "oauth21_stdio"
|
||||
)
|
||||
context.fastmcp_context.set_state("user_email", single_user)
|
||||
context.fastmcp_context.set_state("username", single_user)
|
||||
except Exception as e:
|
||||
logger.debug(f"Error determining stdio single-user session: {e}")
|
||||
|
||||
logger.debug(
|
||||
f"Error determining stdio single-user session: {e}"
|
||||
)
|
||||
|
||||
# Check for MCP session binding
|
||||
if not context.fastmcp_context.get_state("authenticated_user_email") and hasattr(context.fastmcp_context, 'session_id'):
|
||||
if not context.fastmcp_context.get_state(
|
||||
"authenticated_user_email"
|
||||
) and hasattr(context.fastmcp_context, "session_id"):
|
||||
mcp_session_id = context.fastmcp_context.session_id
|
||||
if mcp_session_id:
|
||||
try:
|
||||
from auth.oauth21_session_store import get_oauth21_session_store
|
||||
|
||||
store = get_oauth21_session_store()
|
||||
|
||||
|
||||
# Check if this MCP session is bound to a user
|
||||
bound_user = store.get_user_by_mcp_session(mcp_session_id)
|
||||
if bound_user:
|
||||
logger.debug(f"MCP session bound to {bound_user}")
|
||||
context.fastmcp_context.set_state("authenticated_user_email", bound_user)
|
||||
context.fastmcp_context.set_state("authenticated_via", "mcp_session_binding")
|
||||
context.fastmcp_context.set_state("auth_provider_type", "oauth21_session")
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_user_email", bound_user
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"authenticated_via", "mcp_session_binding"
|
||||
)
|
||||
context.fastmcp_context.set_state(
|
||||
"auth_provider_type", "oauth21_session"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Error checking MCP session binding: {e}")
|
||||
|
||||
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next):
|
||||
"""Extract auth info from token and set in context state"""
|
||||
logger.debug("Processing tool call authentication")
|
||||
|
||||
|
||||
try:
|
||||
await self._process_request_for_auth(context)
|
||||
|
||||
|
||||
logger.debug("Passing to next handler")
|
||||
result = await call_next(context)
|
||||
logger.debug("Handler completed")
|
||||
return result
|
||||
|
||||
|
||||
except Exception as e:
|
||||
# Check if this is an authentication error - don't log traceback for these
|
||||
if "GoogleAuthenticationError" in str(type(e)) or "Access denied: Cannot retrieve credentials" in str(e):
|
||||
if "GoogleAuthenticationError" in str(
|
||||
type(e)
|
||||
) or "Access denied: Cannot retrieve credentials" in str(e):
|
||||
logger.info(f"Authentication check failed: {e}")
|
||||
else:
|
||||
logger.error(f"Error in on_call_tool middleware: {e}", exc_info=True)
|
||||
raise
|
||||
|
||||
|
||||
async def on_get_prompt(self, context: MiddlewareContext, call_next):
|
||||
"""Extract auth info for prompt requests too"""
|
||||
logger.debug("Processing prompt authentication")
|
||||
|
||||
|
||||
try:
|
||||
await self._process_request_for_auth(context)
|
||||
|
||||
|
||||
logger.debug("Passing prompt to next handler")
|
||||
result = await call_next(context)
|
||||
logger.debug("Prompt handler completed")
|
||||
return result
|
||||
|
||||
|
||||
except Exception as e:
|
||||
# Check if this is an authentication error - don't log traceback for these
|
||||
if "GoogleAuthenticationError" in str(type(e)) or "Access denied: Cannot retrieve credentials" in str(e):
|
||||
if "GoogleAuthenticationError" in str(
|
||||
type(e)
|
||||
) or "Access denied: Cannot retrieve credentials" in str(e):
|
||||
logger.info(f"Authentication check failed in prompt: {e}")
|
||||
else:
|
||||
logger.error(f"Error in on_get_prompt middleware: {e}", exc_info=True)
|
||||
|
||||
Reference in New Issue
Block a user