mlflow--mlflow/mlflow/server/security_utils.py

"""
Shared security utilities for MLflow server middleware.

This module contains common functions used by both Flask and FastAPI
security implementations.
"""

import fnmatch
from urllib.parse import urlparse

from mlflow.environment_variables import (
    MLFLOW_SERVER_ALLOWED_HOSTS,
    MLFLOW_SERVER_CORS_ALLOWED_ORIGINS,
)

# Security response messages
INVALID_HOST_MSG = "Invalid Host header - possible DNS rebinding attack detected"
CORS_BLOCKED_MSG = "Cross-origin request blocked"

# HTTP methods that modify state
STATE_CHANGING_METHODS = ["POST", "PUT", "DELETE", "PATCH"]

# Paths exempt from host validation
HEALTH_ENDPOINTS = ["/health", "/version"]

# API path prefixes for MLflow endpoints
API_PATH_PREFIX = "/api/"
AJAX_API_PATH_PREFIX = "/ajax-api/"

# Test-only endpoints that should not have CORS blocking
TEST_ENDPOINTS = ["/test", "/api/test"]

# Localhost addresses
LOCALHOST_VARIANTS = ["localhost", "127.0.0.1", "[::1]", "0.0.0.0"]
CORS_LOCALHOST_HOSTS = ["localhost", "127.0.0.1", "[::1]", "::1"]

# Private IP range start values for 172.16.0.0/12
PRIVATE_172_RANGE_START = 16
PRIVATE_172_RANGE_END = 32

# Regex patterns for localhost origins
LOCALHOST_ORIGIN_PATTERNS = [
    r"^http://localhost(:[0-9]+)?$",
    r"^http://127\.0\.0\.1(:[0-9]+)?$",
    r"^http://\[::1\](:[0-9]+)?$",
]


def get_localhost_addresses() -> list[str]:
    """Get localhost/loopback addresses."""
    return LOCALHOST_VARIANTS


def get_private_ip_patterns() -> list[str]:
    """
    Generate wildcard patterns for private IP ranges.

    These are the standard RFC-defined private address ranges:
    - RFC 1918 (IPv4): 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16
      https://datatracker.ietf.org/doc/html/rfc1918
    - RFC 4193 (IPv6): fc00::/7
      https://datatracker.ietf.org/doc/html/rfc4193

    Additional references:
    - IANA IPv4 Special-Purpose Address Registry:
      https://www.iana.org/assignments/iana-ipv4-special-registry/
    - IANA IPv6 Special-Purpose Address Registry:
      https://www.iana.org/assignments/iana-ipv6-special-registry/
    """
    return [
        "192.168.*",
        "10.*",
        *[f"172.{i}.*" for i in range(PRIVATE_172_RANGE_START, PRIVATE_172_RANGE_END)],
        "fc00:*",
        "fd00:*",
    ]


def get_allowed_hosts_from_env() -> list[str] | None:
    """Get allowed hosts from environment variable."""
    if allowed_hosts_env := MLFLOW_SERVER_ALLOWED_HOSTS.get():
        return [host.strip() for host in allowed_hosts_env.split(",")]
    return None


def get_allowed_origins_from_env() -> list[str] | None:
    """Get allowed CORS origins from environment variable."""
    if allowed_origins_env := MLFLOW_SERVER_CORS_ALLOWED_ORIGINS.get():
        return [origin.strip() for origin in allowed_origins_env.split(",")]
    return None


def is_localhost_origin(origin: str) -> bool:
    """Check if an origin is from localhost."""
    if not origin:
        return False

    try:
        parsed = urlparse(origin)
        hostname = parsed.hostname
        return hostname in CORS_LOCALHOST_HOSTS
    except Exception:
        return False


def should_block_cors_request(origin: str, method: str, allowed_origins: list[str] | None) -> bool:
    """Determine if a CORS request should be blocked."""
    if not origin or method not in STATE_CHANGING_METHODS:
        return False

    if is_localhost_origin(origin):
        return False

    if allowed_origins:
        # If wildcard "*" is in the list, allow all origins
        if "*" in allowed_origins:
            return False

        return not any(
            fnmatch.fnmatch(origin, allowed) if "*" in allowed else origin == allowed
            for allowed in allowed_origins
        )

    return True


def is_api_endpoint(path: str) -> bool:
    """Check if a path is an API endpoint that should have CORS/OPTIONS handling."""
    return (
        path.startswith(API_PATH_PREFIX) or path.startswith(AJAX_API_PATH_PREFIX)
    ) and path not in TEST_ENDPOINTS


def is_allowed_host_header(allowed_hosts: list[str], host: str) -> bool:
    """Validate if the host header matches allowed patterns."""
    if not host:
        return False

    # If wildcard "*" is in the list, allow all hosts
    if "*" in allowed_hosts:
        return True

    return any(
        fnmatch.fnmatch(host, allowed) if "*" in allowed else host == allowed
        for allowed in allowed_hosts
    )


def get_default_allowed_hosts() -> list[str]:
    """Get default allowed hosts patterns."""
    wildcard_hosts = []
    for host in get_localhost_addresses():
        if host.startswith("["):
            # IPv6: escape opening bracket for fnmatch
            escaped = host.replace("[", "[[]", 1)
            wildcard_hosts.append(f"{escaped}:*")
        else:
            wildcard_hosts.append(f"{host}:*")

    return get_localhost_addresses() + wildcard_hosts + get_private_ip_patterns()