memorymap.rs

use std::{alloc::Layout, mem, ops::{Deref, DerefMut}};

use crate::pager::Pagable;
pub struct MemoryMap(Vec<u8>);

impl MemoryMap {
    fn page_round(size: usize) -> usize {
        Layout::from_size_align(size, Vec::<u8>::page_size()).unwrap().size()
    }

    pub fn new(size: usize) -> Result<Self, Box<dyn std::error::Error>> {
        // Round up to the nearest page size
        let rounded_size = Self::page_round(size);
        println!("Rounded {} bytes to {} bytes", size, rounded_size);
        
        let image_ptr = unsafe {
            libc::mmap(
                std::ptr::null_mut(),
                rounded_size,
                libc::PROT_READ | libc::PROT_WRITE,
                libc::MAP_PRIVATE | libc::MAP_ANON,
                -1,
                0,
            )
        };
        // Check for errors
        if image_ptr == libc::MAP_FAILED {
            return Err(format!("Failed to allocate memory: {:?}", std::io::Error::last_os_error()).into());
        }

        println!("Allocated memory at 0x{:x} ({} bytes)", image_ptr as usize, rounded_size);

        let mut image = unsafe { Vec::from_raw_parts(image_ptr as *mut u8, rounded_size, rounded_size) };
        image.fill(0);

        Ok(MemoryMap(image))
    }
}

impl Drop for MemoryMap {
    fn drop(&mut self) {
        let ptr = self.0.as_ptr();
        let len = self.0.len();
        unsafe {
            // Use `libc::munmap` to unmap the memory at `ptr` with `len` bytes.
            libc::munmap(ptr as *mut _, len);
            // Detach the inner Vec<u8> to prevent Rust from dropping it again.
            mem::forget(mem::replace(&mut self.0, Vec::new()));
        }
    }
}

impl Deref for MemoryMap {
    type Target = Vec<u8>;

    fn deref(&self) -> &Vec<u8> {
        &self.0
    }
}

impl DerefMut for MemoryMap {
    fn deref_mut(&mut self) -> &mut Vec<u8> {
        &mut self.0
    }
}

#[cfg(test)]
mod tests {
    use crate::protection::{Protectable, Protection};
    use super::*;

    #[test]
    fn it_works() {
        let mut map = MemoryMap::new(Vec::<u8>::page_size() * 2).unwrap();
        map.protect(0, Protection::build(&[Protection::Read])).unwrap();
        map[Vec::<u8>::page_size()] = 1;
        map[9999999999999] = 1;
        //println!("{}", map[0]);
    }
}

pager.rs

pub trait Pagable<T> {
    fn page_size() -> usize; // Returns the page size
    fn pages(&self) -> usize; // Returns the number of pages
    fn page(&self, page: usize) -> Option<&[T]>; // Returns a slice of the page
    fn page_mut(&mut self, page: usize) -> Option<&mut [T]>; // Returns a slice of the page
    fn offset_page(&self, offset: usize) -> Option<usize>; // Returns the page number for the given offset
}

impl<T> Pagable<T> for Vec<T> {
    // Returns the system page size
    fn page_size() -> usize {
        (unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) } as usize)
    }

    // Returns the number of pages
    fn pages(&self) -> usize {
        self.len() / Self::page_size()
    }

    fn page(&self, page: usize) -> Option<&[T]> {
        let page_size = Self::page_size();
        let offset = page * page_size;
        let end = offset + page_size;

        if end > self.len() {
            return None;
        }

        Some(&self[offset..end])
    }

    fn page_mut(&mut self, page: usize) -> Option<&mut [T]> {
        let page_size = Self::page_size();
        let offset = page * page_size;
        let end = offset + page_size;

        if end > self.len() {
            return None;
        }

        Some(&mut self[offset..end])
    }

    fn offset_page(&self, offset: usize) -> Option<usize> {
        let page_size = Self::page_size();
        let page = offset / page_size;

        if page > self.pages() {
            return None;
        }

        Some(page)
    }
}

protection.rs

use crate::pager::Pagable;

// Protection type enum, raw values are the same as the libc constants
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Protection {
    None = libc::PROT_NONE as isize,
    Read = libc::PROT_READ as isize,
    Write = libc::PROT_WRITE as isize,
    Execute = libc::PROT_EXEC as isize,
}

impl Protection {
    pub fn build(prot: &[Protection]) -> i32 {
        let mut result = 0;
        for p in prot {
            result |= *p as i32;
        }
        result
    }
}

pub trait Protectable {
    // Protect the byte given, or the page containing the byte
    fn protect(&self, offset: usize, prot: i32) -> Result<(), Box<dyn std::error::Error>>;
}

// Implementation for anything that implements the Pagable trait, must be a u8 slice since it's bytes
impl<T: Pagable<u8>> Protectable for T {
    fn protect(&self, offset: usize, prot: i32) -> Result<(), Box<dyn std::error::Error>> {
        // Get the page number for the given pointer
        let page = self.offset_page(offset).ok_or("Pointer is not in the memory map")?;

        // Get the page slice
        let page_slice = self.page(page).ok_or("Pointer is not in the memory map")?;

        // Protect the page
        let result = unsafe {
            libc::mprotect(page_slice.as_ptr() as *mut libc::c_void, Self::page_size(), prot)
        };

        // Check for errors
        if result != 0 {
            return Err(Box::new(std::io::Error::last_os_error()));
        }

        Ok(())
    }
}

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

    /*#[test]
    fn test_protection() {
        let mut image = MemoryMap::new(0x1000).unwrap();
        let page = image.page(0).unwrap();

        // Protect the page
        image.protect(page, Protection::build(&[Protection::Read])).unwrap();

        // Check that the page is protected
        assert_eq!(image.protected(page), true);

        // Unprotect the page
        image.protect(page, Protection::build(&[Protection::Read, Protection::Write])).unwrap();

        // Check that the page is unprotected
        assert_eq!(image.protected(page), false);
    }*/

    #[test]
    fn protection_values() {
        assert_eq!(Protection::None as i32, libc::PROT_NONE);
        assert_eq!(Protection::Read as i32, libc::PROT_READ);
        assert_eq!(Protection::Write as i32, libc::PROT_WRITE);
        assert_eq!(Protection::Execute as i32, libc::PROT_EXEC);
    }

    #[test]
    fn combination_protections() {
        assert_eq!(Protection::build(&[Protection::Read, Protection::Write]), libc::PROT_READ | libc::PROT_WRITE);
        assert_eq!(Protection::build(&[Protection::Read, Protection::Write, Protection::Execute]), libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC);
    }
}