insectbox 0.1.0

use ascon_aead::aead::generic_array::sequence::Split;
use ascon_aead::aead::{AeadInPlace, KeyInit};
use ascon_aead::{Ascon128, Key, Nonce};
use ascon_hash::{AsconHash, Digest};
use rand_chacha::{
    rand_core::{RngCore, SeedableRng},
    ChaCha20Rng,
};
use std::mem::ManuallyDrop;
use std::ops::{Deref, DerefMut};
use std::slice;
use zeroize::Zeroize;

use crate::pagevec::PageVec;
use crate::utils;

unsafe fn derive_secrets(region_ptr: *const u8) -> (Key<Ascon128>, Nonce<Ascon128>) {
    let s1 = slice::from_raw_parts(region_ptr, *utils::PAGE_SIZE);
    let region_p3 = region_ptr.add(*utils::PAGE_SIZE * 2);
    let s2 = slice::from_raw_parts(region_p3, *utils::PAGE_SIZE);
    let hash = <AsconHash as Digest>::new()
        .chain_update(s1)
        .chain_update(s2)
        .finalize();
    // rust magic
    hash.split()
}

fn cleanup(region_ptr: *mut u8) {
    // SAFETY: this memory was allocated using utils::alloc
    unsafe {
        utils::munlock(region_ptr, *utils::PAGE_SIZE * 3);
        utils::memzero(region_ptr, *utils::PAGE_SIZE * 3);
        utils::free(region_ptr, *utils::PAGE_SIZE * 3);
    }
}

