Skip to content

Commit 8d85d37

Browse files
committed
examples and docs
1 parent e1a2807 commit 8d85d37

File tree

4 files changed

+227
-75
lines changed

4 files changed

+227
-75
lines changed

README.md

+112-32
Original file line numberDiff line numberDiff line change
@@ -11,55 +11,134 @@ It uses [zeroize](https://crates.io/crates/zeroize) crate under the hood.
1111
It can work with `bytearray` and `numpy array`.
1212

1313
> [!WARNING]
14-
> **Currently it doens't work in the case of [Copy-on-write fork](https://en.wikipedia.org/wiki/Copy-on-write), you can follow this [issue](https://github.com/radumarias/zeroize-python/issues/1)
15-
> Also by itself it doesn't work if memory is moved or moved to swap file. You can use `crypes` with `libc.mlockall()` to lock the memory, see example below.**
14+
> **In the case of [Copy-on-write fork](https://en.wikipedia.org/wiki/Copy-on-write) you should zeroize the memory before fork the child process, see example below.
15+
> Also by itself it doesn't work if memory is moved or moved to swap file. You can use `crypes` with `libc.mlock()` to lock the memory, see example below.**
1616
17-
# Example
17+
# Examples
18+
19+
## Lock and zeroize memory
1820

1921
```python
20-
import zeroize
22+
from zeroize import zeroize1, zeroize_np
2123
import numpy as np
2224
import ctypes
2325

2426

25-
# Lock memory using ctypes
26-
def lock_memory():
27-
libc = ctypes.CDLL("libc.so.6")
28-
# Lock all current and future pages from being swapped out
29-
libc.mlockall(ctypes.c_int(0x02 | 0x04)) # MCL_CURRENT | MCL_FUTURE
27+
# Load the C standard library
28+
LIBC = ctypes.CDLL("libc.so.6")
29+
MLOCK = LIBC.mlock
30+
MUNLOCK = LIBC.munlock
31+
32+
# Define mlock and munlock argument types
33+
MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
34+
MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
35+
36+
37+
def lock_memory(buffer):
38+
"""Locks the memory of the given buffer."""
39+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
40+
size = len(buffer)
41+
if MLOCK(address, size) != 0:
42+
raise RuntimeError("Failed to lock memory")
43+
44+
45+
def unlock_memory(buffer):
46+
"""Unlocks the memory of the given buffer."""
47+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
48+
size = len(buffer)
49+
if MUNLOCK(address, size) != 0:
50+
raise RuntimeError("Failed to unlock memory")
51+
52+
53+
try:
54+
print("allocate memory")
55+
56+
# regular array
57+
arr = bytearray(b"1234567890")
58+
59+
# numpy array
60+
arr_np = np.array([0] * 10, dtype=np.uint8)
61+
arr_np[:] = arr
62+
assert arr_np.tobytes() == b"1234567890"
63+
64+
print("locking memory")
65+
66+
lock_memory(arr)
67+
lock_memory(arr_np)
68+
69+
print("zeroize'ing...: ")
70+
zeroize1(arr)
71+
zeroize_np(arr_np)
72+
73+
print("checking if is zeroized")
74+
assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
75+
assert all(arr_np == 0)
76+
77+
print("all good, bye!")
78+
finally:
79+
# Unlock the memory
80+
print("unlocking memory")
81+
unlock_memory(arr)
82+
unlock_memory(arr_np)
83+
```
84+
85+
## Zeroing memory before starting child process
86+
87+
This mitigates the problems that appears on [Copy-on-write fork](https://en.wikipedia.org/wiki/Copy-on-write). You need to zeroize the data before forking the child process.
88+
```python
89+
import os
90+
from zeroize import zeroize1, zeroize_np
91+
import numpy as np
92+
import ctypes
93+
3094

95+
# Load the C standard library
96+
LIBC = ctypes.CDLL("libc.so.6")
97+
MLOCK = LIBC.mlock
98+
MUNLOCK = LIBC.munlock
3199

32-
def unlock_memory():
33-
libc = ctypes.CDLL("libc.so.6")
34-
# Unlock all locked pages
35-
libc.munlockall()
100+
# Define mlock and munlock argument types
101+
MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
102+
MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
36103

37104

38-
print("locking memory")
39-
lock_memory()
105+
def lock_memory(buffer):
106+
"""Locks the memory of the given buffer."""
107+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
108+
size = len(buffer)
109+
if MLOCK(address, size) != 0:
110+
raise RuntimeError("Failed to lock memory")
40111

41-
print("allocate memory")
42112

43-
# regular array
44-
arr = bytearray(b"1234567890")
113+
def unlock_memory(buffer):
114+
"""Unlocks the memory of the given buffer."""
115+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
116+
size = len(buffer)
117+
if MUNLOCK(address, size) != 0:
118+
raise RuntimeError("Failed to unlock memory")
45119

46-
# numpy array
47-
arr_np = np.array([0] * 10, dtype=np.uint8)
48-
arr_np[:] = arr
49-
assert arr_np.tobytes() == b"1234567890"
50120

51-
print("zeroize'ing...: ")
52-
zeroize.zeroize1(arr)
53-
zeroize.zeroize_np(arr_np)
121+
try:
122+
sensitive_data = bytearray(b"Sensitive Information")
123+
lock_memory(sensitive_data)
54124

55-
print("checking if is zeroized")
56-
assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
57-
assert all(arr_np == 0)
125+
print("Before zeroization:", sensitive_data)
58126

59-
print("unlocking memory")
60-
unlock_memory()
127+
zeroize1(sensitive_data)
128+
print("After zeroization:", sensitive_data)
61129

62-
print("all good, bye!")
130+
# Forking after zeroization to ensure no sensitive data is copied
131+
pid = os.fork()
132+
if pid == 0:
133+
# This is the child process
134+
print("Child process memory after fork:", sensitive_data)
135+
else:
136+
# This is the parent process
137+
os.wait() # Wait for the child process to exit
138+
finally:
139+
# Unlock the memory
140+
print("unlocking memory")
141+
unlock_memory(sensitive_data)
63142
```
64143

65144
# Building from source
@@ -94,5 +173,6 @@ source .env/bin/activate
94173
pip install maturin
95174
pip install numpy
96175
maturin develop
97-
python main.py
176+
python examples/lock_and_zeroize.py
177+
python examples/zeroize_before_fork.py
98178
```

examples/lock_and_zeroize.py

+61
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
from zeroize import zeroize1, zeroize_np
2+
import numpy as np
3+
import ctypes
4+
5+
6+
# Load the C standard library
7+
LIBC = ctypes.CDLL("libc.so.6")
8+
MLOCK = LIBC.mlock
9+
MUNLOCK = LIBC.munlock
10+
11+
# Define mlock and munlock argument types
12+
MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
13+
MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
14+
15+
16+
def lock_memory(buffer):
17+
"""Locks the memory of the given buffer."""
18+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
19+
size = len(buffer)
20+
if MLOCK(address, size) != 0:
21+
raise RuntimeError("Failed to lock memory")
22+
23+
24+
def unlock_memory(buffer):
25+
"""Unlocks the memory of the given buffer."""
26+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
27+
size = len(buffer)
28+
if MUNLOCK(address, size) != 0:
29+
raise RuntimeError("Failed to unlock memory")
30+
31+
32+
try:
33+
print("allocate memory")
34+
35+
# regular array
36+
arr = bytearray(b"1234567890")
37+
38+
# numpy array
39+
arr_np = np.array([0] * 10, dtype=np.uint8)
40+
arr_np[:] = arr
41+
assert arr_np.tobytes() == b"1234567890"
42+
43+
print("locking memory")
44+
45+
lock_memory(arr)
46+
lock_memory(arr_np)
47+
48+
print("zeroize'ing...: ")
49+
zeroize1(arr)
50+
zeroize_np(arr_np)
51+
52+
print("checking if is zeroized")
53+
assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
54+
assert all(arr_np == 0)
55+
56+
print("all good, bye!")
57+
finally:
58+
# Unlock the memory
59+
print("unlocking memory")
60+
unlock_memory(arr)
61+
unlock_memory(arr_np)

examples/zeroize_before_fork.py

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import os
2+
from zeroize import zeroize1, zeroize_np
3+
import numpy as np
4+
import ctypes
5+
6+
7+
# Load the C standard library
8+
LIBC = ctypes.CDLL("libc.so.6")
9+
MLOCK = LIBC.mlock
10+
MUNLOCK = LIBC.munlock
11+
12+
# Define mlock and munlock argument types
13+
MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
14+
MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
15+
16+
17+
def lock_memory(buffer):
18+
"""Locks the memory of the given buffer."""
19+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
20+
size = len(buffer)
21+
if MLOCK(address, size) != 0:
22+
raise RuntimeError("Failed to lock memory")
23+
24+
25+
def unlock_memory(buffer):
26+
"""Unlocks the memory of the given buffer."""
27+
address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
28+
size = len(buffer)
29+
if MUNLOCK(address, size) != 0:
30+
raise RuntimeError("Failed to unlock memory")
31+
32+
33+
try:
34+
sensitive_data = bytearray(b"Sensitive Information")
35+
lock_memory(sensitive_data)
36+
37+
print("Before zeroization:", sensitive_data)
38+
39+
zeroize1(sensitive_data)
40+
print("After zeroization:", sensitive_data)
41+
42+
# Forking after zeroization to ensure no sensitive data is copied
43+
pid = os.fork()
44+
if pid == 0:
45+
# This is the child process
46+
print("Child process memory after fork:", sensitive_data)
47+
else:
48+
# This is the parent process
49+
os.wait() # Wait for the child process to exit
50+
finally:
51+
# Unlock the memory
52+
print("unlocking memory")
53+
unlock_memory(sensitive_data)
54+

main.py

-43
This file was deleted.

0 commit comments

Comments
 (0)