Use assembly stub to take the address of assembly functions

Go 1.17 is adding a new register-based calling convention [1] ("ABIInternal"),
which used is when calling between Go functions. Assembly functions are still
written using the old ABI ("ABI0"). That is, they still accept arguments on the
stack, and pass arguments to other functions on the stack. The call rules look
approximately like this:

1. Direct call from Go function to Go function: compiler emits direct
   ABIInternal call.
2. Indirect call from Go function to Go function: compiler emits indirect
   ABIInternal call.
3. Direct call from Go function to assembly function: compiler emits direct
   ABI0 call.
4. Indirect call from Go function to assembly function: compiler emits indirect
   ABIInternal call to ABI conversion wrapper function.
5. Direct or indirect call from assembly function to assembly function:
   assembly/linker emits call to original ABI0 function.
6. Direct or indirect call from assembly function to Go function:
   assembly/linker emits ABI0 call to ABI conversion wrapper function.

Case 4 is the interesting one here. Since the compiler can't know the ABI of an
indirect call, all indirect calls are made with ABIInternal. In order to
support indirect ABI0 assembly function calls, a wrapper is generated that
translates ABIInternal arguments to ABI0 arguments, calls the target function,
and then converts results back.

When the address of an ABI0 function is taken from Go code, it evaluates to the
address of this wrapper function rather than the target function so that later
indirect calls will work as expected.

This is normally fine, but gVisor does more than just call some of the assembly
functions we take the address of: either noting the start and end address for
future reference from a signal handler (safecopy), or copying the function text
to a new mapping (platforms).

Both of these fail with wrappers enabled (currently, this is Go tip with
 GOEXPERIMENT=regabiwrappers) because these operations end up operating on the
wrapper instead of the target function.

We work around this issue by taking advantage of case 5: references to assembly
symbols from other assembly functions resolve directly to the desired target
symbol. Thus, rather than using reflect to get the address of a Go reference to
the functions, we create assembly stubs that return the address of the
function. This approach works just as well on current versions of Go, so the
change can be made immediately and doesn't require any build tags.

[1] https://go.googlesource.com/go/+/refs/heads/master/src/cmd/compile/abi-internal.md

PiperOrigin-RevId: 368505655

6 files changed, 54 insertions, 4 deletions

