Repository: tetratelabs/wazero (go)

Issue:
Running into `wasm error: indirect call type mismatch` when using the compiler backend, but not the interpreter
**Describe the bug**
I am running a wasm module that uses `call_indirect` - when using the interpreter it runs successfully, but when using the compiler it fails with `wasm error: indirect call type mismatch` - it also runs successfully on wasmtime, so I think this could be a compiler bug.

**To Reproduce**
I've uploaded the wasm here: [wasm.zip](https://github.com/user-attachments/files/16985989/wasm.zip)

Then using the extism CLI: https://github.com/extism/cli/ (can be installed by cloning the repo and running `go install ./extism` from the root of the repo)
```
EXTISM_ENABLE_WASI_OUTPUT=1 extism call ./plugin.wasm req --allow-host '*' --input "https://extism.org" --wasi --link import_module=import_module.wasm --log-level info
```

Note: the current CLI uses wazero 1.7, but I've also tested using the latest code on `main` with the same result.

**Expected behavior**
The call should return successfully instead of throwing the type mismatch error

**Environment (please complete the relevant information):**
- Go version: 1.22.6
- wazero Version:tested on 1.7 and the latest main (`58488880a334e8bda5be0d715a24d8ddeb34c725`)
- Host architecture: amd64
- Runtime mode: compiler

**Additional context**
If I comment out lines 3558-3581 in `internal/engine/wazevo/frontend/lower.go` it works as expected.



Pull request:
compiler: ResolveImportedFunction should infer the typeID from the importing module
Fixes #2313.

This changes the signature for `ResolveImportedFunction()` so that it takes an additional index into the typeID section of the "current" module, i.e. the subject of the import.

The current behavior is to search for the import recursively into the chain of dependencies; however, this resolves to a typeID that is specific to _that_ module. For instance:

- consider the case of a module `m` declaring an import `m2inst.sym`.
- let be `m2inst.sym` an import with type `t` with ID `tid` 
- let be `m2inst` be linked to an instance of module `m2` with type `t2` and ID `t2id`.

Now: we want `t` to be compatible with `t2`; this we verify at:

https://github.com/tetratelabs/wazero/blob/6bb9899ea4c745a4ae23179a34fac0540c2ce551/internal/wasm/store.go#L441-L447

The current behavior is for `ResolveImportedFunction()` to resolve recursively a reference to a function so that we can then jump precisely to its address at runtime; however, we should _not_ carry on its type ID, and instead use the type ID of the import.

In other words, even if the types are compatible, we should be still referring to type `t` with ID `tid` from `m`, and _not_ type ID `t2id`, because, of course, the index `t2id` might be different from the local index `tid`.

The solution is to add another `descFunc Index` (from the name of the field `Import.DescFunc`) to the signature of  `ResolveImportedFunction()`, so that the typeID is correctly resolved from `moduleEngine.module.TypeIDs[descFunc]` (i.e. the typeID defined in the module that is requiring the imported function).

The tests are adjusted accordingly; notably, this also means that, in order for `m.ResolveImportedFunction()` to work properly, `m` must be correctly initialized with an `m.module` being the module that is performing the import, so the initialization of `moduleEngine` is pushed down, after the declaration of the "`imported`" and "`importing`" modules.

Let me know if you think it is worth adding more tests!

@zshipko please double-check that this is fixing your issue (I have slightly changed my original approach)
@ncruces would you like to double-check this against your sqlite lib?
@anuraaga if you spare any time, would you take a look too? :P






