Skip to content

Commit

Permalink
docs
Browse files Browse the repository at this point in the history
  • Loading branch information
radumarias committed Jun 2, 2024
1 parent 91768e7 commit 2a1cd88
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 5 deletions.
30 changes: 26 additions & 4 deletions main.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,43 @@
import zeroize
import numpy as np
import ctypes


# Lock memory using ctypes
def lock_memory():
libc = ctypes.CDLL("libc.so.6")
# Lock all current and future pages from being swapped out
libc.mlockall(ctypes.c_int(0x02 | 0x04)) # MCL_CURRENT | MCL_FUTURE


def unlock_memory():
libc = ctypes.CDLL("libc.so.6")
# Unlock all locked pages
libc.munlockall()


print("locking memory")
lock_memory()

print("allocate memory")

# regular array
arr = bytearray(b'1234567890')
arr = bytearray(b"1234567890")

# numpy array
arr_np = np.array([0] * 10, dtype=np.uint8)
arr_np[:] = arr
assert arr_np.tobytes() == b'1234567890'
assert arr_np.tobytes() == b"1234567890"

print("zeroize'ing...: ")
zeroize.zeroize1(arr)
zeroize.zeroize_np(arr_np)

print("checking if is zeroized...")
assert arr == bytearray(b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
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("unlocking memory")
unlock_memory()

print("all good, bye!")
20 changes: 19 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
use numpy::{PyArray1, PyArrayMethods};
use pyo3::buffer::PyBuffer;
use pyo3::prelude::*;
use pyo3::types::PyByteArray;
use pyo3::types::{PyByteArray, PyMemoryView};
use zeroize_rs::Zeroize;

/// A Python module implemented in Rust.
#[pymodule]
fn zeroize(_py: Python, m: &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)?)?;
Ok(())
}

Expand All @@ -23,3 +25,19 @@ fn zeroize_np<'py>(arr: &Bound<'py, PyArray1<u8>>) -> PyResult<()> {
unsafe { arr.as_slice_mut().unwrap().zeroize(); }
Ok(())
}

// #[pyfunction]
// fn zeroize_mv<'py>(arr: &PyMemoryView, len: usize) -> PyResult<()> {
// // Get the buffer information
// let buffer = PyBuffer::<u8>::get(arr)?;
//
// // Get the raw mutable pointer and length of the memory view
// let ptr = arr.as_ptr() as *mut u8;
//
// // Create a mutable slice from the raw pointer and length
// let arr: &mut [u8] = unsafe { std::slice::from_raw_parts_mut(ptr, len) };
//
// arr.zeroize();
//
// Ok(())
// }

0 comments on commit 2a1cd88

Please sign in to comment.