hacnho commited on
Commit
b603c7b
·
verified ·
1 Parent(s): fd35a6a

Upload reproduce.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. reproduce.py +82 -0
reproduce.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ from __future__ import annotations
3
+
4
+ import argparse
5
+ import json
6
+ import time
7
+ from pathlib import Path
8
+
9
+ from google.protobuf import descriptor_pb2
10
+
11
+
12
+ ROOT = Path(__file__).resolve().parent
13
+ ARTIFACTS = ROOT / "artifacts"
14
+
15
+
16
+ def vmhwm_kb() -> int | None:
17
+ with open("/proc/self/status", "r", encoding="utf-8", errors="ignore") as fh:
18
+ for line in fh:
19
+ if line.startswith("VmHWM:"):
20
+ return int(line.split()[1])
21
+ return None
22
+
23
+
24
+ def parse_one(path: Path) -> dict:
25
+ start = time.time()
26
+ before = vmhwm_kb()
27
+ msg = descriptor_pb2.ExtensionRangeOptions()
28
+ msg.ParseFromString(path.read_bytes())
29
+ after = vmhwm_kb()
30
+ agg_len = len(msg.uninterpreted_option[0].aggregate_value) if msg.uninterpreted_option else 0
31
+ return {
32
+ "path": str(path),
33
+ "elapsed": time.time() - start,
34
+ "vmhwm_before_kb": before,
35
+ "vmhwm_after_kb": after,
36
+ "uninterpreted_option_len": len(msg.uninterpreted_option),
37
+ "first_aggregate_value_len": agg_len,
38
+ }
39
+
40
+
41
+ def main() -> int:
42
+ parser = argparse.ArgumentParser()
43
+ parser.add_argument(
44
+ "--control",
45
+ type=Path,
46
+ default=ARTIFACTS / "control_one_extensionrange_uninterpreted_option.pb",
47
+ )
48
+ parser.add_argument(
49
+ "--malicious",
50
+ type=Path,
51
+ default=ARTIFACTS / "malicious_extensionrange_uninterpreted_option_5000000.pb",
52
+ )
53
+ parser.add_argument("--json-out", type=Path)
54
+ args = parser.parse_args()
55
+
56
+ control = parse_one(args.control)
57
+ malicious = parse_one(args.malicious)
58
+ summary = {
59
+ "format": "Protocol Buffers",
60
+ "protobuf_version": __import__("google.protobuf").protobuf.__version__,
61
+ "entrypoint": "descriptor_pb2.ExtensionRangeOptions().ParseFromString(...)",
62
+ "control": control,
63
+ "malicious": malicious,
64
+ "delta": {
65
+ "elapsed_ratio": malicious["elapsed"] / max(control["elapsed"], 1e-9),
66
+ "vmhwm_delta_kb": malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"],
67
+ "uninterpreted_options_added": malicious["uninterpreted_option_len"] - control["uninterpreted_option_len"],
68
+ "repro_ok": (
69
+ control["uninterpreted_option_len"] == 1
70
+ and malicious["uninterpreted_option_len"] == 5_000_000
71
+ and malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"] >= 150_000
72
+ ),
73
+ },
74
+ }
75
+ if args.json_out:
76
+ args.json_out.write_text(json.dumps(summary, indent=2) + "\n")
77
+ print(json.dumps(summary, indent=2))
78
+ return 0
79
+
80
+
81
+ if __name__ == "__main__":
82
+ raise SystemExit(main())