transform (gc): create stack slots in callers of external functions

This updates the stack slot pass to include callers of external functions which may access non-argument memory.
Wiithout this change, a use-after-free could occur on WASM when calling a reentrant function or switching to another goroutine.

transform/gc.go

@@ -4,6 +4,11 @@ import (
 	"tinygo.org/x/go-llvm"
 )
 
+// This is somewhat ugly to access through the API.
+// https://github.com/llvm/llvm-project/blob/94ebcfd16dac67486bae624f74e1c5c789448bae/llvm/include/llvm/Support/ModRef.h#L62
+// https://github.com/llvm/llvm-project/blob/94ebcfd16dac67486bae624f74e1c5c789448bae/llvm/include/llvm/Support/ModRef.h#L87
+const shiftExcludeArgMem = 2
+
 // MakeGCStackSlots converts all calls to runtime.trackPointer to explicit
 // stores to stack slots that are scannable by the GC.
 func MakeGCStackSlots(mod llvm.Module) bool {
@@ -36,20 +41,49 @@ func MakeGCStackSlots(mod llvm.Module) bool {
 	defer targetData.Dispose()
 	uintptrType := ctx.IntType(targetData.PointerSize() * 8)
 
-	// Look at *all* functions to see whether they are free of function pointer
+	// All functions that call runtime.alloc needs stack objects.
+	trackFuncs := map[llvm.Value]struct{}{}
+	markParentFunctions(trackFuncs, alloc)
+
+	// External functions may indirectly suspend the goroutine or perform a heap allocation.
+	// Their callers should get stack objects.
+	memAttr := llvm.AttributeKindID("memory")
+	for fn := mod.FirstFunction(); !fn.IsNil(); fn = llvm.NextFunction(fn) {
+		if _, ok := trackFuncs[fn]; ok {
+			continue // already found
+		}
+		if !fn.FirstBasicBlock().IsNil() {
+			// This is not an external function.
+			continue
+		}
+		if fn == trackPointer {
+			// Manually exclude trackPointer.
+			continue
+		}
+
+		mem := fn.GetEnumFunctionAttribute(memAttr)
+		if !mem.IsNil() && mem.GetEnumValue()>>shiftExcludeArgMem == 0 {
+			// This does not access non-argument memory.
+			// Exclude it.
+			continue
+		}
+
+		// The callers need stack objects.
+		markParentFunctions(trackFuncs, fn)
+	}
+
+	// Look at all other functions to see whether they contain function pointer
 	// calls.
 	// This takes less than 5ms for ~100kB of WebAssembly but would perhaps be
 	// faster when written in C++ (to avoid the CGo overhead).
-	funcsWithFPCall := map[llvm.Value]struct{}{}
-	n := 0
 	for fn := mod.FirstFunction(); !fn.IsNil(); fn = llvm.NextFunction(fn) {
-		n++
+		if _, ok := trackFuncs[fn]; ok {
-		if _, ok := funcsWithFPCall[fn]; ok {
 			continue // already found
 		}
-		done := false
+
-		for bb := fn.FirstBasicBlock(); !bb.IsNil() && !done; bb = llvm.NextBasicBlock(bb) {
+	scanBody:
-			for call := bb.FirstInstruction(); !call.IsNil() && !done; call = llvm.NextInstruction(call) {
+		for bb := fn.FirstBasicBlock(); !bb.IsNil(); bb = llvm.NextBasicBlock(bb) {
+			for call := bb.FirstInstruction(); !call.IsNil(); call = llvm.NextInstruction(call) {
 				if call.IsACallInst().IsNil() {
 					continue // only looking at calls
 				}
@@ -57,33 +91,13 @@ func MakeGCStackSlots(mod llvm.Module) bool {
 				if !called.IsAFunction().IsNil() {
 					continue // only looking for function pointers
 				}
-				funcsWithFPCall[fn] = struct{}{}
+				trackFuncs[fn] = struct{}{}
-				markParentFunctions(funcsWithFPCall, fn)
+				markParentFunctions(trackFuncs, fn)
-				done = true
+				break scanBody
 			}
 		}
 	}
 
-	// Determine which functions need stack objects. Many leaf functions don't
-	// need it: it only causes overhead for them.
-	// Actually, in one test it was only able to eliminate stack object from 12%
-	// of functions that had a call to runtime.trackPointer (8 out of 68
-	// functions), so this optimization is not as big as it may seem.
-	allocatingFunctions := map[llvm.Value]struct{}{} // set of allocating functions
-
-	// Work from runtime.alloc and trace all parents to check which functions do
-	// a heap allocation (and thus which functions do not).
-	markParentFunctions(allocatingFunctions, alloc)
-
-	// Also trace all functions that call a function pointer.
-	for fn := range funcsWithFPCall {
-		// Assume that functions that call a function pointer do a heap
-		// allocation as a conservative guess because the called function might
-		// do a heap allocation.
-		allocatingFunctions[fn] = struct{}{}
-		markParentFunctions(allocatingFunctions, fn)
-	}
-
 	// Collect some variables used below in the loop.
 	stackChainStart := mod.NamedGlobal("runtime.stackChainStart")
 	if stackChainStart.IsNil() {
@@ -110,7 +124,7 @@ func MakeGCStackSlots(mod llvm.Module) bool {
 		// Pick the parent function.
 		fn := call.InstructionParent().Parent()
 
-		if _, ok := allocatingFunctions[fn]; !ok {
+		if _, ok := trackFuncs[fn]; !ok {
 			// This function nor any of the functions it calls (recursively)
 			// allocate anything from the heap, so it will not trigger a garbage
 			// collection cycle. Thus, it does not need to track local pointer

transform/testdata/gc-stackslots.ll

@@ -4,6 +4,7 @@ target triple = "wasm32-unknown-unknown-wasm"
 @runtime.stackChainStart = external global ptr
 @someGlobal = global i8 3
 @ptrGlobal = global ptr null
+@arrGlobal = global [8 x i8] zeroinitializer
 
 declare void @runtime.trackPointer(ptr nocapture readonly)
 
@@ -116,3 +117,51 @@ define void @allocAndSave(ptr %x) {
   store ptr %x, ptr @ptrGlobal
   ret void
 }
+
+declare void @"(internal/task).Pause"()
+
+define ptr @getAndPause() {
+	%ptr = call ptr @getPointer()
+	call void @runtime.trackPointer(ptr %ptr)
+	; Calling a function with unknown memory access forces stack slot creation.
+	call void @"(internal/task).Pause"()
+	ret ptr %ptr
+}
+
+; Function Attrs: memory(readwrite)
+declare void @externCallWithMemAttr() #0
+
+define ptr @getAndCallWithMemAttr() {
+	%ptr = call ptr @getPointer()
+	call void @runtime.trackPointer(ptr %ptr)
+	; Calling an external function which may access non-arg memory forces stack slot creation.
+	call void @externCallWithMemAttr()
+	ret ptr %ptr
+}
+
+; Generic function that returns a slice (that must be tracked).
+define {ptr, i32, i32} @getSlice() {
+  ret {ptr, i32, i32} {ptr @someGlobal, i32 8, i32 8}
+}
+
+define i32 @copyToSlice(ptr %src.ptr, i32 %src.len, i32 %src.cap) {
+  %dst = call {ptr, i32, i32} @getSlice()
+  %dst.ptr = extractvalue {ptr, i32, i32} %dst, 0
+  call void @runtime.trackPointer(ptr %dst.ptr)
+  %dst.len = extractvalue {ptr, i32, i32} %dst, 1
+  ; Math intrinsics do not need stack slots.
+  %minLen = call i32 @llvm.umin.i32(i32 %dst.len, i32 %src.len)
+  ; Intrinsics which only access argument memory do not need stack slots.
+  call void @llvm.memmove.p0.p0.i32(ptr %dst.ptr, ptr %src.ptr, i32 %minLen, i1 false)
+  ret i32 %minLen
+}
+
+; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
+declare i32 @llvm.umin.i32(i32, i32) #1
+
+; Function Attrs: nocallback nofree nounwind willreturn memory(argmem: readwrite)
+declare void @llvm.memmove.p0.p0.i32(ptr nocapture writeonly, ptr nocapture readonly, i32, i1 immarg) #2
+
+attributes #0 = { memory(readwrite) }
+attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+attributes #2 = { nocallback nofree nounwind willreturn memory(argmem: readwrite) }

transform/testdata/gc-stackslots.out.ll

@@ -4,6 +4,7 @@ target triple = "wasm32-unknown-unknown-wasm"
 @runtime.stackChainStart = internal global ptr null
 @someGlobal = global i8 3
 @ptrGlobal = global ptr null
+@arrGlobal = global [8 x i8] zeroinitializer
 
 declare void @runtime.trackPointer(ptr nocapture readonly)
 
@@ -166,3 +167,61 @@ define void @allocAndSave(ptr %x) {
   store ptr %1, ptr @runtime.stackChainStart, align 4
   ret void
 }
+
+declare void @"(internal/task).Pause"()
+
+define ptr @getAndPause() {
+  %gc.stackobject = alloca { ptr, i32, ptr }, align 8
+  store { ptr, i32, ptr } { ptr null, i32 1, ptr null }, ptr %gc.stackobject, align 4
+  %1 = load ptr, ptr @runtime.stackChainStart, align 4
+  %2 = getelementptr { ptr, i32, ptr }, ptr %gc.stackobject, i32 0, i32 0
+  store ptr %1, ptr %2, align 4
+  store ptr %gc.stackobject, ptr @runtime.stackChainStart, align 4
+  %ptr = call ptr @getPointer()
+  %3 = getelementptr { ptr, i32, ptr }, ptr %gc.stackobject, i32 0, i32 2
+  store ptr %ptr, ptr %3, align 4
+  call void @"(internal/task).Pause"()
+  store ptr %1, ptr @runtime.stackChainStart, align 4
+  ret ptr %ptr
+}
+
+; Function Attrs: memory(readwrite)
+declare void @externCallWithMemAttr() #0
+
+define ptr @getAndCallWithMemAttr() {
+  %gc.stackobject = alloca { ptr, i32, ptr }, align 8
+  store { ptr, i32, ptr } { ptr null, i32 1, ptr null }, ptr %gc.stackobject, align 4
+  %1 = load ptr, ptr @runtime.stackChainStart, align 4
+  %2 = getelementptr { ptr, i32, ptr }, ptr %gc.stackobject, i32 0, i32 0
+  store ptr %1, ptr %2, align 4
+  store ptr %gc.stackobject, ptr @runtime.stackChainStart, align 4
+  %ptr = call ptr @getPointer()
+  %3 = getelementptr { ptr, i32, ptr }, ptr %gc.stackobject, i32 0, i32 2
+  store ptr %ptr, ptr %3, align 4
+  call void @externCallWithMemAttr()
+  store ptr %1, ptr @runtime.stackChainStart, align 4
+  ret ptr %ptr
+}
+
+define { ptr, i32, i32 } @getSlice() {
+  ret { ptr, i32, i32 } { ptr @someGlobal, i32 8, i32 8 }
+}
+
+define i32 @copyToSlice(ptr %src.ptr, i32 %src.len, i32 %src.cap) {
+  %dst = call { ptr, i32, i32 } @getSlice()
+  %dst.ptr = extractvalue { ptr, i32, i32 } %dst, 0
+  %dst.len = extractvalue { ptr, i32, i32 } %dst, 1
+  %minLen = call i32 @llvm.umin.i32(i32 %dst.len, i32 %src.len)
+  call void @llvm.memmove.p0.p0.i32(ptr %dst.ptr, ptr %src.ptr, i32 %minLen, i1 false)
+  ret i32 %minLen
+}
+
+; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
+declare i32 @llvm.umin.i32(i32, i32) #1
+
+; Function Attrs: nocallback nofree nounwind willreturn memory(argmem: readwrite)
+declare void @llvm.memmove.p0.p0.i32(ptr nocapture writeonly, ptr nocapture readonly, i32, i1 immarg) #2
+
+attributes #0 = { memory(readwrite) }
+attributes #1 = { nocallback nofree nosync nounwind speculatable willreturn memory(none) }
+attributes #2 = { nocallback nofree nounwind willreturn memory(argmem: readwrite) }