From 9b002025d45e36cccaa579e930c7aa8267c77c4d Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Wed, 22 Jan 2025 12:58:47 +0200
Subject: [PATCH] [NTOS:KE/x86] Detect more KeFeatureBits

---
 ntoskrnl/include/internal/i386/ke.h |  8 ++++++++
 ntoskrnl/ke/i386/cpu.c              |  8 ++++++++
 ntoskrnl/ke/i386/kiinit.c           | 11 +++++++++++
 3 files changed, 27 insertions(+)

diff --git a/ntoskrnl/include/internal/i386/ke.h b/ntoskrnl/include/internal/i386/ke.h
index 596c58b418f..3d4675b30d8 100644
--- a/ntoskrnl/include/internal/i386/ke.h
+++ b/ntoskrnl/include/internal/i386/ke.h
@@ -22,6 +22,14 @@ extern "C"
 #define KD_BREAKPOINT_SIZE        sizeof(UCHAR)
 #define KD_BREAKPOINT_VALUE       0xCC
 
+/* CPUID 1 - ECX flags */
+#define X86_FEATURE_SSE3        0x00000001
+#define X86_FEATURE_SSSE3       0x00000200
+#define X86_FEATURE_SSE4_1      0x00080000
+#define X86_FEATURE_SSE4_2      0x00100000
+#define X86_FEATURE_XSAVE       0x04000000
+#define X86_FEATURE_RDRAND      0x40000000
+
 /* CPUID 1 - EDX flags */
 #define X86_FEATURE_FPU         0x00000001 /* x87 FPU is present */
 #define X86_FEATURE_VME         0x00000002 /* Virtual 8086 Extensions are present */
diff --git a/ntoskrnl/ke/i386/cpu.c b/ntoskrnl/ke/i386/cpu.c
index ab639340443..a5d4912db43 100644
--- a/ntoskrnl/ke/i386/cpu.c
+++ b/ntoskrnl/ke/i386/cpu.c
@@ -361,6 +361,14 @@ KiGetFeatureBits(VOID)
             break;
     }
 
+    /* Get some features from ECX */
+    if (CpuInfo.Ecx & X86_FEATURE_SSE3) FeatureBits |= KF_SSE3;
+    if (CpuInfo.Ecx & X86_FEATURE_SSSE3) FeatureBits |= KF_SSSE3;
+    if (CpuInfo.Ecx & X86_FEATURE_SSE4_1) FeatureBits |= KF_SSE4_1;
+    if (CpuInfo.Ecx & X86_FEATURE_SSE4_2) FeatureBits |= KF_SSE4_2;
+    if (CpuInfo.Ecx & X86_FEATURE_XSAVE) FeatureBits |= KF_XSTATE;
+    if (CpuInfo.Ecx & X86_FEATURE_RDRAND) FeatureBits |= KF_RDRAND;
+
     /* Set the current features */
     CpuFeatures = CpuInfo.Edx;
 
diff --git a/ntoskrnl/ke/i386/kiinit.c b/ntoskrnl/ke/i386/kiinit.c
index 4d576e86f8a..671adad67c3 100644
--- a/ntoskrnl/ke/i386/kiinit.c
+++ b/ntoskrnl/ke/i386/kiinit.c
@@ -572,6 +572,17 @@ KiInitializeKernel(IN PKPROCESS InitProcess,
         (KeFeatureBits & KF_3DNOW) ? TRUE: FALSE;
     SharedUserData->ProcessorFeatures[PF_RDTSC_INSTRUCTION_AVAILABLE] =
         (KeFeatureBits & KF_RDTSC) ? TRUE: FALSE;
+    SharedUserData->ProcessorFeatures[PF_RDRAND_INSTRUCTION_AVAILABLE] =
+        (KeFeatureBits & KF_RDRAND) ? TRUE : FALSE;
+    // Note: On x86 we lack support for saving/restoring SSE state
+    SharedUserData->ProcessorFeatures[PF_SSE3_INSTRUCTIONS_AVAILABLE] =
+        (KeFeatureBits & KF_SSE3) ? TRUE : FALSE;
+    SharedUserData->ProcessorFeatures[PF_SSSE3_INSTRUCTIONS_AVAILABLE] =
+        (KeFeatureBits & KF_SSSE3) ? TRUE : FALSE;
+    SharedUserData->ProcessorFeatures[PF_SSE4_1_INSTRUCTIONS_AVAILABLE] =
+        (KeFeatureBits & KF_SSE4_1) ? TRUE : FALSE;
+    SharedUserData->ProcessorFeatures[PF_SSE4_2_INSTRUCTIONS_AVAILABLE] =
+        (KeFeatureBits & KF_SSE4_2) ? TRUE : FALSE;
 
     /* Set up the thread-related fields in the PRCB */
     Prcb->CurrentThread = InitThread;