Skip to main content

rwasm/wasmtime/
context.rs

1use crate::{CallerTr, StoreTr, SyscallHandler, TrapCode, TypedCaller};
2use wasmtime::{AsContext, AsContextMut, StoreLimits};
3
4pub struct WrappedContext<T: 'static> {
5    pub(crate) syscall_handler: SyscallHandler<T>,
6    pub(crate) fuel: Option<u64>,
7    pub(crate) resource_limiter: StoreLimits,
8    pub(crate) data: T,
9}
10
11pub struct WasmtimeCaller<'a, T: 'static> {
12    caller: wasmtime::Caller<'a, WrappedContext<T>>,
13}
14
15impl<'a, T: 'static> WasmtimeCaller<'a, T> {
16    pub fn wrap_typed(caller: wasmtime::Caller<'a, WrappedContext<T>>) -> TypedCaller<'a, T> {
17        TypedCaller::Wasmtime(Self { caller })
18    }
19    pub fn unwrap(self) -> wasmtime::Caller<'a, WrappedContext<T>> {
20        self.caller
21    }
22}
23
24impl<'a, T: 'static> StoreTr<T> for WasmtimeCaller<'a, T> {
25    fn memory_read(&mut self, offset: usize, buffer: &mut [u8]) -> Result<(), TrapCode> {
26        let global_memory = self
27            .caller
28            .get_export("memory")
29            .unwrap_or_else(|| unreachable!("wasmtime: missing memory export, it's not possible"))
30            .into_memory()
31            .unwrap_or_else(|| unreachable!("wasmtime: missing memory export, it's not possible"));
32        global_memory
33            .read(self.caller.as_context(), offset, buffer)
34            .map_err(|_| TrapCode::MemoryOutOfBounds)
35    }
36
37    fn memory_read_into_vec(&mut self, offset: usize, length: usize) -> Result<Vec<u8>, TrapCode> {
38        let mut data = vec![0u8; length];
39        self.memory_read(offset, &mut data)?;
40        Ok(data)
41    }
42
43    fn memory_write(&mut self, offset: usize, buffer: &[u8]) -> Result<(), TrapCode> {
44        let global_memory = self
45            .caller
46            .get_export("memory")
47            .unwrap_or_else(|| unreachable!("wasmtime: missing memory export, it's not possible"))
48            .into_memory()
49            .unwrap_or_else(|| unreachable!("wasmtime: missing memory export, it's not possible"));
50        global_memory
51            .write(self.caller.as_context_mut(), offset, buffer)
52            .map_err(|_| TrapCode::MemoryOutOfBounds)
53    }
54
55    fn data_mut(&mut self) -> &mut T {
56        &mut self.caller.data_mut().data
57    }
58
59    fn data(&self) -> &T {
60        &self.caller.data().data
61    }
62
63    fn try_consume_fuel(&mut self, delta: u64) -> Result<(), TrapCode> {
64        if let Ok(remaining_fuel) = self.caller.get_fuel() {
65            let new_fuel = remaining_fuel
66                .checked_sub(delta)
67                .ok_or(TrapCode::OutOfFuel)?;
68            self.caller
69                .set_fuel(new_fuel)
70                .unwrap_or_else(|_| unreachable!("wasmtime: fuel mode is disabled in wasmtime"));
71        } else if let Some(fuel) = self.caller.data_mut().fuel.as_mut() {
72            *fuel = fuel.checked_sub(delta).ok_or(TrapCode::OutOfFuel)?;
73        }
74        Ok(())
75    }
76
77    fn remaining_fuel(&self) -> Option<u64> {
78        if let Ok(fuel) = self.caller.get_fuel() {
79            Some(fuel)
80        } else {
81            self.caller.data().fuel.as_ref().copied()
82        }
83    }
84
85    fn reset_fuel(&mut self, new_fuel_limit: u64) {
86        let has_fuel_enabled = self.caller.get_fuel().is_ok();
87        if has_fuel_enabled {
88            self.caller.set_fuel(new_fuel_limit).unwrap();
89        } else {
90            self.caller.data_mut().fuel = Some(new_fuel_limit)
91        }
92    }
93}
94
95impl<'a, T: 'static> CallerTr<T> for WasmtimeCaller<'a, T> {}