Merge pull request #1946 from xiaobo55x:dieTramp

PiperOrigin-RevId: 299405663

3 files changed, 41 insertions, 17 deletions

diff --git a/pkg/sentry/platform/kvm/bluepill_arm64.s b/pkg/sentry/platform/kvm/bluepill_arm64.s
index c61700892..04efa0147 100644
--- a/pkg/sentry/platform/kvm/bluepill_arm64.s
+++ b/pkg/sentry/platform/kvm/bluepill_arm64.s
@@ -82,6 +82,8 @@ fallback:
 
 // dieTrampoline: see bluepill.go, bluepill_arm64_unsafe.go for documentation.
 TEXT ·dieTrampoline(SB),NOSPLIT,$0
-	// TODO(gvisor.dev/issue/1249): dieTrampoline supporting for Arm64.
-	MOVD	R9, 8(RSP)
-	BL	·dieHandler(SB)
+	// R0: Fake the old PC as caller
+	// R1: First argument (vCPU)
+	MOVD.P R1, 8(RSP) // R1: First argument (vCPU)
+	MOVD.P R0, 8(RSP) // R0: Fake the old PC as caller
+	B ·dieHandler(SB)
diff --git a/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go b/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
index 2f02c03cf..195331383 100644
--- a/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
+++ b/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
@@ -18,9 +18,30 @@ package kvm
 
 import (
 	"gvisor.dev/gvisor/pkg/sentry/arch"
+	"gvisor.dev/gvisor/pkg/sentry/platform/ring0"
 )
 
+// dieArchSetup initialies the state for dieTrampoline.
+//
+// The arm64 dieTrampoline requires the vCPU to be set in R1, and the last PC
+// to be in R0. The trampoline then simulates a call to dieHandler from the
+// provided PC.
+//
 //go:nosplit
 func dieArchSetup(c *vCPU, context *arch.SignalContext64, guestRegs *userRegs) {
-	// TODO(gvisor.dev/issue/1249): dieTrampoline supporting for Arm64.
+	// If the vCPU is in user mode, we set the stack to the stored stack
+	// value in the vCPU itself. We don't want to unwind the user stack.
+	if guestRegs.Regs.Pstate&ring0.PSR_MODE_MASK == ring0.PSR_MODE_EL0t {
+		regs := c.CPU.Registers()
+		context.Regs[0] = regs.Regs[0]
+		context.Sp = regs.Sp
+		context.Regs[29] = regs.Regs[29] // stack base address
+	} else {
+		context.Regs[0] = guestRegs.Regs.Pc
+		context.Sp = guestRegs.Regs.Sp
+		context.Regs[29] = guestRegs.Regs.Regs[29]
+		context.Pstate = guestRegs.Regs.Pstate
+	}
+	context.Regs[1] = uint64(uintptr(unsafe.Pointer(c)))
+	context.Pc = uint64(dieTrampolineAddr)
 }
diff --git a/pkg/sentry/platform/ring0/aarch64.go b/pkg/sentry/platform/ring0/aarch64.go
index f6da41c27..8122ac6e2 100644
--- a/pkg/sentry/platform/ring0/aarch64.go
+++ b/pkg/sentry/platform/ring0/aarch64.go
@@ -27,26 +27,27 @@ const (
 	_PTE_PGT_BASE = 0x7000
 	_PTE_PGT_SIZE = 0x1000
 
-	_PSR_MODE_EL0t = 0x0
-	_PSR_MODE_EL1t = 0x4
-	_PSR_MODE_EL1h = 0x5
-	_PSR_EL_MASK   = 0xf
-
-	_PSR_D_BIT = 0x200
-	_PSR_A_BIT = 0x100
-	_PSR_I_BIT = 0x80
-	_PSR_F_BIT = 0x40
+	_PSR_D_BIT = 0x00000200
+	_PSR_A_BIT = 0x00000100
+	_PSR_I_BIT = 0x00000080
+	_PSR_F_BIT = 0x00000040
 )
 
 const (
+	// PSR bits
+	PSR_MODE_EL0t = 0x00000000
+	PSR_MODE_EL1t = 0x00000004
+	PSR_MODE_EL1h = 0x00000005
+	PSR_MODE_MASK = 0x0000000f
+
 	// KernelFlagsSet should always be set in the kernel.
-	KernelFlagsSet = _PSR_MODE_EL1h
+	KernelFlagsSet = PSR_MODE_EL1h
 
 	// UserFlagsSet are always set in userspace.
-	UserFlagsSet = _PSR_MODE_EL0t
+	UserFlagsSet = PSR_MODE_EL0t
 
-	KernelFlagsClear = _PSR_EL_MASK
-	UserFlagsClear   = _PSR_EL_MASK
+	KernelFlagsClear = PSR_MODE_MASK
+	UserFlagsClear   = PSR_MODE_MASK
 
 	PsrDefaultSet = _PSR_D_BIT | _PSR_A_BIT | _PSR_I_BIT | _PSR_F_BIT
 )