From af2ce4d08fc60e95626606f1d0d1525c8008db35 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Wed, 6 Dec 2023 21:09:13 +0200
Subject: [PATCH] [NTOS:KE/x64] Implement KiSaveProcessorState /
 KiRestoreProcessorState

---
 ntoskrnl/include/internal/ke.h | 11 ++++++++---
 ntoskrnl/ke/amd64/cpu.c        | 35 ++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/ntoskrnl/include/internal/ke.h b/ntoskrnl/include/internal/ke.h
index ba3a5079a95..649bd1045ab 100644
--- a/ntoskrnl/include/internal/ke.h
+++ b/ntoskrnl/include/internal/ke.h
@@ -1024,9 +1024,14 @@ KiSaveProcessorControlState(
 VOID
 NTAPI
 KiSaveProcessorState(
-    IN PKTRAP_FRAME TrapFrame,
-    IN PKEXCEPTION_FRAME ExceptionFrame
-);
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame);
+
+VOID
+NTAPI
+KiRestoreProcessorState(
+    _Out_ PKTRAP_FRAME TrapFrame,
+    _Out_ PKEXCEPTION_FRAME ExceptionFrame);
 
 VOID
 FASTCALL
diff --git a/ntoskrnl/ke/amd64/cpu.c b/ntoskrnl/ke/amd64/cpu.c
index b2b40b0bf62..d1e9567c1b8 100644
--- a/ntoskrnl/ke/amd64/cpu.c
+++ b/ntoskrnl/ke/amd64/cpu.c
@@ -612,6 +612,41 @@ KiSaveProcessorControlState(OUT PKPROCESSOR_STATE ProcessorState)
     ProcessorState->SpecialRegisters.MsrSyscallMask = __readmsr(X86_MSR_SFMASK);
 }
 
+VOID
+NTAPI
+KiSaveProcessorState(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame)
+{
+    PKPRCB Prcb = KeGetCurrentPrcb();
+
+    /* Save all context */
+    Prcb->ProcessorState.ContextFrame.ContextFlags = CONTEXT_ALL;
+    KeTrapFrameToContext(TrapFrame, ExceptionFrame, &Prcb->ProcessorState.ContextFrame);
+
+    /* Save control registers */
+    KiSaveProcessorControlState(&Prcb->ProcessorState);
+}
+
+VOID
+NTAPI
+KiRestoreProcessorState(
+    _Out_ PKTRAP_FRAME TrapFrame,
+    _Out_ PKEXCEPTION_FRAME ExceptionFrame)
+{
+    PKPRCB Prcb = KeGetCurrentPrcb();
+
+    /* Restore all context */
+    KeContextToTrapFrame(&Prcb->ProcessorState.ContextFrame,
+                         ExceptionFrame,
+                         TrapFrame,
+                         CONTEXT_ALL,
+                         TrapFrame->PreviousMode);
+
+    /* Restore control registers */
+    KiRestoreProcessorControlState(&Prcb->ProcessorState);
+}
+
 VOID
 NTAPI
 KeFlushEntireTb(IN BOOLEAN Invalid,