CVE-2024-0520: Relative Path Traversal (CWE-23)
value: A vulnerability in mlflow/mlflow version 8.2.1 allows for remote code execution due to improper neutralization of special elements used in an OS command ('Command Injection') within the `mlflow.data.http_dataset_source.py` module. Specifically, when loading a dataset from a source URL with an HTTP scheme, the filename extracted from the `Content-Disposition` header or the URL path is used to generate the final file path without proper sanitization. This flaw enables an attacker to control the file path fully by utilizing path traversal or absolute path techniques, such as '../../tmp/poc.txt' or '/tmp/poc.txt', leading to arbitrary file write. Exploiting this vulnerability could allow a malicious user to execute commands on the vulnerable machine, potentially gaining access to data and model information. The issue is fixed in version 2.9.0.
CWE description: The product uses external input to construct a pathname that should be within a restricted directory, but it does not properly neutralize sequences such as ".." that can resolve to a location that is outside of that directory.
Published: 2024-06-06T19:15Z
Repository: https://github.com/mlflow/mlflow
Language: Python
Files changed: mlflow/data/http_dataset_source.py, tests/data/test_http_dataset_source.py
Commit date: 2023-12-04 18:52:51 +0900
Commit message: Sanitize filename in Content-Disposition (#10584)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Diff: diff --git a/mlflow/data/http_dataset_source.py b/mlflow/data/http_dataset_source.py
index cd87b3d1b3..9492f9888d 100644
--- a/mlflow/data/http_dataset_source.py
+++ b/mlflow/data/http_dataset_source.py
@@ -9,10 +9,18 @@ from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.utils.file_utils import create_tmp_dir
 from mlflow.utils.rest_utils import augmented_raise_for_status, cloud_storage_http_request
 
 
+def _is_path(filename: str) -> bool:
+    """
+    Return True if `filename` is a path, False otherwise. For example,
+    "foo/bar" is a path, but "bar" is not.
+    """
+    return os.path.basename(filename) != filename
+
+
 class HTTPDatasetSource(DatasetSource):
     """
     Represents the source of a dataset stored at a web location and referred to
     by an HTTP or HTTPS URL.
     """
@@ -55,10 +63,15 @@ class HTTPDatasetSource(DatasetSource):
         if content_disposition is not None and (
             file_name := next(re.finditer(r"filename=(.+)", content_disposition), None)
         ):
             # NB: If the filename is quoted, unquote it
             basename = file_name[1].strip("'\"")
+            if _is_path(basename):
+                raise MlflowException.invalid_parameter_value(
+                    f"Invalid filename in Content-Disposition header: {basename}. "
+                    "It must be a file name, not a path."
+                )
         elif path is not None and len(posixpath.basename(path)) > 0:
             basename = posixpath.basename(path)
         else:
             basename = "dataset_source"
 
diff --git a/tests/data/test_http_dataset_source.py b/tests/data/test_http_dataset_source.py
index 77c33001b5..b7b4a14dbd 100644
--- a/tests/data/test_http_dataset_source.py
+++ b/tests/data/test_http_dataset_source.py
@@ -126,5 +126,32 @@ def test_source_load_with_content_disposition_header(attachment_filename, expect
         )
         source.load()
         loaded = source.load()
         assert os.path.exists(loaded)
         assert os.path.basename(loaded) == expected_filename
+
+
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "/foo/bar.txt",
+        "./foo/bar.txt",
+        "../foo/bar.txt",
+        "foo/bar.txt",
+    ],
+)
+def test_source_load_with_content_disposition_header_invalid_filename(filename):
+    def download_with_mock_content_disposition_headers(*args, **kwargs):
+        response = cloud_storage_http_request(*args, **kwargs)
+        response.headers["Content-Disposition"] = f"attachment; filename={filename}"
+        return response
+
+    with mock.patch(
+        "mlflow.data.http_dataset_source.cloud_storage_http_request",
+        side_effect=download_with_mock_content_disposition_headers,
+    ):
+        source = HTTPDatasetSource(
+            "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
+        )
+
+        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
+            source.load()

Fixed code: def _is_path(filename: str) -> bool:
    """
    Return True if `filename` is a path, False otherwise. For example,
    "foo/bar" is a path, but "bar" is not.
    """
    return os.path.basename(filename) != filename


            if _is_path(basename):
                raise MlflowException.invalid_parameter_value(
                    f"Invalid filename in Content-Disposition header: {basename}. "
                    "It must be a file name, not a path."
                )


@pytest.mark.parametrize(
    "filename",
    [
        "/foo/bar.txt",
        "./foo/bar.txt",
        "../foo/bar.txt",
        "foo/bar.txt",
    ],
)
def test_source_load_with_content_disposition_header_invalid_filename(filename):
    def download_with_mock_content_disposition_headers(*args, **kwargs):
        response = cloud_storage_http_request(*args, **kwargs)
        response.headers["Content-Disposition"] = f"attachment; filename={filename}"
        return response

    with mock.patch(
        "mlflow.data.http_dataset_source.cloud_storage_http_request",
        side_effect=download_with_mock_content_disposition_headers,
    ):
        source = HTTPDatasetSource(
            "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
        )

        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
            source.load()