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,
};

const ADD_FUNC_INDEX: usize = 0;
const GET_TIME_INDEX: usize = 2;
const GET_RANDOM_INDEX: usize = 3;
const GET_INPUT_INDEX: usize = 4;
const PRINT_CLEVER_HACK: usize = 5;
const SEND_SIGNAL_INDEX: usize = 1;

pub struct HostExternals;
impl HostExternals {
    fn check_signature(&self, index: usize, signature: &Signature) -> bool {
        match index {
            ADD_FUNC_INDEX => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) =
                    (&[ValueType::I32, ValueType::I32], Some(ValueType::I32));
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }

            SEND_SIGNAL_INDEX => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) =
                    (&[ValueType::I32, ValueType::I32], Some(ValueType::I32));
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }

            GET_TIME_INDEX => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) =
                    (&[], Some(ValueType::I32));
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }

            GET_RANDOM_INDEX => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) =
                    (&[], Some(ValueType::I32));
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }

            GET_INPUT_INDEX => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) =
                    (&[], Some(ValueType::I32));
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }

            PRINT_CLEVER_HACK => {
                let (params, ret_ty): (&[ValueType], Option<ValueType>) = (&[ValueType::I64], None);
                if params.len() != signature.params().len() {
                    return false;
                }
                if ret_ty != signature.return_type() {
                    return false;
                }
                for (ty, param) in params.iter().zip(signature.params()) {
                    if *ty != *param {
                        return false;
                    }
                }
                true
            }
            _ => false,
        }
    }
}

impl Externals for HostExternals {
    fn invoke_index(
        &mut self,
        index: usize,
        args: RuntimeArgs,
    ) -> Result<Option<RuntimeValue>, Trap> {
        match index {
            ADD_FUNC_INDEX => {
                let a: u32 = args.nth_checked(0)?;
                let b: u32 = args.nth_checked(1)?;
                let result = a + b;
                trace!("SYSCALL: {} + {} = {}", a, b, result);
                Ok(Some(RuntimeValue::I32(result as i32)))
            }

            SEND_SIGNAL_INDEX => {
                let pid: u32 = args.nth_checked(0)?;
                let signal: u32 = args.nth_checked(1)?;

                trace!("SYSCALL: send signal {} to pid {}", signal, pid);
                let ret = RuntimeValue::I32(0);
                Ok(Some(ret))
            }

            GET_TIME_INDEX => {
                use core::sync::atomic::Ordering::*;
                trace!("SYSCALL: get time");

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

                let ret = RuntimeValue::I32(tick_time.try_into().unwrap());
                Ok(Some(ret))
            }

            GET_RANDOM_INDEX => {
                trace!("SYSCALL: get random");
                let rand = generate_process_pass();

                let ret = RuntimeValue::I32(rand as i32);

                // let ret = RuntimeValue::I32(rand.try_into().unwrap());
                Ok(Some(ret))
            }

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

                let ret = RuntimeValue::I32(input.unwrap_or(0x00 as char) as i32);

                Ok(Some(ret))
            }

            PRINT_CLEVER_HACK => {
                // println!("SYSCALL: print clever hack");

                let combinated = args.nth_checked::<u64>(0).unwrap().to_le_bytes();

                for x in combinated.iter() {
                    let chr = *x as char;
                    print!("{}", chr);
                }

                println!("\n");
                Ok(None)
            }

            _ => {
                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 = match field_name {
            "add" => ADD_FUNC_INDEX,
            "send_signal" => SEND_SIGNAL_INDEX,
            "get_time" => GET_TIME_INDEX,
            "get_random" => GET_RANDOM_INDEX,
            "get_input" => GET_INPUT_INDEX,
            "print_clever_hack" => PRINT_CLEVER_HACK,
            _ => {
                return Err(Error::Instantiation(format!(
                    "Export {} not found",
                    field_name
                )))
            }
        };

        if !self.check_signature(index, signature) {
            return Err(Error::Instantiation(format!(
                "Export {} has a bad signature {:?}",
                field_name, signature
            )));
        }

        trace!("Resolved export {} as func {}", field_name, index);
        Ok(FuncInstance::alloc_host(signature.clone(), index))
    }
}