[NTOS:KE:X64][ASM:X64] Fix delivery of APCs

- Deliver pending APCs on trap exit
- Pass the trapframe of KiApcInterrupt to KiDeliverApcs, not NULL.
- Fix parameter passing from KiSwapContext to KiSwapContextInternal and KiSwapContextResume, so that the ApcBypass parameter is not uninitialized
- Fix return value of KiSwapContextResume to correctly indicate whether we want to have APCs directly delivered or not (when there are non, or when delivery is suppressed)

ntoskrnl/ke/amd64/ctxswitch.S

@@ -11,6 +11,11 @@
 
 #include <ksamd64.inc>
 
+; BOOLEAN
+; KiSwapContextResume(
+;    _In_ KIRQL WaitIrql,
+;    _In_ PKTHREAD OldThread,
+;    _In_ PKTHREAD NewThread)
 EXTERN KiSwapContextResume:PROC
 
 /* FUNCTIONS ****************************************************************/
@@ -118,23 +123,22 @@ PUBLIC KiThreadStartup
  * \brief
  *     The KiSwapContextInternal routine switches context to another thread.
  *
- * \param rcx
+ * \param cl
- *     Pointer to the KTHREAD to which the caller wishes to switch to.
+ *     The IRQL at wich the old thread is suspended
  *
  * \param rdx
  *     Pointer to the KTHREAD to which the caller wishes to switch from.
  *
- * \param r8b
+ * \param r8
- *     APC bypass
+ *     Pointer to the KTHREAD to which the caller wishes to switch to.
  *
  * \return
- *     None.
+ *     The WaitStatus of the Target Thread.
  *
  * \remarks
  *     ...
  *
  *--*/
-PUBLIC KiSwapContextInternal
 .PROC KiSwapContextInternal
 
     push rbp
@@ -143,22 +147,17 @@ PUBLIC KiSwapContextInternal
     .allocstack 6 * 8
     .endprolog
 
-    /* Save APC bypass */
+    /* Save WaitIrql as KSWITCH_FRAME::ApcBypass */
-    mov [rsp + SwApcBypass], r8b
+    mov [rsp + SwApcBypass], cl
 
     /* Save kernel stack of old thread */
     mov [rdx + KTHREAD_KernelStack], rsp
 
-    /* Save new thread in rbp */
-    mov rbp, rcx
-
-    //call KiSwapContextSuspend
-
     /* Load stack of new thread */
-    mov rsp, [rbp + KTHREAD_KernelStack]
+    mov rsp, [r8 + KTHREAD_KernelStack]
 
     /* Reload APC bypass */
-    mov r8b, [rsp + SwApcBypass]
+    mov cl, [rsp + SwApcBypass]
 
     call KiSwapContextResume
 
@@ -178,13 +177,13 @@ PUBLIC KiSwapContextInternal
  *     The KiSwapContext routine switches context to another thread.
  *
  * BOOLEAN
- * KiSwapContext(KIRQL WaitIrql, PKTHREAD CurrentThread);
+ * KiSwapContext(KIRQL WaitIrql, PKTHREAD OldThread);
  *
- * \param WaitIrql <rcx>
+ * \param WaitIrql <cl>
- *     ...
+ *     The IRQL at wich the old thread is suspended
  *
- * \param CurrentThread <rdx>
+ * \param OldThread <rdx>
- *     Pointer to the KTHREAD of the current thread.
+ *     Pointer to the KTHREAD of the previous thread.
  *
  * \return
  *     The WaitStatus of the Target Thread.
@@ -205,7 +204,7 @@ PUBLIC KiSwapContext
     GENERATE_EXCEPTION_FRAME
 
     /* Do the swap with the registers correctly setup */
-    mov rcx, gs:[PcCurrentThread] /* Pointer to the new thread */
+    mov r8, gs:[PcCurrentThread] /* Pointer to the new thread */
     call KiSwapContextInternal
 
     /* Restore the registers from the KEXCEPTION_FRAME */

