Repository: tetratelabs/wazero (go)

Issue:
Less memory copies with Memory.MapSlice(offset, []byte)
Folks are doing more work moving data into the memory sandbox addressable by Wasm.

For example, @knqyf263 is working on [structured encoding][1] and @inkeliz is working on [random-access encoding][2]. Structs aside, the problem exists even with strings. You need to pass them into wasm (currently a copy) and its sandbox shouldn't be violated in the process..

Some alternatives solve this by providing unsafe access to the entire memory (ex a []byte). However, doing this breaks the sandbox and also any access controls you may have around boundaries. Furthermore, it may confuse who owns allocation, and typically that's the wasm code itself (ex TinyGo or rust malloc/free functions). Finally, it cements the underlying structure of memory, making changing it, ex [to pages of byte arrays][3], difficult later.

Another way could be to scope a memory overlay, given an offset and a range. For example, you could call an exported function to allocate memory, then apply an overlay which reads from that instead of the underlying source.

```go
var thing []byte = // something to manage in go

// Reserve linear memory using ex TinyGo's allocator.
results, _ := malloc.Call(ctx, len(thing))
ptr := results[0]
defer free.Call(ctx, ptr)

// map the slice so that until unmap is called, any writes to this region are write-through regardless
// of whether it is updated by Go or Wasm. Since this block is after malloc and before free, TinyGo
// shouldn't attempt to clear it.
unmap := mod.Memory().MapSlice(ptr, thing)
defer unmap()
```

Note this pattern could be accomplished many ways. For example if there is global mutable state,
the size could be indicated by a Global which is allocated during an "init" function. The actual
mapping could be done in a host function, which makes it possible to do in a .start function.
Regardless of how this is done, it could allow more precise handing over control of a region to
the host vs exposing the entire memory.

To do this internally may be trickier as there could be many stacks of these layers, so I'm not
entirely sure we want to even do this (think JIT). However, I figured we can think about it a while
and gather feedback about options.

[1]: #523
[2]: https://github.com/tetratelabs/wazero/issues/523#issuecomment-1117634602
[3]: #288

Pull request description:
Documents Memory.Read as a way to share memory between Go and Wasm
By testing the side-effects of Memory.Read, we ensure users who control
the underlying memory capacity can use the returned slice for
write-through access to Wasm addressible memory. Notably, this allows a
shared fixed length data structure to exist with a pointer on the Go
side and a memory offset on the Wasm side.

Thanks to @inkeliz for this idea and help on the code.

Fixes #525

