insectbox 0.1.0

//! Memory management utilities
//!
//! This module is full of unsafe code and should be used carefully.
//! The available functions are used to build the rest of the crate, but they
//! are exposed as public in case anyone would want to implement anything
//! similar or extend this crate's functionality.

use lazy_static::lazy_static;
#[cfg(windows)]
use windows_sys::Win32;

mod alloc;
mod mem;

pub use alloc::*;
pub use mem::*;

lazy_static! {
    /// The operating system's page size. Evaluated at runtime.
    pub static ref PAGE_SIZE: usize = {
        #[cfg(unix)]
        unsafe {
            libc::sysconf(libc::_SC_PAGESIZE) as usize
        }
        #[cfg(windows)]
        unsafe {
            let mut sysconf = core::mem::MaybeUninit::uninit();
            Win32::System::SystemInformation::GetSystemInfo(sysconf.as_mut_ptr());
            (*sysconf.as_ptr()).dwPageSize as usize
        }
    };
}

/// Memory protection "enum" for [`mprotect`]
///
/// This "enum" is implemented as a module because it is just a wrapper for
/// C constants.
#[cfg(unix)]
#[allow(non_snake_case, non_upper_case_globals)]
pub mod Prot {
    pub type PF = libc::c_int;

    pub const NoAccess: PF = libc::PROT_NONE;
    pub const ReadOnly: PF = libc::PROT_READ;
    pub const ReadWrite: PF = libc::PROT_READ | libc::PROT_WRITE;
}

/// Memory protection "enum" for [`mprotect`]
///
/// This "enum" is implemented as a module because it is just a wrapper for
/// C constants.
#[cfg(windows)]
#[allow(non_snake_case, non_upper_case_globals)]
pub mod Prot {
    use windows_sys::Win32;

    pub type PF = Win32::System::Memory::PAGE_PROTECTION_FLAGS;

    pub const NoAccess: PF = Win32::System::Memory::PAGE_NOACCESS;
    pub const ReadOnly: PF = Win32::System::Memory::PAGE_READONLY;
    pub const ReadWrite: PF = Win32::System::Memory::PAGE_READWRITE;
}

#[cfg(unix)]
#[inline]
unsafe fn __mprotect(start: *mut u8, n: usize, prot: Prot::PF) -> bool {
    libc::mprotect(start.cast(), n, prot) == 0
}

#[cfg(windows)]
#[inline]
unsafe fn __mprotect(start: *mut u8, n: usize, prot: Prot::PF) -> bool {
    let mut dump = std::mem::MaybeUninit::uninit();
    let ret = Win32::System::Memory::VirtualProtect(start.cast(), n, prot, dump.as_mut_ptr());
    dump.assume_init_drop();
    ret != 0
}

/// Protect memory pages with a protection flag
///
/// This function protects all pages spanned by the range
/// [`start`, `start + n - 1`) with a flag `prot` from [`Prot`].
///
/// # Safety
/// * `start` has to be aligned to a page boundary
/// * The entire range MUST be part of a single allocation made by
/// [`alloc`](crate::utils::alloc::alloc)
/// * `n` has to be bigger than one (`n > 1`)
///
/// # Note
/// Even if setting `n` to 2 would protect a single page just fine, it is
/// not guaranteed. This was conlcuded from testing - at least one contrary
/// example was found. A better practice is to set `n` to `k * PAGE_SIZE()`
/// for `k` equal to the number of pages one would wish to protect.
pub unsafe fn mprotect(start: *mut u8, n: usize, prot: Prot::PF) -> bool {
    __mprotect(start, n, prot)
}

/// Prevent the memory from being paged to the disk
///
/// This function locks all pages spanned by the range
/// [`start`, `start + n - 1`)
///
/// In addition on `linux`, `freebsd` and `dragonfly` targets advise the
/// kernel not to include the memory area in core dumps.
///
/// # Safety
/// * The entire range MUST be part of a single allocation made by
/// [`alloc`](crate::utils::alloc::alloc)
/// * `n` has to be bigger than one (`n > 1`)
pub unsafe fn mlock(start: *mut u8, n: usize) -> bool {
    #[cfg(unix)]
    {
        #[cfg(target_os = "linux")]
        libc::madvise(start.cast(), n, libc::MADV_DONTDUMP);

        #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
        libc::madvise(start.cast(), n, libc::MADV_NOCORE);

        libc::mlock(start.cast(), n) == 0
    }
    #[cfg(windows)]
    {
        Win32::System::Memory::VirtualLock(start.cast(), n) != 0
    }
}

/// Reverse the changes made by [`mlock`]
///
/// This function unlocks all pages spanned by the range
/// [`start`, `start + n - 1`), which were previously locked using [`mlock`].
///
/// # Safety
/// * The entire range MUST be part of a single allocation made by
/// [`alloc`](crate::utils::alloc::alloc)
/// * `n` has to be bigger than one (`n > 1`)
pub unsafe fn munlock(start: *mut u8, n: usize) -> bool {
    #[cfg(unix)]
    {
        #[cfg(target_os = "linux")]
        libc::madvise(start.cast(), n, libc::MADV_DODUMP);

        #[cfg(any(target_os = "freebsd", target_os = "dragonfly"))]
        libc::madvise(start.cast(), n, libc::MADV_DOCORE);

        libc::munlock(start.cast(), n) == 0
    }
    #[cfg(windows)]
    {
        Win32::System::Memory::VirtualUnlock(start.cast(), n) != 0
    }
}