Patch:
diff --git a/internal/engine/interpreter/interpreter.go b/internal/engine/interpreter/interpreter.go
index ee0b453c..5b5e6e9d 100644
--- a/internal/engine/interpreter/interpreter.go
+++ b/internal/engine/interpreter/interpreter.go
@@ -487,7 +487,7 @@ func (e *engine) setLabelAddress(op *uint64, label label, labelAddressResolution
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (e *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (e *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	imported := importedModuleEngine.(*moduleEngine)
 	e.functions[index] = imported.functions[indexInImportedModule]
 }
diff --git a/internal/engine/wazevo/module_engine.go b/internal/engine/wazevo/module_engine.go
index efa1b9bb..8811feed 100644
--- a/internal/engine/wazevo/module_engine.go
+++ b/internal/engine/wazevo/module_engine.go
@@ -237,7 +237,7 @@ func (m *moduleEngine) putLocalMemory() {
 }
 
 // ResolveImportedFunction implements wasm.ModuleEngine.
-func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
+func (m *moduleEngine) ResolveImportedFunction(index, descFunc, indexInImportedModule wasm.Index, importedModuleEngine wasm.ModuleEngine) {
 	executableOffset, moduleCtxOffset, typeIDOffset := m.parent.offsets.ImportedFunctionOffset(index)
 	importedME := importedModuleEngine.(*moduleEngine)
 
@@ -245,12 +245,12 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 		indexInImportedModule -= wasm.Index(len(importedME.importedFunctions))
 	} else {
 		imported := &importedME.importedFunctions[indexInImportedModule]
-		m.ResolveImportedFunction(index, imported.indexInModule, imported.me)
+		m.ResolveImportedFunction(index, descFunc, imported.indexInModule, imported.me)
 		return // Recursively resolve the imported function.
 	}
 
 	offset := importedME.parent.functionOffsets[indexInImportedModule]
-	typeID := getTypeIDOf(indexInImportedModule, importedME.module)
+	typeID := m.module.TypeIDs[descFunc]
 	executable := &importedME.parent.executable[offset]
 	// Write functionInstance.
 	binary.LittleEndian.PutUint64(m.opaque[executableOffset:], uint64(uintptr(unsafe.Pointer(executable))))
@@ -261,28 +261,6 @@ func (m *moduleEngine) ResolveImportedFunction(index, indexInImportedModule wasm
 	m.importedFunctions[index] = importedFunction{me: importedME, indexInModule: indexInImportedModule}
 }
 
-func getTypeIDOf(funcIndex wasm.Index, m *wasm.ModuleInstance) wasm.FunctionTypeID {
-	source := m.Source
-
-	var typeIndex wasm.Index
-	if funcIndex >= source.ImportFunctionCount {
-		funcIndex -= source.ImportFunctionCount
-		typeIndex = source.FunctionSection[funcIndex]
-	} else {
-		var cnt wasm.Index
-		for i := range source.ImportSection {
-			if source.ImportSection[i].Type == wasm.ExternTypeFunc {
-				if cnt == funcIndex {
-					typeIndex = source.ImportSection[i].DescFunc
-					break
-				}
-				cnt++
-			}
-		}
-	}
-	return m.TypeIDs[typeIndex]
-}
-
 // ResolveImportedMemory implements wasm.ModuleEngine.
 func (m *moduleEngine) ResolveImportedMemory(importedModuleEngine wasm.ModuleEngine) {
 	importedME := importedModuleEngine.(*moduleEngine)
diff --git a/internal/wasm/engine.go b/internal/wasm/engine.go
index 61a342ef..8c387e9e 100644
--- a/internal/wasm/engine.go
+++ b/internal/wasm/engine.go
@@ -44,9 +44,10 @@ type ModuleEngine interface {
 
 	// ResolveImportedFunction is used to add imported functions needed to make this ModuleEngine fully functional.
 	// 	- `index` is the function Index of this imported function.
+	// 	- `descFunc` is the type Index in Module.TypeSection of this imported function. It corresponds to Import.DescFunc.
 	// 	- `indexInImportedModule` is the function Index of the imported function in the imported module.
 	//	- `importedModuleEngine` is the ModuleEngine for the imported ModuleInstance.
-	ResolveImportedFunction(index, indexInImportedModule Index, importedModuleEngine ModuleEngine)
+	ResolveImportedFunction(index, descFunc, indexInImportedModule Index, importedModuleEngine ModuleEngine)
 
 	// ResolveImportedMemory is called when this module imports a memory from another module.
 	ResolveImportedMemory(importedModuleEngine ModuleEngine)
diff --git a/internal/wasm/store.go b/internal/wasm/store.go
index dda6e5b6..c7909c67 100644
--- a/internal/wasm/store.go
+++ b/internal/wasm/store.go
@@ -446,7 +446,7 @@ func (m *ModuleInstance) resolveImports(ctx context.Context, module *Module) (er
 					return
 				}
 
-				m.Engine.ResolveImportedFunction(i.IndexPerType, imported.Index, importedModule.Engine)
+				m.Engine.ResolveImportedFunction(i.IndexPerType, i.DescFunc, imported.Index, importedModule.Engine)
 			case ExternTypeTable:
 				expected := i.DescTable
 				importedTable := importedModule.Tables[imported.Index]


Test patch:
diff --git a/internal/engine/wazevo/module_engine_test.go b/internal/engine/wazevo/module_engine_test.go
index 45234e19..7e6bbf33 100644
--- a/internal/engine/wazevo/module_engine_test.go
+++ b/internal/engine/wazevo/module_engine_test.go
@@ -152,28 +152,19 @@ func TestModuleEngine_setupOpaque(t *testing.T) {
 }
 
 func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
-	const begin = 5000
-	m := &moduleEngine{
-		opaque:            make([]byte, 10000),
-		importedFunctions: make([]importedFunction, 4),
-		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
-			ImportedFunctionsBegin: begin,
-		}},
-	}
-
 	var op1, op2 byte = 0xaa, 0xbb
-	im1 := &moduleEngine{
+	importing := &moduleEngine{
 		opaquePtr: &op1,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
 			functionOffsets: []int{1, 5, 10},
 		},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 0, 111, 222, 333},
+			TypeIDs: []wasm.FunctionTypeID{0, 0, 0, 888, 111, 222, 333},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{4, 5, 6}},
 		},
 	}
