From ad36554973d2c3758b221da13bd5919d138eda8d Mon Sep 17 00:00:00 2001 From: Phil Cummins Date: Mon, 21 Oct 2024 12:55:00 +0200 Subject: [PATCH 1/3] adds stubwasi.rs --- src/lib.rs | 126 +--------------------------------------------- src/stubwasi.rs | 130 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 124 deletions(-) create mode 100644 src/stubwasi.rs diff --git a/src/lib.rs b/src/lib.rs index e561d64..e18f353 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,12 +19,6 @@ use { str, }, summary::{Escape, Locations, Summary}, - wasm_convert::IntoValType, - wasm_encoder::{ - CodeSection, ExportKind, ExportSection, Function, FunctionSection, Instruction, Module, - TypeSection, - }, - wasmparser::{FuncType, Parser, Payload, TypeRef}, wasmtime::{ component::{Component, Instance, Linker, ResourceTable, ResourceType}, Config, Engine, Store, @@ -43,6 +37,7 @@ pub mod command; mod prelink; #[cfg(feature = "pyo3")] mod python; +mod stubwasi; mod summary; #[cfg(test)] mod test; @@ -282,16 +277,12 @@ pub async fn componentize( // Link all the libraries (including any native extensions) into a single component. let mut linker = wit_component::Linker::default().validate(true); - let mut wasi_imports = HashMap::new(); for Library { name, module, dl_openable, } in &libraries { - if stub_wasi { - add_wasi_imports(module, &mut wasi_imports)?; - } linker = linker.library(name, module, *dl_openable)?; } @@ -306,51 +297,7 @@ pub async fn componentize( let component = linker.encode()?; let stubbed_component = if stub_wasi { - // When `stub_wasi` is `true`, we apply the pre-initialization snapshot to an alternate version of the - // component -- one where the WASI imports have been stubbed out. - - let mut linker = wit_component::Linker::default().validate(true); - - for Library { - name, - module, - dl_openable, - } in &libraries - { - linker = linker.library(name, module, *dl_openable)?; - } - - for (module, imports) in &wasi_imports { - linker = linker.adapter(module, &make_stub_adapter(module, imports))?; - } - - let component = linker.encode()?; - - // As of this writing, `wit_component::Linker` generates a component such that the first module is the - // `main` one, followed by any adapters, followed by any libraries, followed by the `init` module, which is - // finally followed by any shim modules. Given that the stubbed component may contain more adapters than - // the non-stubbed version, we need to tell `component-init` how to translate module indexes from the - // former to the latter. - // - // TODO: this is pretty fragile in that it could silently break if `wit_component::Linker`'s implementation - // changes. Can we make it more robust? - - let old_adapter_count = 1; - let new_adapter_count = u32::try_from(wasi_imports.len()).unwrap(); - assert!(new_adapter_count >= old_adapter_count); - - Some((component, move |index: u32| { - if index == 0 { - // `main` module - 0 - } else if index <= new_adapter_count { - // adapter module - old_adapter_count - } else { - // one of the other kinds of module - index + old_adapter_count - new_adapter_count - } - })) + stubwasi::link_stub_modules(libraries) } else { None }; @@ -672,72 +619,3 @@ fn add_wasi_and_stubs( Ok(()) } - -fn add_wasi_imports<'a>( - module: &'a [u8], - imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>, -) -> Result<()> { - let mut types = Vec::new(); - for payload in Parser::new(0).parse_all(module) { - match payload? { - Payload::TypeSection(reader) => { - types = reader - .into_iter_err_on_gc_types() - .collect::, _>>()?; - } - - Payload::ImportSection(reader) => { - for import in reader { - let import = import?; - - if import.module == "wasi_snapshot_preview1" - || import.module.starts_with("wasi:") - { - if let TypeRef::Func(ty) = import.ty { - imports - .entry(import.module) - .or_default() - .insert(import.name, types[usize::try_from(ty).unwrap()].clone()); - } else { - bail!("encountered non-function import from WASI namespace") - } - } - } - break; - } - - _ => {} - } - } - - Ok(()) -} - -fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec { - let mut types = TypeSection::new(); - let mut functions = FunctionSection::new(); - let mut exports = ExportSection::new(); - let mut code = CodeSection::new(); - - for (index, (name, ty)) in stubs.iter().enumerate() { - let index = u32::try_from(index).unwrap(); - types.ty().function( - ty.params().iter().map(|&v| IntoValType(v).into()), - ty.results().iter().map(|&v| IntoValType(v).into()), - ); - functions.function(index); - exports.export(name, ExportKind::Func, index); - let mut function = Function::new([]); - function.instruction(&Instruction::Unreachable); - function.instruction(&Instruction::End); - code.function(&function); - } - - let mut module = Module::new(); - module.section(&types); - module.section(&functions); - module.section(&exports); - module.section(&code); - - module.finish() -} diff --git a/src/stubwasi.rs b/src/stubwasi.rs new file mode 100644 index 0000000..a83a4f0 --- /dev/null +++ b/src/stubwasi.rs @@ -0,0 +1,130 @@ +use std::collections::HashMap; + +use anyhow::bail; +use wasm_convert::IntoValType; +use wasm_encoder::{ + CodeSection, ExportKind, ExportSection, Function, FunctionSection, Instruction as Ins, Module, + TypeSection, +}; +use wasmparser::{FuncType, Parser, Payload, TypeRef}; + +use crate::Library; + +pub fn link_stub_modules(libraries: Vec) -> Option<(Vec, impl Fn(u32) -> u32)> { + let mut wasi_imports = HashMap::new(); + let mut linker = wit_component::Linker::default().validate(true); + + for Library { + name, + module, + dl_openable, + } in &libraries + { + add_wasi_imports(module, &mut wasi_imports).unwrap(); + linker = linker.library(name, module, *dl_openable).unwrap(); + } + + for (module, imports) in &wasi_imports { + linker = linker + .adapter(module, &make_stub_adapter(module, imports)) + .unwrap(); + } + + let component = linker.encode().unwrap(); + + // As of this writing, `wit_component::Linker` generates a component such that the first module is the + // `main` one, followed by any adapters, followed by any libraries, followed by the `init` module, which is + // finally followed by any shim modules. Given that the stubbed component may contain more adapters than + // the non-stubbed version, we need to tell `component-init` how to translate module indexes from the + // former to the latter. + // + // TODO: this is pretty fragile in that it could silently break if `wit_component::Linker`'s implementation + // changes. Can we make it more robust? + + let old_adapter_count = 1; + let new_adapter_count = u32::try_from(wasi_imports.len()).unwrap(); + assert!(new_adapter_count >= old_adapter_count); + + Some((component, move |index: u32| { + if index == 0 { + // `main` module + 0 + } else if index <= new_adapter_count { + // adapter module + old_adapter_count + } else { + // one of the other kinds of module + index + old_adapter_count - new_adapter_count + } + })) +} + +fn add_wasi_imports<'a>( + module: &'a [u8], + imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>, +) -> Result<(), anyhow::Error> { + let mut types = Vec::new(); + for payload in Parser::new(0).parse_all(module) { + match payload? { + Payload::TypeSection(reader) => { + types = reader + .into_iter_err_on_gc_types() + .collect::, _>>() + .unwrap(); + } + + Payload::ImportSection(reader) => { + for import in reader { + let import = import?; + + if import.module == "wasi_snapshot_preview1" + || import.module.starts_with("wasi:") + { + if let TypeRef::Func(ty) = import.ty { + imports + .entry(import.module) + .or_default() + .insert(import.name, types[usize::try_from(ty).unwrap()].clone()); + } else { + bail!("encountered non-function import from WASI namespace") + } + } + } + break; + } + + _ => {} + } + } + + Ok(()) +} + +fn make_stub_adapter(_module: &str, stubs: &HashMap<&str, FuncType>) -> Vec { + let mut types = TypeSection::new(); + let mut functions = FunctionSection::new(); + let mut exports = ExportSection::new(); + let mut code = CodeSection::new(); + + for (index, (name, ty)) in stubs.iter().enumerate() { + let index = u32::try_from(index).unwrap(); + types.ty().function( + ty.params().iter().map(|&v| IntoValType(v).into()), + ty.results().iter().map(|&v| IntoValType(v).into()), + ); + functions.function(index); + exports.export(name, ExportKind::Func, index); + let mut function = Function::new([]); + function.instruction(&Ins::Unreachable); + function.instruction(&Ins::End); + code.function(&function); + } + + let mut module = Module::new(); + module.section(&types); + module.section(&functions); + module.section(&exports); + module.section(&code); + + module.finish() +} From 830a75bb3ebc5a9f5803768c19395aaaad1b329b Mon Sep 17 00:00:00 2001 From: Phil Cummins Date: Mon, 28 Oct 2024 12:24:04 +0100 Subject: [PATCH 2/3] stubwasi.link_stub_modules returns Result --- src/stubwasi.rs | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/src/stubwasi.rs b/src/stubwasi.rs index a83a4f0..902fcb1 100644 --- a/src/stubwasi.rs +++ b/src/stubwasi.rs @@ -1,6 +1,6 @@ use std::collections::HashMap; -use anyhow::bail; +use anyhow::{bail, Error}; use wasm_convert::IntoValType; use wasm_encoder::{ CodeSection, ExportKind, ExportSection, Function, FunctionSection, Instruction as Ins, Module, @@ -10,7 +10,9 @@ use wasmparser::{FuncType, Parser, Payload, TypeRef}; use crate::Library; -pub fn link_stub_modules(libraries: Vec) -> Option<(Vec, impl Fn(u32) -> u32)> { +pub fn link_stub_modules( + libraries: Vec, +) -> Result, impl Fn(u32) -> u32)>, Error> { let mut wasi_imports = HashMap::new(); let mut linker = wit_component::Linker::default().validate(true); @@ -20,17 +22,15 @@ pub fn link_stub_modules(libraries: Vec) -> Option<(Vec, impl Fn(u3 dl_openable, } in &libraries { - add_wasi_imports(module, &mut wasi_imports).unwrap(); - linker = linker.library(name, module, *dl_openable).unwrap(); + add_wasi_imports(module, &mut wasi_imports)?; + linker = linker.library(name, module, *dl_openable)?; } for (module, imports) in &wasi_imports { - linker = linker - .adapter(module, &make_stub_adapter(module, imports)) - .unwrap(); + linker = linker.adapter(module, &make_stub_adapter(module, imports))?; } - let component = linker.encode().unwrap(); + let component = linker.encode()?; // As of this writing, `wit_component::Linker` generates a component such that the first module is the // `main` one, followed by any adapters, followed by any libraries, followed by the `init` module, which is @@ -42,10 +42,10 @@ pub fn link_stub_modules(libraries: Vec) -> Option<(Vec, impl Fn(u3 // changes. Can we make it more robust? let old_adapter_count = 1; - let new_adapter_count = u32::try_from(wasi_imports.len()).unwrap(); + let new_adapter_count = u32::try_from(wasi_imports.len())?; assert!(new_adapter_count >= old_adapter_count); - Some((component, move |index: u32| { + Ok(Some((component, move |index: u32| { if index == 0 { // `main` module 0 @@ -56,21 +56,20 @@ pub fn link_stub_modules(libraries: Vec) -> Option<(Vec, impl Fn(u3 // one of the other kinds of module index + old_adapter_count - new_adapter_count } - })) + }))) } fn add_wasi_imports<'a>( module: &'a [u8], imports: &mut HashMap<&'a str, HashMap<&'a str, FuncType>>, -) -> Result<(), anyhow::Error> { +) -> Result<(), Error> { let mut types = Vec::new(); for payload in Parser::new(0).parse_all(module) { match payload? { Payload::TypeSection(reader) => { types = reader .into_iter_err_on_gc_types() - .collect::, _>>() - .unwrap(); + .collect::, _>>()?; } Payload::ImportSection(reader) => { From 317850782e47679c02b8ae434e55522c1dc046c1 Mon Sep 17 00:00:00 2001 From: Phil Cummins Date: Mon, 28 Oct 2024 12:31:14 +0100 Subject: [PATCH 3/3] fix lib.rs --- src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/lib.rs b/src/lib.rs index e18f353..24bd389 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -297,7 +297,7 @@ pub async fn componentize( let component = linker.encode()?; let stubbed_component = if stub_wasi { - stubwasi::link_stub_modules(libraries) + stubwasi::link_stub_modules(libraries)? } else { None };