[NTOS:KE/x64] Implement processor freeze code

ntoskrnl/include/internal/amd64/ke.h

@@ -480,6 +480,11 @@ VOID
 KiUserCallbackExit(
     _In_ PKTRAP_FRAME TrapFrame);
 
+BOOLEAN
+KiProcessorFreezeHandler(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

ntoskrnl/include/internal/ke.h

@@ -982,6 +982,16 @@ VOID
 NTAPI
 KiInitMachineDependent(VOID);
 
+VOID
+NTAPI
+KxFreezeExecution(
+    VOID);
+
+VOID
+NTAPI
+KxThawExecution(
+    VOID);
+
 BOOLEAN
 NTAPI
 KeFreezeExecution(IN PKTRAP_FRAME TrapFrame,

ntoskrnl/ke/amd64/freeze.c

@@ -0,0 +1,161 @@
+/*
+ * PROJECT:     ReactOS Kernel
+ * LICENSE:     GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
+ * PURPOSE:     Processor freeze support for x64
+ * COPYRIGHT:   Copyright 2023 Timo Kreuzer <timo.kreuzer@reactos.org>
+ */
+
+/* INCLUDES *******************************************************************/
+
+#include <ntoskrnl.h>
+#define NDEBUG
+#include <debug.h>
+
+/* NOT INCLUDES ANYMORE ******************************************************/
+
+PKPRCB KiFreezeOwner;
+
+/* FUNCTIONS *****************************************************************/
+
+BOOLEAN
+KiProcessorFreezeHandler(
+    _In_ PKTRAP_FRAME TrapFrame,
+    _In_ PKEXCEPTION_FRAME ExceptionFrame)
+{
+    PKPRCB CurrentPrcb = KeGetCurrentPrcb();
+
+    /* Make sure this is a freeze request */
+    if (CurrentPrcb->IpiFrozen != IPI_FROZEN_STATE_TARGET_FREEZE)
+    {
+        /* Not a freeze request, return FALSE to signal it is unhandled */
+        return FALSE;
+    }
+
+    /* We are frozen now */
+    CurrentPrcb->IpiFrozen = IPI_FROZEN_STATE_FROZEN;
+
+    /* Save the processor state */
+    KiSaveProcessorState(TrapFrame, ExceptionFrame);
+
+    /* Wait for the freeze owner to release us */
+    while (CurrentPrcb->IpiFrozen != IPI_FROZEN_STATE_THAW)
+    {
+        YieldProcessor();
+        KeMemoryBarrier();
+    }
+
+    /* Restore the processor state */
+    KiRestoreProcessorState(TrapFrame, ExceptionFrame);
+
+    /* We are running again now */
+    CurrentPrcb->IpiFrozen = IPI_FROZEN_STATE_RUNNING;
+
+    /* Return TRUE to signal that we handled the freeze */
+    return TRUE;
+}
+
+VOID
+NTAPI
+KxFreezeExecution(
+    VOID)
+{
+    PKPRCB CurrentPrcb = KeGetCurrentPrcb();
+
+    /* Avoid blocking on recursive debug action */
+    if (KiFreezeOwner == CurrentPrcb)
+    {
+        return;
+    }
+
+    /* Try to acquire the freeze owner */
+    while (InterlockedCompareExchangePointer(&KiFreezeOwner, CurrentPrcb, NULL))
+    {
+        /* Someone else was faster. We expect an NMI to freeze any time.
+           Spin here until the freeze owner is available. */
+        while (KiFreezeOwner != NULL)
+        {
+            YieldProcessor();
+            KeMemoryBarrier();
+        }
+    }
+
+    /* We are the owner now */
+    CurrentPrcb->IpiFrozen = IPI_FROZEN_STATE_OWNER;
+
+    /* Loop all processors */
+    for (ULONG i = 0; i < KeNumberProcessors; i++)
+    {
+        PKPRCB TargetPrcb = KiProcessorBlock[i];
+        if (TargetPrcb != CurrentPrcb)
+        {
+            /* Nobody else is allowed to change IpiFrozen, except the freeze owner */
+            ASSERT(TargetPrcb->IpiFrozen == IPI_FROZEN_STATE_RUNNING);
+
+            /* Request target to freeze */
+            TargetPrcb->IpiFrozen = IPI_FROZEN_STATE_TARGET_FREEZE;
+        }
+    }
+
+    /* Send the freeze IPI */
+    KiIpiSend(KeActiveProcessors & ~CurrentPrcb->SetMember, IPI_FREEZE);
+
+    /* Wait for all targets to be frozen */
+    for (ULONG i = 0; i < KeNumberProcessors; i++)
+    {
+        PKPRCB TargetPrcb = KiProcessorBlock[i];
+        if (TargetPrcb != CurrentPrcb)
+        {
+            /* Wait for the target to be frozen */
+            while (TargetPrcb->IpiFrozen != IPI_FROZEN_STATE_FROZEN)
+            {
+                YieldProcessor();
+                KeMemoryBarrier();
+            }
+        }
+    }
+
+    /* All targets are frozen, we can continue */
+}
+
+VOID
+NTAPI
+KxThawExecution(
+    VOID)
+{
+    PKPRCB CurrentPrcb = KeGetCurrentPrcb();
+
+    /* Loop all processors */
+    for (ULONG i = 0; i < KeNumberProcessors; i++)
+    {
+        PKPRCB TargetPrcb = KiProcessorBlock[i];
+        if (TargetPrcb != CurrentPrcb)
+        {
+            /* Make sure they are still frozen */
+            ASSERT(TargetPrcb->IpiFrozen == IPI_FROZEN_STATE_FROZEN);
+
+            /* Request target to thaw */
+            TargetPrcb->IpiFrozen = IPI_FROZEN_STATE_THAW;
+        }
+    }
+
+    /* Wait for all targets to be running */
+    for (ULONG i = 0; i < KeNumberProcessors; i++)
+    {
+        PKPRCB TargetPrcb = KiProcessorBlock[i];
+        if (TargetPrcb != CurrentPrcb)
+        {
+            /* Wait for the target to be running again */
+            while (TargetPrcb->IpiFrozen != IPI_FROZEN_STATE_RUNNING)
+            {
+                YieldProcessor();
+                KeMemoryBarrier();
+            }
+        }
+    }
+
+    /* We are running again now */
+    CurrentPrcb->IpiFrozen = IPI_FROZEN_STATE_RUNNING;
+
+    /* Release the freeze owner */
+    InterlockedExchangePointer(&KiFreezeOwner, NULL);
+}

ntoskrnl/ke/amd64/traphandler.c

@@ -89,6 +89,14 @@ KiNmiInterruptHandler(
     _In_ PKTRAP_FRAME TrapFrame,
     _In_ PKEXCEPTION_FRAME ExceptionFrame)
 {
+    /* Check if this is a freeze */
+    if (KiProcessorFreezeHandler(TrapFrame, ExceptionFrame))
+    {
+        /* NMI was handled */
+        return;
+    }
+
+    /* Handle the NMI */
     KiHandleNmi();
 }
 

ntoskrnl/ke/freeze.c

@@ -49,7 +49,8 @@ KeFreezeExecution(IN PKTRAP_FRAME TrapFrame,
 #endif
 
 #ifdef CONFIG_SMP
-    // TODO: Add SMP support.
+    /* Architecture specific freeze code */
+    KxFreezeExecution();
 #endif
 
     /* Save the old IRQL to be restored on unfreeze */
@@ -64,7 +65,8 @@ NTAPI
 KeThawExecution(IN BOOLEAN Enable)
 {
 #ifdef CONFIG_SMP
-    // TODO: Add SMP support.
+    /* Architecture specific thaw code */
+    KxThawExecution();
 #endif
 
     /* Clear the freeze flag */

ntoskrnl/ke/i386/freeze.c

@@ -0,0 +1,28 @@
+/*
+ * PROJECT:     ReactOS Kernel
+ * LICENSE:     GPL-2.0-or-later (https://spdx.org/licenses/GPL-2.0-or-later)
+ * PURPOSE:     Processor freeze support for i386
+ * COPYRIGHT:
+ */
+
+/* INCLUDES *******************************************************************/
+
+#include <ntoskrnl.h>
+#define NDEBUG
+#include <debug.h>
+
+VOID
+NTAPI
+KxFreezeExecution(
+    VOID)
+{
+    UNIMPLEMENTED;
+}
+
+VOID
+NTAPI
+KxThawExecution(
+    VOID)
+{
+    UNIMPLEMENTED;
+}

ntoskrnl/ntos.cmake

@@ -316,6 +316,7 @@ if(ARCH STREQUAL "i386")
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/cpu.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/context.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/exp.c
+        ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/freeze.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/irqobj.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/kiinit.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/i386/ldt.c
@@ -350,6 +351,7 @@ elseif(ARCH STREQUAL "amd64")
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/context.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/cpu.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/except.c
+        ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/freeze.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/interrupt.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/ipi.c
         ${REACTOS_SOURCE_DIR}/ntoskrnl/ke/amd64/irql.c

sdk/include/ndk/amd64/ketypes.h

@@ -300,6 +300,21 @@ typedef enum
 #define IPI_PACKET_READY        8
 #define IPI_SYNCH_REQUEST       16
 
+//
+// Flags for KPRCB::IpiFrozen
+//
+// Values shown with !ipi extension in WinDbg:
+// 0 = [Running], 1 = [Unknown], 2 = [Frozen], 3 = [Thaw], 4 = [Freeze Owner]
+// 5 = [Target Freeze], 6-15 = [Unknown]
+// 0x20 = [Active] (flag)
+//
+#define IPI_FROZEN_STATE_RUNNING 0
+#define IPI_FROZEN_STATE_FROZEN 2
+#define IPI_FROZEN_STATE_THAW 3
+#define IPI_FROZEN_STATE_OWNER 4
+#define IPI_FROZEN_STATE_TARGET_FREEZE 5
+#define IPI_FROZEN_FLAG_ACTIVE 0x20
+
 //
 // PRCB Flags
 //