Repository: tetratelabs/wazero
Pull request #737 (2022-08-06 09:46:07)

Fixes function calls / memory access with wasm-implemented host functions 
Signed-off-by: Takeshi Yoneda <takeshi@tetrate.io>


Patch:
diff --git a/config.go b/config.go
index 99bef862..cee38053 100644
--- a/config.go
+++ b/config.go
@@ -565,12 +565,15 @@ type ModuleConfig interface {
 	// See WithNanosleep
 	WithSysNanosleep() ModuleConfig
 
-	// WithRandSource configures a source of random bytes. Defaults to crypto/rand.Reader.
+	// WithRandSource configures a source of random bytes. Defaults to return a
+	// deterministic source. You might override this with crypto/rand.Reader
 	//
-	// This reader is most commonly used by the functions like "random_get" in "wasi_snapshot_preview1" or "seed" in
-	// AssemblyScript standard "env" although it could be used by functions imported from other modules.
+	// This reader is most commonly used by the functions like "random_get" in
+	// "wasi_snapshot_preview1", "seed" in AssemblyScript standard "env", and
+	// "getRandomData" when runtime.GOOS is "js".
 	//
-	// Note: The caller is responsible to close any io.Reader they supply: It is not closed on api.Module Close.
+	// Note: The caller is responsible to close any io.Reader they supply: It
+	// is not closed on api.Module Close.
 	WithRandSource(io.Reader) ModuleConfig
 }
 
diff --git a/internal/engine/compiler/impl_amd64.go b/internal/engine/compiler/impl_amd64.go
index 6b26c6d1..4a1e8242 100644
--- a/internal/engine/compiler/impl_amd64.go
+++ b/internal/engine/compiler/impl_amd64.go
@@ -4509,7 +4509,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, functionAddres
 	// Since call frame stack pointer is the index for callEngine.callFrameStack slice,
 	// here we get the actual offset in bytes via shifting callFrameStackPointerRegister by callFrameDataSizeMostSignificantSetBit.
 	// That is valid because the size of callFrame struct is a power of 2 (see TestVerifyOffsetValue), which means
-	// multiplying withe the size of struct equals shifting by its most significant bit.
+	// multiplying with the size of struct equals shifting by its most significant bit.
 	c.assembler.CompileConstToRegister(amd64.SHLQ, int64(callFrameDataSizeMostSignificantSetBit), callFrameStackPointerRegister)
 
 	// At this point, callFrameStackPointerRegister holds the offset in call frame slice in bytes,
