From e9bdc76c02bd6ad4c6af574fa1dd31578bec012e Mon Sep 17 00:00:00 2001 From: Konstantin Bogomolov Date: Wed, 17 Jan 2024 15:08:05 -0800 Subject: [PATCH] Exclude AMX extended state from being xsave/xrstor'd. For now we are going to completely disable using AMX, so we will always subtract extended state size reserved from AMX from the rest of the extended state size, and hardcode the AMX XCR0 bits to be always off. Fixes #9750. PiperOrigin-RevId: 599302059 --- pkg/cpuid/cpuid_amd64.go | 39 +++++++++++++++++-- pkg/cpuid/native_amd64.go | 3 ++ pkg/cpuid/native_amd64.s | 13 +++++++ pkg/ring0/entry_amd64.s | 23 +++++++---- pkg/sentry/arch/fpu/fpu_amd64.s | 11 +++++- .../systrap/sysmsg/sighandler_amd64.c | 2 +- .../systrap/sysmsg/syshandler_amd64.S | 8 ++-- .../platform/systrap/sysmsg/sysmsg_offsets.h | 7 ++++ 8 files changed, 87 insertions(+), 19 deletions(-) diff --git a/pkg/cpuid/cpuid_amd64.go b/pkg/cpuid/cpuid_amd64.go index 7a7277d22..16c55dff4 100644 --- a/pkg/cpuid/cpuid_amd64.go +++ b/pkg/cpuid/cpuid_amd64.go @@ -361,8 +361,22 @@ func (fs FeatureSet) Intel() bool { // If xSaveInfo isn't supported, cpuid will not fault but will // return bogus values. var ( - xsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ebx - maxXsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ecx + xsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ebx + maxXsaveSize = native(In{Eax: uint32(xSaveInfo)}).Ecx + amxTileCfgSize = native(In{Eax: uint32(xSaveInfo), Ecx: 17}).Eax + amxTileDataSize = native(In{Eax: uint32(xSaveInfo), Ecx: 18}).Eax +) + +const ( + // XCR0AMXMask are the bits that enable xsave to operate on AMX TILECFG + // and TILEDATA. + // + // Note: TILECFG and TILEDATA are always either both enabled or both + // disabled. + // + // See Intel® 64 and IA-32 Architectures Software Developer’s Manual Vol.1 + // section 13.3 for details. + XCR0AMXMask = uint64((1 << 17) | (1 << 18)) ) // ExtendedStateSize returns the number of bytes needed to save the "extended @@ -370,13 +384,17 @@ var ( // 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), 64 + return uint(xsaveSize) - fs.AMXExtendedStateSize(), 64 } // If we don't support xsave, we fall back to fxsave, which requires @@ -384,15 +402,28 @@ func (fs FeatureSet) ExtendedStateSize() (size, align uint) { return 512, 16 } +// AMXExtendedStateSize returns the number of bytes within the "extended state" +// area that is used for AMX. +func (fs FeatureSet) AMXExtendedStateSize() uint { + xcr0 := xgetbv(0) + if (xcr0 & XCR0AMXMask) != 0 { + return uint(amxTileCfgSize + amxTileDataSize) + } + return 0 +} + // ValidXCR0Mask returns the valid bits in control register XCR0. // +// Always exclude AMX bits, because we do not support it. +// TODO(gvisor.dev/issues/9896): Implement AMX Support. +// //go:nosplit func (fs FeatureSet) ValidXCR0Mask() uint64 { if !fs.HasFeature(X86FeatureXSAVE) { return 0 } ax, _, _, dx := fs.query(xSaveInfo) - return uint64(dx)<<32 | uint64(ax) + return (uint64(dx)<<32 | uint64(ax)) ^ XCR0AMXMask } // UseXsave returns the choice of fp state saving instruction. diff --git a/pkg/cpuid/native_amd64.go b/pkg/cpuid/native_amd64.go index eaf77511d..ac2fcbbcc 100644 --- a/pkg/cpuid/native_amd64.go +++ b/pkg/cpuid/native_amd64.go @@ -215,6 +215,9 @@ func readMaxCPUFreq() { } +// xgetbv reads an extended control register. +func xgetbv(reg uintptr) uint64 + // archInitialize initializes hostFeatureSet. func archInitialize() { hostFeatureSet = FeatureSet{ diff --git a/pkg/cpuid/native_amd64.s b/pkg/cpuid/native_amd64.s index dd21b4bdd..04a1433a9 100644 --- a/pkg/cpuid/native_amd64.s +++ b/pkg/cpuid/native_amd64.s @@ -23,3 +23,16 @@ TEXT ·native(SB),NOSPLIT|NOFRAME,$0-24 MOVL CX, ret_Ecx+16(FP) MOVL DX, ret_Edx+20(FP) RET + +// xgetbv reads an extended control register. +// +// The code corresponds to: +// +// xgetbv +// +TEXT ·xgetbv(SB),NOSPLIT|NOFRAME,$0-16 + MOVQ reg+0(FP), CX + BYTE $0x0f; BYTE $0x01; BYTE $0xd0; + MOVL AX, ret+8(FP) + MOVL DX, ret+12(FP) + RET diff --git a/pkg/ring0/entry_amd64.s b/pkg/ring0/entry_amd64.s index d2f870a52..ad49618b1 100644 --- a/pkg/ring0/entry_amd64.s +++ b/pkg/ring0/entry_amd64.s @@ -88,6 +88,13 @@ #define PTRACE_FS_BASE 168 // +checkoffset linux PtraceRegs.Fs_base #define PTRACE_GS_BASE 176 // +checkoffset linux PtraceRegs.Gs_base +// The value for XCR0 is defined to xsave/xrstor everything except for AMX +// regions. +// TODO(gvisor.dev/issues/9896): Implement AMX Support. +#define XCR0_AMX_MASK ((1 << 17) | (1 << 18)) +#define XCR0_EAX (0xffffffff ^ XCR0_AMX_MASK) +#define XCR0_EDX 0xffffffff + // Saves a register set. // // This is a macro because it may need to executed in contents where a stack is @@ -229,8 +236,8 @@ TEXT ·doSwitchToUser(SB),NOSPLIT,$16-48 // 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. - MOVL $0xffffffff, AX - MOVL $0xffffffff, DX + MOVL $XCR0_EAX, AX + MOVL $XCR0_EDX, DX BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x2f // XRSTOR64 0(DI) JMP fprestore_done no_xrstor: @@ -275,8 +282,8 @@ done_sysret_or_iret: JZ no_xsave // Use xsave/xsaveopt to save all extended state. // We save everything unconditionally by setting RFBM to all 1's. - MOVL $0xffffffff, AX - MOVL $0xffffffff, DX + MOVL $XCR0_EAX, AX + MOVL $XCR0_EDX, DX TESTB CX, CX JZ no_xsaveopt BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x37; // XSAVEOPT64 0(DI) @@ -513,8 +520,8 @@ kernel: JZ no_xsave // Use xsave/xsaveopt to save all extended state. // We save everything unconditionally by setting RFBM to all 1's. - MOVL $0xffffffff, AX - MOVL $0xffffffff, DX + MOVL $XCR0_EAX, AX + MOVL $XCR0_EDX, DX TESTB CX, CX JZ no_xsaveopt BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x37; // XSAVEOPT64 0(DI) @@ -645,8 +652,8 @@ kernel: JZ no_xsave // Use xsave/xsaveopt to save all extended state. // We save everything unconditionally by setting RFBM to all 1's. - MOVL $0xffffffff, AX - MOVL $0xffffffff, DX + MOVL $XCR0_EAX, AX + MOVL $XCR0_EDX, DX TESTB CX, CX JZ no_xsaveopt BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x37; // XSAVEOPT64 0(DI) diff --git a/pkg/sentry/arch/fpu/fpu_amd64.s b/pkg/sentry/arch/fpu/fpu_amd64.s index 7ce63092c..61d3676ca 100644 --- a/pkg/sentry/arch/fpu/fpu_amd64.s +++ b/pkg/sentry/arch/fpu/fpu_amd64.s @@ -22,6 +22,13 @@ // FXSAVE/XSAVE area. (Intel SDM Vol. 1, Table 10-2 "Format of an FXSAVE Area") #define MXCSR_OFFSET 24 +// The value for XCR0 is defined to xsave/xrstor everything except for AMX +// regions. +// TODO(gvisor.dev/issues/9896): Implement AMX Support. +#define XCR0_AMX_MASK ((1 << 17) | (1 << 18)) +#define XCR0_EAX (0xffffffff ^ XCR0_AMX_MASK) +#define XCR0_EDX 0xffffffff + // initX86FPState initializes floating point state. // // func initX86FPState(data *FloatingPointData, useXsave bool) @@ -66,8 +73,8 @@ TEXT ·initX86FPState(SB), $24-9 MOVL $MXCSR_DEFAULT, MXCSR_OFFSET(DI) // Initialize registers with XRSTOR. - MOVL $0xffffffff, AX - MOVL $0xffffffff, DX + MOVL $XCR0_EAX, AX + MOVL $XCR0_EDX, DX BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x2f // XRSTOR64 0(DI) // Now that all the state has been reset, write it back out to the diff --git a/pkg/sentry/platform/systrap/sysmsg/sighandler_amd64.c b/pkg/sentry/platform/systrap/sysmsg/sighandler_amd64.c index 0637f8ea5..184c9cac5 100644 --- a/pkg/sentry/platform/systrap/sysmsg/sighandler_amd64.c +++ b/pkg/sentry/platform/systrap/sysmsg/sighandler_amd64.c @@ -480,7 +480,7 @@ static void prep_fpstate_for_sigframe(void *buf, uint32_t user_size, sw_bytes->magic1 = FP_XSTATE_MAGIC1; sw_bytes->extended_size = user_size + FP_XSTATE_MAGIC2_SIZE; - sw_bytes->xfeatures = ~(0ULL); + sw_bytes->xfeatures = ~(0ULL) ^ (XCR0_AMX_MASK); sw_bytes->xstate_size = user_size; *(uint32_t *)(buf + user_size) = use_xsave ? FP_XSTATE_MAGIC2 : 0; } diff --git a/pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S b/pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S index 86855bd91..bc0d60663 100644 --- a/pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S +++ b/pkg/sentry/platform/systrap/sysmsg/syshandler_amd64.S @@ -122,8 +122,8 @@ restored_gs: // - the memory address of the thread_context is loaded in %rcx. .macro save_fpstate lea offsetof_thread_context_fpstate(%rcx), %rdi - movl $0xffffffff, %eax - movl $0xffffffff, %edx + movl $XCR0_EAX, %eax + movl $XCR0_EDX, %edx movl __export_arch_state+offsetof_arch_state_xsave_mode(%rip), %esi cmpl $XSAVE_MODE_XSAVEOPT, %esi jl use_xsave @@ -153,8 +153,8 @@ fpu_saved: cmpl $XSAVE_MODE_FXSAVE, %eax jz use_fxrstor use_xrstor: - movl $0xffffffff, %eax - movl $0xffffffff, %edx + movl $XCR0_EAX, %eax + movl $XCR0_EDX, %edx xrstor (%rdi) jmp fpu_restored use_fxrstor: diff --git a/pkg/sentry/platform/systrap/sysmsg/sysmsg_offsets.h b/pkg/sentry/platform/systrap/sysmsg/sysmsg_offsets.h index 4689f0dde..51fa2ed31 100644 --- a/pkg/sentry/platform/systrap/sysmsg/sysmsg_offsets.h +++ b/pkg/sentry/platform/systrap/sysmsg/sysmsg_offsets.h @@ -20,6 +20,13 @@ // for the pkg/sentry/platform/systrap/usertrap package. #define FAULT_OPCODE 0x06 +// The value for XCR0 is defined to xsave/xrstor everything except for AMX +// regions. +// TODO(gvisor.dev/issues/9896): Implement AMX Support. +#define XCR0_AMX_MASK ((1 << 17) | (1 << 18)) +#define XCR0_EAX (0xffffffff ^ XCR0_AMX_MASK) +#define XCR0_EDX 0xffffffff + // LINT.IfChange #define MAX_FPSTATE_LEN 3584 // Note: To be explicit, 2^12 = 4096; if ALLOCATED_SIZEOF_THREAD_CONTEXT_STRUCT