Patch:
diff --git a/api/wasm.go b/api/wasm.go
index 34b987d3..a0ba7a66 100644
--- a/api/wasm.go
+++ b/api/wasm.go
@@ -215,6 +215,20 @@ type Memory interface {
 	ReadFloat64Le(ctx context.Context, offset uint32) (float64, bool)
 
 	// Read reads byteCount bytes from the underlying buffer at the offset or returns false if out of range.
+	//
+	// This returns a view of the underlying memory, not a copy. This means any writes to the slice returned are visible
+	// to Wasm, and any updates from Wasm are visible reading the returned slice.
+	//
+	// For example:
+	//	buf, _ = memory.Read(ctx, offset, byteCount)
+	//	buf[1] = 'a' // writes through to memory, meaning Wasm code see 'a' at that position.
+	//
+	// If you don't desire this behavior, make a copy of the returned slice before affecting it.
+	//
+	// Note: The returned slice is no longer shared on a capacity change. For example, `buf = append(buf, 'a')` might result
+	// in a slice that is no longer shared. The same exists Wasm side. For example, if Wasm changes its memory capacity,
+	// ex via "memory.grow"), the host slice is no longer shared. Those who need a stable view must set Wasm memory
+	// min=max, or use wazero.RuntimeConfig WithMemoryCapacityPages to ensure max is always allocated.
 	Read(ctx context.Context, offset, byteCount uint32) ([]byte, bool)
 
 	// WriteByte writes a single byte to the underlying buffer at the offset in or returns false if out of range.
diff --git a/internal/wasm/binary/data.go b/internal/wasm/binary/data.go
index c1a2ac76..74e08089 100644
--- a/internal/wasm/binary/data.go
+++ b/internal/wasm/binary/data.go
@@ -12,7 +12,7 @@ import (
 // dataSegmentPrefix represents three types of data segments.
 //
 // https://www.w3.org/TR/2022/WD-wasm-core-2-20220419/binary/modules.html#data-section
-type dataSegmentPrefix = byte
+type dataSegmentPrefix = uint32
 
 const (
 	// dataSegmentPrefixActive is the prefix for the version 1.0 compatible data segment, which is classified as "active" in 2.0.
@@ -24,7 +24,7 @@ const (
 )
 
 func decodeDataSegment(r *bytes.Reader, enabledFeatures wasm.Features) (*wasm.DataSegment, error) {
-	dataSegmentPrefx, err := r.ReadByte()
+	dataSegmentPrefx, _, err := leb128.DecodeUint32(r)
 	if err != nil {
 		return nil, fmt.Errorf("read data segment prefix: %w", err)
 	}
diff --git a/internal/wasm/binary/element.go b/internal/wasm/binary/element.go
index b869abf6..258618d7 100644
--- a/internal/wasm/binary/element.go
+++ b/internal/wasm/binary/element.go
@@ -99,7 +99,7 @@ const (
 )
 
 func decodeElementSegment(r *bytes.Reader, enabledFeatures wasm.Features) (*wasm.ElementSegment, error) {
-	prefix, err := r.ReadByte()
+	prefix, _, err := leb128.DecodeUint32(r)
 	if err != nil {
 		return nil, fmt.Errorf("read element prefix: %w", err)
 	}
diff --git a/internal/wasm/memory.go b/internal/wasm/memory.go
index 50542624..9ef179b8 100644
--- a/internal/wasm/memory.go
+++ b/internal/wasm/memory.go
@@ -37,6 +37,18 @@ type MemoryInstance struct {
 	Min, Cap, Max uint32
 }
 
+// NewMemoryInstance creates a new instance based on the parameters in the SectionIDMemory.
+func NewMemoryInstance(memSec *Memory) *MemoryInstance {
+	min := MemoryPagesToBytesNum(memSec.Min)
+	capacity := MemoryPagesToBytesNum(memSec.Cap)
+	return &MemoryInstance{
+		Buffer: make([]byte, min, capacity),
+		Min:    memSec.Min,
+		Cap:    memSec.Cap,
+		Max:    memSec.Max,
+	}
+}
+
 // Size implements the same method as documented on api.Memory.
 func (m *MemoryInstance) Size(_ context.Context) uint32 {
 	// Note: If you use the context.Context param, don't forget to coerce nil to context.Background()!
diff --git a/internal/wasm/module.go b/internal/wasm/module.go
index 0eb1a9d3..f2b06514 100644
--- a/internal/wasm/module.go
+++ b/internal/wasm/module.go
@@ -559,14 +559,7 @@ func paramNames(localNames IndirectNameMap, funcIdx uint32, paramLen int) []stri
 func (m *Module) buildMemory() (mem *MemoryInstance) {
 	memSec := m.MemorySection
 	if memSec != nil {
-		min := MemoryPagesToBytesNum(memSec.Min)
-		capacity := MemoryPagesToBytesNum(memSec.Cap)
-		mem = &MemoryInstance{
-			Buffer: make([]byte, min, capacity),
-			Min:    memSec.Min,
-			Cap:    memSec.Cap,
-			Max:    memSec.Max,
-		}
+		mem = NewMemoryInstance(memSec)
 	}
 	return
 }


Test patch:
diff --git a/internal/integration_test/engine/adhoc_test.go b/internal/integration_test/engine/adhoc_test.go
index 766ca4c7..8c3ae106 100644
--- a/internal/integration_test/engine/adhoc_test.go
+++ b/internal/integration_test/engine/adhoc_test.go
@@ -15,6 +15,8 @@ import (
 	"github.com/tetratelabs/wazero/sys"
 )
 
+const memoryCapacityPages = 2
+
 // testCtx is an arbitrary, non-default context. Non-nil also prevents linter errors.
 var testCtx = context.WithValue(context.Background(), struct{}{}, "arbitrary")
 
@@ -43,6 +45,9 @@ func TestEngineInterpreter(t *testing.T) {
 }
 
 func runAllTests(t *testing.T, tests map[string]func(t *testing.T, r wazero.Runtime), config wazero.RuntimeConfig) {
+	config.WithMemoryCapacityPages(func(minPages uint32, maxPages *uint32) uint32 {
+		return memoryCapacityPages
+	})
 	for name, testf := range tests {
 		name := name   // pin
 		testf := testf // pin
@@ -431,6 +436,13 @@ func testMemOps(t *testing.T, r wazero.Runtime) {
   (export "size" (func $size))
   (export "grow" (func $grow))
   (export "memory" (memory 0))
+
+  (func $store (param $offset i32)
+	local.get 0    ;; memory offset
+	i64.const 1
+	i64.store
+  )
+  (export "store" (func $store))
 )`))
 	require.NoError(t, err)
 	defer memory.Close(testCtx)
@@ -444,16 +456,38 @@ func testMemOps(t *testing.T, r wazero.Runtime) {
 	require.Zero(t, results[0])
 	require.Zero(t, memory.ExportedMemory("memory").Size(testCtx))
 
+	// Any offset should be out of bounds error even when it is less than memory capacity(=memoryCapacityPages).
+	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
+	require.Error(t, err) // Out of bounds error.
+
 	// Try to grow the memory by one page
 	results, err = memory.ExportedFunction("grow").Call(testCtx, 1)
 	require.NoError(t, err)
 	require.Zero(t, results[0]) // should succeed and return the old size in pages.
 
+	// Any offset larger than the current size should be out of of bounds error even when it is less than memory capacity.
+	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8)
+	require.Error(t, err) // Out of bounds error.
+
 	// Check the size command works!
 	results, err = memory.ExportedFunction("size").Call(testCtx)
 	require.NoError(t, err)
 	require.Equal(t, uint64(1), results[0])                        // 1 page
 	require.Equal(t, uint32(65536), memory.Memory().Size(testCtx)) // 64KB
+
+	// Grow again so that the memory size matches memory capacity.
+	results, err = memory.ExportedFunction("grow").Call(testCtx, 1)
+	require.NoError(t, err)
+	require.Equal(t, uint64(1), results[0])
+
+	// Verify the size matches cap.
+	results, err = memory.ExportedFunction("size").Call(testCtx)
+	require.NoError(t, err)
+	require.Equal(t, uint64(memoryCapacityPages), results[0])
+
+	// Now the store instruction at the memory capcity bound should suceed.
+	_, err = memory.ExportedFunction("store").Call(testCtx, wasm.MemoryPagesToBytesNum(memoryCapacityPages)-8) // i64.store needs 8 bytes from offset.
+	require.NoError(t, err)
 }
 
 func testMultipleInstantiation(t *testing.T, r wazero.Runtime) {
diff --git a/internal/testing/enginetest/enginetest.go b/internal/testing/enginetest/enginetest.go
index 1ee2181e..f2efc7f8 100644
--- a/internal/testing/enginetest/enginetest.go
+++ b/internal/testing/enginetest/enginetest.go
@@ -506,6 +506,121 @@ wasm stack trace:
 	}
 }
 
+// RunTestModuleEngine_Memory shows that the byte slice returned from api.Memory Read is not a copy, rather a re-slice
+// of the underlying memory. This allows both host and Wasm to see each other's writes, unless one side changes the
+// capacity of the slice.
+//
+// Known cases that change the slice capacity:
+// * Host code calls append on a byte slice returned by api.Memory Read
+// * Wasm code calls wasm.OpcodeMemoryGrowName and this changes the capacity (by default, it will).
+func RunTestModuleEngine_Memory(t *testing.T, et EngineTester) {
+	e := et.NewEngine(wasm.Features20220419)
+
+	wasmPhrase := "Well, that'll be the day when you say goodbye."
+	wasmPhraseSize := uint32(len(wasmPhrase))
+
+	// Define a basic function which defines one parameter. This is used to test results when incorrect arity is used.
+	one := uint32(1)
+	m := &wasm.Module{
+		TypeSection:     []*wasm.FunctionType{{Params: []api.ValueType{api.ValueTypeI32}}, {}},
+		FunctionSection: []wasm.Index{0, 1},
+		MemorySection:   &wasm.Memory{Min: 1, Cap: 1, Max: 2},
+		DataSection: []*wasm.DataSegment{
+			{
+				OffsetExpression: nil, // passive
+				Init:             []byte(wasmPhrase),
+			},
+		},
+		DataCountSection: &one,
+		CodeSection: []*wasm.Code{
+			{Body: []byte{ // "grow"
+				wasm.OpcodeLocalGet, 0, // how many pages to grow (param)
+				wasm.OpcodeMemoryGrow, 0, // memory index zero
+				wasm.OpcodeDrop, // drop the previous page count (or -1 if grow failed)
+				wasm.OpcodeEnd,
+			}},
+			{Body: []byte{ // "init"
+				wasm.OpcodeI32Const, 0, // target offset
+				wasm.OpcodeI32Const, 0, // source offset
+				wasm.OpcodeI32Const, byte(wasmPhraseSize), // len
+				wasm.OpcodeMiscPrefix, wasm.OpcodeMiscMemoryInit, 0, 0, // segment 0, memory 0
+				wasm.OpcodeEnd,
+			}},
+		},
+	}
+	// Compile the Wasm into wazeroir
+	err := e.CompileModule(testCtx, m)
+	require.NoError(t, err)
+
+	// Assign memory to the module instance
+	module := &wasm.ModuleInstance{
+		Name:          t.Name(),
+		Memory:        wasm.NewMemoryInstance(m.MemorySection),
+		DataInstances: []wasm.DataInstance{m.DataSection[0].Init},
+	}
+	var memory api.Memory = module.Memory
+
+	// To use functions, we need to instantiate them (associate them with a ModuleInstance).
+	grow := getFunctionInstance(m, 0, module)
+	addFunction(module, "grow", grow)
+	init := getFunctionInstance(m, 1, module)
+	addFunction(module, "init", init)
+
+	// Compile the module
+	me, err := e.NewModuleEngine(module.Name, m, nil, module.Functions, nil, nil)
+	init.Module.Engine = me
+	require.NoError(t, err)
+	linkModuleToEngine(module, me)
+
+	buf, ok := memory.Read(testCtx, 0, wasmPhraseSize)
+	require.True(t, ok)
+	require.Equal(t, make([]byte, wasmPhraseSize), buf)
+
+	// Initialize the memory using Wasm. This copies the test phrase.
+	_, err = me.Call(testCtx, module.CallCtx, init)
+	require.NoError(t, err)
+
+	// We expect the same []byte read earlier to now include the phrase in wasm.
+	require.Equal(t, wasmPhrase, string(buf))
+
+	hostPhrase := "Goodbye, cruel world. I'm off to join the circus." // Intentionally slightly longer.
+	hostPhraseSize := uint32(len(hostPhrase))
+
+	// Copy over the buffer, which should stop at the current length.
+	copy(buf, hostPhrase)
+	require.Equal(t, "Goodbye, cruel world. I'm off to join the circ", string(buf))
+
+	// The underlying memory should be updated. This proves that Memory.Read returns a re-slice, not a copy, and that
+	// programs can rely on this (for example, to update shared state in Wasm and view that in Go and visa versa).
+	buf2, ok := memory.Read(testCtx, 0, wasmPhraseSize)
+	require.True(t, ok)
+	require.Equal(t, buf, buf2)
+
+	// Now, append to the buffer we got from Wasm. As this changes capacity, it should result in a new byte slice.
+	buf = append(buf, 'u', 's', '.')
+	require.Equal(t, hostPhrase, string(buf))
+
+	// To prove the above, we re-read the memory and should not see the appended bytes (rather zeros instead).
+	buf2, ok = memory.Read(testCtx, 0, hostPhraseSize)
+	require.True(t, ok)
+	hostPhraseTruncated := "Goodbye, cruel world. I'm off to join the circ" + string([]byte{0, 0, 0})
+	require.Equal(t, hostPhraseTruncated, string(buf2))
+
+	// Now, we need to prove the other direction, that when Wasm changes the capacity, the host's buffer is unaffected.
+	_, err = me.Call(testCtx, module.CallCtx, grow, 1)
+	require.NoError(t, err)
+
+	// The host buffer should still contain the same bytes as before grow
+	require.Equal(t, hostPhraseTruncated, string(buf2))
+
+	// Re-initialize the memory in wasm, which overwrites the region.
+	_, err = me.Call(testCtx, module.CallCtx, init)
+	require.NoError(t, err)
+
+	// The host was not affected because it is a different slice due to "memory.grow" affecting the underlying memory.
+	require.Equal(t, hostPhraseTruncated, string(buf2))
+}
+
 const (
 	wasmFnName               = "wasm_div_by"
 	hostFnName               = "host_div_by"
diff --git a/internal/wasm/interpreter/interpreter_test.go b/internal/wasm/interpreter/interpreter_test.go
index dad6b7fa..20670c1c 100644
--- a/internal/wasm/interpreter/interpreter_test.go
+++ b/internal/wasm/interpreter/interpreter_test.go
@@ -98,6 +98,10 @@ func TestInterpreter_ModuleEngine_Call_Errors(t *testing.T) {
 	enginetest.RunTestModuleEngine_Call_Errors(t, et)
 }
 
+func TestInterpreter_ModuleEngine_Memory(t *testing.T) {
+	enginetest.RunTestModuleEngine_Memory(t, et)
+}
+
 func TestInterpreter_NonTrappingFloatToIntConversion(t *testing.T) {
 	_0x80000000 := uint32(0x80000000)
 	_0xffffffff := uint32(0xffffffff)
diff --git a/internal/wasm/jit/engine_test.go b/internal/wasm/jit/engine_test.go
index d3bd389e..f163838b 100644
--- a/internal/wasm/jit/engine_test.go
+++ b/internal/wasm/jit/engine_test.go
@@ -165,6 +165,11 @@ func TestJIT_ModuleEngine_Call_Errors(t *testing.T) {
 	enginetest.RunTestModuleEngine_Call_Errors(t, et)
 }
 
+func TestJIT_ModuleEngine_Memory(t *testing.T) {
+	requireSupportedOSArch(t)
+	enginetest.RunTestModuleEngine_Memory(t, et)
+}
+
 func requireSupportedOSArch(t *testing.T) {
 	if runtime.GOARCH != "amd64" && runtime.GOARCH != "arm64" {
 		t.Skip()