From 61b6cd9d285abc83fe388c93d5a734bc1529ef26 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 6 Jan 2020 01:52:52 +0200
Subject: [PATCH] [NTOS:KE] Implement KiGetTrapContextInternal

---
 ntoskrnl/include/internal/amd64/ke.h |  5 +++++
 ntoskrnl/ke/amd64/context.c          | 22 ++++++++++++++++++++++
 ntoskrnl/ke/amd64/trap.S             | 23 +++++++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/ntoskrnl/include/internal/amd64/ke.h b/ntoskrnl/include/internal/amd64/ke.h
index f3112871d62..f9af4818b54 100644
--- a/ntoskrnl/include/internal/amd64/ke.h
+++ b/ntoskrnl/include/internal/amd64/ke.h
@@ -473,6 +473,11 @@ KiGetUserModeStackAddress(void)
     return &PsGetCurrentThread()->Tcb.TrapFrame->Rsp;
 }
 
+VOID
+KiGetTrapContext(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _Out_ PCONTEXT Context);
+
 VOID
 KiSetTrapContext(
     _Out_ PKTRAP_FRAME TrapFrame,
diff --git a/ntoskrnl/ke/amd64/context.c b/ntoskrnl/ke/amd64/context.c
index 3c345f89052..06ceb662ba4 100644
--- a/ntoskrnl/ke/amd64/context.c
+++ b/ntoskrnl/ke/amd64/context.c
@@ -292,6 +292,28 @@ KeTrapFrameToContext(IN PKTRAP_FRAME TrapFrame,
     if (OldIrql < APC_LEVEL) KeLowerIrql(OldIrql);
 }
 
+VOID
+RtlGetUnwindContext(
+    _Out_ PCONTEXT Context,
+    _In_ DWORD64 TargetFrame);
+
+VOID
+KiGetTrapContextInternal(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _Out_ PCONTEXT Context)
+{
+    ULONG64 TargetFrame;
+
+    /* Get the volatile register context from the trap frame */
+    KeTrapFrameToContext(TrapFrame, NULL, Context);
+
+    /* The target frame is MAX_SYSCALL_PARAM_SIZE bytes before the trap frame */
+    TargetFrame = (ULONG64)TrapFrame - MAX_SYSCALL_PARAM_SIZE;
+
+    /* Get the nonvolatiles on the stack */
+    RtlGetUnwindContext(Context, TargetFrame);
+}
+
 VOID
 RtlSetUnwindContext(
     _In_ PCONTEXT Context,
diff --git a/ntoskrnl/ke/amd64/trap.S b/ntoskrnl/ke/amd64/trap.S
index 6f1c2a23521..7b07b333e73 100644
--- a/ntoskrnl/ke/amd64/trap.S
+++ b/ntoskrnl/ke/amd64/trap.S
@@ -1137,6 +1137,29 @@ KiConvertToGuiThreadFailed:
 
 ENDFUNC
 
+EXTERN KiGetTrapContextInternal:PROC
+
+/*
+ * VOID
+ * KiGetTrapContext(
+ *     _Out_ PKTRAP_FRAME TrapFrame,
+ *     _In_ PCONTEXT Context);
+ */
+PUBLIC KiGetTrapContext
+.PROC KiGetTrapContext
+
+    /* Generate a KEXCEPTION_FRAME on the stack */
+    GENERATE_EXCEPTION_FRAME
+
+    call KiGetTrapContextInternal
+
+    /* Restore the registers from the KEXCEPTION_FRAME */
+    RESTORE_EXCEPTION_STATE
+
+    /* Return */
+    ret
+
+.ENDP
 
 EXTERN KiSetTrapContextInternal:PROC