EnigmaConsultant/huntr-poc-mlflow-paramspec-deserialization-dos

MLflow model-signature param-schema deserialization DoS (unhandled KeyError/TypeError in ParamSpec.from_json_dict)

Target: mlflow/mlflow — verified against PyPI mlflow==3.14.0 Environment: CPython 3.13.12, numpy 2.5.1 (Kali Linux) Class: Uncaught exception / loader denial-of-service (non-RCE). Triggered at pure model load time — no inference, no allocation. Attack surface: the MLmodel metadata file, which travels inside every model directory and is fully attacker-controlled. The victim only has to load / inspect the model.

---

Root cause

In mlflow/types/schema.py, ParamSpec.from_json_dict (line 1292) resolves the declared parameter type with no error handling:

# mlflow/types/schema.py  (ParamSpec.from_json_dict, ~line 1288-1296)
dtype = kwargs.get("type") or kwargs.get("dtype")
dtype = Object.from_json_dict(**kwargs) if dtype == OBJECT_TYPE else DataType[dtype]
return cls(
    name=str(kwargs["name"]),
    dtype=dtype,
    default=kwargs["default"],
    ...
)

DataType[dtype] is an Enum.__getitem__ lookup:

• For an unknown type string, it raises builtins.KeyError.
• For a non-string / unhashable value (e.g. a JSON list), it raises builtins.TypeError: unhashable type: '...'.

Neither is an MlflowException, and nothing on the Model.load() / mlflow.pyfunc.load_model() path catches it, so the loader crashes with a raw Python exception.

This is a guarding inconsistency, not a design decision

The same DataType[...] lookup is defensively wrapped elsewhere in the exact same file:

# mlflow/types/schema.py  (ParamSpec.__init__, ~line 1135-1142)
try:
    self._dtype = DataType[dtype] if isinstance(dtype, str) else dtype
except KeyError:
    supported_types = [t.name for t in DataType if t.name != "binary"]
    raise MlflowException.invalid_parameter_value(
        f"Unsupported type '{dtype}', expected instance of DataType or "
        f"one of {supported_types}",
    )

ColSpec.__init__ (lines ~741-747) has the identical guard. But from_json_dict resolves the dtype BEFORE calling the constructor, so it bypasses ParamSpec.__init__'s guard entirely. The deserialization entrypoint is the one place the guard is missing.

Load path reaching the bug

Model.load()                                  mlflow/models/model.py:821
  -> Model.from_dict()                        mlflow/models/model.py:831
    -> ModelSignature.from_dict()             mlflow/models/signature.py:160
      -> ParamSchema.from_json()              mlflow/types/schema.py:1348
        -> ParamSpec.from_json_dict()         mlflow/types/schema.py:1292  <-- crash

mlflow.pyfunc.load_model() materializes the same signature metadata and crashes identically.

---

Distinct from prior / same-session MLflow findings

• signature-recursion DoS — unbounded recursion in Array/Object/Map.from_json_dict. Different root cause (recursion depth), different code.
• tensorspec dtype DoS — huge structured-dtype itemsize allocation at enforce time. Different root cause (allocation), fires at inference, not load.
• pyfunc artifacts path traversal — filesystem traversal, unrelated code.

This one is a missing error-handling gap on the params deserialization path (line 1292), firing at pure load time with no inference and no allocation. Distinct root cause, distinct location.

No known CVE covers the param-schema deserialization path as of packaging.

---

Proof of Concept

Verified end-to-end against a genuine, loadable sklearn model:

