coreml-uncontrolled-mem-poc

coremltools CWE-789: Uncontrolled Memory Allocation in MIL Proto Loader

Status: CONFIRMED — READY TO SUBMIT

Severity: High (P2) — DoS via OOM

Target

• Repo: apple/coremltools
• Platform: huntr.com
• Format: .mlpackage (MIL protobuf)

Root Cause

get_proto_dim() in milproto/helper.py returns dim.constant.size (a uint64 from attacker-controlled protobuf) with zero bounds validation. This flows directly into np.array(value).reshape(shape) with no interposing check.

# helper.py — get_proto_dim() — NO bounds check:
def get_proto_dim(dim):
    if dim.WhichOneof("dimension") == "constant":
        return dim.constant.size  # uint64, attacker-controlled

# load.py:309-330 — shape applied to numpy allocation:
shape = valuetype.get_shape()          # attacker value, no check
value = np.array(value).reshape(shape) # OOM when shape=(2147483647,)

Trigger APIs

All call _convert_model_spec_to_pymil_prog() → milproto.load.load():

• ct.optimize.coreml.get_weights_metadata(mlmodel)
• ct.optimize.coreml.linear_quantize_weights(mlmodel)
• ct.optimize.coreml.palettize_weights(mlmodel)
• ct.optimize.coreml.prune_weights(mlmodel)

Vulnerable Files

PoC File

• poc_coremltools_cwe789.py — builds malicious .mlpackage and triggers OOM

Reproduction

pip install coremltools
python poc_coremltools_cwe789.py

Fix

# helper.py — get_proto_dim():
MAX_DIM = 1 << 31  # 2 billion elements
if size > MAX_DIM:
    raise ValueError(f"Tensor dimension {size} exceeds maximum {MAX_DIM}")
return size