[NTOS:KE] Implement KeFreezeExecution and KeThawExecution for SMP

DarkFire01 · Dec 3, 2023 · 105f6f4 · 105f6f4
1 parent d8ba592
commit 105f6f4
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 6 deletions.
diff --git a/ntoskrnl/include/internal/ke.h b/ntoskrnl/include/internal/ke.h
@@ -11,6 +11,11 @@ extern "C"
 
 /* INTERNAL KERNEL TYPES ****************************************************/
 
+/* Matches windbg expectation */
+#define IPI_FROZEN_RUNNING 0
+#define IPI_FROZEN_HALTED 2
+#define IPI_FROZEN_THAWING 3
+
 typedef struct _WOW64_PROCESS
 {
   PVOID Wow64;

diff --git a/ntoskrnl/kd64/kdapi.c b/ntoskrnl/kd64/kdapi.c
@@ -1260,6 +1260,27 @@ KdpNotSupported(IN PDBGKD_MANIPULATE_STATE64 State)
                  &KdpContext);
 }
 
+KCONTINUE_STATUS
+NTAPI
+KeSwitchFrozenProcessor(IN USHORT ProcessorNumber)
+{
+    PKPRCB Prcb, TargetPrcb;
+    Prcb = KeGetCurrentPrcb();
+
+    if (ProcessorNumber <= KeNumberProcessors)
+    {
+        KdpDprintf("Processor Switch triggered Procesor: %d\n", ProcessorNumber);
+        TargetPrcb = KiProcessorBlock[ProcessorNumber];
+        InterlockedExchange((LONG*)&TargetPrcb->IpiFrozen, IPI_FROZEN_RUNNING);
+        InterlockedExchange((LONG*)&Prcb->IpiFrozen, IPI_FROZEN_HALTED);
+        return ContinueSuccess;
+    }
+    else
+    {
+        return ContinueProcessorReselected;
+    }
+}
+
 KCONTINUE_STATUS
 NTAPI
 KdpSendWaitContinue(IN ULONG PacketType,
@@ -1268,6 +1289,7 @@ KdpSendWaitContinue(IN ULONG PacketType,
                     IN OUT PCONTEXT Context)
 {
     STRING Data, Header;
+    KCONTINUE_STATUS StatusCheck;
     DBGKD_MANIPULATE_STATE64 ManipulateState;
     ULONG Length;
     KDSTATUS RecvCode;
@@ -1469,10 +1491,10 @@ KdpSendWaitContinue(IN ULONG PacketType,
                 break;
 
             case DbgKdSwitchProcessor:
-
-                /* TODO */
-                KdpDprintf("Processor Switch support is unimplemented!\n");
-                KdpNotSupported(&ManipulateState);
+                KdRestore(FALSE);
+                StatusCheck = KeSwitchFrozenProcessor(ManipulateState.Processor);
+                KdSave(FALSE);
+                return StatusCheck;
                 break;
 
             case DbgKdPageInApi:

diff --git a/ntoskrnl/ke/freeze.c b/ntoskrnl/ke/freeze.c
@@ -19,13 +19,55 @@
 KIRQL KiOldIrql;
 ULONG KiFreezeFlag;
 
+/* PRIVATE FUNCTIONS ***********************************************************/
+
+VOID
+NTAPI
+KiFreezeTargetExecution(_In_ PKTRAP_FRAME TrapFrame,
+                        _In_ PKEXCEPTION_FRAME ExceptionFrame)
+{
+    EXCEPTION_RECORD ExceptionRecord;
+    PKPRCB Prcb;
+    Prcb = KeGetCurrentPrcb();
+
+    if (TrapFrame)
+        KiSaveProcessorState(TrapFrame, ExceptionFrame);
+    /* Multiple processors can write this value */
+    InterlockedExchange((LONG*)&Prcb->IpiFrozen, IPI_FROZEN_HALTED);
+
+    /* Wait for triggering AP to give the go ahead to thaw */
+    while (Prcb->IpiFrozen != IPI_FROZEN_THAWING)
+    {
+        /* We only continue on if we are thawing, otherwise something else has happened! */
+        if (Prcb->IpiFrozen == IPI_FROZEN_RUNNING)
+        {
+            RtlZeroMemory(&ExceptionRecord, sizeof(EXCEPTION_RECORD));
+            ExceptionRecord.ExceptionAddress = (PVOID)Prcb->ProcessorState.ContextFrame.Eip;
+            KdpSwitchProcessor(&ExceptionRecord,
+                               (PCONTEXT)&Prcb->ProcessorState, FALSE);
+        }
+    }
+
+    KeFlushCurrentTb();
+    /* Notify AP we're running once again */
+    InterlockedExchange((LONG*)&Prcb->IpiFrozen, IPI_FROZEN_RUNNING);
+}
+
 /* FUNCTIONS ******************************************************************/
 
 BOOLEAN
 NTAPI
 KeFreezeExecution(IN PKTRAP_FRAME TrapFrame,
                   IN PKEXCEPTION_FRAME ExceptionFrame)
 {
+#ifdef CONFIG_SMP
+    KAFFINITY TargetAffinity;
+    PKPRCB TargetPrcb;
+    KAFFINITY Current;
+    PKPRCB Prcb;
+    LONG i;
+#endif
+
     BOOLEAN Enable;
     KIRQL OldIrql;
 
@@ -49,7 +91,27 @@ KeFreezeExecution(IN PKTRAP_FRAME TrapFrame,
 #endif
 
 #ifdef CONFIG_SMP
-    // TODO: Add SMP support.
+    Prcb = KeGetCurrentPrcb();
+    TargetAffinity = KeActiveProcessors;
+    TargetAffinity &= ~Prcb->SetMember;
+    if (TargetAffinity)
+    {
+        for (i = 0, Current = 1; i < KeNumberProcessors; i++, Current <<= 1)
+        {
+            if (TargetAffinity & Current)
+            {
+                /* stop target processor */
+                KiIpiSend(Current, IPI_FREEZE);
+                TargetPrcb = KiProcessorBlock[i];
+
+                /* Await for this processor to be frozen*/
+                while (TargetPrcb->IpiFrozen != IPI_FROZEN_HALTED)
+                {
+                    /* Do nothing, we're trying to synch */
+                }
+            }
+        }
+    }
 #endif
 
     /* Save the old IRQL to be restored on unfreeze */
@@ -64,7 +126,31 @@ NTAPI
 KeThawExecution(IN BOOLEAN Enable)
 {
 #ifdef CONFIG_SMP
-    // TODO: Add SMP support.
+    KAFFINITY TargetAffinity;
+    PKPRCB TargetPrcb;
+    KAFFINITY Current;
+    PKPRCB Prcb;
+    LONG i;
+
+    Prcb = KeGetCurrentPrcb();
+    TargetAffinity = KeActiveProcessors;
+    TargetAffinity &= ~Prcb->SetMember;
+
+    /* Loop through every processor */
+    for (i = 0, Current = 1; i < KeNumberProcessors; i++, Current <<= 1)
+    {
+        if (TargetAffinity & Current)
+        {
+            TargetPrcb = KiProcessorBlock[i];
+
+            /* Multiple processors can write this value */
+            InterlockedExchange((LONG*)&TargetPrcb->IpiFrozen, IPI_FROZEN_THAWING);
+            while (Prcb->IpiFrozen != IPI_FROZEN_RUNNING)
+            {
+                /* Do nothing we're waiting for ready */
+            }
+        }
+    }
 #endif
 
     /* Clear the freeze flag */