kvm: RSEQ: Add RSEQ support

This change adds RSEQ support for platform/kvm.

To support RSEQ, the KVM platform must provide two key capabilities:
 (1) A stable, unique CPU ID.
 (2) A way to detect when a thread as been preempted.

This implementation provides the necessary support as follows:
CPU ID:
  - platform/kvm now advertises the KVM vCPU ID as the cpu_id

Preemption Detection:
  - Compares the last context the CPU ran against the current context being
 scheduled.
  - Compares the context's rseqCPU and CPU ID retrieved by the platform

To facilitate this, several new methods are introduced to the platform
interface and implemented by platform/kvm:
 - HasCpuNumbers()
 - NumCPUs()
 - DetectsCPUPreemption()
 - PreemptCpu()
 - PreemptAllCpus()

A new user-facing KVM specific option, use_cpu_numbers, is introduced to
effectively feature guard RSEQ. When use_cpu_numbers is true, vCPU IDs are
exposed to the guest, enabling RSEQ support. Setting it to false reverts to
the previous behavior and effectively disables RSEQ for applications that
rely on it.

PiperOrigin-RevId: 823228013

Author: tranji-cloud

pkg/sentry/kernel/kernel.go

@@ -490,6 +490,11 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 	k.cpuClockTickerWakeCh = make(chan struct{}, 1)
 	k.cpuClockTickerStopCond.L = &k.runningTasksMu
 	k.applicationCores = args.ApplicationCores
+	if args.UseHostCores && k.HasCPUNumbers() {
+		args.UseHostCores = false
+		log.Infof("UseHostCores enabled but the platform implements HasCPUNumbers(): setting UseHostCores to false")
+	}
+
 	if args.UseHostCores {
 		k.useHostCores = true
 		maxCPU, err := hostcpu.MaxPossibleCPU()
@@ -502,6 +507,15 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 			k.applicationCores = minAppCores
 		}
 	}
+
+	if k.HasCPUNumbers() {
+		if k.applicationCores < uint(k.NumCPUs()) {
+			log.Infof("ApplicationCores is less than NumCPUs: %d < %d", k.applicationCores, k.NumCPUs())
+			log.Infof("Setting applicationCores to NumCPUs: %d", k.NumCPUs())
+			k.applicationCores = uint(k.NumCPUs())
+		}
+	}
+
 	k.extraAuxv = args.ExtraAuxv
 	k.vdso = args.Vdso
 	k.vdsoParams = args.VdsoParams

pkg/sentry/kernel/rseq.go

@@ -20,7 +20,6 @@ import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
-	"gvisor.dev/gvisor/pkg/sentry/hostcpu"
 	"gvisor.dev/gvisor/pkg/usermem"
 )
 
