From 2827bb400f324196af5c253d151086dea1a63dcb Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Sat, 18 Mar 2023 11:11:06 +0100
Subject: [PATCH] Implement ExAcquireRundownProtection() and
 ExReleaseRundownProtection() routines

---
 sdk/xtdk/exfuncs.h    | 28 ++++++++++++
 sdk/xtdk/extypes.h    |  8 ++++
 sdk/xtdk/xtkmapi.h    |  1 +
 xtoskrnl/ex/rundown.c | 99 +++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 136 insertions(+)
 create mode 100644 sdk/xtdk/exfuncs.h

diff --git a/sdk/xtdk/exfuncs.h b/sdk/xtdk/exfuncs.h
new file mode 100644
index 0000000..385cf7e
--- /dev/null
+++ b/sdk/xtdk/exfuncs.h
@@ -0,0 +1,28 @@
+/**
+ * PROJECT:         ExectOS
+ * COPYRIGHT:       See COPYING.md in the top level directory
+ * FILE:            sdk/xtdk/exfuncs.h
+ * DESCRIPTION:     Kernel executive routines
+ * DEVELOPERS:      Rafal Kupiec <belliash@codingworkshop.eu.org>
+ */
+
+#ifndef __XTDK_EXFUNCS_H
+#define __XTDK_EXFUNCS_H
+
+#include <xtbase.h>
+#include <xttypes.h>
+
+
+XTFASTCALL
+BOOLEAN
+ExAcquireRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
+XTFASTCALL
+VOID
+ExInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
+XTFASTCALL
+VOID
+ExReleaseRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
+#endif /* __XTDK_EXFUNCS_H */
diff --git a/sdk/xtdk/extypes.h b/sdk/xtdk/extypes.h
index 371c6d1..8c37a9d 100644
--- a/sdk/xtdk/extypes.h
+++ b/sdk/xtdk/extypes.h
@@ -11,6 +11,7 @@
 
 #include <xtbase.h>
 #include <xttypes.h>
+#include <ketypes.h>
 
 
 /* Executive rundown protection structure definition */
@@ -23,4 +24,11 @@ typedef struct _EX_RUNDOWN_REFERENCE
     };
 } EX_RUNDOWN_REFERENCE, *PEX_RUNDOWN_REFERENCE;
 
+/* Executive rundown wait block definition */
+typedef struct _EX_RUNDOWN_WAIT_BLOCK
+{
+    ULONG_PTR Count;
+    KEVENT WakeEvent;
+} EX_RUNDOWN_WAIT_BLOCK, *PEX_RUNDOWN_WAIT_BLOCK;
+
 #endif /* __XTDK_EXTYPES_H */
diff --git a/sdk/xtdk/xtkmapi.h b/sdk/xtdk/xtkmapi.h
index db672dc..9002b25 100644
--- a/sdk/xtdk/xtkmapi.h
+++ b/sdk/xtdk/xtkmapi.h
@@ -42,6 +42,7 @@
 #include ARCH_HEADER(mmtypes.h)
 
 /* XT routines */
+#include <exfuncs.h>
 #include <hlfuncs.h>
 #include <kefuncs.h>
 #include <mmfuncs.h>
diff --git a/xtoskrnl/ex/rundown.c b/xtoskrnl/ex/rundown.c
index 193398e..8a22722 100644
--- a/xtoskrnl/ex/rundown.c
+++ b/xtoskrnl/ex/rundown.c
@@ -9,6 +9,54 @@
 #include <xtos.h>
 
 
+/**
+ * Acquires the rundown protection for given descriptor.
+ *
+ * @param Descriptor
+ *        Supplies a pointer to the rundown block descriptor.
+ *
+ * @return This routine returns TRUE if protection acquired successfully, or FALSE otherwise.
+ *
+ * @since NT 5.1
+ */
+XTFASTCALL
+BOOLEAN
+ExAcquireRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
+{
+    ULONG_PTR CurrentValue, NewValue;
+
+    /* Get current value */
+    CurrentValue = Descriptor->Count;
+
+    /* Main loop execution */
+    while(TRUE)
+    {
+        /* Make sure protection is not active yet */
+        if(CurrentValue & 0x1)
+        {
+            /* Already active, nothing to do */
+            return FALSE;
+        }
+
+        /* Attempt to increment the usage count */
+        NewValue = CurrentValue + 2;
+
+        /* Exchange the value */
+        NewValue = (ULONG_PTR)RtlInterlockedCompareExchangePointer(&Descriptor->Ptr, (PVOID)NewValue,
+                                                                   (PVOID)CurrentValue);
+
+        /* Make sure protection acquired */
+        if(NewValue == CurrentValue)
+        {
+            /* Successfully acquired protection */
+            return TRUE;
+        }
+
+        /* Update value and try once again */
+        CurrentValue = NewValue;
+    }
+}
+
 /**
  * Initializes the rundown protection descriptor.
  *
@@ -26,3 +74,54 @@ ExInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
     /* Reset descriptor counter */
     Descriptor->Count = 0;
 }
+
+/**
+ * Releases the rundown protection for given descriptor.
+ *
+ * @param Descriptor
+ *        Supplies a pointer to the descriptor to be initialized.
+ *
+ * @return This routine does not return any value.
+ *
+ * @since NT 5.1
+ */
+XTFASTCALL
+VOID
+ExReleaseRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
+{
+    ULONG_PTR CurrentValue, NewValue;
+    PEX_RUNDOWN_WAIT_BLOCK WaitBlock;
+
+    CurrentValue = Descriptor->Count;
+
+    while(TRUE)
+    {
+        if(CurrentValue & 0x1)
+        {
+            WaitBlock = (PEX_RUNDOWN_WAIT_BLOCK)(CurrentValue & ~0x1);
+
+            if(!RtlInterlockedDecrementLongPtr((PLONG_PTR)&WaitBlock->Count))
+            {
+                KeSetEvent(&WaitBlock->WakeEvent, 0, FALSE);
+            }
+
+            break;
+        }
+        else
+        {
+            /* Attempt to decrement the usage count */
+            NewValue = CurrentValue - 2;
+
+            /* Exchange the value */
+            NewValue = (ULONG_PTR)RtlInterlockedCompareExchangePointer(&Descriptor->Ptr, (PVOID)NewValue,
+                                                                       (PVOID)CurrentValue);
+
+            if(NewValue == CurrentValue)
+            {
+                break;
+            }
+
+            CurrentValue = NewValue;
+        }
+    }
+}