lethe 0.8.0

A secure, free, cross-platform and open-source drive wiping utility
use std::ptr::slice_from_raw_parts_mut;

pub(crate) struct AlignedBuffer {
    ptr: *mut u8,
    layout: std::alloc::Layout,
}

impl AlignedBuffer {
    pub(crate) fn new(size: usize, align: usize) -> Self {
        unsafe {
            let buf_layout = std::alloc::Layout::from_size_align_unchecked(size, align);
            let buf_ptr = std::alloc::alloc(buf_layout);
            AlignedBuffer {
                ptr: buf_ptr,
                layout: buf_layout,
            }
        }
    }

    pub(crate) fn fill(&mut self, value: u8) -> () {
        unsafe { self.ptr.write_bytes(value, self.layout.size()) }
    }

    pub(crate) fn as_mut_slice(&self) -> &mut [u8] {
        unsafe { &mut *slice_from_raw_parts_mut(self.ptr, self.layout.size()) }
    }
}

impl Drop for AlignedBuffer {
    fn drop(&mut self) {
        unsafe { std::alloc::dealloc(self.ptr as *mut u8, self.layout) }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_aligned_allocation() {
        let size = 65536;
        let align = 4096;

        let buf = AlignedBuffer::new(size, align);

        assert_eq!(buf.as_mut_slice().len(), size);
        assert_eq!(buf.ptr as usize % align, 0);
    }

    #[test]
    fn test_vec_fill() {
        let mut buf = AlignedBuffer::new(1024, 1024);

        buf.fill(0xff);
        assert_eq!(buf.as_mut_slice().iter().filter(|x| **x != 0xff).count(), 0);

        buf.fill(0x11);
        assert_eq!(buf.as_mut_slice().iter().filter(|x| **x != 0x11).count(), 0);
    }
}