Prevents direct use of host functions via ExportedFunction (#2259)

Signed-off-by: Takeshi Yoneda <t.y.mathetake@gmail.com>

api/wasm.go

@@ -151,9 +151,13 @@ type Module interface {
 
 	// ExportedFunction returns a function exported from this module or nil if it wasn't.
 	//
-	// Note: The default wazero.ModuleConfig attempts to invoke `_start`, which
-	// in rare cases can close the module. When in doubt, check IsClosed prior
-	// to invoking a function export after instantiation.
+	// # Notes
+	//   - The default wazero.ModuleConfig attempts to invoke `_start`, which
+	//     in rare cases can close the module. When in doubt, check IsClosed prior
+	//     to invoking a function export after instantiation.
+	//   - The semantics of host functions assumes the existence of an "importing module" because, for example, the host function needs access to
+	//     the memory of the importing module. Therefore, direct use of ExportedFunction is forbidden for host modules.
+	//     Practically speaking, it is usually meaningless to directly call a host function from Go code as it is already somewhere in Go code.
 	ExportedFunction(name string) Function
 
 	// ExportedFunctionDefinitions returns all the exported function

builder.go

@@ -179,6 +179,9 @@ type HostFunctionBuilder interface {
 //     are deferred until Compile.
 //   - Functions are indexed in order of calls to NewFunctionBuilder as
 //     insertion ordering is needed by ABI such as Emscripten (invoke_*).
+//   - The semantics of host functions assumes the existence of an "importing module" because, for example, the host function needs access to
+//     the memory of the importing module. Therefore, direct use of ExportedFunction is forbidden for host modules.
+//     Practically speaking, it is usually meaningless to directly call a host function from Go code as it is already somewhere in Go code.
 type HostModuleBuilder interface {
 	// Note: until golang/go#5860, we can't use example tests to embed code in interface godocs.
 
@@ -341,12 +344,24 @@ func (b *hostModuleBuilder) Compile(ctx context.Context) (CompiledModule, error)
 	return c, nil
 }
 
+// hostModuleInstance is a wrapper around api.Module that prevents calling ExportedFunction.
+type hostModuleInstance struct{ api.Module }
+
+// ExportedFunction implements api.Module ExportedFunction.
+func (h hostModuleInstance) ExportedFunction(name string) api.Function {
+	panic("calling ExportedFunction is forbidden on host modules. See the note on ExportedFunction interface")
+}
+
 // Instantiate implements HostModuleBuilder.Instantiate
 func (b *hostModuleBuilder) Instantiate(ctx context.Context) (api.Module, error) {
 	if compiled, err := b.Compile(ctx); err != nil {
 		return nil, err
 	} else {
 		compiled.(*compiledModule).closeWithModule = true
-		return b.r.InstantiateModule(ctx, compiled, NewModuleConfig())
+		m, err := b.r.InstantiateModule(ctx, compiled, NewModuleConfig())
+		if err != nil {
+			return nil, err
+		}
+		return hostModuleInstance{m}, nil
 	}
 }

builder_test.go

@@ -368,6 +368,10 @@ func TestNewHostModuleBuilder_Instantiate(t *testing.T) {
 	m, err := r.NewHostModuleBuilder("env").Instantiate(testCtx)
 	require.NoError(t, err)
 
+	// Calling ExportedFunction should fail.
+	err = require.CapturePanic(func() { m.ExportedFunction("any") })
+	require.EqualError(t, err, `calling ExportedFunction is forbidden on host modules. See the note on ExportedFunction interface`)
+
 	// If this was instantiated, it would be added to the store under the same name
 	require.Equal(t, r.Module("env"), m)
 

internal/wasm/store_module_list.go

@@ -3,8 +3,6 @@ package wasm
 import (
 	"errors"
 	"fmt"
-
-	"github.com/tetratelabs/wazero/api"
 )
 
 // deleteModule makes the moduleName available for instantiation again.
@@ -88,7 +86,7 @@ func (s *Store) registerModule(m *ModuleInstance) error {
 }
 
 // Module implements wazero.Runtime Module
-func (s *Store) Module(moduleName string) api.Module {
+func (s *Store) Module(moduleName string) *ModuleInstance {
 	m, err := s.module(moduleName)
 	if err != nil {
 		return nil

runtime.go

@@ -197,7 +197,13 @@ func (r *runtime) Module(moduleName string) api.Module {
 	if len(moduleName) == 0 {
 		return nil
 	}
-	return r.store.Module(moduleName)
+	m := r.store.Module(moduleName)
+	if m == nil {
+		return nil
+	} else if m.Source.IsHostModule {
+		return hostModuleInstance{m}
+	}
+	return m
 }
 
 // CompileModule implements Runtime.CompileModule