import logging from http import HTTPStatus from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from starlette.types import ASGIApp from mlflow.environment_variables import ( MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE, MLFLOW_SERVER_X_FRAME_OPTIONS, ) from mlflow.server.security_utils import ( CORS_BLOCKED_MSG, HEALTH_ENDPOINTS, INVALID_HOST_MSG, LOCALHOST_ORIGIN_PATTERNS, get_allowed_hosts_from_env, get_allowed_origins_from_env, get_default_allowed_hosts, is_allowed_host_header, is_api_endpoint, should_block_cors_request, ) from mlflow.tracing.constant import TRACE_RENDERER_ASSET_PATH _logger = logging.getLogger(__name__) class HostValidationMiddleware: """Middleware to validate Host headers using fnmatch patterns.""" def __init__(self, app: ASGIApp, allowed_hosts: list[str]): self.app = app self.allowed_hosts = allowed_hosts async def __call__(self, scope, receive, send): if scope["type"] != "http": return await self.app(scope, receive, send) if scope["path"] in HEALTH_ENDPOINTS: return await self.app(scope, receive, send) headers = dict(scope.get("headers", [])) host = headers.get(b"host", b"").decode("utf-8") if not is_allowed_host_header(self.allowed_hosts, host): _logger.warning(f"Rejected request with invalid Host header: {host}") async def send_403(message): if message["type"] == "http.response.start": message["status"] = 403 message["headers"] = [(b"content-type", b"text/plain")] await send(message) await send_403({"type": "http.response.start", "status": 403, "headers": []}) await send({"type": "http.response.body", "body": INVALID_HOST_MSG.encode()}) return return await self.app(scope, receive, send) class SecurityHeadersMiddleware: """Middleware to add security headers to all responses.""" def __init__(self, app: ASGIApp): self.app = app self.x_frame_options = MLFLOW_SERVER_X_FRAME_OPTIONS.get() async def __call__(self, scope, receive, send): if scope["type"] != "http": return await self.app(scope, receive, send) async def send_wrapper(message): if message["type"] == "http.response.start": headers = dict(message.get("headers", [])) headers[b"x-content-type-options"] = b"nosniff" # Skip X-Frame-Options for notebook renderer to allow iframe embedding in notebooks path = scope.get("path", "") is_notebook_renderer = path.startswith(TRACE_RENDERER_ASSET_PATH) if ( self.x_frame_options and self.x_frame_options.upper() != "NONE" and not is_notebook_renderer ): headers[b"x-frame-options"] = self.x_frame_options.upper().encode() if ( scope["method"] == "OPTIONS" and message.get("status") == 200 and is_api_endpoint(scope["path"]) ): message["status"] = HTTPStatus.NO_CONTENT message["headers"] = list(headers.items()) await send(message) await self.app(scope, receive, send_wrapper) class CORSBlockingMiddleware: """Middleware to actively block cross-origin state-changing requests.""" def __init__(self, app: ASGIApp, allowed_origins: list[str]): self.app = app self.allowed_origins = allowed_origins async def __call__(self, scope, receive, send): if scope["type"] != "http": return await self.app(scope, receive, send) if not is_api_endpoint(scope["path"]): return await self.app(scope, receive, send) method = scope["method"] headers = dict(scope["headers"]) origin = headers.get(b"origin", b"").decode("utf-8") if should_block_cors_request(origin, method, self.allowed_origins): _logger.warning(f"Blocked cross-origin request from {origin}") await send({ "type": "http.response.start", "status": HTTPStatus.FORBIDDEN, "headers": [[b"content-type", b"text/plain"]], }) await send({ "type": "http.response.body", "body": CORS_BLOCKED_MSG.encode(), }) return await self.app(scope, receive, send) def get_allowed_hosts() -> list[str]: """Get list of allowed hosts from environment or defaults.""" return get_allowed_hosts_from_env() or get_default_allowed_hosts() def get_allowed_origins() -> list[str]: """Get list of allowed CORS origins from environment or defaults.""" return get_allowed_origins_from_env() or [] def init_fastapi_security(app: FastAPI) -> None: """ Initialize security middleware for FastAPI application. This configures: - Host header validation (DNS rebinding protection) via TrustedHostMiddleware - CORS protection via CORSMiddleware - Security headers via custom middleware Args: app: FastAPI application instance. """ if MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE.get() == "true": return app.add_middleware(SecurityHeadersMiddleware) allowed_origins = get_allowed_origins() if allowed_origins and "*" in allowed_origins: _logger.warning( "MLFLOW_SERVER_CORS_ALLOWED_ORIGINS=* is set; disabling credentialed CORS. " "Wildcard origins with credentials is a CORS spec violation. " "Set an explicit origin allowlist to enable credentialed cross-origin requests." ) # CORSBlockingMiddleware is intentionally skipped here: the admin explicitly # opted into wildcard CORS, so server-side rejection of cross-origin requests # would defeat the purpose. The browser-side guardrail is that # allow_credentials=False causes browsers to strip cookies/Authorization # headers on cross-origin requests, so authenticated endpoints still 401. app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=False, allow_methods=["*"], allow_headers=["*"], expose_headers=["*"], ) else: # Use CORSBlockingMiddleware for blocking CORS requests on the server side, # and CORSMiddleware for responding to OPTIONS requests. app.add_middleware(CORSBlockingMiddleware, allowed_origins=allowed_origins) app.add_middleware( CORSMiddleware, allow_origins=allowed_origins, allow_origin_regex="|".join(LOCALHOST_ORIGIN_PATTERNS), allow_credentials=True, allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "PATCH"], allow_headers=["*"], expose_headers=["*"], ) allowed_hosts = get_allowed_hosts() if allowed_hosts and "*" not in allowed_hosts: app.add_middleware(HostValidationMiddleware, allowed_hosts=allowed_hosts)