Perform explicit CPUID and FP state compatibility checks on restore

PiperOrigin-RevId: 245341004
Change-Id: Ic4d581039d034a8ae944b43e45e84eb2c3973657

3 files changed, 110 insertions, 13 deletions

diff --git a/pkg/cpuid/cpuid.go b/pkg/cpuid/cpuid.go
index 64e2e68f1..61441150e 100644
--- a/pkg/cpuid/cpuid.go
+++ b/pkg/cpuid/cpuid.go
@@ -446,6 +446,20 @@ const (
 	extendedFeatures                                       // Returns some extended feature bits in edx and ecx.
 )
 
+// These are the extended floating point state features. They are used to
+// enumerate floating point features in XCR0, XSTATE_BV, etc.
+const (
+	XSAVEFeatureX87         = 1 << 0
+	XSAVEFeatureSSE         = 1 << 1
+	XSAVEFeatureAVX         = 1 << 2
+	XSAVEFeatureBNDREGS     = 1 << 3
+	XSAVEFeatureBNDCSR      = 1 << 4
+	XSAVEFeatureAVX512op    = 1 << 5
+	XSAVEFeatureAVX512zmm0  = 1 << 6
+	XSAVEFeatureAVX512zmm16 = 1 << 7
+	XSAVEFeaturePKRU        = 1 << 9
+)
+
 var cpuFreqMHz float64
 
 // x86FeaturesFromString includes features from x86FeatureStrings and
@@ -561,6 +575,26 @@ func (fs *FeatureSet) Intel() bool {
 	return fs.VendorID == "GenuineIntel"
 }
 
