[NTOS:KE/x64] Implement KiSetTrapContext

KiSetTrapContext is an asm wrapper around RtlSetUnwindContext, which first stores an exception frame to assure that all non-volatile registers were put on the stack, then calls RtlSetUnwindContext to update their first saving positions on the stack and finally restore the exception frame to potentially load any updated registers, that haven't been saved elsewhere on the stack.

ntoskrnl/ke/amd64/context.c

@@ -292,3 +292,29 @@ KeTrapFrameToContext(IN PKTRAP_FRAME TrapFrame,
     if (OldIrql < APC_LEVEL) KeLowerIrql(OldIrql);
 }
 
+VOID
+RtlSetUnwindContext(
+    _In_ PCONTEXT Context,
+    _In_ DWORD64 TargetFrame);
+
+VOID
+KiSetTrapContextInternal(
+    _Out_ PKTRAP_FRAME TrapFrame,
+    _In_ PCONTEXT Context,
+    _In_ KPROCESSOR_MODE RequestorMode)
+{
+    ULONG64 TargetFrame;
+
+    /* Save the volatile register context in the trap frame */
+    KeContextToTrapFrame(Context,
+                         NULL,
+                         TrapFrame,
+                         Context->ContextFlags,
+                         RequestorMode);
+
+    /* The target frame is MAX_SYSCALL_PARAM_SIZE bytes before the trap frame */
+    TargetFrame = (ULONG64)TrapFrame - MAX_SYSCALL_PARAM_SIZE ;
+
+    /* Set the nonvolatiles on the stack */
+    RtlSetUnwindContext(Context, TargetFrame);
+}

ntoskrnl/ke/amd64/trap.S

@@ -730,9 +730,6 @@ FUNC KiInterruptDispatch
     ExitTrap (TF_SAVE_ALL or TF_SEND_EOI)
 ENDFUNC
 
-
-#define MAX_SYSCALL_PARAM_SIZE (16 * 8)
-
 EXTERN KiSystemCallHandler:PROC
 
 /*! \name KiSystemCallEntry64
@@ -1049,6 +1046,33 @@ KiConvertToGuiThreadFailed:
 
 ENDFUNC
 
+
+EXTERN KiSetTrapContextInternal:PROC
+
+/*
+ * VOID
+ * KiSetTrapContext(
+ *     _Out_ PKTRAP_FRAME TrapFrame,
+ *     _In_ PCONTEXT Context,
+ *     _In_ KPROCESSOR_MODE RequestorMode);
+ */
+PUBLIC KiSetTrapContext
+.PROC KiSetTrapContext
+    
+    /* Generate a KEXCEPTION_FRAME on the stack */
+    GENERATE_EXCEPTION_FRAME
+
+    call KiSetTrapContextInternal
+
+    /* Restore the registers from the KEXCEPTION_FRAME */
+    RESTORE_EXCEPTION_STATE
+
+    /* Return */
+    ret
+
+.ENDP
+
+
 /*
  * VOID
  * KiDeliverApc(

sdk/include/asm/ksamd64.template.h

@@ -1071,3 +1071,5 @@ OFFSET(KINTERRUPT_DispatchCount, KINTERRUPT, DispatchCount),
 OFFSET(KINTERRUPT_TrapFrame, KINTERRUPT, TrapFrame),
 OFFSET(KINTERRUPT_DispatchCode, KINTERRUPT, DispatchCode),
 
+HEADER("Misc definitions"),
+CONSTANT(MAX_SYSCALL_PARAM_SIZE),

sdk/include/ndk/amd64/asm.h

@@ -254,6 +254,9 @@ Author:
 #define CR0_CD                                  HEX(40000000)
 #define CR0_PG                                  HEX(80000000)
 
+/* Number of bytes reserved for syscall parameters */
+#define MAX_SYSCALL_PARAM_SIZE (16 * 8)
+
 #ifdef _ASM_
 //
 // CR4