use sodiumoxide for mlock

radumarias · Jun 23, 2024 · b303bec · b303bec
1 parent e695866
commit b303bec
Show file tree

Hide file tree

Showing 4 changed files with 110 additions and 11 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -14,6 +14,7 @@ zeroize_rs = { version = "1.8.1", package = "zeroize"}
 numpy = "0.21"
 memsec = "0.7.0"
 region = "3.0.2"
+sodiumoxide = "0.2.7"
 
 [profile.release]
 panic = "abort"
diff --git a/src/lib.rs b/src/lib.rs
@@ -25,13 +25,10 @@ fn zeroize1(arr: &Bound<'_, PyAny>, py: Python<'_>) -> PyResult<()> {
 #[pyfunction]
 fn mlock(arr: &Bound<'_, PyAny>, py: Python<'_>) -> PyResult<()> {
     let arr = as_array_mut(arr, py)?;
-    unsafe {
-        if !_mlock(arr.as_mut_ptr(), arr.len()) {
-            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
-                "mlock failed",
-            ));
-        }
-    }
+    // if !_mlock(arr.as_mut_ptr(), arr.len()) {
+    sodiumoxide::utils::mlock(arr).map_err(|_|PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+        "mlock failed",
+    ))?;
     Ok(())
 }
 

diff --git a/tests/test_zeroize.py b/tests/test_zeroize.py
@@ -61,8 +61,7 @@ def unlock_memory(address, size):
     0.5,
     1,
     2,
-#     2.97,
-    4,
+    2.83,
 ]
 
 
@@ -111,9 +110,11 @@ def test_zeroize1_sizes(self):
                 print(f"Length of the array: {size2}")
 
                 print("lock arr")
-                lock_memory(address, size)
+                mlock(arr)
+#                 lock_memory(address, size)
                 print("lock arr2")
-                lock_memory(address2, size2)
+                mlock(arr2)
+#                 lock_memory(address2, size2)
 
                 zeroize1(arr)
                 zeroize1(arr_np)