From d6874fe7b9f7f265b110381247de77f984edc8cb Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Wed, 25 Dec 2024 14:44:21 +0200
Subject: [PATCH] [NTOS:KE/x64] Implement KiSave/RestoreXState

---
 ntoskrnl/CMakeLists.txt                    |  5 +++
 ntoskrnl/include/internal/amd64/intrin_i.h | 48 ++++++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/ntoskrnl/CMakeLists.txt b/ntoskrnl/CMakeLists.txt
index 330296a50b9..ffa2c64633b 100644
--- a/ntoskrnl/CMakeLists.txt
+++ b/ntoskrnl/CMakeLists.txt
@@ -28,6 +28,11 @@ if(CMAKE_C_COMPILER_ID STREQUAL "Clang")
     target_sources(libntoskrnl PRIVATE $<TARGET_OBJECTS:memcmp>)
 endif()
 
+if((CMAKE_C_COMPILER_ID STREQUAL "Clang") OR (CMAKE_C_COMPILER_ID STREQUAL "GNU"))
+    # Enable xsave feature
+    add_compile_options(-mxsave)
+endif()
+
 list(APPEND PCH_SKIP_SOURCE
     guid.c)
 
diff --git a/ntoskrnl/include/internal/amd64/intrin_i.h b/ntoskrnl/include/internal/amd64/intrin_i.h
index 07b73fac030..1e44a9695ee 100644
--- a/ntoskrnl/include/internal/amd64/intrin_i.h
+++ b/ntoskrnl/include/internal/amd64/intrin_i.h
@@ -59,6 +59,54 @@ KiInitGdtEntry(PKGDTENTRY64 Entry, ULONG64 Base, ULONG Size, UCHAR Type, UCHAR D
     Entry->MustBeZero = 0;
 }
 
+extern ULONG64 KeFeatureBits;
+
+FORCEINLINE
+VOID
+KiSaveXState(
+    _Out_ PVOID Buffer,
+    _In_ ULONG64 ComponentMask)
+{
+    ComponentMask &= ~XSTATE_MASK_LEGACY_SSE;
+    if (KeFeatureBits & KF_XSAVES)
+    {
+        _xsaves64(Buffer, ComponentMask);
+    }
+    else if (KeFeatureBits & KF_XSAVEOPT)
+    {
+        _xsaveopt64(Buffer, ComponentMask);
+    }
+    else if (KeFeatureBits & KF_XSTATE)
+    {
+        _xsave64(Buffer, ComponentMask);
+    }
+    else
+    {
+        _fxsave64(Buffer);
+    }
+}
+
+FORCEINLINE
+VOID
+KiRestoreXState(
+    _In_ PVOID Buffer,
+    _In_ ULONG64 ComponentMask)
+{
+    ComponentMask &= ~XSTATE_MASK_LEGACY_SSE;
+    if (KeFeatureBits & KF_XSAVES)
+    {
+        _xrstors64(Buffer, ComponentMask);
+    }
+    else if (KeFeatureBits & KF_XSTATE)
+    {
+        _xrstor64(Buffer, ComponentMask);
+    }
+    else
+    {
+        _fxrstor64(Buffer);
+    }
+}
+
 #if defined(__GNUC__)
 
 static __inline__ __attribute__((always_inline)) void __lgdt(void *Source)