gVisor: Turn arch.Context interface references to struct pointers.

Since `arch.Context64` is the only concrete implementation of
`arch.Context`, this does not functionally change anything.
While `arch.Context64` *can* mean multiple things in the codebase, since
the `arch` module uses conditional compilation to create different builds
for AMD64 vs ARM64, the `Context64` type is still singular within each
possible version this can compile to.

This avoids the overhead involved in calling interface
functions in Go:
https://github.com/teh-cmc/go-internals/blob/master/chapter2_interfaces/README.md#dynamic-dispatch

This overhead is particularly painful for functions like `SyscallSaveOrig`
which is a no-op on AMD64.

On KVM, this reduces syscall latency by 4~5%:

```
name       old wall_ns/op  new wall_ns/op  delta
Getpid           602 ± 4%        579 ± 3%  -3.94%  (p=0.000 n=66+195)
GetpidOpt        604 ± 3%        573 ± 3%  -5.10%  (p=0.000 n=71+194)

name       old cpu_ns/op   new cpu_ns/op   delta
Getpid           604 ± 5%        581 ± 0%  -3.81%  (p=0.000 n=75+130)
GetpidOpt        604 ± 2%        574 ± 3%  -4.95%  (p=0.000 n=71+196)
```

PiperOrigin-RevId: 461742893

pkg/sentry/arch/arch.go

@@ -50,14 +50,22 @@ func (a Arch) String() string {
 	}
 }
 
-// Context provides architecture-dependent information for a specific thread.
+// contextInterface provides architecture-dependent information for a thread.
+// This is currently not referenced, because there exists only one concrete
+// implementation of this interface (*Context64), which we reference directly
+// wherever this interface could otherwise be used in order to avoid the
+// overhead involved in calling functions on interfaces in Go.
+// This interface is still useful in order to see the entire
+// architecture-dependent call surface it must support, as this is difficult
+// to follow across the rest of this module due to the conditional compilation
+// of the files that make it up.
 //
 // NOTE(b/34169503): Currently we use uintptr here to refer to a generic native
 // register value. While this will work for the foreseeable future, it isn't
 // strictly correct. We may want to create some abstraction that makes this
 // more clear or enables us to store values of arbitrary widths. This is
 // particularly true for RegisterMap().
