[NTOS:KE:X64] Implement KiConvertToGuiThread, KeSwitchKernelStack and support for win32k syscalls in KiSystemCallHandler

ntoskrnl/ke/amd64/stubs.c

@@ -8,6 +8,7 @@
 /* INCLUDES ******************************************************************/
 
 #include <ntoskrnl.h>
+#include <fltkernel.h>
 
 #define NDEBUG
 #include <debug.h>
@@ -17,6 +18,10 @@ KiRetireDpcListInDpcStack(
     PKPRCB Prcb,
     PVOID DpcStack);
 
+NTSTATUS
+KiConvertToGuiThread(
+    VOID);
+
 VOID
 NTAPI
 KiDpcInterruptHandler(VOID)
@@ -86,13 +91,66 @@ KeZeroPages(IN PVOID Address,
     RtlZeroMemory(Address, Size);
 }
 
+PVOID
+KiSwitchKernelStackHelper(
+    LONG_PTR StackOffset,
+    PVOID OldStackBase);
+
 PVOID
 NTAPI
 KeSwitchKernelStack(PVOID StackBase, PVOID StackLimit)
 {
-    UNIMPLEMENTED;
-    __debugbreak();
-    return NULL;
+    PKTHREAD CurrentThread;
+    PVOID OldStackBase;
+    LONG_PTR StackOffset;
+    SIZE_T StackSize;
+
+    /* Get the current thread */
+    CurrentThread = KeGetCurrentThread();
+
+    /* Save the old stack base */
+    OldStackBase = CurrentThread->StackBase;
+
+    /* Get the size of the current stack */
+    StackSize = (ULONG_PTR)CurrentThread->StackBase - CurrentThread->StackLimit;
+    ASSERT(StackSize <= (ULONG_PTR)StackBase - (ULONG_PTR)StackLimit);
+
+    /* Copy the current stack contents to the new stack */
+    RtlCopyMemory((PUCHAR)StackBase - StackSize,
+                  (PVOID)CurrentThread->StackLimit,
+                  StackSize);
+
+    /* Calculate the offset between the old and the new stack */
+    StackOffset = (PUCHAR)StackBase - (PUCHAR)CurrentThread->StackBase;
+
+    /* Disable interrupts while messing with the stack */
+    _disable();
+
+    /* Set the new trap frame */
+    CurrentThread->TrapFrame = (PKTRAP_FRAME)Add2Ptr(CurrentThread->TrapFrame,
+                                                     StackOffset);
+
+    /* Set the new initial stack */
+    CurrentThread->InitialStack = Add2Ptr(CurrentThread->InitialStack,
+                                          StackOffset);
+
+    /* Set the new stack limits */
+    CurrentThread->StackBase = StackBase;
+    CurrentThread->StackLimit = (ULONG_PTR)StackLimit;
+    CurrentThread->LargeStack = TRUE;
+
+    /* Adjust the PCR fields */
+    __addgsqword(FIELD_OFFSET(KPCR, NtTib.StackBase), StackOffset);
+    __addgsqword(FIELD_OFFSET(KIPCR, Prcb.RspBase), StackOffset);
+
+    /* Finally switch RSP to the new stack.
+       We pass OldStackBase to make sure it is not lost. */
+    OldStackBase = KiSwitchKernelStackHelper(StackOffset, OldStackBase);
+
+    /* Reenable interrupts */
+    _enable();
+
+    return OldStackBase;
 }
 
 NTSTATUS
@@ -313,9 +371,7 @@ KiSystemCallHandler(
     PULONG64 KernelParams, UserParams;
     ULONG ServiceNumber, Offset, Count;
     ULONG64 UserRsp;
-
-    DPRINT("Syscall #%ld\n", TrapFrame->Rax);
-    //__debugbreak();
+    NTSTATUS Status;
 
     /* Increase system call count */
     __addgsdword(FIELD_OFFSET(KIPCR, Prcb.KeSystemCalls), 1);
@@ -355,6 +411,40 @@ KiSystemCallHandler(
     /* Get descriptor table */
     DescriptorTable = (PVOID)((ULONG_PTR)Thread->ServiceTable + Offset);
 
+    /* Validate the system call number */
+    if (ServiceNumber >= DescriptorTable->Limit)
+    {
+        /* Check if this is a GUI call */
+        if (!(Offset & SERVICE_TABLE_TEST))
+        {
+            /* Fail the call */
+            TrapFrame->Rax = STATUS_INVALID_SYSTEM_SERVICE;
+            return (PVOID)NtSyscallFailure;
+        }
+
+        /* Convert us to a GUI thread -- must wrap in ASM to get new EBP */
+        Status = KiConvertToGuiThread();
+
+        /* Reload trap frame and descriptor table pointer from new stack */
+        TrapFrame = *(volatile PVOID*)&Thread->TrapFrame;
+        DescriptorTable = (PVOID)(*(volatile ULONG_PTR*)&Thread->ServiceTable + Offset);
+
+        if (!NT_SUCCESS(Status))
+        {
+            /* Set the last error and fail */
+            TrapFrame->Rax = Status;
+            return (PVOID)NtSyscallFailure;
+        }
+
+        /* Validate the system call number again */
+        if (ServiceNumber >= DescriptorTable->Limit)
+        {
+            /* Fail the call */
+            TrapFrame->Rax = STATUS_INVALID_SYSTEM_SERVICE;
+            return (PVOID)NtSyscallFailure;
+        }
+    }
+
     /* Get stack bytes and calculate argument count */
     Count = DescriptorTable->Number[ServiceNumber] / 8;
 

ntoskrnl/ke/amd64/trap.S

@@ -20,6 +20,7 @@ EXTERN KiGeneralProtectionFaultHandler:PROC
 EXTERN KiXmmExceptionHandler:PROC
 EXTERN KiDeliverApc:PROC
 EXTERN KiDpcInterruptHandler:PROC
+EXTERN PsConvertToGuiThread:PROC
 
 #ifdef _WINKD_
 EXTERN KdSetOwedBreakpoints:PROC
@@ -892,6 +893,28 @@ FUNC KiZwSystemService
 
 ENDFUNC
 
+PUBLIC KiConvertToGuiThread
+FUNC KiConvertToGuiThread
+
+    push rbp
+    .pushreg rbp
+    sub rsp, 32
+    .allocstack 32
+    .endprolog
+
+    /* Call the worker function */
+    call PsConvertToGuiThread
+
+    /* Adjust rsp */
+    add rsp, 32
+
+    /* Restore rbp */
+    pop rbp
+
+    /* return to the caller */
+    ret
+
+ENDFUNC
 
 KiExitToUserApc:
     int 3
@@ -934,6 +957,28 @@ KiInitializeSegments:
     mov gs, ax
     ret
 
+/*!
+ * VOID
+ * KiSwitchKernelStackHelper(
+ *     LONG_PTR StackOffset,
+ *     PVOID OldStackBase);
+ */
+PUBLIC KiSwitchKernelStackHelper
+KiSwitchKernelStackHelper:
+    
+    /* Pop return address from the current stack */
+    pop rax
+
+    /* Switch to new stack */
+    lea rsp, [rsp + rcx]
+
+    /* Push return address on the new stack */
+    push rax
+
+    /* Return on new stack */
+    mov rax, rdx
+    ret;
+
 
 #ifdef _MSC_VER
 #undef lgdt