[NTOS:KE/x64] Handle NMI vs swapgs race condition

ntoskrnl/include/internal/amd64/intrin_i.h

@@ -91,6 +91,10 @@ static __inline__ __attribute__((always_inline)) void __str(unsigned short *Dest
 	__asm__ __volatile__("str %0" : : "m"(*Destination) : "memory");
 }
 
+static __inline__ __attribute__((always_inline)) void __swapgs(void)
+{
+	__asm__ __volatile__("swapgs" : : : "memory");
+}
 
 #elif defined(_MSC_VER)
 
@@ -106,6 +110,7 @@ void __ltr(unsigned short Source);
 
 void __str(unsigned short *Destination);
 
+void __swapgs(void);
 
 #else
 #error Unknown compiler for inline assembler

ntoskrnl/ke/amd64/traphandler.c

@@ -89,15 +89,39 @@ KiNmiInterruptHandler(
     _In_ PKTRAP_FRAME TrapFrame,
     _In_ PKEXCEPTION_FRAME ExceptionFrame)
 {
+    BOOLEAN ManualSwapGs = FALSE;
+
+    /* Check if the NMI came from kernel mode */
+    if ((TrapFrame->SegCs & MODE_MASK) == 0)
+    {
+        /* Check if GS base is already kernel mode. This is needed, because
+           we might be interrupted during an interrupt/exception from user-mode
+           before the swapgs instruction. */
+        if ((LONG64)__readmsr(MSR_GS_BASE) >= 0)
+        {
+            /* Swap GS to kernel */
+            __swapgs();
+            ManualSwapGs = TRUE;
+        }
+    }
+
     /* Check if this is a freeze */
     if (KiProcessorFreezeHandler(TrapFrame, ExceptionFrame))
     {
         /* NMI was handled */
-        return;
+        goto Exit;
     }
 
     /* Handle the NMI */
     KiHandleNmi();
+
+Exit:
+    /* Check if we need to swap GS back */
+    if (ManualSwapGs)
+    {
+        /* Swap GS back to user */
+        __swapgs();
+    }
 }
 
 #define MAX_SYSCALL_PARAMS 16