Remove unsoundness of widening store borrows

crates/wasmtime/src/runtime/externals/table.rs

@@ -3,7 +3,7 @@ use crate::runtime::RootedGcRefImpl;
 use crate::runtime::vm::{
     self, GcStore, SendSyncPtr, TableElementType, VMFuncRef, VMGcRef, VMStore,
 };
-use crate::store::{AutoAssertNoGc, StoreInstanceId, StoreOpaque};
+use crate::store::{AutoAssertNoGc, StoreInstanceId, StoreOpaque, StoreResourceLimiter};
 use crate::trampoline::generate_table_export;
 use crate::{
     AnyRef, AsContext, AsContextMut, ExnRef, ExternRef, Func, HeapType, Ref, RefType,
@@ -94,7 +94,8 @@ impl Table {
     /// # }
     /// ```
     pub fn new(mut store: impl AsContextMut, ty: TableType, init: Ref) -> Result<Table> {
-        vm::one_poll(Table::_new(store.as_context_mut().0, ty, init))
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        vm::one_poll(Table::_new(store, limiter.as_mut(), ty, init))
             .expect("must use `new_async` when async resource limiters are in use")
     }
 
@@ -112,11 +113,17 @@ impl Table {
         ty: TableType,
         init: Ref,
     ) -> Result<Table> {
-        Table::_new(store.as_context_mut().0, ty, init).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Table::_new(store, limiter.as_mut(), ty, init).await
     }
 
-    async fn _new(store: &mut StoreOpaque, ty: TableType, init: Ref) -> Result<Table> {
-        let table = generate_table_export(store, &ty).await?;
+    async fn _new(
+        store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
+        ty: TableType,
+        init: Ref,
+    ) -> Result<Table> {
+        let table = generate_table_export(store, limiter, &ty).await?;
         table._fill(store, 0, init, ty.minimum())?;
         Ok(table)
     }

crates/wasmtime/src/runtime/gc/enabled/arrayref.rs

@@ -1,7 +1,7 @@
 //! Working with GC `array` objects.
 
-use crate::runtime::vm::VMGcRef;
-use crate::store::StoreId;
+use crate::runtime::vm::{VMGcRef, VMStore};
+use crate::store::{StoreId, StoreResourceLimiter};
 use crate::vm::{self, VMArrayRef, VMGcHeader};
 use crate::{AnyRef, FieldType};
 use crate::{
@@ -297,9 +297,15 @@ impl ArrayRef {
         elem: &Val,
         len: u32,
     ) -> Result<Rooted<ArrayRef>> {
-        let store = store.as_context_mut().0;
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
         assert!(!store.async_support());
-        vm::assert_ready(Self::_new_async(store, allocator, elem, len))
+        vm::assert_ready(Self::_new_async(
+            store,
+            limiter.as_mut(),
+            allocator,
+            elem,
+            len,
+        ))
     }
 
     /// Asynchronously allocate a new `array` of the given length, with every
@@ -341,17 +347,19 @@ impl ArrayRef {
         elem: &Val,
         len: u32,
     ) -> Result<Rooted<ArrayRef>> {
-        Self::_new_async(store.as_context_mut().0, allocator, elem, len).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Self::_new_async(store, limiter.as_mut(), allocator, elem, len).await
     }
 
     pub(crate) async fn _new_async(
         store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         allocator: &ArrayRefPre,
         elem: &Val,
         len: u32,
     ) -> Result<Rooted<ArrayRef>> {
         store
-            .retry_after_gc_async((), |store, ()| {
+            .retry_after_gc_async(limiter, (), |store, ()| {
                 Self::new_from_iter(store, allocator, RepeatN(elem, len))
             })
             .await
@@ -445,9 +453,14 @@ impl ArrayRef {
         allocator: &ArrayRefPre,
         elems: &[Val],
     ) -> Result<Rooted<ArrayRef>> {
-        let store = store.as_context_mut().0;
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
         assert!(!store.async_support());
-        vm::assert_ready(Self::_new_fixed_async(store, allocator, elems))
+        vm::assert_ready(Self::_new_fixed_async(
+            store,
+            limiter.as_mut(),
+            allocator,
+            elems,
+        ))
     }
 
     /// Asynchronously allocate a new `array` containing the given elements.
@@ -491,16 +504,18 @@ impl ArrayRef {
         allocator: &ArrayRefPre,
         elems: &[Val],
     ) -> Result<Rooted<ArrayRef>> {
-        Self::_new_fixed_async(store.as_context_mut().0, allocator, elems).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Self::_new_fixed_async(store, limiter.as_mut(), allocator, elems).await
     }
 
     pub(crate) async fn _new_fixed_async(
         store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         allocator: &ArrayRefPre,
         elems: &[Val],
     ) -> Result<Rooted<ArrayRef>> {
         store
-            .retry_after_gc_async((), |store, ()| {
+            .retry_after_gc_async(limiter, (), |store, ()| {
                 Self::new_from_iter(store, allocator, elems.iter())
             })
             .await

crates/wasmtime/src/runtime/gc/enabled/exnref.rs

@@ -1,7 +1,7 @@
 //! Implementation of `exnref` in Wasmtime.
 
-use crate::runtime::vm::VMGcRef;
-use crate::store::StoreId;
+use crate::runtime::vm::{VMGcRef, VMStore};
+use crate::store::{StoreId, StoreResourceLimiter};
 use crate::vm::{self, VMExnRef, VMGcHeader};
 use crate::{
     AsContext, AsContextMut, GcRefImpl, GcRootIndex, HeapType, ManuallyRooted, RefType, Result,
@@ -205,9 +205,15 @@ impl ExnRef {
         tag: &Tag,
         fields: &[Val],
     ) -> Result<Rooted<ExnRef>> {
-        let store = store.as_context_mut().0;
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
         assert!(!store.async_support());
-        vm::assert_ready(Self::_new_async(store, allocator, tag, fields))
+        vm::assert_ready(Self::_new_async(
+            store,
+            limiter.as_mut(),
+            allocator,
+            tag,
+            fields,
+        ))
     }
 
     /// Asynchronously allocate a new exception object and get a
@@ -245,18 +251,20 @@ impl ExnRef {
         tag: &Tag,
         fields: &[Val],
     ) -> Result<Rooted<ExnRef>> {
-        Self::_new_async(store.as_context_mut().0, allocator, tag, fields).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Self::_new_async(store, limiter.as_mut(), allocator, tag, fields).await
     }
 
     pub(crate) async fn _new_async(
         store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         allocator: &ExnRefPre,
         tag: &Tag,
         fields: &[Val],
     ) -> Result<Rooted<ExnRef>> {
         Self::type_check_tag_and_fields(store, allocator, tag, fields)?;
         store
-            .retry_after_gc_async((), |store, ()| {
+            .retry_after_gc_async(limiter, (), |store, ()| {
                 Self::new_unchecked(store, allocator, tag, fields)
             })
             .await

crates/wasmtime/src/runtime/gc/enabled/externref.rs

@@ -2,11 +2,11 @@
 
 use super::{AnyRef, RootedGcRefImpl};
 use crate::prelude::*;
-use crate::runtime::vm::{self, VMGcRef};
+use crate::runtime::vm::{self, VMGcRef, VMStore};
 use crate::{
     AsContextMut, GcHeapOutOfMemory, GcRefImpl, GcRootIndex, HeapType, ManuallyRooted, RefType,
     Result, Rooted, StoreContext, StoreContextMut, ValRaw, ValType, WasmTy,
-    store::{AutoAssertNoGc, StoreOpaque},
+    store::{AutoAssertNoGc, StoreOpaque, StoreResourceLimiter},
 };
 use core::any::Any;
 use core::mem;
@@ -210,13 +210,13 @@ impl ExternRef {
     /// Panics if the `context` is configured for async; use
     /// [`ExternRef::new_async`][crate::ExternRef::new_async] to perform
     /// asynchronous allocation instead.
-    pub fn new<T>(mut context: impl AsContextMut, value: T) -> Result<Rooted<ExternRef>>
+    pub fn new<T>(mut store: impl AsContextMut, value: T) -> Result<Rooted<ExternRef>>
     where
         T: 'static + Any + Send + Sync,
     {
-        let ctx = context.as_context_mut().0;
-        assert!(!ctx.async_support());
-        vm::assert_ready(Self::_new_async(ctx, value))
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        assert!(!store.async_support());
+        vm::assert_ready(Self::_new_async(store, limiter.as_mut(), value))
     }
 
     /// Asynchronously allocates a new `ExternRef` wrapping the given value.
@@ -295,16 +295,17 @@ impl ExternRef {
     /// [`ExternRef::new`][crate::ExternRef::new] to perform synchronous
     /// allocation instead.
     #[cfg(feature = "async")]
-    pub async fn new_async<T>(mut context: impl AsContextMut, value: T) -> Result<Rooted<ExternRef>>
+    pub async fn new_async<T>(mut store: impl AsContextMut, value: T) -> Result<Rooted<ExternRef>>
     where
         T: 'static + Any + Send + Sync,
     {
-        let ctx = context.as_context_mut().0;
-        Self::_new_async(ctx, value).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Self::_new_async(store, limiter.as_mut(), value).await
     }
 
     pub(crate) async fn _new_async<T>(
         store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         value: T,
     ) -> Result<Rooted<ExternRef>>
     where
@@ -315,7 +316,7 @@ impl ExternRef {
         let value: Box<dyn Any + Send + Sync> = Box::new(value);
 
         let gc_ref = store
-            .retry_after_gc_async(value, |store, value| {
+            .retry_after_gc_async(limiter, value, |store, value| {
                 store
                     .require_gc_store_mut()?
                     .alloc_externref(value)

crates/wasmtime/src/runtime/gc/enabled/structref.rs

@@ -2,13 +2,13 @@
 
 use crate::runtime::vm::VMGcRef;
 use crate::store::StoreId;
-use crate::vm::{self, VMGcHeader, VMStructRef};
+use crate::vm::{self, VMGcHeader, VMStore, VMStructRef};
 use crate::{AnyRef, FieldType};
 use crate::{
     AsContext, AsContextMut, EqRef, GcHeapOutOfMemory, GcRefImpl, GcRootIndex, HeapType,
     ManuallyRooted, RefType, Rooted, StructType, Val, ValRaw, ValType, WasmTy,
     prelude::*,
-    store::{AutoAssertNoGc, StoreContextMut, StoreOpaque},
+    store::{AutoAssertNoGc, StoreContextMut, StoreOpaque, StoreResourceLimiter},
 };
 use core::mem::{self, MaybeUninit};
 use wasmtime_environ::{GcLayout, GcStructLayout, VMGcKind, VMSharedTypeIndex};
@@ -231,9 +231,9 @@ impl StructRef {
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
-        let store = store.as_context_mut().0;
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
         assert!(!store.async_support());
-        vm::assert_ready(Self::_new_async(store, allocator, fields))
+        vm::assert_ready(Self::_new_async(store, limiter.as_mut(), allocator, fields))
     }
 
     /// Asynchronously allocate a new `struct` and get a reference to it.
@@ -264,17 +264,19 @@ impl StructRef {
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
-        Self::_new_async(store.as_context_mut().0, allocator, fields).await
+        let (mut limiter, store) = store.as_context_mut().0.resource_limiter_and_store_opaque();
+        Self::_new_async(store, limiter.as_mut(), allocator, fields).await
     }
 
     pub(crate) async fn _new_async(
         store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
         Self::type_check_fields(store, allocator, fields)?;
         store
-            .retry_after_gc_async((), |store, ()| {
+            .retry_after_gc_async(limiter, (), |store, ()| {
                 Self::new_unchecked(store, allocator, fields)
             })
             .await

crates/wasmtime/src/runtime/instance.rs

@@ -2,9 +2,11 @@ use crate::linker::{Definition, DefinitionType};
 use crate::prelude::*;
 use crate::runtime::vm::{
     self, Imports, ModuleRuntimeInfo, VMFuncRef, VMFunctionImport, VMGlobalImport, VMMemoryImport,
-    VMTableImport, VMTagImport,
+    VMStore, VMTableImport, VMTagImport,
+};
+use crate::store::{
+    AllocateInstanceKind, InstanceId, StoreInstanceId, StoreOpaque, StoreResourceLimiter,
 };
-use crate::store::{AllocateInstanceKind, InstanceId, StoreInstanceId, StoreOpaque};
 use crate::types::matching;
 use crate::{
     AsContextMut, Engine, Export, Extern, Func, Global, Memory, Module, ModuleExport, SharedMemory,
@@ -248,9 +250,12 @@ impl Instance {
         module: &Module,
         imports: Imports<'_>,
     ) -> Result<Instance> {
-        // SAFETY: the safety contract of `new_raw` is the same as this
-        // function.
-        let (instance, start) = unsafe { Instance::new_raw(store.0, module, imports).await? };
+        let (instance, start) = {
+            let (mut limiter, store) = store.0.resource_limiter_and_store_opaque();
+            // SAFETY: the safety contract of `new_raw` is the same as this
+            // function.
+            unsafe { Instance::new_raw(store, limiter.as_mut(), module, imports).await? }
+        };
         if let Some(start) = start {
             if store.0.async_support() {
                 #[cfg(feature = "async")]
@@ -285,6 +290,7 @@ impl Instance {
     /// provided as well.
     async unsafe fn new_raw(
         store: &mut StoreOpaque,
+        mut limiter: Option<&mut StoreResourceLimiter<'_>>,
         module: &Module,
         imports: Imports<'_>,
     ) -> Result<(Instance, Option<FuncIndex>)> {
@@ -295,7 +301,7 @@ impl Instance {
 
         // Allocate the GC heap, if necessary.
         if module.env_module().needs_gc_heap {
-            store.ensure_gc_store().await?;
+            store.ensure_gc_store(limiter.as_deref_mut()).await?;
         }
 
         let compiled_module = module.compiled_module();
@@ -313,6 +319,7 @@ impl Instance {
         let id = unsafe {
             store
                 .allocate_instance(
+                    limiter.as_deref_mut(),
                     AllocateInstanceKind::Module(module_id),
                     &ModuleRuntimeInfo::Module(module.clone()),
                     imports,
@@ -349,7 +356,7 @@ impl Instance {
             .features()
             .contains(WasmFeatures::BULK_MEMORY);
 
-        vm::initialize_instance(store, id, compiled_module.module(), bulk_memory).await?;
+        vm::initialize_instance(store, limiter, id, compiled_module.module(), bulk_memory).await?;
 
         Ok((instance, compiled_module.module().start_func))
     }