table: ensures references alive regardless of closures (#2139)

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

internal/integration_test/engine/adhoc_test.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -71,6 +72,8 @@ var tests = map[string]testCase{
 	"many params many results / main / listener":                       {f: testManyParamsResultsMainListener},
 	"many params many results / call_many_consts_and_pick_last_vector": {f: testManyParamsResultsCallManyConstsAndPickLastVector},
 	"many params many results / call_many_consts_and_pick_last_vector / listener": {f: testManyParamsResultsCallManyConstsAndPickLastVectorListener},
+	"close table importing module": {f: testCloseTableImportingModule},
+	"close table exporting module": {f: testCloseTableExportingModule},
 }
 
 func TestEngineCompiler(t *testing.T) {
@@ -2174,3 +2177,187 @@ wasm stack trace:
 	.$0()
 	... maybe followed by omitted frames`, err.Error())
 }
+
+// testCloseTableExportingModule tests the situation where the module instance that
+// is the initial owner of the table is closed and then try to call call_indirect on the table.
+//
+// This is in practice extremely edge case and shouldn't occur in real world, but in any way, seg fault should not occur.
+func testCloseTableExportingModule(t *testing.T, r wazero.Runtime) {
+	exportingBin := binaryencoding.EncodeModule(&wasm.Module{
+		ExportSection: []wasm.Export{
+			{Name: "t", Type: wasm.ExternTypeTable, Index: 0},
+		},
+		TableSection: []wasm.Table{{Type: wasm.RefTypeFuncref, Min: 10}},
+		NameSection:  &wasm.NameSection{ModuleName: "exporting"},
+		ElementSection: []wasm.ElementSegment{
+			{
+				OffsetExpr: wasm.ConstantExpression{
+					Opcode: wasm.OpcodeI32Const,
+					Data:   leb128.EncodeInt32(5),
+				}, TableIndex: 0, Type: wasm.RefTypeFuncref, Mode: wasm.ElementModeActive,
+				// Set the function 0, 1 at table offset 5.
+				Init: []wasm.Index{0, 1},
+			},
+		},
+		TypeSection:     []wasm.FunctionType{{Results: []wasm.ValueType{i32}}},
+		FunctionSection: []wasm.Index{0, 0},
+		CodeSection: []wasm.Code{
+			{Body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeI32Const, 2, wasm.OpcodeEnd}},
+		},
+	})
+
+	mainBin := binaryencoding.EncodeModule(&wasm.Module{
+		ImportSection: []wasm.Import{{
+			Type:      wasm.ExternTypeTable,
+			Module:    "exporting",
+			Name:      "t",
+			DescTable: wasm.Table{Type: wasm.RefTypeFuncref, Min: 10},
+		}},
+		TypeSection: []wasm.FunctionType{
+			{Results: []wasm.ValueType{i32}},                                // Type for functions in the table.
+			{Params: []wasm.ValueType{i32}, Results: []wasm.ValueType{i32}}, // Type for the main function.
+		},
+		ExportSection: []wasm.Export{
+			{Name: "", Type: wasm.ExternTypeFunc, Index: 0},
+		},
+		FunctionSection: []wasm.Index{1},
+		CodeSection: []wasm.Code{
+			{Body: []byte{
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeCallIndirect, 0, 0,
+				wasm.OpcodeEnd,
+			}},
+		},
+	})
+
+	ctx := context.Background()
+	exportingMod, err := r.Instantiate(ctx, exportingBin)
+	require.NoError(t, err)
+
+	mainMod, err := r.Instantiate(ctx, mainBin)
+	require.NoError(t, err)
+
+	main := mainMod.ExportedFunction("")
+	require.NotNil(t, main)
+
+	err = exportingMod.Close(ctx)
+	require.NoError(t, err)
+
+	// Trigger GC to make sure the module instance that is the initial owner of the table is collected.
+	runtime.GC()
+
+	// Call call_indirect multiple times, should be safe.
+	for i := 0; i < 10; i++ {
+		_, err = main.Call(ctx, 0)
+		// Null function call
+		require.Error(t, err)
+
+		res, err := main.Call(ctx, 5)
+		require.NoError(t, err)
+		require.Equal(t, uint64(1), res[0])
+
+		res, err = main.Call(ctx, 6)
+		require.NoError(t, err)
+		require.Equal(t, uint64(2), res[0])
+		time.Sleep(time.Millisecond * 10)
+	}
+}
+
+// testCloseTableImportingModule is similar to testCloseTableExportingModule, but the module
+// importing the table sets the function reference in the table, and then the module instance
+// that is the initial owner of the table will call call_indirect on the table.
+//
+// This is in practice extremely edge case and shouldn't occur in real world, but in any way, seg fault should not occur.
+func testCloseTableImportingModule(t *testing.T, r wazero.Runtime) {
+	exportingBin := binaryencoding.EncodeModule(&wasm.Module{
+		ExportSection: []wasm.Export{
+			{Name: "t", Type: wasm.ExternTypeTable, Index: 0},
+			{Name: "main", Type: wasm.ExternTypeFunc, Index: 0},
+		},
+		TableSection: []wasm.Table{{Type: wasm.RefTypeFuncref, Min: 10}},
+		NameSection:  &wasm.NameSection{ModuleName: "exporting"},
+		TypeSection: []wasm.FunctionType{
+			{Results: []wasm.ValueType{i32}},                                // Type for functions in the table.
+			{Params: []wasm.ValueType{i32}, Results: []wasm.ValueType{i32}}, // Type for the main function.
+		},
+		FunctionSection: []wasm.Index{1},
+		CodeSection: []wasm.Code{
+			{Body: []byte{
+				wasm.OpcodeLocalGet, 0,
+				wasm.OpcodeCallIndirect, 0, 0,
+				wasm.OpcodeEnd,
+			}},
+		},
+	})
+
+	importingBin := binaryencoding.EncodeModule(&wasm.Module{
+		NameSection: &wasm.NameSection{ModuleName: "importing"},
+		ImportSection: []wasm.Import{{
+			Type:      wasm.ExternTypeTable,
+			Module:    "exporting",
+			Name:      "t",
+			DescTable: wasm.Table{Type: wasm.RefTypeFuncref, Min: 10},
+		}},
+		TypeSection: []wasm.FunctionType{
+			{Results: []wasm.ValueType{i32}}, // Type for functions in the table.
+		},
+		ElementSection: []wasm.ElementSegment{
+			{
+				OffsetExpr: wasm.ConstantExpression{
+					Opcode: wasm.OpcodeI32Const,
+					Data:   leb128.EncodeInt32(5),
+				}, TableIndex: 0, Type: wasm.RefTypeFuncref, Mode: wasm.ElementModeActive,
+				// Set the function 0, 1 at table offset 5.
+				Init: []wasm.Index{0, 1},
+			},
+		},
+		FunctionSection: []wasm.Index{0, 0},
+		CodeSection: []wasm.Code{
+			{Body: []byte{wasm.OpcodeI32Const, 1, wasm.OpcodeEnd}},
+			{Body: []byte{wasm.OpcodeI32Const, 2, wasm.OpcodeEnd}},
+		},
+	})
+
+	ctx := context.Background()
+	exportingMod, err := r.Instantiate(ctx, exportingBin)
+	require.NoError(t, err)
+
+	instantiateClose(t, r, ctx, importingBin)
+
+	// Trigger GC to make sure the module instance that is the initial owner of the references in the table is collected.
+	time.Sleep(time.Millisecond * 10)
+	runtime.GC()
+	time.Sleep(time.Millisecond * 10)
+
+	main := exportingMod.ExportedFunction("main")
+	require.NotNil(t, main)
+
+	// Call call_indirect multiple times, should be safe.
+	for i := 0; i < 10; i++ {
+		_, err = main.Call(ctx, 0)
+		// Null function call
+		require.Error(t, err)
+
+		res, err := main.Call(ctx, 5)
+		require.NoError(t, err)
+		require.Equal(t, uint64(1), res[0])
+
+		res, err = main.Call(ctx, 6)
+		require.NoError(t, err)
+		require.Equal(t, uint64(2), res[0])
+		runtime.GC()
+		time.Sleep(time.Millisecond * 10)
+	}
+}
+
+func instantiateClose(t *testing.T, r wazero.Runtime, ctx context.Context, bin []byte) {
+	compiled, err := r.CompileModule(ctx, bin)
+	require.NoError(t, err)
+	m, err := r.InstantiateModule(ctx, compiled, wazero.NewModuleConfig())
+	require.NoError(t, err)
+	err = m.Close(ctx)
+	require.NoError(t, err)
+	require.NoError(t, compiled.Close(ctx))
+	runtime.GC()
+}

internal/wasm/store.go

@@ -360,6 +360,12 @@ func (s *Store) instantiate(
 	m.buildGlobals(module, m.Engine.FunctionInstanceReference)
 	m.buildMemory(module)
 	m.Exports = module.Exports
+	for _, exp := range m.Exports {
+		if exp.Type == ExternTypeTable {
+			t := m.Tables[exp.Index]
+			t.involvingModuleInstances = append(t.involvingModuleInstances, m)
+		}
+	}
 
 	// As of reference types proposal, data segment validation must happen after instantiation,
 	// and the side effect must persist even if there's out of bounds error after instantiation.
@@ -447,6 +453,12 @@ func (m *ModuleInstance) resolveImports(module *Module) (err error) {
 					}
 				}
 				m.Tables[i.IndexPerType] = importedTable
+				importedTable.involvingModuleInstancesMutex.Lock()
+				if len(importedTable.involvingModuleInstances) == 0 {
+					panic("BUG: involvingModuleInstances must not be nil when it's imported")
+				}
+				importedTable.involvingModuleInstances = append(importedTable.involvingModuleInstances, m)
+				importedTable.involvingModuleInstancesMutex.Unlock()
 			case ExternTypeMemory:
 				expected := i.DescMem
 				importedMemory := importedModule.MemoryInstance

internal/wasm/table.go

@@ -3,6 +3,7 @@ package wasm
 import (
 	"fmt"
 	"math"
+	"sync"
 
 	"github.com/tetratelabs/wazero/api"
 	"github.com/tetratelabs/wazero/internal/leb128"
@@ -131,6 +132,16 @@ type TableInstance struct {
 
 	// Type is either RefTypeFuncref or RefTypeExternRef.
 	Type RefType
+
+	// The following is only used when the table is exported.
+
+	// involvingModuleInstances is a set of module instances which are involved in the table instance.
+	// This is critical for safety purpose because once a table is imported, it can hold any reference to
+	// any function in the owner and importing module instances. Therefore, these module instance,
+	// transitively the compiled modules, must be alive as long as the table instance is alive.
+	involvingModuleInstances []*ModuleInstance
+	// involvingModuleInstancesMutex is a mutex to protect involvingModuleInstances.
+	involvingModuleInstancesMutex sync.RWMutex
 }
 
 // ElementInstance represents an element instance in a module.

internal/wasm/table_test.go

@@ -21,7 +21,7 @@ func Test_resolveImports_table(t *testing.T) {
 
 	t.Run("ok", func(t *testing.T) {
 		max := uint32(10)
-		tableInst := &TableInstance{Max: &max}
+		tableInst := &TableInstance{Max: &max, involvingModuleInstances: []*ModuleInstance{{}}}
 		s := newStore()
 		s.nameToModule[moduleName] = &ModuleInstance{
 			Tables:     []*TableInstance{tableInst},
@@ -36,6 +36,7 @@ func Test_resolveImports_table(t *testing.T) {
 		})
 		require.NoError(t, err)
 		require.Equal(t, m.Tables[0], tableInst)
+		require.Equal(t, m.Tables[0].involvingModuleInstances[1], m)
 	})
 	t.Run("minimum size mismatch", func(t *testing.T) {
 		s := newStore()