diff --git a/config.go b/config.go index 7cb1c4b6..6aadcfc7 100644 --- a/config.go +++ b/config.go @@ -7,6 +7,7 @@ import ( "io" "io/fs" "math" + "strings" "github.com/tetratelabs/wazero/internal/wasm" "github.com/tetratelabs/wazero/internal/wasm/interpreter" @@ -143,6 +144,11 @@ type ModuleConfig struct { preopens map[uint32]*wasm.FileEntry // preopenPaths allow overwriting of existing paths. preopenPaths map[string]uint32 + // replacedImports holds the latest state of WithImport + // Note: Key is NUL delimited as import module and name can both include any UTF-8 characters. + replacedImports map[string][2]string + // replacedImportModules holds the latest state of WithImportModule + replacedImportModules map[string]string } func NewModuleConfig() *ModuleConfig { @@ -170,6 +176,60 @@ func (c *ModuleConfig) WithName(name string) *ModuleConfig { return c } +// WithImport replaces a specific import module and name with a new one. This allows you to break up a monolithic +// module imports, such as "env". This can also help reduce cyclic dependencies. +// +// For example, if a module was compiled with one module owning all imports: +// (import "js" "tbl" (table $tbl 4 funcref)) +// (import "js" "increment" (func $increment (result i32))) +// (import "js" "decrement" (func $decrement (result i32))) +// (import "js" "wasm_increment" (func $wasm_increment (result i32))) +// (import "js" "wasm_decrement" (func $wasm_decrement (result i32))) +// +// Use this function to import "increment" and "decrement" from the module "go" and other imports from "wasm": +// config.WithImportModule("js", "wasm") +// config.WithImport("wasm", "increment", "go", "increment") +// config.WithImport("wasm", "decrement", "go", "decrement") +// +// Upon instantiation, imports resolve as if they were compiled like so: +// (import "wasm" "tbl" (table $tbl 4 funcref)) +// (import "go" "increment" (func $increment (result i32))) +// (import "go" "decrement" (func $decrement (result i32))) +// (import "wasm" "wasm_increment" (func $wasm_increment (result i32))) +// (import "wasm" "wasm_decrement" (func $wasm_decrement (result i32))) +// +// Note: Any WithImport instructions happen in order, after any WithImportModule instructions. +func (c *ModuleConfig) WithImport(oldModule, oldName, newModule, newName string) *ModuleConfig { + if c.replacedImports == nil { + c.replacedImports = map[string][2]string{} + } + var builder strings.Builder + builder.WriteString(oldModule) + builder.WriteByte(0) // delimit with NUL as module and name can be any UTF-8 characters. + builder.WriteString(oldName) + c.replacedImports[builder.String()] = [2]string{newModule, newName} + return c +} + +// WithImportModule replaces every import with oldModule with newModule. This is helpful for modules who have +// transitioned to a stable status since the underlying wasm was compiled. +// +// For example, if a module was compiled like below, with an old module for WASI: +// (import "wasi_unstable" "args_get" (func (param i32, i32) (result i32))) +// +// Use this function to update it to the current version: +// config.WithImportModule("wasi_unstable", wasi.ModuleSnapshotPreview1) +// +// See WithImport for a comprehensive example. +// Note: Any WithImportModule instructions happen in order, before any WithImport instructions. +func (c *ModuleConfig) WithImportModule(oldModule, newModule string) *ModuleConfig { + if c.replacedImportModules == nil { + c.replacedImportModules = map[string]string{} + } + c.replacedImportModules[oldModule] = newModule + return c +} + // WithStartFunctions configures the functions to call after the module is instantiated. Defaults to "_start". // // Note: If any function doesn't exist, it is skipped. However, all functions that do exist are called in order. @@ -352,3 +412,53 @@ func (c *ModuleConfig) toSysContext() (sys *wasm.SysContext, err error) { return wasm.NewSysContext(math.MaxUint32, c.args, environ, c.stdin, c.stdout, c.stderr, preopens) } + +func (c *ModuleConfig) replaceImports(module *wasm.Module) *wasm.Module { + if (c.replacedImportModules == nil && c.replacedImports == nil) || module.ImportSection == nil { + return module + } + + changed := false + + ret := *module // shallow copy + replacedImports := make([]*wasm.Import, len(module.ImportSection)) + copy(replacedImports, module.ImportSection) + + // First, replace any import.Module + for oldModule, newModule := range c.replacedImportModules { + for i, imp := range replacedImports { + if imp.Module == oldModule { + changed = true + cp := *imp // shallow copy + cp.Module = newModule + replacedImports[i] = &cp + } else { + replacedImports[i] = imp + } + } + } + + // Now, replace any import.Module+import.Name + for oldImport, newImport := range c.replacedImports { + for i, imp := range replacedImports { + nulIdx := strings.IndexByte(oldImport, 0) + oldModule := oldImport[0:nulIdx] + oldName := oldImport[nulIdx+1:] + if imp.Module == oldModule && imp.Name == oldName { + changed = true + cp := *imp // shallow copy + cp.Module = newImport[0] + cp.Name = newImport[1] + replacedImports[i] = &cp + } else { + replacedImports[i] = imp + } + } + } + + if !changed { + return module + } + ret.ImportSection = replacedImports + return &ret +} diff --git a/config_test.go b/config_test.go index 6765909f..0aefe38b 100644 --- a/config_test.go +++ b/config_test.go @@ -124,6 +124,386 @@ func TestRuntimeConfig_FeatureToggle(t *testing.T) { } } +func TestModuleConfig(t *testing.T) { + tests := []struct { + name string + with func(*ModuleConfig) *ModuleConfig + expected *ModuleConfig + }{ + { + name: "WithName", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithName("wazero") + }, + expected: &ModuleConfig{ + name: "wazero", + }, + }, + { + name: "WithName - empty", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithName("") + }, + expected: &ModuleConfig{}, + }, + { + name: "WithImport", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "abort", "assemblyscript", "abort") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}}, + }, + }, + { + name: "WithImport - empty to non-empty - module", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("", "abort", "assemblyscript", "abort") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"\000abort": {"assemblyscript", "abort"}}, + }, + }, + { + name: "WithImport - non-empty to empty - module", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "abort", "", "abort") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"", "abort"}}, + }, + }, + { + name: "WithImport - empty to non-empty - name", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "", "assemblyscript", "abort") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000": {"assemblyscript", "abort"}}, + }, + }, + { + name: "WithImport - non-empty to empty - name", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "abort", "assemblyscript", "") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", ""}}, + }, + }, + { + name: "WithImport - override", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "abort", "assemblyscript", "abort"). + WithImport("env", "abort", "go", "exit") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"go", "exit"}}, + }, + }, + { + name: "WithImport - twice", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImport("env", "abort", "assemblyscript", "abort"). + WithImport("wasi_unstable", "proc_exit", "wasi_snapshot_preview1", "proc_exit") + }, + expected: &ModuleConfig{ + replacedImports: map[string][2]string{ + "env\000abort": {"assemblyscript", "abort"}, + "wasi_unstable\000proc_exit": {"wasi_snapshot_preview1", "proc_exit"}, + }, + }, + }, + { + name: "WithImportModule", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImportModule("env", "assemblyscript") + }, + expected: &ModuleConfig{ + replacedImportModules: map[string]string{"env": "assemblyscript"}, + }, + }, + { + name: "WithImportModule - empty to non-empty", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImportModule("", "assemblyscript") + }, + expected: &ModuleConfig{ + replacedImportModules: map[string]string{"": "assemblyscript"}, + }, + }, + { + name: "WithImportModule - non-empty to empty", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImportModule("env", "") + }, + expected: &ModuleConfig{ + replacedImportModules: map[string]string{"env": ""}, + }, + }, + { + name: "WithImportModule - override", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImportModule("env", "assemblyscript"). + WithImportModule("env", "go") + }, + expected: &ModuleConfig{ + replacedImportModules: map[string]string{"env": "go"}, + }, + }, + { + name: "WithImportModule - twice", + with: func(c *ModuleConfig) *ModuleConfig { + return c.WithImportModule("env", "go"). + WithImportModule("wasi_unstable", "wasi_snapshot_preview1") + }, + expected: &ModuleConfig{ + replacedImportModules: map[string]string{ + "env": "go", + "wasi_unstable": "wasi_snapshot_preview1", + }, + }, + }, + } + for _, tt := range tests { + tc := tt + + t.Run(tc.name, func(t *testing.T) { + input := &ModuleConfig{} + rc := tc.with(input) + require.Equal(t, tc.expected, rc) + }) + } +} + +func TestModuleConfig_replaceImports(t *testing.T) { + tests := []struct { + name string + config *ModuleConfig + input *wasm.Module + expected *wasm.Module + expectSame bool + }{ + { + name: "no config, no imports", + config: &ModuleConfig{}, + input: &wasm.Module{}, + expected: &wasm.Module{}, + expectSame: true, + }, + { + name: "no config", + config: &ModuleConfig{}, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasi_snapshot_preview1", Name: "args_sizes_get", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasi_snapshot_preview1", Name: "fd_write", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + expectSame: true, + }, + { + name: "replacedImportModules", + config: &ModuleConfig{ + replacedImportModules: map[string]string{"wasi_unstable": "wasi_snapshot_preview1"}, + }, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasi_unstable", Name: "args_sizes_get", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasi_unstable", Name: "fd_write", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + expected: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasi_snapshot_preview1", Name: "args_sizes_get", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasi_snapshot_preview1", Name: "fd_write", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + }, + { + name: "replacedImportModules doesn't match", + config: &ModuleConfig{ + replacedImportModules: map[string]string{"env": ""}, + }, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasi_snapshot_preview1", Name: "args_sizes_get", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasi_snapshot_preview1", Name: "fd_write", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + expectSame: true, + }, + { + name: "replacedImports", + config: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}}, + }, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "env", Name: "abort", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "env", Name: "seed", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + expected: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "assemblyscript", Name: "abort", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "env", Name: "seed", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + }, + { + name: "replacedImports don't match", + config: &ModuleConfig{ + replacedImports: map[string][2]string{"env\000abort": {"assemblyscript", "abort"}}, + }, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasi_snapshot_preview1", Name: "args_sizes_get", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasi_snapshot_preview1", Name: "fd_write", + Type: wasm.ExternTypeFunc, + DescFunc: 2, + }, + }, + }, + expectSame: true, + }, + { + name: "replacedImportModules and replacedImports", + config: &ModuleConfig{ + replacedImportModules: map[string]string{"js": "wasm"}, + replacedImports: map[string][2]string{ + "wasm\000increment": {"go", "increment"}, + "wasm\000decrement": {"go", "decrement"}, + }, + }, + input: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "js", Name: "tbl", + Type: wasm.ExternTypeTable, + DescTable: &wasm.Table{Min: 4}, + }, + { + Module: "js", Name: "increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "js", Name: "decrement", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "js", Name: "wasm_increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "js", Name: "wasm_increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + }, + }, + expected: &wasm.Module{ + ImportSection: []*wasm.Import{ + { + Module: "wasm", Name: "tbl", + Type: wasm.ExternTypeTable, + DescTable: &wasm.Table{Min: 4}, + }, + { + Module: "go", Name: "increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "go", Name: "decrement", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasm", Name: "wasm_increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + { + Module: "wasm", Name: "wasm_increment", + Type: wasm.ExternTypeFunc, + DescFunc: 0, + }, + }, + }, + }, + } + for _, tt := range tests { + tc := tt + + t.Run(tc.name, func(t *testing.T) { + actual := tc.config.replaceImports(tc.input) + if tc.expectSame { + require.Same(t, tc.input, actual) + } else { + require.NotSame(t, tc.input, actual) + require.Equal(t, tc.expected, actual) + } + }) + } +} + func TestModuleConfig_toSysContext(t *testing.T) { testFS := fstest.MapFS{} testFS2 := fstest.MapFS{} diff --git a/examples/file_system_test.go b/examples/file_system_test.go index f1c585e7..29c7cd49 100644 --- a/examples/file_system_test.go +++ b/examples/file_system_test.go @@ -42,19 +42,15 @@ func Test_Cat(t *testing.T) { // Combine the above into our baseline config, overriding defaults (which discard stdout and have no file system). config := wazero.NewModuleConfig().WithStdout(stdoutBuf).WithFS(rooted) - // Compile the `cat` module. - code, err := r.CompileModule(catWasm) - require.NoError(t, err) - // Instantiate WASI, which implements system I/O such as console output. wm, err := wasi.InstantiateSnapshotPreview1(r) require.NoError(t, err) defer wm.Close() - // InstantiateModuleWithConfig by default runs the "_start" function which is what TinyGo compiles "main" to. + // InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to. // * Set the program name (arg[0]) to "cat" and add args to write "cat.go" to stdout twice. // * We use both "/cat.go" and "./cat.go" because WithFS by default maps the workdir "." to "/". - cat, err := r.InstantiateModuleWithConfig(code, config.WithArgs("cat", "/cat.go", "./cat.go")) + cat, err := r.InstantiateModuleFromCodeWithConfig(catWasm, config.WithArgs("cat", "/cat.go", "./cat.go")) require.NoError(t, err) defer cat.Close() diff --git a/examples/host_func_test.go b/examples/host_func_test.go index de99debb..ae432789 100644 --- a/examples/host_func_test.go +++ b/examples/host_func_test.go @@ -60,10 +60,6 @@ func Test_hostFunc(t *testing.T) { _, err := r.NewModuleBuilder("env").ExportFunction("get_random_bytes", getRandomBytes).Instantiate() require.NoError(t, err) - // Compile the `hostFunc` module. - code, err := r.CompileModule(hostFuncWasm) - require.NoError(t, err) - // Configure stdout (console) to write to a buffer. stdout := bytes.NewBuffer(nil) config := wazero.NewModuleConfig().WithStdout(stdout) @@ -73,8 +69,8 @@ func Test_hostFunc(t *testing.T) { require.NoError(t, err) defer wm.Close() - // InstantiateModuleWithConfig runs the "_start" function which is what TinyGo compiles "main" to. - module, err := r.InstantiateModuleWithConfig(code, config) + // InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to. + module, err := r.InstantiateModuleFromCodeWithConfig(hostFuncWasm, config) require.NoError(t, err) defer module.Close() diff --git a/examples/replace_test.go b/examples/replace_test.go new file mode 100644 index 00000000..ad780072 --- /dev/null +++ b/examples/replace_test.go @@ -0,0 +1,42 @@ +package examples + +import ( + "testing" + + "github.com/stretchr/testify/require" + + "github.com/tetratelabs/wazero" + "github.com/tetratelabs/wazero/api" +) + +// Test_Replace shows how you can replace a module import when it doesn't match instantiated modules. +func Test_Replace(t *testing.T) { + r := wazero.NewRuntime() + + // Instantiate a function that closes the module under "assemblyscript.abort". + host, err := r.NewModuleBuilder("assemblyscript"). + ExportFunction("abort", func(m api.Module, messageOffset, fileNameOffset, line, col uint32) { + _ = m.CloseWithExitCode(255) + }).Instantiate() + require.NoError(t, err) + defer host.Close() + + // Compile code that needs the function "env.abort". + code, err := r.CompileModule([]byte(`(module + (import "env" "abort" (func $~lib/builtins/abort (param i32 i32 i32 i32))) + + (export "abort" (func 0)) ;; exports the import for testing +)`)) + require.NoError(t, err) + + // Instantiate the module, replacing the import "env.abort" with "assemblyscript.abort". + mod, err := r.InstantiateModuleWithConfig(code, wazero.NewModuleConfig(). + WithName(t.Name()). + WithImport("env", "abort", "assemblyscript", "abort")) + require.NoError(t, err) + defer mod.Close() + + // Since the above worked, the exported function closes the module. + _, err = mod.ExportedFunction("abort").Call(nil, 0, 0, 0, 0) + require.EqualError(t, err, `module "Test_Replace" closed with exit_code(255)`) +} diff --git a/examples/stdio_test.go b/examples/stdio_test.go index 360eb3e6..83b5f55e 100644 --- a/examples/stdio_test.go +++ b/examples/stdio_test.go @@ -19,10 +19,6 @@ var stdioWasm []byte func Test_stdio(t *testing.T) { r := wazero.NewRuntime() - // Compile the `stdioWasm` module. - code, err := r.CompileModule(stdioWasm) - require.NoError(t, err) - // Configure standard I/O (ex stdout) to write to buffers instead of no-op. stdinBuf := bytes.NewBuffer([]byte("WASI\n")) stdoutBuf := bytes.NewBuffer(nil) @@ -34,8 +30,8 @@ func Test_stdio(t *testing.T) { require.NoError(t, err) defer wm.Close() - // InstantiateModuleWithConfig runs the "_start" function which is what TinyGo compiles "main" to. - module, err := r.InstantiateModuleWithConfig(code, config) + // InstantiateModuleFromCodeWithConfig runs the "_start" function which is what TinyGo compiles "main" to. + module, err := r.InstantiateModuleFromCodeWithConfig(stdioWasm, config) require.NoError(t, err) defer module.Close() diff --git a/wasm.go b/wasm.go index 422e1dbf..220a3bc6 100644 --- a/wasm.go +++ b/wasm.go @@ -57,6 +57,16 @@ type Runtime interface { // source multiple times, use CompileModule as InstantiateModule avoids redundant decoding and/or compilation. InstantiateModuleFromCode(source []byte) (api.Module, error) + // InstantiateModuleFromCodeWithConfig is a convenience function that chains CompileModule to + // InstantiateModuleWithConfig. + // + // Ex. To only change the module name: + // wasm, _ := wazero.NewRuntime().InstantiateModuleFromCodeWithConfig(source, wazero.NewModuleConfig(). + // WithName("wasm") + // ) + // defer wasm.Close() + InstantiateModuleFromCodeWithConfig(source []byte, config *ModuleConfig) (api.Module, error) + // InstantiateModule instantiates the module namespace or errs if the configuration was invalid. // // Ex. @@ -164,6 +174,15 @@ func (r *runtime) InstantiateModuleFromCode(source []byte) (api.Module, error) { } } +// InstantiateModuleFromCodeWithConfig implements Runtime.InstantiateModuleFromCodeWithConfig +func (r *runtime) InstantiateModuleFromCodeWithConfig(source []byte, config *ModuleConfig) (api.Module, error) { + if code, err := r.CompileModule(source); err != nil { + return nil, err + } else { + return r.InstantiateModuleWithConfig(code, config) + } +} + // InstantiateModule implements Runtime.InstantiateModule func (r *runtime) InstantiateModule(code *CompiledCode) (mod api.Module, err error) { return r.InstantiateModuleWithConfig(code, NewModuleConfig()) @@ -181,7 +200,9 @@ func (r *runtime) InstantiateModuleWithConfig(code *CompiledCode, config *Module name = code.module.NameSection.ModuleName } - mod, err = r.store.Instantiate(r.ctx, code.module, name, sys) + module := config.replaceImports(code.module) + + mod, err = r.store.Instantiate(r.ctx, module, name, sys) if err != nil { return } diff --git a/wasm_test.go b/wasm_test.go index d69287e0..d1eb994d 100644 --- a/wasm_test.go +++ b/wasm_test.go @@ -288,10 +288,7 @@ func TestFunction_Context(t *testing.T) { defer closer() // nolint // Instantiate the module and get the export of the above hostFn - code, err := r.CompileModule(source) - require.NoError(t, err) - - module, err := r.InstantiateModuleWithConfig(code, NewModuleConfig().WithName(t.Name())) + module, err := r.InstantiateModuleFromCodeWithConfig(source, NewModuleConfig().WithName(t.Name())) require.NoError(t, err) defer module.Close()