insectbox 0.1.0

use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};

use crate::utils;

/// A "secure" box type, which shields the inner data from buffer overreads
///
/// [`SecBox<T>`] wraps a `*mut T` and provides a safe interface using 
/// functions defined in the [`utils`] module. In particular:
/// * a [`SecBox<T>`] owns an entire memory page
/// * the page is protected using [`mprotect`](utils::mprotect) by the flag 
/// [`Prot::NoAccess`](utils::Prot)
/// * the memory will not be paged to the disk because of 
/// [`mlock`](utils::mlock)
///
/// After creating a [`SecBox<T>`] it is in a locked state and can be unlocked 
/// to yield a [`InsecBox<T>`] instance. [`InsecBox<T>`] implements [`Deref`] 
/// and other traits which allow for deref coercion, but it does not implement 
/// any associated functions. To close access to the underlying data again use 
/// [`SecBox::secure`]. This approach eliminates unnecessary imports.
///
/// Both [`SecBox<T>`] and [`InsecBox<T>`] implement [`Drop`], which clears the 
/// memory they were using and sets it to zero, erasing all leftover data.
///
/// # Note
///
/// This is not an allocator, thus storing types like [`Vec<T>`] or [`String`] 
/// does not make any sense (the [`SecBox<T>`] will just store the fat pointer 
/// part of such structs).
///
/// # Examples
///
/// Accessing data:
/// ```
/// use insectbox::SecBox;
///
/// let sb = SecBox::new(b"I'm in a SecBox :)".to_owned());
/// // at this point we can only unlock such a box
/// let sb = sb.unlock();
/// // because of deref coercion we can extract the data and do stuff with it
/// println!("{:?}", sb.as_ref());
/// // after finishing we can secure the data again
/// let sb = SecBox::secure(sb);
/// // it will be dropped now, if this was the intended end of lifetime for sb
/// // we could have left it as an InsecBox<T>
/// ```
///
/// Constructing from a constructor:
/// ```
/// use insectbox::SecBox;
///
/// fn make_arr<const N: usize>() -> [u8; N] {
///     [0x69; N]
/// }
///
/// # fn main() {
/// // will panic if size_of::<T>() > PAGE_SIZE
/// // AVOID large numbers
/// let sb = SecBox::construct(make_arr::<16>);
/// let sb = sb.unlock();
///
/// assert_eq!(&[0x69; 16], sb.as_slice());
/// # }
/// ```
#[derive(Debug)]
pub struct SecBox<T> {
    ptr: *mut T,
}

impl<T> SecBox<T> {
    /// Make a new [`SecBox<T>`] by moving `t` into the memory owned by the 
    /// secure box.
    ///
    /// # Panics
    /// This function panics if `size_of::<T>() > PAGE_SIZE` or if the memory 
    /// allocation failed.
    #[must_use]
    #[inline]
    pub fn new(val: T) -> SecBox<T> {
        Self::try_new(val).expect("SecBox memory allocation failed!")
    }

    /// Make a new [`SecBox<T>`] by calling `f` and assigning the result to 
    /// the inner pointer.
    ///
    /// # Panics
    /// This function panics if `size_of::<T>() > PAGE_SIZE` or if the memory 
    /// allocation failed.
    #[must_use]
    #[inline]
    pub fn construct<F: Fn() -> T>(f: F) -> SecBox<T> {
        Self::try_construct(f).expect("SecBox memory allocation failed!")
    }

    /// Make a new [`SecBox<T>`] by moving `t` into the memory owned by the 
    /// secure box. Returns `None` on failure.
    #[must_use]
    pub fn try_new(val: T) -> Option<SecBox<T>> {
        if size_of::<T>() > *utils::PAGE_SIZE {
            return None;
        }
        let region = utils::alloc(*utils::PAGE_SIZE)?;
        let ptr = unsafe {
            utils::memzero(region.as_ptr(), *utils::PAGE_SIZE);
            let region = region.as_ptr() as *mut T;
            // SAFETY: we are aligned already and we have enough space
            region.write_unaligned(val);
            utils::mlock(region.cast(), *utils::PAGE_SIZE);
            utils::mprotect(region.cast(), *utils::PAGE_SIZE, utils::Prot::NoAccess);
            region
        };
        Some(Self { ptr })
    }

