ableOS_v1Change/kernel/src/holeybytes/kernel_services/mem_serve.rs

use {
    crate::holeybytes::{kernel_services::block_read, Vm},
    alloc::alloc::{alloc, dealloc},
    core::alloc::Layout,
    log::{debug, info},
};

pub enum MemoryServiceError {
    InvalidMemoryFormat,
}
#[derive(Debug)]
pub enum MemoryQuotaType {
    NoQuota = 0,
    SoftQuota = 1,
    HardQuota = 2,
    KillQuota = 3,
}

fn alloc_page(vm: &mut Vm, _mem_addr: u64, _length: usize) -> Result<(), MemoryServiceError> {
    let ptr = unsafe { alloc(Layout::from_size_align_unchecked(4096, 8)) };
    info!("Block address: {:?}", ptr);
    vm.registers[1] = hbvm::value::Value(ptr as u64);
    vm.registers[2] = hbvm::value::Value(4096);
    Ok(())
}

#[inline(always)]
unsafe fn memset(mut dest: *mut u8, src: *const u8, count: usize, size: usize) {
    const BLOCK_SIZE: usize = 64;
    let mut remaining = count * size;

    if remaining < 16 {
        src.copy_to_nonoverlapping(dest, remaining);
        return;
    }

    let mut buffer = [0u8; BLOCK_SIZE];
    let mut buffer_size = size;
    src.copy_to_nonoverlapping(buffer.as_mut_ptr(), size);

    while core::intrinsics::likely(buffer_size * 2 <= BLOCK_SIZE) {
        buffer
            .as_mut_ptr()
            .copy_to_nonoverlapping(buffer.as_mut_ptr().add(buffer_size), buffer_size);
        buffer_size *= 2;
    }

    let buffer_ptr = buffer.as_ptr() as *const u64;

    while (dest as usize) & 7 != 0 && remaining >= 8 {
        buffer.as_ptr().copy_to_nonoverlapping(dest, 1);
        dest = dest.add(1);
        remaining -= 1;
    }

    while core::intrinsics::likely(remaining >= 8) {
        *(dest as *mut u64) = *buffer_ptr;
        dest = dest.add(8);
        remaining -= 8;
    }

    if remaining > 0 {
        buffer.as_ptr().copy_to_nonoverlapping(dest, remaining);
    }
}

#[inline(always)]
pub fn memory_msg_handler(
    vm: &mut Vm,
    mem_addr: u64,
    length: usize,
) -> Result<(), MemoryServiceError> {
    let msg_vec = block_read(mem_addr, length);
    let msg_type = msg_vec[0];
    match msg_type {
        0 => unsafe {
            let page_count = msg_vec[1];

            let ptr = alloc(Layout::from_size_align_unchecked(
                page_count as usize * 4096,
                8,
            ));

            log::debug!("Allocating {} pages @ {:x}", page_count, ptr as u64);

            vm.registers[1] = hbvm::value::Value(ptr as u64);
            log::debug!("Kernel ptr: {:x}", ptr as u64);
        },

        1 => unsafe {
            let page_count = msg_vec[1];

            let mptr_raw: [u8; 8] = msg_vec[2..10].try_into().unwrap();
            let mptr: u64 = u64::from_le_bytes(mptr_raw);
            log::debug!("Deallocating {} pages @ {:x}", page_count, mptr);

            dealloc(
                mptr as *mut u8,
                Layout::from_size_align_unchecked(page_count as usize * 4096, 8),
            )
        },
        2 => {
            use MemoryQuotaType::*;
            let quota_type = match msg_vec[1] {
                0 => NoQuota,
                1 => SoftQuota,
                2 => HardQuota,
                3 => KillQuota,
                _ => NoQuota,
            };
            let hid_raw: [u8; 8] = msg_vec[2..10].try_into().unwrap();
            let hid: u64 = u64::from_le_bytes(hid_raw);
            let pid_raw: [u8; 8] = msg_vec[10..18].try_into().unwrap();
            let pid: u64 = u64::from_le_bytes(pid_raw);
            debug!(
                "Setting HID-{:x}:PID-{:x}'s quota type to {:?}",
                hid, pid, quota_type
            )
        }
        3 => {
            let page_count = msg_vec[1];
            log::debug!(" {} pages", page_count);
        }
        4 => unsafe {
            let count = u32::from_le_bytes(msg_vec[1..5].try_into().unwrap_unchecked()) as usize;
            let src = u64::from_le_bytes(msg_vec[5..13].try_into().unwrap_unchecked()) as *const u8;
            let dest = u64::from_le_bytes(msg_vec[13..21].try_into().unwrap_unchecked()) as *mut u8;

            src.copy_to_nonoverlapping(dest, count);
        },
        5 => unsafe {
            let count = u32::from_le_bytes(msg_vec[1..5].try_into().unwrap_unchecked()) as usize;
            let size = u32::from_le_bytes(msg_vec[5..9].try_into().unwrap_unchecked()) as usize;
            let src = u64::from_le_bytes(msg_vec[9..17].try_into().unwrap_unchecked()) as *const u8;
            let dest = u64::from_le_bytes(msg_vec[17..25].try_into().unwrap_unchecked()) as *mut u8;

            memset(dest, src, count, size);
        },
        _ => {
            log::debug!("Unknown memory service message type: {}", msg_type);
        }
    }

    Ok(())
}