wasmtime 42.0.2

#[cfg(feature = "component-model-async")]
use crate::runtime::vm::VMStore;
use crate::runtime::vm::component::{ComponentInstance, OwnedComponentInstance};
use crate::store::{StoreData, StoreId, StoreOpaque};
use crate::{AsContext, AsContextMut, Store, StoreContextMut};
#[cfg(feature = "component-model-async")]
use alloc::vec::Vec;
use core::pin::Pin;
use wasmtime_environ::PrimaryMap;
use wasmtime_environ::component::RuntimeComponentInstanceIndex;

/// Default amount of fuel allowed for all guest-to-host calls in the component
/// model.
///
/// This is the maximal amount of data which will be copied from the guest to
/// the host by default. This is set large enough as to not be hit all that
/// often in theory but also small enough such that if left unconfigured on a
/// host doesn't mean that it's automatically susceptible to DoS for example.
const DEFAULT_HOSTCALL_FUEL: usize = 128 << 20;

pub struct ComponentStoreData {
    instances: PrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,

    /// Whether an instance belonging to this store has trapped.
    trapped: bool,

    /// Fuel to be used for each time the guest calls the host or transfers data
    /// to the host.
    ///
    /// Caps the size of the allocations made on the host to this amount
    /// effectively.
    hostcall_fuel: usize,
}

impl Default for ComponentStoreData {
    fn default() -> Self {
        Self {
            instances: PrimaryMap::default(),
            trapped: false,
            hostcall_fuel: DEFAULT_HOSTCALL_FUEL,
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct ComponentInstanceId(u32);
wasmtime_environ::entity_impl!(ComponentInstanceId);

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct RuntimeInstance {
    pub instance: ComponentInstanceId,
    pub index: RuntimeComponentInstanceIndex,
}

impl StoreOpaque {
    pub(crate) fn trapped(&self) -> bool {
        self.store_data().components.trapped
    }

    pub(crate) fn set_trapped(&mut self) {
        self.store_data_mut().components.trapped = true;
    }
}

impl StoreData {
    pub(crate) fn push_component_instance(
        &mut self,
        data: OwnedComponentInstance,
    ) -> ComponentInstanceId {
        let expected = data.get().id();
        let ret = self.components.instances.push(Some(data));
        assert_eq!(expected, ret);
        ret
    }
}

impl ComponentStoreData {
    pub fn next_component_instance_id(&self) -> ComponentInstanceId {
        self.instances.next_key()
    }

    #[cfg(feature = "component-model-async")]
    pub(crate) fn drop_fibers_and_futures(store: &mut dyn VMStore) {
        let mut fibers = Vec::new();
        let mut futures = Vec::new();
        store
            .concurrent_state_mut()
            .take_fibers_and_futures(&mut fibers, &mut futures);

        for mut fiber in fibers {
            fiber.dispose(store);
        }

        crate::component::concurrent::tls::set(store, move || drop(futures));
    }

    #[cfg(feature = "component-model-async")]
    pub(crate) fn assert_instance_states_empty(&mut self) {
        for (_, instance) in self.instances.iter_mut() {
            let Some(instance) = instance.as_mut() else {
                continue;
            };

            assert!(
                instance
                    .get_mut()
                    .instance_states()
                    .0
                    .iter_mut()
                    .all(|(_, state)| state.handle_table().is_empty()
                        && state.concurrent_state().pending_is_empty())
            );
        }
    }
}

impl StoreData {
    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
        self.components.instances[id].as_ref().unwrap().get()
    }

    pub(crate) fn component_instance_mut(
        &mut self,
        id: ComponentInstanceId,
    ) -> Pin<&mut ComponentInstance> {
        self.components.instances[id].as_mut().unwrap().get_mut()
    }
}

impl StoreOpaque {
    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
        self.store_data().component_instance(id)
    }

    #[cfg(feature = "component-model-async")]
    pub(crate) fn component_instance_mut(
        &mut self,
        id: ComponentInstanceId,
    ) -> Pin<&mut ComponentInstance> {
        self.store_data_mut().component_instance_mut(id)
    }

    pub(crate) fn hostcall_fuel(&self) -> usize {
        self.store_data().components.hostcall_fuel
    }

    pub(crate) fn set_hostcall_fuel(&mut self, fuel: usize) {
        self.store_data_mut().components.hostcall_fuel = fuel;
    }
}

/// A type used to represent an allocated `ComponentInstance` located within a
/// store.
///
/// This type is held in various locations as a "safe index" into a store. This
/// encapsulates a `StoreId` which owns the instance as well as the index within
/// the store's list of which instance it's pointing to.
///
/// This type can notably be used to index into a `StoreOpaque` to project out
/// the `ComponentInstance` that is associated with this id.
#[repr(C)] // used by reference in the C API
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct StoreComponentInstanceId {
    store_id: StoreId,
    instance: ComponentInstanceId,
}

impl StoreComponentInstanceId {
    pub(crate) fn new(
        store_id: StoreId,
        instance: ComponentInstanceId,
    ) -> StoreComponentInstanceId {
        StoreComponentInstanceId { store_id, instance }
    }

