[NTOS:KE] Use callout stacks for user callback stack expansion (optional)

This is a Windows 7+ feature, which allows deeper recursion of user mode callbacks for win32k. It is required for some modern applications, like latest versions of Chrome.

ntoskrnl/include/internal/ke.h

@@ -1084,6 +1084,16 @@ NTAPI
 KiGetStackPointer(
     VOID);
 
+VOID
+NTAPI
+KiReapCallbackStacks(
+    _In_ PKTHREAD Thread);
+
+VOID
+NTAPI
+KiRemoveThreadCalloutStack(
+    _In_ PKTHREAD Thread);
+
 #ifdef __cplusplus
 } // extern "C"
 

ntoskrnl/ke/amd64/usercall.c

@@ -8,6 +8,7 @@
 /* INCLUDES ******************************************************************/
 
 #include <ntoskrnl.h>
+#include <minwin/ntosifs.h>
 #define NDEBUG
 #include <debug.h>
 
@@ -146,7 +147,6 @@ KiUserModeCallout(
     KTRAP_FRAME CallbackTrapFrame;
     PKIPCR Pcr;
     ULONG_PTR InitialStack;
-    NTSTATUS Status;
 
     /* Get the current thread */
     CurrentThread = KeGetCurrentThread();
@@ -165,7 +165,7 @@ KiUserModeCallout(
     if ((InitialStack - KERNEL_STACK_SIZE) < CurrentThread->StackLimit)
     {
         /* We don't, we'll have to grow our stack */
-        Status = MmGrowKernelStack((PVOID)InitialStack);
+        NTSTATUS Status = MmGrowKernelStack((PVOID)InitialStack);
 
         /* Quit if we failed */
         if (!NT_SUCCESS(Status)) return Status;
@@ -194,7 +194,7 @@ KiUserModeCallout(
     /* Get PCR */
     Pcr = (PKIPCR)KeGetPcr();
 
-    /* Set user-mode dispatcher address as EIP */
+    /* Set user-mode dispatcher address as RIP */
     Pcr->TssBase->Rsp0 = InitialStack;
     Pcr->Prcb.RspBase = InitialStack;
     CallbackTrapFrame.Rip = (ULONG_PTR)KeUserCallbackDispatcher;
@@ -230,6 +230,98 @@ KiSetupUserCalloutFrame(
 #endif
 }
 
+#if ENABLE_CALLBACK_STACKS
+
+typedef struct _KUSER_CALLOUT_PARAMETER
+{
+    PVOID *Result;
+    PULONG ResultLength;
+    NTSTATUS CallbackStatus;
+} KUSER_CALLOUT_PARAMETER, *PKUSER_CALLOUT_PARAMETER;
+
+VOID
+NTAPI
+KiCallUserModeHelper(
+    _In_ PVOID Parameter)
+{
+    PKUSER_CALLOUT_PARAMETER CalloutParameter = (PKUSER_CALLOUT_PARAMETER)Parameter;
+    NTSTATUS Status;
+
+    /* Call the callback */
+    Status = KiCallUserMode(CalloutParameter->Result,
+                            CalloutParameter->ResultLength);
+
+    /* Store the callback status */
+    CalloutParameter->CallbackStatus = Status;
+}
+
+NTSTATUS
+NTAPI
+KiCallUserModeEx(
+    _Out_ PVOID *Result,
+    _Out_ PULONG ResultLength)
+{
+    KUSER_CALLOUT_PARAMETER CalloutParameter;
+    NTSTATUS Status;
+
+#if (NTDDI_VERSION >= NTDDI_WIN8)
+    Thread->CallbackNestingLevel++;
+#endif // (NTDDI_VERSION >= NTDDI_WIN8)
+
+    /* Setup the callout parameter */
+    CalloutParameter.Result = Result;
+    CalloutParameter.ResultLength = ResultLength;
+
+    /* Call the callback and make sure the stack is large enough */
+    Status = KeExpandKernelStackAndCallout(KiCallUserModeHelper,
+                                           &CalloutParameter,
+                                           KERNEL_STACK_SIZE);
+    if (NT_SUCCESS(Status))
+    {
+        /* Return the callback status */
+        Status = CalloutParameter.CallbackStatus;
+    }
+
+#if (NTDDI_VERSION >= NTDDI_WIN8)
+    Thread->CallbackNestingLevel--;
+#endif // (NTDDI_VERSION >= NTDDI_WIN8)
+
+    /* Return the status */
+    return Status;
+}
+
+VOID
+NTAPI
+KiReapCallbackStacks(
+    _In_ PKTHREAD Thread)
+{
+    ASSERT(Thread != KeGetCurrentThread());
+
+    /* Loop while we have a callback stack */
+    while (Thread->CallbackStack)
+    {
+        /* Get the current callout frame */
+        PKCALLOUT_FRAME CalloutFrame = Thread->CallbackStack;
+
+        /* Restore the previous callback stack */
+        Thread->CallbackStack = (PVOID)CalloutFrame->CallbackStack;
+
+        /* Check whether the saved trapframe is outside of the current stack */
+        if (((ULONG_PTR)CalloutFrame->TrapFrame < (ULONG_PTR)Thread->StackLimit) ||
+            ((ULONG_PTR)CalloutFrame->TrapFrame > (ULONG_PTR)Thread->StackBase))
+        {
+            /* We were on a callout stack, remove it */
+            KiRemoveThreadCalloutStack(Thread);
+        }
+
+#if (NTDDI_VERSION >= NTDDI_WIN8)
+        Thread->CallbackNestingLevel--;
+#endif // (NTDDI_VERSION >= NTDDI_WIN8)
+    }
+}
+
+#endif // ENABLE_CALLBACK_STACK
+
 NTSTATUS
 NTAPI
 KeUserModeCallback(
@@ -291,7 +383,11 @@ KeUserModeCallback(
 
         /* Jump to user mode */
         *UserStackPointer = (ULONG_PTR)CalloutFrame;
+#ifdef USE_CALLOUT_STACK
+        CallbackStatus = KiCallUserModeEx(Result, ResultLength);
+#else
         CallbackStatus = KiCallUserMode(Result, ResultLength);
+#endif
         if (CallbackStatus != STATUS_CALLBACK_POP_STACK)
         {
 #ifdef _M_IX86

ntoskrnl/ke/i386/usercall.c

@@ -358,6 +358,14 @@ KiUserModeCallout(PKCALLOUT_FRAME CalloutFrame)
     KiServiceExit(CallbackTrapFrame, 0);
 }
 
+VOID
+NTAPI
+KiReapCallbackStacks(
+    _In_ PKTHREAD Thread)
+{
+    __debugbreak();
+}
+
 /*++
  * @name NtCallbackReturn
  *

ntoskrnl/ntos.cmake

@@ -295,7 +295,7 @@ list(APPEND SOURCE
     ${REACTOS_SOURCE_DIR}/ntoskrnl/wmi/wmi.c
     ${REACTOS_SOURCE_DIR}/ntoskrnl/wmi/wmidrv.c)
 
-if(DLL_EXPORT_VERSION GREATER_EQUAL 0x600)
+if(DLL_EXPORT_VERSION GREATER_EQUAL 0x600 OR ENABLE_CALLBACK_STACKS)
 list(APPEND SOURCE
     ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/callout.c
 )

ntoskrnl/ps/kill.c

@@ -184,8 +184,14 @@ PspReapRoutine(IN PVOID Context)
             /* Get the first Thread Entry */
             Thread = CONTAINING_RECORD(NextEntry, ETHREAD, ReaperLink);
 
+            /* Reap callback stacks */
+            if (Thread->Tcb.CalloutActive)
+            {
+                KiReapCallbackStacks(&Thread->Tcb);
+            }
+
             /* Delete this entry's kernel stack */
-            MmDeleteKernelStack((PVOID)Thread->Tcb.StackBase,
+            MmDeleteKernelStack(Thread->Tcb.StackBase,
                                 Thread->Tcb.LargeStack);
             Thread->Tcb.InitialStack = NULL;
 
@@ -400,8 +406,14 @@ PspDeleteThread(IN PVOID ObjectBody)
     /* Check if we have a stack */
     if (Thread->Tcb.InitialStack)
     {
+        /* Reap callback stacks */
+        if (Thread->Tcb.CalloutActive)
+        {
+            KiReapCallbackStacks(&Thread->Tcb);
+        }
+
         /* Release it */
-        MmDeleteKernelStack((PVOID)Thread->Tcb.StackBase,
+        MmDeleteKernelStack(Thread->Tcb.StackBase,
                             Thread->Tcb.LargeStack);
     }
 

sdk/cmake/config.cmake

@@ -117,3 +117,11 @@ set(USE_DUMMY_PSEH FALSE CACHE BOOL
 
 set(DLL_EXPORT_VERSION "0x502" CACHE STRING
 "The NT version the user mode DLLs target.")
+
+if(ARCH STREQUAL "amd64")
+set(ENABLE_CALLBACK_STACKS TRUE CACHE BOOL
+"Whether to enable the use of callout stacks in user mode callbacks.")
+else()
+set(ENABLE_CALLBACK_STACKS FALSE CACHE BOOL
+"Whether to enable the use of callout stacks in user mode callbacks.")
+endif()