CVE-2023-6976: Unrestricted Upload of File with Dangerous Type (CWE-434)
value: This vulnerability is capable of writing arbitrary files into arbitrary locations on the remote filesystem in the context of the server process.
CWE description: The product allows the upload or transfer of dangerous file types that are automatically processed within its environment.
CVSS v3 base score: 8.8
Published: 2023-12-20T06:15Z
Repository: https://github.com/mlflow/mlflow
Language: Python
Files changed: mlflow/server/handlers.py, mlflow/store/artifact/http_artifact_repo.py, mlflow/utils/uri.py
Commit date: 2023-12-13 11:47:56 +0900
Commit message: Use validated path after path validation  (#10666)

Signed-off-by: harupy <17039389+harupy@users.noreply.github.com>
Diff: diff --git a/mlflow/server/handlers.py b/mlflow/server/handlers.py
index cd7f7397ad..d41e4db378 100644
--- a/mlflow/server/handlers.py
+++ b/mlflow/server/handlers.py
@@ -546,11 +546,11 @@ def get_artifact_handler():
 
     query_string = request.query_string.decode("utf-8")
     request_dict = parser.parse(query_string, normalized=True)
     run_id = request_dict.get("run_id") or request_dict.get("run_uuid")
     path = request_dict["path"]
-    validate_path_is_safe(path)
+    path = validate_path_is_safe(path)
     run = _get_tracking_store().get_run(run_id)
 
     if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
         artifact_repo = _get_artifact_repo_mlflow_artifacts()
         artifact_path = _get_proxied_run_artifact_destination_path(
@@ -939,11 +939,11 @@ def _list_artifacts():
         },
     )
     response_message = ListArtifacts.Response()
     if request_message.HasField("path"):
         path = request_message.path
-        validate_path_is_safe(path)
+        path = validate_path_is_safe(path)
     else:
         path = None
     run_id = request_message.run_id or request_message.run_uuid
     run = _get_tracking_store().get_run(run_id)
 
@@ -1234,11 +1234,11 @@ def upload_artifact_handler():
     if not path:
         raise MlflowException(
             message="Request must specify path.",
             error_code=INVALID_PARAMETER_VALUE,
         )
