Repository: mlflow/mlflow
Issue category: Security Vulnerability

[BUG] Security Vulnerability
Please check it here https://huntr.com/bounties/e3d7a994-bfd6-4772-ac9b-9aee1aa16a5f/


Functions to edit:
mlflow/store/artifact/local_artifact_repo.py:LocalArtifactRepository.download_artifacts
mlflow/store/artifact/local_artifact_repo.py:LocalArtifactRepository._download_file

Patch:
diff --git a/mlflow/store/artifact/local_artifact_repo.py b/mlflow/store/artifact/local_artifact_repo.py
index f7d78e90fd0ac..42a0dd98938cc 100644
--- a/mlflow/store/artifact/local_artifact_repo.py
+++ b/mlflow/store/artifact/local_artifact_repo.py
@@ -9,6 +9,7 @@
     mkdir,
     relative_path_to_artifact_path,
 )
+from mlflow.utils.uri import validate_path_is_safe
 
 
 class LocalArtifactRepository(ArtifactRepository):
@@ -74,8 +75,9 @@ def download_artifacts(self, artifact_path, dst_path=None):
         """
         if dst_path:
             return super().download_artifacts(artifact_path, dst_path)
-        # NOTE: The artifact_path is expected to be in posix format.
+        # NOTE: The artifact_path is expected to be a relative path in posix format.
         # Posix paths work fine on windows but just in case we normalize it here.
+        artifact_path = validate_path_is_safe(artifact_path)
         local_artifact_path = os.path.join(self.artifact_dir, os.path.normpath(artifact_path))
         if not os.path.exists(local_artifact_path):
             raise OSError(f"No such file or directory: '{local_artifact_path}'")
@@ -100,8 +102,9 @@ def list_artifacts(self, path=None):
             return []
 
     def _download_file(self, remote_file_path, local_path):
-        # NOTE: The remote_file_path is expected to be in posix format.
+        # NOTE: The remote_file_path is expected to be a relative path in posix format.
         # Posix paths work fine on windows but just in case we normalize it here.
+        remote_file_path = validate_path_is_safe(remote_file_path)
         remote_file_path = os.path.join(self.artifact_dir, os.path.normpath(remote_file_path))
         shutil.copy2(remote_file_path, local_path)
 

Test patch: diff --git a/tests/store/artifact/test_local_artifact_repo.py b/tests/store/artifact/test_local_artifact_repo.py
index ed0c7693848b8..0cea539777dfb 100644
--- a/tests/store/artifact/test_local_artifact_repo.py
+++ b/tests/store/artifact/test_local_artifact_repo.py
@@ -200,3 +200,8 @@ def test_delete_artifacts(local_artifact_repo):
 
 def test_delete_artifacts_with_nonexistent_path_succeeds(local_artifact_repo):
     local_artifact_repo.delete_artifacts("nonexistent")
+
+
+def test_download_artifacts_invalid_remote_file_path(local_artifact_repo):
+    with pytest.raises(MlflowException, match="Invalid path"):
+        local_artifact_repo.download_artifacts("/absolute/path/to/file")
diff --git a/tests/utils/test_promptlab_utils.py b/tests/utils/test_promptlab_utils.py
index 326a9dea032e9..fc71bea53ce42 100644
--- a/tests/utils/test_promptlab_utils.py
+++ b/tests/utils/test_promptlab_utils.py
@@ -104,4 +104,4 @@ def test_create_promptlab_run(store):
     # try to load the model
     import mlflow.pyfunc
 
-    mlflow.pyfunc.load_model(os.path.join(artifact_location, "model"))
+    mlflow.pyfunc.load_model(f"{artifact_location}/model")