CVE-2023-6568: Improper Neutralization of Input During Web Page Generation ('Cross-site Scripting') (CWE-79)
value: A reflected Cross-Site Scripting (XSS) vulnerability exists in the mlflow/mlflow repository, specifically within the handling of the Content-Type header in POST requests. An attacker can inject malicious JavaScript code into the Content-Type header, which is then improperly reflected back to the user without adequate sanitization or escaping, leading to arbitrary JavaScript execution in the context of the victim's browser. The vulnerability is present in the mlflow/server/auth/__init__.py file, where the user-supplied Content-Type header is directly injected into a Python formatted string and returned to the user, facilitating the XSS attack.
CWE description: The product does not neutralize or incorrectly neutralizes user-controllable input before it is placed in output that is used as a web page that is served to other users.
CVSS v3 base score: 6.1
Published: 2023-12-07T05:15Z
Repository: https://github.com/mlflow/mlflow
Language: Python
Files changed: docs/source/auth/rest-api.rst, docs/source/rest-api.rst, mlflow/server/auth/__init__.py, mlflow/server/handlers.py, mlflow/server/validation.py, tests/server/test_handlers.py
Commit date: 2023-11-30 09:38:30 +0900
Commit message: Add Content-Type header validation (#10526)

Signed-off-by: B-Step62 <yuki.watanabe@databricks.com>
Diff: diff --git a/docs/source/auth/rest-api.rst b/docs/source/auth/rest-api.rst
index baf6d2779..2bbba4b7d 100644
--- a/docs/source/auth/rest-api.rst
+++ b/docs/source/auth/rest-api.rst
@@ -10,10 +10,12 @@ The MLflow Authentication REST API allows you to create, get, update and delete
 experiment permissions and registered model permissions.
 The API is hosted under the ``/api`` route on the MLflow tracking server. For example, to list
 experiments on a tracking server hosted at ``http://localhost:5000``, access
 ``http://localhost:5000/api/2.0/mlflow/users/create``.
 
+.. important::
+    The MLflow REST API requires content type ``application/json`` for all POST requests.
 
 .. contents:: Table of Contents
     :local:
     :depth: 1
 
diff --git a/docs/source/rest-api.rst b/docs/source/rest-api.rst
index bbf97c4ac..5989a6c73 100755
--- a/docs/source/rest-api.rst
+++ b/docs/source/rest-api.rst
@@ -9,18 +9,20 @@ REST API
 The MLflow REST API allows you to create, list, and get experiments and runs, and log parameters, metrics, and artifacts.
 The API is hosted under the ``/api`` route on the MLflow tracking server. For example, to search for
 experiments on a tracking server hosted at ``http://localhost:5000``, make a POST request to
 ``http://localhost:5000/api/2.0/mlflow/experiments/search``.
 
+.. important::
+    The MLflow REST API requires content type ``application/json`` for all POST requests.
+
 .. contents:: Table of Contents
     :local:
     :depth: 1
 
 ===========================
 
 
-
 .. _mlflowMlflowServicecreateExperiment:
 
 Create Experiment
 =================
 
diff --git a/mlflow/server/auth/__init__.py b/mlflow/server/auth/__init__.py
index e281ad750..6c46569f9 100644
--- a/mlflow/server/auth/__init__.py
+++ b/mlflow/server/auth/__init__.py
@@ -760,11 +760,15 @@ def create_user():
         password = _get_request_param("password")
 
         user = store.create_user(username, password)
         return make_response({"user": user.to_json()})
     else:
-        return make_response(f"Invalid content type: '{content_type}'", 400)
+        message = (
+            "Invalid content type. Must be one of: "
+            "application/x-www-form-urlencoded, application/json"
+        )
+        return make_response(message, 400)
 
 
 @catch_mlflow_exception
 def get_user():
     username = _get_request_param("username")
diff --git a/mlflow/server/handlers.py b/mlflow/server/handlers.py
index e15ee3653..cb2961c36 100644
--- a/mlflow/server/handlers.py
+++ b/mlflow/server/handlers.py
@@ -89,10 +89,11 @@ from mlflow.protos.service_pb2 import (
     SetExperimentTag,
     SetTag,
     UpdateExperiment,
     UpdateRun,
 )
+from mlflow.server.validation import _validate_content_type
 from mlflow.store.artifact.artifact_repo import MultipartUploadMixin
 from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
 from mlflow.store.db.db_types import DATABASE_ENGINES
 from mlflow.tracking._model_registry import utils as registry_utils
 from mlflow.tracking._model_registry.registry import ModelRegistryStoreRegistry
@@ -401,10 +402,11 @@ def _validate_param_against_schema(schema, param, value, proto_parsing_succeeded
 
     return None
 
 
 def _get_request_json(flask_request=request):
+    _validate_content_type(flask_request, ["application/json"])
     return flask_request.get_json(force=True, silent=True)
 
 
 def _get_request_message(request_message, flask_request=request, schema=None):
     from querystring_parser import parser
@@ -1110,10 +1112,11 @@ def get_metric_history_bulk_handler():
 
 @catch_mlflow_exception
 @_disable_if_artifacts_only
 def search_datasets_handler():
     MAX_EXPERIMENT_IDS_PER_REQUEST = 20
+    _validate_content_type(request, ["application/json"])
     experiment_ids = request.json.get("experiment_ids", [])
     if not experiment_ids:
         raise MlflowException(
             message="SearchDatasets request must specify at least one experiment_id.",
             error_code=INVALID_PARAMETER_VALUE,
@@ -1177,10 +1180,12 @@ def create_promptlab_run_handler():
             raise MlflowException(
                 message=f"CreatePromptlabRun request must specify {arg_name}.",
                 error_code=INVALID_PARAMETER_VALUE,
             )
 
+    _validate_content_type(request, ["application/json"])
+
     args = request.json
     experiment_id = args.get("experiment_id")
     assert_arg_exists("experiment_id", experiment_id)
     run_name = args.get("run_name", None)
     tags = args.get("tags", [])
diff --git a/mlflow/server/validation.py b/mlflow/server/validation.py
new file mode 100644
index 000000000..4c9f64628
--- /dev/null
+++ b/mlflow/server/validation.py
@@ -0,0 +1,31 @@
+from typing import List
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+
+def _validate_content_type(flask_request, allowed_content_types: List[str]):
+    """
+    Validates that the request content type is one of the allowed content types.
+
+    :param flask_request: Flask request object (flask.request)
+    :param allowed_content_types: A list of allowed content types
+    """
+    if flask_request.method not in ["POST", "PUT"]:
+        return
+
+    if flask_request.content_type is None:
+        raise MlflowException(
+            message="Bad Request. Content-Type header is missing.",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    # Remove any parameters e.g. "application/json; charset=utf-8" -> "application/json"
+    content_type = flask_request.content_type.split(";")[0]
+    if content_type not in allowed_content_types:
+        message = f"Bad Request. Content-Type must be one of {allowed_content_types}."
+
+        raise MlflowException(
+            message=message,
+            error_code=INVALID_PARAMETER_VALUE,
+        )
diff --git a/tests/server/test_handlers.py b/tests/server/test_handlers.py
index 3b7632387..1bc5bd047 100644
--- a/tests/server/test_handlers.py
+++ b/tests/server/test_handlers.py
@@ -154,25 +154,37 @@ def test_all_model_registry_endpoints_available():
 
 
 def test_can_parse_json():
     request = mock.MagicMock()
     request.method = "POST"
+    request.content_type = "application/json"
     request.get_json = mock.MagicMock()
     request.get_json.return_value = {"name": "hello"}
     msg = _get_request_message(CreateExperiment(), flask_request=request)
     assert msg.name == "hello"
 
 
 def test_can_parse_post_json_with_unknown_fields():
     request = mock.MagicMock()
     request.method = "POST"
+    request.content_type = "application/json"
     request.get_json = mock.MagicMock()
     request.get_json.return_value = {"name": "hello", "WHAT IS THIS FIELD EVEN": "DOING"}
     msg = _get_request_message(CreateExperiment(), flask_request=request)
     assert msg.name == "hello"
 
 
+def test_can_parse_post_json_with_content_type_params():
+    request = mock.MagicMock()
+    request.method = "POST"
+    request.content_type = "application/json; charset=utf-8"
+    request.get_json = mock.MagicMock()
+    request.get_json.return_value = {"name": "hello"}
+    msg = _get_request_message(CreateExperiment(), flask_request=request)
+    assert msg.name == "hello"
+
+
 def test_can_parse_get_json_with_unknown_fields():
     request = mock.MagicMock()
     request.method = "GET"
     request.query_string = b"name=hello&superDuperUnknown=field"
     msg = _get_request_message(CreateExperiment(), flask_request=request)
@@ -182,16 +194,37 @@ def test_can_parse_get_json_with_unknown_fields():
 # Previous versions of the client sent a doubly string encoded JSON blob,
 # so this test ensures continued compliance with such clients.
 def test_can_parse_json_string():
     request = mock.MagicMock()
     request.method = "POST"
+    request.content_type = "application/json"
     request.get_json = mock.MagicMock()
     request.get_json.return_value = '{"name": "hello2"}'
     msg = _get_request_message(CreateExperiment(), flask_request=request)
     assert msg.name == "hello2"
 
 
+def test_can_block_post_request_with_invalid_content_type():
+    request = mock.MagicMock()
+    request.method = "POST"
+    request.content_type = "text/plain"
+    request.get_json = mock.MagicMock()
+    request.get_json.return_value = {"name": "hello"}
+    with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"):
+        _get_request_message(CreateExperiment(), flask_request=request)
+
+
+def test_can_block_post_request_with_missing_content_type():
+    request = mock.MagicMock()
+    request.method = "POST"
+    request.content_type = None
+    request.get_json = mock.MagicMock()
+    request.get_json.return_value = {"name": "hello"}
+    with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"):
+        _get_request_message(CreateExperiment(), flask_request=request)
+
+
 def test_search_runs_default_view_type(mock_get_request_message, mock_tracking_store):
     """
     Search Runs default view type is filled in as ViewType.ACTIVE_ONLY
     """
     mock_get_request_message.return_value = SearchRuns(experiment_ids=["0"])

Vulnerable code: 
        return make_response(f"Invalid content type: '{content_type}'", 400)
Fixed code: .. important::
    The MLflow REST API requires content type ``application/json`` for all POST requests.
.. important::
    The MLflow REST API requires content type ``application/json`` for all POST requests.

        message = (
            "Invalid content type. Must be one of: "
            "application/x-www-form-urlencoded, application/json"
        )
        return make_response(message, 400)
from mlflow.server.validation import _validate_content_type
    _validate_content_type(flask_request, ["application/json"])
    _validate_content_type(request, ["application/json"])
    _validate_content_type(request, ["application/json"])

from typing import List

from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE


def _validate_content_type(flask_request, allowed_content_types: List[str]):
    """
    Validates that the request content type is one of the allowed content types.

    :param flask_request: Flask request object (flask.request)
    :param allowed_content_types: A list of allowed content types
    """
    if flask_request.method not in ["POST", "PUT"]:
        return

    if flask_request.content_type is None:
        raise MlflowException(
            message="Bad Request. Content-Type header is missing.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    # Remove any parameters e.g. "application/json; charset=utf-8" -> "application/json"
    content_type = flask_request.content_type.split(";")[0]
    if content_type not in allowed_content_types:
        message = f"Bad Request. Content-Type must be one of {allowed_content_types}."

        raise MlflowException(
            message=message,
            error_code=INVALID_PARAMETER_VALUE,
        )
    request.content_type = "application/json"
    request.content_type = "application/json"
def test_can_parse_post_json_with_content_type_params():
    request = mock.MagicMock()
    request.method = "POST"
    request.content_type = "application/json; charset=utf-8"
    request.get_json = mock.MagicMock()
    request.get_json.return_value = {"name": "hello"}
    msg = _get_request_message(CreateExperiment(), flask_request=request)
    assert msg.name == "hello"


    request.content_type = "application/json"
def test_can_block_post_request_with_invalid_content_type():
    request = mock.MagicMock()
    request.method = "POST"
    request.content_type = "text/plain"
    request.get_json = mock.MagicMock()
    request.get_json.return_value = {"name": "hello"}
    with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"):
        _get_request_message(CreateExperiment(), flask_request=request)


def test_can_block_post_request_with_missing_content_type():
    request = mock.MagicMock()
    request.method = "POST"
    request.content_type = None
    request.get_json = mock.MagicMock()
    request.get_json.return_value = {"name": "hello"}
    with pytest.raises(MlflowException, match=r"Bad Request. Content-Type"):
        _get_request_message(CreateExperiment(), flask_request=request)