1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
//! Module for resolving memory and function imports from WebAssembly module into state machine.
//!
//! Subset of functionality borrowed from tomaka/redshirt:
//! <https://github.com/tomaka/redshirt/blob/4df506f68821353a7fd67bb94c4223df6b683e1b/kernel/core/src/scheduler/vm.rs>
//!

#![allow(clippy::type_complexity)]

use alloc::string::String;
use core::{cell::RefCell, convert::TryFrom as _};

use wasmi::{Module, Signature};

#[derive(Debug)]
pub enum NewErr {
    /// Error in the interpreter.
    Interpreter(wasmi::Error),
    /// If a "memory" symbol is provided, it must be a memory.
    MemoryIsntMemory,
    /// A memory object has both been imported and exported.
    MultipleMemoriesNotSupported,
    /// If a "__indirect_function_table" symbol is provided, it must be a table.
    IndirectTableIsntTable,
}

pub struct ProcessStateMachine {
    /// Original module, with resolved imports.
    pub module: wasmi::ModuleRef,

    /// Memory of the module instantiation.
    ///
    /// Right now we only support one unique `Memory` object per process. This is it.
    /// Contains `None` if the process doesn't export any memory object, which means it doesn't use
    /// any memory.
    pub memory: Option<wasmi::MemoryRef>,

    /// Table of the indirect function calls.
    ///
    /// In WASM, function pointers are in reality indices in a table called
    /// `__indirect_function_table`. This is this table, if it exists.
    pub indirect_table: Option<wasmi::TableRef>,
}

impl ProcessStateMachine {
    /// Creates a new process state machine from the given module.
    ///
    /// The closure is called for each import that the module has. It must assign a number to each
    /// import, or return an error if the import can't be resolved. When the VM calls one of these
    /// functions, this number will be returned back in order for the user to know how to handle
    /// the call.
    ///
    /// A single main thread (whose user data is passed by parameter) is automatically created and
    /// is paused at the start of the "_start" function of the module.
    pub fn new(
        module: &Module,
        mut symbols: impl FnMut(&str, &str, &Signature) -> Result<usize, ()>,
    ) -> Result<Self, NewErr> {
        struct ImportResolve<'a> {
            func: RefCell<&'a mut dyn FnMut(&str, &str, &Signature) -> Result<usize, ()>>,
            memory: RefCell<&'a mut Option<wasmi::MemoryRef>>,
        }

        impl<'a> wasmi::ImportResolver for ImportResolve<'a> {
            fn resolve_func(
                &self,
                module_name: &str,
                field_name: &str,
                signature: &wasmi::Signature,
            ) -> Result<wasmi::FuncRef, wasmi::Error> {
                let closure = &mut **self.func.borrow_mut();
                let index = match closure(module_name, field_name, signature) {
                    Ok(i) => i,
                    Err(_) => {
                        return Err(wasmi::Error::Instantiation(
                            format!("Couldn't resolve `{module_name}`:`{field_name}`")
                        ))
                    }
                };

                Ok(wasmi::FuncInstance::alloc_host(signature.clone(), index))
            }

            fn resolve_global(
                &self,
                _module_name: &str,
                _field_name: &str,
                _global_type: &wasmi::GlobalDescriptor,
            ) -> Result<wasmi::GlobalRef, wasmi::Error> {
                Err(wasmi::Error::Instantiation(String::from(
                    "Importing globals is not supported yet",
                )))
            }

            fn resolve_memory(
                &self,
                _module_name: &str,
                _field_name: &str,
                memory_type: &wasmi::MemoryDescriptor,
            ) -> Result<wasmi::MemoryRef, wasmi::Error> {
                let mut mem = self.memory.borrow_mut();
                if mem.is_some() {
                    return Err(wasmi::Error::Instantiation(String::from(
                        "Only one memory object is supported yet",
                    )));
                }

                let new_mem = wasmi::MemoryInstance::alloc(
                    wasmi::memory_units::Pages(usize::try_from(memory_type.initial()).unwrap()),
                    memory_type
                        .maximum()
                        .map(|p| wasmi::memory_units::Pages(usize::try_from(p).unwrap())),
                )
                .unwrap();
                **mem = Some(new_mem.clone());
                Ok(new_mem)
            }

            fn resolve_table(
                &self,
                _module_name: &str,
                _field_name: &str,
                _table_type: &wasmi::TableDescriptor,
            ) -> Result<wasmi::TableRef, wasmi::Error> {
                Err(wasmi::Error::Instantiation(String::from(
                    "Importing tables is not supported yet",
                )))
            }
        }

        let (not_started, imported_memory) = {
            let mut imported_memory = None;
            let resolve = ImportResolve {
                func: RefCell::new(&mut symbols),
                memory: RefCell::new(&mut imported_memory),
            };
            let not_started =
                wasmi::ModuleInstance::new(module, &resolve).map_err(NewErr::Interpreter)?;
            (not_started, imported_memory)
        };

        // TODO: WASM has a special "start" instruction that can be used to designate a function
        // that must be executed before the module is considered initialized. It is unclear whether
        // this is intended to be a function that for example initializes global variables, or if
        // this is an equivalent of "_start". In practice, Rust never seems to generate such as
        // "start" instruction, so for now we ignore it. The code below panics if there is such
        // a "start" item, so we will fortunately not blindly run into troubles.
        let module = not_started.assert_no_start();

        let memory = if let Some(imported_mem) = imported_memory {
            if module
                .export_by_name("memory")
                .map_or(false, |m| m.as_memory().is_some())
            {
                return Err(NewErr::MultipleMemoriesNotSupported);
            }
            Some(imported_mem)
        } else if let Some(mem) = module.export_by_name("memory") {
            if let Some(mem) = mem.as_memory() {
                Some(mem.clone())
            } else {
                return Err(NewErr::MemoryIsntMemory);
            }
        } else {
            None
        };

        let indirect_table = if let Some(tbl) = module.export_by_name("__indirect_function_table") {
            if let Some(tbl) = tbl.as_table() {
                Some(tbl.clone())
            } else {
                return Err(NewErr::IndirectTableIsntTable);
            }
        } else {
            None
        };

        let state_machine = ProcessStateMachine {
            module,
            memory,
            indirect_table,
        };

        Ok(state_machine)
    }
}