-    validate_path_is_safe(path)
+    path = validate_path_is_safe(path)
 
     if request.content_length and request.content_length > 10 * 1024 * 1024:
         raise MlflowException(
             message="Artifact size is too large. Max size is 10MB.",
             error_code=INVALID_PARAMETER_VALUE,
@@ -1658,11 +1658,11 @@ def get_model_version_artifact_handler():
     query_string = request.query_string.decode("utf-8")
     request_dict = parser.parse(query_string, normalized=True)
     name = request_dict.get("name")
     version = request_dict.get("version")
     path = request_dict["path"]
-    validate_path_is_safe(path)
+    path = validate_path_is_safe(path)
     artifact_uri = _get_model_registry_store().get_model_version_download_uri(name, version)
     if _is_servable_proxied_run_artifact_root(artifact_uri):
         artifact_repo = _get_artifact_repo_mlflow_artifacts()
         artifact_path = _get_proxied_run_artifact_destination_path(
             proxied_artifact_root=artifact_uri,
@@ -1884,11 +1884,11 @@ def _get_model_version_by_alias():
 def _download_artifact(artifact_path):
     """
     A request handler for `GET /mlflow-artifacts/artifacts/<artifact_path>` to download an artifact
     from `artifact_path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
     tmp_dir = tempfile.TemporaryDirectory()
     artifact_repo = _get_artifact_repo_mlflow_artifacts()
     dst = artifact_repo.download_artifacts(artifact_path, tmp_dir.name)
 
     # Ref: https://stackoverflow.com/a/24613980/6943581
@@ -1909,11 +1909,11 @@ def _download_artifact(artifact_path):
 def _upload_artifact(artifact_path):
     """
     A request handler for `PUT /mlflow-artifacts/artifacts/<artifact_path>` to upload an artifact
     to `artifact_path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
     head, tail = posixpath.split(artifact_path)
     with tempfile.TemporaryDirectory() as tmp_dir:
         tmp_path = os.path.join(tmp_dir, tail)
         with open(tmp_path, "wb") as f:
             chunk_size = 1024 * 1024  # 1 MB
@@ -1935,15 +1935,11 @@ def _list_artifacts_mlflow_artifacts():
     """
     A request handler for `GET /mlflow-artifacts/artifacts?path=<value>` to list artifacts in `path`
     (a relative path from the root artifact directory).
     """
     request_message = _get_request_message(ListArtifactsMlflowArtifacts())
-    if request_message.HasField("path"):
-        validate_path_is_safe(request_message.path)
-        path = request_message.path
-    else:
-        path = None
+    path = validate_path_is_safe(request_message.path) if request_message.HasField("path") else None
     artifact_repo = _get_artifact_repo_mlflow_artifacts()
     files = []
     for file_info in artifact_repo.list_artifacts(path):
         basename = posixpath.basename(file_info.path)
         new_file_info = FileInfo(basename, file_info.is_dir, file_info.file_size)
@@ -1960,11 +1956,11 @@ def _list_artifacts_mlflow_artifacts():
 def _delete_artifact_mlflow_artifacts(artifact_path):
     """
     A request handler for `DELETE /mlflow-artifacts/artifacts?path=<value>` to delete artifacts in
     `path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
     _get_request_message(DeleteArtifact())
     artifact_repo = _get_artifact_repo_mlflow_artifacts()
     artifact_repo.delete_artifacts(artifact_path)
     response_message = DeleteArtifact.Response()
     response = Response(mimetype="application/json")
@@ -1982,11 +1978,11 @@ def _validate_support_multipart_upload(artifact_repo):
 def _create_multipart_upload_artifact(artifact_path):
     """
     A request handler for `POST /mlflow-artifacts/mpu/create` to create a multipart upload
     to `artifact_path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
 
     request_message = _get_request_message(
         CreateMultipartUpload(),
         schema={
             "path": [_assert_required, _assert_string],
@@ -2015,11 +2011,11 @@ def _create_multipart_upload_artifact(artifact_path):
 def _complete_multipart_upload_artifact(artifact_path):
     """
     A request handler for `POST /mlflow-artifacts/mpu/complete` to complete a multipart upload
     to `artifact_path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
 
     request_message = _get_request_message(
         CompleteMultipartUpload(),
         schema={
             "path": [_assert_required, _assert_string],
@@ -2048,11 +2044,11 @@ def _complete_multipart_upload_artifact(artifact_path):
 def _abort_multipart_upload_artifact(artifact_path):
     """
     A request handler for `POST /mlflow-artifacts/mpu/abort` to abort a multipart upload
     to `artifact_path` (a relative path from the root artifact directory).
     """
-    validate_path_is_safe(artifact_path)
+    artifact_path = validate_path_is_safe(artifact_path)
 
     request_message = _get_request_message(
         AbortMultipartUpload(),
         schema={
             "path": [_assert_required, _assert_string],
diff --git a/mlflow/store/artifact/http_artifact_repo.py b/mlflow/store/artifact/http_artifact_repo.py
index 5f349aedc0..a44b37824f 100644
--- a/mlflow/store/artifact/http_artifact_repo.py
+++ b/mlflow/store/artifact/http_artifact_repo.py
@@ -87,13 +87,13 @@ class HttpArtifactRepository(ArtifactRepository, MultipartUploadMixin):
         host_creds = _get_default_host_creds(url)
         resp = http_request(host_creds, endpoint, "GET", params=params)
         augmented_raise_for_status(resp)
         file_infos = []
         for f in resp.json().get("files", []):
-            validate_path_is_safe(f["path"])
+            validated_path = validate_path_is_safe(f["path"])
             file_info = FileInfo(
-                posixpath.join(path, f["path"]) if path else f["path"],
+                posixpath.join(path, validated_path) if path else validated_path,
                 f["is_dir"],
                 int(f["file_size"]) if ("file_size" in f) else None,
             )
             file_infos.append(file_info)
 
diff --git a/mlflow/utils/uri.py b/mlflow/utils/uri.py
index b1a526bed5..8d54df5994 100644
--- a/mlflow/utils/uri.py
+++ b/mlflow/utils/uri.py
@@ -451,10 +451,12 @@ def validate_path_is_safe(path):
         or pathlib.PurePosixPath(path).is_absolute()
         or (is_windows() and len(path) >= 2 and path[1] == ":")
     ):
         raise exc
 
+    return path
+
 
 def validate_query_string(query):
     query = _decode(query)
     # Block query strings contain any traveral path (../) because they
     # could be resolved as part of the path and allow path traversal.

Vulnerable code:     validate_path_is_safe(path)
        validate_path_is_safe(path)
    validate_path_is_safe(path)
    validate_path_is_safe(path)
    validate_path_is_safe(artifact_path)
    validate_path_is_safe(artifact_path)
    if request_message.HasField("path"):
        validate_path_is_safe(request_message.path)
        path = request_message.path
    else:
        path = None
    validate_path_is_safe(artifact_path)
    validate_path_is_safe(artifact_path)
    validate_path_is_safe(artifact_path)
    validate_path_is_safe(artifact_path)
            validate_path_is_safe(f["path"])
                posixpath.join(path, f["path"]) if path else f["path"],
Fixed code:     path = validate_path_is_safe(path)
        path = validate_path_is_safe(path)
    path = validate_path_is_safe(path)
    path = validate_path_is_safe(path)
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = validate_path_is_safe(artifact_path)
    path = validate_path_is_safe(request_message.path) if request_message.HasField("path") else None
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = validate_path_is_safe(artifact_path)
            validated_path = validate_path_is_safe(f["path"])
                posixpath.join(path, validated_path) if path else validated_path,
    return path