-type Context interface {
+type contextInterface interface {
 	// Arch returns the architecture for this Context.
 	Arch() Arch
 
@@ -78,7 +86,7 @@ type Context interface {
 	Width() uint
 
 	// Fork creates a clone of the context.
-	Fork() Context
+	Fork() *Context64
 
 	// SyscallNo returns the syscall number.
 	SyscallNo() uintptr
@@ -228,6 +236,9 @@ type Context interface {
 	FullRestore() bool
 }
 
+// Compile-time assertion that Context64 implements contextInterface.
+var _ = (contextInterface)((*Context64)(nil))
+
 // MmapDirection is a search direction for mmaps.
 type MmapDirection int
 

pkg/sentry/arch/arch_aarch64.go

@@ -251,10 +251,10 @@ func (s *State) FullRestore() bool {
 }
 
 // New returns a new architecture context.
-func New(arch Arch) Context {
+func New(arch Arch) *Context64 {
 	switch arch {
 	case ARM64:
-		return &context64{
+		return &Context64{
 			State{
 				fpState: fpu.NewState(),
 			},

pkg/sentry/arch/arch_amd64.go

@@ -101,66 +101,67 @@ const (
 	minMmapRand64 = (1 << 26) * hostarch.PageSize
 )
 
-// context64 represents an AMD64 context.
+// Context64 represents an AMD64 context.
 //
 // +stateify savable
-type context64 struct {
+type Context64 struct {
 	State
 }
 
 // Arch implements Context.Arch.
-func (c *context64) Arch() Arch {
+func (c *Context64) Arch() Arch {
 	return AMD64
 }
 
-func (c *context64) FloatingPointData() *fpu.State {
+// FloatingPointData returns the state of the floating-point unit.
+func (c *Context64) FloatingPointData() *fpu.State {
 	return &c.State.fpState
 }
 
 // Fork returns an exact copy of this context.
-func (c *context64) Fork() Context {
-	return &context64{
+func (c *Context64) Fork() *Context64 {
+	return &Context64{
 		State: c.State.Fork(),
 	}
 }
 
 // Return returns the current syscall return value.
-func (c *context64) Return() uintptr {
+func (c *Context64) Return() uintptr {
 	return uintptr(c.Regs.Rax)
 }
 
 // SetReturn sets the syscall return value.
-func (c *context64) SetReturn(value uintptr) {
+func (c *Context64) SetReturn(value uintptr) {
 	c.Regs.Rax = uint64(value)
 }
 
 // IP returns the current instruction pointer.
-func (c *context64) IP() uintptr {
+func (c *Context64) IP() uintptr {
 	return uintptr(c.Regs.Rip)
 }
 
 // SetIP sets the current instruction pointer.
-func (c *context64) SetIP(value uintptr) {
+func (c *Context64) SetIP(value uintptr) {
 	c.Regs.Rip = uint64(value)
 }
 
 // Stack returns the current stack pointer.
-func (c *context64) Stack() uintptr {
+func (c *Context64) Stack() uintptr {
 	return uintptr(c.Regs.Rsp)
 }
 
 // SetStack sets the current stack pointer.
-func (c *context64) SetStack(value uintptr) {
+func (c *Context64) SetStack(value uintptr) {
 	c.Regs.Rsp = uint64(value)
 }
 
 // TLS returns the current TLS pointer.
-func (c *context64) TLS() uintptr {
+func (c *Context64) TLS() uintptr {
 	return uintptr(c.Regs.Fs_base)
 }
 
 // SetTLS sets the current TLS pointer. Returns false if value is invalid.
-func (c *context64) SetTLS(value uintptr) bool {
+func (c *Context64) SetTLS(value uintptr) bool {
 	if !isValidSegmentBase(uint64(value)) {
 		return false
 	}
@@ -171,23 +172,23 @@ func (c *context64) SetTLS(value uintptr) bool {
 }
 
 // SetOldRSeqInterruptedIP implements Context.SetOldRSeqInterruptedIP.
-func (c *context64) SetOldRSeqInterruptedIP(value uintptr) {
+func (c *Context64) SetOldRSeqInterruptedIP(value uintptr) {
 	c.Regs.R10 = uint64(value)
 }
 
 // Native returns the native type for the given val.
-func (c *context64) Native(val uintptr) marshal.Marshallable {
+func (c *Context64) Native(val uintptr) marshal.Marshallable {
 	v := primitive.Uint64(val)
 	return &v
 }
 
 // Value returns the generic val for the given native type.
-func (c *context64) Value(val marshal.Marshallable) uintptr {
+func (c *Context64) Value(val marshal.Marshallable) uintptr {
 	return uintptr(*val.(*primitive.Uint64))
 }
 
 // Width returns the byte width of this architecture.
-func (c *context64) Width() uint {
+func (c *Context64) Width() uint {
 	return 8
 }
 
@@ -197,7 +198,7 @@ func mmapRand(max uint64) hostarch.Addr {
 }
 
 // NewMmapLayout implements Context.NewMmapLayout consistently with Linux.
-func (c *context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (MmapLayout, error) {
+func (c *Context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (MmapLayout, error) {
 	min, ok := min.RoundUp()
 	if !ok {
 		return MmapLayout{}, unix.EINVAL
@@ -263,7 +264,7 @@ func (c *context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (M
 }
 
 // PIELoadAddress implements Context.PIELoadAddress.
-func (c *context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
+func (c *Context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
 	base := preferredPIELoadAddr
 	max, ok := base.AddLength(maxMmapRand64)
 	if !ok {
@@ -285,7 +286,7 @@ func (c *context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
 const userStructSize = 928
 
 // PtracePeekUser implements Context.PtracePeekUser.
-func (c *context64) PtracePeekUser(addr uintptr) (marshal.Marshallable, error) {
+func (c *Context64) PtracePeekUser(addr uintptr) (marshal.Marshallable, error) {
 	if addr&7 != 0 || addr >= userStructSize {
 		return nil, unix.EIO
 	}
@@ -303,7 +304,7 @@ func (c *context64) PtracePeekUser(addr uintptr) (marshal.Marshallable, error) {
 }
 
 // PtracePokeUser implements Context.PtracePokeUser.
-func (c *context64) PtracePokeUser(addr, data uintptr) error {
+func (c *Context64) PtracePokeUser(addr, data uintptr) error {
 	if addr&7 != 0 || addr >= userStructSize {
 		return unix.EIO
 	}

pkg/sentry/arch/arch_arm64.go

@@ -75,20 +75,20 @@ const (
 	minMmapRand64 = (1 << 18) * hostarch.PageSize
 )
 
-// context64 represents an ARM64 context.
+// Context64 represents an ARM64 context.
 //
 // +stateify savable
-type context64 struct {
+type Context64 struct {
 	State
 	sigFPState []fpu.State // fpstate to be restored on sigreturn.
 }
 
 // Arch implements Context.Arch.
-func (c *context64) Arch() Arch {
+func (c *Context64) Arch() Arch {
 	return ARM64
 }
 
-func (c *context64) copySigFPState() []fpu.State {
+func (c *Context64) copySigFPState() []fpu.State {
 	var sigfps []fpu.State
 	for _, s := range c.sigFPState {
 		sigfps = append(sigfps, s.Fork())
@@ -97,8 +97,8 @@ func (c *context64) copySigFPState() []fpu.State {
 }
 
 // Fork returns an exact copy of this context.
-func (c *context64) Fork() Context {
-	return &context64{
+func (c *Context64) Fork() *Context64 {
+	return &Context64{
 		State:      c.State.Fork(),
 		sigFPState: c.copySigFPState(),
 	}
@@ -116,42 +116,42 @@ func (c *context64) Fork() Context {
 // R30: the link register.
 
 // Return returns the current syscall return value.
-func (c *context64) Return() uintptr {
+func (c *Context64) Return() uintptr {
 	return uintptr(c.Regs.Regs[0])
 }
 
 // SetReturn sets the syscall return value.
-func (c *context64) SetReturn(value uintptr) {
+func (c *Context64) SetReturn(value uintptr) {
 	c.Regs.Regs[0] = uint64(value)
 }
 
 // IP returns the current instruction pointer.
-func (c *context64) IP() uintptr {
+func (c *Context64) IP() uintptr {
 	return uintptr(c.Regs.Pc)
 }
 
 // SetIP sets the current instruction pointer.
-func (c *context64) SetIP(value uintptr) {
+func (c *Context64) SetIP(value uintptr) {
 	c.Regs.Pc = uint64(value)
 }
 
 // Stack returns the current stack pointer.
-func (c *context64) Stack() uintptr {
+func (c *Context64) Stack() uintptr {
 	return uintptr(c.Regs.Sp)
 }
 
 // SetStack sets the current stack pointer.
-func (c *context64) SetStack(value uintptr) {
+func (c *Context64) SetStack(value uintptr) {
 	c.Regs.Sp = uint64(value)
 }
 
 // TLS returns the current TLS pointer.
-func (c *context64) TLS() uintptr {
+func (c *Context64) TLS() uintptr {
 	return uintptr(c.Regs.TPIDR_EL0)
 }
 
 // SetTLS sets the current TLS pointer. Returns false if value is invalid.
-func (c *context64) SetTLS(value uintptr) bool {
+func (c *Context64) SetTLS(value uintptr) bool {
 	if value >= uintptr(maxAddr64) {
 		return false
 	}
@@ -161,23 +161,23 @@ func (c *context64) SetTLS(value uintptr) bool {
 }
 
 // SetOldRSeqInterruptedIP implements Context.SetOldRSeqInterruptedIP.
-func (c *context64) SetOldRSeqInterruptedIP(value uintptr) {
+func (c *Context64) SetOldRSeqInterruptedIP(value uintptr) {
 	c.Regs.Regs[3] = uint64(value)
 }
 
 // Native returns the native type for the given val.
-func (c *context64) Native(val uintptr) marshal.Marshallable {
+func (c *Context64) Native(val uintptr) marshal.Marshallable {
 	v := primitive.Uint64(val)
 	return &v
 }
 
 // Value returns the generic val for the given native type.
-func (c *context64) Value(val marshal.Marshallable) uintptr {
+func (c *Context64) Value(val marshal.Marshallable) uintptr {
 	return uintptr(*val.(*primitive.Uint64))
 }
 
 // Width returns the byte width of this architecture.
-func (c *context64) Width() uint {
+func (c *Context64) Width() uint {
 	return 8
 }
 
@@ -187,7 +187,7 @@ func mmapRand(max uint64) hostarch.Addr {
 }
 
 // NewMmapLayout implements Context.NewMmapLayout consistently with Linux.
-func (c *context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (MmapLayout, error) {
+func (c *Context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (MmapLayout, error) {
 	min, ok := min.RoundUp()
 	if !ok {
 		return MmapLayout{}, unix.EINVAL
@@ -253,7 +253,7 @@ func (c *context64) NewMmapLayout(min, max hostarch.Addr, r *limits.LimitSet) (M
 }
 
 // PIELoadAddress implements Context.PIELoadAddress.
-func (c *context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
+func (c *Context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
 	base := preferredPIELoadAddr
 	max, ok := base.AddLength(maxMmapRand64)
 	if !ok {
@@ -272,17 +272,18 @@ func (c *context64) PIELoadAddress(l MmapLayout) hostarch.Addr {
 }
 
 // PtracePeekUser implements Context.PtracePeekUser.
-func (c *context64) PtracePeekUser(addr uintptr) (marshal.Marshallable, error) {
+func (c *Context64) PtracePeekUser(addr uintptr) (marshal.Marshallable, error) {
 	// TODO(gvisor.dev/issue/1239): Full ptrace supporting for Arm64.
 	return c.Native(0), nil
 }
 
 // PtracePokeUser implements Context.PtracePokeUser.
-func (c *context64) PtracePokeUser(addr, data uintptr) error {
+func (c *Context64) PtracePokeUser(addr, data uintptr) error {
 	// TODO(gvisor.dev/issue/1239): Full ptrace supporting for Arm64.
 	return nil
 }
 
-func (c *context64) FloatingPointData() *fpu.State {
+// FloatingPointData returns the state of the floating-point unit.
+func (c *Context64) FloatingPointData() *fpu.State {
 	return &c.State.fpState
 }

pkg/sentry/arch/arch_x86.go

@@ -390,10 +390,10 @@ func (s *State) FullRestore() bool {
 }
 
 // New returns a new architecture context.
-func New(arch Arch) Context {
+func New(arch Arch) *Context64 {
 	switch arch {
 	case AMD64:
-		return &context64{
+		return &Context64{
 			State{
 				fpState: fpu.NewState(),
 			},

pkg/sentry/arch/signal_amd64.go

@@ -114,7 +114,7 @@ const (
 
 // SignalSetup implements Context.SignalSetup. (Compare to Linux's
 // arch/x86/kernel/signal.c:__setup_rt_frame().)
-func (c *context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.SignalInfo, alt *linux.SignalStack, sigset linux.SignalSet, featureSet cpuid.FeatureSet) error {
+func (c *Context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.SignalInfo, alt *linux.SignalStack, sigset linux.SignalSet, featureSet cpuid.FeatureSet) error {
 	// "The 128-byte area beyond the location pointed to by %rsp is considered
 	// to be reserved and shall not be modified by signal or interrupt
 	// handlers. ... leaf functions may use this area for their entire stack
@@ -265,7 +265,7 @@ func (c *context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.Sig
 
 // SignalRestore implements Context.SignalRestore. (Compare to Linux's
 // arch/x86/kernel/signal.c:sys_rt_sigreturn().)
-func (c *context64) SignalRestore(st *Stack, rt bool, featureSet cpuid.FeatureSet) (linux.SignalSet, linux.SignalStack, error) {
+func (c *Context64) SignalRestore(st *Stack, rt bool, featureSet cpuid.FeatureSet) (linux.SignalSet, linux.SignalStack, error) {
 	// Copy out the stack frame.
 	var uc UContext64
 	if _, err := uc.CopyIn(st, StackBottomMagic); err != nil {

pkg/sentry/arch/signal_arm64.go

@@ -74,7 +74,7 @@ type UContext64 struct {
 }
 
 // SignalSetup implements Context.SignalSetup.
-func (c *context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.SignalInfo, alt *linux.SignalStack, sigset linux.SignalSet, featureSet cpuid.FeatureSet) error {
+func (c *Context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.SignalInfo, alt *linux.SignalStack, sigset linux.SignalSet, featureSet cpuid.FeatureSet) error {
 	sp := st.Bottom
 
 	// Construct the UContext64 now since we need its size.
@@ -139,7 +139,7 @@ func (c *context64) SignalSetup(st *Stack, act *linux.SigAction, info *linux.Sig
 }
 
 // SignalRestore implements Context.SignalRestore.
-func (c *context64) SignalRestore(st *Stack, rt bool, featureSet cpuid.FeatureSet) (linux.SignalSet, linux.SignalStack, error) {
+func (c *Context64) SignalRestore(st *Stack, rt bool, featureSet cpuid.FeatureSet) (linux.SignalSet, linux.SignalStack, error) {
 	// Copy out the stack frame.
 	var uc UContext64
 	if _, err := uc.CopyIn(st, StackBottomMagic); err != nil {

pkg/sentry/arch/stack.go

@@ -31,7 +31,7 @@ type Stack struct {
 	// Our arch info.
 	// We use this for automatic Native conversion of hostarch.Addrs during
 	// Push() and Pop().
-	Arch Context
+	Arch *Context64
 
 	// The interface used to actually copy user memory.
 	IO usermem.IO

pkg/sentry/arch/syscalls_amd64.go

@@ -23,11 +23,11 @@ const restartSyscallNr = uintptr(219)
 // syscall handler(doSyscall()).
 //
 // Noop on x86.
-func (c *context64) SyscallSaveOrig() {
+func (c *Context64) SyscallSaveOrig() {
 }
 
 // SyscallNo returns the syscall number according to the 64-bit convention.
-func (c *context64) SyscallNo() uintptr {
+func (c *Context64) SyscallNo() uintptr {
 	return uintptr(c.Regs.Orig_rax)
 }
 
@@ -36,7 +36,7 @@ func (c *context64) SyscallNo() uintptr {
 // Due to the way addresses are mapped for the sentry this binary *must* be
 // built in 64-bit mode. So we can just assume the syscall numbers that come
 // back match the expected host system call numbers.
-func (c *context64) SyscallArgs() SyscallArguments {
+func (c *Context64) SyscallArgs() SyscallArguments {
 	return SyscallArguments{
 		SyscallArgument{Value: uintptr(c.Regs.Rdi)},
 		SyscallArgument{Value: uintptr(c.Regs.Rsi)},
@@ -48,13 +48,13 @@ func (c *context64) SyscallArgs() SyscallArguments {
 }
 
 // RestartSyscall implements Context.RestartSyscall.
-func (c *context64) RestartSyscall() {
+func (c *Context64) RestartSyscall() {
 	c.Regs.Rip -= SyscallWidth
 	c.Regs.Rax = c.Regs.Orig_rax
 }
 
 // RestartSyscallWithRestartBlock implements Context.RestartSyscallWithRestartBlock.
-func (c *context64) RestartSyscallWithRestartBlock() {
+func (c *Context64) RestartSyscallWithRestartBlock() {
 	c.Regs.Rip -= SyscallWidth
 	c.Regs.Rax = uint64(restartSyscallNr)
 }

pkg/sentry/arch/syscalls_arm64.go

@@ -26,12 +26,12 @@ const restartSyscallNr = uintptr(128)
 // is saved to the pt_regs.orig_x0 in kernel code. But currently, the orig_x0
 // was not accessible to the userspace application, so we have to do the same
 // operation in the sentry code to save the R0 value into the App context.
-func (c *context64) SyscallSaveOrig() {
+func (c *Context64) SyscallSaveOrig() {
 	c.OrigR0 = c.Regs.Regs[0]
 }
 
 // SyscallNo returns the syscall number according to the 64-bit convention.
-func (c *context64) SyscallNo() uintptr {
+func (c *Context64) SyscallNo() uintptr {
 	return uintptr(c.Regs.Regs[8])
 }
 
@@ -50,7 +50,7 @@ func (c *context64) SyscallNo() uintptr {
 // R19...R28: callee-saved registers.
 // R29: the frame pointer.
 // R30: the link register.
-func (c *context64) SyscallArgs() SyscallArguments {
+func (c *Context64) SyscallArgs() SyscallArguments {
 	return SyscallArguments{
 		SyscallArgument{Value: uintptr(c.OrigR0)},
 		SyscallArgument{Value: uintptr(c.Regs.Regs[1])},
@@ -65,7 +65,7 @@ func (c *context64) SyscallArgs() SyscallArguments {
 // Prepare for system call restart, OrigR0 will be restored to R0.
 // Please see the linux code as reference:
 // arch/arm64/kernel/signal.c:do_signal()
-func (c *context64) RestartSyscall() {
+func (c *Context64) RestartSyscall() {
 	c.Regs.Pc -= SyscallWidth
 	// R0 will be backed up into OrigR0 when entering doSyscall().
 	// Please see the linux code as reference:
@@ -75,7 +75,7 @@ func (c *context64) RestartSyscall() {
 }
 
 // RestartSyscallWithRestartBlock implements Context.RestartSyscallWithRestartBlock.
-func (c *context64) RestartSyscallWithRestartBlock() {
+func (c *Context64) RestartSyscallWithRestartBlock() {
 	c.Regs.Pc -= SyscallWidth
 	c.Regs.Regs[0] = uint64(c.OrigR0)
 	c.Regs.Regs[8] = uint64(restartSyscallNr)

pkg/sentry/kernel/task_image.go

@@ -40,7 +40,7 @@ type TaskImage struct {
 	Name string
 
 	// Arch is the architecture-specific context (registers, etc.)
-	Arch arch.Context
+	Arch *arch.Context64
 
 	// MemoryManager is the task's address space.
 	MemoryManager *mm.MemoryManager
@@ -65,7 +65,7 @@ func (image *TaskImage) release() {
 }
 
 // Fork returns a duplicate of image. The copied TaskImage always has an
-// independent arch.Context. If shareAddressSpace is true, the copied
+// independent arch.Context64. If shareAddressSpace is true, the copied
 // TaskImage shares an address space with the original; otherwise, the copied
 // TaskImage has an independent address space that is initially a duplicate
 // of the original's.
@@ -96,11 +96,11 @@ func (image *TaskImage) Fork(ctx context.Context, k *Kernel, shareAddressSpace b
 	return newImage, nil
 }
 
-// Arch returns t's arch.Context.
+// Arch returns t's arch.Context64.
 //
 // Preconditions: The caller must be running on the task goroutine, or t.mu
 // must be locked.
-func (t *Task) Arch() arch.Context {
+func (t *Task) Arch() *arch.Context64 {
 	return t.image.Arch
 }
 

pkg/sentry/loader/elf.go

@@ -573,14 +573,14 @@ func loadParsedELF(ctx context.Context, m *mm.MemoryManager, f fsbridge.File, in
 
 // loadInitialELF loads f into mm.
 //
-// It creates an arch.Context for the ELF and prepares the mm for this arch.
+// It creates an arch.Context64 for the ELF and prepares the mm for this arch.
 //
 // It does not load the ELF interpreter, or return any auxv entries.
 //
 // Preconditions:
 //   - f is an ELF file.
 //   - f is the first ELF loaded into m.
-func loadInitialELF(ctx context.Context, m *mm.MemoryManager, fs cpuid.FeatureSet, f fsbridge.File) (loadedELF, arch.Context, error) {
+func loadInitialELF(ctx context.Context, m *mm.MemoryManager, fs cpuid.FeatureSet, f fsbridge.File) (loadedELF, *arch.Context64, error) {
 	info, err := parseHeader(ctx, f)
 	if err != nil {
 		ctx.Infof("Failed to parse initial ELF: %v", err)
@@ -593,7 +593,7 @@ func loadInitialELF(ctx context.Context, m *mm.MemoryManager, fs cpuid.FeatureSe
 		return loadedELF{}, nil, linuxerr.ENOEXEC
 	}
 
-	// Create the arch.Context now so we can prepare the mmap layout before
+	// Create the arch.Context64 now so we can prepare the mmap layout before
 	// mapping anything.
 	ac := arch.New(info.arch)
 
@@ -647,7 +647,7 @@ func loadInterpreterELF(ctx context.Context, m *mm.MemoryManager, f fsbridge.Fil
 // path and argv.
 //
 // Preconditions: args.File is an ELF file.
-func loadELF(ctx context.Context, args LoadArgs) (loadedELF, arch.Context, error) {
+func loadELF(ctx context.Context, args LoadArgs) (loadedELF, *arch.Context64, error) {
 	bin, ac, err := loadInitialELF(ctx, args.MemoryManager, args.Features, args.File)
 	if err != nil {
 		ctx.Infof("Error loading binary: %v", err)

pkg/sentry/loader/loader.go

@@ -130,7 +130,7 @@ func checkIsRegularFile(ctx context.Context, file fsbridge.File, filename string
 }
 
 // allocStack allocates and maps a stack in to any available part of the address space.
-func allocStack(ctx context.Context, m *mm.MemoryManager, a arch.Context) (*arch.Stack, error) {
+func allocStack(ctx context.Context, m *mm.MemoryManager, a *arch.Context64) (*arch.Stack, error) {
 	ar, err := m.MapStack(ctx)
 	if err != nil {
 		return nil, err
@@ -154,10 +154,10 @@ const (
 //
 // It returns:
 //   - loadedELF, description of the loaded binary
-//   - arch.Context matching the binary arch
+//   - arch.Context64 matching the binary arch
 //   - fs.Dirent of the binary file
 //   - Possibly updated args.Argv
-func loadExecutable(ctx context.Context, args LoadArgs) (loadedELF, arch.Context, fsbridge.File, []string, error) {
+func loadExecutable(ctx context.Context, args LoadArgs) (loadedELF, *arch.Context64, fsbridge.File, []string, error) {
 	for i := 0; i < maxLoaderAttempts; i++ {
 		if args.File == nil {
 			var err error
@@ -230,7 +230,7 @@ func loadExecutable(ctx context.Context, args LoadArgs) (loadedELF, arch.Context
 // Preconditions:
 //   - The Task MemoryManager is empty.
 //   - Load is called on the Task goroutine.
-func Load(ctx context.Context, args LoadArgs, extraAuxv []arch.AuxEntry, vdso *VDSO) (abi.OS, arch.Context, string, *syserr.Error) {
+func Load(ctx context.Context, args LoadArgs, extraAuxv []arch.AuxEntry, vdso *VDSO) (abi.OS, *arch.Context64, string, *syserr.Error) {
 	// Load the executable itself.
 	loaded, ac, file, newArgv, err := loadExecutable(ctx, args)
 	if err != nil {

pkg/sentry/mm/lifecycle.go

@@ -42,10 +42,10 @@ func NewMemoryManager(p platform.Platform, mfp pgalloc.MemoryFileProvider, sleep
 	}
 }
 
-// SetMmapLayout initializes mm's layout from the given arch.Context.
+// SetMmapLayout initializes mm's layout from the given arch.Context64.
 //
 // Preconditions: mm contains no mappings and is not used concurrently.
-func (mm *MemoryManager) SetMmapLayout(ac arch.Context, r *limits.LimitSet) (arch.MmapLayout, error) {
+func (mm *MemoryManager) SetMmapLayout(ac *arch.Context64, r *limits.LimitSet) (arch.MmapLayout, error) {
 	layout, err := ac.NewMmapLayout(mm.p.MinUserAddress(), mm.p.MaxUserAddress(), r)
 	if err != nil {
 		return arch.MmapLayout{}, err

pkg/sentry/platform/cpuid_amd64.go

@@ -43,7 +43,7 @@ type emulationContext struct {
 }
 
 // TryCPUIDEmulate checks for a CPUID instruction and performs emulation.
-func TryCPUIDEmulate(ctx context.Context, mm MemoryManager, ac arch.Context) bool {
+func TryCPUIDEmulate(ctx context.Context, mm MemoryManager, ac *arch.Context64) bool {
 	s := ac.StateData()
 	inst := make([]byte, len(arch.CPUIDInstruction))
 	tasklessCtx := emulationContext{

pkg/sentry/platform/cpuid_arm64.go

@@ -23,6 +23,6 @@ import (
 )
 
 // TryCPUIDEmulate always returns false: there is no cpuid.
-func TryCPUIDEmulate(ctx context.Context, mm MemoryManager, ac arch.Context) bool {
+func TryCPUIDEmulate(ctx context.Context, mm MemoryManager, ac *arch.Context64) bool {
 	return false
 }

pkg/sentry/platform/kvm/context.go

@@ -45,7 +45,7 @@ type tryCPUIDError struct{}
 func (tryCPUIDError) Error() string { return "cpuid emulation failed" }
 
 // Switch runs the provided context in the given address space.
-func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac arch.Context, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	localAS := as.(*addressSpace)
 
@@ -125,4 +125,4 @@ func (c *context) Release() {}
 func (c *context) FullStateChanged() {}
 
 // PullFullState implements platform.Context.PullFullState.
-func (c *context) PullFullState(as platform.AddressSpace, ac arch.Context) {}
+func (c *context) PullFullState(as platform.AddressSpace, ac *arch.Context64) {}

pkg/sentry/platform/platform.go

@@ -180,7 +180,7 @@ type MemoryManager interface {
 
 // Context represents the execution context for a single thread.
 type Context interface {
-	// Switch resumes execution of the thread specified by the arch.Context
+	// Switch resumes execution of the thread specified by the arch.Context64
 	// in the provided address space. This call will block while the thread
 	// is executing.
 	//
@@ -207,7 +207,7 @@ type Context interface {
 	//		concurrent call to Switch().
 	//
 	//	- ErrContextCPUPreempted: See the definition of that error for details.
-	Switch(ctx context.Context, mm MemoryManager, ac arch.Context, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error)
+	Switch(ctx context.Context, mm MemoryManager, ac *arch.Context64, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error)
 
 	// PullFullState() pulls a full state of the application thread.
 	//
@@ -221,7 +221,7 @@ type Context interface {
 	// PullFullState() to load all registers and FPU state.
 	//
 	// Preconditions: The caller must be running on the task goroutine.
-	PullFullState(as AddressSpace, ac arch.Context)
+	PullFullState(as AddressSpace, ac *arch.Context64)
 
 	// FullStateChanged() indicates that a thread state has been changed by
 	// the Sentry. This happens in case of the rt_sigreturn, execve, etc.

pkg/sentry/platform/ptrace/ptrace.go

@@ -105,7 +105,7 @@ func (*PTrace) NewContext(ctx pkgcontext.Context) platform.Context {
 }
 
 // Switch runs the provided context in the given address space.
-func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac arch.Context, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+func (c *context) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, cpu int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	s := as.(*subprocess)
 restart:
@@ -197,7 +197,7 @@ func (c *context) Release() {}
 func (c *context) FullStateChanged() {}
 
 // PullFullState implements platform.Context.PullFullState.
-func (c *context) PullFullState(as platform.AddressSpace, ac arch.Context) {}
+func (c *context) PullFullState(as platform.AddressSpace, ac *arch.Context64) {}
 
 // PTrace represents a collection of ptrace subprocesses.
 type PTrace struct {

pkg/sentry/platform/ptrace/subprocess.go

@@ -501,7 +501,7 @@ func (t *thread) NotifyInterrupt() {
 // switchToApp is called from the main SwitchToApp entrypoint.
 //
 // This function returns true on a system call, false on a signal.
-func (s *subprocess) switchToApp(c *context, ac arch.Context) bool {
+func (s *subprocess) switchToApp(c *context, ac *arch.Context64) bool {
 	// Lock the thread for ptrace operations.
 	runtime.LockOSThread()
 	defer runtime.UnlockOSThread()