Repository: tetratelabs/wazero
Pull request #357 (2022-03-11 00:28:16)

Prevent callFrame and value stack from allocated on gorutine stack
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>

mathetake: I hope the comments on the change explain well! 

Patch:
diff --git a/internal/wasm/jit/engine.go b/internal/wasm/jit/engine.go
index 8ff3730d..cb56e6f1 100644
--- a/internal/wasm/jit/engine.go
+++ b/internal/wasm/jit/engine.go
@@ -468,8 +468,34 @@ func newEngine() *engine {
 	return &engine{compiledFunctions: map[*wasm.FunctionInstance]*compiledFunction{}}
 }
 
-// TODO: better make them configurable?
-const (
+// Do not make these variables as constants, otherwise there would be
+// dangerous memory accees from native code.
+//
+// Background: Go has a mechanism called "goroutine stack-shrink" where Go
+// runtime shrinks Gorotuine's stack when it is GCing. Shrinking means that
+// all the contents on the goroutine stack will be relocated by runtime,
+// Therefore, the memory address of these contents change undeterministically.
+// Not only shrinks, but also Goruntime grows the goroutine stack at any point
+// of function call entries, which also might end up relcating contents.
+//
+// On the other hand, we hold pointers to the data region of value stack and
+// callframe stack slices and use these raw pointers from native code.
+// Therefore, it is dangerous if these two stacks are allocated on stack
+// as these stack's address might be changed by Goruntime which we cannot
+// detect.
+//
+// By declaring theses values as `var`, slices created via `make([]..., var)`
+// will never be allocated on stack [1]. This means accessing these slices via
+// raw pointers is safe: As of version 1.18, Go's garbage collectornever relocates
+// heap-allocated objects (aka no compilation of memory [2]).
+//
+// On Go upgrades, re-validate heap-allocation via `go build -gcflags='-m' ./internal/wasm/jit/...`.
+//
+// [1] https://github.com/golang/go/blob/68ecdc2c70544c303aa923139a5f16caf107d955/src/cmd/compile/internal/escape/utils.go#L206-L208
+// [2] https://github.com/golang/go/blob/68ecdc2c70544c303aa923139a5f16caf107d955/src/runtime/mgc.go#L9
+// [3] https://mayurwadekar2.medium.com/escape-analysis-in-golang-ee40a1c064c1
+// [4] https://medium.com/@yulang.chu/go-stack-or-heap-2-slices-which-keep-in-stack-have-limitation-of-size-b3f3adfd6190
+var (
 	initialValueStackSize     = 64
 	initialCallFrameStackSize = 16
 )
@@ -499,12 +525,6 @@ func (ce *callEngine) popValue() (ret uint64) {
 	return
 }
 
-func (ce *callEngine) resetValueStackElement0Address() {
-	valueStackHeader := (*reflect.SliceHeader)(unsafe.Pointer(&ce.valueStack))
-	ce.globalContext.valueStackElement0Address = valueStackHeader.Data
-	ce.globalContext.valueStackLen = uint64(valueStackHeader.Len)
-}
-
 func (ce *callEngine) pushValue(v uint64) {
 	ce.valueStack[ce.valueStackTopIndex()] = v
 	ce.valueStackContext.stackPointer++
@@ -603,11 +623,6 @@ jitentry:
 				frame.String(), ce.valueStackContext.stackBasePointer, ce.valueStackContext.stackPointer)
 		}
 
-		// Seems like even not using "append", accessing the slice might lead to relocation by Go runtime on amd64.
-		// For safety, we reset the cached addresses and length here.
-		// TODO: investigate the root cause.
-		ce.resetValueStackElement0Address()
-
 		// Call into the JIT code.
 		jitcall(frame.returnAddress, uintptr(unsafe.Pointer(ce)))
 
@@ -678,6 +693,9 @@ func (ce *callEngine) builtinFunctionGrowValueStack(stackPointerCeil uint64) {
 	top := ce.valueStackContext.stackBasePointer + ce.valueStackContext.stackPointer
 	copy(newStack[:top], ce.valueStack[:top])
 	ce.valueStack = newStack
+	valueStackHeader := (*reflect.SliceHeader)(unsafe.Pointer(&ce.valueStack))
+	ce.globalContext.valueStackElement0Address = valueStackHeader.Data
+	ce.globalContext.valueStackLen = uint64(valueStackHeader.Len)
 }
 
 var callStackCeiling = uint64(buildoptions.CallStackCeiling)


Test patch:
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index a0a9bf0c..e18b3997 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -283,3 +283,96 @@ func TestRelease(t *testing.T) {
 		})
 	}
 }
