fpu: verify that a restored fpu state can be loaded by xrstor

PiperOrigin-RevId: 538886033

pkg/safecopy/BUILD

@@ -15,9 +15,12 @@ go_library(
         "memcpy_amd64.s",
         "memcpy_arm64.s",
         "safecopy.go",
+        "safecopy_amd64_unsafe.go",
+        "safecopy_arm64.go",
         "safecopy_unsafe.go",
         "sighandler_amd64.s",
         "sighandler_arm64.s",
+        "xrstor_amd64.s",
     ],
     visibility = ["//:sandbox"],
     deps = [
@@ -34,5 +37,26 @@ go_test(
         "safecopy_test.go",
     ],
     library = ":safecopy",
-    deps = ["@org_golang_x_sys//unix:go_default_library"],
+    deps = [
+        "@org_golang_x_sys//unix:go_default_library",
+    ],
+)
+
+go_test(
+    name = "safecopy_x_test",
+    srcs = [
+        "safecopy_amd64_test.go",
+        "safecopy_arm64_test.go",
+    ],
+    # NOTE: It seems that bazel code generation does not properly parse tags
+    # when run via the architecture transition for nogo. This should be fixed
+    # at some point in the future, but for now we can simply skip nogo analysis
+    # on the test itself. It still applies to the core library.
+    nogo = False,
+    tags = ["not_run:arm"],
+    deps = [
+        "//pkg/cpuid",
+        "//pkg/safecopy",
+        "//pkg/sentry/arch/fpu",
+    ],
 )

pkg/safecopy/safecopy.go