-	im2 := &moduleEngine{
+	imported := &moduleEngine{
 		opaquePtr: &op2,
 		parent: &compiledModule{
 			executables:     &executables{executable: make([]byte, 1000)},
@@ -185,10 +176,20 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, im1)
-	m.ResolveImportedFunction(1, 0, im2)
-	m.ResolveImportedFunction(2, 2, im1)
-	m.ResolveImportedFunction(3, 1, im1)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 4, 0, importing)
+	m.ResolveImportedFunction(1, 3, 0, imported)
+	m.ResolveImportedFunction(2, 6, 2, importing)
+	m.ResolveImportedFunction(3, 5, 1, importing)
 
 	for i, tc := range []struct {
 		index      int
@@ -196,10 +197,10 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{index: 0, op: &op1, executable: &im1.parent.executable[1], expTypeID: 111},
-		{index: 1, op: &op2, executable: &im2.parent.executable[50], expTypeID: 999},
-		{index: 2, op: &op1, executable: &im1.parent.executable[10], expTypeID: 333},
-		{index: 3, op: &op1, executable: &im1.parent.executable[5], expTypeID: 222},
+		{index: 0, op: &op1, executable: &importing.parent.executable[1], expTypeID: 111},
+		{index: 1, op: &op2, executable: &imported.parent.executable[50], expTypeID: 888},
+		{index: 2, op: &op1, executable: &importing.parent.executable[10], expTypeID: 333},
+		{index: 3, op: &op1, executable: &importing.parent.executable[5], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
 			buf := m.opaque[begin+wazevoapi.FunctionInstanceSize*tc.index:]
@@ -216,15 +217,6 @@ func TestModuleEngine_ResolveImportedFunction(t *testing.T) {
 }
 
 func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
-	const begin = 5000
-	m := &moduleEngine{
-		opaque:            make([]byte, 10000),
-		importedFunctions: make([]importedFunction, 4),
-		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
-			ImportedFunctionsBegin: begin,
-		}},
-	}
-
 	var importingOp, importedOp byte = 0xaa, 0xbb
 	imported := &moduleEngine{
 		opaquePtr: &importedOp,
@@ -245,13 +237,23 @@ func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
 		},
 		importedFunctions: []importedFunction{{me: imported, indexInModule: 0}},
 		module: &wasm.ModuleInstance{
-			TypeIDs: []wasm.FunctionTypeID{0, 222, 0},
+			TypeIDs: []wasm.FunctionTypeID{999, 222, 0},
 			Source:  &wasm.Module{FunctionSection: []wasm.Index{1}},
 		},
 	}
 
