Coverage for /pythoncovmergedfiles/medio/medio/src/onnx/onnx/fuzz/fuzz_compose.py: 70%

1###### Coverage stub 

2import atexit 

3import coverage 

4cov = coverage.coverage(data_file='.coverage', cover_pylib=True) 

5cov.start() 

6# Register an exist handler that will print coverage 

7def exit_handler(): 

8 cov.stop() 

9 cov.save() 

10atexit.register(exit_handler) 

11####### End of coverage stub 

12# Copyright (c) ONNX Project Contributors 

13# SPDX-License-Identifier: Apache-2.0 

14"""Atheris fuzz harness for onnx.compose. 

15 

16Two input paths are exercised per iteration, selected by a fuzzer-controlled 

17toggle byte read from the *tail* of the input: 

18 

19* Raw -> a 4-byte big-endian length prefix splits the remaining bytes 

20 into m1 | m2, each parsed via onnx.load_model_from_string. Catches 

21 protobuf parser bugs and bugs reachable only through crafted serialized 

22 models the structured builder will not produce. 

23 

24* Structured -> helper.make_model from FuzzedDataProvider builds two small 

25 linear graphs with predictable input/output names, so merge_models 

26 reaches its connect_io/add_prefix logic on most iterations rather than 

27 only when the parser happens to accept a random byte string. 

28 

29Further bits of the toggle byte select prefix1/prefix2 (add_prefix 

30collision-resolution path) and a randomized (vs. derived) io_map, so a 

31single harness covers merge_models, merge_graphs, check_overlapping_names, 

32the recursive connect_io subgraph rewrite, and add_prefix. 

33""" 

34 

35 

36import random 

37import struct 

38import sys 

39 

40import atheris 

41 

42with atheris.instrument_imports(): 

43 from onnx import ( 

44 TensorProto, 

45 checker, 

46 compose, 

47 defs, 

48 helper, 

49 load_model_from_string, 

50 ) 

51 

52_UNARY = ("Relu", "Sigmoid", "Tanh", "Abs", "Neg", "Identity") 

53 

54 

55def _value_info(name): 

56 return helper.make_tensor_value_info(name, TensorProto.FLOAT, ["N"]) 

57 

58 

59def _build_model(fdp, tag, opset): 

60 """Build a small linear graph with predictable input/output names. 

61 

62 Naming the sole input/output `In<tag>`/`Out<tag>` lets the derived 

63 io_map (m1's outputs -> m2's inputs) connect on most iterations, so 

64 merge_models' rewrite logic is reached without relying on the raw 

65 protobuf parser to produce compatible names. `opset` is shared between 

66 both models: merge_models rejects mismatched opset_import up front, so 

67 picking it independently per model would make the structured path 

68 rarely reach the merge/connect logic it exists to exercise. 

69 """ 

70 n_ops = fdp.ConsumeIntInRange(0, 4) 

71 last = f"In{tag}" 

72 nodes = [] 

73 for i in range(n_ops): 

74 op = _UNARY[fdp.ConsumeIntInRange(0, len(_UNARY) - 1)] 

75 nxt = f"v{tag}_{i}" 

76 nodes.append(helper.make_node(op, [last], [nxt])) 

77 last = nxt 

78 out_name = f"Out{tag}" 

79 nodes.append(helper.make_node("Identity", [last], [out_name])) 

80 graph = helper.make_graph( 

81 nodes, f"g{tag}", [_value_info(f"In{tag}")], [_value_info(out_name)] 

82 ) 

83 return helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)]) 

84 

85 

86def _derived_io_map(m1, m2): 

87 outs = [o.name for o in m1.graph.output] 

88 ins = [i.name for i in m2.graph.input] 

89 return list(zip(outs, ins, strict=False)) 

90 

91 

92def _random_io_map(data, m1, m2): 

93 outs = [o.name for o in m1.graph.output] 

94 ins = [i.name for i in m2.graph.input] 

95 if not outs or not ins: 

96 return [] 

97 rng = random.Random(data) 

98 rng.shuffle(outs) 

99 rng.shuffle(ins) 

100 n = rng.randint(0, min(len(outs), len(ins))) 

101 return list(zip(outs[:n], ins[:n], strict=True)) 

102 

103 

104def TestOneInput(data): 

105 # Toggles live in the trailing byte; see module docstring and 

106 # onnx/fuzz/README.md's "fuzz_compose.py toggle byte" table. 

107 if len(data) < 2: 

108 return 

109 toggles = data[-1] 

110 body = data[:-1] 

111 use_prefix = bool(toggles & 0x01) 

112 use_structured = bool(toggles & 0x04) 

113 use_random_io_map = bool(toggles & 0x08) 

114 # bits 0x02, 0x10..0x80 are reserved for future toggles; mutations 

115 # against them are harmless until claimed. 

116 

117 try: 

118 if use_structured: 

119 fdp = atheris.FuzzedDataProvider(body) 

120 opset = fdp.ConsumeIntInRange(7, defs.onnx_opset_version()) 

121 m1 = _build_model(fdp, 1, opset) 

122 m2 = _build_model(fdp, 2, opset) 

123 else: 

124 if len(body) < 4: 

125 return 

126 (n1,) = struct.unpack(">I", body[:4]) 

127 rest = body[4:] 

128 if n1 > len(rest): 

129 return 

130 m1 = load_model_from_string(rest[:n1]) 

131 m2 = load_model_from_string(rest[n1:]) 

132 

133 io_map = ( 

134 _random_io_map(body, m1, m2) 

135 if use_random_io_map 

136 else _derived_io_map(m1, m2) 

137 ) 

138 

139 prefix1 = "g1_" if use_prefix else None 

140 prefix2 = "g2_" if use_prefix else None 

141 merged = compose.merge_models(m1, m2, io_map, prefix1=prefix1, prefix2=prefix2) 

142 checker.check_model(merged, full_check=True) 

143 except Exception: 

144 # Malformed fuzz inputs raise a broad set of expected exceptions 

145 # (ValidationError, TypeError, ValueError, DecodeError, ...). 

146 # Real bugs surface as crashes, hangs, or sanitizer reports. 

147 return 

148 

149 

150def main(): 

151 atheris.instrument_all() 

152 atheris.Setup(sys.argv, TestOneInput, enable_python_coverage=True) 

153 atheris.Fuzz() 

154 

155 

156if __name__ == "__main__": 

157 main()