CVE-2023-6753: Improper Limitation of a Pathname to a Restricted Directory ('Path Traversal') (CWE-22)
value: Path Traversal in GitHub repository mlflow/mlflow prior to 2.9.2.
CWE description: The product uses external input to construct a pathname that is intended to identify a file or directory that is located underneath a restricted parent directory, but the product does not properly neutralize special elements within the pathname that can cause the pathname to resolve to a location that is outside of the restricted directory.
CVSS v3 base score: 8.8
Published: 2023-12-13T00: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-08 10:15:35 -0500
Commit message: Fix for windows traversal attack (#10647)

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
Signed-off-by: Ben Wilson <39283302+BenWilson2@users.noreply.github.com>
Co-authored-by: Harutaka Kawamura <hkawamura0130@gmail.com>
Diff: diff --git a/mlflow/data/http_dataset_source.py b/mlflow/data/http_dataset_source.py
index 9492f9888d..71da81ae8f 100644
--- a/mlflow/data/http_dataset_source.py
+++ b/mlflow/data/http_dataset_source.py
@@ -1,7 +1,6 @@
 import os
-import posixpath
 import re
 from typing import Any, Dict
 from urllib.parse import urlparse
 
 from mlflow.data.dataset_source import DatasetSource
@@ -39,10 +38,27 @@ class HTTPDatasetSource(DatasetSource):
 
     @staticmethod
     def _get_source_type() -> str:
         return "http"
 
+    def _extract_filename(self, response) -> str:
+        """
+        Extracts a filename from the Content-Disposition header or the URL's path.
+        """
+        if content_disposition := response.headers.get("Content-Disposition"):
+            for match in re.finditer(r"filename=(.+)", content_disposition):
+                filename = match[1].strip("'\"")
+                if _is_path(filename):
+                    raise MlflowException.invalid_parameter_value(
+                        f"Invalid filename in Content-Disposition header: {filename}. "
+                        "It must be a file name, not a path."
+                    )
+                return filename
+
+        # Extract basename from URL if no valid filename in Content-Disposition
+        return os.path.basename(urlparse(self.url).path)
+
     def load(self, dst_path=None) -> str:
         """
         Downloads the dataset source to the local filesystem.
 
         :param dst_path: Path of the local filesystem destination directory to which to download the
@@ -56,25 +72,13 @@ class HTTPDatasetSource(DatasetSource):
             url=self.url,
             stream=True,
         )
         augmented_raise_for_status(resp)
 
-        path = urlparse(self.url).path
-        content_disposition = resp.headers.get("Content-Disposition")
-        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 = self._extract_filename(resp)
+
+        if not basename:
             basename = "dataset_source"
 
         if dst_path is None:
             dst_path = create_tmp_dir()
 
diff --git a/tests/data/test_http_dataset_source.py b/tests/data/test_http_dataset_source.py
index b7b4a14dbd..fb0705a7e6 100644
--- a/tests/data/test_http_dataset_source.py
+++ b/tests/data/test_http_dataset_source.py
@@ -6,10 +6,11 @@ import pandas as pd
 import pytest
 
 from mlflow.data.dataset_source_registry import get_dataset_source_from_json, resolve_dataset_source
 from mlflow.data.http_dataset_source import HTTPDatasetSource
 from mlflow.exceptions import MlflowException
+from mlflow.utils.os import is_windows
 from mlflow.utils.rest_utils import cloud_storage_http_request
 
 
 def test_source_to_and_from_json():
     url = "http://mywebsite.com/path/to/my/dataset.txt"
@@ -153,5 +154,32 @@ def test_source_load_with_content_disposition_header_invalid_filename(filename):
             "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()
+
+
+@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
+@pytest.mark.parametrize(
+    "filename",
+    [
+        r"..\..\poc.txt",
+        r"Users\User\poc.txt",
+    ],
+)
+def test_source_load_with_content_disposition_header_invalid_filename_windows(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"
+        )
+
+        # Expect an MlflowException for invalid filenames
+        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
+            source.load()

Vulnerable code: import posixpath
        path = urlparse(self.url).path
        content_disposition = resp.headers.get("Content-Disposition")
        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:
Fixed code:     def _extract_filename(self, response) -> str:
        """
        Extracts a filename from the Content-Disposition header or the URL's path.
        """
        if content_disposition := response.headers.get("Content-Disposition"):
            for match in re.finditer(r"filename=(.+)", content_disposition):
                filename = match[1].strip("'\"")
                if _is_path(filename):
                    raise MlflowException.invalid_parameter_value(
                        f"Invalid filename in Content-Disposition header: {filename}. "
                        "It must be a file name, not a path."
                    )
                return filename

        # Extract basename from URL if no valid filename in Content-Disposition
        return os.path.basename(urlparse(self.url).path)

        basename = self._extract_filename(resp)

        if not basename:
from mlflow.utils.os import is_windows


@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
@pytest.mark.parametrize(
    "filename",
    [
        r"..\..\poc.txt",
        r"Users\User\poc.txt",
    ],
)
def test_source_load_with_content_disposition_header_invalid_filename_windows(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"
        )

        # Expect an MlflowException for invalid filenames
        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
            source.load()