From dc404fdaef7f031bd6c6d7fe901398190b98247f Mon Sep 17 00:00:00 2001 From: Taylor Wilsdon Date: Tue, 12 Aug 2025 08:45:24 -0400 Subject: [PATCH] oauth cleanup --- auth/oauth_common_handlers.py | 21 ++++++------ core/server.py | 61 +++++++++++++---------------------- 2 files changed, 32 insertions(+), 50 deletions(-) diff --git a/auth/oauth_common_handlers.py b/auth/oauth_common_handlers.py index e532b64..f9ab4ea 100644 --- a/auth/oauth_common_handlers.py +++ b/auth/oauth_common_handlers.py @@ -30,7 +30,7 @@ logger = logging.getLogger(__name__) async def handle_oauth_authorize(request: Request): """Common handler for OAuth authorization proxy.""" origin = request.headers.get("origin") - + if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) return JSONResponse(content={}, headers=cors_headers) @@ -69,7 +69,7 @@ async def handle_oauth_authorize(request: Request): async def handle_proxy_token_exchange(request: Request): """Common handler for OAuth token exchange proxy with comprehensive error handling.""" origin = request.headers.get("origin") - + if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) return JSONResponse(content={}, headers=cors_headers) @@ -200,7 +200,7 @@ async def handle_proxy_token_exchange(request: Request): "Cache-Control": "no-store" } response_headers.update(cors_headers) - + return JSONResponse( status_code=response.status, content=response_data, @@ -227,7 +227,7 @@ async def handle_oauth_protected_resource(request: Request): Handle OAuth protected resource metadata requests. """ origin = request.headers.get("origin") - + # Handle preflight if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) @@ -261,7 +261,7 @@ async def handle_oauth_protected_resource(request: Request): "Cache-Control": "public, max-age=3600" } response_headers.update(cors_headers) - + return JSONResponse( content=metadata, headers=response_headers @@ -273,13 +273,13 @@ async def handle_oauth_authorization_server(request: Request): Handle OAuth authorization server metadata. """ origin = request.headers.get("origin") - + if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) return JSONResponse(content={}, headers=cors_headers) config = get_oauth_config() - + # Get authorization server metadata from centralized config # Pass scopes directly to keep all metadata generation in one place metadata = config.get_authorization_server_metadata(scopes=get_current_scopes()) @@ -293,7 +293,7 @@ async def handle_oauth_authorization_server(request: Request): "Cache-Control": "public, max-age=3600" } response_headers.update(cors_headers) - + return JSONResponse( content=metadata, headers=response_headers @@ -303,7 +303,7 @@ async def handle_oauth_authorization_server(request: Request): async def handle_oauth_client_config(request: Request): """Common handler for OAuth client configuration.""" origin = request.headers.get("origin") - + if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) return JSONResponse(content={}, headers=cors_headers) @@ -327,7 +327,6 @@ async def handle_oauth_client_config(request: Request): "client_uri": config.base_url, "redirect_uris": [ f"{config.base_url}/oauth2callback", - "http://localhost:5173/auth/callback" ], "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], @@ -346,7 +345,7 @@ async def handle_oauth_client_config(request: Request): async def handle_oauth_register(request: Request): """Common handler for OAuth dynamic client registration with comprehensive error handling.""" origin = request.headers.get("origin") - + if request.method == "OPTIONS": cors_headers = get_development_cors_headers(origin) return JSONResponse(content={}, headers=cors_headers) diff --git a/core/server.py b/core/server.py index 9885093..463e023 100644 --- a/core/server.py +++ b/core/server.py @@ -24,7 +24,6 @@ from core.config import ( get_oauth_redirect_uri as get_oauth_redirect_uri_for_current_mode, ) -# Try to import GoogleRemoteAuthProvider for FastMCP 2.11.1+ try: from auth.google_remote_auth_provider import GoogleRemoteAuthProvider GOOGLE_REMOTE_AUTH_AVAILABLE = True @@ -37,7 +36,6 @@ logger = logging.getLogger(__name__) _auth_provider: Optional[Union[GoogleWorkspaceAuthProvider, GoogleRemoteAuthProvider]] = None -# --- Middleware Definitions --- session_middleware = Middleware(MCPSessionMiddleware) # Custom FastMCP that adds secure middleware stack for OAuth 2.1 @@ -55,7 +53,6 @@ class SecureFastMCP(FastMCP): logger.info("Added middleware stack: Session Management") return app -# --- Server Instance --- server = SecureFastMCP( name="google_workspace", auth=None, @@ -86,7 +83,7 @@ def configure_server_for_http(): # Use centralized OAuth configuration from auth.oauth_config import get_oauth_config config = get_oauth_config() - + # Check if OAuth 2.1 is enabled via centralized config oauth21_enabled = config.is_oauth21_enabled() @@ -95,17 +92,23 @@ def configure_server_for_http(): logger.warning("⚠️ OAuth 2.1 enabled but OAuth credentials not configured") return - if GOOGLE_REMOTE_AUTH_AVAILABLE: - logger.info("🔐 OAuth 2.1 enabled with automatic OAuth 2.0 fallback for legacy clients") - try: - _auth_provider = GoogleRemoteAuthProvider() - server.auth = _auth_provider - set_auth_provider(_auth_provider) - logger.debug("OAuth 2.1 authentication enabled") - except Exception as e: - logger.error(f"Failed to initialize GoogleRemoteAuthProvider: {e}", exc_info=True) - else: - logger.error("OAuth 2.1 is enabled, but GoogleRemoteAuthProvider is not available.") + if not GOOGLE_REMOTE_AUTH_AVAILABLE: + logger.error("CRITICAL: OAuth 2.1 enabled but FastMCP 2.11.1+ is not properly installed.") + logger.error("Please run: uv sync --frozen") + raise RuntimeError( + "OAuth 2.1 requires FastMCP 2.11.1+ with RemoteAuthProvider support. " + "Please reinstall dependencies using 'uv sync --frozen'." + ) + + logger.info("🔐 OAuth 2.1 enabled with automatic OAuth 2.0 fallback for legacy clients") + try: + _auth_provider = GoogleRemoteAuthProvider() + server.auth = _auth_provider + set_auth_provider(_auth_provider) + logger.debug("OAuth 2.1 authentication enabled") + except Exception as e: + logger.error(f"Failed to initialize GoogleRemoteAuthProvider: {e}", exc_info=True) + raise else: logger.info("OAuth 2.0 mode - Server will use legacy authentication.") server.auth = None @@ -114,7 +117,6 @@ def get_auth_provider() -> Optional[Union[GoogleWorkspaceAuthProvider, GoogleRem """Gets the global authentication provider instance.""" return _auth_provider -# --- Custom Routes --- @server.custom_route("/health", methods=["GET"]) async def health_check(request: Request): try: @@ -187,20 +189,19 @@ async def oauth2_callback(request: Request) -> HTMLResponse: logger.error(f"Error processing OAuth callback: {str(e)}", exc_info=True) return create_server_error_response(str(e)) -# --- Tools --- @server.tool() async def start_google_auth(service_name: str, user_google_email: str = USER_GOOGLE_EMAIL) -> str: """ Manually initiate Google OAuth authentication flow. - - NOTE: This tool should typically NOT be called directly. The authentication system + + NOTE: This tool should typically NOT be called directly. The authentication system automatically handles credential checks and prompts for authentication when needed. Only use this tool if: 1. You need to re-authenticate with different credentials 2. You want to proactively authenticate before using other tools 3. The automatic authentication flow failed and you need to retry - - In most cases, simply try calling the Google Workspace tool you need - it will + + In most cases, simply try calling the Google Workspace tool you need - it will automatically handle authentication if required. """ if not user_google_email: @@ -221,21 +222,3 @@ async def start_google_auth(service_name: str, user_google_email: str = USER_GOO logger.error(f"Failed to start Google authentication flow: {e}", exc_info=True) return f"**Error:** An unexpected error occurred: {e}" -# OAuth 2.1 Discovery Endpoints - register manually when OAuth 2.1 is enabled but GoogleRemoteAuthProvider is not available -# These will only be registered if MCP_ENABLE_OAUTH21=true and we're in fallback mode -if os.getenv("MCP_ENABLE_OAUTH21", "false").lower() == "true" and not GOOGLE_REMOTE_AUTH_AVAILABLE: - from auth.oauth_common_handlers import ( - handle_oauth_authorize, - handle_proxy_token_exchange, - handle_oauth_protected_resource, - handle_oauth_authorization_server, - handle_oauth_client_config, - handle_oauth_register - ) - - server.custom_route("/.well-known/oauth-protected-resource", methods=["GET", "OPTIONS"])(handle_oauth_protected_resource) - server.custom_route("/.well-known/oauth-authorization-server", methods=["GET", "OPTIONS"])(handle_oauth_authorization_server) - server.custom_route("/.well-known/oauth-client", methods=["GET", "OPTIONS"])(handle_oauth_client_config) - server.custom_route("/oauth2/authorize", methods=["GET", "OPTIONS"])(handle_oauth_authorize) - server.custom_route("/oauth2/token", methods=["POST", "OPTIONS"])(handle_proxy_token_exchange) - server.custom_route("/oauth2/register", methods=["POST", "OPTIONS"])(handle_oauth_register)