diff --git a/pkg/sentry/platform/kvm/bluepill.go b/pkg/sentry/platform/kvm/bluepill.go
index fd1131638..bb9967b9f 100644
--- a/pkg/sentry/platform/kvm/bluepill.go
+++ b/pkg/sentry/platform/kvm/bluepill.go
@@ -16,7 +16,6 @@ package kvm
 
 import (
 	"fmt"
-	"reflect"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/ring0"
@@ -36,6 +35,14 @@ func sighandler()
 // dieArchSetup and the assembly implementation for dieTrampoline.
 func dieTrampoline()
 
+// Return the start address of the functions above.
+//
+// In Go 1.17+, Go references to assembly functions resolve to an ABIInternal
+// wrapper function rather than the function itself. We must reference from
+// assembly to get the ABI0 (i.e., primary) address.
+func addrOfSighandler() uintptr
+func addrOfDieTrampoline() uintptr
+
 var (
 	// bounceSignal is the signal used for bouncing KVM.
 	//
@@ -87,10 +94,10 @@ func (c *vCPU) die(context *arch.SignalContext64, msg string) {
 
 func init() {
 	// Install the handler.
-	if err := safecopy.ReplaceSignalHandler(bluepillSignal, reflect.ValueOf(sighandler).Pointer(), &savedHandler); err != nil {
+	if err := safecopy.ReplaceSignalHandler(bluepillSignal, addrOfSighandler(), &savedHandler); err != nil {
 		panic(fmt.Sprintf("Unable to set handler for signal %d: %v", bluepillSignal, err))
 	}
 
 	// Extract the address for the trampoline.
-	dieTrampolineAddr = reflect.ValueOf(dieTrampoline).Pointer()
+	dieTrampolineAddr = addrOfDieTrampoline()
 }
diff --git a/pkg/sentry/platform/kvm/bluepill_amd64.s b/pkg/sentry/platform/kvm/bluepill_amd64.s
index 025ea93b5..953024600 100644
--- a/pkg/sentry/platform/kvm/bluepill_amd64.s
+++ b/pkg/sentry/platform/kvm/bluepill_amd64.s
@@ -81,8 +81,20 @@ fallback:
 	MOVQ ·savedHandler(SB), AX
 	JMP AX
 
+// func addrOfSighandler() uintptr
+TEXT ·addrOfSighandler(SB), $0-8
+	MOVQ $·sighandler(SB), AX
+	MOVQ AX, ret+0(FP)
+	RET
+
 // dieTrampoline: see bluepill.go, bluepill_amd64_unsafe.go for documentation.
 TEXT ·dieTrampoline(SB),NOSPLIT,$0
 	PUSHQ BX // First argument (vCPU).
 	PUSHQ AX // Fake the old RIP as caller.
 	JMP ·dieHandler(SB)
+
+// func addrOfDieTrampoline() uintptr
+TEXT ·addrOfDieTrampoline(SB), $0-8
+	MOVQ $·dieTrampoline(SB), AX
+	MOVQ AX, ret+0(FP)
+	RET
diff --git a/pkg/sentry/platform/kvm/bluepill_arm64.s b/pkg/sentry/platform/kvm/bluepill_arm64.s
index 09c7e88e5..308f2a951 100644
--- a/pkg/sentry/platform/kvm/bluepill_arm64.s
+++ b/pkg/sentry/platform/kvm/bluepill_arm64.s
@@ -92,6 +92,12 @@ fallback:
 	MOVD	·savedHandler(SB), R7
 	B	(R7)
 
+// func addrOfSighandler() uintptr
+TEXT ·addrOfSighandler(SB), $0-8
+	MOVD	$·sighandler(SB), R0
+	MOVD	R0, ret+0(FP)
+	RET
+
 // dieTrampoline: see bluepill.go, bluepill_arm64_unsafe.go for documentation.
 TEXT ·dieTrampoline(SB),NOSPLIT,$0
 	// R0: Fake the old PC as caller
@@ -99,3 +105,9 @@ TEXT ·dieTrampoline(SB),NOSPLIT,$0
 	MOVD.P R1, 8(RSP) // R1: First argument (vCPU)
 	MOVD.P R0, 8(RSP) // R0: Fake the old PC as caller
 	B ·dieHandler(SB)
+
+// func addrOfDieTrampoline() uintptr
+TEXT ·addrOfDieTrampoline(SB), $0-8
+	MOVD	$·dieTrampoline(SB), R0
+	MOVD	R0, ret+0(FP)
+	RET
diff --git a/pkg/sentry/platform/ptrace/stub_amd64.s b/pkg/sentry/platform/ptrace/stub_amd64.s
index 16f9c523e..d5c3f901f 100644
--- a/pkg/sentry/platform/ptrace/stub_amd64.s
+++ b/pkg/sentry/platform/ptrace/stub_amd64.s
@@ -109,6 +109,12 @@ parent_dead:
 	SYSCALL
 	HLT
 
+// func addrOfStub() uintptr
+TEXT ·addrOfStub(SB), $0-8
+	MOVQ $·stub(SB), AX
+	MOVQ AX, ret+0(FP)
+	RET
+
 // stubCall calls the stub function at the given address with the given PPID.
 //
 // This is a distinct function because stub, above, may be mapped at any
diff --git a/pkg/sentry/platform/ptrace/stub_arm64.s b/pkg/sentry/platform/ptrace/stub_arm64.s
index 6162df02a..4664cd4ad 100644
--- a/pkg/sentry/platform/ptrace/stub_arm64.s
+++ b/pkg/sentry/platform/ptrace/stub_arm64.s
@@ -102,6 +102,12 @@ parent_dead:
 	SVC
 	HLT
 
+// func addrOfStub() uintptr
+TEXT ·addrOfStub(SB), $0-8
+	MOVD	$·stub(SB), R0
+	MOVD	R0, ret+0(FP)
+	RET
+
 // stubCall calls the stub function at the given address with the given PPID.
 //
 // This is a distinct function because stub, above, may be mapped at any
diff --git a/pkg/sentry/platform/ptrace/stub_unsafe.go b/pkg/sentry/platform/ptrace/stub_unsafe.go
index 5c9b7784f..1fbdea898 100644
--- a/pkg/sentry/platform/ptrace/stub_unsafe.go
+++ b/pkg/sentry/platform/ptrace/stub_unsafe.go
@@ -26,6 +26,13 @@ import (
 // stub is defined in arch-specific assembly.
 func stub()
 
+// addrOfStub returns the start address of stub.
+//
+// In Go 1.17+, Go references to assembly functions resolve to an ABIInternal
+// wrapper function rather than the function itself. We must reference from
+// assembly to get the ABI0 (i.e., primary) address.
+func addrOfStub() uintptr
+
 // stubCall calls the stub at the given address with the given pid.
 func stubCall(addr, pid uintptr)
 
@@ -41,7 +48,7 @@ func unsafeSlice(addr uintptr, length int) (slice []byte) {
 // stubInit initializes the stub.
 func stubInit() {
 	// Grab the existing stub.
-	stubBegin := reflect.ValueOf(stub).Pointer()
+	stubBegin := addrOfStub()
 	stubLen := int(safecopy.FindEndAddress(stubBegin) - stubBegin)
 	stubSlice := unsafeSlice(stubBegin, stubLen)
 	mapLen := uintptr(stubLen)