    /// Make a new [`SecBox<T>`] by calling `f` and assigning the result to 
    /// the inner pointer. Returns `None` on failure.
    #[must_use]
    pub fn try_construct<F: Fn() -> T>(f: F) -> Option<SecBox<T>> {
        if size_of::<T>() > *utils::PAGE_SIZE {
            return None;
        }
        let region = utils::alloc(*utils::PAGE_SIZE)?;
        let ptr = unsafe {
            utils::memzero(region.as_ptr(), *utils::PAGE_SIZE);
            let region = region.as_ptr() as *mut T;
            // SAFETY: we are aligned already and we have enough space
            *region = f();
            utils::mlock(region.cast(), *utils::PAGE_SIZE);
            utils::mprotect(region.cast(), *utils::PAGE_SIZE, utils::Prot::NoAccess);
            region
        };
        Some(Self { ptr })
    }

    /// Consumes `self` and returns an instance of [`InsecBox<T>`] which can 
    /// be used to access the inner data.
    #[must_use]
    pub fn unlock(self) -> InsecBox<T> {
        let u = ManuallyDrop::new(self);
        // SAFETY: this memory was allocated using utils::alloc
        unsafe {
            utils::mprotect(u.ptr.cast(), *utils::PAGE_SIZE, utils::Prot::ReadWrite);
        }
        InsecBox { ptr: u.ptr }
    }

    /// Takes ownership of an [`InsecBox<T>`] and protects it again as a
    /// [`SecBox<T>`].
    #[must_use]
    pub fn secure(s: InsecBox<T>) -> SecBox<T> {
        let s = ManuallyDrop::new(s);
        // SAFETY: this memory was allocated using utils::alloc
        unsafe {
            utils::mprotect(s.ptr.cast(), *utils::PAGE_SIZE, utils::Prot::NoAccess);
        }
        SecBox { ptr: s.ptr }
    }
}

impl<T> Drop for SecBox<T> {
    fn drop(&mut self) {
        let region = self.ptr as *mut u8;
        // SAFETY: this memory was allocated using utils::alloc
        unsafe {
            utils::mprotect(region, *utils::PAGE_SIZE, utils::Prot::ReadWrite);
            utils::munlock(region, *utils::PAGE_SIZE);
            utils::memzero(region, *utils::PAGE_SIZE);
            utils::free(region, *utils::PAGE_SIZE);
        }
    }
}

/// An unprotected pointer to the data inside a [`SecBox<T>`]
///
/// If an [`InsecBox<T>`] exists, it means that the data stored in a secure 
/// box is no longer protected. Since [`InsecBox<T>`] is owned, it upholds all 
/// of Rust's aliasing rules. Because of that the following example doesn't 
/// compile:
/// ```compile_fail
/// use insectbox::SecBox;
///
/// let sb = SecBox::new([5_u8; 32]);
/// let sb = sb.unlock();
/// let slice = &sb[..16];
/// let sb = SecBox::secure(sb);
/// println!("{:?}", slice);
/// ```
///
/// Look at [`SecBox<T>`] documentation for more details.
#[derive(Debug)]
pub struct InsecBox<T> {
    ptr: *mut T,
}

impl<T> Drop for InsecBox<T> {
    fn drop(&mut self) {
        let region = self.ptr as *mut u8;
        // SAFETY: this memory was allocated using utils::alloc
        unsafe {
            utils::munlock(region, *utils::PAGE_SIZE);
            utils::memzero(region, *utils::PAGE_SIZE);
            utils::free(region, *utils::PAGE_SIZE);
        }
    }
}

// For all the following traits
// SAFETY: self.ptr is aligned, and points to a value of T
impl<T> Deref for InsecBox<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.ptr }
    }
}

impl<T> DerefMut for InsecBox<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.ptr }
    }
}

impl<T> AsRef<T> for InsecBox<T>
where
    <InsecBox<T> as Deref>::Target: AsRef<T>,
{
    fn as_ref(&self) -> &T {
        self.deref().as_ref()
    }
}

impl<T> AsMut<T> for InsecBox<T>
where
    <InsecBox<T> as Deref>::Target: AsMut<T>,
{
    fn as_mut(&mut self) -> &mut T {
        self.deref_mut().as_mut()
    }
}

impl<T> std::borrow::Borrow<T> for InsecBox<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

impl<T> std::borrow::BorrowMut<T> for InsecBox<T> {
    fn borrow_mut(&mut self) -> &mut T {
        &mut **self
    }
}