@@ -4637,6 +4637,7 @@ func (c *amd64Compiler) compileCallFunctionImpl(index wasm.Index, functionAddres
 	if err := c.compileModuleContextInitialization(); err != nil {
 		return err
 	}
+
 	// Due to the change to callEngine.valueStackContext.stackBasePointer.
 	c.compileReservedStackBasePointerInitialization()
 
@@ -4926,7 +4927,7 @@ func (c *amd64Compiler) compileReservedStackBasePointerInitialization() {
 }
 
 func (c *amd64Compiler) compileReservedMemoryPointerInitialization() {
-	if c.ir.HasMemory {
+	if c.ir.HasMemory || c.ir.UsesMemory {
 		c.assembler.CompileMemoryToRegister(amd64.MOVQ,
 			amd64ReservedRegisterForCallEngine, callEngineModuleContextMemoryElement0AddressOffset,
 			amd64ReservedRegisterForMemory,
@@ -4965,13 +4966,6 @@ func (c *amd64Compiler) compileMaybeGrowValueStack() error {
 // callEngine.ModuleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *amd64Compiler) compileModuleContextInitialization() error {
-	if c.ir.IsHostFunction {
-		// If this is the host function, our semantic is that host functions use the caller's context.
-		// Therefore, we can skip all the initialization steps here as at this point, the caller's memory, etc
-		// are already set on the callEngine.
-		return nil
-	}
-
 	// amd64CallingConventionModuleInstanceAddressRegister holds the module instance's address
 	// so mark it used so that it won't be used as a free register until the module context initialization finishes.
 	c.locationStack.markRegisterUsed(amd64CallingConventionModuleInstanceAddressRegister)
@@ -5061,7 +5055,7 @@ func (c *amd64Compiler) compileModuleContextInitialization() error {
 		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister2,
 			amd64ReservedRegisterForCallEngine, callEngineModuleContextMemorySliceLenOffset)
 
-		// Set elemnt zero address.
+		// Set element zero address.
 		c.assembler.CompileMemoryToRegister(amd64.MOVQ, tmpRegister, memoryInstanceBufferOffset, tmpRegister2)
 		c.assembler.CompileRegisterToMemory(amd64.MOVQ, tmpRegister2,
 			amd64ReservedRegisterForCallEngine, callEngineModuleContextMemoryElement0AddressOffset)
diff --git a/internal/engine/compiler/impl_arm64.go b/internal/engine/compiler/impl_arm64.go
index f370174a..928f360a 100644
--- a/internal/engine/compiler/impl_arm64.go
+++ b/internal/engine/compiler/impl_arm64.go
@@ -4117,7 +4117,7 @@ func (c *arm64Compiler) compileReservedStackBasePointerRegisterInitialization()
 }
 
 func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
-	if c.ir.HasMemory {
+	if c.ir.HasMemory || c.ir.UsesMemory {
 		// "arm64ReservedRegisterForMemory = ce.MemoryElement0Address"
 		c.assembler.CompileMemoryToRegister(
 			arm64.LDRD,
@@ -4131,12 +4131,6 @@ func (c *arm64Compiler) compileReservedMemoryRegisterInitialization() {
 // ce.moduleContext.ModuleInstanceAddress.
 // This is called in two cases: in function preamble, and on the return from (non-Go) function calls.
 func (c *arm64Compiler) compileModuleContextInitialization() error {
-	if c.ir.IsHostFunction {
-		// If this is the host function, our semantic is that host functions use the caller's context.
-		// Therefore, we can skip all the initialization steps here as at this point, the caller's memory, etc
-		// are already set on the callEngine.
-		return nil
-	}
 
 	c.markRegisterUsed(arm64CallingConventionModuleInstanceAddressRegister)
 	defer c.markRegisterUnused(arm64CallingConventionModuleInstanceAddressRegister)
diff --git a/internal/platform/crypto.go b/internal/platform/crypto.go
new file mode 100644
index 00000000..c141f00f
--- /dev/null
+++ b/internal/platform/crypto.go
@@ -0,0 +1,17 @@
+package platform
+
+import (
+	"io"
+	"math/rand"
+)
+
+// seed is a fixed seed value for NewFakeRandSource.
+//
+// Trivia: While arbitrary, 42 was chosen as it is the "Ultimate Answer" in
+// the Douglas Adams novel "The Hitchhiker's Guide to the Galaxy."
+const seed = int64(42)
+
+// NewFakeRandSource returns a deterministic source of random values.
+func NewFakeRandSource() io.Reader {
+	return rand.New(rand.NewSource(seed))
+}
diff --git a/internal/sys/sys.go b/internal/sys/sys.go
index 663773e0..7b1e405c 100644
--- a/internal/sys/sys.go
+++ b/internal/sys/sys.go
@@ -2,7 +2,6 @@ package sys
 
 import (
 	"context"
-	"crypto/rand"
 	"errors"
 	"fmt"
 	"io"
@@ -118,7 +117,7 @@ func (c *Context) FS(ctx context.Context) *FSContext {
 	return c.fsc
 }
 
-// RandSource is a source of random bytes and defaults to crypto/rand.Reader.
+// RandSource is a source of random bytes and defaults to a deterministic source.
 // see wazero.ModuleConfig WithRandSource
 func (c *Context) RandSource() io.Reader {
 	return c.randSource
@@ -153,8 +152,10 @@ func NewContext(
 	stdin io.Reader,
 	stdout, stderr io.Writer,
 	randSource io.Reader,
-	walltime *sys.Walltime, walltimeResolution sys.ClockResolution,
-	nanotime *sys.Nanotime, nanotimeResolution sys.ClockResolution,
+	walltime *sys.Walltime,
+	walltimeResolution sys.ClockResolution,
+	nanotime *sys.Nanotime,
+	nanotimeResolution sys.ClockResolution,
 	nanosleep *sys.Nanosleep,
 	fs fs.FS,
 ) (sysCtx *Context, err error) {
@@ -187,7 +188,7 @@ func NewContext(
 	}
 
 	if randSource == nil {
-		sysCtx.randSource = rand.Reader
+		sysCtx.randSource = platform.NewFakeRandSource()
 	} else {
 		sysCtx.randSource = randSource
 	}


Test patch:
diff --git a/assemblyscript/assemblyscript_test.go b/assemblyscript/assemblyscript_test.go
index a23895fc..7f10dcfe 100644
--- a/assemblyscript/assemblyscript_test.go
+++ b/assemblyscript/assemblyscript_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	_ "embed"
+	"encoding/hex"
 	"errors"
 	"io"
 	"strings"
@@ -123,8 +124,7 @@ func TestAbort_Error(t *testing.T) {
 }
 
 func TestSeed(t *testing.T) {
-	b := []byte{0, 1, 2, 3, 4, 5, 6, 7}
-	mod, r, log := requireProxyModule(t, NewFunctionExporter(), wazero.NewModuleConfig().WithRandSource(bytes.NewReader(b)))
+	mod, r, log := requireProxyModule(t, NewFunctionExporter(), wazero.NewModuleConfig())
 	defer r.Close(testCtx)
 
 	ret, err := mod.ExportedFunction(functionSeed).Call(testCtx)
@@ -132,11 +132,11 @@ func TestSeed(t *testing.T) {
 	require.Equal(t, `
 --> proxy.seed()
 	==> env.~lib/builtins/seed()
-	<== (7.949928895127363e-275)
-<-- (7.949928895127363e-275)
+	<== (4.958153677776298e-175)
+<-- (4.958153677776298e-175)
 `, "\n"+log.String())
 
-	require.Equal(t, b, u64.LeBytes(ret[0]))
+	require.Equal(t, "538c7f96b164bf1b", hex.EncodeToString(u64.LeBytes(ret[0])))
 }
 
 func TestSeed_error(t *testing.T) {
diff --git a/config_test.go b/config_test.go
index a17b2711..729450fe 100644
--- a/config_test.go
+++ b/config_test.go
@@ -2,6 +2,7 @@ package wazero
 
 import (
 	"context"
+	"crypto/rand"
 	"io"
 	"io/fs"
 	"math"
@@ -479,6 +480,23 @@ func TestModuleConfig_toSysContext(t *testing.T) {
 				testFS2, // fs
 			),
 		},
+		{
+			name:  "WithRandSource",
+			input: base.WithRandSource(rand.Reader),
+			expected: requireSysContext(t,
+				math.MaxUint32, // max
+				nil,            // args
+				nil,            // environ
+				nil,            // stdin
+				nil,            // stdout
+				nil,            // stderr
+				rand.Reader,    // randSource
+				&wt, 1,         // walltime, walltimeResolution
+				&nt, 1, // nanotime, nanotimeResolution
+				nil, // nanosleep
+				nil, // fs
+			),
+		},
 	}
 
 	for _, tt := range tests {
diff --git a/internal/sys/sys_test.go b/internal/sys/sys_test.go
index b4fc8b5c..1779dd36 100644
--- a/internal/sys/sys_test.go
+++ b/internal/sys/sys_test.go
@@ -3,7 +3,6 @@ package sys
 import (
 	"bytes"
 	"context"
-	"crypto/rand"
 	"io"
 	"testing"
 	"time"
@@ -55,7 +54,7 @@ func TestDefaultSysContext(t *testing.T) {
 	require.Zero(t, sysCtx.Nanotime(testCtx)) // See above on functions.
 	require.Equal(t, sys.ClockResolution(1), sysCtx.NanotimeResolution())
 	require.Equal(t, &ns, sysCtx.nanosleep)
-	require.Equal(t, rand.Reader, sysCtx.RandSource())
+	require.Equal(t, platform.NewFakeRandSource(), sysCtx.RandSource())
 	require.Equal(t, NewFSContext(testfs.FS{}), sysCtx.FS(testCtx))
 }
 
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 08283031..f5777983 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -790,15 +790,22 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 
 	callReadMem := &wasm.Code{ // shows indirect calls still use the same memory
 		IsHostFunction: true,
-		Body:           []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd},
+		Body: []byte{
+			wasm.OpcodeCall, 1,
+			// On the return from the another host function,
+			// we should still be able to access the memory.
+			wasm.OpcodeI32Const, 0,
+			wasm.OpcodeI32Load, 0x2, 0x0,
+			wasm.OpcodeEnd,
+		},
 	}
 	hostModule := &wasm.Module{
 		TypeSection:     []*wasm.FunctionType{ft},
 		FunctionSection: []wasm.Index{0, 0},
-		CodeSection:     []*wasm.Code{readMem, callReadMem},
+		CodeSection:     []*wasm.Code{callReadMem, readMem},
 		ExportSection: []*wasm.Export{
-			{Name: readMemName, Type: wasm.ExternTypeFunc, Index: 0},
-			{Name: callReadMemName, Type: wasm.ExternTypeFunc, Index: 1},
+			{Name: callReadMemName, Type: wasm.ExternTypeFunc, Index: 0},
+			{Name: readMemName, Type: wasm.ExternTypeFunc, Index: 1},
 		},
 		NameSection: &wasm.NameSection{
 			ModuleName:    "host",
@@ -832,8 +839,8 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 			{Name: callImportCallReadMemName, Type: wasm.ExternTypeFunc, Index: 3},
 		},
 		CodeSection: []*wasm.Code{
-			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Calling the index 0 = readMemFn.
-			{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}}, // Calling the index 1 = callReadMemFn.
+			{Body: []byte{wasm.OpcodeCall, 0, wasm.OpcodeEnd}}, // Calling the index 0 = callReadMemFn.
+			{Body: []byte{wasm.OpcodeCall, 1, wasm.OpcodeEnd}}, // Calling the index 1 = readMemFn.
 		},
 		NameSection: &wasm.NameSection{
 			ModuleName: "importing",
@@ -854,7 +861,7 @@ func setupCallMemTests(t *testing.T, e wasm.Engine, readMem *wasm.Code, fnlf exp
 	importing := &wasm.ModuleInstance{Name: importingModule.NameSection.ModuleName, TypeIDs: []wasm.FunctionTypeID{0}}
 	importingFunctions := importing.BuildFunctions(importingModule, buildListeners(fnlf, importingModule))
 	// Note: adds imported functions readMemFn and callReadMemFn at index 0 and 1.
-	importing.Functions = append([]*wasm.FunctionInstance{readMemFn, callReadMemFn}, importingFunctions...)
+	importing.Functions = append([]*wasm.FunctionInstance{callReadMemFn, readMemFn}, importingFunctions...)
 	importing.BuildExports(importingModule.ExportSection)
 
 	// Compile the importing module
diff --git a/wasi_snapshot_preview1/random_test.go b/wasi_snapshot_preview1/random_test.go
index ee778461..53b42147 100644
--- a/wasi_snapshot_preview1/random_test.go
+++ b/wasi_snapshot_preview1/random_test.go
@@ -12,8 +12,7 @@ import (
 )
 
 func Test_randomGet(t *testing.T) {
-	mod, r, log := requireProxyModule(t, wazero.NewModuleConfig().
-		WithRandSource(deterministicRandomSource()))
+	mod, r, log := requireProxyModule(t, wazero.NewModuleConfig())
 	defer r.Close(testCtx)
 
 	expectedMemory := []byte{
@@ -42,8 +41,7 @@ func Test_randomGet(t *testing.T) {
 }
 
 func Test_randomGet_Errors(t *testing.T) {
-	mod, r, log := requireProxyModule(t, wazero.NewModuleConfig().
-		WithRandSource(deterministicRandomSource()))
+	mod, r, log := requireProxyModule(t, wazero.NewModuleConfig())
 	defer r.Close(testCtx)
 
 	memorySize := mod.Memory().Size(testCtx)
diff --git a/wasi_snapshot_preview1/wasi_test.go b/wasi_snapshot_preview1/wasi_test.go
index 7f990edb..02d1843c 100644
--- a/wasi_snapshot_preview1/wasi_test.go
+++ b/wasi_snapshot_preview1/wasi_test.go
@@ -4,8 +4,6 @@ import (
 	"bytes"
 	"context"
 	_ "embed"
-	"io"
-	"math/rand"
 	"testing"
 
 	"github.com/tetratelabs/wazero"
@@ -16,12 +14,6 @@ import (
 	"github.com/tetratelabs/wazero/internal/testing/require"
 )
 
-const seed = int64(42) // fixed seed value
-
-var deterministicRandomSource = func() io.Reader {
-	return rand.New(rand.NewSource(seed))
-}
-
 // testCtx is an arbitrary, non-default context. Non-nil also prevents linter errors.
 var testCtx = context.WithValue(context.Background(), struct{}{}, "arbitrary")