ableos/ableos/src/wasm_jumploader/host_functions.rs

use crate::{arch::generate_process_pass, rhai_shell::KEYBUFF};
use wasmi::{
    Error, Externals, FuncInstance, FuncRef, ModuleImportResolver, RuntimeArgs, RuntimeValue,
    Signature, Trap, ValueType,
};

macro_rules! wasmi_type {
    (bool) => {
        ValueType::I32
    };
    (u8) => {
        ValueType::I32
    };
    (u16) => {
        ValueType::I32
    };
    (u32) => {
        ValueType::I32
    };
    (u64) => {
        ValueType::I64
    };
    (i8) => {
        ValueType::I32
    };
    (i16) => {
        ValueType::I32
    };
    (i32) => {
        ValueType::I32
    };
    (i64) => {
        ValueType::I64
    };
    (f32) => {
        ValueType::F32
    };
    (f64) => {
        ValueType::F64
    };
}

macro_rules! wasmi_return_type {
    () => {
        None
    };
    ($type: ident) => {
        Some(wasmi_type!($type))
    };
}

macro_rules! host_externals {
    (
        $($index: literal: $name: ident ($($arg_name: ident: $arg_type: ident),* $(,)?) $(-> $return_type: ident)? $block: block)*
    ) => {
        pub struct HostExternals;
        impl Externals for HostExternals {
            fn invoke_index(
                &mut self,
                index: usize,
                args: RuntimeArgs
            ) -> Result<Option<RuntimeValue>, Trap> {
                match index {
                    $(
                        $index => match args.as_ref() {
                            [$($arg_name),*] => {
                                $(
                                    let $arg_name: $arg_type = (*$arg_name)
                                        .try_into()
                                        .ok_or(wasmi::TrapKind::UnexpectedSignature)?;
                                )*
                                $block
                            },
                            _ => return Err(wasmi::TrapKind::UnexpectedSignature.into()),
                        }
                    ),*
                    _ => {
                        error!("Unimplemented function at {index}");
                        Err(Trap::new(wasmi::TrapKind::Unreachable))
                    }
                }
            }
        }

        impl ModuleImportResolver for HostExternals {
            fn resolve_func(&self, field_name: &str, signature: &Signature) -> Result<FuncRef, Error> {
                let (index, param_types, return_type): (usize, &[ValueType], Option<ValueType>) =
                    match field_name {
                        $(
                            stringify!($name) => (
                                $index,
                                &[$(wasmi_type!($arg_type)),*],
                                wasmi_return_type!($($return_type)?),
                            ),
                        )*
                        _ => {
                            return Err(Error::Instantiation(format!(
                                "Export {field_name} not found",
                            )));
                        },
                    };

                if !(param_types.len() != signature.params().len() || return_type != signature.return_type())
                && param_types
                    .iter()
                    .zip(signature.params())
                    .find(|(ty, param)| ty != param)
                    .is_none()
                {
                    trace!("Resolved export {field_name} at func {index}");
                    Ok(FuncInstance::alloc_host(signature.clone(), index))
                } else {
                    Err(Error::Instantiation(format!(
                        "Export {field_name} has a bad signature {signature:?}",
                    )))
                }
            }
        }
    };
}

host_externals! {
    0: add(a: u32, b: u32) -> u32 {
        let result = a + b;
        trace!("SYSCALL: {} + {} = {}", a, b, result);
        Ok(Some(result.into()))
    }

    1: send_signal(pid: u32, signal: u32) -> i32 {
        trace!("SYSCALL: send signal {} to pid {}", signal, pid);
        Ok(Some(0.into()))
    }

    2: get_time() -> u32 {
        use core::sync::atomic::Ordering::*;
        trace!("SYSCALL: get time");

        x86_64::instructions::interrupts::disable();
        let tick_time = crate::kmain::TICK.load(Relaxed);
        x86_64::instructions::interrupts::enable();

        let ret: u32 = tick_time.try_into().unwrap();
        Ok(Some(ret.into()))
    }

    3: get_random() -> i32 {
        trace!("SYSCALL: get random");
        let rand = generate_process_pass() as i32;
        Ok(Some(rand.into()))
    }

    4: get_input() -> i32 {
        let input = None;
        x86_64::instructions::interrupts::without_interrupts(|| KEYBUFF.lock().pop());
        if let Some(chr) = input {
            trace!("SYSCALL: input: {}", chr);
        }

        let ret = input.unwrap_or(0x00 as char) as i32;
        Ok(Some(ret.into()))
    }

    5: print_char(chr: u8) {
        trace!("SYSCALL: print: {}", chr);
        print!("{}", char::from(chr));
        Ok(None)
    }
}