From 294eb31cfd00571f09449859a4177b0271f91635 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Tue, 14 Nov 2023 16:51:04 +0200
Subject: [PATCH] [NTOS:KE/x64] Handle extended processor state on context
 switch

---
 ntoskrnl/ke/amd64/thrdini.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/ntoskrnl/ke/amd64/thrdini.c b/ntoskrnl/ke/amd64/thrdini.c
index a1f9d614508..e4cbdff084b 100644
--- a/ntoskrnl/ke/amd64/thrdini.c
+++ b/ntoskrnl/ke/amd64/thrdini.c
@@ -167,9 +167,21 @@ KiSwapContextResume(
     PKPROCESS OldProcess, NewProcess;
 
     /* Setup ring 0 stack pointer */
-    Pcr->TssBase->Rsp0 = (ULONG64)NewThread->InitialStack; // FIXME: NPX save area?
+    Pcr->TssBase->Rsp0 = (ULONG64)NewThread->InitialStack;
     Pcr->Prcb.RspBase = Pcr->TssBase->Rsp0;
 
+    /* Save old thread's extended state */
+    if (OldThread->NpxState != 0)
+    {
+        KiSaveXState(OldThread->StateSaveArea, OldThread->NpxState);
+    }
+
+    /* Load new thread's extended state */
+    if (NewThread->NpxState != 0)
+    {
+        KiRestoreXState(NewThread->StateSaveArea, NewThread->NpxState);
+    }
+
     /* Now we are the new thread. Check if it's in a new process */
     OldProcess = OldThread->ApcState.Process;
     NewProcess = NewThread->ApcState.Process;