From e8fa736c7a5449154e4b0f4d5cc03de396a50496 Mon Sep 17 00:00:00 2001
From: Aiken Harris <harraiken91@gmail.com>
Date: Fri, 10 Jul 2026 08:03:58 +0200
Subject: [PATCH] Implement RAII guard for critical regions

---
 xtoskrnl/includes/ke/guard.hh | 94 ++++++++++++++++++++++++-----------
 xtoskrnl/ob/lifecycl.cc       |  3 +-
 xtoskrnl/ps/psmgr.cc          |  2 +-
 3 files changed, 66 insertions(+), 33 deletions(-)

diff --git a/xtoskrnl/includes/ke/guard.hh b/xtoskrnl/includes/ke/guard.hh
index fe51dda..f1ec90e 100644
--- a/xtoskrnl/includes/ke/guard.hh
+++ b/xtoskrnl/includes/ke/guard.hh
@@ -16,6 +16,40 @@
 /* Kernel Library */
 namespace KE
 {
+    class CriticalRegionGuard
+    {
+        private:
+            BOOLEAN Owned;
+            PKTHREAD SystemThread;
+
+        public:
+            CriticalRegionGuard(IN PKTHREAD Thread = NULLPTR,
+                                IN BOOLEAN Acquire = TRUE)
+            {
+                Owned = Acquire;
+                SystemThread = Thread;
+                if(Owned)
+                {
+                    if(!SystemThread)
+                    {
+                        SystemThread = KE::Processor::GetCurrentThread();
+                    }
+                    KE::KThread::EnterCriticalRegion(SystemThread);
+                }
+            }
+
+            ~CriticalRegionGuard()
+            {
+                if(Owned)
+                {
+                    KE::KThread::LeaveCriticalRegion(SystemThread);
+                }
+            }
+
+            CriticalRegionGuard(const CriticalRegionGuard&) = delete;
+            CriticalRegionGuard& operator=(const CriticalRegionGuard&) = delete;
+    };
+
     class QueuedSpinLockGuard
     {
         private:
@@ -47,36 +81,6 @@ namespace KE
             QueuedSpinLockGuard& operator=(const QueuedSpinLockGuard&) = delete;
     };
 
-    class SystemQueuedSpinLockGuard
-    {
-        private:
-            KSPIN_LOCK_QUEUE_LEVEL QueuedLockLevel;
-            BOOLEAN Owned;
-
-        public:
-            SystemQueuedSpinLockGuard(IN OUT KSPIN_LOCK_QUEUE_LEVEL LockLevel,
-                                      IN BOOLEAN Acquire = TRUE)
-            {
-                QueuedLockLevel = LockLevel;
-                Owned = Acquire;
-                if(Owned)
-                {
-                    KE::SpinLock::AcquireQueuedSpinLock(QueuedLockLevel);
-                }
-            }
-
-            ~SystemQueuedSpinLockGuard()
-            {
-                if(Owned)
-                {
-                    KE::SpinLock::ReleaseQueuedSpinLock(QueuedLockLevel);
-                }
-            }
-
-            SystemQueuedSpinLockGuard(const SystemQueuedSpinLockGuard&) = delete;
-            SystemQueuedSpinLockGuard& operator=(const SystemQueuedSpinLockGuard&) = delete;
-    };
-
     class SpinLockGuard
     {
         private:
@@ -106,6 +110,36 @@ namespace KE
             SpinLockGuard(const SpinLockGuard&) = delete;
             SpinLockGuard& operator=(const SpinLockGuard&) = delete;
     };
+
+    class SystemQueuedSpinLockGuard
+    {
+        private:
+            KSPIN_LOCK_QUEUE_LEVEL QueuedLockLevel;
+            BOOLEAN Owned;
+
+        public:
+            SystemQueuedSpinLockGuard(IN OUT KSPIN_LOCK_QUEUE_LEVEL LockLevel,
+                                      IN BOOLEAN Acquire = TRUE)
+            {
+                QueuedLockLevel = LockLevel;
+                Owned = Acquire;
+                if(Owned)
+                {
+                    KE::SpinLock::AcquireQueuedSpinLock(QueuedLockLevel);
+                }
+            }
+
+            ~SystemQueuedSpinLockGuard()
+            {
+                if(Owned)
+                {
+                    KE::SpinLock::ReleaseQueuedSpinLock(QueuedLockLevel);
+                }
+            }
+
+            SystemQueuedSpinLockGuard(const SystemQueuedSpinLockGuard&) = delete;
+            SystemQueuedSpinLockGuard& operator=(const SystemQueuedSpinLockGuard&) = delete;
+    };
 }
 
 #endif /* __XTOSKRNL_KE_GUARD_HH */
diff --git a/xtoskrnl/ob/lifecycl.cc b/xtoskrnl/ob/lifecycl.cc
index 1352a8d..4cccc31 100644
--- a/xtoskrnl/ob/lifecycl.cc
+++ b/xtoskrnl/ob/lifecycl.cc
@@ -840,7 +840,7 @@ OB::LifeCycle::DeleteObject(IN PVOID Object,
     if(CreatorInfo && !RTL::LinkedList::ListEmpty(&CreatorInfo->TypeList))
     {
         /* Acquire exclusive access to the object type tracking list */
-        KE::KThread::EnterCriticalRegion();
+        KE::CriticalRegionGuard CriticalRegion;
         KE::PushLock::AcquireExclusivePushLock(&ObjectType->TypeLock);
 
         /* Unlink the object from the creator type list */
@@ -848,7 +848,6 @@ OB::LifeCycle::DeleteObject(IN PVOID Object,
 
         /* Release the object type tracking list */
         KE::PushLock::ReleaseExclusivePushLock(&ObjectType->TypeLock);
-        KE::KThread::LeaveCriticalRegion();
     }
 
     /* Check if the object has an allocated name buffer */
diff --git a/xtoskrnl/ps/psmgr.cc b/xtoskrnl/ps/psmgr.cc
index 2952dd4..54d2452 100644
--- a/xtoskrnl/ps/psmgr.cc
+++ b/xtoskrnl/ps/psmgr.cc
@@ -57,7 +57,7 @@ PS::ProcessManager::InitializeProcessManager(VOID)
     RTL::Memory::ZeroMemory(&ObjectTypeInitializer, sizeof(ObjectTypeInitializer));
 
     /* Initialize common attributes */
-    ObjectTypeInitializer.InvalidAttributes = OBJECT_PERMANENT | OBJECT_EXCLUSIVE | OBJECT_OPENIF;
+    ObjectTypeInitializer.InvalidAttributes = OBJECT_EXCLUSIVE | OBJECT_OPENIF | OBJECT_PERMANENT;
     ObjectTypeInitializer.Length = sizeof(ObjectTypeInitializer);
     ObjectTypeInitializer.PoolType = NonPagedPool;
     ObjectTypeInitializer.SecurityRequired = TRUE;