@@ -50,7 +49,7 @@ type OldRSeqCriticalRegion struct {
 
 // RSeqAvailable returns true if t supports (old and new) restartable sequences.
 func (t *Task) RSeqAvailable() bool {
-	return t.k.useHostCores && t.k.Platform.DetectsCPUPreemption()
+	return (t.k.useHostCores || t.k.Platform.HasCPUNumbers()) && t.k.Platform.DetectsCPUPreemption()
 }
 
 // SetRSeq registers addr as this thread's rseq structure.
@@ -201,7 +200,7 @@ func (t *Task) rseqUpdateCPU() error {
 		return nil
 	}
 
-	t.rseqCPU = int32(hostcpu.GetCPU())
+	t.rseqCPU = t.CPU()
 
 	// Update both CPUs, even if one fails.
 	rerr := t.rseqCopyOutCPU()

pkg/sentry/kernel/task_run.go

@@ -24,7 +24,6 @@ import (
 	"gvisor.dev/gvisor/pkg/goid"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/refs"
-	"gvisor.dev/gvisor/pkg/sentry/hostcpu"
 	"gvisor.dev/gvisor/pkg/sentry/ktime"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
 	"gvisor.dev/gvisor/pkg/sentry/platform"
@@ -207,7 +206,7 @@ func (app *runApp) execute(t *Task) taskRunState {
 	if t.rseqPreempted {
 		t.rseqPreempted = false
 		if t.rseqAddr != 0 || t.oldRSeqCPUAddr != 0 {
-			t.rseqCPU = int32(hostcpu.GetCPU())
+			t.rseqCPU = t.CPU()
 			if err := t.rseqCopyOutCPU(); err != nil {
 				t.Debugf("Failed to copy CPU to %#x for rseq: %v", t.rseqAddr, err)
 				t.forceSignal(linux.SIGSEGV, false)

pkg/sentry/kernel/task_sched.go

@@ -365,7 +365,7 @@ func (t *Task) SetCPUMask(mask sched.CPUSet) error {
 		return linuxerr.EINVAL
 	}
 
-	if t.k.useHostCores {
+	if t.k.useHostCores || t.k.Platform.HasCPUNumbers() {
 		// No-op; pretend the mask was immediately changed back.
 		return nil
 	}
@@ -383,6 +383,10 @@ func (t *Task) SetCPUMask(mask sched.CPUSet) error {
 
 // CPU returns the cpu id for a given task.
 func (t *Task) CPU() int32 {
+	if t.k.Platform.HasCPUNumbers() {
+		return t.p.LastCPUNumber()
+	}
+
 	if t.k.useHostCores {
 		return int32(hostcpu.GetCPU())
 	}

pkg/sentry/platform/kvm/config.go

@@ -23,6 +23,11 @@ type Config struct {
 	// kernel.InitKernelArgs. It is necessary to forward it to KVM in order
 	// to initialize the correct amount of vCPUs.
 	ApplicationCores int
+
+	// UseCPUNums use KVM CPU numbers as CPU numbers in the sentry.
+	// This is necessary to support features like rseq and KVM based
+	// preemption.
+	UseCPUNums bool
 }
 
 func (*machine) applyConfig(config *Config) error { return nil }

pkg/sentry/platform/kvm/context.go

@@ -16,6 +16,7 @@ package kvm
 
 import (
 	"gvisor.dev/gvisor/pkg/abi/linux"
+	"gvisor.dev/gvisor/pkg/atomicbitops"
 	pkgcontext "gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/hostarch"
 	"gvisor.dev/gvisor/pkg/ring0"
@@ -36,6 +37,9 @@ type platformContext struct {
 
 	// interrupt is the interrupt platformContext.
 	interrupt interrupt.Forwarder
+
+	// lastUsedCPU is the last CPU ID used by this platformContext.
+	lastUsedCPU atomicbitops.Int32
 }
 
 // tryCPUIDError indicates that CPUID emulation should occur.
@@ -45,7 +49,7 @@ type tryCPUIDError struct{}
 func (tryCPUIDError) Error() string { return "cpuid emulation failed" }
 
 // Switch runs the provided platformContext in the given address space.
-func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, _ int32) (*linux.SignalInfo, hostarch.AccessType, error) {
+func (c *platformContext) Switch(ctx pkgcontext.Context, mm platform.MemoryManager, ac *arch.Context64, rseqCPU int32) (*linux.SignalInfo, hostarch.AccessType, error) {
 	as := mm.AddressSpace()
 	localAS := as.(*addressSpace)
 
@@ -58,6 +62,20 @@ restart:
 		c.machine.Put(cpu) // Already preempted.
 		return nil, hostarch.NoAccess, platform.ErrContextInterrupt
 	}
+	// If this CPU was last used to run a different context
+	// or if this context last ran on a different CPU, then we've
+	// been preempted.
+	last := cpu.lastCtx.Swap(c)
+	c.lastUsedCPU.Store(int32(cpu.id))
+	preempted := rseqCPU >= 0 && (last != c || rseqCPU != int32(cpu.id))
+	if preempted {
+		// Release resources.
+		c.machine.Put(cpu)
+
+		// All done.
+		c.interrupt.Disable()
+		return nil, hostarch.NoAccess, platform.ErrContextCPUPreempted
+	}
 
 	// Set the active address space.
 	//
@@ -136,3 +154,8 @@ func (c *platformContext) PullFullState(as platform.AddressSpace, ac *arch.Conte
 
 // PrepareSleep implements platform.Context.platform.Context.
 func (*platformContext) PrepareSleep() {}
+
+// LastCPUNumber implements platform.Context.LastCPUNumber.
+func (c *platformContext) LastCPUNumber() int32 {
+	return c.lastUsedCPU.Load()
+}

pkg/sentry/platform/kvm/kvm.go

@@ -62,8 +62,6 @@ type runData struct {
 
 // KVM represents a lightweight VM context.
 type KVM struct {
-	platform.NoCPUPreemptionDetection
-
 	// KVM never changes mm_structs.
 	platform.UseHostProcessMemoryBarrier
 
@@ -180,6 +178,41 @@ func (k *KVM) ConcurrencyCount() int {
 	return k.machine.maxVCPUs
 }
 
+// HasCPUNumbers implements platform.Platform.HasCPUNumbers.
+func (k *KVM) HasCPUNumbers() bool {
+	return k.machine.useCPUNums
+}
+
+// NumCPUs implements platform.Platform.NumCPUs.
+func (k *KVM) NumCPUs() int32 {
+	if !k.HasCPUNumbers() {
+		panic("platform is not configured to use CPU numbers")
+	}
+	return int32(k.machine.maxVCPUs)
+}
+
+// DetectsCPUPreemption implements platform.Platform.DetectsCPUPreemption.
+func (k *KVM) DetectsCPUPreemption() bool {
+	return true
+}
+
+// PreemptAllCPUs implements platform.Platform.PreemptAllCPUs.
+func (k *KVM) PreemptAllCPUs() error {
+	for _, c := range k.machine.vCPUsByID {
+		c.lastCtx.Store(nil)
+		c.BounceToHost()
+	}
+	return nil
+}
+
+// PreemptCPU implements platform.Platform.PreemptCPU.
+func (k *KVM) PreemptCPU(cpu int32) error {
+	c := k.machine.vCPUsByID[cpu]
+	c.lastCtx.Store(nil)
+	c.BounceToHost()
+	return nil
+}
+
 // NewContext returns an interruptible context.
 func (k *KVM) NewContext(pkgcontext.Context) platform.Context {
 	return &platformContext{
@@ -192,6 +225,7 @@ type constructor struct{}
 func (*constructor) New(opts platform.Options) (platform.Platform, error) {
 	return New(opts.DeviceFile, Config{
 		ApplicationCores: opts.ApplicationCores,
+		UseCPUNums:       opts.UseCPUNums,
 	})
 }
 

pkg/sentry/platform/kvm/machine.go

@@ -89,6 +89,9 @@ type machine struct {
 
 	// usedSlots is the set of used physical addresses (not sorted).
 	usedSlots []uintptr
+
+	// useCPUNums indicates whether to enable the use vCPU numbers as CPU numbers.
+	useCPUNums bool
 }
 
 const (
@@ -216,6 +219,9 @@ type vCPU struct {
 
 	// dieState holds state related to vCPU death.
 	dieState dieState
+
+	// lastCtx is the last context that was scheduled on this vCPU
+	lastCtx atomic.Pointer[platformContext]
 }
 
 type dieState struct {
@@ -275,6 +281,7 @@ func newMachine(vm int, config *Config) (*machine, error) {
 	m := &machine{
 		fd:               vm,
 		applicationCores: config.ApplicationCores,
+		useCPUNums:       config.UseCPUNums,
 	}
 	m.available.L = &m.mu
 

pkg/sentry/platform/platform.go

@@ -96,6 +96,17 @@ type Platform interface {
 	// NewContext returns a new execution context.
 	NewContext(context.Context) Context
 
+	// PreemptAllCPUs causes all concurrent calls to Context.Switch() on the given CPU, as well
+	// as the first following call to Context.Switch() for each Context, to
+	// return ErrContextCPUPreempted.
+	//
+	// Precondition(s): cpu must be in the range [0, NumCPUs()).
+	//
+	// PreemptCPU is only supported if DetectsCPUPremption() && HasCPUNumbers() == true.
+	// Platforms for which this does not hold may panic if PreemptCPU is
+	// called.
+	PreemptCPU(cpu int32) error
+
 	// PreemptAllCPUs causes all concurrent calls to Context.Switch(), as well
 	// as the first following call to Context.Switch() for each Context, to
 	// return ErrContextCPUPreempted.
@@ -121,6 +132,12 @@ type Platform interface {
 	// in parallel. Concurrent calls to Context.Switch() beyond
 	// ConcurrencyCount() may block until previous calls have returned.
 	ConcurrencyCount() int
+
+	// HasCPUNumbers returns true if the platform assigns CPU numbers to contexts.
+	HasCPUNumbers() bool
+
+	// NumCPUs returns the number of CPUs on the platform.
+	NumCPUs() int32
 }
 
 // NoCPUPreemptionDetection implements Platform.DetectsCPUPreemption and
@@ -137,6 +154,25 @@ func (NoCPUPreemptionDetection) PreemptAllCPUs() error {
 	panic("This platform does not support CPU preemption detection")
 }
 
+// NoCPUNumbers implements Platform.HasCPUNumbers for platforms that do
+// not support it.
+type NoCPUNumbers struct{}
+
+// HasCPUNumbers implements Platform.HasCPUNumbers.
+func (NoCPUNumbers) HasCPUNumbers() bool {
+	return false
+}
+
+// NumCPUs implements Platform.NumCPUs.
+func (NoCPUNumbers) NumCPUs() int32 {
+	panic("platform does not support CPU numbers")
+}
+
+// PreemptCPU implements Platform.PreemptCPU.
+func (NoCPUNumbers) PreemptCPU(cpu int32) error {
+	panic("platform does not support preempting a specific CPU")
+}
+
 // UseHostGlobalMemoryBarrier implements Platform.HaveGlobalMemoryBarrier and
 // Platform.GlobalMemoryBarrier by invoking equivalent functionality on the
 // host.
@@ -264,6 +300,16 @@ type Context interface {
 	// PrepareSleep() is called when the thread switches to the
 	// interruptible sleep state.
 	PrepareSleep()
+
+	// LastCPUNumber returns the last CPU number that this context was running on.
+	// If the context never ran on a CPU, it may return any valid CPU number, as long as the first
+	// call to Switch will detect that the CPU number is incorrect and return ErrContextCPUPreempted.
+	LastCPUNumber() int32
+}
+
+// LastCPUNumber implements Context.LastCPUNumber.
+func (NoCPUNumbers) LastCPUNumber() int32 {
+	panic("context does not support last CPU number")
 }
 
 // ContextError is one of the possible errors returned by Context.Switch().
@@ -538,6 +584,11 @@ type Options struct {
 	// ApplicationCores is used by KVM to determine the correct amount of
 	// vCPUs to create.
 	ApplicationCores int
+
+	// UseCPUNums is used by KVM to determine whether to use KVM CPU numbers
+	// as CPU numbers in the sentry. This is necessary to support features like
+	// rseq
+	UseCPUNums bool
 }
 
 // Constructor represents a platform type.

pkg/sentry/platform/ptrace/ptrace.go

@@ -74,6 +74,7 @@ var (
 
 type context struct {
 	archContext
+	platform.NoCPUNumbers
 
 	// signalInfo is the signal info, if and when a signal is received.
 	signalInfo linux.SignalInfo
@@ -214,6 +215,7 @@ type PTrace struct {
 	platform.MMapMinAddr
 	platform.NoCPUPreemptionDetection
 	platform.UseHostGlobalMemoryBarrier
+	platform.NoCPUNumbers
 }
 
 // New returns a new ptrace-based implementation of the platform interface.