Fixed validation of instantiation before start fn; add component builtin test

Author: arjunr2

crates/wasmtime/src/runtime/component/instance.rs

@@ -9,7 +9,8 @@ use crate::instance::OwnedImports;
 use crate::linker::DefinitionType;
 use crate::prelude::*;
 #[cfg(feature = "rr")]
-use crate::rr::{RecordBuffer, component_events::InstantiationEvent};
+use crate::rr::RecordBuffer;
+use crate::rr::component_hooks;
 use crate::runtime::vm::component::{
     CallContexts, ComponentInstance, ResourceTables, TypedResource, TypedResourceIndex,
 };
@@ -1080,7 +1081,6 @@ impl<'a> Instantiator<'a> {
 
     /// Convenience helper to return the instance ID of the `ComponentInstance` that's
     /// being instantiated
-    #[cfg(feature = "rr")]
     fn id(&self) -> ComponentInstanceId {
         self.id
     }
@@ -1194,12 +1194,13 @@ impl<T: 'static> InstancePre<T> {
             .increment_component_instance_count()?;
         let mut instantiator = Instantiator::new(&self.component, store.0, &self.imports);
 
+        // Record/replay hooks
         store.0.validate_rr_config()?;
-        #[cfg(feature = "rr")]
-        store.0.record_event(|| InstantiationEvent {
-            component: *self.component.checksum(),
-            instance: instantiator.id(),
-        })?;
+        component_hooks::record_and_replay_validate_instantiation(
+            &mut store,
+            *self.component.checksum(),
+            instantiator.id(),
+        )?;
 
         instantiator.run(&mut store).await.map_err(|e| {
             store

crates/wasmtime/src/runtime/instance.rs

@@ -1,7 +1,6 @@
 use crate::linker::{Definition, DefinitionType};
 use crate::prelude::*;
-#[cfg(feature = "rr")]
-use crate::rr::core_events::InstantiationEvent;
+use crate::rr::core_hooks;
 use crate::runtime::vm::{
     self, Imports, ModuleRuntimeInfo, VMFuncRef, VMFunctionImport, VMGlobalImport, VMMemoryImport,
     VMStore, VMTableImport, VMTagImport,
@@ -262,18 +261,18 @@ impl Instance {
             unsafe { Instance::new_raw(store, limiter.as_mut(), module, imports).await? }
         };
 
+        // Record/replay hooks
         store.0.validate_rr_config()?;
         if !from_component {
-            #[cfg(feature = "rr")]
-            {
-                // Components already record instantiation, so do not record their internal modules
-                rr_validate_module_unexported_memory(module)?;
-                store.0.record_event(|| InstantiationEvent {
-                    module: *module.checksum(),
-                    instance: instance.id(),
-                })?;
-            }
+            // Components already record instantiation, so do not record their internal modules
+            core_hooks::rr_validate_module_unexported_memory(&module)?;
+            core_hooks::record_and_replay_validate_instantiation(
+                store,
+                *module.checksum(),
+                instance.id(),
+            )?;
         }
+
         if let Some(start) = start {
             if store.0.async_support() {
                 #[cfg(feature = "async")]
@@ -1008,22 +1007,3 @@ fn typecheck<I>(
     }
     Ok(())
 }
-
-/// Ensure that memories are not exported memories in Core wasm modules when
-/// recording is enabled
-#[cfg(feature = "rr")]
-fn rr_validate_module_unexported_memory(module: &Module) -> Result<()> {
-    // Check for exported memories when recording is enabled.
-    if module.engine().is_recording()
-        && module.exports().any(|export| {
-            if let crate::ExternType::Memory(_) = export.ty() {
-                true
-            } else {
-                false
-            }
-        })
-    {
-        bail!("Cannot support recording for core wasm modules when a memory is exported");
-    }
-    Ok(())
-}

crates/wasmtime/src/runtime/rr/core.rs

@@ -576,14 +576,20 @@ mod tests {
     use wasmtime_environ::FuncIndex;
 
     impl ReplayBuffer {
+        /// Pop the next replay event and calls `f` with a expected event type
+        ///
+        /// ## Errors
+        ///
+        /// See [`next_event_typed`](Replayer::next_event_typed)
+        #[inline]
         fn next_event_and<T, F>(&mut self, f: F) -> Result<(), ReplayError>
         where
             T: TryFrom<RREvent>,
             ReplayError: From<<T as TryFrom<RREvent>>::Error>,
             F: FnOnce(T) -> Result<(), ReplayError>,
         {
-            let event = self.next_event_typed::<T>()?;
-            f(event)
+            let call_event = self.next_event_typed()?;
+            Ok(f(call_event)?)
         }
     }
 

crates/wasmtime/src/runtime/rr/hooks.rs

@@ -1,15 +1,17 @@
-/// Component specific RR hooks that use `component-model` feature gating
-#[cfg(feature = "component-model")]
-pub mod component_hooks;
-/// Core RR hooks
-pub mod core_hooks;
-
 use crate::{FuncType, WasmFuncOrigin};
+#[cfg(feature = "rr")]
+use crate::{StoreContextMut, rr::ReplayHostContext};
 #[cfg(feature = "component-model")]
 use alloc::sync::Arc;
 #[cfg(feature = "component-model")]
 use wasmtime_environ::component::{ComponentTypes, TypeFuncIndex};
 
+/// Component specific RR hooks that use `component-model` feature gating
+#[cfg(feature = "component-model")]
+pub mod component_hooks;
+/// Core RR hooks
+pub mod core_hooks;
+
 /// Wasm function type information for RR hooks
 pub enum RRWasmFuncType<'a> {
     /// Core RR hooks to be performed
@@ -27,3 +29,29 @@ pub enum RRWasmFuncType<'a> {
     #[cfg(feature = "component-model")]
     None,
 }
+
+/// Obtain the replay host context from the store.
+///
+/// SAFETY: The store's data is always of type `ReplayHostContext` when created by
+/// the replay driver. As an additional guarantee, we assert that replay is indeed
+/// truly enabled.
+#[cfg(feature = "rr")]
+unsafe fn replay_data_from_store<'a, T: 'static>(
+    store: &StoreContextMut<'a, T>,
+) -> &'a ReplayHostContext {
+    assert!(store.0.replay_enabled());
+    let raw_ptr: *const T = store.data();
+    unsafe { &*(raw_ptr as *const ReplayHostContext) }
+}
+
+/// Same as [replay_data_from_store], but mutable
+///
+/// SAFETY: See [replay_data_from_store]
+#[cfg(feature = "rr")]
+unsafe fn replay_data_from_store_mut<'a, T: 'static>(
+    store: &mut StoreContextMut<'a, T>,
+) -> &'a mut ReplayHostContext {
+    assert!(store.0.replay_enabled());
+    let raw_ptr: *mut T = store.data_mut();
+    unsafe { &mut *(raw_ptr as *mut ReplayHostContext) }
+}

crates/wasmtime/src/runtime/rr/hooks/component_hooks.rs

@@ -1,21 +1,23 @@
+#[cfg(feature = "rr")]
+use super::replay_data_from_store_mut;
 use crate::ValRaw;
 use crate::component::{ComponentInstanceId, func::LowerContext};
 #[cfg(feature = "rr")]
-use crate::rr::common_events::{HostFuncEntryEvent, WasmFuncReturnEvent};
+use crate::rr::common_events::{HostFuncEntryEvent, HostFuncReturnEvent, WasmFuncReturnEvent};
 #[cfg(feature = "rr")]
 use crate::rr::component_events::{
-    LowerFlatEntryEvent, LowerFlatReturnEvent, LowerMemoryEntryEvent, LowerMemoryReturnEvent,
-    MemorySliceWriteEvent, PostReturnEvent, WasmFuncBeginEvent, WasmFuncEntryEvent,
+    InstantiationEvent, LowerFlatEntryEvent, LowerFlatReturnEvent, LowerMemoryEntryEvent,
+    LowerMemoryReturnEvent, MemorySliceWriteEvent, PostReturnEvent, WasmFuncBeginEvent,
+    WasmFuncEntryEvent,
 };
 #[cfg(feature = "rr")]
-use crate::rr::{
-    RRFuncArgVals, RecordBuffer, Recorder, ResultEvent, common_events::HostFuncReturnEvent,
-};
+use crate::rr::{RRFuncArgVals, RecordBuffer, Recorder, ResultEvent, Validate};
 use crate::store::StoreOpaque;
 use crate::{StoreContextMut, prelude::*};
 use alloc::sync::Arc;
 use core::mem::MaybeUninit;
 use core::ops::{Deref, DerefMut};
+use wasmtime_environ::WasmChecksum;
 use wasmtime_environ::component::{ComponentTypes, ExportIndex, InterfaceType, TypeFuncIndex};
 #[cfg(all(feature = "rr"))]
 use wasmtime_environ::component::{MAX_FLAT_PARAMS, MAX_FLAT_RESULTS};
@@ -206,6 +208,31 @@ where
     lower_result
 }
 
+/// Hook for recording a component instantiation event and validating the
+/// instantiation during replay.
+#[inline]
+pub fn record_and_replay_validate_instantiation<T>(
+    store: &mut StoreContextMut<'_, T>,
+    component: WasmChecksum,
+    instance: ComponentInstanceId,
+) -> Result<()> {
+    #[cfg(feature = "rr")]
+    {
+        store.0.record_event(|| InstantiationEvent {
+            component,
+            instance,
+        })?;
+        if store.0.replay_enabled() {
+            let replay_data = unsafe { replay_data_from_store_mut(store) };
+            replay_data.take_current_component_instantiation().expect(
+                "replay driver should have set component instantiate data before trying to validate it",
+            ).validate(&InstantiationEvent { component, instance })?;
+        }
+    }
+    let _ = (store, component, instance);
+    Ok(())
+}
+
 #[cfg(feature = "rr")]
 #[inline(always)]
 fn create_host_func_entry_event(

crates/wasmtime/src/runtime/rr/hooks/core_hooks.rs

@@ -1,14 +1,18 @@
+#[cfg(feature = "rr")]
+use super::{replay_data_from_store, replay_data_from_store_mut};
 use crate::rr::FlatBytes;
 #[cfg(feature = "rr")]
 use crate::rr::{
-    RREvent, RRFuncArgVals, ReplayError, ReplayHostContext, Replayer, ResultEvent,
+    RREvent, RRFuncArgVals, ReplayError, Replayer, ResultEvent, Validate,
     common_events::HostFuncEntryEvent, common_events::HostFuncReturnEvent,
-    common_events::WasmFuncReturnEvent, core_events::WasmFuncEntryEvent,
+    common_events::WasmFuncReturnEvent, core_events::InstantiationEvent,
+    core_events::WasmFuncEntryEvent,
 };
-use crate::store::StoreOpaque;
-use crate::{Caller, FuncType, StoreContextMut, ValRaw, WasmFuncOrigin, prelude::*};
+use crate::store::{InstanceId, StoreOpaque};
+use crate::{Caller, FuncType, Module, StoreContextMut, ValRaw, WasmFuncOrigin, prelude::*};
 #[cfg(feature = "rr")]
 use wasmtime_environ::EntityIndex;
+use wasmtime_environ::WasmChecksum;
 
 /// Record and replay hook operation for core wasm function entry events
 ///
@@ -121,7 +125,7 @@ where
     Ok(())
 }
 
-/// Replay hook operation for host function return events
+/// Replay hook operation for host function return events.
 #[inline]
 pub fn replay_host_func_return<T, U: 'static>(
     args: &mut [T],
@@ -146,14 +150,9 @@ where
                 RREvent::CoreWasmFuncEntry(event) => {
                     let entity = EntityIndex::from(event.origin.index);
 
-                    // SAFETY: The store's data is always of type `ReplayHostContext<T>` when created by
-                    // the replay driver. As an additional guarantee, we assert that replay is indeed
-                    // truly enabled.
-                    assert!(caller.store.0.replay_enabled());
-                    let replay_data = unsafe {
-                        let raw_ptr: *const U = caller.store.data();
-                        &*(raw_ptr as *const ReplayHostContext)
-                    };
+                    // Unwrapping the `replay_buffer_mut()` above ensures that we are in replay mode
+                    // passing the safety contract for `replay_data_from_store`
+                    let replay_data = unsafe { replay_data_from_store(&caller.store) };
 
                     // Grab the correct module instance
                     let instance = replay_data.get_module_instance(event.origin.instance)?;
@@ -194,3 +193,48 @@ where
     let _ = (args, caller);
     Ok(())
 }
+
+/// Hook for recording a module instantiation event and validating the
+/// instantiation during replay.
+pub fn record_and_replay_validate_instantiation<T: 'static>(
+    store: &mut StoreContextMut<'_, T>,
+    module: WasmChecksum,
+    instance: InstanceId,
+) -> Result<()> {
+    #[cfg(feature = "rr")]
+    {
+        store
+            .0
+            .record_event(|| InstantiationEvent { module, instance })?;
+        if store.0.replay_enabled() {
+            let replay_data = unsafe { replay_data_from_store_mut(store) };
+            replay_data.take_current_module_instantiation().expect(
+                "replay driver should have set module instantiate data before trying to validate it",
+            ).validate(&InstantiationEvent { module, instance })?;
+        }
+    }
+    let _ = (store, module, instance);
+    Ok(())
+}
+
+/// Ensure that memories are not exported memories in Core wasm modules when
+/// recording is enabled.
+pub fn rr_validate_module_unexported_memory(module: &Module) -> Result<()> {
+    // Check for exported memories when recording is enabled.
+    #[cfg(feature = "rr")]
+    {
+        if module.engine().is_recording()
+            && module.exports().any(|export| {
+                if let crate::ExternType::Memory(_) = export.ty() {
+                    true
+                } else {
+                    false
+                }
+            })
+        {
+            bail!("Cannot support recording for core wasm modules when a memory is exported");
+        }
+    }
+    let _ = module;
+    Ok(())
+}