Skip to main content

wasmer_wasix/runtime/
mod.rs

1pub mod module_cache;
2pub mod package_loader;
3pub mod resolver;
4pub mod task_manager;
5
6use self::module_cache::CacheError;
7pub use self::task_manager::{SpawnType, VirtualTaskManager};
8use module_cache::HashedModuleData;
9use wasmer_types::{CompilationProgressCallback, ModuleHash};
10
11use std::{
12    borrow::Cow,
13    fmt,
14    ops::Deref,
15    sync::{Arc, Mutex},
16};
17
18use anyhow::Context as _;
19use futures::future::BoxFuture;
20use virtual_mio::block_on;
21use virtual_net::{DynVirtualNetworking, VirtualNetworking};
22use wasmer::{Engine, Module, RuntimeError};
23use wasmer_wasix_types::wasi::ExitCode;
24
25#[cfg(feature = "journal")]
26use crate::journal::{DynJournal, DynReadableJournal};
27use crate::{
28    SpawnError, WasiTtyState,
29    bin_factory::BinaryPackageCommand,
30    http::{DynHttpClient, HttpClient},
31    os::TtyBridge,
32    runtime::{
33        module_cache::{
34            ModuleCache, ThreadLocalCache,
35            progress::{ModuleLoadProgress, ModuleLoadProgressReporter},
36        },
37        package_loader::{PackageLoader, UnsupportedPackageLoader},
38        resolver::{BackendSource, MultiSource, Source},
39    },
40};
41
42/// Opaque per-instantiation state, created by
43/// [`InstantiationHook::additional_imports`] and handed back to
44/// [`InstantiationHook::configure_new_instance`] for the instance built with
45/// those imports.
46///
47/// Hooks put the data they need to carry between the two phases in with
48/// [`InstantiationState::new`] and get it back out with
49/// [`InstantiationState::take`]. Callers only pass the value along, unmodified.
50// Carrying the state through the instantiation, instead of parking it in the
51// hook, is what makes concurrent instantiations safe: an implementation never
52// has to guess which pending instantiation a configure_new_instance call
53// belongs to, so two threads cold-starting the same module in different
54// stores cannot receive each other's state.
55#[derive(Default)]
56pub struct InstantiationState {
57    state: Option<Box<dyn std::any::Any + Send>>,
58}
59
60impl InstantiationState {
61    /// State that carries no data, for hooks that need nothing from the import
62    /// phase.
63    pub fn empty() -> Self {
64        Self { state: None }
65    }
66
67    /// Carries `state` from the import phase to the instance setup phase.
68    pub fn new<T: std::any::Any + Send>(state: T) -> Self {
69        Self {
70            state: Some(Box::new(state)),
71        }
72    }
73
74    /// Whether this state carries no data.
75    pub fn is_empty(&self) -> bool {
76        self.state.is_none()
77    }
78
79    /// Takes back the data stored by [`InstantiationState::new`].
80    ///
81    /// Fails if the state is empty or holds a different type, both of which
82    /// mean it did not come from the matching import phase.
83    pub fn take<T: std::any::Any + Send>(self) -> anyhow::Result<T> {
84        let state = self
85            .state
86            .context("missing instantiation state from the import phase")?;
87        state
88            .downcast::<T>()
89            .map(|state| *state)
90            .map_err(|_| anyhow::anyhow!("instantiation state does not belong to this hook"))
91    }
92}
93
94impl fmt::Debug for InstantiationState {
95    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
96        if self.is_empty() {
97            f.write_str("InstantiationState::empty()")
98        } else {
99            f.write_str("InstantiationState(..)")
100        }
101    }
102}
103
104/// A hook into the instantiation of WASIX module instances.
105///
106/// Registered on [`PluggableRuntime::with_instantiation_hook`] or
107/// [`OverriddenRuntime::with_instantiation_hook`], and invoked once per
108/// instance the runtime creates (process bootstrap, thread spawn, dynamically
109/// linked side module).
110///
111/// Both methods have a no-op default, so an implementation only needs to
112/// provide the phases it cares about.
113// Keeping both phases on one trait is what lets each hook route its own
114// InstantiationState from its import phase to its own setup phase when
115// several hooks are registered on the same runtime.
116pub trait InstantiationHook: fmt::Debug + Send + Sync + 'static {
117    /// Creates additional imports for an instance about to be created in
118    /// `store`.
119    ///
120    /// The returned [`InstantiationState`] is handed back to
121    /// [`InstantiationHook::configure_new_instance`] for the instance built
122    /// with these imports. If instantiation fails, the state is dropped.
123    fn additional_imports(
124        &self,
125        module: &wasmer::Module,
126        store: &mut wasmer::StoreMut,
127    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
128        let _ = (module, store);
129        Ok((wasmer::Imports::new(), InstantiationState::empty()))
130    }
131
132    /// Configures an instantiated instance before initialization/startup.
133    ///
134    /// `state` is the [`InstantiationState`] this hook returned from the
135    /// [`InstantiationHook::additional_imports`] call whose imports the
136    /// instance was created with.
137    fn configure_new_instance(
138        &self,
139        module: &wasmer::Module,
140        store: &mut wasmer::StoreMut,
141        instance: &wasmer::Instance,
142        imported_memory: Option<&wasmer::Memory>,
143        state: InstantiationState,
144    ) -> anyhow::Result<()> {
145        let _ = (module, store, instance, imported_memory, state);
146        Ok(())
147    }
148}
149
150impl<H: InstantiationHook + ?Sized> InstantiationHook for Arc<H> {
151    fn additional_imports(
152        &self,
153        module: &wasmer::Module,
154        store: &mut wasmer::StoreMut,
155    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
156        (**self).additional_imports(module, store)
157    }
158
159    fn configure_new_instance(
160        &self,
161        module: &wasmer::Module,
162        store: &mut wasmer::StoreMut,
163        instance: &wasmer::Instance,
164        imported_memory: Option<&wasmer::Memory>,
165        state: InstantiationState,
166    ) -> anyhow::Result<()> {
167        (**self).configure_new_instance(module, store, instance, imported_memory, state)
168    }
169}
170
171/// Adapts an import-creation closure to [`InstantiationHook`].
172struct ImportsOnlyHook<F>(F);
173
174impl<F> fmt::Debug for ImportsOnlyHook<F> {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        f.write_str("ImportsOnlyHook(..)")
177    }
178}
179
180impl<F> InstantiationHook for ImportsOnlyHook<F>
181where
182    F: Fn(&wasmer::Module, &mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports>
183        + Send
184        + Sync
185        + 'static,
186{
187    fn additional_imports(
188        &self,
189        module: &wasmer::Module,
190        store: &mut wasmer::StoreMut,
191    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
192        Ok(((self.0)(module, store)?, InstantiationState::empty()))
193    }
194}
195
196/// Adapts an instance-setup closure to [`InstantiationHook`].
197struct InstanceSetupHook<F>(F);
198
199impl<F> fmt::Debug for InstanceSetupHook<F> {
200    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
201        f.write_str("InstanceSetupHook(..)")
202    }
203}
204
205impl<F> InstantiationHook for InstanceSetupHook<F>
206where
207    F: Fn(
208            &wasmer::Module,
209            &mut wasmer::StoreMut,
210            &wasmer::Instance,
211            Option<&wasmer::Memory>,
212        ) -> anyhow::Result<()>
213        + Send
214        + Sync
215        + 'static,
216{
217    fn configure_new_instance(
218        &self,
219        module: &wasmer::Module,
220        store: &mut wasmer::StoreMut,
221        instance: &wasmer::Instance,
222        imported_memory: Option<&wasmer::Memory>,
223        _state: InstantiationState,
224    ) -> anyhow::Result<()> {
225        (self.0)(module, store, instance, imported_memory)
226    }
227}
228
229#[derive(Clone)]
230pub enum TaintReason {
231    UnknownWasiVersion,
232    NonZeroExitCode(ExitCode),
233    RuntimeError(RuntimeError),
234    DlSymbolResolutionFailed(String),
235}
236
237/// The input to load a module.
238///
239/// Exists because the semantics for resolving modules can vary between
240/// different sources.
241///
242/// All variants are wrapped in `Cow` to allow for zero-copy usage when possible.
243#[allow(clippy::large_enum_variant)]
244pub enum ModuleInput<'a> {
245    /// Raw bytes.
246    Bytes(Cow<'a, [u8]>),
247    /// Pre-hashed module data.
248    Hashed(Cow<'a, HashedModuleData>),
249    /// A binary package command.
250    Command(Cow<'a, BinaryPackageCommand>),
251}
252
253impl<'a> ModuleInput<'a> {
254    /// Convert to an owned version of the module input.
255    pub fn to_owned(&'a self) -> ModuleInput<'static> {
256        // The manual code below is needed due to compiler issues with the lifetime.
257        match self {
258            Self::Bytes(Cow::Borrowed(b)) => {
259                let v: Vec<u8> = (*b).to_owned();
260                let c: Cow<'static, [u8]> = Cow::from(v);
261                ModuleInput::Bytes(c)
262            }
263            Self::Bytes(Cow::Owned(b)) => ModuleInput::Bytes(Cow::Owned((*b).clone())),
264            Self::Hashed(Cow::Borrowed(h)) => ModuleInput::Hashed(Cow::Owned((*h).clone())),
265            Self::Hashed(Cow::Owned(h)) => ModuleInput::Hashed(Cow::Owned(h.clone())),
266            Self::Command(Cow::Borrowed(c)) => ModuleInput::Command(Cow::Owned((*c).clone())),
267            Self::Command(Cow::Owned(c)) => ModuleInput::Command(Cow::Owned(c.clone())),
268        }
269    }
270
271    /// Get the module hash.
272    ///
273    /// NOTE: may be expensive, depending on the variant.
274    pub fn hash(&self) -> ModuleHash {
275        match self {
276            Self::Bytes(b) => {
277                // Hash on the fly
278                ModuleHash::new(b)
279            }
280            Self::Hashed(hashed) => *hashed.hash(),
281            Self::Command(cmd) => *cmd.hash(),
282        }
283    }
284
285    /// Get the raw WebAssembly bytes.
286    pub fn wasm(&self) -> &[u8] {
287        match self {
288            Self::Bytes(b) => b,
289            Self::Hashed(hashed) => hashed.wasm().as_ref(),
290            Self::Command(cmd) => cmd.atom_ref().as_ref(),
291        }
292    }
293
294    /// Convert to a `HashedModuleData`.
295    ///
296    /// May involve cloning and hashing.
297    pub fn to_hashed(&self) -> HashedModuleData {
298        match self {
299            Self::Bytes(b) => HashedModuleData::new(b.as_ref()),
300            Self::Hashed(hashed) => hashed.as_ref().clone(),
301            Self::Command(cmd) => HashedModuleData::from_command(cmd),
302        }
303    }
304}
305
306/// Runtime components used when running WebAssembly programs.
307///
308/// Think of this as the "System" in "WebAssembly Systems Interface".
309#[allow(unused_variables)]
310pub trait Runtime
311where
312    Self: fmt::Debug,
313{
314    /// Provides access to all the networking related functions such as sockets.
315    fn networking(&self) -> &DynVirtualNetworking;
316
317    /// Retrieve the active [`VirtualTaskManager`].
318    fn task_manager(&self) -> &Arc<dyn VirtualTaskManager>;
319
320    /// A package loader.
321    fn package_loader(&self) -> Arc<dyn PackageLoader + Send + Sync> {
322        Arc::new(UnsupportedPackageLoader)
323    }
324
325    /// A cache for compiled modules.
326    fn module_cache(&self) -> Arc<dyn ModuleCache + Send + Sync> {
327        // Return a cache that uses a thread-local variable. This isn't ideal
328        // because it allows silently sharing state, possibly between runtimes.
329        //
330        // That said, it means people will still get *some* level of caching
331        // because each cache returned by this default implementation will go
332        // through the same thread-local variable.
333        Arc::new(ThreadLocalCache::default())
334    }
335
336    /// The package registry.
337    fn source(&self) -> Arc<dyn Source + Send + Sync>;
338
339    /// Get a [`wasmer::Engine`] for module compilation.
340    fn engine(&self) -> Engine {
341        Engine::default()
342    }
343
344    /// Create a new [`wasmer::Store`].
345    fn new_store(&self) -> wasmer::Store {
346        cfg_select! {
347            feature = "sys" => {
348                wasmer::Store::new(self.engine())
349            }
350            _ => {
351                wasmer::Store::default()
352            }
353        }
354    }
355
356    /// Create additional imports for a new WASIX instance in the provided store.
357    ///
358    /// This callback may be invoked multiple times (e.g. process bootstrap,
359    /// thread spawn), so implementations should create imports that are valid
360    /// for the given store each time.
361    ///
362    /// The returned [`InstantiationState`] is per-instantiation state that the
363    /// caller must pass to [`Runtime::configure_new_instance`] once the
364    /// instance built from these imports exists. If instantiation fails, the
365    /// state is simply dropped.
366    fn additional_imports(
367        &self,
368        _module: &wasmer::Module,
369        _store: &mut wasmer::StoreMut,
370    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
371        Ok((wasmer::Imports::new(), InstantiationState::empty()))
372    }
373
374    /// Configure an instantiated instance before initialization/startup.
375    ///
376    /// `state` must be the [`InstantiationState`] returned by the
377    /// [`Runtime::additional_imports`] call whose imports this instance was
378    /// created with.
379    fn configure_new_instance(
380        &self,
381        _module: &wasmer::Module,
382        _store: &mut wasmer::StoreMut,
383        _instance: &wasmer::Instance,
384        _imported_memory: Option<&wasmer::Memory>,
385        _state: InstantiationState,
386    ) -> anyhow::Result<()> {
387        Ok(())
388    }
389
390    /// Get a custom HTTP client
391    fn http_client(&self) -> Option<&DynHttpClient> {
392        None
393    }
394
395    /// Get access to the TTY used by the environment.
396    fn tty(&self) -> Option<&(dyn TtyBridge + Send + Sync)> {
397        None
398    }
399
400    /// The primary way to load a module given a module input.
401    ///
402    /// The engine to use can be optionally provided, otherwise the most appropriate engine
403    /// should be selected.
404    ///
405    /// An optional progress reporter callback can be provided to report progress during module loading.
406    fn resolve_module<'a>(
407        &'a self,
408        input: ModuleInput<'a>,
409        engine: Option<&Engine>,
410        on_progress: Option<ModuleLoadProgressReporter>,
411    ) -> BoxFuture<'a, Result<Module, SpawnError>> {
412        let data = input.to_hashed();
413
414        let engine = if let Some(e) = engine {
415            e.clone()
416        } else {
417            match &input {
418                ModuleInput::Bytes(_) => self.engine(),
419                ModuleInput::Hashed(_) => self.engine(),
420                ModuleInput::Command(cmd) => self.engine(),
421            }
422        };
423
424        let module_cache = self.module_cache();
425
426        let task = async move { load_module(&engine, &module_cache, input, on_progress).await };
427        Box::pin(task)
428    }
429
430    /// Sync variant of [`Self::resolve_module`].
431    fn resolve_module_sync(
432        &self,
433        input: ModuleInput<'_>,
434        engine: Option<&Engine>,
435        on_progress: Option<ModuleLoadProgressReporter>,
436    ) -> Result<Module, SpawnError> {
437        block_on(self.resolve_module(input, engine, on_progress))
438    }
439
440    /// Load the module for a command.
441    ///
442    /// Will load the module from the cache if possible, otherwise will compile.
443    ///
444    /// NOTE: This always be preferred over [`Self::load_module`] to avoid
445    /// re-hashing the module!
446    #[deprecated(since = "0.601.0", note = "Use `resolve_module` instead")]
447    fn load_command_module(
448        &self,
449        cmd: &BinaryPackageCommand,
450    ) -> BoxFuture<'_, Result<Module, SpawnError>> {
451        self.resolve_module(ModuleInput::Command(Cow::Owned(cmd.clone())), None, None)
452    }
453
454    /// Sync version of [`Self::load_command_module`].
455    #[deprecated(since = "0.601.0", note = "Use `resolve_module_sync` instead")]
456    fn load_command_module_sync(&self, cmd: &BinaryPackageCommand) -> Result<Module, SpawnError> {
457        block_on(self.resolve_module(ModuleInput::Command(Cow::Borrowed(cmd)), None, None))
458    }
459
460    /// Load a WebAssembly module from raw bytes.
461    ///
462    /// Will load the module from the cache if possible, otherwise will compile.
463    #[deprecated(since = "0.601.0", note = "Use `resolve_module` instead")]
464    fn load_module<'a>(&'a self, wasm: &'a [u8]) -> BoxFuture<'a, Result<Module, SpawnError>> {
465        self.resolve_module(ModuleInput::Bytes(Cow::Borrowed(wasm)), None, None)
466    }
467
468    /// Synchronous version of [`Self::load_module`].
469    #[deprecated(
470        since = "0.601.0",
471        note = "Use `load_command_module` or `load_hashed_module` instead - this method can have high overhead"
472    )]
473    fn load_module_sync(&self, wasm: &[u8]) -> Result<Module, SpawnError> {
474        block_on(self.resolve_module(ModuleInput::Bytes(Cow::Borrowed(wasm)), None, None))
475    }
476
477    /// Load a WebAssembly module from pre-hashed data.
478    ///
479    /// Will load the module from the cache if possible, otherwise will compile.
480    fn load_hashed_module(
481        &self,
482        module: HashedModuleData,
483        engine: Option<&Engine>,
484    ) -> BoxFuture<'_, Result<Module, SpawnError>> {
485        self.resolve_module(ModuleInput::Hashed(Cow::Owned(module)), engine, None)
486    }
487
488    /// Synchronous version of [`Self::load_hashed_module`].
489    fn load_hashed_module_sync(
490        &self,
491        wasm: HashedModuleData,
492        engine: Option<&Engine>,
493    ) -> Result<Module, SpawnError> {
494        block_on(self.resolve_module(ModuleInput::Hashed(Cow::Owned(wasm)), engine, None))
495    }
496
497    /// Callback thats invokes whenever the instance is tainted, tainting can occur
498    /// for multiple reasons however the most common is a panic within the process
499    fn on_taint(&self, _reason: TaintReason) {}
500
501    /// The list of all read-only journals which will be used to restore the state of the
502    /// runtime at a particular point in time
503    #[cfg(feature = "journal")]
504    fn read_only_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynReadableJournal>> + 'a> {
505        Box::new(std::iter::empty())
506    }
507
508    /// The list of writable journals which will be appended to
509    #[cfg(feature = "journal")]
510    fn writable_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynJournal>> + 'a> {
511        Box::new(std::iter::empty())
512    }
513
514    /// The snapshot capturer takes and restores snapshots of the WASM process at specific
515    /// points in time by reading and writing log entries
516    #[cfg(feature = "journal")]
517    fn active_journal(&self) -> Option<&'_ DynJournal> {
518        None
519    }
520}
521
522pub type DynRuntime = dyn Runtime + Send + Sync;
523
524/// Load a Webassembly module, trying to use a pre-compiled version if possible.
525///
526// This function exists to provide a reusable baseline implementation for
527// implementing [`Runtime::load_module`], so custom logic can be added on top.
528#[tracing::instrument(level = "debug", skip_all)]
529pub async fn load_module(
530    engine: &Engine,
531    module_cache: &(dyn ModuleCache + Send + Sync),
532    input: ModuleInput<'_>,
533    on_progress: Option<ModuleLoadProgressReporter>,
534) -> Result<Module, crate::SpawnError> {
535    let wasm_hash = input.hash();
536
537    let result = if let Some(on_progress) = &on_progress {
538        module_cache
539            .load_with_progress(wasm_hash, engine, on_progress.clone())
540            .await
541    } else {
542        module_cache.load(wasm_hash, engine).await
543    };
544
545    match result {
546        Ok(module) => return Ok(module),
547        Err(CacheError::NotFound) => {}
548        Err(other) => {
549            tracing::warn!(
550                %wasm_hash,
551                error=&other as &dyn std::error::Error,
552                "Unable to load the cached module",
553            );
554        }
555    }
556
557    let res = if let Some(progress) = on_progress {
558        #[allow(unused_variables)]
559        let p = CompilationProgressCallback::new(move |p| {
560            progress.notify(ModuleLoadProgress::CompilingModule(p))
561        });
562        #[cfg(feature = "sys")]
563        {
564            if engine.is_sys() {
565                use wasmer::sys::NativeEngineExt;
566                engine.new_module_with_progress(input.wasm(), p)
567            } else {
568                Module::new(&engine, input.wasm())
569            }
570        }
571        #[cfg(not(feature = "sys"))]
572        {
573            Module::new(&engine, input.wasm())
574        }
575    } else {
576        Module::new(&engine, input.wasm())
577    };
578
579    let module = res.map_err(|err| crate::SpawnError::CompileError {
580        module_hash: wasm_hash,
581        error: err,
582    })?;
583
584    // TODO: pass a [`HashedModule`] struct that is safe by construction.
585    if let Err(e) = module_cache.save(wasm_hash, engine, &module).await {
586        tracing::warn!(
587            %wasm_hash,
588            error=&e as &dyn std::error::Error,
589            "Unable to cache the compiled module",
590        );
591    }
592
593    Ok(module)
594}
595
596#[derive(Debug, Default)]
597pub struct DefaultTty {
598    state: Mutex<WasiTtyState>,
599}
600
601impl TtyBridge for DefaultTty {
602    fn reset(&self) {
603        let mut state = self.state.lock().unwrap();
604        state.echo = false;
605        state.line_buffered = false;
606        state.line_feeds = false
607    }
608
609    fn tty_get(&self) -> WasiTtyState {
610        let state = self.state.lock().unwrap();
611        state.clone()
612    }
613
614    fn tty_set(&self, tty_state: WasiTtyState) {
615        let mut state = self.state.lock().unwrap();
616        *state = tty_state;
617    }
618}
619
620#[derive(Debug, Clone)]
621pub struct PluggableRuntime {
622    pub rt: Arc<dyn VirtualTaskManager>,
623    pub networking: DynVirtualNetworking,
624    pub http_client: Option<DynHttpClient>,
625    pub package_loader: Arc<dyn PackageLoader + Send + Sync>,
626    pub source: Arc<dyn Source + Send + Sync>,
627    pub engine: Engine,
628    pub module_cache: Arc<dyn ModuleCache + Send + Sync>,
629    pub tty: Option<Arc<dyn TtyBridge + Send + Sync>>,
630    #[cfg(feature = "journal")]
631    pub read_only_journals: Vec<Arc<DynReadableJournal>>,
632    #[cfg(feature = "journal")]
633    pub writable_journals: Vec<Arc<DynJournal>>,
634    pub instantiation_hooks: Vec<Arc<dyn InstantiationHook>>,
635}
636
637impl PluggableRuntime {
638    pub fn new(rt: Arc<dyn VirtualTaskManager>) -> Self {
639        // TODO: the cfg flags below should instead be handled by separate implementations.
640        cfg_select! {
641            feature = "host-vnet" => {
642                let networking = Arc::new(virtual_net::host::LocalNetworking::default());
643            }
644            _ => {
645                let networking = Arc::new(virtual_net::UnsupportedVirtualNetworking::default());
646            }
647        }
648        let http_client =
649            crate::http::default_http_client().map(|client| Arc::new(client) as DynHttpClient);
650
651        let loader = UnsupportedPackageLoader;
652
653        let mut source = MultiSource::default();
654        if let Some(client) = &http_client {
655            source.add_source(BackendSource::new(
656                BackendSource::WASMER_PROD_ENDPOINT.parse().unwrap(),
657                client.clone(),
658            ));
659        }
660
661        Self {
662            rt,
663            networking,
664            http_client,
665            engine: Default::default(),
666            tty: None,
667            source: Arc::new(source),
668            package_loader: Arc::new(loader),
669            module_cache: Arc::new(module_cache::in_memory()),
670            #[cfg(feature = "journal")]
671            read_only_journals: Vec::new(),
672            #[cfg(feature = "journal")]
673            writable_journals: Vec::new(),
674            instantiation_hooks: Vec::new(),
675        }
676    }
677
678    pub fn set_networking_implementation<I>(&mut self, net: I) -> &mut Self
679    where
680        I: VirtualNetworking + Sync,
681    {
682        self.networking = Arc::new(net);
683        self
684    }
685
686    pub fn set_engine(&mut self, engine: Engine) -> &mut Self {
687        self.engine = engine;
688        self
689    }
690
691    pub fn set_tty(&mut self, tty: Arc<dyn TtyBridge + Send + Sync>) -> &mut Self {
692        self.tty = Some(tty);
693        self
694    }
695
696    pub fn set_module_cache(
697        &mut self,
698        module_cache: impl ModuleCache + Send + Sync + 'static,
699    ) -> &mut Self {
700        self.module_cache = Arc::new(module_cache);
701        self
702    }
703
704    pub fn set_source(&mut self, source: impl Source + Send + 'static) -> &mut Self {
705        self.source = Arc::new(source);
706        self
707    }
708
709    pub fn set_package_loader(
710        &mut self,
711        package_loader: impl PackageLoader + 'static,
712    ) -> &mut Self {
713        self.package_loader = Arc::new(package_loader);
714        self
715    }
716
717    pub fn set_http_client(
718        &mut self,
719        client: impl HttpClient + Send + Sync + 'static,
720    ) -> &mut Self {
721        self.http_client = Some(Arc::new(client));
722        self
723    }
724
725    #[cfg(feature = "journal")]
726    pub fn add_read_only_journal(&mut self, journal: Arc<DynReadableJournal>) -> &mut Self {
727        self.read_only_journals.push(journal);
728        self
729    }
730
731    #[cfg(feature = "journal")]
732    pub fn add_writable_journal(&mut self, journal: Arc<DynJournal>) -> &mut Self {
733        self.writable_journals.push(journal);
734        self
735    }
736
737    /// Registers a hook that only creates additional imports.
738    pub fn with_additional_imports(
739        &mut self,
740        imports: impl Fn(&wasmer::Module, &mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports>
741        + Send
742        + Sync
743        + 'static,
744    ) -> &mut Self {
745        self.with_instantiation_hook(ImportsOnlyHook(imports))
746    }
747
748    /// Registers a hook that only configures newly created instances.
749    pub fn with_instance_setup(
750        &mut self,
751        callback: impl Fn(
752            &wasmer::Module,
753            &mut wasmer::StoreMut,
754            &wasmer::Instance,
755            Option<&wasmer::Memory>,
756        ) -> anyhow::Result<()>
757        + Send
758        + Sync
759        + 'static,
760    ) -> &mut Self {
761        self.with_instantiation_hook(InstanceSetupHook(callback))
762    }
763
764    /// Registers a hook that takes part in both phases of instantiation, so it
765    /// can carry [`InstantiationState`] from its imports to its instance setup.
766    pub fn with_instantiation_hook(&mut self, hook: impl InstantiationHook) -> &mut Self {
767        self.instantiation_hooks.push(Arc::new(hook));
768        self
769    }
770}
771
772/// Runs the import phase of `hooks`, returning the merged imports and the
773/// per-hook states, aligned by index with `hooks`.
774fn run_import_hooks(
775    hooks: &[Arc<dyn InstantiationHook>],
776    module: &wasmer::Module,
777    store: &mut wasmer::StoreMut,
778) -> anyhow::Result<(wasmer::Imports, Vec<InstantiationState>)> {
779    let mut imports = wasmer::Imports::new();
780    let mut states = Vec::with_capacity(hooks.len());
781    for hook in hooks {
782        let (hook_imports, state) = hook.additional_imports(module, store)?;
783        imports.extend(&hook_imports);
784        states.push(state);
785    }
786    Ok((imports, states))
787}
788
789/// Composite state used by [`OverriddenRuntime`] to carry the inner
790/// runtime's state alongside its own hooks' states.
791struct OverriddenInstantiationState {
792    inner: InstantiationState,
793    own: Vec<InstantiationState>,
794}
795
796/// Runs the setup phase of `hooks`, handing each hook the state it produced
797/// during the import phase.
798fn run_setup_hooks(
799    hooks: &[Arc<dyn InstantiationHook>],
800    states: Vec<InstantiationState>,
801    module: &wasmer::Module,
802    store: &mut wasmer::StoreMut,
803    instance: &wasmer::Instance,
804    imported_memory: Option<&wasmer::Memory>,
805) -> anyhow::Result<()> {
806    anyhow::ensure!(
807        states.len() == hooks.len(),
808        "instance setup state does not match the registered instantiation hooks \
809         (got {} states for {} hooks)",
810        states.len(),
811        hooks.len(),
812    );
813    for (hook, state) in hooks.iter().zip(states) {
814        hook.configure_new_instance(module, store, instance, imported_memory, state)?;
815    }
816    Ok(())
817}
818
819impl Runtime for PluggableRuntime {
820    fn networking(&self) -> &DynVirtualNetworking {
821        &self.networking
822    }
823
824    fn http_client(&self) -> Option<&DynHttpClient> {
825        self.http_client.as_ref()
826    }
827
828    fn package_loader(&self) -> Arc<dyn PackageLoader + Send + Sync> {
829        Arc::clone(&self.package_loader)
830    }
831
832    fn source(&self) -> Arc<dyn Source + Send + Sync> {
833        Arc::clone(&self.source)
834    }
835
836    fn engine(&self) -> Engine {
837        self.engine.clone()
838    }
839
840    fn new_store(&self) -> wasmer::Store {
841        wasmer::Store::new(self.engine.clone())
842    }
843
844    fn task_manager(&self) -> &Arc<dyn VirtualTaskManager> {
845        &self.rt
846    }
847
848    fn tty(&self) -> Option<&(dyn TtyBridge + Send + Sync)> {
849        self.tty.as_deref()
850    }
851
852    fn module_cache(&self) -> Arc<dyn ModuleCache + Send + Sync> {
853        self.module_cache.clone()
854    }
855
856    fn additional_imports(
857        &self,
858        module: &wasmer::Module,
859        store: &mut wasmer::StoreMut,
860    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
861        if self.instantiation_hooks.is_empty() {
862            return Ok((wasmer::Imports::new(), InstantiationState::empty()));
863        }
864        let (imports, states) = run_import_hooks(&self.instantiation_hooks, module, store)?;
865        Ok((imports, InstantiationState::new(states)))
866    }
867
868    fn configure_new_instance(
869        &self,
870        module: &wasmer::Module,
871        store: &mut wasmer::StoreMut,
872        instance: &wasmer::Instance,
873        imported_memory: Option<&wasmer::Memory>,
874        state: InstantiationState,
875    ) -> anyhow::Result<()> {
876        if self.instantiation_hooks.is_empty() {
877            return Ok(());
878        }
879        let states = state
880            .take::<Vec<InstantiationState>>()
881            .context("invalid instance setup state from additional_imports")?;
882        run_setup_hooks(
883            &self.instantiation_hooks,
884            states,
885            module,
886            store,
887            instance,
888            imported_memory,
889        )
890    }
891
892    #[cfg(feature = "journal")]
893    fn read_only_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynReadableJournal>> + 'a> {
894        Box::new(self.read_only_journals.iter().cloned())
895    }
896
897    #[cfg(feature = "journal")]
898    fn writable_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynJournal>> + 'a> {
899        Box::new(self.writable_journals.iter().cloned())
900    }
901
902    #[cfg(feature = "journal")]
903    fn active_journal(&self) -> Option<&DynJournal> {
904        self.writable_journals.iter().last().map(|a| a.as_ref())
905    }
906}
907
908/// Runtime that allows for certain things to be overridden
909/// such as the active journals
910#[derive(Clone, Debug)]
911pub struct OverriddenRuntime {
912    inner: Arc<DynRuntime>,
913    task_manager: Option<Arc<dyn VirtualTaskManager>>,
914    networking: Option<DynVirtualNetworking>,
915    http_client: Option<DynHttpClient>,
916    package_loader: Option<Arc<dyn PackageLoader + Send + Sync>>,
917    source: Option<Arc<dyn Source + Send + Sync>>,
918    engine: Option<Engine>,
919    module_cache: Option<Arc<dyn ModuleCache + Send + Sync>>,
920    tty: Option<Arc<dyn TtyBridge + Send + Sync>>,
921    instantiation_hooks: Vec<Arc<dyn InstantiationHook>>,
922    #[cfg(feature = "journal")]
923    pub read_only_journals: Option<Vec<Arc<DynReadableJournal>>>,
924    #[cfg(feature = "journal")]
925    pub writable_journals: Option<Vec<Arc<DynJournal>>>,
926}
927
928impl OverriddenRuntime {
929    pub fn new(inner: Arc<DynRuntime>) -> Self {
930        Self {
931            inner,
932            task_manager: None,
933            networking: None,
934            http_client: None,
935            package_loader: None,
936            source: None,
937            engine: None,
938            module_cache: None,
939            tty: None,
940            instantiation_hooks: Vec::new(),
941            #[cfg(feature = "journal")]
942            read_only_journals: None,
943            #[cfg(feature = "journal")]
944            writable_journals: None,
945        }
946    }
947
948    pub fn with_task_manager(mut self, task_manager: Arc<dyn VirtualTaskManager>) -> Self {
949        self.task_manager.replace(task_manager);
950        self
951    }
952
953    pub fn with_networking(mut self, networking: DynVirtualNetworking) -> Self {
954        self.networking.replace(networking);
955        self
956    }
957
958    pub fn with_http_client(mut self, http_client: DynHttpClient) -> Self {
959        self.http_client.replace(http_client);
960        self
961    }
962
963    pub fn with_package_loader(
964        mut self,
965        package_loader: Arc<dyn PackageLoader + Send + Sync>,
966    ) -> Self {
967        self.package_loader.replace(package_loader);
968        self
969    }
970
971    pub fn with_source(mut self, source: Arc<dyn Source + Send + Sync>) -> Self {
972        self.source.replace(source);
973        self
974    }
975
976    pub fn with_engine(mut self, engine: Engine) -> Self {
977        self.engine.replace(engine);
978        self
979    }
980
981    pub fn with_module_cache(mut self, module_cache: Arc<dyn ModuleCache + Send + Sync>) -> Self {
982        self.module_cache.replace(module_cache);
983        self
984    }
985
986    pub fn with_tty(mut self, tty: Arc<dyn TtyBridge + Send + Sync>) -> Self {
987        self.tty.replace(tty);
988        self
989    }
990
991    /// Registers a hook that only creates additional imports.
992    pub fn with_additional_imports(
993        self,
994        imports: impl Fn(&wasmer::Module, &mut wasmer::StoreMut) -> anyhow::Result<wasmer::Imports>
995        + Send
996        + Sync
997        + 'static,
998    ) -> Self {
999        self.with_instantiation_hook(ImportsOnlyHook(imports))
1000    }
1001
1002    /// Registers a hook that only configures newly created instances.
1003    pub fn with_instance_setup(
1004        self,
1005        callback: impl Fn(
1006            &wasmer::Module,
1007            &mut wasmer::StoreMut,
1008            &wasmer::Instance,
1009            Option<&wasmer::Memory>,
1010        ) -> anyhow::Result<()>
1011        + Send
1012        + Sync
1013        + 'static,
1014    ) -> Self {
1015        self.with_instantiation_hook(InstanceSetupHook(callback))
1016    }
1017
1018    /// Registers a hook that takes part in both phases of instantiation, so it
1019    /// can carry [`InstantiationState`] from its imports to its instance setup.
1020    pub fn with_instantiation_hook(mut self, hook: impl InstantiationHook) -> Self {
1021        self.instantiation_hooks.push(Arc::new(hook));
1022        self
1023    }
1024
1025    #[cfg(feature = "journal")]
1026    pub fn with_read_only_journals(mut self, journals: Vec<Arc<DynReadableJournal>>) -> Self {
1027        self.read_only_journals.replace(journals);
1028        self
1029    }
1030
1031    #[cfg(feature = "journal")]
1032    pub fn with_writable_journals(mut self, journals: Vec<Arc<DynJournal>>) -> Self {
1033        self.writable_journals.replace(journals);
1034        self
1035    }
1036}
1037
1038impl Runtime for OverriddenRuntime {
1039    fn networking(&self) -> &DynVirtualNetworking {
1040        if let Some(net) = self.networking.as_ref() {
1041            net
1042        } else {
1043            self.inner.networking()
1044        }
1045    }
1046
1047    fn task_manager(&self) -> &Arc<dyn VirtualTaskManager> {
1048        if let Some(rt) = self.task_manager.as_ref() {
1049            rt
1050        } else {
1051            self.inner.task_manager()
1052        }
1053    }
1054
1055    fn source(&self) -> Arc<dyn Source + Send + Sync> {
1056        if let Some(source) = self.source.clone() {
1057            source
1058        } else {
1059            self.inner.source()
1060        }
1061    }
1062
1063    fn package_loader(&self) -> Arc<dyn PackageLoader + Send + Sync> {
1064        if let Some(loader) = self.package_loader.clone() {
1065            loader
1066        } else {
1067            self.inner.package_loader()
1068        }
1069    }
1070
1071    fn module_cache(&self) -> Arc<dyn ModuleCache + Send + Sync> {
1072        if let Some(cache) = self.module_cache.clone() {
1073            cache
1074        } else {
1075            self.inner.module_cache()
1076        }
1077    }
1078
1079    fn engine(&self) -> Engine {
1080        if let Some(engine) = self.engine.clone() {
1081            engine
1082        } else {
1083            self.inner.engine()
1084        }
1085    }
1086
1087    fn new_store(&self) -> wasmer::Store {
1088        if let Some(engine) = self.engine.clone() {
1089            wasmer::Store::new(engine)
1090        } else {
1091            self.inner.new_store()
1092        }
1093    }
1094
1095    fn additional_imports(
1096        &self,
1097        module: &wasmer::Module,
1098        store: &mut wasmer::StoreMut,
1099    ) -> anyhow::Result<(wasmer::Imports, InstantiationState)> {
1100        let (mut imports, inner_state) = self.inner.additional_imports(module, store)?;
1101        if self.instantiation_hooks.is_empty() && inner_state.is_empty() {
1102            return Ok((imports, InstantiationState::empty()));
1103        }
1104        let (own_imports, own_states) = run_import_hooks(&self.instantiation_hooks, module, store)?;
1105        imports.extend(&own_imports);
1106        Ok((
1107            imports,
1108            InstantiationState::new(OverriddenInstantiationState {
1109                inner: inner_state,
1110                own: own_states,
1111            }),
1112        ))
1113    }
1114
1115    fn configure_new_instance(
1116        &self,
1117        module: &wasmer::Module,
1118        store: &mut wasmer::StoreMut,
1119        instance: &wasmer::Instance,
1120        imported_memory: Option<&wasmer::Memory>,
1121        state: InstantiationState,
1122    ) -> anyhow::Result<()> {
1123        let state = if state.is_empty() {
1124            anyhow::ensure!(
1125                self.instantiation_hooks.is_empty(),
1126                "missing instance setup state from additional_imports"
1127            );
1128            OverriddenInstantiationState {
1129                inner: InstantiationState::empty(),
1130                own: Vec::new(),
1131            }
1132        } else {
1133            state
1134                .take::<OverriddenInstantiationState>()
1135                .context("invalid instance setup state from additional_imports")?
1136        };
1137        self.inner
1138            .configure_new_instance(module, store, instance, imported_memory, state.inner)?;
1139        run_setup_hooks(
1140            &self.instantiation_hooks,
1141            state.own,
1142            module,
1143            store,
1144            instance,
1145            imported_memory,
1146        )
1147    }
1148
1149    fn http_client(&self) -> Option<&DynHttpClient> {
1150        if let Some(client) = self.http_client.as_ref() {
1151            Some(client)
1152        } else {
1153            self.inner.http_client()
1154        }
1155    }
1156
1157    fn tty(&self) -> Option<&(dyn TtyBridge + Send + Sync)> {
1158        if let Some(tty) = self.tty.as_ref() {
1159            Some(tty.deref())
1160        } else {
1161            self.inner.tty()
1162        }
1163    }
1164
1165    #[cfg(feature = "journal")]
1166    fn read_only_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynReadableJournal>> + 'a> {
1167        if let Some(journals) = self.read_only_journals.as_ref() {
1168            Box::new(journals.iter().cloned())
1169        } else {
1170            self.inner.read_only_journals()
1171        }
1172    }
1173
1174    #[cfg(feature = "journal")]
1175    fn writable_journals<'a>(&'a self) -> Box<dyn Iterator<Item = Arc<DynJournal>> + 'a> {
1176        if let Some(journals) = self.writable_journals.as_ref() {
1177            Box::new(journals.iter().cloned())
1178        } else {
1179            self.inner.writable_journals()
1180        }
1181    }
1182
1183    #[cfg(feature = "journal")]
1184    fn active_journal(&self) -> Option<&'_ DynJournal> {
1185        if let Some(journals) = self.writable_journals.as_ref() {
1186            journals.iter().last().map(|a| a.as_ref())
1187        } else {
1188            self.inner.active_journal()
1189        }
1190    }
1191}
1192
1193#[cfg(test)]
1194mod tests {
1195    use super::InstantiationState;
1196
1197    #[test]
1198    fn instantiation_state_round_trips_the_hook_data() {
1199        let state = InstantiationState::new(42u32);
1200        assert!(!state.is_empty());
1201        assert_eq!(state.take::<u32>().unwrap(), 42);
1202    }
1203
1204    #[test]
1205    fn empty_instantiation_state_carries_nothing() {
1206        let state = InstantiationState::empty();
1207        assert!(state.is_empty());
1208        let err = state.take::<u32>().unwrap_err();
1209        assert!(err.to_string().contains("missing instantiation state"));
1210    }
1211
1212    #[test]
1213    fn instantiation_state_from_another_hook_is_rejected() {
1214        // What a hook receiving state that isn't its own must see, rather than
1215        // silently operating on another instantiation's data.
1216        let state = InstantiationState::new("some other hook's state");
1217        let err = state.take::<u32>().unwrap_err();
1218        assert!(err.to_string().contains("does not belong to this hook"));
1219    }
1220}