ableos-idl/ableos/src/wasm_jumploader/host_functions.rs

use core::arch;

use wasmi::{
    Error, Externals, FuncInstance, FuncRef, ModuleImportResolver, RuntimeArgs, RuntimeValue,
    Signature, Trap, ValueType,
};

pub struct HostExternals {}

const ADD_FUNC_INDEX: usize = 0;
const SEND_SIGNAL_INDEX: usize = 1;
const GET_TIME_INDEX: usize = 2;
const GET_RANDOM_INDEX: usize = 3;

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;
                println!("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)?;

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

            GET_TIME_INDEX => {
                use core::sync::atomic::Ordering::*;
                println!("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 => {
                println!("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))
            }
            _ => panic!("Unimplemented function at {}", index),
        }
    }
}
use crate::arch::generate_process_pass;
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;
                    }
                }
                return 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;
                    }
                }
                return 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;
                    }
                }
                return 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;
                    }
                }
                return true;
            }
            _ => false,
        }
    }
}

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,
            _ => {
                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))

        /*Ok(FuncInstance::alloc_host(
            Signature::new(&[ValueType::I32, ValueType::I32][..], Some(ValueType::I32)),
            index,
        ))*/
    }
}
use crate::wasm_jumploader::host_functions::ValueType::I32;