test_parse.py

"""
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()