1. Train a real sklearn.LogisticRegression, infer a signature with params ({"temperature": 0.5}), and mlflow.sklearn.save_model. This produces a real model whose MLmodel signature.params = [{"name":"temperature","type":"double",...}].
2. Baseline: Model.load / mlflow.pyfunc.load_model / mlflow.sklearn.load_model all return OK.
3. Attack (KeyError): edit only the MLmodel signature.params JSON to declare an unknown type — [{"name":"temperature","type":"__evil__","default":0.5}]. Reloading raises an UNCAUGHT builtins.KeyError: '__evil__' from the loader.
4. Attack-2 (TypeError): set the params type to a JSON list ["x"] — both loaders raise UNCAUGHT builtins.TypeError: unhashable type: 'list'.
5. Negative control: putting the exact same invalid type string "__evil__" into an inputs ColSpec instead of params yields a graceful mlflow.exceptions.MlflowException (Unsupported type '__evil__', expected instance of DataType or one of [...]) across all three loaders — proving MLflow intends graceful handling and the params path is an unguarded gap, not by-design.

Impact: a ~1 KB malicious MLmodel metadata file DoS-crashes any code path that materializes model signature metadata — Model.load, pyfunc.load_model, and downstream tracking-server / model-registry / UI inspection — with an unhandled exception.

Repro scripts: e2e_repro.py, fuzz_load.py, indep_verify.py.

---

Captured evidence (verbatim)

$ python3 e2e_repro.py
mlflow 3.14.0 | numpy 2.5.1 | py 3.13.12

--- original signature block ---
params: [{"name": "temperature", "default": 0.5, "shape": null, "type": "double"}]

=== BASELINE (untampered real model) ===
[baseline] Model.load: OK
[baseline] pyfunc.load_model: OK
[baseline] sklearn.load_model: OK

=== ATTACK: malicious params type string in MLmodel (KeyError) ===
[attack-keyerror] Model.load: builtins.KeyError (UNCAUGHT crash): '__evil__'
[attack-keyerror] pyfunc.load_model: builtins.KeyError (UNCAUGHT crash): '__evil__'

=== ATTACK 2: malicious params dtype as non-hashable list (TypeError) ===
[attack-typeerror] Model.load: builtins.TypeError (UNCAUGHT crash): unhashable type: 'list'
[attack-typeerror] pyfunc.load_model: builtins.TypeError (UNCAUGHT crash): unhashable type: 'list'

=== NEGATIVE CONTROL: same invalid type in COL-SPEC inputs -> graceful ===
[negctl-colspec] Model.load: mlflow.exceptions.MlflowException (graceful MlflowException): Unsupported type '__evil__', expected instance of DataType or one of ['boolean', 'integer'
[negctl-colspec] pyfunc.load_model: mlflow.exceptions.MlflowException (graceful MlflowException): Unsupported type '__evil__'...

Full traceback of the primary crash (Model.load, invalid params type)

Traceback (most recent call last):
  File "e2e_repro.py", line 77, in <module>
    Model.load(mpath)
  File ".../mlflow/models/model.py", line 821, in load
    return cls.from_dict(model_dict)
  File ".../mlflow/models/model.py", line 831, in from_dict
    signature = ModelSignature.from_dict(model_dict["signature"])
  File ".../mlflow/models/signature.py", line 160, in from_dict
    params = ParamSchema.from_json(x) if (x := signature_dict.get("params")) else None
  File ".../mlflow/types/schema.py", line 1348, in from_json
    return cls([ParamSpec.from_json_dict(**x) for x in json.loads(json_str)])
  File ".../mlflow/types/schema.py", line 1292, in from_json_dict
    dtype = Object.from_json_dict(**kwargs) if dtype == OBJECT_TYPE else DataType[dtype]
  File "/usr/lib/python3.13/enum.py", line 794, in __getitem__
    return cls._member_map_[name]
KeyError: '__evil__'

---

Suggested fix

Wrap the DataType[dtype] resolution in ParamSpec.from_json_dict in the same try/except KeyError -> MlflowException guard already used in ParamSpec.__init__ and ColSpec.__init__, and validate that dtype is a hashable string before the lookup (to catch the TypeError/unhashable case). Alternatively, route resolution through the constructor so the existing guard applies.