From 8bde231f1e7da7085ba836168c5f09be9b9ae82e Mon Sep 17 00:00:00 2001 From: Ryan Levick Date: Tue, 21 May 2024 20:59:16 +0200 Subject: [PATCH] Implement `Component::define_unknown_imports_as_traps` (#8672) * Implement Component::define_unknown_imports_as_traps This will search through a components imports, find any imports that have not yet been defined in the linker and add a definition which will trap upon being called. * Address PR feedback Signed-off-by: Ryan Levick * Small refactoring Signed-off-by: Ryan Levick --------- Signed-off-by: Ryan Levick --- .../wasmtime/src/runtime/component/linker.rs | 71 +++++++++++++++++++ crates/wasmtime/src/runtime/linker.rs | 3 +- tests/all/component_model/linker.rs | 36 ++++++++++ 3 files changed, 108 insertions(+), 2 deletions(-) diff --git a/crates/wasmtime/src/runtime/component/linker.rs b/crates/wasmtime/src/runtime/component/linker.rs index 59487bb65508..6abe76045346 100644 --- a/crates/wasmtime/src/runtime/component/linker.rs +++ b/crates/wasmtime/src/runtime/component/linker.rs @@ -336,6 +336,73 @@ impl Linker { .instantiate_async(store) .await } + + /// Implement any imports of the given [`Component`] with a function which traps. + /// + /// By default a [`Linker`] will error when unknown imports are encountered when instantiating a [`Component`]. + /// This changes this behavior from an instant error to a trap that will happen if the import is called. + pub fn define_unknown_imports_as_traps(&mut self, component: &Component) -> Result<()> { + use wasmtime_environ::component::ComponentTypes; + use wasmtime_environ::component::TypeDef; + // Recursively stub out all imports of the component with a function that traps. + fn stub_item( + linker: &mut LinkerInstance, + item_name: &str, + item_def: &TypeDef, + parent_instance: Option<&str>, + types: &ComponentTypes, + ) -> Result<()> { + // Skip if the item isn't an instance and has already been defined in the linker. + if !matches!(item_def, TypeDef::ComponentInstance(_)) && linker.get(item_name).is_some() + { + return Ok(()); + } + + match item_def { + TypeDef::ComponentFunc(_) => { + let fully_qualified_name = parent_instance + .map(|parent| format!("{}#{}", parent, item_name)) + .unwrap_or_else(|| item_name.to_owned()); + linker.func_new(&item_name, move |_, _, _| { + bail!("unknown import: `{fully_qualified_name}` has not been defined") + })?; + } + TypeDef::ComponentInstance(i) => { + let instance = &types[*i]; + let mut linker_instance = linker.instance(item_name)?; + for (export_name, export) in instance.exports.iter() { + stub_item( + &mut linker_instance, + export_name, + export, + Some(item_name), + types, + )?; + } + } + TypeDef::Resource(_) => { + let ty = crate::component::ResourceType::host::<()>(); + linker.resource(item_name, ty, |_, _| Ok(()))?; + } + TypeDef::Component(_) | TypeDef::Module(_) => { + bail!("unable to define {} imports as traps", item_def.desc()) + } + _ => {} + } + Ok(()) + } + + for (_, (import_name, import_type)) in &component.env_component().import_types { + stub_item( + &mut self.root(), + import_name, + import_type, + None, + component.types(), + )?; + } + Ok(()) + } } impl LinkerInstance<'_, T> { @@ -616,6 +683,10 @@ impl LinkerInstance<'_, T> { self.map .insert(name, &mut self.strings, self.allow_shadowing, item) } + + fn get(&self, name: &str) -> Option<&Definition> { + self.map.get(name, &self.strings) + } } impl NameMap { diff --git a/crates/wasmtime/src/runtime/linker.rs b/crates/wasmtime/src/runtime/linker.rs index 6600f39bbdf0..f02c969b8bb1 100644 --- a/crates/wasmtime/src/runtime/linker.rs +++ b/crates/wasmtime/src/runtime/linker.rs @@ -256,8 +256,7 @@ impl Linker { /// Implement any imports of the given [`Module`] with a function which traps. /// /// By default a [`Linker`] will error when unknown imports are encountered - /// in a command module while using [`Linker::module`]. Use this function - /// when + /// in a command module while using [`Linker::module`]. /// /// This method can be used to allow unknown imports from command modules. /// diff --git a/tests/all/component_model/linker.rs b/tests/all/component_model/linker.rs index 73030cacff66..d7c6e4916420 100644 --- a/tests/all/component_model/linker.rs +++ b/tests/all/component_model/linker.rs @@ -145,3 +145,39 @@ fn linker_substituting_types_issue_8003() -> Result<()> { } Ok(()) } + +#[test] +fn linker_defines_unknown_imports_as_traps() -> Result<()> { + let engine = Engine::default(); + let mut linker = Linker::<()>::new(&engine); + + let component = Component::new( + &engine, + r#"(component + (import "foo" (func)) + (import "bar" (instance (export "baz" (func)))) + (import "qux" (type (sub resource))) + )"#, + )?; + linker.define_unknown_imports_as_traps(&component)?; + let mut store = Store::new(&engine, ()); + let _ = linker.instantiate(&mut store, &component)?; + + Ok(()) +} + +#[test] +fn linker_fails_to_define_unknown_core_module_imports_as_traps() -> Result<()> { + let engine = Engine::default(); + let mut linker = Linker::<()>::new(&engine); + + let component = Component::new( + &engine, + r#"(component + (import "foo" (core module)) + )"#, + )?; + assert!(linker.define_unknown_imports_as_traps(&component).is_err()); + + Ok(()) +}