CVE-2023-4033 (Python, published 2023-08-01T01:15Z)
Description: value: OS Command Injection in GitHub repository mlflow/mlflow prior to 2.6.0.
Severity: nan (CVSS v3 base score: 7.8)
Weakness: CWE-78 \u2014 Improper Neutralization of Special Elements used in an OS Command ('OS Command Injection')
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.
Repository: https://github.com/mlflow/mlflow
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 message: Fix potential issues with PyFuncBackend in cli (#9053)

Signed-off-by: Serena Ruan <serena.rxy@gmail.com>

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,
            )