    #[inline]
    pub fn assert_belongs_to(&self, store: StoreId) {
        self.store_id.assert_belongs_to(store)
    }

    #[inline]
    pub(crate) fn store_id(&self) -> StoreId {
        self.store_id
    }

    #[inline]
    pub(crate) fn instance(&self) -> ComponentInstanceId {
        self.instance
    }

    /// Looks up the `vm::ComponentInstance` within `store` that this id points
    /// to.
    ///
    /// # Panics
    ///
    /// Panics if `self` does not belong to `store`.
    pub(crate) fn get<'a>(&self, store: &'a StoreOpaque) -> &'a ComponentInstance {
        self.assert_belongs_to(store.id());
        store.component_instance(self.instance)
    }

    /// Mutable version of `get` above.
    ///
    /// # Panics
    ///
    /// Panics if `self` does not belong to `store`.
    pub(crate) fn get_mut<'a>(&self, store: &'a mut StoreOpaque) -> Pin<&'a mut ComponentInstance> {
        self.from_data_get_mut(store.store_data_mut())
    }

    /// Return a mutable `ComponentInstance` and a `ModuleRegistry`
    /// from the store.
    ///
    /// # Panics
    ///
    /// Panics if `self` does not belong to `store`.
    #[cfg(feature = "component-model-async")]
    pub(crate) fn get_mut_and_registry<'a>(
        &self,
        store: &'a mut StoreOpaque,
    ) -> (
        Pin<&'a mut ComponentInstance>,
        &'a crate::module::ModuleRegistry,
    ) {
        let (store_data, registry) = store.store_data_mut_and_registry();
        let instance = self.from_data_get_mut(store_data);
        (instance, registry)
    }

    /// Same as `get_mut`, but borrows less of a store.
    pub(crate) fn from_data_get_mut<'a>(
        &self,
        store: &'a mut StoreData,
    ) -> Pin<&'a mut ComponentInstance> {
        self.assert_belongs_to(store.id());
        store.component_instance_mut(self.instance)
    }
}

impl<T> Store<T> {
    /// Returns the amount of "hostcall fuel" used for guest-to-host component
    /// calls.
    ///
    /// This is either the default amount if it hasn't been configured or
    /// returns the last value passed to [`Store::set_hostcall_fuel`].
    ///
    /// See [`Store::set_hostcall_fuel`] `for more details.
    pub fn hostcall_fuel(&self) -> usize {
        self.as_context().0.hostcall_fuel()
    }

    /// Sets the amount of "hostcall fuel" used for guest-to-host component
    /// calls.
    ///
    /// Whenever the guest calls the host it often wants to transfer some data
    /// as well, such as strings or lists. This configured fuel value can be
    /// used to limit the amount of data that the host allocates on behalf of
    /// the guest. This is a DoS mitigation mechanism to prevent a malicious
    /// guest from causing the host to allocate an unbounded amount of memory
    /// for example.
    ///
    /// Fuel is considered distinct for each host call. The host is responsible
    /// for ensuring it retains a proper amount of data between host calls if
    /// applicable. The `fuel` provided here will be the initial value for each
    /// time the guest calls the host.
    ///
    /// The `fuel` value here should roughly corresponds to the maximal number
    /// of bytes that the guest may transfer to the host in one call.
    ///
    /// Note that data transferred from the host to the guest is not limited
    /// because it's already resident on the host itself. Only data from the
    /// guest to the host is limited.
    ///
    /// The default value for this is 128 MiB.
    pub fn set_hostcall_fuel(&mut self, fuel: usize) {
        self.as_context_mut().set_hostcall_fuel(fuel)
    }
}

impl<T> StoreContextMut<'_, T> {
    /// See [`Store::hostcall_fuel`].
    pub fn hostcall_fuel(&self) -> usize {
        self.0.hostcall_fuel()
    }

    /// See [`Store::set_hostcall_fuel`].
    pub fn set_hostcall_fuel(&mut self, fuel: usize) {
        self.0.set_hostcall_fuel(fuel)
    }
}