File size: 3,676 Bytes
33b540b | 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 | """
Test how coremltools.utils.load_spec() (and, as a control, raw
Model_pb2.Model().ParseFromString()) behaves when fed the deeply-nested
payloads produced by build_payloads.py, under whatever protobuf
backend/version is active in the current Python process/venv.
Prints, for each payload file:
- the resolved protobuf backend (api_implementation.Type())
- the installed protobuf version
- the outcome of load_spec(): OK / DecodeError (manageable) /
RecursionError (NOT caught by protobuf, escapes to caller) /
"process crashed" (detected by the driver script wrapping this one,
since a native crash can't be caught here)
Run with the venv's python whose backend/version you want to test, e.g.:
verify/venv-protobuf/Scripts/python.exe test_parse.py out/pipeline_depth1000.mlmodel
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python verify/venv-cve-old/Scripts/python.exe test_parse.py out/pipeline_depth5000.mlmodel
"""
import sys
import threading
def main():
if len(sys.argv) < 2:
print("usage: test_parse.py <payload.mlmodel> [--raw]")
sys.exit(2)
path = sys.argv[1]
raw_mode = "--raw" in sys.argv[2:]
from google.protobuf.internal import api_implementation
import google.protobuf
print(f"[INFO] protobuf version = {google.protobuf.__version__}")
print(f"[INFO] backend (api_implementation.Type()) = {api_implementation.Type()}")
with open(path, "rb") as f:
data = f.read()
print(f"[INFO] payload = {path} ({len(data)} bytes)")
# Run the parse attempt in a worker thread with a DEFAULT-size stack
# (deliberately NOT enlarged) to faithfully reproduce what happens to a
# normal application thread that calls load_spec()/ParseFromString() on
# an attacker-supplied file -- we are not trying to help the parser
# survive, we are trying to observe what a real caller would experience.
result = {}
def worker():
try:
if raw_mode:
from coremltools.proto import Model_pb2
m = Model_pb2.Model()
m.ParseFromString(data)
result["ok"] = True
result["mode"] = "raw ParseFromString"
else:
import coremltools.models.utils as u
spec = u.load_spec(_write_tmp(data))
result["ok"] = True
result["mode"] = "load_spec"
except RecursionError as e:
result["ok"] = False
result["exc"] = "RecursionError"
result["msg"] = str(e)
except Exception as e:
# google.protobuf.message.DecodeError and anything else
result["ok"] = False
result["exc"] = type(e).__name__
result["msg"] = str(e)
t = threading.Thread(target=worker)
t.start()
t.join()
if not result:
# Thread died without setting result -> only possible if the whole
# process was about to be torn down by a native fault; in practice
# a hard native crash kills the entire process before we get here,
# so this branch is mostly theoretical / a safety net.
print("[RESULT] UNKNOWN (worker thread produced no result)")
return
if result.get("ok"):
print(f"[RESULT] OK: parsed successfully via {result['mode']} (no size/depth limit triggered)")
else:
print(f"[RESULT] EXCEPTION: {result['exc']}: {result['msg'][:200]}")
_tmp_counter = [0]
def _write_tmp(data):
import os
import tempfile
fd, path = tempfile.mkstemp(suffix=".mlmodel")
with os.fdopen(fd, "wb") as f:
f.write(data)
return path
if __name__ == "__main__":
main()
|