have Rust methods accept PyAny and convert inside

Author: radumarias

Cargo.lock

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

Cargo.toml

@@ -12,133 +12,89 @@ It can work with `bytearray` and `numpy array`.
 
 > [!WARNING]
 > **In the case of [Copy-on-write fork](https://en.wikipedia.org/wiki/Copy-on-write) you need to zeroize the memory before forking the child process, see example below.  
-> Also by itself it doesn't work if memory is moved or moved to swap. You can use `zeroize.mlock()` and `zeroize.mlock_np()` to lock the memory, max size you can lock is 4MB, at least on Linux, see example below.**
+> Also by itself it doesn't work if memory is moved or moved to swap. You can use `zeroize.mlock()` and `zeroize.mlock()` to lock the memory, max size you can lock is 4MB, at least on Linux, see example below.**
 
 # Examples
 
 ## Lock and zeroize memory
 
 ```python
-from zeroize import zeroize1, zeroize_np
-import numpy as np
-import ctypes
-
-
-# Load the C standard library
-LIBC = ctypes.CDLL("libc.so.6")
-MLOCK = LIBC.mlock
-MUNLOCK = LIBC.munlock
-
-# Define mlock and munlock argument types
-MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
-MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
-
-
-def lock_memory(buffer):
-    """Locks the memory of the given buffer."""
-    address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
-    size = len(buffer)
-    if MLOCK(address, size) != 0:
-        raise RuntimeError("Failed to lock memory")
+"""By itself it doesn't work if memory is moved or moved to swap. You can use `crypes` with `libc.mlock()` to lock the memory"""
 
+from zeroize import zeroize1, mlock, munlock
+import numpy as np
 
-def unlock_memory(buffer):
-    """Unlocks the memory of the given buffer."""
-    address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
-    size = len(buffer)
-    if MUNLOCK(address, size) != 0:
-        raise RuntimeError("Failed to unlock memory")
 
+if __name__ == "__main__":
+    try:
+        print("allocate memory")
 
-try:
-    print("allocate memory")
+        # regular array
+        # max size you can lock is 4MB, at least on Linux
+        arr = bytearray(b"1234567890")
 
-    # regular array
-    arr = bytearray(b"1234567890")
+        # numpy array
+        # max size you can lock is 4MB, at least on Linux
+        arr_np = np.array([0] * 10, dtype=np.uint8)
+        arr_np[:] = arr
+        assert arr_np.tobytes() == b"1234567890"
 
-    # numpy array
-    arr_np = np.array([0] * 10, dtype=np.uint8)
-    arr_np[:] = arr
-    assert arr_np.tobytes() == b"1234567890"
+        print("locking memory")
 
-    print("locking memory")
+        mlock(arr)
+        mlock(arr_np)
 
-    lock_memory(arr)
-    lock_memory(arr_np)
+        print("zeroize'ing...: ")
+        zeroize1(arr)
+        zeroize1(arr_np)
 
-    print("zeroize'ing...: ")
-    zeroize1(arr)
-    zeroize_np(arr_np)
+        print("checking if is zeroized")
+        assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
+        assert all(arr_np == 0)
 
-    print("checking if is zeroized")
-    assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
-    assert all(arr_np == 0)
+        print("all good, bye!")
 
-    print("all good, bye!")
-finally:
-    # Unlock the memory
-    print("unlocking memory")
-    unlock_memory(arr)
-    unlock_memory(arr_np)
+    finally:
+        # Unlock the memory
+        print("unlocking memory")
+        munlock(arr)
+        munlock(arr_np)
 ```
 
-## Zeroing memory before starting child process
+## Zeroing memory before forking child process
 
 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.
 ```python
-import os
-from zeroize import zeroize1, zeroize_np
-import numpy as np
-import ctypes
-
-
-# Load the C standard library
-LIBC = ctypes.CDLL("libc.so.6")
-MLOCK = LIBC.mlock
-MUNLOCK = LIBC.munlock
-
-# Define mlock and munlock argument types
-MLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
-MUNLOCK.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
-
-
-def lock_memory(buffer):
-    """Locks the memory of the given buffer."""
-    address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
-    size = len(buffer)
-    if MLOCK(address, size) != 0:
-        raise RuntimeError("Failed to lock memory")
+""" In the case of [Copy-on-write fork](https://en.wikipedia.org/wiki/Copy-on-write) you need to zeroize the memory before forking the child process. """
 
+import os
+from zeroize import zeroize1, mlock, munlock
 
-def unlock_memory(buffer):
-    """Unlocks the memory of the given buffer."""
-    address = ctypes.addressof(ctypes.c_char.from_buffer(buffer))
-    size = len(buffer)
-    if MUNLOCK(address, size) != 0:
-        raise RuntimeError("Failed to unlock memory")
 
+if __name__ == "__main__":
+    try:
+        # max size you can lock is 4MB, at least on Linux
+        sensitive_data = bytearray(b"Sensitive Information")
+        mlock(sensitive_data)
 
-try:
-    sensitive_data = bytearray(b"Sensitive Information")
-    lock_memory(sensitive_data)
+        print("Before zeroization:", sensitive_data)
 
-    print("Before zeroization:", sensitive_data)
+        zeroize1(sensitive_data)
+        print("After zeroization:", sensitive_data)
 
-    zeroize1(sensitive_data)
-    print("After zeroization:", sensitive_data)
+        # Forking after zeroization to ensure no sensitive data is copied
+        pid = os.fork()
+        if pid == 0:
+            # This is the child process
+            print("Child process memory after fork:", sensitive_data)
+        else:
+            # This is the parent process
+            os.wait()  # Wait for the child process to exit
 
-    # Forking after zeroization to ensure no sensitive data is copied
-    pid = os.fork()
-    if pid == 0:
-        # This is the child process
-        print("Child process memory after fork:", sensitive_data)
-    else:
-        # This is the parent process
-        os.wait()  # Wait for the child process to exit
-finally:
-    # Unlock the memory
-    print("unlocking memory")
-    unlock_memory(sensitive_data)
+    finally:
+        # Unlock the memory
+        print("unlocking memory")
+        munlock(sensitive_data)
 ```
 
 # Building from source

README.md

@@ -1,6 +1,6 @@
 """By itself it doesn't work if memory is moved or moved to swap. You can use `crypes` with `libc.mlock()` to lock the memory"""
 
-from zeroize import zeroize1, zeroize_np, mlock, munlock, mlock_np, munlock_np
+from zeroize import zeroize1, mlock, munlock
 import numpy as np
 
 
@@ -21,11 +21,11 @@
         print("locking memory")
 
         mlock(arr)
-        mlock_np(arr_np)
+        mlock(arr_np)
 
         print("zeroize'ing...: ")
         zeroize1(arr)
-        zeroize_np(arr_np)
+        zeroize1(arr_np)
 
         print("checking if is zeroized")
         assert arr == bytearray(b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00")
@@ -37,4 +37,4 @@
         # Unlock the memory
         print("unlocking memory")
         munlock(arr)
-        munlock_np(arr_np)
+        munlock(arr_np)

examples/lock_and_zeroize.py

@@ -24,78 +24,66 @@ static mut INITIALIZED: bool = false;
 #[pymodule]
 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(())
 }
 
 
 #[pyfunction]
-fn zeroize1<'py>(arr: &Bound<'py, PyByteArray>) -> PyResult<()> {
-    unsafe { arr.as_bytes_mut().zeroize(); }
+fn zeroize1<'py>(arr: &Bound<'py, PyAny>) -> PyResult<()> {
+    as_array(arr)?.zeroize();
     Ok(())
 }
 
 #[pyfunction]
