From bb840068163c9c32889d8a35aaed076444f67000 Mon Sep 17 00:00:00 2001 From: Konstantin Bogomolov Date: Thu, 18 Jan 2024 20:09:06 -0800 Subject: [PATCH] Fixup AMX workaround for ptrace. SETREGSET/GETREGSET expect AMX portions of fpstate to always be used. For this reason we need to allocate enough memory for this to happen, even if we never populate the AMX portions within initX86FPState. PiperOrigin-RevId: 599702181 --- pkg/cpuid/cpuid_amd64.go | 8 ++------ pkg/ring0/entry_amd64.s | 7 +------ pkg/ring0/lib_amd64.go | 3 ++- pkg/sentry/arch/fpu/fpu_amd64.go | 11 +++++++++-- pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go | 3 ++- 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/pkg/cpuid/cpuid_amd64.go b/pkg/cpuid/cpuid_amd64.go index 16c55dff4..f5fd4ae34 100644 --- a/pkg/cpuid/cpuid_amd64.go +++ b/pkg/cpuid/cpuid_amd64.go @@ -384,17 +384,13 @@ const ( // Extended state includes floating point registers, and other cpu state that's // not associated with the normal task context. // -// We do not support enabling AMX within gVisor, therefore always exclude -// AMXExtendedStateSize. -// TODO(gvisor.dev/issues/9896): Implement AMX Support. -// // Note: the return value matches the size of signal FP state frames. // Look at check_xstate_in_sigframe() in the kernel sources for more details. // //go:nosplit func (fs FeatureSet) ExtendedStateSize() (size, align uint) { if fs.UseXsave() { - return uint(xsaveSize) - fs.AMXExtendedStateSize(), 64 + return uint(xsaveSize), 64 } // If we don't support xsave, we fall back to fxsave, which requires @@ -423,7 +419,7 @@ func (fs FeatureSet) ValidXCR0Mask() uint64 { return 0 } ax, _, _, dx := fs.query(xSaveInfo) - return (uint64(dx)<<32 | uint64(ax)) ^ XCR0AMXMask + return (uint64(dx)<<32 | uint64(ax)) &^ XCR0AMXMask } // UseXsave returns the choice of fp state saving instruction. diff --git a/pkg/ring0/entry_amd64.s b/pkg/ring0/entry_amd64.s index ad49618b1..218eb3fea 100644 --- a/pkg/ring0/entry_amd64.s +++ b/pkg/ring0/entry_amd64.s @@ -233,9 +233,7 @@ TEXT ·doSwitchToUser(SB),NOSPLIT,$16-48 MOVB ·hasXSAVE(SB), BX TESTB BX, BX JZ no_xrstor - // Use xrstor to restore all available fp state. For now, we restore - // everything unconditionally by setting the implicit operand edx:eax - // (the "requested feature bitmap") to all 1's. + // Use xrstor to restore all available fp state. MOVL $XCR0_EAX, AX MOVL $XCR0_EDX, DX BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x2f // XRSTOR64 0(DI) @@ -281,7 +279,6 @@ done_sysret_or_iret: TESTB BX, BX JZ no_xsave // Use xsave/xsaveopt to save all extended state. - // We save everything unconditionally by setting RFBM to all 1's. MOVL $XCR0_EAX, AX MOVL $XCR0_EDX, DX TESTB CX, CX @@ -519,7 +516,6 @@ kernel: TESTB BX, BX JZ no_xsave // Use xsave/xsaveopt to save all extended state. - // We save everything unconditionally by setting RFBM to all 1's. MOVL $XCR0_EAX, AX MOVL $XCR0_EDX, DX TESTB CX, CX @@ -651,7 +647,6 @@ kernel: TESTB BX, BX JZ no_xsave // Use xsave/xsaveopt to save all extended state. - // We save everything unconditionally by setting RFBM to all 1's. MOVL $XCR0_EAX, AX MOVL $XCR0_EDX, DX TESTB CX, CX diff --git a/pkg/ring0/lib_amd64.go b/pkg/ring0/lib_amd64.go index 585a5fa2a..73f774667 100644 --- a/pkg/ring0/lib_amd64.go +++ b/pkg/ring0/lib_amd64.go @@ -119,7 +119,8 @@ func Init(fs cpuid.FeatureSet) { hasFSGSBASE = fs.HasFeature(cpuid.X86FeatureFSGSBase) validXCR0Mask = uintptr(fs.ValidXCR0Mask()) if hasXSAVE { - localXCR0 = xgetbv(0) + XCR0AMXMask := uintptr((1 << 17) | (1 << 18)) + localXCR0 = xgetbv(0) &^ XCR0AMXMask } } diff --git a/pkg/sentry/arch/fpu/fpu_amd64.go b/pkg/sentry/arch/fpu/fpu_amd64.go index 6d4bfdfc5..a7e4bf3c3 100644 --- a/pkg/sentry/arch/fpu/fpu_amd64.go +++ b/pkg/sentry/arch/fpu/fpu_amd64.go @@ -70,8 +70,13 @@ const ( func initX86FPState(data *byte, useXsave bool) func newX86FPStateSlice() State { - size, align := cpuid.HostFeatureSet().ExtendedStateSize() - capacity := size + FP_XSTATE_MAGIC2_SIZE + maxsize, align := cpuid.HostFeatureSet().ExtendedStateSize() + // We need capacity to be large enough to hold AMX bytes because of + // ptrace. PTRACE_SETREGSET/GETREGSET assume that AMX portions should + // always be used. + // TODO(gvisor.dev/issues/9896): Implement AMX Support. + capacity := maxsize + FP_XSTATE_MAGIC2_SIZE + size := maxsize - cpuid.HostFeatureSet().AMXExtendedStateSize() // Always use at least 4096 bytes. // // For the KVM platform, this state is a fixed 4096 bytes, so make sure @@ -129,6 +134,8 @@ func InitHostState() { hostXCR0Mask = featureSet.ValidXCR0Mask() hostUseXsave = featureSet.UseXsave() hostFPSize, _ = featureSet.ExtendedStateSize() + // TODO(gvisor.dev/issues/9896): Implement AMX Support. + hostFPSize = hostFPSize - featureSet.AMXExtendedStateSize() }) } diff --git a/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go b/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go index 8d14eeeda..4732391ed 100644 --- a/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go +++ b/pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go @@ -52,7 +52,8 @@ func (s *ArchState) Init() { fs := cpuid.HostFeatureSet() fpLenUint, _ := fs.ExtendedStateSize() - s.fpLen = uint32(fpLenUint) + // TODO(gvisor.dev/issues/9896): Implement AMX Support. + s.fpLen = uint32(fpLenUint - fs.AMXExtendedStateSize()) if fs.UseXsaveopt() { s.xsaveMode = xsaveopt } else if fs.UseXsave() {