cmd/cgo: add #cgo noescape/nocallback annotations

When passing pointers of Go objects from Go to C, the cgo command generate
_Cgo_use(pN) for the unsafe.Pointer type arguments, so that the Go compiler
will escape these object to heap.

Since the C function may callback to Go, then the Go stack might grow/shrink,
that means the pointers that the C function have will be invalid.

After adding the #cgo noescape/nocallback annotations for a C function,
the cgo command won't generate _Cgo_use(pN), and the Go compiler won't
force the object escape to heap, that means the C function won't callback
to Go, if it do callback to Go, the Go process will crash.

Fixes #56378

Author: doujiang24

src/cmd/cgo/gcc.go

@@ -24,6 +24,7 @@ import (
 	"math"
 	"os"
 	"os/exec"
+	"regexp"
 	"strconv"
 	"strings"
 	"unicode"
@@ -49,6 +50,8 @@ var nameToC = map[string]string{
 
 var incomplete = "_cgopackage.Incomplete"
 
+var cgoRx = regexp.MustCompile(`^#cgo\s+(nocallback|noescape)\s+(\S+)\s*$`)
+
 // cname returns the C name to use for C.s.
 // The expansions are listed in nameToC and also
 // struct_foo becomes "struct foo", and similarly for
@@ -73,18 +76,30 @@ func cname(s string) string {
 	return s
 }
 
-// DiscardCgoDirectives processes the import C preamble, and discards
-// all #cgo CFLAGS and LDFLAGS directives, so they don't make their
-// way into _cgo_export.h.
-func (f *File) DiscardCgoDirectives() {
+// ProcessCgoDirectives processes the import C preamble:
+//  1. discards all #cgo CFLAGS and LDFLAGS directives, so they don't make their
+//     way into _cgo_export.h.
+//  2. parse the nocallback and noescape directives.
+func (f *File) ProcessCgoDirectives() {
 	linesIn := strings.Split(f.Preamble, "\n")
 	linesOut := make([]string, 0, len(linesIn))
+	f.NoCallbacks = make(map[string]struct{})
+	f.NoEscapes = make(map[string]struct{})
 	for _, line := range linesIn {
 		l := strings.TrimSpace(line)
 		if len(l) < 5 || l[:4] != "#cgo" || !unicode.IsSpace(rune(l[4])) {
 			linesOut = append(linesOut, line)
 		} else {
 			linesOut = append(linesOut, "")
+
+			match := cgoRx.FindStringSubmatch(line)
+			if match != nil {
+				if match[1] == "nocallback" {
+					f.NoCallbacks[match[2]] = struct{}{}
+				} else {
+					f.NoEscapes[match[2]] = struct{}{}
+				}
+			}
 		}
 	}
 	f.Preamble = strings.Join(linesOut, "\n")

src/cmd/cgo/main.go

@@ -48,6 +48,8 @@ type Package struct {
 	Preamble    string          // collected preamble for _cgo_export.h
 	typedefs    map[string]bool // type names that appear in the types of the objects we're interested in
 	typedefList []typedefInfo
+	noCallbacks map[string]struct{} // C function names that with #cgo nocallback directive
+	noEscapes   map[string]struct{} // C function names that with #cgo noescape directive
 }
 
 // A typedefInfo is an element on Package.typedefList: a typedef name
@@ -59,16 +61,18 @@ type typedefInfo struct {
 
 // A File collects information about a single Go input file.
 type File struct {
-	AST      *ast.File           // parsed AST
-	Comments []*ast.CommentGroup // comments from file
-	Package  string              // Package name
-	Preamble string              // C preamble (doc comment on import "C")
-	Ref      []*Ref              // all references to C.xxx in AST
-	Calls    []*Call             // all calls to C.xxx in AST
-	ExpFunc  []*ExpFunc          // exported functions for this file
-	Name     map[string]*Name    // map from Go name to Name
-	NamePos  map[*Name]token.Pos // map from Name to position of the first reference
-	Edit     *edit.Buffer
+	AST         *ast.File           // parsed AST
+	Comments    []*ast.CommentGroup // comments from file
+	Package     string              // Package name
+	Preamble    string              // C preamble (doc comment on import "C")
+	Ref         []*Ref              // all references to C.xxx in AST
+	Calls       []*Call             // all calls to C.xxx in AST
+	ExpFunc     []*ExpFunc          // exported functions for this file
+	Name        map[string]*Name    // map from Go name to Name
+	NamePos     map[*Name]token.Pos // map from Name to position of the first reference
+	Edit        *edit.Buffer
+	NoCallbacks map[string]struct{} // C function names that with #cgo nocallback directive
+	NoEscapes   map[string]struct{} // C function names that with #cgo noescape directive
 }
 
 func (f *File) offset(p token.Pos) int {
@@ -368,7 +372,7 @@ func main() {
 		f := new(File)
 		f.Edit = edit.NewBuffer(b)
 		f.ParseGo(input, b)
-		f.DiscardCgoDirectives()
+		f.ProcessCgoDirectives()
 		fs[i] = f
 	}
 
@@ -481,6 +485,22 @@ func (p *Package) Record(f *File) {
 		}
 	}
 
+	// merge nocallback & noescape
+	if p.noCallbacks == nil {
+		p.noCallbacks = f.NoCallbacks
+	} else {
+		for k, v := range f.NoCallbacks {
+			p.noCallbacks[k] = v
+		}
+	}
+	if p.noEscapes == nil {
+		p.noEscapes = f.NoEscapes
+	} else {
+		for k, v := range f.NoEscapes {
+			p.noEscapes[k] = v
+		}
+	}
+
 	if f.ExpFunc != nil {
 		p.ExpFunc = append(p.ExpFunc, f.ExpFunc...)
 		p.Preamble += "\n" + f.Preamble

src/cmd/cgo/out.go

@@ -103,6 +103,8 @@ func (p *Package) writeDefs() {
 		fmt.Fprintf(fgo2, "var _Cgo_always_false bool\n")
 		fmt.Fprintf(fgo2, "//go:linkname _Cgo_use runtime.cgoUse\n")
 		fmt.Fprintf(fgo2, "func _Cgo_use(interface{})\n")
+		fmt.Fprintf(fgo2, "//go:linkname _Cgo_no_callback runtime.cgoNoCallback\n")
+		fmt.Fprintf(fgo2, "func _Cgo_no_callback(bool)\n")
 	}
 
 	typedefNames := make([]string, 0, len(typedef))
@@ -610,6 +612,17 @@ func (p *Package) writeDefsFunc(fgo2 io.Writer, n *Name, callsMalloc *bool) {
 		arg = "uintptr(unsafe.Pointer(&r1))"
 	}
 
+	_, noCallback := p.noCallbacks[n.C]
+	_, noEscape := p.noEscapes[n.C]
+	noCgoCallback := false
+	if noEscape && noCallback {
+		noCgoCallback = true
+	}
+	if noCgoCallback {
+		// disable cgocallback, will check it in runtime.
+		fmt.Fprintf(fgo2, "\t_Cgo_no_callback(true)\n")
+	}
+
 	prefix := ""
 	if n.AddError {
 		prefix = "errno := "
@@ -618,13 +631,19 @@ func (p *Package) writeDefsFunc(fgo2 io.Writer, n *Name, callsMalloc *bool) {
 	if n.AddError {
 		fmt.Fprintf(fgo2, "\tif errno != 0 { r2 = syscall.Errno(errno) }\n")
 	}
-	fmt.Fprintf(fgo2, "\tif _Cgo_always_false {\n")
-	if d.Type.Params != nil {
-		for i := range d.Type.Params.List {
-			fmt.Fprintf(fgo2, "\t\t_Cgo_use(p%d)\n", i)
+	if noCgoCallback {
+		fmt.Fprintf(fgo2, "\t_Cgo_no_callback(false)\n")
+	} else {
+		// skip _Cgo_use when noescape & nocallback both exist,
+		// so that the compiler won't force to escape them to heap.
+		fmt.Fprintf(fgo2, "\tif _Cgo_always_false {\n")
+		if d.Type.Params != nil {
+			for i := range d.Type.Params.List {
+				fmt.Fprintf(fgo2, "\t\t_Cgo_use(p%d)\n", i)
+			}
 		}
+		fmt.Fprintf(fgo2, "\t}\n")
 	}
-	fmt.Fprintf(fgo2, "\t}\n")
 	fmt.Fprintf(fgo2, "\treturn\n")
 	fmt.Fprintf(fgo2, "}\n")
 }

src/cmd/go/internal/modindex/build.go

@@ -20,6 +20,7 @@ import (
 	"io"
 	"io/fs"
 	"path/filepath"
+	"regexp"
 	"sort"
 	"strings"
 	"unicode"
@@ -108,6 +109,8 @@ type Context struct {
 	OpenFile func(path string) (io.ReadCloser, error)
 }
 
+var cgoRx = regexp.MustCompile(`^#cgo\s+(?:nocallback|noescape)\s+(?:\S+)\s*$`)
+
 // joinPath calls ctxt.JoinPath (if not nil) or else filepath.Join.
 func (ctxt *Context) joinPath(elem ...string) string {
 	if f := ctxt.JoinPath; f != nil {
@@ -622,6 +625,10 @@ func (ctxt *Context) saveCgo(filename string, di *build.Package, text string) er
 			continue
 		}
 
+		if cgoRx.FindStringSubmatch(line) != nil {
+			continue
+		}
+
 		// Split at colon.
 		line, argstr, ok := strings.Cut(strings.TrimSpace(line[4:]), ":")
 		if !ok {

src/go/build/build.go

@@ -23,6 +23,7 @@ import (
 	"os/exec"
 	pathpkg "path"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"sort"
 	"strconv"
@@ -113,6 +114,8 @@ type Context struct {
 	OpenFile func(path string) (io.ReadCloser, error)
 }
 
+var cgoRx = regexp.MustCompile(`^#cgo\s+(?:nocallback|noescape)\s+(?:\S+)\s*$`)
+
 // joinPath calls ctxt.JoinPath (if not nil) or else filepath.Join.
 func (ctxt *Context) joinPath(elem ...string) string {
 	if f := ctxt.JoinPath; f != nil {
@@ -1687,6 +1690,10 @@ func (ctxt *Context) saveCgo(filename string, di *Package, cg *ast.CommentGroup)
 			continue
 		}
 
+		if cgoRx.FindStringSubmatch(line) != nil {
+			continue
+		}
+
 		// Split at colon.
 		line, argstr, ok := strings.Cut(strings.TrimSpace(line[4:]), ":")
 		if !ok {

src/runtime/cgo.go

@@ -61,3 +61,8 @@ func cgoUse(any) { throw("cgoUse should not be called") }
 var cgoAlwaysFalse bool
 
 var cgo_yield = &_cgo_yield
+
+func cgoNoCallback(v bool) {
+	g := getg()
+	g.nocgocallback = v
+}

src/runtime/cgocall.go

@@ -216,6 +216,10 @@ func cgocallbackg(fn, frame unsafe.Pointer, ctxt uintptr) {
 		exit(2)
 	}
 
+	if gp.nocgocallback {
+		throw("runtime: disable callback from C")
+	}
+
 	// The call from C is on gp.m's g0 stack, so we must ensure
 	// that we stay on that M. We have to do this before calling
 	// exitsyscall, since it would otherwise be free to move us to

src/runtime/crash_cgo_test.go

@@ -749,6 +749,18 @@ func TestNeedmDeadlock(t *testing.T) {
 	}
 }
 
+func TestCgoNoCallback(t *testing.T) {
+	switch runtime.GOOS {
+	case "plan9", "windows":
+		t.Skipf("no signals on %s", runtime.GOOS)
+	}
+	got := runTestProg(t, "testprogcgo", "CgoNoCallback")
+	want := "runtime: disable callback from C"
+	if !strings.Contains(got, want) {
+		t.Fatalf("did not see %q in output: %q", want, got)
+	}
+}
+
 func TestCgoTracebackGoroutineProfile(t *testing.T) {
 	output := runTestProg(t, "testprogcgo", "GoroutineProfile")
 	want := "OK\n"

src/runtime/runtime2.go

@@ -473,6 +473,7 @@ type g struct {
 
 	raceignore    int8  // ignore race detection events
 	tracking      bool  // whether we're tracking this G for sched latency statistics
+	nocgocallback bool  // whether disable callback from C
 	trackingSeq   uint8 // used to decide whether to track this G
 	trackingStamp int64 // timestamp of when the G last started being tracked
 	runnableTime  int64 // the amount of time spent runnable, cleared when running, only used when tracking

src/runtime/testdata/testprogcgo/CgoNoCallback.go

@@ -0,0 +1,40 @@
+// Copyright 2020 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build !plan9 && !windows
+// +build !plan9,!windows
+
+package main
+
+// #cgo noescape/nocallback annotations for a C function means it should not callback to Go.
+// But it do callback to go in this test, Go should crash here.
+
+/*
+#cgo nocallback runCShouldNotCallback
+#cgo noescape runCShouldNotCallback
+
+extern void CallbackToGo();
+
+static void runCShouldNotCallback() {
+	CallbackToGo();
+}
+*/
+import "C"
+
+import (
+	"fmt"
+)
+
+func init() {
+	register("CgoNoCallback", CgoNoCallback)
+}
+
+//export CallbackToGo
+func CallbackToGo() {
+}
+
+func CgoNoCallback() {
+	C.runCShouldNotCallback()
+	fmt.Println("OK")
+}