use libsodium-sys for mlock/munlock

Author: radumarias

Cargo.lock

@@ -11,4 +11,6 @@ crate-type = ["cdylib"]
 [dependencies]
 pyo3 = "0.21.2"
 zeroize_rs = { version = "1.8.1", package = "zeroize"}
-numpy = "0.21"
+numpy = "0.21"
+libsodium-sys = "0.2.7"
+libc = "0.2.155"

Cargo.toml

@@ -1,15 +1,35 @@
+use std::mem;
+use std::sync::Once;
+
+use libc::{self, size_t};
+use libsodium_sys::{
+    sodium_init
+    , sodium_mlock
+    , sodium_munlock,
+};
 use numpy::{PyArray1, PyArrayMethods};
-use pyo3::buffer::PyBuffer;
 use pyo3::prelude::*;
-use pyo3::types::{PyByteArray, PyMemoryView};
+use pyo3::types::{PyByteArray, PyCFunction};
 use zeroize_rs::Zeroize;
 
+/// The global [`sync::Once`] that ensures we only perform
+/// library initialization one time.
+static INIT: Once = Once::new();
+
+/// A flag that returns whether this library has been safely
+/// initialized.
+static mut INITIALIZED: bool = false;
+
 /// A Python module implemented in Rust.
 #[pymodule]
-fn zeroize(_py: Python, m: &PyModule) -> PyResult<()> {
+fn zeroize<'py>(py: Python, m: &Bound<'py, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(zeroize1, m)?)?;
     m.add_function(wrap_pyfunction!(zeroize_np, m)?)?;
     // m.add_function(wrap_pyfunction!(zeroize_mv, m)?)?;
+    m.add_function(wrap_pyfunction!(mlock, m)?)?;
+    m.add_function(wrap_pyfunction!(munlock, m)?)?;
+    m.add_function(wrap_pyfunction!(mlock_np, m)?)?;
+    m.add_function(wrap_pyfunction!(munlock_np, m)?)?;
     Ok(())
 }
 
@@ -26,6 +46,58 @@ fn zeroize_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
     Ok(())
 }
 
+#[pyfunction]
+fn mlock<'py>(arr: &Bound<'py, PyByteArray>) -> PyResult<()> {
+    unsafe {
+        if !init() {
+            panic!("libsodium failed to initialize")
+        }
+        if !_mlock(arr.as_bytes_mut().as_mut_ptr()) {
+            panic!("mlock failed")
+        }
+    }
+    Ok(())
+}
+
+#[pyfunction]
+fn mlock_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
+    unsafe {
+        if !init() {
+            panic!("libsodium failed to initialize")
+        }
+        if !_mlock(arr.as_slice_mut().unwrap().as_mut_ptr()) {
+            panic!("mlock failed")
+        }
+    }
+    Ok(())
+}
+
+#[pyfunction]
+fn munlock<'py>(arr: &Bound<'py, PyByteArray>) -> PyResult<()> {
+    unsafe {
+        if !init() {
+            panic!("libsodium failed to initialize")
+        }
+        if !_munlock(arr.as_bytes_mut().as_mut_ptr()) {
+            panic!("mlock failed")
+        }
+    }
+    Ok(())
+}
+
+#[pyfunction]
+fn munlock_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
+    unsafe {
+        if !init() {
+            panic!("libsodium failed to initialize")
+        }
+        if !_munlock(arr.as_slice_mut().unwrap().as_mut_ptr()) {
+            panic!("mlock failed")
+        }
+    }
+    Ok(())
+}
+
 // #[pyfunction]
 // fn zeroize_mv<'py>(arr: &PyMemoryView, len: usize) -> PyResult<()> {
 //     // Get the buffer information
@@ -41,3 +113,61 @@ fn zeroize_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
 //
 //     Ok(())
 // }
