CVE-2023-4033: Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection') (CWE-78)
value: OS Command Injection in GitHub repository mlflow/mlflow prior to 2.6.0.
CWE description: The product constructs all or part of an OS command using externally-influenced input from an upstream component, but it does not neutralize or incorrectly neutralizes special elements that could modify the intended OS command when it is sent to a downstream component.
CVSS v3 base score: 7.8
Published: 2023-08-01T01:15Z
Repository: https://github.com/mlflow/mlflow
Language: Python
Files changed: mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py, mlflow/pyfunc/backend.py, mlflow/pyfunc/scoring_server/__init__.py, tests/models/test_cli.py
Commit date: 2023-07-18 09:31:13 +0800
Commit message: Fix potential issues with PyFuncBackend in cli (#9053)

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>
Diff: diff --git a/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py b/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py
new file mode 100644
index 000000000..84e30c1a6
--- /dev/null
+++ b/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py
@@ -0,0 +1,29 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+import argparse
+
+from mlflow.pyfunc.scoring_server import _predict
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-uri", required=True)
+    parser.add_argument("--input-path", required=False)
+    parser.add_argument("--output-path", required=False)
+    parser.add_argument("--content-type", required=True)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    _predict(
+        model_uri=args.model_uri,
+        input_path=args.input_path if args.input_path else None,
+        output_path=args.output_path if args.output_path else None,
+        content_type=args.content_type,
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/mlflow/pyfunc/backend.py b/mlflow/pyfunc/backend.py
index 5e5cc9235..bb8e8bf33 100644
--- a/mlflow/pyfunc/backend.py
+++ b/mlflow/pyfunc/backend.py
@@ -1,10 +1,11 @@
 import logging
 import os
 import pathlib
 import subprocess
 import posixpath
+import shlex
 import sys
 import warnings
 import ctypes
 import signal
 from pathlib import Path
@@ -22,10 +23,11 @@ from mlflow.models.container import ENABLE_MLSERVER
 from mlflow.pyfunc import ENV, scoring_server, mlserver, _extract_conda_env
 
 from mlflow.utils.conda import get_or_create_conda_env, get_conda_bin_executable
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
 from mlflow.utils import env_manager as _EnvManager
+from mlflow.pyfunc import _mlflow_pyfunc_backend_predict
 from mlflow.utils.file_utils import (
     path_to_local_file_uri,
     get_or_create_tmp_dir,
     get_or_create_nfs_tmp_dir,
 )
@@ -141,24 +143,23 @@ class PyFuncBackend(FlavorBackend):
         # NB: Absolute windows paths do not work with mlflow apis, use file uri to ensure
         # platform compatibility.
         local_uri = path_to_local_file_uri(local_path)
 
         if self._env_manager != _EnvManager.LOCAL:
-            command = (
-                'python -c "from mlflow.pyfunc.scoring_server import _predict; _predict('
-                "model_uri={model_uri}, "
-                "input_path={input_path}, "
-                "output_path={output_path}, "
-                "content_type={content_type})"
-                '"'
-            ).format(
-                model_uri=repr(local_uri),
-                input_path=repr(input_path),
-                output_path=repr(output_path),
-                content_type=repr(content_type),
-            )
-            return self.prepare_env(local_path).execute(command)
+            predict_cmd = [
+                "python",
+                _mlflow_pyfunc_backend_predict.__file__,
+                "--model-uri",
+                str(local_uri),
+                "--content-type",
+                shlex.quote(str(content_type)),
+            ]
+            if input_path:
+                predict_cmd += ["--input-path", shlex.quote(str(input_path))]
+            if output_path:
+                predict_cmd += ["--output-path", shlex.quote(str(output_path))]
+            return self.prepare_env(local_path).execute(" ".join(predict_cmd))
         else:
             scoring_server._predict(local_uri, input_path, output_path, content_type)
 
     def serve(
         self,
diff --git a/mlflow/pyfunc/scoring_server/__init__.py b/mlflow/pyfunc/scoring_server/__init__.py
index 3f378814b..159704eea 100644
--- a/mlflow/pyfunc/scoring_server/__init__.py
+++ b/mlflow/pyfunc/scoring_server/__init__.py
@@ -15,10 +15,11 @@ Defines four endpoints:
 from typing import Tuple, Dict
 import flask
 import json
 import logging
 import os
+import shlex
 import sys
 import traceback
 
 from mlflow.environment_variables import MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT
 
@@ -29,10 +30,11 @@ from mlflow.environment_variables import MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT
 # ALl of the mlflow dependencies below need to be backwards compatible.
 from mlflow.exceptions import MlflowException
 from mlflow.types import Schema
 from mlflow.utils import reraise
 from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.os import is_windows
 from mlflow.utils.proto_json_utils import (
     NumpyEncoder,
     dataframe_from_parsed_json,
     _get_jsonable_obj,
     parse_tf_serving_input,
@@ -326,18 +328,20 @@ def _serve(model_uri, port, host):
 def get_cmd(
     model_uri: str, port: int = None, host: int = None, timeout: int = None, nworkers: int = None
 ) -> Tuple[str, Dict[str, str]]:
     local_uri = path_to_local_file_uri(model_uri)
     timeout = timeout or MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get()
+
     # NB: Absolute windows paths do not work with mlflow apis, use file uri to ensure
     # platform compatibility.
-    if os.name != "nt":
+    if not is_windows():
         args = [f"--timeout={timeout}"]
         if port and host:
-            args.append(f"-b {host}:{port}")
+            address = shlex.quote(f"{host}:{port}")
+            args.append(f"-b {address}")
         elif host:
-            args.append(f"-b {host}")
+            args.append(f"-b {shlex.quote(host)}")
 
         if nworkers:
             args.append(f"-w {nworkers}")
 
         command = (
@@ -345,11 +349,11 @@ def get_cmd(
             " -- mlflow.pyfunc.scoring_server.wsgi:app"
         )
     else:
         args = []
         if host:
-            args.append(f"--host={host}")
+            args.append(f"--host={shlex.quote(host)}")
 
         if port:
             args.append(f"--port={port}")
 
         command = (
diff --git a/tests/models/test_cli.py b/tests/models/test_cli.py
index 8554eda67..3a7a07183 100644
--- a/tests/models/test_cli.py
+++ b/tests/models/test_cli.py
@@ -26,18 +26,20 @@ from mlflow.utils.conda import _get_conda_env_name
 import mlflow.models.cli as models_cli
 
 from mlflow.environment_variables import MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import ErrorCode, BAD_REQUEST
+from mlflow.pyfunc.backend import PyFuncBackend
 from mlflow.pyfunc.scoring_server import (
     CONTENT_TYPE_JSON,
     CONTENT_TYPE_CSV,
 )
 from mlflow.utils.file_utils import TempDir
 from mlflow.utils.environment import _mlflow_conda_env
 from mlflow.utils import env_manager as _EnvManager
 from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.process import ShellCommandException
 from tests.helper_functions import (
     pyfunc_build_image,
     pyfunc_serve_from_docker_image,
     pyfunc_serve_from_docker_image_with_env_override,
     RestEndpoint,
@@ -172,11 +174,11 @@ def test_serve_gunicorn_opts(iris_data, sk_model):
 def test_predict(iris_data, sk_model):
     with TempDir(chdr=True) as tmp:
         with mlflow.start_run() as active_run:
             mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
             model_uri = f"runs:/{active_run.info.run_id}/model"
-        model_registry_uri = "models:/{name}/{stage}".format(name="impredicting", stage="None")
+        model_registry_uri = "models:/impredicting/None"
         input_json_path = tmp.path("input.json")
         input_csv_path = tmp.path("input.csv")
         output_json_path = tmp.path("output.json")
         x, _ = iris_data
         with open(input_json_path, "w") as f:
@@ -329,10 +331,177 @@ def test_predict(iris_data, sk_model):
         actual = actual[actual.columns[0]].values
         expected = sk_model.predict(x)
         assert all(expected == actual)
 
 
+def test_predict_check_content_type(iris_data, sk_model, tmp_path):
+    with mlflow.start_run():
+        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
+    model_registry_uri = "models:/impredicting/None"
+    input_json_path = tmp_path / "input.json"
+    input_csv_path = tmp_path / "input.csv"
+    output_json_path = tmp_path / "output.json"
+
+    x, _ = iris_data
+    with input_json_path.open("w") as f:
+        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)
+
+    pd.DataFrame(x).to_csv(input_csv_path, index=False)
+
+    # Throw errors for invalid content_type
+    prc = subprocess.run(
+        [
+            "mlflow",
+            "models",
+            "predict",
+            "-m",
+            model_registry_uri,
+            "-i",
+            input_json_path,
+            "-o",
+            output_json_path,
+            "-t",
+            "invalid",
+            "--env-manager",
+            "local",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_with_tracking_uri(),
+        check=False,
+    )
+    assert prc.returncode != 0
+    assert "Unknown content type" in prc.stderr.decode("utf-8")
+
+
+def test_predict_check_input_path(iris_data, sk_model, tmp_path):
+    with mlflow.start_run():
+        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
+    model_registry_uri = "models:/impredicting/None"
+    input_json_path = tmp_path / "input with space.json"
+    input_csv_path = tmp_path / "input.csv"
+    output_json_path = tmp_path / "output.json"
+
+    x, _ = iris_data
+    with input_json_path.open("w") as f:
+        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)
+
+    pd.DataFrame(x).to_csv(input_csv_path, index=False)
+
+    # Valid input path with space
+    prc = subprocess.run(
+        [
+            "mlflow",
+            "models",
+            "predict",
+            "-m",
+            model_registry_uri,
+            "-i",
+            f"{input_json_path}",
+            "-o",
+            output_json_path,
+            "--env-manager",
+            "local",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_with_tracking_uri(),
+        check=False,
+        text=True,
+    )
+    assert prc.returncode == 0
+
+    # Throw errors for invalid input_path
+    prc = subprocess.run(
+        [
+            "mlflow",
+            "models",
+            "predict",
+            "-m",
+            model_registry_uri,
+            "-i",
+            f'{input_json_path}"; echo ThisIsABug! "',
+            "-o",
+            output_json_path,
+            "--env-manager",
+            "local",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_with_tracking_uri(),
+        check=False,
+        text=True,
+    )
+    assert prc.returncode != 0
+    assert "ThisIsABug!" not in prc.stdout
+    assert "FileNotFoundError" in prc.stderr
+
+    prc = subprocess.run(
+        [
+            "mlflow",
+            "models",
+            "predict",
+            "-m",
+            model_registry_uri,
+            "-i",
+            f'{input_csv_path}"; echo ThisIsABug! "',
+            "-o",
+            output_json_path,
+            "-t",
+            "csv",
+            "--env-manager",
+            "local",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_with_tracking_uri(),
+        check=False,
+        text=True,
+    )
+    assert prc.returncode != 0
+    assert "ThisIsABug!" not in prc.stdout
+    assert "FileNotFoundError" in prc.stderr
+
+
+def test_predict_check_output_path(iris_data, sk_model, tmp_path):
+    with mlflow.start_run():
+        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
+    model_registry_uri = "models:/impredicting/None"
+    input_json_path = tmp_path / "input.json"
+    input_csv_path = tmp_path / "input.csv"
+    output_json_path = tmp_path / "output.json"
+
+    x, _ = iris_data
+    with input_json_path.open("w") as f:
+        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)
+
+    pd.DataFrame(x).to_csv(input_csv_path, index=False)
+
+    prc = subprocess.run(
+        [
+            "mlflow",
+            "models",
+            "predict",
+            "-m",
+            model_registry_uri,
+            "-i",
+            input_json_path,
+            "-o",
+            f'{output_json_path}"; echo ThisIsABug! "',
+            "--env-manager",
+            "local",
+        ],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        env=env_with_tracking_uri(),
+        check=False,
+        text=True,
+    )
+    assert prc.returncode == 0
+    assert "ThisIsABug!" not in prc.stdout
+
+
 def test_prepare_env_passes(sk_model):
     if no_conda:
         pytest.skip("This test requires conda.")
 
     with TempDir(chdr=True):
@@ -572,10 +741,29 @@ def test_env_manager_unsupported_value():
             ["--model-uri", "model", "--env-manager", "abc"],
             catch_exceptions=False,
         )
 
 
+def test_host_invalid_value():
+    class MyModel(mlflow.pyfunc.PythonModel):
+        def predict(self, ctx, model_input):
+            return model_input
+
+    with mlflow.start_run():
+        model_info = mlflow.pyfunc.log_model(
+            python_model=MyModel(), artifact_path="test_model", registered_model_name="model"
+        )
+
+    with mock.patch("mlflow.models.cli.get_flavor_backend", return_value=PyFuncBackend({})):
+        with pytest.raises(ShellCommandException, match=r"Non-zero exit code: 1"):
+            CliRunner().invoke(
+                models_cli.serve,
+                ["--model-uri", model_info.model_uri, "--host", "localhost & echo BUG"],
+                catch_exceptions=False,
+            )
+
+
 def test_change_conda_env_root_location(tmp_path, sk_model):
     env_root1_path = tmp_path / "root1"
     env_root1_path.mkdir()
 
     env_root2_path = tmp_path / "root2"

Vulnerable code:             command = (
                'python -c "from mlflow.pyfunc.scoring_server import _predict; _predict('
                "model_uri={model_uri}, "
                "input_path={input_path}, "
                "output_path={output_path}, "
                "content_type={content_type})"
                '"'
            ).format(
                model_uri=repr(local_uri),
                input_path=repr(input_path),
                output_path=repr(output_path),
                content_type=repr(content_type),
            )
            return self.prepare_env(local_path).execute(command)
    if os.name != "nt":
            args.append(f"-b {host}:{port}")
            args.append(f"-b {host}")
            args.append(f"--host={host}")
        model_registry_uri = "models:/{name}/{stage}".format(name="impredicting", stage="None")
Fixed code: """
This script should be executed in a fresh python interpreter process using `subprocess`.
"""
import argparse

from mlflow.pyfunc.scoring_server import _predict


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-uri", required=True)
    parser.add_argument("--input-path", required=False)
    parser.add_argument("--output-path", required=False)
    parser.add_argument("--content-type", required=True)
    return parser.parse_args()


def main():
    args = parse_args()
    _predict(
        model_uri=args.model_uri,
        input_path=args.input_path if args.input_path else None,
        output_path=args.output_path if args.output_path else None,
        content_type=args.content_type,
    )


if __name__ == "__main__":
    main()
import shlex
from mlflow.pyfunc import _mlflow_pyfunc_backend_predict
            predict_cmd = [
                "python",
                _mlflow_pyfunc_backend_predict.__file__,
                "--model-uri",
                str(local_uri),
                "--content-type",
                shlex.quote(str(content_type)),
            ]
            if input_path:
                predict_cmd += ["--input-path", shlex.quote(str(input_path))]
            if output_path:
                predict_cmd += ["--output-path", shlex.quote(str(output_path))]
            return self.prepare_env(local_path).execute(" ".join(predict_cmd))
import shlex
from mlflow.utils.os import is_windows

    if not is_windows():
            address = shlex.quote(f"{host}:{port}")
            args.append(f"-b {address}")
            args.append(f"-b {shlex.quote(host)}")
            args.append(f"--host={shlex.quote(host)}")
from mlflow.pyfunc.backend import PyFuncBackend
from mlflow.utils.process import ShellCommandException
        model_registry_uri = "models:/impredicting/None"
def test_predict_check_content_type(iris_data, sk_model, tmp_path):
    with mlflow.start_run():
        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
    model_registry_uri = "models:/impredicting/None"
    input_json_path = tmp_path / "input.json"
    input_csv_path = tmp_path / "input.csv"
    output_json_path = tmp_path / "output.json"

    x, _ = iris_data
    with input_json_path.open("w") as f:
        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)

    pd.DataFrame(x).to_csv(input_csv_path, index=False)

    # Throw errors for invalid content_type
    prc = subprocess.run(
        [
            "mlflow",
            "models",
            "predict",
            "-m",
            model_registry_uri,
            "-i",
            input_json_path,
            "-o",
            output_json_path,
            "-t",
            "invalid",
            "--env-manager",
            "local",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env_with_tracking_uri(),
        check=False,
    )
    assert prc.returncode != 0
    assert "Unknown content type" in prc.stderr.decode("utf-8")


def test_predict_check_input_path(iris_data, sk_model, tmp_path):
    with mlflow.start_run():
        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
    model_registry_uri = "models:/impredicting/None"
    input_json_path = tmp_path / "input with space.json"
    input_csv_path = tmp_path / "input.csv"
    output_json_path = tmp_path / "output.json"

    x, _ = iris_data
    with input_json_path.open("w") as f:
        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)

    pd.DataFrame(x).to_csv(input_csv_path, index=False)

    # Valid input path with space
    prc = subprocess.run(
        [
            "mlflow",
            "models",
            "predict",
            "-m",
            model_registry_uri,
            "-i",
            f"{input_json_path}",
            "-o",
            output_json_path,
            "--env-manager",
            "local",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env_with_tracking_uri(),
        check=False,
        text=True,
    )
    assert prc.returncode == 0

    # Throw errors for invalid input_path
    prc = subprocess.run(
        [
            "mlflow",
            "models",
            "predict",
            "-m",
            model_registry_uri,
            "-i",
            f'{input_json_path}"; echo ThisIsABug! "',
            "-o",
            output_json_path,
            "--env-manager",
            "local",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env_with_tracking_uri(),
        check=False,
        text=True,
    )
    assert prc.returncode != 0
    assert "ThisIsABug!" not in prc.stdout
    assert "FileNotFoundError" in prc.stderr

    prc = subprocess.run(
        [
            "mlflow",
            "models",
            "predict",
            "-m",
            model_registry_uri,
            "-i",
            f'{input_csv_path}"; echo ThisIsABug! "',
            "-o",
            output_json_path,
            "-t",
            "csv",
            "--env-manager",
            "local",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env_with_tracking_uri(),
        check=False,
        text=True,
    )
    assert prc.returncode != 0
    assert "ThisIsABug!" not in prc.stdout
    assert "FileNotFoundError" in prc.stderr


def test_predict_check_output_path(iris_data, sk_model, tmp_path):
    with mlflow.start_run():
        mlflow.sklearn.log_model(sk_model, "model", registered_model_name="impredicting")
    model_registry_uri = "models:/impredicting/None"
    input_json_path = tmp_path / "input.json"
    input_csv_path = tmp_path / "input.csv"
    output_json_path = tmp_path / "output.json"

    x, _ = iris_data
    with input_json_path.open("w") as f:
        json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f)

    pd.DataFrame(x).to_csv(input_csv_path, index=False)

    prc = subprocess.run(
        [
            "mlflow",
            "models",
            "predict",
            "-m",
            model_registry_uri,
            "-i",
            input_json_path,
            "-o",
            f'{output_json_path}"; echo ThisIsABug! "',
            "--env-manager",
            "local",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        env=env_with_tracking_uri(),
        check=False,
        text=True,
    )
    assert prc.returncode == 0
    assert "ThisIsABug!" not in prc.stdout


def test_host_invalid_value():
    class MyModel(mlflow.pyfunc.PythonModel):
        def predict(self, ctx, model_input):
            return model_input

    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            python_model=MyModel(), artifact_path="test_model", registered_model_name="model"
        )

    with mock.patch("mlflow.models.cli.get_flavor_backend", return_value=PyFuncBackend({})):
        with pytest.raises(ShellCommandException, match=r"Non-zero exit code: 1"):
            CliRunner().invoke(
                models_cli.serve,
                ["--model-uri", model_info.model_uri, "--host", "localhost & echo BUG"],
                catch_exceptions=False,
            )