@@ -129,6 +129,7 @@ func initializeAddresses() {
 	compareAndSwapUint32End = FindEndAddress(compareAndSwapUint32Begin)
 	loadUint32Begin = addrOfLoadUint32()
 	loadUint32End = FindEndAddress(loadUint32Begin)
+	initializeArchAddresses()
 }
 
 func init() {

pkg/safecopy/safecopy_amd64_test.go

@@ -0,0 +1,49 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build amd64 || i386
+// +build amd64 i386
+
+package safecopy_x_test
+
+import (
+	"testing"
+	"unsafe"
+
+	"gvisor.dev/gvisor/pkg/cpuid"
+	"gvisor.dev/gvisor/pkg/safecopy"
+	"gvisor.dev/gvisor/pkg/sentry/arch/fpu"
+)
+
+func TestCheckXstateFault(t *testing.T) {
+	cpuid.Initialize()
+	state := fpu.NewState()
+	state.SetMXCSR(0xffffff) // Invalid value
+	err := safecopy.CheckXstate(state.BytePointer())
+	if want := (safecopy.SegvError{uintptr(unsafe.Pointer(state.BytePointer()))}); err != want {
+		t.Errorf("Unexpected error: got %v, want %v", err, want)
+	}
+}
+
+func TestCheckXstateSuccess(t *testing.T) {
+	cpuid.Initialize()
+	if !cpuid.HostFeatureSet().UseXsave() {
+		t.Skip("xsave isn't supported")
+	}
+	state := fpu.NewState()
+	err := safecopy.CheckXstate(state.BytePointer())
+	if err != nil {
+		t.Errorf("Unexpected error: %v", err)
+	}
+}

pkg/safecopy/safecopy_amd64_unsafe.go

@@ -0,0 +1,42 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build amd64 || i386
+// +build amd64 i386
+
+package safecopy
+
+import (
+	"unsafe"
+)
+
+var (
+	checkXstateBegin uintptr
+	checkXstateEnd   uintptr
+)
+
+func initializeArchAddresses() {
+	checkXstateBegin = addrOfCheckXstate()
+	checkXstateEnd = FindEndAddress(checkXstateBegin)
+}
+
+//go:noescape
+func checkXstate(addr uintptr) (fault uintptr, sig int32, mxcsr uint32, cw uint16)
+func addrOfCheckXstate() uintptr
+
+// CheckXstate verifies that xstate can be restored by the xrstor instruction.
+func CheckXstate(state *byte) error {
+	_, sig, _, _ := checkXstate(uintptr(unsafe.Pointer(state)))
+	return errorFromFaultSignal(uintptr(unsafe.Pointer(state)), sig)
+}

pkg/safecopy/safecopy_arm64.go

@@ -0,0 +1,21 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build arm64
+// +build arm64
+
+package safecopy
+
+func initializeArchAddresses() {
+}

pkg/safecopy/safecopy_arm64_test.go

@@ -0,0 +1,25 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build arm64
+// +build arm64
+
+package safecopy_x_test
+
+import (
+	"testing"
+)
+
+func TestDummy(t *testing.T) {
+}

pkg/safecopy/sighandler_amd64.s

@@ -110,6 +110,14 @@ not_casuint32:
 	JMP handle_fault
 
 not_loaduint32:
+	CMPQ CX, ·checkXstateBegin(SB)
+	JB not_checkXstate
+	CMPQ CX, ·checkXstateEnd(SB)
+	JAE not_checkXstate
+
+	LEAQ handleCheckXstateFault(SB), CX
+	JMP handle_fault
+not_checkXstate:
 original_handler:
 	// Jump to the previous signal handler, which is likely the golang one.
 	XORQ CX, CX

pkg/safecopy/xrstor_amd64.s

@@ -0,0 +1,59 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "textflag.h"
+
+// handleCheckXstateFault returns (the value stored in AX, the value stored in DI).
+// Control is transferred to it when checkXstate below receives SIGSEGV or SIGBUS,
+// with the faulting address stored in AX and the signal number stored in DI.
+//
+// It must have the same frame configuration as memcpy so that it can undo any
+// potential call frame set up by the assembler.
+TEXT handleCheckXstateFault(SB), NOSPLIT|NOFRAME, $0-26
+	MOVQ	AX, addr+8(FP)
+	MOVL	DI, sig+16(FP)
+
+	LDMXCSR	mxcsr+20(FP)
+	BYTE $0xDB; BYTE $0xE2; // FNCLEX
+	FLDCW cw+24(FP)
+	RET
+
+
+// ·checkXstate verifies that the specified floating point state can be loaded.
+TEXT ·checkXstate(SB),NOSPLIT|NOFRAME,$0-26
+	// Store 0 as the returned signal number. If we run to completion,
+	// this is the value the caller will see; if a signal is received,
+	// handleMemcpyFault will store a different value in this address.
+	MOVL	$0, sig+16(FP)
+	// MXCSR and the x87 control word are the only floating point state
+	// that is callee-save and thus we must save.
+	STMXCSR mxcsr+20(FP)
+	FSTCW	cw+24(FP)
+
+	MOVQ addr+0(FP), DI
+	MOVL $0xffffffff, AX
+	MOVL $0xffffffff, DX
+	XRSTOR (DI)
+
+	// Restore MXCSR and the x87 control word.
+	LDMXCSR	mxcsr+20(FP)
+	BYTE $0xDB; BYTE $0xE2; // FNCLEX
+	FLDCW cw+24(FP)
+	RET
+
+// func addrOfCheckXstate() uintptr
+TEXT ·addrOfCheckXstate(SB), $0-8
+	MOVQ	$·checkXstate(SB), AX
+	MOVQ	AX, ret+0(FP)
+	RET

pkg/sentry/arch/fpu/BUILD

@@ -21,6 +21,7 @@ go_library(
         "//pkg/cpuid",
         "//pkg/errors/linuxerr",
         "//pkg/hostarch",
+        "//pkg/safecopy",
         "//pkg/sync",
         "@org_golang_x_sys//unix:go_default_library",
     ],

pkg/sentry/arch/fpu/fpu_amd64.go

@@ -18,12 +18,14 @@
 package fpu
 
 import (
+	"fmt"
 	"io"
 
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/cpuid"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/hostarch"
+	"gvisor.dev/gvisor/pkg/safecopy"
 	"gvisor.dev/gvisor/pkg/sync"
 )
 
@@ -302,6 +304,11 @@ func (s *State) SetMXCSR(mxcsr uint32) {
 	hostarch.ByteOrder.PutUint32((*s)[mxcsrOffset:], mxcsr)
 }
 
+// GetMXCSR gets the MXCSR control/status register in the state.
+func (s *State) GetMXCSR() uint32 {
+	return hostarch.ByteOrder.Uint32((*s)[mxcsrOffset:])
+}
+
 // BytePointer returns a pointer to the first byte of the state.
 //
 //go:nosplit
@@ -339,25 +346,35 @@ func (s *State) AfterLoad() {
 	// 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) < len(old) {
-		// What do we support?
-		supportedBV := fxsaveBV
-		if fs := cpuid.HostFeatureSet(); fs.UseXsave() {
-			supportedBV = fs.ValidXCR0Mask()
-		}
+	// 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 = hostarch.ByteOrder.Uint64(old[xstateBVOffset:])
-		}
+	// What was in use?
+	savedBV := fxsaveBV
+	if len(old) >= xstateBVOffset+8 {
+		savedBV = hostarch.ByteOrder.Uint64(old[xstateBVOffset:])
+	}
 
-		// Supported features must be a superset of saved features.
-		if savedBV&^supportedBV != 0 {
-			panic(ErrLoadingState{supportedFeatures: supportedBV, savedFeatures: savedBV})
-		}
+	// Supported features must be a superset of saved features.
+	if savedBV&^supportedBV != 0 {
+		panic(ErrLoadingState{supportedFeatures: supportedBV, savedFeatures: savedBV})
 	}
 
 	// Copy to the new, aligned location.
 	copy(*s, old)
+
+	mxcsrBefore := s.GetMXCSR()
+	sanitizeMXCSR(*s)
+	mxcsrAfter := s.GetMXCSR()
+	if mxcsrBefore != mxcsrAfter {
+		panic(fmt.Sprintf("incompatible mxcsr value: %x (%x)", mxcsrBefore, mxcsrAfter))
+	}
+	if fs := cpuid.HostFeatureSet(); fs.UseXsave() {
+		if err := safecopy.CheckXstate(s.BytePointer()); err != nil {
+			panic(fmt.Sprintf("incompatible state: %s (%#v)", err, *s))
+		}
+	}
 }