wasi: ensure the context used for _start is consistent (#273)

This adds `StoreConfig.Context` to centralize assignment of the initial
context used implicitly by the WebAssembly 1.0 (MVP) start function and
also the WASI snapshot-01 "_start" exported function. This also
backfills tests and comments around propagation.

Signed-off-by: Adrian Cole <adrian@tetrate.io>

internal/wasi/wasi_test.go

@@ -1,6 +1,7 @@
 package internalwasi
 
 import (
+	"context"
 	_ "embed"
 	"errors"
 	"fmt"
@@ -642,7 +643,7 @@ func instantiateWasmStore(t *testing.T, wasiFunction, wasiImport, moduleName str
 )`, wasiFunction, wasiImport)))
 	require.NoError(t, err)
 
-	store := wasm.NewStore(interpreter.NewEngine())
+	store := wasm.NewStore(context.Background(), interpreter.NewEngine())
 
 	snapshotPreview1Functions := SnapshotPreview1Functions(opts...)
 	goFunc := snapshotPreview1Functions[wasiFunction]

internal/wasm/host.go

@@ -12,10 +12,10 @@ import (
 // compile time check to ensure ModuleContext implements publicwasm.ModuleContext
 var _ publicwasm.ModuleContext = &ModuleContext{}
 
-func NewModuleContext(s *Store, instance *ModuleInstance) *ModuleContext {
+func NewModuleContext(ctx context.Context, engine Engine, instance *ModuleInstance) *ModuleContext {
 	return &ModuleContext{
-		Engine: s.Engine,
-		ctx:    context.Background(),
+		ctx:    ctx,
+		Engine: engine,
 		memory: instance.Memory,
 		Module: instance,
 	}
@@ -23,14 +23,14 @@ func NewModuleContext(s *Store, instance *ModuleInstance) *ModuleContext {
 
 // ModuleContext implements wasm.ModuleContext and wasm.Module
 type ModuleContext struct {
+	// ctx is the default context, exposed as wasm.ModuleContext Context
+	ctx context.Context
 	// Engine is exported for wazero.MakeWasmFunc
 	Engine Engine
 	// Module is exported for wazero.MakeWasmFunc
 	Module *ModuleInstance
 	// memory is exposed as wasm.ModuleContext Memory
 	memory publicwasm.Memory
-	// ctx is exposed as wasm.ModuleContext Context
-	ctx context.Context
 }
 
 // WithContext allows overriding context without re-allocation when the result would be the same.

internal/wasm/host_test.go

@@ -454,7 +454,7 @@ func TestFunction_Call(t *testing.T) {
 	name := "test"
 	fn := "fn"
 	engine := &nopEngine{}
-	s := NewStore(engine)
+	s := NewStore(context.Background(), engine)
 	m := &ModuleInstance{
 		Name: name,
 		Exports: map[string]*ExportInstance{
@@ -472,7 +472,7 @@ func TestFunction_Call(t *testing.T) {
 			},
 		},
 	}
-	ctx := NewModuleContext(s, m)
+	ctx := NewModuleContext(context.Background(), s.Engine, m)
 	s.ModuleInstances[name] = m
 	s.ModuleContexts[name] = ctx
 

internal/wasm/interpreter/interpreter_test.go

@@ -1,6 +1,7 @@
 package interpreter
 
 import (
+	"context"
 	"reflect"
 	"testing"
 
@@ -65,15 +66,7 @@ func TestInterpreter_CallHostFunc(t *testing.T) {
 		}}
 
 		// When calling a host func directly, there may be no stack. This ensures the module's memory is used.
-		it.callHostFunc(newModuleContext(&it, module), it.functions[0])
+		it.callHostFunc(wasm.NewModuleContext(context.Background(), &it, module), it.functions[0])
 		require.Same(t, memory, ctxMemory)
 	})
 }
-
-func newModuleContext(engine wasm.Engine, module *wasm.ModuleInstance) *wasm.ModuleContext {
-	ctx := wasm.NewModuleContext(&wasm.Store{
-		Engine:          engine,
-		ModuleInstances: map[string]*wasm.ModuleInstance{"test": module},
-	}, module)
-	return ctx
-}

internal/wasm/store.go

@@ -2,6 +2,7 @@ package internalwasm
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"math"
@@ -27,6 +28,9 @@ type (
 	Store struct {
 		// The following fields are wazero-specific fields of Store.
 
+		// ctx is the default context used for function calls
+		ctx context.Context
+
 		// Engine is a global context for a Store which is in responsible for compilation and execution of Wasm modules.
 		Engine Engine
 
@@ -231,8 +235,9 @@ func (m *ModuleInstance) GetExport(name string, kind ExportKind) (*ExportInstanc
 	return exp, nil
 }
 
-func NewStore(engine Engine) *Store {
+func NewStore(ctx context.Context, engine Engine) *Store {
 	return &Store{
+		ctx:                    ctx,
 		ModuleInstances:        map[string]*ModuleInstance{},
 		ModuleContexts:         map[string]*ModuleContext{},
 		TypeIDs:                map[string]FunctionTypeID{},
@@ -318,8 +323,8 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 	}
 
 	// Build the default context for calls to this module
-	ctx := NewModuleContext(s, instance)
-	s.ModuleContexts[name] = ctx
+	modCtx := NewModuleContext(s.ctx, s.Engine, instance)
+	s.ModuleContexts[name] = modCtx
 
 	// Now we are safe to finalize the state.
 	rollbackFuncs = nil
@@ -327,11 +332,11 @@ func (s *Store) Instantiate(module *Module, name string) (*ModuleContext, error)
 	// Execute the start function.
 	if module.StartSection != nil {
 		funcIdx := *module.StartSection
-		if _, err = s.Engine.Call(ctx, instance.Functions[funcIdx]); err != nil {
+		if _, err = s.Engine.Call(modCtx, instance.Functions[funcIdx]); err != nil {
 			return nil, fmt.Errorf("module[%s] start function failed: %w", name, err)
 		}
 	}
-	return ctx, nil
+	return modCtx, nil
 }
 
 // ModuleExports implements wasm.Store ModuleExports

internal/wasm/store_test.go

@@ -2,6 +2,7 @@ package internalwasm
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"math"
 	"os"
@@ -17,7 +18,7 @@ import (
 func TestStore_GetModuleInstance(t *testing.T) {
 	name := "test"
 
-	s := NewStore(nopEngineInstance)
+	s := NewStore(context.Background(), nopEngineInstance)
 
 	m1 := s.getModuleInstance(name)
 	require.Equal(t, m1, s.ModuleInstances[name])
@@ -28,7 +29,7 @@ func TestStore_GetModuleInstance(t *testing.T) {
 }
 
 func TestStore_AddHostFunction(t *testing.T) {
-	s := NewStore(nopEngineInstance)
+	s := NewStore(context.Background(), nopEngineInstance)
 
 	hf, err := NewGoFunc("fn", func(wasm.ModuleContext) {
 	})
@@ -57,7 +58,7 @@ func TestStore_AddHostFunction(t *testing.T) {
 }
 
 func TestStore_ExportImportedHostFunction(t *testing.T) {
-	s := NewStore(nopEngineInstance)
+	s := NewStore(context.Background(), nopEngineInstance)
 
 	hf, err := NewGoFunc("host_fn", func(wasm.ModuleContext) {
 	})
@@ -87,7 +88,7 @@ func TestStore_ExportImportedHostFunction(t *testing.T) {
 func TestStore_BuildFunctionInstances_FunctionNames(t *testing.T) {
 	name := "test"
 
-	s := NewStore(nopEngineInstance)
+	s := NewStore(context.Background(), nopEngineInstance)
 	mi := s.getModuleInstance(name)
 
 	zero := Index(0)
@@ -134,7 +135,7 @@ func (e *nopEngine) Compile(_ *FunctionInstance) error {
 
 func TestStore_addHostFunction(t *testing.T) {
 	t.Run("too many functions", func(t *testing.T) {
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		const max = 10
 		s.maximumFunctionAddress = max
 		s.Functions = make([]*FunctionInstance, max)
@@ -142,7 +143,7 @@ func TestStore_addHostFunction(t *testing.T) {
 		require.Error(t, err)
 	})
 	t.Run("ok", func(t *testing.T) {
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		for i := 0; i < 10; i++ {
 			f := &FunctionInstance{FunctionKind: FunctionKindGoNoContext}
 			require.Len(t, s.Functions, i)
@@ -161,7 +162,7 @@ func TestStore_addHostFunction(t *testing.T) {
 
 func TestStore_getTypeInstance(t *testing.T) {
 	t.Run("too many functions", func(t *testing.T) {
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		const max = 10
 		s.maximumFunctionTypes = max
 		s.TypeIDs = make(map[string]FunctionTypeID)
@@ -180,7 +181,7 @@ func TestStore_getTypeInstance(t *testing.T) {
 		} {
 			tc := tc
 			t.Run(tc.String(), func(t *testing.T) {
-				s := NewStore(nopEngineInstance)
+				s := NewStore(context.Background(), nopEngineInstance)
 				actual, err := s.getTypeInstance(tc)
 				require.NoError(t, err)
 
@@ -196,7 +197,7 @@ func TestStore_getTypeInstance(t *testing.T) {
 func TestStore_buildGlobalInstances(t *testing.T) {
 	t.Run("too many globals", func(t *testing.T) {
 		// Setup a store to have the reasonably low max on globals for testing.
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		const max = 10
 		s.maximumGlobals = max
 
@@ -205,7 +206,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 		require.Error(t, err)
 	})
 	t.Run("invalid constant expression", func(t *testing.T) {
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 
 		// Empty constant expression is invalid.
 		m := &Module{GlobalSection: []*Global{{Init: &ConstantExpression{}}}}
@@ -214,7 +215,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 	})
 
 	t.Run("global type mismatch", func(t *testing.T) {
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		m := &Module{GlobalSection: []*Global{{
 			// Global with i32.const initial value, but with type specified as f64 must be error.
 			Init: &ConstantExpression{Opcode: OpcodeI32Const, Data: []byte{0}},
@@ -233,7 +234,7 @@ func TestStore_buildGlobalInstances(t *testing.T) {
 
 		m := &Module{GlobalSection: []*Global{global}}
 
-		s := NewStore(nopEngineInstance)
+		s := NewStore(context.Background(), nopEngineInstance)
 		target := &ModuleInstance{}
 		_, err = s.buildGlobalInstances(m, target)
 		require.NoError(t, err)
@@ -328,7 +329,7 @@ func TestStore_executeConstExpression(t *testing.T) {
 		t.Run("global index out of range", func(t *testing.T) {
 			// Data holds the index in leb128 and this time the value exceeds len(globals) (=0).
 			expr := &ConstantExpression{Data: []byte{1}, Opcode: OpcodeGlobalGet}
-			globals := []*GlobalInstance{}
+			var globals []*GlobalInstance
 			_, _, err := executeConstExpression(globals, expr)
 			require.Error(t, err)
 		})

store.go

@@ -1,6 +1,7 @@
 package wazero
 
 import (
+	"context"
 	"fmt"
 
 	internalwasm "github.com/tetratelabs/wazero/internal/wasm"
@@ -23,24 +24,42 @@ func NewEngineJIT() *Engine { // TODO: compiler?
 
 // StoreConfig allows customization of a Store via NewStoreWithConfig
 type StoreConfig struct {
+	// Context is the default context used to initialize the module. Defaults to context.Background.
+	//
+	// Notes:
+	// * If the Module defines a start function, this is used to invoke it.
+	// * This is the outer-most ancestor of wasm.ModuleContext Context() during wasm.HostFunction invocations.
+	// * This is the default context of wasm.Function when callers pass nil.
+	//
+	// See https://www.w3.org/TR/wasm-core-1/#start-function%E2%91%A0
+	Context context.Context
 	// Engine defaults to NewEngineInterpreter
 	Engine *Engine
 }
 
 func NewStore() wasm.Store {
-	return internalwasm.NewStore(interpreter.NewEngine())
+	return internalwasm.NewStore(context.Background(), interpreter.NewEngine())
 }
 
 // NewStoreWithConfig returns a store with the given configuration.
 func NewStoreWithConfig(config *StoreConfig) wasm.Store {
+	ctx := config.Context
+	if ctx == nil {
+		ctx = context.Background()
+	}
 	engine := config.Engine
 	if engine == nil {
 		engine = NewEngineInterpreter()
 	}
-	return internalwasm.NewStore(engine.e)
+	return internalwasm.NewStore(ctx, engine.e)
 }
 
 // InstantiateModule instantiates the module namespace or errs if the configuration was invalid.
+//
+// Ex.
+//	exports, _ := wazero.InstantiateModule(wazero.NewStore(), mod)
+//
+// Note: StoreConfig.Context is used for any WebAssembly 1.0 (MVP) Start Function.
 func InstantiateModule(store wasm.Store, module *Module) (wasm.ModuleExports, error) {
 	internal, ok := store.(*internalwasm.Store)
 	if !ok {

store_test.go

@@ -0,0 +1,101 @@
+package wazero
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tetratelabs/wazero/wasm"
+)
+
+func TestFunction_Context(t *testing.T) {
+	type key string
+	storeCtx := context.WithValue(context.Background(), key("wa"), "zero")
+	config := &StoreConfig{Context: storeCtx}
+
+	notStoreCtx := context.WithValue(context.Background(), key("wazer"), "o")
+
+	tests := []struct {
+		name     string
+		ctx      context.Context
+		expected context.Context
+	}{
+		{
+			name:     "nil defaults to store context",
+			ctx:      nil,
+			expected: storeCtx,
+		},
+		{
+			name:     "set overrides store context",
+			ctx:      notStoreCtx,
+			expected: notStoreCtx,
+		},
+	}
+
+	for _, tt := range tests {
+		tc := tt
+
+		t.Run(tc.name, func(t *testing.T) {
+			store := NewStoreWithConfig(config)
+
+			// Define a host function so that we can catch the context propagated from a module function call
+			functionName := "fn"
+			expectedResult := uint64(math.MaxUint64)
+			hostFn := func(ctx wasm.ModuleContext) uint64 {
+				require.Equal(t, tc.expected, ctx.Context())
+				return expectedResult
+			}
+			mod := requireImportAndExportFunction(t, store, hostFn, functionName)
+
+			// Instantiate the module and get the export of the above hostFn
+			exports, err := InstantiateModule(store, mod)
+			require.NoError(t, err)
+			fn, ok := exports.Function(functionName)
+			require.True(t, ok)
+
+			// This fails if the function wasn't invoked, or had an unexpected context.
+			results, err := fn(tc.ctx)
+			require.NoError(t, err)
+			require.Equal(t, expectedResult, results[0])
+		})
+	}
+}
+
+func TestInstantiateModule_UsesStoreContext(t *testing.T) {
+	type key string
+	config := &StoreConfig{Context: context.WithValue(context.Background(), key("wa"), "zero")}
+	store := NewStoreWithConfig(config)
+
+	// Define a function that will be set as the start function
+	var calledStart bool
+	start := func(ctx wasm.ModuleContext) {
+		calledStart = true
+		require.Equal(t, config.Context, ctx.Context())
+	}
+	_, err := ExportHostFunctions(store, "", map[string]interface{}{"start": start})
+	require.NoError(t, err)
+
+	mod, err := DecodeModuleText([]byte(`(module $store_test.go
+	(import "" "start" (func $start))
+	(start $start)
+)`))
+	require.NoError(t, err)
+
+	// Instantiate the module, which calls the start function. This will fail if the context wasn't as intended.
+	_, err = InstantiateModule(store, mod)
+	require.NoError(t, err)
+	require.True(t, calledStart)
+}
+
+// requireImportAndExportFunction re-exports a host function because only host functions can see the propagated context.
+func requireImportAndExportFunction(t *testing.T, store wasm.Store, hostFn func(ctx wasm.ModuleContext) uint64, functionName string) *Module {
+	_, err := ExportHostFunctions(store, "host", map[string]interface{}{functionName: hostFn})
+	require.NoError(t, err)
+	wat := fmt.Sprintf(`(module (import "host" "%[1]s" (func (result i64))) (export "%[1]s" (func 0)))`, functionName)
+	mod, err := DecodeModuleText([]byte(wat))
+	require.NoError(t, err)
+	return mod
+}

tests/spectest/spec_test.go

@@ -268,7 +268,7 @@ func runTest(t *testing.T, newEngine func() wasm.Engine) {
 		wastName := filepath.Base(base.SourceFile)
 
 		t.Run(wastName, func(t *testing.T) {
-			store := wasm.NewStore(newEngine())
+			store := wasm.NewStore(context.Background(), newEngine())
 			addSpectestModule(t, store)
 
 			var lastInstanceName string

wasi.go

@@ -84,6 +84,7 @@ func WASISnapshotPreview1WithConfig(c *WASIConfig) map[string]interface{} {
 // * "_start" is an exported nullary function and does not export "_initialize"
 // * "memory" is an exported memory.
 //
+// Note: "_start" is invoked in the StoreConfig.Context.
 // Note: Exporting "__indirect_function_table" is mentioned as required, but not enforced here.
 // Note: The wasm.Functions return value does not restrict exports after "_start" as allowed in the specification.
 // Note: All TinyGo Wasm are WASI commands. They initialize memory on "_start" and import "fd_write" to implement panic.
@@ -102,9 +103,9 @@ func StartWASICommand(store wasm.Store, module *Module) (wasm.ModuleExports, err
 		return nil, err
 	}
 
-	ctx := internal.ModuleContexts[module.name]
-	start, _ := ctx.Function(internalwasi.FunctionStart)
-	if _, err = start(ctx.Context()); err != nil {
+	exports := internal.ModuleContexts[module.name]
+	start, _ := exports.Function(internalwasi.FunctionStart)
+	if _, err = start(exports.Context()); err != nil {
 		return nil, fmt.Errorf("module[%s] function[%s] failed: %w", module.name, internalwasi.FunctionStart, err)
 	}
 	return ret, nil

wasi_test.go

@@ -0,0 +1,42 @@
+package wazero
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tetratelabs/wazero/wasi"
+	"github.com/tetratelabs/wazero/wasm"
+)
+
+func TestStartWASICommand_UsesStoreContext(t *testing.T) {
+	type key string
+	config := &StoreConfig{Context: context.WithValue(context.Background(), key("wa"), "zero")}
+	store := NewStoreWithConfig(config)
+
+	// Define a function that will be re-exported as the WASI function: _start
+	var calledStart bool
+	start := func(ctx wasm.ModuleContext) {
+		calledStart = true
+		require.Equal(t, config.Context, ctx.Context())
+	}
+	_, err := ExportHostFunctions(store, "", map[string]interface{}{"start": start})
+	require.NoError(t, err)
+
+	mod, err := DecodeModuleText([]byte(`(module $wasi_test.go
+	(import "" "start" (func $start))
+	(memory 1)
+	(export "_start" (func $start))
+	(export "memory" (memory 0))
+)`))
+	require.NoError(t, err)
+
+	_, err = ExportHostFunctions(store, wasi.ModuleSnapshotPreview1, WASISnapshotPreview1())
+	require.NoError(t, err)
+
+	// Start the module as a WASI command. This will fail if the context wasn't as intended.
+	_, err = StartWASICommand(store, mod)
+	require.NoError(t, err)
+	require.True(t, calledStart)
+}

wasm/wasm.go

@@ -24,10 +24,19 @@ type ModuleExports interface {
 }
 
 // Function is an advanced API allowing efficient invocation of WebAssembly 1.0 (MVP) functions, given predefined
-// knowledge about the function signature. An error is returned for any failure looking up or invoking the function including
-// signature mismatch.
+// knowledge about the function signature. An error is returned for any failure looking up or invoking the function
+// including signature mismatch.
 //
-// Web Assembly 1.0 (MVP) Value Type Conversion:
+// If the `ctx` is nil, it defaults to the same context as the module was initialized with.
+//
+// To ensure context propagation in a HostFunction, use or derive `ctx` from ModuleContext.Context:
+//
+//	hostFunction := func(ctx wasm.ModuleContext, offset, byteCount uint32) uint32 {
+//		fn, _ = ctx.Function("__read")
+//		results, err := fn(ctx.Context(), offset, byteCount)
+//	--snip--
+//
+// The following describes how remaining parameters map to Web Assembly 1.0 (MVP) Value Types:
 //  * I32 - uint64(uint32,int32,int64)
 //  * I64 - uint64
 //  * F32 - EncodeF32 DecodeF32 from float32
@@ -43,8 +52,6 @@ type ModuleExports interface {
 //	results, _ := fn(ctx, wasm.EncodeF64(input))
 //	result := wasm.DecodeF64(result[0])
 //
-// Note: The ctx parameter will be the outer-most ancestor of ModuleContext.Context
-// ctx will default to context.Background() is nil is passed.
 // See https://www.w3.org/TR/wasm-core-1/#binary-valtype
 type Function func(ctx context.Context, params ...uint64) ([]uint64, error)