ntoskrnl/ke/amd64/thrdini.c

@@ -137,9 +137,9 @@ KiInitializeContextThread(IN PKTHREAD Thread,
 
 BOOLEAN
 KiSwapContextResume(
-    IN PKTHREAD NewThread,
+    _In_ BOOLEAN ApcBypass,
-    IN PKTHREAD OldThread,
+    _In_ PKTHREAD OldThread,
-    IN BOOLEAN ApcBypass)
+    _In_ PKTHREAD NewThread)
 {
     PKIPCR Pcr = (PKIPCR)KeGetPcr();
     PKPROCESS OldProcess, NewProcess;
@@ -187,14 +187,15 @@ KiSwapContextResume(
     if (NewThread->ApcState.KernelApcPending)
     {
         /* Are APCs enabled? */
-        if (!NewThread->SpecialApcDisable)
+        if ((NewThread->SpecialApcDisable == 0) &&
+            (ApcBypass == 0))
         {
-            /* Request APC delivery */
+            /* Return TRUE to indicate that we want APCs to be delivered */
-            if (!ApcBypass)
-                HalRequestSoftwareInterrupt(APC_LEVEL);
-            else
             return TRUE;
         }
+
+        /* Request an APC interrupt to be delivered later */
+        HalRequestSoftwareInterrupt(APC_LEVEL);
     }
 
     /* Return stating that no kernel APCs are pending*/

ntoskrnl/ke/amd64/trap.S

@@ -592,7 +592,7 @@ PUBLIC KiApcInterrupt
     mov cl, [rbp + KTRAP_FRAME_SegCs] // ProcessorMode
     and cl, 1
     mov rdx, 0                        // ExceptionFrame
-    mov r8, rdx                       // TrapFrame
+    mov r8, rbp                       // TrapFrame
     call KiDeliverApc
 
     /* Disable interrupts */
@@ -1000,8 +1000,47 @@ KiConvertToGuiThreadFailed:
 
 ENDFUNC
 
-KiExitToUserApc:
+;
-    int 3
+; VOID
+; KiDeliverApc(
+;     _In_ KPROCESSOR_MODE DeliveryMode,
+;     _In_ PKEXCEPTION_FRAME ExceptionFrame,
+;     _In_ PKTRAP_FRAME TrapFrame);
+;
+EXTERN KiDeliverApc:PROC
+
+PUBLIC KiInitiateUserApc
+.PROC KiInitiateUserApc
+
+    ; Generate a KEXCEPTION_FRAME on the stack
+    GENERATE_EXCEPTION_FRAME
+
+    ; Raise IRQL to APC_LEVEL
+    mov rax, APC_LEVEL
+    mov cr8, rax
+
+    ; Enable interrupts
+    sti
+
+    ; Get the current trap frame
+    mov rax, gs:[PcCurrentThread]
+    mov r8, [rax + KTHREAD_TrapFrame]
+
+    ; Call the C function
+    mov ecx, 1
+    mov rdx, rsp
+    call KiDeliverApc
+
+    ; Disable interrupts again
+    cli
+
+    ; Restore the registers from the KEXCEPTION_FRAME
+    RESTORE_EXCEPTION_STATE
+
+    ; Return
+    ret
+
+.ENDP
 
 
 PUBLIC KiInitializeSegments

sdk/include/asm/trapamd64.inc

@@ -150,6 +150,7 @@ ENDM
 MACRO(ExitTrap, Flags)
     LOCAL kernel_mode_return
     LOCAL IntsEnabled
+    LOCAL NoUserApc
 
 #if DBG
         /* Check previous irql */
@@ -181,8 +182,9 @@ MACRO(ExitTrap, Flags)
         /* Load current thread into r10 */
         mov r10, gs:[PcCurrentThread]
         cmp byte ptr [r10 + KTHREAD_UserApcPending], 0
-        jne KiExitToUserApc
+        je NoUserApc
-
+        call KiInitiateUserApc
+    NoUserApc:
     endif
 
 #if DBG