-fn zeroize_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
-    unsafe { arr.as_slice_mut().unwrap().zeroize(); }
-    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")
-        }
+fn mlock<'py>(arr: &Bound<'py, PyAny>) -> PyResult<()> {
+    if !init() {
+        return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+            "libsodium failed to initialize",
+        ));
     }
-    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")
+        if !_mlock(as_array(arr)?.as_mut_ptr()) {
+            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+                "mlock failed",
+            ));
         }
     }
     Ok(())
 }
 
 #[pyfunction]
-fn munlock<'py>(arr: &Bound<'py, PyByteArray>) -> PyResult<()> {
+fn munlock<'py>(arr: &Bound<'py, PyAny>) -> PyResult<()> {
+    if !init() {
+        return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+            "libsodium failed to initialize",
+        ));
+    }
     unsafe {
-        if !init() {
-            panic!("libsodium failed to initialize")
-        }
-        if !_munlock(arr.as_bytes_mut().as_mut_ptr()) {
-            panic!("mlock failed")
+        if !_munlock(as_array(arr)?.as_mut_ptr()) {
+            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+                "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")
+fn as_array<'a>(arr: &'a Bound<PyAny>) -> PyResult<&'a mut [u8]> {
+    let arr = unsafe {
+        if let Ok(arr) = arr.downcast::<PyByteArray>() {
+            arr.as_bytes_mut()
+        } else if let Ok(arr) = arr.downcast::<PyArray1<u8>>() {
+            arr.as_slice_mut().unwrap()
+        } else {
+            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
+                "Expected a PyByteArray or PyArray1<u8>",
+            ));
         }
-    }
-    Ok(())
+    };
+    Ok(arr)
 }
 
 // #[pyfunction]
@@ -121,7 +109,7 @@ fn munlock_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
 /// not* be used.
 ///
 /// Calling it multiple times is a no-op.
-pub(crate) fn init() -> bool {
+fn init() -> bool {
     unsafe {
         INIT.call_once(|| {
             // NOTE: Calls to transmute fail to compile if the source