hacnho commited on
Commit
c152b3a
·
verified ·
1 Parent(s): 67ccde0

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
+ import msgpack
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
+ blob = path.read_bytes()
26
+ start = time.time()
27
+ before = vmhwm_kb()
28
+ obj = msgpack.unpackb(blob, raw=False)
29
+ after = vmhwm_kb()
30
+ return {
31
+ "path": str(path),
32
+ "elapsed": time.time() - start,
33
+ "vmhwm_before_kb": before,
34
+ "vmhwm_after_kb": after,
35
+ "python_type": type(obj).__name__,
36
+ "length": len(obj),
37
+ }
38
+
39
+
40
+ def main() -> int:
41
+ parser = argparse.ArgumentParser()
42
+ parser.add_argument(
43
+ "--control",
44
+ type=Path,
45
+ default=ARTIFACTS / "control_bin32_same_size.msgpack",
46
+ )
47
+ parser.add_argument(
48
+ "--malicious",
49
+ type=Path,
50
+ default=ARTIFACTS / "malicious_array32_empty_strings_20000000.msgpack",
51
+ )
52
+ parser.add_argument("--json-out", type=Path)
53
+ args = parser.parse_args()
54
+
55
+ control = parse_one(args.control)
56
+ malicious = parse_one(args.malicious)
57
+ summary = {
58
+ "format": "MessagePack (.msgpack)",
59
+ "msgpack_version": msgpack.__version__,
60
+ "entrypoint": "msgpack.unpackb(..., raw=False)",
61
+ "control": control,
62
+ "malicious": malicious,
63
+ "delta": {
64
+ "elapsed_ratio": malicious["elapsed"] / max(control["elapsed"], 1e-9),
65
+ "vmhwm_delta_kb": malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"],
66
+ "same_size_files": args.control.stat().st_size == args.malicious.stat().st_size,
67
+ "repro_ok": (
68
+ control["python_type"] == "bytes"
69
+ and malicious["python_type"] == "list"
70
+ and malicious["length"] == 20_000_000
71
+ and malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"] >= 120_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())