Repository: tetratelabs/wazero
Pull request #2189 (2024-04-17 02:43:05)

experimental: reflect custom allocated buffer capacity to internal buffer
/cc @ncruces 

When trying to use the custom allocator interface, I was getting crashes since the capacity wasn't being applied to the internal `m.Buffer` itself

```
 [recovered]
	panic: runtime error: slice bounds out of range [::266304] with capacity 196608 (recovered by wazero)
wasm stack trace:
	wasi_snapshot_preview1.fd_write(i32,i32,i32,i32) i32
	.$21(i32,i32,i32) i32
	.$22(i32,i32,i32) i32
	.$31(i32,i32,i32)
	.$787(i32)
	.$783(i32,i32,i32)
	.$806(i32,i32,i32) i32

Go runtime stack trace:
goroutine 37 [running]:
runtime/debug.Stack()
	/opt/homebrew/Cellar/go/1.22.2/libexec/src/runtime/debug/stack.go:24 +0x6c
github.com/tetratelabs/wazero/internal/wasmdebug.(*stackTrace).FromRecovered(0xc000a0c798, {0x101482be0, 0xc000212048})
	/Users/anuraag/go/pkg/mod/github.com/tetratelabs/wazero@v1.7.1/internal/wasmdebug/debug.go:139 +0x1a8
github.com/tetratelabs/wazero/internal/engine/wazevo.(*callEngine).callWithStack.func1()
	/Users/anuraag/go/pkg/mod/github.com/tetratelabs/wazero@v1.7.1/internal/engine/wazevo/call_engine.go:256 +0x4b0
panic({0x101482be0?, 0xc000212048?})
	/opt/homebrew/Cellar/go/1.22.2/libexec/src/runtime/panic.go:770 +0x124
github.com/tetratelabs/wazero/internal/wasm.(*MemoryInstance).Read(0xc00013c7e0, 0x41030, 0x10)
	/Users/anuraag/go/pkg/mod/github.com/tetratelabs/wazero@v1.7.1/internal/wasm/memory.go:160 +0x104
github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1.writev({0x1014aedd0, 0xc00013c7e0}, 0x41030, 0x2, 0xc00002ece8)
```


Patch:
diff --git a/internal/wasm/memory.go b/internal/wasm/memory.go
index ac2ee2ed..5cc5012d 100644
--- a/internal/wasm/memory.go
+++ b/internal/wasm/memory.go
@@ -241,7 +241,7 @@ func (m *MemoryInstance) Grow(delta uint32) (result uint32, ok bool) {
 			// We assume grow is called under a guest lock.
 			// But the memory length is accessed elsewhere,
 			// so use atomic to make the new length visible across threads.
-			atomicStoreLength(&m.Buffer, uintptr(len(buffer)))
+			atomicStoreLengthAndCap(&m.Buffer, uintptr(len(buffer)), uintptr(cap(buffer)))
 			m.Cap = memoryBytesNumToPages(uint64(cap(buffer)))
 		} else {
 			m.Buffer = buffer
@@ -294,6 +294,15 @@ func PagesToUnitOfBytes(pages uint32) string {
 
 // Below are raw functions used to implement the api.Memory API:
 
+// Uses atomic write to update the length of a slice.
+func atomicStoreLengthAndCap(slice *[]byte, length uintptr, cap uintptr) {
+	slicePtr := (*reflect.SliceHeader)(unsafe.Pointer(slice))
+	capPtr := (*uintptr)(unsafe.Pointer(&slicePtr.Cap))
+	atomic.StoreUintptr(capPtr, cap)
+	lenPtr := (*uintptr)(unsafe.Pointer(&slicePtr.Len))
+	atomic.StoreUintptr(lenPtr, length)
+}
+
 // Uses atomic write to update the length of a slice.
 func atomicStoreLength(slice *[]byte, length uintptr) {
 	slicePtr := (*reflect.SliceHeader)(unsafe.Pointer(slice))


Test patch:
diff --git a/internal/wasm/memory_test.go b/internal/wasm/memory_test.go
index bba90513..542591f5 100644
--- a/internal/wasm/memory_test.go
+++ b/internal/wasm/memory_test.go
@@ -828,6 +828,14 @@ func Test_atomicStoreLength(t *testing.T) {
 	require.Equal(t, 15, len(slice))
 }
 
+func Test_atomicStoreLengthAndCap(t *testing.T) {
+	// Doesn't verify atomicity, but at least we're updating the correct thing.
+	slice := make([]byte, 10, 20)
+	atomicStoreLengthAndCap(&slice, 12, 18)
+	require.Equal(t, 12, len(slice))
+	require.Equal(t, 18, cap(slice))
+}
+
 func TestNewMemoryInstance_Shared(t *testing.T) {
 	tests := []struct {
 		name string