+// ErrIncompatible is returned by FeatureSet.HostCompatible if fs is not a
+// subset of the host feature set.
+type ErrIncompatible struct {
+	message string
+}
+
+// Error implements error.
+func (e ErrIncompatible) Error() string {
+	return e.message
+}
+
+// CheckHostCompatible returns nil if fs is a subset of the host feature set.
+func (fs *FeatureSet) CheckHostCompatible() error {
+	hfs := HostFeatureSet()
+	if diff := fs.Subtract(hfs); diff != nil {
+		return ErrIncompatible{fmt.Sprintf("CPU feature set %v incompatible with host feature set %v (missing: %v)", fs.FlagsString(false), hfs.FlagsString(false), diff)}
+	}
+	return nil
+}
+
 // Helper to convert 3 regs into 12-byte vendor ID.
 func vendorIDFromRegs(bx, cx, dx uint32) string {
 	bytes := make([]byte, 0, 12)
diff --git a/pkg/sentry/arch/arch_state_x86.go b/pkg/sentry/arch/arch_state_x86.go
index 604bd08a6..01949049d 100644
--- a/pkg/sentry/arch/arch_state_x86.go
+++ b/pkg/sentry/arch/arch_state_x86.go
@@ -15,14 +15,31 @@
 package arch
 
 import (
-	"sync"
+	"fmt"
 	"syscall"
 
-	"gvisor.googlesource.com/gvisor/pkg/log"
+	"gvisor.googlesource.com/gvisor/pkg/cpuid"
+	"gvisor.googlesource.com/gvisor/pkg/sentry/usermem"
 )
 
-// warnOnce is used to warn about truncated state only once.
-var warnOnce sync.Once
+// ErrFloatingPoint indicates a failed restore due to unusable floating point
+// state.
+type ErrFloatingPoint struct {
+	// supported is the supported floating point state.
+	supported uint64
+
+	// saved is the saved floating point state.
+	saved uint64
+}
+
+// Error returns a sensible description of the restore error.
+func (e ErrFloatingPoint) Error() string {
+	return fmt.Sprintf("floating point state contains unsupported features; supported: %#x saved: %#x", e.supported, e.saved)
+}
+
+// XSTATE_BV does not exist if FXSAVE is used, but FXSAVE implicitly saves x87
+// and SSE state, so this is the equivalent XSTATE_BV value.
+const fxsaveBV uint64 = cpuid.XSAVEFeatureX87 | cpuid.XSAVEFeatureSSE
 
 // afterLoad is invoked by stateify.
 func (s *State) afterLoad() {
@@ -33,7 +50,8 @@ func (s *State) afterLoad() {
 	// state that may be saved by the new CPU. Even if extraneous new state
 	// is saved, the state we care about is guaranteed to be a subset of
 	// new state. Later optimizations can use less space when using a
-	// smaller state component bitmap. Intel SDM section 13 has more info.
+	// smaller state component bitmap. Intel SDM Volume 1 Chapter 13 has
+	// more info.
 	s.x86FPState = newX86FPState()
 
 	// x86FPState always contains all the FP state supported by the host.
@@ -41,15 +59,30 @@ func (s *State) afterLoad() {
 	// which we cannot restore.
 	//
 	// The x86 FP state areas are backwards compatible, so we can simply
-	// truncate the additional floating point state. Applications should
-	// not depend on the truncated state because it should relate only to
-	// features that were not exposed in the app FeatureSet.
+	// truncate the additional floating point state.
+	//
+	// Applications should not depend on the truncated state because it
+	// should relate only to features that were not exposed in the app
+	// FeatureSet. However, because we do not *prevent* them from using
+	// this state, we must verify here that there is no in-use state
+	// (according to XSTATE_BV) which we do not support.
 	if len(s.x86FPState) < len(old) {
-		warnOnce.Do(func() {
-			// This will occur on every instance of state, don't
-			// bother warning more than once.
-			log.Infof("dropping %d bytes of floating point state; the application should not depend on this state", len(old)-len(s.x86FPState))
-		})
+		// What do we support?
+		supportedBV := fxsaveBV
+		if fs := cpuid.HostFeatureSet(); fs.UseXsave() {
+			supportedBV = fs.ValidXCR0Mask()
+		}
+
+		// What was in use?
+		savedBV := fxsaveBV
+		if len(old) >= xstateBVOffset+8 {
+			savedBV = usermem.ByteOrder.Uint64(old[xstateBVOffset:])
+		}
+
+		// Supported features must be a superset of saved features.
+		if savedBV&^supportedBV != 0 {
+			panic(ErrFloatingPoint{supported: supportedBV, saved: savedBV})
+		}
 	}
 
 	// Copy to the new, aligned location.
diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index ee6334509..a1b2d7161 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/pkg/sentry/kernel/kernel.go
@@ -337,6 +337,17 @@ func (k *Kernel) SaveTo(w io.Writer) error {
 		return fmt.Errorf("failed to invalidate unsavable mappings: %v", err)
 	}
 
+	// Save the CPUID FeatureSet before the rest of the kernel so we can
+	// verify its compatibility on restore before attempting to restore the
+	// entire kernel, which may fail on an incompatible machine.
+	//
+	// N.B. This will also be saved along with the full kernel save below.
+	cpuidStart := time.Now()
+	if err := state.Save(w, k.FeatureSet(), nil); err != nil {
+		return err
+	}
+	log.Infof("CPUID save took [%s].", time.Since(cpuidStart))
+
 	// Save the kernel state.
 	kernelStart := time.Now()
 	var stats state.Stats
@@ -469,6 +480,25 @@ func (k *Kernel) LoadFrom(r io.Reader, net inet.Stack) error {
 
 	initAppCores := k.applicationCores
 
+	// Load the pre-saved CPUID FeatureSet.
+	//
+	// N.B. This was also saved along with the full kernel below, so we
+	// don't need to explicitly install it in the Kernel.
+	cpuidStart := time.Now()
+	var features cpuid.FeatureSet
+	if err := state.Load(r, &features, nil); err != nil {
+		return err
+	}
+	log.Infof("CPUID load took [%s].", time.Since(cpuidStart))
+
+	// Verify that the FeatureSet is usable on this host. We do this before
+	// Kernel load so that the explicit CPUID mismatch error has priority
+	// over floating point state restore errors that may occur on load on
+	// an incompatible machine.
+	if err := features.CheckHostCompatible(); err != nil {
+		return err
+	}
+
 	// Load the kernel state.
 	kernelStart := time.Now()
 	var stats state.Stats