/// An "encrypted" box type, which encrypts the inner data
///
/// In particular:
/// * a [`CryptBox<T>`] owns three entire memory pages (12 KiB on modern 
/// systems, BEWARE!)
/// * the inner data is encrypted using Ascon128
/// * the pages are 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 [`CryptBox<T>`] it is encrypted and can be decrypted
/// to yield a [`PlainBox<T>`] instance. [`PlainBox<T>`] implements [`Deref`] 
/// and other traits which allow for deref coercion, but it does not implement 
/// any associated functions. To encrypt the underlying data again use 
/// [`CryptBox::encrypt`]. This approach eliminates unnecessary imports.
///
/// Both [`CryptBox<T>`] and [`PlainBox<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 [`CryptBox<T>`] will just store the fat pointer 
/// part of such structs).
///
/// # Examples
///
/// Accessing data:
/// ```
/// use insectbox::CryptBox;
///
/// let cb = CryptBox::new(b"I'm in a CryptBox :)".to_owned());
/// // at this point we can only decrypt such a box
/// let cb = cb.decrypt();
/// // because of deref coercion we can extract the data and do stuff with it
/// println!("{:?}", cb.as_ref());
/// // after finishing we can encrypt the data again
/// let cb = CryptBox::encrypt(cb);
/// // it will be dropped now, if this was the intended end of lifetime for cb
/// // we could have left it as a PlainBox<T>
/// ```
///
/// Constructing from a constructor:
/// ```
/// use insectbox::CryptBox;
///
/// fn make_arr<const N: usize>() -> [u8; N] {
///     [0x69; N]
/// }
///
/// # fn main() {
/// // will panic if size_of::<T>() + 16 > PAGE_SIZE
/// // AVOID large numbers
/// let cb = CryptBox::construct(make_arr::<32>);
/// let cb = cb.decrypt();
///
/// assert_eq!(&[0x69; 32], cb.as_slice());
/// # }
/// ```
#[derive(Debug)]
pub struct CryptBox<T> {
    data: PageVec,
    _marker: std::marker::PhantomData<T>,
}

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

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

    /// Make a new [`CryptBox<T>`] by moving `t` into the memory owned by the 
    /// box. Returns `None` on failure.
    #[must_use]
    pub fn try_new(val: T) -> Option<CryptBox<T>> {
        if size_of::<T>() + 16 > *utils::PAGE_SIZE {
            return None;
        }
        let region_ptr = utils::alloc(*utils::PAGE_SIZE * 3)?;
        let (mut data, mut key, mut nonce) = unsafe {
            let region_slice =
                slice::from_raw_parts_mut(region_ptr.as_ptr(), *utils::PAGE_SIZE * 3);
            let mut rng = ChaCha20Rng::from_os_rng();
            rng.fill_bytes(region_slice);
            let (key, nonce) = derive_secrets(region_ptr.as_ptr());

            let region_data = region_ptr.add(*utils::PAGE_SIZE).as_ptr() as *mut T;
            region_data.write_unaligned(val);
            (
                PageVec {
                    ptr: region_data as *mut u8,
                    len: size_of::<T>(),
                },
                key,
                nonce,
            )
        };
        let cipher = Ascon128::new(&key);
        match cipher.encrypt_in_place(&nonce, b"", &mut data) {
            Ok(_) => (),
            Err(_) => {
                cleanup(region_ptr.as_ptr());
                return None;
            }
        }
        unsafe {
            utils::mlock(region_ptr.as_ptr(), *utils::PAGE_SIZE * 3);
            utils::mprotect(
                region_ptr.as_ptr(),
                *utils::PAGE_SIZE * 3,
                utils::Prot::NoAccess,
            );
        }
        key.zeroize();
        nonce.zeroize();
        Some(Self {
            data,
            _marker: std::marker::PhantomData,
        })
    }

    /// Make a new [`CryptBox<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<CryptBox<T>> {
        if size_of::<T>() + 16 > *utils::PAGE_SIZE {
            return None;
        }
        let region_ptr = utils::alloc(*utils::PAGE_SIZE * 3)?;
        let (mut data, mut key, mut nonce) = unsafe {
            let region_slice =
                slice::from_raw_parts_mut(region_ptr.as_ptr(), *utils::PAGE_SIZE * 3);
            let mut rng = ChaCha20Rng::from_os_rng();
            rng.fill_bytes(region_slice);
            let (key, nonce) = derive_secrets(region_ptr.as_ptr());

            let region_data = region_ptr.add(*utils::PAGE_SIZE).as_ptr() as *mut T;
            *region_data = f();
            (
                PageVec {
                    ptr: region_data as *mut u8,
                    len: size_of::<T>(),
                },
                key,
                nonce,
            )
        };
        let cipher = Ascon128::new(&key);
        match cipher.encrypt_in_place(&nonce, b"", &mut data) {
            Ok(_) => (),
            Err(_) => {
                cleanup(region_ptr.as_ptr());
                return None;
            }
        }
        unsafe {
            utils::mlock(region_ptr.as_ptr(), *utils::PAGE_SIZE * 3);
            utils::mprotect(
                region_ptr.as_ptr(),
                *utils::PAGE_SIZE * 3,
                utils::Prot::NoAccess,
            );
        }
        key.zeroize();
        nonce.zeroize();
        Some(Self {
            data,
            _marker: std::marker::PhantomData,
        })
    }

    /// Consumes `self` and returns an instance of [`PlainBox<T>`] which can 
    /// be used to access the inner data.
    ///
    /// # Panics
    /// This will panic if there was a decryption failure. The allocation will 
    /// be freed, unless it was the underlying memory which caused the failure.
    #[must_use]
    pub fn decrypt(self) -> PlainBox<T> {
        let mut s = ManuallyDrop::new(self);
        unsafe {
            let region_ptr = s.data.ptr.sub(*utils::PAGE_SIZE);
            utils::mprotect(region_ptr, *utils::PAGE_SIZE * 3, utils::Prot::ReadWrite);
            let (mut key, mut nonce) = derive_secrets(region_ptr);
            let cipher = Ascon128::new(&key);
            match cipher.decrypt_in_place(&nonce, b"", &mut s.data) {
                Ok(_) => (),
                Err(_) => {
                    cleanup(region_ptr);
                    panic!("CryptBox decryption failure!");
                }
            }
            key.zeroize();
            nonce.zeroize();
            let region_data = s.data.ptr as *mut T;
            PlainBox { ptr: region_data }
        }
    }

    /// Takes ownership of a [`PlainBox<T>`] and encrypts it again as a
    /// [`CryptBox<T>`].
    #[must_use]
    pub fn encrypt(s: PlainBox<T>) -> CryptBox<T> {
        let s = ManuallyDrop::new(s);
        let mut data = PageVec {
            ptr: s.ptr as *mut u8,
            len: size_of::<T>(),
        };
        unsafe {
            let region_ptr = data.ptr.sub(*utils::PAGE_SIZE);
            let (mut key, mut nonce) = derive_secrets(region_ptr);
            let cipher = Ascon128::new(&key);
            match cipher.encrypt_in_place(&nonce, b"", &mut data) {
                Ok(_) => (),
                Err(_) => {
                    cleanup(region_ptr);
                    panic!("CryptBox decryption failure!");
                }
            }
            utils::mprotect(region_ptr, *utils::PAGE_SIZE * 3, utils::Prot::NoAccess);
            key.zeroize();
            nonce.zeroize();
        }
        Self {
            data,
            _marker: std::marker::PhantomData,
        }
    }
}

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

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

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

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

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

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

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

impl<T> AsMut<T> for PlainBox<T>
where
    <PlainBox<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 PlainBox<T> {
    fn borrow(&self) -> &T {
        &**self
    }
}

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