- use array len to mlock/munlock

- reduce size to 4MB in tests, as this is how much we can lock

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zeroize"
-version = "0.4.0"
+version = "0.4.2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

README.md

@@ -6,7 +6,7 @@
 
 Securely clear secrets from memory. Built on stable Rust primitives which guarantee memory is zeroed using an operation will not be 'optimized away' by the compiler.
 
-It uses [zeroize](https://crates.io/crates/zeroize) crate under the hood to zeroize and [libsodium-sys](https://crates.io/crates/libsodium-sys) for `mlock()` and `munlock()`.  
+It uses [zeroize](https://crates.io/crates/zeroize) crate under the hood to zeroize and [libsodium-sys](https://crates.io/crates/libsodium-sys) for `mlock()` and `munlock()`. **Maximum you can mlock is 4MB**.  
 It can work with `bytearray` and `numpy array`.
 
 > [!WARNING]
@@ -29,11 +29,11 @@ if __name__ == "__main__":
         print("allocate memory")
 
         # regular array
-        # max size you can lock is 4MB, at least on Linux
+        # Maximum you can mlock is 4MB
         arr = bytearray(b"1234567890")
 
         # numpy array
-        # max size you can lock is 4MB, at least on Linux
+        # Maximum you can mlock is 4MB
         arr_np = np.array([0] * 10, dtype=np.uint8)
         arr_np[:] = arr
         assert arr_np.tobytes() == b"1234567890"
@@ -72,7 +72,7 @@ from zeroize import zeroize1, mlock, munlock
 
 if __name__ == "__main__":
     try:
-        # max size you can lock is 4MB, at least on Linux
+        # Maximum you can mlock is 4MB
         sensitive_data = bytearray(b"Sensitive Information")
         mlock(sensitive_data)
 
@@ -89,6 +89,8 @@ if __name__ == "__main__":
         else:
             # This is the parent process
             os.wait()  # Wait for the child process to exit
+
+        print("all good, bye!")
 
     finally:
         # Unlock the memory

examples/lock_and_zeroize.py

@@ -9,9 +9,11 @@
         print("allocate memory")
 
         # regular array
+        # Maximum you can mlock is 4MB
         arr = bytearray(b"1234567890")
 
         # numpy array
+        # Maximum you can mlock is 4MB
         arr_np = np.array([0] * 10, dtype=np.uint8)
         arr_np[:] = arr
         assert arr_np.tobytes() == b"1234567890"

examples/zeroize_before_fork.py

@@ -6,7 +6,7 @@
 
 if __name__ == "__main__":
     try:
-        # max size you can lock is 4MB, at least on Linux
+        # Maximum you can mlock is 4MB
         sensitive_data = bytearray(b"Sensitive Information")
         mlock(sensitive_data)
 
@@ -23,6 +23,8 @@
         else:
             # This is the parent process
             os.wait()  # Wait for the child process to exit
+
+        print("all good, bye!")
 
     finally:
         # Unlock the memory

src/lib.rs

@@ -11,8 +11,7 @@ use libsodium_sys::{
 };
 use numpy::{PyArray1, PyArrayMethods};
 use pyo3::prelude::*;
-use pyo3::pybacked::PyBackedBytes;
-use pyo3::types::{PyByteArray, PyBytes, PyCFunction};
+use pyo3::types::{PyByteArray, PyBytes};
 use zeroize_rs::Zeroize;
 
 /// The global [`sync::Once`] that ensures we only perform
@@ -48,7 +47,8 @@ fn mlock<'py>(arr: &Bound<'py, PyAny>) -> PyResult<()> {
         ));
     }
     unsafe {
-        if !_mlock(as_array_mut(arr)?.as_mut_ptr()) {
+        let arr = as_array_mut(arr)?;
+        if !_mlock(arr.as_mut_ptr(), arr.len()) {
             return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                 "mlock failed",
             ));
@@ -65,7 +65,8 @@ fn munlock<'py>(arr: &Bound<'py, PyAny>) -> PyResult<()> {
         ));
     }
     unsafe {
-        if !_munlock(as_array_mut(arr)?.as_mut_ptr()) {
+        let arr = as_array_mut(arr)?;
+        if !_munlock(arr.as_mut_ptr(), arr.len()) {
             return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                 "mlock failed",
             ));
@@ -159,13 +160,13 @@ fn init() -> bool {
 }
 
 /// Calls the platform's underlying `mlock(2)` implementation.
-unsafe fn _mlock<T>(ptr: *mut T) -> bool {
-    sodium_mlock(ptr.cast(), mem::size_of::<T>()) == 0
+unsafe fn _mlock<T>(ptr: *mut T, len: usize) -> bool {
+    sodium_mlock(ptr.cast(), len) == 0
 }
 
 /// Calls the platform's underlying `munlock(2)` implementation.
-unsafe fn _munlock<T>(ptr: *mut T) -> bool {
-    sodium_munlock(ptr.cast(), mem::size_of::<T>()) == 0
+unsafe fn _munlock<T>(ptr: *mut T, len: usize) -> bool {
+    sodium_munlock(ptr.cast(), len) == 0
 }
 
 #[cfg(test)]

tests/test_zeroize.py

@@ -12,11 +12,6 @@
     1,
     2,
     4,
-    8,
-    16,
-    32,
-    64,
-    128,
 ]
 
 
@@ -43,8 +38,9 @@ def test_zeroize1_sizes(self):
             try:
                 arr = bytearray(int(size * 1024 * 1024))
                 arr_np = np.random.randint(
-                    0, 256, int(size * 1024 * 1024), dtype=np.uint8
+                    0, 256, int(size), dtype=np.uint8
                 )
+                print(f"Testing size: {size} MB")
                 mlock(arr)
                 mlock(arr_np)
                 zeroize1(arr)