-	m.ResolveImportedFunction(0, 0, importing)
-	m.ResolveImportedFunction(1, 1, importing)
+	const begin = 5000
+	m := &moduleEngine{
+		opaque:            make([]byte, 10000),
+		importedFunctions: make([]importedFunction, 4),
+		parent: &compiledModule{offsets: wazevoapi.ModuleContextOffsetData{
+			ImportedFunctionsBegin: begin,
+		}},
+		module: importing.module,
+	}
+
+	m.ResolveImportedFunction(0, 0, 0, importing)
+	m.ResolveImportedFunction(1, 1, 1, importing)
 
 	for i, tc := range []struct {
 		index      int
@@ -259,7 +261,7 @@ func TestModuleEngine_ResolveImportedFunction_recursive(t *testing.T) {
 		executable *byte
 		expTypeID  wasm.FunctionTypeID
 	}{
-		{index: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 111},
+		{index: 0, op: &importedOp, executable: &imported.parent.executable[10], expTypeID: 999},
 		{index: 1, op: &importingOp, executable: &importing.parent.executable[500], expTypeID: 222},
 	} {
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -312,26 +314,6 @@ func Test_functionInstance_offsets(t *testing.T) {
 	require.Equal(t, typeID, ptr+16)
 }
 
-func Test_getTypeIDOf(t *testing.T) {
-	m := &wasm.ModuleInstance{
-		TypeIDs: []wasm.FunctionTypeID{111, 222, 333, 444},
-		Source: &wasm.Module{
-			ImportFunctionCount: 1,
-			ImportSection: []wasm.Import{
-				{Type: wasm.ExternTypeMemory},
-				{Type: wasm.ExternTypeTable},
-				{Type: wasm.ExternTypeFunc, DescFunc: 3},
-			},
-			FunctionSection: []wasm.Index{2, 1, 0},
-		},
-	}
-
-	require.Equal(t, wasm.FunctionTypeID(444), getTypeIDOf(0, m))
-	require.Equal(t, wasm.FunctionTypeID(333), getTypeIDOf(1, m))
-	require.Equal(t, wasm.FunctionTypeID(222), getTypeIDOf(2, m))
-	require.Equal(t, wasm.FunctionTypeID(111), getTypeIDOf(3, m))
-}
-
 func Test_newAlignedOpaque(t *testing.T) {
 	for i := 0; i < 100; i++ {
 		s := 16 * (i + 10)
diff --git a/internal/wasm/store_test.go b/internal/wasm/store_test.go
index 08297591..a54ef182 100644
--- a/internal/wasm/store_test.go
+++ b/internal/wasm/store_test.go
@@ -341,7 +341,7 @@ func TestStore_Instantiate_Errors(t *testing.T) {
 				importedModuleName: {{Type: ExternTypeFunc, Module: importedModuleName, Name: "fn", DescFunc: 0}},
 				"non-exist":        {{Name: "fn", DescFunc: 0}},
 			},
-		}, importingModuleName, nil, nil)
+		}, importingModuleName, nil, []FunctionTypeID{0})
 		require.EqualError(t, err, "module[non-exist] not instantiated")
 	})
 
@@ -485,7 +485,7 @@ func (e *mockModuleEngine) FunctionInstanceReference(i Index) Reference {
 }
 
 // ResolveImportedFunction implements the same method as documented on wasm.ModuleEngine.
-func (e *mockModuleEngine) ResolveImportedFunction(index, importedIndex Index, _ ModuleEngine) {
+func (e *mockModuleEngine) ResolveImportedFunction(index, _, importedIndex Index, _ ModuleEngine) {
 	e.resolveImportsCalled[index] = importedIndex
 }
 
@@ -742,7 +742,7 @@ func Test_resolveImports(t *testing.T) {
 				},
 			}
 
-			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module}
+			m := &ModuleInstance{Engine: &mockModuleEngine{resolveImportsCalled: map[Index]Index{}}, s: s, Source: module, TypeIDs: []FunctionTypeID{0, 1}}
 			err := m.resolveImports(context.Background(), module)
 			require.NoError(t, err)