chore: import upstream snapshot with attribution

tests/server/test_security.py

@@ -0,0 +1,421 @@
+import logging
+
+import pytest
+from fastapi import FastAPI
+from flask import Flask
+from starlette.testclient import TestClient
+from werkzeug.test import Client
+
+from mlflow.server import security
+from mlflow.server.fastapi_security import init_fastapi_security
+from mlflow.server.security_utils import is_allowed_host_header, is_api_endpoint
+
+
+def test_default_allowed_hosts():
+    hosts = security.get_allowed_hosts()
+    assert "localhost" in hosts
+    assert "127.0.0.1" in hosts
+    assert "[::1]" in hosts
+    assert "localhost:*" in hosts
+    assert "127.0.0.1:*" in hosts
+    assert "[[]::1]:*" in hosts
+    assert "192.168.*" in hosts
+    assert "10.*" in hosts
+
+
+def test_custom_allowed_hosts(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("MLFLOW_SERVER_ALLOWED_HOSTS", "example.com,app.example.com")
+    hosts = security.get_allowed_hosts()
+    assert "example.com" in hosts
+    assert "app.example.com" in hosts
+
+
+@pytest.mark.parametrize(
+    ("host_header", "expected_status", "expected_error"),
+    [
+        ("localhost", 200, None),
+        ("127.0.0.1", 200, None),
+        ("evil.attacker.com", 403, b"Invalid Host header"),
+    ],
+)
+def test_dns_rebinding_protection(
+    test_app, host_header, expected_status, expected_error, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("MLFLOW_SERVER_ALLOWED_HOSTS", "localhost,127.0.0.1")
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.get("/test", headers={"Host": host_header})
+    assert response.status_code == expected_status
+    if expected_error:
+        assert expected_error in response.data
+
+
+@pytest.mark.parametrize(
+    ("method", "origin", "expected_status", "expected_cors_header"),
+    [
+        ("POST", "http://localhost:3000", 200, "http://localhost:3000"),
+        ("POST", "http://evil.com", 403, None),
+        ("POST", None, 200, None),
+        ("GET", "http://evil.com", 200, None),
+    ],
+)
+def test_cors_protection(
+    test_app, method, origin, expected_status, expected_cors_header, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv(
+        "MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", "http://localhost:3000,https://app.example.com"
+    )
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    headers = {"Origin": origin} if origin else {}
+    response = getattr(client, method.lower())("/api/2.0/mlflow/experiments/list", headers=headers)
+    assert response.status_code == expected_status
+
+    if expected_cors_header:
+        assert response.headers.get("Access-Control-Allow-Origin") == expected_cors_header
+
+
+def test_wildcard_cors_disables_credentials(
+    test_app, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+    monkeypatch.setenv("MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", "*")
+    # The "mlflow" logger sets propagate=False, so caplog's root handler does not
+    # see records. Attach caplog's handler directly to the security logger.
+    security_logger = logging.getLogger("mlflow.server.security")
+    security_logger.addHandler(caplog.handler)
+    try:
+        with caplog.at_level("WARNING", logger="mlflow.server.security"):
+            security.init_security_middleware(test_app)
+    finally:
+        security_logger.removeHandler(caplog.handler)
+    assert any("disabling credentialed CORS" in record.message for record in caplog.records)
+
+    client = Client(test_app)
+    response = client.post(
+        "/api/2.0/mlflow/experiments/list", headers={"Origin": "http://evil.com"}
+    )
+    assert response.status_code == 200
+    # The load-bearing security guarantee: no Access-Control-Allow-Credentials: true.
+    # Without that header, browsers strip cookies/Authorization on cross-origin
+    # requests, so an attacker page cannot ride the victim's authenticated session.
+    assert response.headers.get("Access-Control-Allow-Credentials") != "true"
+
+
+@pytest.mark.parametrize(
+    ("origin", "expected_cors_header"),
+    [
+        ("http://localhost:3000", "http://localhost:3000"),
+        ("http://evil.com", None),
+    ],
+)
+def test_preflight_options_request(
+    test_app, origin, expected_cors_header, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", "http://localhost:3000")
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.options(
+        "/api/2.0/mlflow/experiments/list",
+        headers={
+            "Origin": origin,
+            "Access-Control-Request-Method": "POST",
+            "Access-Control-Request-Headers": "Content-Type",
+        },
+    )
+    assert response.status_code == 204
+
+    if expected_cors_header:
+        assert response.headers.get("Access-Control-Allow-Origin") == expected_cors_header
+
+
+def test_security_headers(test_app):
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.get("/test")
+    assert response.headers.get("X-Content-Type-Options") == "nosniff"
+    assert response.headers.get("X-Frame-Options") == "SAMEORIGIN"
+
+
+def test_disable_security_middleware(test_app, monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE", "true")
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.get("/test")
+    assert "X-Content-Type-Options" not in response.headers
+    assert "X-Frame-Options" not in response.headers
+
+    response = client.get("/test", headers={"Host": "evil.com"})
+    assert response.status_code == 200
+
+
+def test_x_frame_options_configuration(monkeypatch: pytest.MonkeyPatch):
+    app = Flask(__name__)
+
+    @app.route("/test")
+    def test():
+        return "OK"
+
+    monkeypatch.setenv("MLFLOW_SERVER_X_FRAME_OPTIONS", "DENY")
+    security.init_security_middleware(app)
+    client = Client(app)
+    response = client.get("/test")
+    assert response.headers.get("X-Frame-Options") == "DENY"
+
+    app2 = Flask(__name__)
+
+    @app2.route("/test")
+    def test2():
+        return "OK"
+
+    # Reset for the second app
+    monkeypatch.setenv("MLFLOW_SERVER_X_FRAME_OPTIONS", "NONE")
+    security.init_security_middleware(app2)
+    client = Client(app2)
+    response = client.get("/test")
+    assert "X-Frame-Options" not in response.headers
+
+
+def test_notebook_trace_renderer_skips_x_frame_options(monkeypatch: pytest.MonkeyPatch):
+    from mlflow.tracing.constant import TRACE_RENDERER_ASSET_PATH
+
+    app = Flask(__name__)
+
+    @app.route(f"{TRACE_RENDERER_ASSET_PATH}/index.html")
+    def notebook_renderer():
+        return "<html>trace renderer</html>"
+
+    @app.route(f"{TRACE_RENDERER_ASSET_PATH}/js/main.js")
+    def notebook_renderer_js():
+        return "console.log('trace renderer');"
+
+    @app.route("/static-files/other-page.html")
+    def other_page():
+        return "<html>other page</html>"
+
+    # Set X-Frame-Options to DENY to test that it's skipped for notebook renderer
+    monkeypatch.setenv("MLFLOW_SERVER_X_FRAME_OPTIONS", "DENY")
+    security.init_security_middleware(app)
+    client = Client(app)
+
+    response = client.get(f"{TRACE_RENDERER_ASSET_PATH}/index.html")
+    assert response.status_code == 200
+    assert "X-Frame-Options" not in response.headers
+
+    response = client.get(f"{TRACE_RENDERER_ASSET_PATH}/js/main.js")
+    assert response.status_code == 200
+    assert "X-Frame-Options" not in response.headers
+
+    response = client.get("/static-files/other-page.html")
+    assert response.status_code == 200
+    assert response.headers.get("X-Frame-Options") == "DENY"
+
+
+@pytest.mark.parametrize(
+    ("allowed_hosts", "host_header", "expected_status"),
+    [
+        ("*", "any.domain.com", 200),
+        ("*.example.com", "app.example.com", 200),
+        ("*.example.com", "sub.app.example.com", 200),
+        ("*.example.com", "evil.com", 403),
+    ],
+)
+def test_wildcard_hosts(
+    test_app, allowed_hosts, host_header, expected_status, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("MLFLOW_SERVER_ALLOWED_HOSTS", allowed_hosts)
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.get("/test", headers={"Host": host_header})
+    assert response.status_code == expected_status
+
+
+@pytest.mark.parametrize(
+    ("allowed_origins", "origin", "expected_status"),
+    [
+        ("*", "http://any.domain.com", 200),
+        ("http://*.example.com", "http://app.example.com", 200),
+        ("http://*.example.com", "http://sub.app.example.com", 200),
+        ("http://*.example.com", "http://evil.com", 403),
+    ],
+)
+def test_wildcard_origins(
+    test_app, allowed_origins, origin, expected_status, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", allowed_origins)
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.post("/api/2.0/mlflow/experiments/list", headers={"Origin": origin})
+    assert response.status_code == expected_status
+
+
+@pytest.mark.parametrize(
+    ("endpoint", "host_header", "expected_status"),
+    [
+        ("/health", "evil.com", 200),
+        ("/test", "evil.com", 403),
+    ],
+)
+def test_endpoint_security_bypass(
+    test_app, endpoint, host_header, expected_status, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("MLFLOW_SERVER_ALLOWED_HOSTS", "localhost")
+    security.init_security_middleware(test_app)
+    client = Client(test_app)
+
+    response = client.get(endpoint, headers={"Host": host_header})
+    assert response.status_code == expected_status
+
+
+@pytest.mark.parametrize(
+    ("hostname", "expected_valid"),
+    [
+        ("192.168.1.1", True),
+        ("10.0.0.1", True),
+        ("172.16.0.1", True),
+        ("127.0.0.1", True),
+        ("localhost", True),
+        ("[::1]", True),
+        ("192.168.1.1:8080", True),
+        ("[::1]:8080", True),
+        ("evil.com", False),
+    ],
+)
+def test_host_validation(hostname, expected_valid):
+    hosts = security.get_allowed_hosts()
+    assert is_allowed_host_header(hosts, hostname) == expected_valid
+
+
+@pytest.mark.parametrize(
+    ("env_var", "env_value", "expected_result"),
+    [
+        (
+            "MLFLOW_SERVER_CORS_ALLOWED_ORIGINS",
+            "http://app1.com,http://app2.com",
+            ["http://app1.com", "http://app2.com"],
+        ),
+        ("MLFLOW_SERVER_ALLOWED_HOSTS", "app1.com,app2.com:8080", ["app1.com", "app2.com:8080"]),
+    ],
+)
+def test_environment_variable_configuration(
+    env_var, env_value, expected_result, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv(env_var, env_value)
+    if "ORIGINS" in env_var:
+        result = security.get_allowed_origins()
+        for expected in expected_result:
+            assert expected in result
+    else:
+        result = security.get_allowed_hosts()
+        for expected in expected_result:
+            assert expected in result
+
+
+@pytest.mark.parametrize(
+    ("path", "expected"),
+    [
+        ("/api/2.0/mlflow/experiments/list", True),
+        ("/ajax-api/2.0/mlflow/experiments/list", True),
+        ("/ajax-api/3.0/mlflow/runs/search", True),
+        ("/api/test", False),
+        ("/test", False),
+        ("/health", False),
+        ("/static/index.html", False),
+    ],
+)
+def test_is_api_endpoint(path, expected):
+    assert is_api_endpoint(path) == expected
+
+
+@pytest.mark.parametrize(
+    ("origin", "expect_cors_header"),
+    [
+        ("http://localhost:3000", True),
+        ("http://127.0.0.1:5000", True),
+        ("http://[::1]:8080", True),
+        ("http://evil.com", False),
+    ],
+)
+def test_fastapi_cors_allows_localhost_origins(fastapi_client, origin, expect_cors_header):
+    response = fastapi_client.get(
+        "/api/2.0/mlflow/experiments/list", headers={"Host": "localhost", "Origin": origin}
+    )
+    if expect_cors_header:
+        assert response.headers.get("access-control-allow-origin") == origin
+    else:
+        assert response.headers.get("access-control-allow-origin") is None
+
+
+def test_fastapi_wildcard_cors_disables_credentials(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+):
+    monkeypatch.setenv("MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", "*")
+    # The "mlflow" logger sets propagate=False, so caplog's root handler does not
+    # see records. Attach caplog's handler directly to the security logger.
+    security_logger = logging.getLogger("mlflow.server.fastapi_security")
+    security_logger.addHandler(caplog.handler)
+
+    app = FastAPI()
+
+    @app.api_route("/api/2.0/mlflow/experiments/list", methods=["GET", "POST", "OPTIONS"])
+    async def api_endpoint():
+        return {"ok": True}
+
+    try:
+        with caplog.at_level("WARNING", logger="mlflow.server.fastapi_security"):
+            init_fastapi_security(app)
+    finally:
+        security_logger.removeHandler(caplog.handler)
+    assert any("disabling credentialed CORS" in record.message for record in caplog.records)
+
+    client = TestClient(app, raise_server_exceptions=False)
+    response = client.post(
+        "/api/2.0/mlflow/experiments/list",
+        headers={"Host": "localhost", "Origin": "http://evil.com"},
+    )
+    assert response.status_code == 200
+    assert response.headers.get("access-control-allow-credentials") != "true"
+
+    # Browsers read Access-Control-Allow-Credentials from the preflight (OPTIONS)
+    # response before deciding whether to send a credentialed request. Verify the
+    # preflight does not advertise credential support either.
+    preflight = client.options(
+        "/api/2.0/mlflow/experiments/list",
+        headers={
+            "Host": "localhost",
+            "Origin": "http://evil.com",
+            "Access-Control-Request-Method": "POST",
+        },
+    )
+    assert preflight.headers.get("access-control-allow-credentials") != "true"
+
+
+def test_fastapi_cors_allows_configured_origin(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setenv("MLFLOW_SERVER_CORS_ALLOWED_ORIGINS", "https://trusted.com")
+
+    app = FastAPI()
+
+    @app.api_route("/api/2.0/mlflow/experiments/list", methods=["GET", "POST", "OPTIONS"])
+    async def api_endpoint():
+        return {"ok": True}
+
+    init_fastapi_security(app)
+    client = TestClient(app, raise_server_exceptions=False)
+
+    response = client.get(
+        "/api/2.0/mlflow/experiments/list",
+        headers={"Host": "localhost", "Origin": "https://trusted.com"},
+    )
+    assert response.headers.get("access-control-allow-origin") == "https://trusted.com"
+
+    response = client.get(
+        "/api/2.0/mlflow/experiments/list",
+        headers={"Host": "localhost", "Origin": "http://evil.com"},
+    )
+    assert response.headers.get("access-control-allow-origin") is None