extism_runtime/
plugin.rs

1use std::collections::BTreeMap;
2
3use crate::*;
4
5/// Plugin contains everything needed to execute a WASM function
6pub struct Plugin {
7    /// All modules that were provided to the linker
8    pub modules: BTreeMap<String, Module>,
9
10    /// Used to define functions and create new instances
11    pub linker: Linker<Internal>,
12    pub store: Store<Internal>,
13
14    /// Instance provides the ability to call functions in a module
15    pub instance: Option<Instance>,
16    pub instance_pre: InstancePre<Internal>,
17
18    /// Keep track of the number of times we're instantiated, this exists
19    /// to avoid issues with memory piling up since `Instance`s are only
20    /// actually cleaned up along with a `Store`
21    instantiations: usize,
22
23    /// The ID used to identify this plugin with the `Timer`
24    pub timer_id: uuid::Uuid,
25
26    /// A handle used to cancel execution of a plugin
27    pub(crate) cancel_handle: sdk::ExtismCancelHandle,
28
29    /// Runtime determines any initialization functions needed
30    /// to run a module
31    pub(crate) runtime: Option<Runtime>,
32}
33
34impl InternalExt for Plugin {
35    fn store(&self) -> &Store<Internal> {
36        &self.store
37    }
38
39    fn store_mut(&mut self) -> &mut Store<Internal> {
40        &mut self.store
41    }
42
43    fn linker(&self) -> &Linker<Internal> {
44        &self.linker
45    }
46
47    fn linker_mut(&mut self) -> &mut Linker<Internal> {
48        &mut self.linker
49    }
50
51    fn linker_and_store(&mut self) -> (&mut Linker<Internal>, &mut Store<Internal>) {
52        (&mut self.linker, &mut self.store)
53    }
54}
55
56const EXPORT_MODULE_NAME: &str = "env";
57
58fn profiling_strategy() -> ProfilingStrategy {
59    match std::env::var("EXTISM_PROFILE").as_deref() {
60        Ok("perf") => ProfilingStrategy::PerfMap,
61        Ok(x) => {
62            log::warn!("Invalid value for EXTISM_PROFILE: {x}");
63            ProfilingStrategy::None
64        }
65        Err(_) => ProfilingStrategy::None,
66    }
67}
68
69fn calculate_available_memory(
70    available_pages: &mut Option<u32>,
71    modules: &BTreeMap<String, Module>,
72) -> anyhow::Result<()> {
73    let available_pages = match available_pages {
74        Some(p) => p,
75        None => return Ok(()),
76    };
77
78    let max_pages = *available_pages;
79    let mut fail_memory_check = false;
80    let mut total_memory_needed = 0;
81    for (name, module) in modules.iter() {
82        if name == "env" {
83            continue;
84        }
85        let mut memories = 0;
86        for export in module.exports() {
87            if let Some(memory) = export.ty().memory() {
88                memories += 1;
89                let memory_max = memory.maximum();
90                match memory_max {
91                    None => anyhow::bail!("Unbounded memory in module {name}, when `memory.max_pages` is set in the manifest all modules \
92                                           must have a maximum bound set on an exported memory"),
93                    Some(m) => {
94                        total_memory_needed += m;
95                        if !fail_memory_check {
96                            continue;
97                        }
98
99                        *available_pages = available_pages.saturating_sub(m as u32);
100                        if *available_pages == 0 {
101                            fail_memory_check = true;
102                        }
103                    }
104                }
105            }
106        }
107
108        if memories == 0 {
109            anyhow::bail!("No memory exported from module {name}, when `memory.max_pages` is set in the manifest all modules must \
110                           have a maximum bound set on an exported memory");
111        }
112    }
113
114    if fail_memory_check {
115        anyhow::bail!("Not enough memory configured to run the provided plugin, `memory.max_pages` is set to {max_pages} in the manifest \
116                       but {total_memory_needed} pages are needed by the plugin");
117    }
118
119    Ok(())
120}
121
122impl Plugin {
123    /// Create a new plugin from the given WASM code
124    pub fn new<'a>(
125        wasm: impl AsRef<[u8]>,
126        imports: impl IntoIterator<Item = &'a Function>,
127        with_wasi: bool,
128    ) -> Result<Plugin, Error> {
129        // Create a new engine, if the `EXITSM_DEBUG` environment variable is set
130        // then we enable debug info
131        let engine = Engine::new(
132            Config::new()
133                .epoch_interruption(true)
134                .debug_info(std::env::var("EXTISM_DEBUG").is_ok())
135                .profiler(profiling_strategy())
136                .cache_config_load_default()?,
137        )?;
138        let mut imports = imports.into_iter();
139        let (manifest, modules) = Manifest::new(&engine, wasm.as_ref())?;
140
141        // Calculate how much memory is available based on the value of `max_pages` and the exported
142        // memory of the modules. An error will be returned if a module doesn't have an exported memory
143        // or there is no maximum set for a module's exported memory.
144        let mut available_pages = manifest.as_ref().memory.max_pages;
145        calculate_available_memory(&mut available_pages, &modules)?;
146        log::trace!("Available pages: {available_pages:?}");
147
148        let mut store = Store::new(
149            &engine,
150            Internal::new(manifest, with_wasi, available_pages)?,
151        );
152        store.epoch_deadline_callback(|_internal| Ok(wasmtime::UpdateDeadline::Continue(1)));
153
154        if available_pages.is_some() {
155            store.limiter(|internal| internal.memory_limiter.as_mut().unwrap());
156        }
157        let mut linker = Linker::new(&engine);
158        linker.allow_shadowing(true);
159
160        // If wasi is enabled then add it to the linker
161        if with_wasi {
162            wasmtime_wasi::add_to_linker(&mut linker, |x: &mut Internal| {
163                &mut x.wasi.as_mut().unwrap().ctx
164            })?;
165
166            #[cfg(feature = "nn")]
167            wasmtime_wasi_nn::add_to_linker(&mut linker, |x: &mut Internal| {
168                &mut x.wasi.as_mut().unwrap().nn
169            })?;
170        }
171        // Get the `main` module, or the last one if `main` doesn't exist
172        let (main_name, main) = modules.get("main").map(|x| ("main", x)).unwrap_or_else(|| {
173            let entry = modules.iter().last().unwrap();
174            (entry.0.as_str(), entry.1)
175        });
176
177        // Define PDK functions
178        macro_rules! define_funcs {
179            ($m:expr, { $($name:ident($($args:expr),*) $(-> $($r:expr),*)?);* $(;)?}) => {
180                match $m {
181                $(
182                    concat!("extism_", stringify!($name)) => {
183                        let t = FuncType::new([$($args),*], [$($($r),*)?]);
184                        linker.func_new(EXPORT_MODULE_NAME, concat!("extism_", stringify!($name)), t, pdk::$name)?;
185                        continue
186                    }
187                )*
188                    _ => ()
189                }
190            };
191        }
192
193        // Add builtins
194        for (name, module) in modules.iter() {
195            if name != main_name {
196                linker.module(&mut store, name, module)?;
197            }
198            for import in module.imports() {
199                let module_name = import.module();
200                let name = import.name();
201                use wasmtime::ValType::*;
202
203                if module_name == EXPORT_MODULE_NAME {
204                    define_funcs!(name,  {
205                        config_get(I64) -> I64;
206                        var_get(I64) -> I64;
207                        var_set(I64, I64);
208                        http_request(I64, I64) -> I64;
209                        http_status_code() -> I32;
210                        log_warn(I64);
211                        log_info(I64);
212                        log_debug(I64);
213                        log_error(I64);
214                    });
215                }
216            }
217        }
218
219        for f in &mut imports {
220            let name = f.name().to_string();
221            let ns = f.namespace().unwrap_or(EXPORT_MODULE_NAME);
222            linker.func_new(ns, &name, f.ty().clone(), unsafe {
223                &*std::sync::Arc::as_ptr(&f.f)
224            })?;
225        }
226
227        let instance_pre = linker.instantiate_pre(&main)?;
228        let timer_id = uuid::Uuid::new_v4();
229        let mut plugin = Plugin {
230            modules,
231            linker,
232            instance: None,
233            instance_pre,
234            store,
235            runtime: None,
236            timer_id,
237            cancel_handle: sdk::ExtismCancelHandle {
238                id: timer_id,
239                epoch_timer_tx: None,
240            },
241            instantiations: 0,
242        };
243
244        plugin.internal_mut().store = &mut plugin.store;
245        plugin.internal_mut().linker = &mut plugin.linker;
246        Ok(plugin)
247    }
248
249    pub(crate) fn reset_store(&mut self) -> Result<(), Error> {
250        self.instance = None;
251        if self.instantiations > 5 {
252            let (main_name, main) = self
253                .modules
254                .get("main")
255                .map(|x| ("main", x))
256                .unwrap_or_else(|| {
257                    let entry = self.modules.iter().last().unwrap();
258                    (entry.0.as_str(), entry.1)
259                });
260
261            let engine = self.store.engine().clone();
262            let internal = self.internal();
263            self.store = Store::new(
264                &engine,
265                Internal::new(
266                    internal.manifest.clone(),
267                    internal.wasi.is_some(),
268                    internal.available_pages,
269                )?,
270            );
271            self.store
272                .epoch_deadline_callback(|_internal| Ok(UpdateDeadline::Continue(1)));
273
274            if self.internal().available_pages.is_some() {
275                self.store
276                    .limiter(|internal| internal.memory_limiter.as_mut().unwrap());
277            }
278
279            for (name, module) in self.modules.iter() {
280                if name != main_name {
281                    self.linker.module(&mut self.store, name, module)?;
282                }
283            }
284            self.instantiations = 0;
285            self.instance_pre = self.linker.instantiate_pre(&main)?;
286
287            let store = &mut self.store as *mut _;
288            let linker = &mut self.linker as *mut _;
289            let internal = self.internal_mut();
290            internal.store = store;
291            internal.linker = linker;
292        }
293
294        Ok(())
295    }
296
297    pub(crate) fn instantiate(&mut self) -> Result<(), Error> {
298        self.instance = Some(self.instance_pre.instantiate(&mut self.store)?);
299        self.instantiations += 1;
300        if let Some(limiter) = &mut self.internal_mut().memory_limiter {
301            limiter.reset();
302        }
303        self.detect_runtime();
304        self.initialize_runtime()?;
305        Ok(())
306    }
307
308    /// Get a function by name
309    pub fn get_func(&mut self, function: impl AsRef<str>) -> Option<Func> {
310        if let None = &self.instance {
311            if let Err(e) = self.instantiate() {
312                error!("Unable to instantiate: {e}");
313                return None;
314            }
315        }
316
317        if let Some(instance) = &mut self.instance {
318            instance.get_func(&mut self.store, function.as_ref())
319        } else {
320            None
321        }
322    }
323
324    /// Store input in memory and initialize `Internal` pointer
325    pub(crate) fn set_input(&mut self, input: *const u8, mut len: usize) -> Result<(), Error> {
326        if input.is_null() {
327            len = 0;
328        }
329
330        {
331            let store = &mut self.store as *mut _;
332            let linker = &mut self.linker as *mut _;
333            let internal = self.internal_mut();
334            internal.store = store;
335            internal.linker = linker;
336        }
337
338        if len > 0 {
339            let bytes = unsafe { std::slice::from_raw_parts(input, len) };
340            trace!("Input size: {}", bytes.len());
341
342            if let Some(f) = self.linker.get(&mut self.store, "env", "extism_reset") {
343                f.into_func().unwrap().call(&mut self.store, &[], &mut [])?;
344            } else {
345                error!("Call to extism_reset failed");
346            }
347
348            let offs = self.memory_alloc_bytes(bytes)?;
349
350            if let Some(f) = self.linker.get(&mut self.store, "env", "extism_input_set") {
351                f.into_func().unwrap().call(
352                    &mut self.store,
353                    &[Val::I64(offs as i64), Val::I64(len as i64)],
354                    &mut [],
355                )?;
356            }
357        }
358
359        Ok(())
360    }
361
362    /// Determine if wasi is enabled
363    pub fn has_wasi(&self) -> bool {
364        self.internal().wasi.is_some()
365    }
366
367    fn detect_runtime(&mut self) {
368        // Check for Haskell runtime initialization functions
369        // Initialize Haskell runtime if `hs_init` is present,
370        // by calling the `hs_init` export
371        if let Some(init) = self.get_func("hs_init") {
372            let reactor_init = if let Some(init) = self.get_func("_initialize") {
373                if init.typed::<(), ()>(&self.store()).is_err() {
374                    trace!(
375                        "_initialize function found with type {:?}",
376                        init.ty(self.store())
377                    );
378                    None
379                } else {
380                    trace!("WASI reactor module detected");
381                    Some(init)
382                }
383            } else {
384                None
385            };
386            self.runtime = Some(Runtime::Haskell { init, reactor_init });
387            return;
388        }
389
390        // Check for `__wasm_call_ctors` or `_initialize`, this is used by WASI to
391        // initialize certain interfaces.
392        let init = if let Some(init) = self.get_func("__wasm_call_ctors") {
393            if init.typed::<(), ()>(&self.store()).is_err() {
394                trace!(
395                    "__wasm_call_ctors function found with type {:?}",
396                    init.ty(self.store())
397                );
398                return;
399            }
400            trace!("WASI runtime detected");
401            init
402        } else if let Some(init) = self.get_func("_initialize") {
403            if init.typed::<(), ()>(&self.store()).is_err() {
404                trace!(
405                    "_initialize function found with type {:?}",
406                    init.ty(self.store())
407                );
408                return;
409            }
410            trace!("Reactor module detected");
411            init
412        } else {
413            return;
414        };
415
416        self.runtime = Some(Runtime::Wasi { init });
417
418        trace!("No runtime detected");
419    }
420
421    pub(crate) fn initialize_runtime(&mut self) -> Result<(), Error> {
422        let mut store = &mut self.store;
423        if let Some(runtime) = &self.runtime {
424            trace!("Plugin::initialize_runtime");
425            match runtime {
426                Runtime::Haskell { init, reactor_init } => {
427                    if let Some(reactor_init) = reactor_init {
428                        reactor_init.call(&mut store, &[], &mut [])?;
429                    }
430                    let mut results = vec![Val::null(); init.ty(&store).results().len()];
431                    init.call(
432                        &mut store,
433                        &[Val::I32(0), Val::I32(0)],
434                        results.as_mut_slice(),
435                    )?;
436                    debug!("Initialized Haskell language runtime");
437                }
438                Runtime::Wasi { init } => {
439                    init.call(&mut store, &[], &mut [])?;
440                    debug!("Initialied WASI runtime");
441                }
442            }
443        }
444
445        Ok(())
446    }
447
448    /// Start the timer for a Plugin - this is used for both timeouts
449    /// and cancellation
450    pub(crate) fn start_timer(
451        &mut self,
452        tx: &std::sync::mpsc::SyncSender<TimerAction>,
453    ) -> Result<(), Error> {
454        let duration = self
455            .internal()
456            .manifest
457            .as_ref()
458            .timeout_ms
459            .map(std::time::Duration::from_millis);
460        self.cancel_handle.epoch_timer_tx = Some(tx.clone());
461        self.store_mut().set_epoch_deadline(1);
462        self.store
463            .epoch_deadline_callback(|_internal| Err(Error::msg("timeout")));
464        let engine: Engine = self.store().engine().clone();
465        tx.send(TimerAction::Start {
466            id: self.timer_id,
467            duration,
468            engine,
469        })?;
470        Ok(())
471    }
472
473    /// Send TimerAction::Stop
474    pub(crate) fn stop_timer(&mut self) -> Result<(), Error> {
475        if let Some(tx) = &self.cancel_handle.epoch_timer_tx {
476            tx.send(TimerAction::Stop { id: self.timer_id })?;
477        }
478        self.store
479            .epoch_deadline_callback(|_internal| Ok(wasmtime::UpdateDeadline::Continue(1)));
480        Ok(())
481    }
482}
483
484// Enumerates the supported PDK language runtimes
485#[derive(Clone)]
486pub(crate) enum Runtime {
487    Haskell {
488        init: Func,
489        reactor_init: Option<Func>,
490    },
491    Wasi {
492        init: Func,
493    },
494}
extism_runtime/plugin.rs

extism_runtime/
plugin.rs
