smlflg's picture
Initial public upload from Projekte/Connector (part 16)
bf06a0d verified
Raw History Blame Contribute Delete
30.2 kB
import asyncio
import hashlib
import logging
import os
from typing import List, Optional
from importlib import metadata
from fastapi.responses import HTMLResponse, JSONResponse, FileResponse
from starlette.applications import Starlette
from starlette.datastructures import MutableHeaders
from starlette.types import Scope, Receive, Send
from starlette.requests import Request
from starlette.middleware import Middleware
from fastmcp import FastMCP
from fastmcp.server.auth.providers.google import GoogleProvider
from auth.oauth21_session_store import get_oauth21_session_store, set_auth_provider
from auth.google_auth import handle_auth_callback, start_auth_flow, check_client_secrets
from auth.oauth_config import is_oauth21_enabled, is_external_oauth21_provider
from auth.mcp_session_middleware import MCPSessionMiddleware
from auth.oauth_responses import (
create_error_response,
create_success_response,
create_server_error_response,
)
from auth.auth_info_middleware import AuthInfoMiddleware
from auth.scopes import BASE_SCOPES, SCOPES, get_current_scopes # noqa
from core.config import (
USER_GOOGLE_EMAIL,
get_transport_mode,
set_transport_mode as _set_transport_mode,
get_oauth_redirect_uri as get_oauth_redirect_uri_for_current_mode,
)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
_auth_provider: Optional[GoogleProvider] = None
_legacy_callback_registered = False
session_middleware = Middleware(MCPSessionMiddleware)
class WellKnownCacheControlMiddleware:
"""Force no-cache headers for OAuth well-known discovery endpoints."""
def __init__(self, app):
self.app = app
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
path = scope.get("path", "")
is_oauth_well_known = (
path == "/.well-known/oauth-authorization-server"
or path.startswith("/.well-known/oauth-authorization-server/")
or path == "/.well-known/oauth-protected-resource"
or path.startswith("/.well-known/oauth-protected-resource/")
)
if not is_oauth_well_known:
await self.app(scope, receive, send)
return
async def send_with_no_cache_headers(message):
if message["type"] == "http.response.start":
headers = MutableHeaders(raw=message.setdefault("headers", []))
headers["Cache-Control"] = "no-store, must-revalidate"
headers["ETag"] = f'"{_compute_scope_fingerprint()}"'
await send(message)
await self.app(scope, receive, send_with_no_cache_headers)
well_known_cache_control_middleware = Middleware(WellKnownCacheControlMiddleware)
def _compute_scope_fingerprint() -> str:
"""Compute a short hash of the current scope configuration for cache-busting."""
scopes_str = ",".join(sorted(get_current_scopes()))
return hashlib.sha256(scopes_str.encode()).hexdigest()[:12]
# Custom FastMCP that adds secure middleware stack for OAuth 2.1
class SecureFastMCP(FastMCP):
def http_app(self, **kwargs) -> "Starlette":
"""Override to add secure middleware stack for OAuth 2.1."""
app = super().http_app(**kwargs)
# Add middleware in order (first added = outermost layer)
app.user_middleware.insert(0, well_known_cache_control_middleware)
# Session Management - extracts session info for MCP context
app.user_middleware.insert(1, session_middleware)
# Rebuild middleware stack
app.middleware_stack = app.build_middleware_stack()
logger.info("Added middleware stack: WellKnownCacheControl, Session Management")
return app
async def list_tools(self, *, run_middleware: bool = True):
"""Override to mark user_google_email as optional when USER_GOOGLE_EMAIL is set.
In single-user / self-hosted mode the env var provides the default email, so
callers (agents, MCP adapters) should not be required to supply it. We patch
the JSON schema returned by list_tools to remove 'user_google_email' from the
``required`` array and inject the env-var value as the ``default``. The
runtime still resolves the email correctly via the service decorator.
"""
tools = list(await super().list_tools(run_middleware=run_middleware))
if not USER_GOOGLE_EMAIL or is_oauth21_enabled():
return tools
patched = []
for tool in tools:
schema = dict(tool.parameters)
required = list(schema.get("required", []))
if "user_google_email" in required:
required = [r for r in required if r != "user_google_email"]
props = {k: dict(v) for k, v in schema.get("properties", {}).items()}
if "user_google_email" in props:
props["user_google_email"]["default"] = USER_GOOGLE_EMAIL
schema = dict(schema, required=required, properties=props)
patched.append(tool.model_copy(update={"parameters": schema}))
else:
patched.append(tool)
return patched
async def call_tool(self, name: str, arguments: Optional[dict], *args, **kwargs):
"""Inject user_google_email before pydantic validates the call arguments.
When USER_GOOGLE_EMAIL is configured and OAuth 2.1 is not active, callers
(agents, adapters) are allowed to omit user_google_email. FastMCP validates
arguments against the function signature BEFORE calling the tool, so we must
inject the default BEFORE that validation step.
"""
arguments = arguments or {}
if (
not is_oauth21_enabled()
and USER_GOOGLE_EMAIL
and "user_google_email" not in arguments
):
arguments = {**arguments, "user_google_email": USER_GOOGLE_EMAIL}
return await super().call_tool(name, arguments, *args, **kwargs)
# Build server instructions with user email context for single-user mode
_server_instructions = None
if USER_GOOGLE_EMAIL:
_server_instructions = f"""Connected Google account: {USER_GOOGLE_EMAIL}
When using Google Workspace tools, always use `{USER_GOOGLE_EMAIL}` as the `user_google_email` parameter. Do not ask the user for their email address."""
logger.info(f"Server instructions configured for user: {USER_GOOGLE_EMAIL}")
server = SecureFastMCP(
name="google_workspace",
auth=None,
instructions=_server_instructions,
)
# Add the AuthInfo middleware to inject authentication into FastMCP context
auth_info_middleware = AuthInfoMiddleware()
server.add_middleware(auth_info_middleware)
def _parse_bool_env(value: str) -> bool:
"""Parse environment variable string to boolean."""
return value.lower() in ("1", "true", "yes", "on")
def _parse_allowed_redirect_uris(value: Optional[str]) -> Optional[List[str]]:
"""Parse a comma-separated list of OAuth client redirect URIs.
Returns a list of non-empty, trimmed URIs, or None when the input is
empty/None. Returning None preserves FastMCP's default behaviour of
accepting any client-supplied redirect URI during DCR — callers that
want to lock down registration must supply a non-empty list.
Patterns supported by FastMCP's matcher (see
``fastmcp.server.auth.redirect_validation``) include ``*`` for port
and path globs (e.g. ``http://localhost:*/callback``) and ``*.example.com``
for subdomain wildcards.
"""
if not value:
return None
uris = [u.strip() for u in value.split(",") if u.strip()]
return uris or None
def set_transport_mode(mode: str):
"""Sets the transport mode for the server."""
_set_transport_mode(mode)
logger.info(f"Transport: {mode}")
def _ensure_legacy_callback_route() -> None:
global _legacy_callback_registered
if _legacy_callback_registered:
return
server.custom_route("/oauth2callback", methods=["GET"])(legacy_oauth2_callback)
_legacy_callback_registered = True
def configure_server_for_http():
"""
Configures the authentication provider for HTTP transport.
This must be called BEFORE server.run().
"""
global _auth_provider
transport_mode = get_transport_mode()
if transport_mode != "streamable-http":
return
# 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()
if oauth21_enabled:
if not config.is_configured():
logger.warning("OAuth 2.1 enabled but OAuth credentials not configured")
return
def validate_and_derive_jwt_key(
jwt_signing_key_override: str | None, client_secret: str | None
) -> bytes:
"""Validate JWT signing key override and derive the final JWT key."""
if jwt_signing_key_override:
if len(jwt_signing_key_override) < 12:
logger.warning(
"OAuth 2.1: FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY is less than 12 characters; "
"use a longer secret to improve key derivation strength."
)
return derive_jwt_key(
low_entropy_material=jwt_signing_key_override,
salt="fastmcp-jwt-signing-key",
)
if client_secret:
return derive_jwt_key(
high_entropy_material=client_secret,
salt="fastmcp-jwt-signing-key",
)
raise ValueError(
"Public client OAuth 2.1 requires FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY "
"when GOOGLE_OAUTH_CLIENT_SECRET is not set."
)
try:
# Import common dependencies for storage backends
from key_value.aio.wrappers.encryption import FernetEncryptionWrapper
from cryptography.fernet import Fernet
from fastmcp.server.auth.jwt_issuer import derive_jwt_key
provider_valid_scopes: List[str] = sorted(get_current_scopes())
provider_required_scopes: List[str] = sorted(BASE_SCOPES)
client_storage = None
jwt_signing_key_override = (
os.getenv("FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY", "").strip()
or None
)
storage_backend = (
os.getenv("WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND", "")
.strip()
.lower()
)
valkey_host = os.getenv("WORKSPACE_MCP_OAUTH_PROXY_VALKEY_HOST", "").strip()
# Determine storage backend: valkey, disk, memory (default)
use_valkey = storage_backend == "valkey" or bool(valkey_host)
use_disk = storage_backend == "disk"
if use_valkey:
try:
from key_value.aio.stores.valkey import ValkeyStore
valkey_port_raw = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_PORT", "6379"
).strip()
valkey_db_raw = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_DB", "0"
).strip()
valkey_port = int(valkey_port_raw)
valkey_db = int(valkey_db_raw)
valkey_use_tls_raw = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_USE_TLS", ""
).strip()
valkey_use_tls = (
_parse_bool_env(valkey_use_tls_raw)
if valkey_use_tls_raw
else valkey_port == 6380
)
valkey_request_timeout_ms_raw = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_REQUEST_TIMEOUT_MS", ""
).strip()
valkey_connection_timeout_ms_raw = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_CONNECTION_TIMEOUT_MS", ""
).strip()
valkey_request_timeout_ms = (
int(valkey_request_timeout_ms_raw)
if valkey_request_timeout_ms_raw
else None
)
valkey_connection_timeout_ms = (
int(valkey_connection_timeout_ms_raw)
if valkey_connection_timeout_ms_raw
else None
)
valkey_username = (
os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_USERNAME", ""
).strip()
or None
)
valkey_password = (
os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_VALKEY_PASSWORD", ""
).strip()
or None
)
if not valkey_host:
valkey_host = "localhost"
client_storage = ValkeyStore(
host=valkey_host,
port=valkey_port,
db=valkey_db,
username=valkey_username,
password=valkey_password,
)
# Configure TLS and timeouts on the underlying Glide client config.
# ValkeyStore currently doesn't expose these settings directly.
glide_config = getattr(client_storage, "_client_config", None)
if glide_config is not None:
glide_config.use_tls = valkey_use_tls
is_remote_host = valkey_host not in {"localhost", "127.0.0.1"}
if valkey_request_timeout_ms is None and (
valkey_use_tls or is_remote_host
):
# Glide defaults to 250ms if unset; increase for remote/TLS endpoints.
valkey_request_timeout_ms = 5000
if valkey_request_timeout_ms is not None:
glide_config.request_timeout = valkey_request_timeout_ms
if valkey_connection_timeout_ms is None and (
valkey_use_tls or is_remote_host
):
valkey_connection_timeout_ms = 10000
if valkey_connection_timeout_ms is not None:
from glide_shared.config import (
AdvancedGlideClientConfiguration,
)
glide_config.advanced_config = (
AdvancedGlideClientConfiguration(
connection_timeout=valkey_connection_timeout_ms
)
)
jwt_signing_key = validate_and_derive_jwt_key(
jwt_signing_key_override, config.client_secret
)
storage_encryption_key = derive_jwt_key(
high_entropy_material=jwt_signing_key.decode(),
salt="fastmcp-storage-encryption-key",
)
client_storage = FernetEncryptionWrapper(
key_value=client_storage,
fernet=Fernet(key=storage_encryption_key),
)
logger.info(
"OAuth 2.1: Using ValkeyStore for FastMCP OAuth proxy client_storage (host=%s, port=%s, db=%s, tls=%s)",
valkey_host,
valkey_port,
valkey_db,
valkey_use_tls,
)
if valkey_request_timeout_ms is not None:
logger.info(
"OAuth 2.1: Valkey request timeout set to %sms",
valkey_request_timeout_ms,
)
if valkey_connection_timeout_ms is not None:
logger.info(
"OAuth 2.1: Valkey connection timeout set to %sms",
valkey_connection_timeout_ms,
)
logger.info(
"OAuth 2.1: Applied Fernet encryption wrapper to Valkey client_storage (key derived from FASTMCP_SERVER_AUTH_GOOGLE_JWT_SIGNING_KEY or GOOGLE_OAUTH_CLIENT_SECRET)."
)
except ImportError as exc:
logger.warning(
"OAuth 2.1: Valkey client_storage requested but Valkey dependencies are not installed (%s). "
"Install 'workspace-mcp[valkey]' (or 'py-key-value-aio[valkey]', which includes 'valkey-glide') "
"or unset WORKSPACE_MCP_OAUTH_PROXY_STORAGE_BACKEND/WORKSPACE_MCP_OAUTH_PROXY_VALKEY_HOST.",
exc,
)
except ValueError as exc:
logger.warning(
"OAuth 2.1: Invalid Valkey configuration; falling back to default storage (%s).",
exc,
)
elif use_disk:
try:
from core.storage import make_sanitized_file_store
disk_directory = os.getenv(
"WORKSPACE_MCP_OAUTH_PROXY_DISK_DIRECTORY", ""
).strip()
if not disk_directory:
# Default to FASTMCP_HOME/oauth-proxy or ~/.fastmcp/oauth-proxy
fastmcp_home = os.getenv("FASTMCP_HOME", "").strip()
if fastmcp_home:
disk_directory = os.path.join(fastmcp_home, "oauth-proxy")
else:
disk_directory = os.path.expanduser(
"~/.fastmcp/oauth-proxy"
)
client_storage = make_sanitized_file_store(disk_directory)
jwt_signing_key = validate_and_derive_jwt_key(
jwt_signing_key_override, config.client_secret
)
storage_encryption_key = derive_jwt_key(
high_entropy_material=jwt_signing_key.decode(),
salt="fastmcp-storage-encryption-key",
)
client_storage = FernetEncryptionWrapper(
key_value=client_storage,
fernet=Fernet(key=storage_encryption_key),
)
logger.info(
"OAuth 2.1: Using FileTreeStore for FastMCP OAuth proxy client_storage (directory=%s)",
disk_directory,
)
except ImportError as exc:
logger.warning(
"OAuth 2.1: Disk storage requested but dependencies not available (%s). "
"Falling back to default storage.",
exc,
)
elif storage_backend == "memory":
from key_value.aio.stores.memory import MemoryStore
client_storage = MemoryStore()
logger.info(
"OAuth 2.1: Using MemoryStore for FastMCP OAuth proxy client_storage"
)
# else: client_storage remains None, FastMCP uses its default
# Ensure JWT signing key is always derived for all storage backends
if "jwt_signing_key" not in locals():
jwt_signing_key = validate_and_derive_jwt_key(
jwt_signing_key_override, config.client_secret
)
# Check if external OAuth provider is configured
if config.is_external_oauth21_provider():
# External OAuth mode: use custom provider that handles ya29.* access tokens
from auth.external_oauth_provider import ExternalOAuthProvider
provider = ExternalOAuthProvider(
client_id=config.client_id,
client_secret=config.client_secret,
base_url=config.get_oauth_base_url(),
redirect_path=config.redirect_path,
required_scopes=provider_valid_scopes,
resource_server_url=config.get_oauth_base_url(),
)
server.auth = provider
logger.info("OAuth 2.1 enabled with EXTERNAL provider mode")
logger.info(
"Expecting Authorization bearer tokens in tool call headers"
)
logger.info(
"Protected resource metadata points to Google's authorization server"
)
else:
# Standard OAuth 2.1 mode: use FastMCP's GoogleProvider
allowed_client_redirect_uris = _parse_allowed_redirect_uris(
os.getenv("WORKSPACE_MCP_ALLOWED_CLIENT_REDIRECT_URIS")
)
if allowed_client_redirect_uris:
logger.info(
"OAuth 2.1: restricting DCR client redirect URIs to allowlist: %s",
allowed_client_redirect_uris,
)
provider = GoogleProvider(
client_id=config.client_id,
client_secret=config.client_secret,
base_url=config.get_oauth_base_url(),
redirect_path=config.redirect_path,
required_scopes=provider_required_scopes,
valid_scopes=provider_valid_scopes,
client_storage=client_storage,
jwt_signing_key=jwt_signing_key,
allowed_client_redirect_uris=allowed_client_redirect_uris,
)
if provider.client_registration_options is not None:
# Keep protocol-level auth limited to base identity scopes, but
# allow dynamically registered MCP clients to request any scope
# needed by enabled tools during subsequent authorization flows.
provider.client_registration_options.default_scopes = (
provider_valid_scopes
)
# Enable protocol-level auth
server.auth = provider
logger.info(
"OAuth 2.1 enabled using FastMCP GoogleProvider with protocol-level auth"
)
# Always set auth provider for token validation in middleware
set_auth_provider(provider)
_auth_provider = provider
except Exception as exc:
logger.error(
"Failed to initialize FastMCP GoogleProvider: %s", exc, exc_info=True
)
raise
else:
logger.info("OAuth 2.0 mode - Server will use legacy authentication.")
server.auth = None
_auth_provider = None
set_auth_provider(None)
_ensure_legacy_callback_route()
def get_auth_provider() -> Optional[GoogleProvider]:
"""Gets the global authentication provider instance."""
return _auth_provider
@server.custom_route("/", methods=["GET"])
@server.custom_route("/health", methods=["GET"])
async def health_check(request: Request):
try:
version = metadata.version("workspace-mcp")
except metadata.PackageNotFoundError:
version = "dev"
return JSONResponse(
{
"status": "healthy",
"service": "workspace-mcp",
"version": version,
"transport": get_transport_mode(),
}
)
@server.custom_route("/attachments/{file_id}", methods=["GET"])
async def serve_attachment(request: Request):
"""Serve a stored attachment file."""
from core.attachment_storage import get_attachment_storage
file_id = request.path_params["file_id"]
storage = get_attachment_storage()
metadata = storage.get_attachment_metadata(file_id)
if not metadata:
return JSONResponse(
{"error": "Attachment not found or expired"}, status_code=404
)
file_path = storage.get_attachment_path(file_id)
if not file_path:
return JSONResponse({"error": "Attachment file not found"}, status_code=404)
return FileResponse(
path=str(file_path),
filename=metadata["filename"],
media_type=metadata["mime_type"],
)
async def legacy_oauth2_callback(request: Request) -> HTMLResponse:
state = request.query_params.get("state")
code = request.query_params.get("code")
error = request.query_params.get("error")
if error:
msg = (
f"Authentication failed: Google returned an error: {error}. State: {state}."
)
logger.error(msg)
return create_error_response(msg)
if not code:
msg = "Authentication failed: No authorization code received from Google."
logger.error(msg)
return create_error_response(msg)
try:
error_message = check_client_secrets()
if error_message:
return create_server_error_response(error_message)
logger.info("OAuth callback: Received authorization code.")
mcp_session_id = None
if hasattr(request, "state") and hasattr(request.state, "session_id"):
mcp_session_id = request.state.session_id
verified_user_id, credentials = await handle_auth_callback(
scopes=get_current_scopes(),
authorization_response=str(request.url),
redirect_uri=get_oauth_redirect_uri_for_current_mode(),
session_id=mcp_session_id,
)
logger.info(
f"OAuth callback: Successfully authenticated user: {verified_user_id}."
)
try:
store = get_oauth21_session_store()
store.store_session(
user_email=verified_user_id,
access_token=credentials.token,
refresh_token=credentials.refresh_token,
token_uri=credentials.token_uri,
client_id=credentials.client_id,
client_secret=credentials.client_secret,
scopes=credentials.scopes,
expiry=credentials.expiry,
session_id=f"google-{state}",
mcp_session_id=mcp_session_id,
)
logger.info(
f"Stored Google credentials in OAuth 2.1 session store for {verified_user_id}"
)
except Exception as e:
logger.error(f"Failed to store credentials in OAuth 2.1 store: {e}")
return create_success_response(verified_user_id)
except Exception as e:
logger.error(f"Error processing OAuth callback: {str(e)}", exc_info=True)
return create_server_error_response(str(e))
@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 is a legacy OAuth 2.0 tool and is disabled when OAuth 2.1 is enabled.
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
automatically handle authentication if required.
"""
if is_oauth21_enabled():
if is_external_oauth21_provider():
return (
"start_google_auth is disabled when OAuth 2.1 is enabled. "
"Provide a valid OAuth 2.1 bearer token in the Authorization header "
"and retry the original tool."
)
return (
"start_google_auth is disabled when OAuth 2.1 is enabled. "
"Authenticate through your MCP client's OAuth 2.1 flow and retry the "
"original tool."
)
if not user_google_email:
raise ValueError("user_google_email must be provided.")
error_message = check_client_secrets()
if error_message:
return f"**Authentication Error:** {error_message}"
try:
transport_mode = get_transport_mode()
if transport_mode == "stdio":
# Only stdio legacy OAuth depends on the standalone callback server.
from auth.oauth_callback_server import ensure_oauth_callback_available
from auth.oauth_config import get_oauth_config
config = get_oauth_config()
success, error_msg = await asyncio.to_thread(
ensure_oauth_callback_available,
transport_mode,
config.port,
config.base_uri,
)
if not success:
error_detail = f" ({error_msg})" if error_msg else ""
return f"**Error:** Cannot initiate OAuth flow - callback server unavailable{error_detail}"
auth_message = await start_auth_flow(
user_google_email=user_google_email,
service_name=service_name,
redirect_uri=get_oauth_redirect_uri_for_current_mode(),
)
return auth_message
except Exception as e:
logger.error(f"Failed to start Google authentication flow: {e}", exc_info=True)
return f"**Error:** An unexpected error occurred: {e}"