[NTOS:KE] Improve KiSystemCallEntry64

ntoskrnl/ke/amd64/stubs.c

@@ -260,10 +260,7 @@ NtSyscallFailure(void)
 
 PVOID
 KiSystemCallHandler(
-    _In_ ULONG64 ReturnAddress,
+    VOID)
-    _In_ ULONG64 P2,
-    _In_ ULONG64 P3,
-    _In_ ULONG64 P4)
 {
     PKTRAP_FRAME TrapFrame;
     PKSERVICE_TABLE_DESCRIPTOR DescriptorTable;
@@ -275,12 +272,6 @@ KiSystemCallHandler(
     /* Get a pointer to the trap frame */
     TrapFrame = (PKTRAP_FRAME)((PULONG64)_AddressOfReturnAddress() + 1 + MAX_SYSCALL_PARAMS);
 
-    /* Save some values in the trap frame */
-    TrapFrame->Rip = ReturnAddress;
-    TrapFrame->Rdx = P2;
-    TrapFrame->R8 = P3;
-    TrapFrame->R9 = P4;
-
     /* Increase system call count */
     __addgsdword(FIELD_OFFSET(KIPCR, Prcb.KeSystemCalls), 1);
 
@@ -388,11 +379,11 @@ KiSystemCallHandler(
                 break;
 
             default:
-                __debugbreak();
+                ASSERT(FALSE);
                 break;
         }
     }
-    _SEH2_EXCEPT(1)
+    _SEH2_EXCEPT(EXCEPTION_EXECUTE_HANDLER)
     {
         TrapFrame->Rax = _SEH2_GetExceptionCode();
         return (PVOID)NtSyscallFailure;

ntoskrnl/ke/amd64/trap.S

@@ -747,10 +747,10 @@ EXTERN KiSystemCallHandler:PROC
  *--*/
 PUBLIC KiSystemCallEntry64
 .PROC KiSystemCallEntry64
-
+    /* The unwind info pretends we have a machine frame */
-    /* Old stack pointer is in rcx, lie and say we saved it in rbp */
+    .PUSHFRAME
-    .setframe rbp, 0
+    .ALLOCSTACK (KTRAP_FRAME_LENGTH + MAX_SYSCALL_PARAM_SIZE - MachineFrameLength)
-    .endprolog
+    .ENDPROLOG
 
     /* Swap gs to kernel, so we can access the PCR */
     swapgs
@@ -763,29 +763,26 @@ PUBLIC KiSystemCallEntry64
 
     /* Allocate a TRAP_FRAME and space for parameters */
     sub rsp, (KTRAP_FRAME_LENGTH + MAX_SYSCALL_PARAM_SIZE)
-#if DBG
-    /* Save rbp and load it with the old stack pointer */
-    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rbp], rbp
-    mov rbp, gs:[PcUserRsp]
-#endif
 
-    /* Save important registers in the trap frame */
+    /* Save volatile registers in the trap frame */
     mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rax], rax
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rip], rcx
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rdx], rdx
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_R8], r8
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_R9], r9
     mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rcx], r10
     mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_EFlags], r11
 
+    /* Store user stack pointer in the trap frame */
+    mov rax, gs:[PcUserRsp]
+    mov [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rsp], rax
+
     /* Set sane segments */
     mov ax, (KGDT64_R3_DATA or RPL_MASK)
     mov ds, ax
     mov es, ax
 
-.ENDP
+GLOBAL_LABEL KiSystemCall64Again
-
-.PROC KiSystemCall64Again
-
-    /* Old stack pointer is in rcx, lie and say we saved it in rbp */
-    .setframe rbp, 0
-    .endprolog
 
     /* Call the C-handler (will enable interrupts) */
     call KiSystemCallHandler
@@ -797,19 +794,9 @@ PUBLIC KiSystemCallEntry64
     mov r9, [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_R9]
     call rax
 
-.ENDP
+GLOBAL_LABEL KiSystemServiceExit
-
-PUBLIC KiSystemServiceExit
-.PROC KiSystemServiceExit
-
-    /* Old stack pointer is in rcx, lie and say we saved it in rbp */
-    .setframe rbp, 0
-    .endprolog
 
 #if DBG
-    /* Restore rbp */
-    mov rbp, [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rbp]
-
     test dword ptr [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_EFlags], HEX(200)
     jnz IntsEnabled
     int 3
@@ -838,6 +825,12 @@ no_user_apc_pending:
     /* Load user mode stack (It was copied to the trap frame) */
     mov rsp, [rsp + MAX_SYSCALL_PARAM_SIZE + KTRAP_FRAME_Rsp]
 
+    /* r8 points to the user stack */
+    mov r8, rsp
+
+    /* r9 matches rbp */
+    mov r9, rbp
+
     /* Swap gs back to user */
     swapgs