+
+/// Initialized libsodium. This function *must* be called at least once
+/// prior to using any of the other functions in this library, and
+/// callers *must* verify that it returns `true`. If it returns `false`,
+/// libsodium was unable to be properly set up and this library *must
+/// not* be used.
+///
+/// Calling it multiple times is a no-op.
+pub(crate) fn init() -> bool {
+    unsafe {
+        INIT.call_once(|| {
+            // NOTE: Calls to transmute fail to compile if the source
+            // and destination type have a different size. We (ab)use
+            // this fact to statically assert the size of types at
+            // compile-time.
+            //
+            // We assume that we can freely cast between rust array
+            // sizes and [`libc::size_t`]. If that's not true, DO NOT
+            // COMPILE.
+            #[allow(clippy::useless_transmute)]
+                let _ = std::mem::transmute::<usize, size_t>(0);
+
+            let mut failure = false;
+
+            // sodium_init returns 0 on success, -1 on failure, and 1 if
+            // the library is already initialized; someone else might
+            // have already initialized it before us, so we only care
+            // about failure
+            failure |= sodium_init() == -1;
+
+            INITIALIZED = !failure;
+        });
+
+        INITIALIZED
+    }
+}
+
+/// 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
+}
+
+/// 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
+}
+
+#[cfg(test)]
+mod test {
+    use zeroize_rs::Zeroize;
+
+    #[test]
+    fn test_zeroize() {
+        let mut arr = [1, 2, 3, 4, 5];
+        arr.zeroize();
+        assert_eq!(arr, [0, 0, 0, 0, 0]);
+    }
+}

src/lib.rs

@@ -1,5 +1,5 @@
 import unittest
-import zeroize
+from zeroize import zeroize1, zeroize_np, mlock, munlock, mlock_np, munlock_np
 import numpy as np
 import ctypes
 import platform
@@ -121,52 +121,52 @@ class TestStringMethods(unittest.TestCase):
     def test_zeroize1(self):
         try:
             arr = bytearray(b"1234567890")
-            lock_memory(arr)
-            zeroize.zeroize1(arr)
+            mlock(arr)
+            zeroize1(arr)
             self.assertEqual(
                 arr, bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
             )
 
         finally:
-            unlock_memory(arr)
+            munlock(arr)
 
     def test_zeroize_np(self):
         try:
             arr = np.array([0] * 10, dtype=np.uint8)
-            lock_memory(arr)
+            mlock_np(arr)
             arr[:] = bytes(b"1234567890")
-            zeroize.zeroize_np(arr)
+            zeroize_np(arr)
             self.assertEqual(True, all(arr == 0))
 
         finally:
-            unlock_memory(arr)
+            munlock_np(arr)
 
     def test_zeroize1_sizes(self):
         for size in SIZES_MB:
             try:
                 arr = bytearray(int(size * 1024 * 1024))
-                lock_memory(arr)
-                zeroize.zeroize1(arr)
+                mlock(arr)
+                zeroize1(arr)
                 self.assertEqual(arr, bytearray(int(size * 1024 * 1024)))
 
             finally:
-                unlock_memory(arr)
+                munlock(arr)
 
     def test_zeroize_np_sizes(self):
         for size in SIZES_MB:
             try:
                 array_size = int(size * 1024 * 1024)
                 random_array = np.random.randint(0, 256, array_size, dtype=np.uint8)
-                lock_memory(random_array)
-                zeroize.zeroize_np(random_array)
+                mlock_np(random_array)
+                zeroize_np(random_array)
                 self.assertEqual(True, all(random_array == 0))
             finally:
-                unlock_memory(random_array)
+                munlock_np(random_array)
 
 
 if __name__ == "__main__":
-    if os_name == "Windows":
-        # Enable the privilege to lock memory
-        enable_lock_memory_privilege()
+    # if os_name == "Windows":
+        # # Enable the privilege to lock memory
+        # enable_lock_memory_privilege()
 
     unittest.main()