From 29721ac552cd104116bbc197b8df1782c72f7710 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Wed, 11 Dec 2024 16:47:28 +0200
Subject: [PATCH] [NTOS:KE/x64] Fix unwinding in KiUserApcDispatcher

This adds a MACHINE_FRAME to the user mode APC dispatcher stack frame.
---
 dll/ntdll/dispatch/amd64/dispatch.S | 26 ++++++++++++++++++++++++++
 ntoskrnl/ke/amd64/usercall.c        | 18 ++++++++++++------
 sdk/include/ndk/amd64/ketypes.h     |  9 +++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/dll/ntdll/dispatch/amd64/dispatch.S b/dll/ntdll/dispatch/amd64/dispatch.S
index ea08389c919..60b3f22bb58 100644
--- a/dll/ntdll/dispatch/amd64/dispatch.S
+++ b/dll/ntdll/dispatch/amd64/dispatch.S
@@ -33,6 +33,32 @@ PUBLIC LdrInitializeThunk
 
 PUBLIC KiUserApcDispatcher
 .PROC KiUserApcDispatcher
+
+    /* The stack is set up with a UAPC_FRAME, which ends with a MACHINE_FRAME */
+    .PUSHFRAME
+    .ALLOCSTACK CONTEXT_FRAME_LENGTH
+
+    /* The stack points to a CONTEXT structure.
+       Create unwind ops for all nonvolatile registers */
+    .SAVEREG rbx, CxRbx
+    .SAVEREG rbp, CxRbp
+    .SAVEREG rsi, CxRsi
+    .SAVEREG rdi, CxRdi
+    .SAVEREG r12, CxR12
+    .SAVEREG r13, CxR13
+    .SAVEREG r14, CxR14
+    .SAVEREG r15, CxR15
+    .SAVEXMM128 xmm6, CxXmm6
+    .SAVEXMM128 xmm7, CxXmm7
+    .SAVEXMM128 xmm8, CxXmm8
+    .SAVEXMM128 xmm9, CxXmm9
+    .SAVEXMM128 xmm10, CxXmm10
+    .SAVEXMM128 xmm11, CxXmm11
+    .SAVEXMM128 xmm12, CxXmm12
+    .SAVEXMM128 xmm13, CxXmm13
+    .SAVEXMM128 xmm14, CxXmm14
+    .SAVEXMM128 xmm15, CxXmm15
+
     .endprolog
     /* We enter with a 16 byte aligned stack */
 
diff --git a/ntoskrnl/ke/amd64/usercall.c b/ntoskrnl/ke/amd64/usercall.c
index 81f7a81fdfc..bab35e9e40e 100644
--- a/ntoskrnl/ke/amd64/usercall.c
+++ b/ntoskrnl/ke/amd64/usercall.c
@@ -16,8 +16,8 @@
  *
  *  \brief
  *      Prepares the current trap frame (which must have come from user mode)
- *      with the ntdll.KiUserApcDispatcher entrypoint, copying a CONTEXT
- *      record with the context from the old trap frame to the threads user
+ *      with the ntdll.KiUserApcDispatcher entrypoint, copying a UAPC_FRAME
+ *      structure with the context from the old trap frame to the threads user
  *      mode stack.
  *
  *  \param ExceptionFrame - Pointer to the Exception Frame
@@ -53,20 +53,22 @@ KiInitializeUserApc(
     _In_ PVOID SystemArgument1,
     _In_ PVOID SystemArgument2)
 {
+    PUAPC_FRAME ApcFrame;
     PCONTEXT Context;
     EXCEPTION_RECORD ExceptionRecord;
 
     /* Sanity check, that the trap frame is from user mode */
     ASSERT((TrapFrame->SegCs & MODE_MASK) != KernelMode);
 
-    /* Align the user tack to 16 bytes and allocate space for a CONTEXT structure */
-    Context = (PCONTEXT)ALIGN_DOWN_POINTER_BY(TrapFrame->Rsp, 16) - 1;
+    /* Allocate a 16 byte aligned UAPC_FRAME structure on the user stack */
+    ApcFrame = (PUAPC_FRAME)ALIGN_DOWN_POINTER_BY(TrapFrame->Rsp - sizeof(*ApcFrame), 16);
+    Context = &ApcFrame->Context;
 
     /* Protect with SEH */
     _SEH2_TRY
     {
-        /* Probe the context */
-        ProbeForWrite(Context, sizeof(CONTEXT), 16);
+        /* Probe the user mode APC frame */
+        ProbeForWrite(ApcFrame, sizeof(*ApcFrame), 16);
 
         /* Convert the current trap frame to a context */
         Context->ContextFlags = CONTEXT_FULL | CONTEXT_DEBUG_REGISTERS;
@@ -77,6 +79,10 @@ KiInitializeUserApc(
         Context->P2Home = (ULONG64)SystemArgument1;
         Context->P3Home = (ULONG64)SystemArgument2;
         Context->P4Home = (ULONG64)NormalRoutine;
+
+        /* Set up the machine frame for unwinding */
+        ApcFrame->MachineFrame.Rip = TrapFrame->Rip;
+        ApcFrame->MachineFrame.Rsp = TrapFrame->Rsp;
     }
     _SEH2_EXCEPT(ExceptionRecord = *_SEH2_GetExceptionInformation()->ExceptionRecord, EXCEPTION_EXECUTE_HANDLER)
     {
diff --git a/sdk/include/ndk/amd64/ketypes.h b/sdk/include/ndk/amd64/ketypes.h
index 410e0b55cb0..ea9031df667 100644
--- a/sdk/include/ndk/amd64/ketypes.h
+++ b/sdk/include/ndk/amd64/ketypes.h
@@ -1073,6 +1073,15 @@ typedef struct _UCALLOUT_FRAME
     MACHINE_FRAME MachineFrame;
 } UCALLOUT_FRAME, *PUCALLOUT_FRAME; // size = 0x0058
 
+//
+// User side APC dispatcher frame
+//
+typedef struct _UAPC_FRAME
+{
+    CONTEXT Context;
+    MACHINE_FRAME MachineFrame;
+} UAPC_FRAME, *PUAPC_FRAME;
+
 //
 // Stack frame layout for KiUserExceptionDispatcher
 // The name is totally made up