+
+// Ensures that value stack and callframe stack are allocated on heap which
+// allows us to safely access to their data region from native code.
+// See comments on initialValueStackSize and initialCallFrameStackSize.
+func TestSliceAllocatedOnHeap(t *testing.T) {
+	e := newEngine()
+	store := wasm.NewStore(context.Background(), e, wasm.Features20191205)
+
+	const hostModuleName = "env"
+	const hostFnName = "grow_and_shrink_goroutine_stack"
+	hm, err := wasm.NewHostModule(hostModuleName, map[string]interface{}{hostFnName: func() {
+		// This function aggressively grow the goroutine stack by recursively
+		// calling the function many times.
+		var callNum = 1000
+		var growGoroutineStack func()
+		growGoroutineStack = func() {
+			if callNum != 0 {
+				callNum--
+				growGoroutineStack()
+			}
+		}
+		growGoroutineStack()
+
+		// Trigger relocation of goroutine stack because at this point we have majority of
+		// goroutine stack unused after recursive call.
+		runtime.GC()
+	}})
+	require.NoError(t, err)
+
+	_, err = store.Instantiate(hm, hostModuleName)
+	require.NoError(t, err)
+
+	const valueStackCorruption = "value_stack_corruption"
+	const callStackCorruption = "call_stack_corruption"
+	const expectedReturnValue = 0x1
+	m := &wasm.Module{
+		TypeSection: []*wasm.FunctionType{
+			{Params: []wasm.ValueType{}, Results: []wasm.ValueType{wasm.ValueTypeI32}},
+			{Params: []wasm.ValueType{}, Results: []wasm.ValueType{}},
+		},
+		FunctionSection: []wasm.Index{
+			wasm.Index(0),
+			wasm.Index(0),
+			wasm.Index(0),
+		},
+		CodeSection: []*wasm.Code{
+			{
+				// value_stack_corruption
+				Body: []byte{
+					wasm.OpcodeCall, 0, // Call host function to shrink Goroutine stack
+					// We expect this value is returned, but if the stack is allocated on
+					// goroutine stack, we write this expected value into the old-location of
+					// stack.
+					wasm.OpcodeI32Const, expectedReturnValue,
+					wasm.OpcodeEnd,
+				},
+			},
+			{
+				// call_stack_corruption
+				Body: []byte{
+					wasm.OpcodeCall, 3, // Call the wasm function below.
+					// At this point, call stack's memory looks like [call_stack_corruption, index3]
+					// With this function call it should end up [call_stack_corruption, host func]
+					// but if the callframe stack is allocated on goroutine stack, we exit the native code
+					// with  [call_stack_corruption, index3] (old call frame stack) with HostCall status code,
+					// and end up trying to call index3 as a host function which results in nil pointer exception.
+					wasm.OpcodeCall, 0,
+					wasm.OpcodeI32Const, expectedReturnValue,
+					wasm.OpcodeEnd,
+				},
+			},
+			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}},
+		},
+		ImportSection: []*wasm.Import{{Module: hostModuleName, Name: hostFnName, DescFunc: 1}},
+		ExportSection: map[string]*wasm.Export{
+			valueStackCorruption: {Type: wasm.ExternTypeFunc, Index: 1, Name: valueStackCorruption},
+			callStackCorruption:  {Type: wasm.ExternTypeFunc, Index: 2, Name: callStackCorruption},
+		},
+	}
+
+	mi, err := store.Instantiate(m, "")
+	require.NoError(t, err)
+
+	for _, fnName := range []string{valueStackCorruption, callStackCorruption} {
+		fnName := fnName
+		t.Run(fnName, func(t *testing.T) {
+			ret, err := mi.ExportedFunction(fnName).Call(nil)
+			require.NoError(t, err)
+
+			require.Equal(t, uint32(expectedReturnValue), uint32(ret[0]))
+		})
+	}
+}
diff --git a/internal/wasm/jit/jit_arm64_test.go b/internal/wasm/jit/jit_arm64_test.go
index d970f4aa..68acdd98 100644
--- a/internal/wasm/jit/jit_arm64_test.go
+++ b/internal/wasm/jit/jit_arm64_test.go
@@ -300,7 +300,6 @@ func TestArm64Compiler_releaseRegisterToStack(t *testing.T) {
 
 			// Run native code after growing the value stack.
 			env.callEngine().builtinFunctionGrowValueStack(tc.stackPointer)
-			env.callEngine().resetValueStackElement0Address()
 			env.exec(code)
 
 			// JIT status must be returned and stack pointer must end up the specified one.
@@ -383,7 +382,6 @@ func TestArm64Compiler_compileLoadValueOnStackToRegister(t *testing.T) {
 
 			// Run native code after growing the value stack, and place the original value.
 			env.callEngine().builtinFunctionGrowValueStack(tc.stackPointer)
-			env.callEngine().resetValueStackElement0Address()
 			env.stack()[tc.stackPointer] = val
 			env.exec(code)