From ba833422b0c9fb0729732b7107ec6511c07e42ab Mon Sep 17 00:00:00 2001
From: belliash <belliash@codingworkshop.eu.org>
Date: Sun, 5 Nov 2023 09:50:04 +0100
Subject: [PATCH] Implement ExCompleteRundownProtection() and
 ExReInitializeRundownProtection() routines and add stub for
 ExWaitForRundownProtectionRelease()

---
 sdk/xtdk/exfuncs.h     | 12 ++++++++++
 sdk/xtdk/extypes.h     |  3 +++
 xtoskrnl/ex/rundown.c  | 51 ++++++++++++++++++++++++++++++++++++++++++
 xtoskrnl/xtoskrnl.spec | 10 ++++++++-
 4 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/sdk/xtdk/exfuncs.h b/sdk/xtdk/exfuncs.h
index 95cff23..8928661 100644
--- a/sdk/xtdk/exfuncs.h
+++ b/sdk/xtdk/exfuncs.h
@@ -18,12 +18,24 @@ XTFASTCALL
 BOOLEAN
 ExAcquireRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
 
+XTFASTCALL
+VOID
+ExCompleteRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
 XTFASTCALL
 VOID
 ExInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
 
+XTFASTCALL
+VOID
+ExReInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
 XTFASTCALL
 VOID
 ExReleaseRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor);
 
+XTFASTCALL
+VOID
+ExWaitForRundownProtectionRelease(IN PEX_RUNDOWN_REFERENCE Descriptor);
+
 #endif /* __XTDK_EXFUNCS_H */
diff --git a/sdk/xtdk/extypes.h b/sdk/xtdk/extypes.h
index 8c37a9d..195d7b6 100644
--- a/sdk/xtdk/extypes.h
+++ b/sdk/xtdk/extypes.h
@@ -14,6 +14,9 @@
 #include <ketypes.h>
 
 
+/* Rundown protection flags */
+#define EX_RUNDOWN_ACTIVE                               0x1
+
 /* Executive rundown protection structure definition */
 typedef struct _EX_RUNDOWN_REFERENCE
 {
diff --git a/xtoskrnl/ex/rundown.c b/xtoskrnl/ex/rundown.c
index 07ac56d..f5f5392 100644
--- a/xtoskrnl/ex/rundown.c
+++ b/xtoskrnl/ex/rundown.c
@@ -57,6 +57,23 @@ ExAcquireRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
     }
 }
 
+/**
+ * Marks the rundown descriptor as completed.
+ *
+ * @param Descriptor
+ *        Supplies a pointer to the descriptor to be completed.
+ *
+ * @return This routine does not return any value.
+ *
+ * @since NT 5.1
+ */
+XTFASTCALL
+VOID
+ExCompleteRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
+{
+    RtlAtomicExchangePointer(&Descriptor->Ptr, (PVOID)EX_RUNDOWN_ACTIVE);
+}
+
 /**
  * Initializes the rundown protection descriptor.
  *
@@ -75,6 +92,23 @@ ExInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
     Descriptor->Count = 0;
 }
 
+/**
+ * Reinitializes the rundown protection structure after it has been completed.
+ *
+ * @param Descriptor
+ *        Supplies a pointer to the descriptor to be reinitialized.
+ *
+ * @return This routine does not return any value.
+ *
+ * @since NT 5.1
+ */
+XTFASTCALL
+VOID
+ExReInitializeRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
+{
+    RtlAtomicExchangePointer(&Descriptor->Ptr, NULL);
+}
+
 /**
  * Releases the rundown protection for given descriptor.
  *
@@ -125,3 +159,20 @@ ExReleaseRundownProtection(IN PEX_RUNDOWN_REFERENCE Descriptor)
         }
     }
 }
+
+/**
+ * Waits until rundown protection calls are completed.
+ *
+ * @param Descriptor
+ *        Supplies a pointer to the rundown block descriptor.
+ *
+ * @return This routine does not return any value.
+ *
+ * @since NT 5.1
+ */
+XTFASTCALL
+VOID
+ExWaitForRundownProtectionRelease(IN PEX_RUNDOWN_REFERENCE Descriptor)
+{
+    UNIMPLEMENTED;
+}
diff --git a/xtoskrnl/xtoskrnl.spec b/xtoskrnl/xtoskrnl.spec
index 9ecfd35..699ff09 100644
--- a/xtoskrnl/xtoskrnl.spec
+++ b/xtoskrnl/xtoskrnl.spec
@@ -1,6 +1,10 @@
+# XTOS exports
 @ fastcall ExAcquireRundownProtection(ptr)
+@ fastcall ExCompleteRundownProtection(ptr)
 @ fastcall ExInitializeRundownProtection(ptr)
+@ fastcall ExReInitializeRundownProtection(ptr)
 @ fastcall ExReleaseRundownProtection(ptr)
+@ fastcall ExWaitForRundownProtectionRelease(ptr)
 @ cdecl HlIoPortInByte(ptr)
 @ cdecl HlIoPortInLong(ptr)
 @ cdecl HlIoPortInShort(ptr)
@@ -29,4 +33,8 @@
 @ cdecl RtlWideStringConcatenate(wstr wstr long)
 @ cdecl RtlWideStringLength(wstr long)
 @ cdecl RtlWideStringTokenize(wstr wstr wstr)
-@ stdcall RtlZeroMemory(ptr long)
\ No newline at end of file
+@ stdcall RtlZeroMemory(ptr long)
+
+
+# NT compatibilty layer exports
+@ fastcall ExRundownCompleted(ptr) ExCompleteRundownProtection