wasm,wasi: make sure buffers returned by malloc are not freed until f… (#3148)

* wasm,wasi: make sure buffers returned by malloc are not freed until free is called

Makefile

@@ -395,9 +395,9 @@ tinygo-bench-fast:
 
 # Same thing, except for wasi rather than the current platform.
 tinygo-test-wasi:
-	$(TINYGO) test -target wasi $(TEST_PACKAGES_FAST) $(TEST_PACKAGES_SLOW)
+	$(TINYGO) test -target wasi $(TEST_PACKAGES_FAST) $(TEST_PACKAGES_SLOW) ./tests/runtime_wasi
 tinygo-test-wasi-fast:
-	$(TINYGO) test -target wasi $(TEST_PACKAGES_FAST)
+	$(TINYGO) test -target wasi $(TEST_PACKAGES_FAST) ./tests/runtime_wasi
 tinygo-bench-wasi:
 	$(TINYGO) test -target wasi -bench . $(TEST_PACKAGES_FAST) $(TEST_PACKAGES_SLOW)
 tinygo-bench-wasi-fast:

src/runtime/arch_tinygowasm.go

@@ -61,31 +61,54 @@ func growHeap() bool {
 	return true
 }
 
-// The below functions override the default allocator of wasi-libc.
-// Most functions are defined but unimplemented to make sure that if there is
-// any code using them, they will get an error instead of (incorrectly) using
-// the wasi-libc dlmalloc heap implementation instead. If they are needed by any
-// program, they can certainly be implemented.
+// The below functions override the default allocator of wasi-libc. This ensures
+// code linked from other languages can allocate memory without colliding with
+// our GC allocations.
+
+var allocs = make(map[uintptr][]byte)
 
 //export malloc
 func libc_malloc(size uintptr) unsafe.Pointer {
-	return alloc(size, nil)
+	buf := make([]byte, size)
+	ptr := unsafe.Pointer(&buf[0])
+	allocs[uintptr(ptr)] = buf
+	return ptr
 }
 
 //export free
 func libc_free(ptr unsafe.Pointer) {
-	free(ptr)
+	if ptr == nil {
+		return
+	}
+	if _, ok := allocs[uintptr(ptr)]; ok {
+		delete(allocs, uintptr(ptr))
+	} else {
+		panic("free: invalid pointer")
+	}
 }
 
 //export calloc
 func libc_calloc(nmemb, size uintptr) unsafe.Pointer {
-	// Note: we could be even more correct here and check that nmemb * size
-	// doesn't overflow. However the current implementation should normally work
-	// fine.
-	return alloc(nmemb*size, nil)
+	// No difference between calloc and malloc.
+	return libc_malloc(nmemb * size)
 }
 
 //export realloc
-func libc_realloc(ptr unsafe.Pointer, size uintptr) unsafe.Pointer {
-	return realloc(ptr, size)
+func libc_realloc(oldPtr unsafe.Pointer, size uintptr) unsafe.Pointer {
+	// It's hard to optimize this to expand the current buffer with our GC, but
+	// it is theoretically possible. For now, just always allocate fresh.
+	buf := make([]byte, size)
+
+	if oldPtr != nil {
+		if oldBuf, ok := allocs[uintptr(oldPtr)]; ok {
+			copy(buf, oldBuf)
+			delete(allocs, uintptr(oldPtr))
+		} else {
+			panic("realloc: invalid pointer")
+		}
+	}
+
+	ptr := unsafe.Pointer(&buf[0])
+	allocs[uintptr(ptr)] = buf
+	return ptr
 }

tests/runtime_wasi/malloc_test.go

@@ -0,0 +1,127 @@
+//go:build tinygo.wasm
+// +build tinygo.wasm
+
+package runtime_wasi
+
+import (
+	"reflect"
+	"runtime"
+	"strconv"
+	"testing"
+	"unsafe"
+)
+
+//export malloc
+func libc_malloc(size uintptr) unsafe.Pointer
+
+//export free
+func libc_free(ptr unsafe.Pointer)
+
+//export calloc
+func libc_calloc(nmemb, size uintptr) unsafe.Pointer
+
+//export realloc
+func libc_realloc(ptr unsafe.Pointer, size uintptr) unsafe.Pointer
+
+func getFilledBuffer_malloc() uintptr {
+	ptr := libc_malloc(5)
+	fillPanda(ptr)
+	return uintptr(ptr)
+}
+
+func getFilledBuffer_calloc() uintptr {
+	ptr := libc_calloc(2, 5)
+	fillPanda(ptr)
+	*(*byte)(unsafe.Add(ptr, 5)) = 'b'
+	*(*byte)(unsafe.Add(ptr, 6)) = 'e'
+	*(*byte)(unsafe.Add(ptr, 7)) = 'a'
+	*(*byte)(unsafe.Add(ptr, 8)) = 'r'
+	*(*byte)(unsafe.Add(ptr, 9)) = 's'
+	return uintptr(ptr)
+}
+
+func getFilledBuffer_realloc() uintptr {
+	origPtr := getFilledBuffer_malloc()
+	ptr := libc_realloc(unsafe.Pointer(origPtr), 9)
+	*(*byte)(unsafe.Add(ptr, 5)) = 'b'
+	*(*byte)(unsafe.Add(ptr, 6)) = 'e'
+	*(*byte)(unsafe.Add(ptr, 7)) = 'a'
+	*(*byte)(unsafe.Add(ptr, 8)) = 'r'
+	return uintptr(ptr)
+}
+
+func getFilledBuffer_reallocNil() uintptr {
+	ptr := libc_realloc(nil, 5)
+	fillPanda(ptr)
+	return uintptr(ptr)
+}
+
+func fillPanda(ptr unsafe.Pointer) {
+	*(*byte)(unsafe.Add(ptr, 0)) = 'p'
+	*(*byte)(unsafe.Add(ptr, 1)) = 'a'
+	*(*byte)(unsafe.Add(ptr, 2)) = 'n'
+	*(*byte)(unsafe.Add(ptr, 3)) = 'd'
+	*(*byte)(unsafe.Add(ptr, 4)) = 'a'
+}
+
+func checkFilledBuffer(t *testing.T, ptr uintptr, content string) {
+	t.Helper()
+	buf := *(*string)(unsafe.Pointer(&reflect.StringHeader{
+		Data: ptr,
+		Len:  uintptr(len(content)),
+	}))
+	if buf != content {
+		t.Errorf("expected %q, got %q", content, buf)
+	}
+}
+
+func TestMallocFree(t *testing.T) {
+	tests := []struct {
+		name      string
+		getBuffer func() uintptr
+		content   string
+	}{
+		{
+			name:      "malloc",
+			getBuffer: getFilledBuffer_malloc,
+			content:   "panda",
+		},
+		{
+			name:      "calloc",
+			getBuffer: getFilledBuffer_calloc,
+			content:   "pandabears",
+		},
+		{
+			name:      "realloc",
+			getBuffer: getFilledBuffer_realloc,
+			content:   "pandabear",
+		},
+		{
+			name:      "realloc nil",
+			getBuffer: getFilledBuffer_reallocNil,
+			content:   "panda",
+		},
+	}
+
+	for _, tc := range tests {
+		tt := tc
+		t.Run(tt.name, func(t *testing.T) {
+			bufPtr := tt.getBuffer()
+			// Don't use defer to free the buffer as it seems to cause the GC to track it.
+
+			// Churn GC, the pointer should still be valid until free is called.
+			for i := 0; i < 1000; i++ {
+				a := "hello" + strconv.Itoa(i)
+				// Some conditional logic to ensure optimization doesn't remove the loop completely.
+				if len(a) < 0 {
+					break
+				}
+				runtime.GC()
+			}
+
+			checkFilledBuffer(t, bufPtr, tt.content)
+
+			libc_free(unsafe.Pointer(bufPtr))
+		})
+	}
+}