diff --git a/NTOSKRNL/KE/threadawarespin.cpp b/NTOSKRNL/KE/threadawarespin.cpp new file mode 100644 index 0000000..d8eb962 --- /dev/null +++ b/NTOSKRNL/KE/threadawarespin.cpp @@ -0,0 +1,144 @@ +/* + * PROJECT: Alcyone System Kernel + * LICENSE: BSD Clause 3 + * PURPOSE: Thread Aware Spin Lock + * NT KERNEL: 5.11.9360 + * COPYRIGHT: 2023-2029 Dibymartanda Samanta <> + */ + + /* + Thread Aware Spinlock, designed keeping following things in mind: + * + *Uses a ticket system (Ticket and Serving counters) to ensure first-come, first-served order + *Tracks the current owner of the lock (OwnerThread), ensuring easier debugging of deadlock + *Reentrant capability: A thread can "acquire" the lock multiple times if it already owns it, preventing self deadlocking + *DeadLock Prevention: Detects and prevents attempts by a thread to re-acquire a lock it already holds, prevent non owning thread from releasing a lock. + * + */ + + class KeThreadAwareSpinLockImpl +{ +public: + volatile LONG Ticket; + volatile LONG Serving; + PKTHREAD OwnerThread; + KPRIORITY OriginalPriority; + BOOLEAN IsBoosted; + + KeThreadAwareSpinLockImpl() + : Ticket(0), Serving(0), OwnerThread(nullptr), + OriginalPriority(0), IsBoosted(FALSE) {} + + void Initialize() + { + Ticket = 0; + Serving = 0; + OwnerThread = nullptr; + OriginalPriority = 0; + IsBoosted = FALSE; + } + + void Acquire() + { + LONG CurrentTicket = {0}; + PKTHREAD CurrentThread = KeGetCurrentThread(); + + if (OwnerThread == CurrentThread) + { + /* Already owned return the thread */ + return; + } + + CurrentTicket = InterlockedIncrement(&Ticket) - 1; + + while (CurrentTicket != Serving) + { + YieldProcessor(); + } + + OwnerThread = CurrentThread; + BoostPriority(); + } + + void Release() + { + PKTHREAD CurrentThread = KeGetCurrentThread(); + + if (OwnerThread != CurrentThread) + { + ASSERT(FALSE); + return; + } + + RestorePriority(); + OwnerThread = nullptr; + InterlockedIncrement(&Serving); + } + +private: + void BoostPriority() +{ + if (!IsBoosted) + { + OriginalPriority = KeGetPriorityThread(OwnerThread); + + /*Boost to the highest priority below real-time range*/ + KPRIORITY BoostPriority = HIGH_PRIORITY; + + /* Ensure we don't lower the priority if it's already higher*/ + if (OriginalPriority < BoostPriority) + { + KeSetPriorityThread(OwnerThread, BoostPriority); + IsBoosted = TRUE; + } + } + + void RestorePriority() + { + if (IsBoosted) + { + KeSetPriorityThread(OwnerThread, OriginalPriority); + IsBoosted = FALSE; + } + } +}; + +typedef KeThreadAwareSpinLockImpl* PKETHREADAWARESPINLOCK; + +/* Exported Function */ +VOID +NTAPI +KeThreadAwareSpinLockInitialize( + _Out_ PKETHREADAWARESPINLOCK SpinLock) +{ + if (SpinLock) + { + SpinLock->Initialize(); + } +} + +VOID +NTAPI +KeThreadAwareSpinLockAcquire( + _Inout_ PKETHREADAWARESPINLOCK SpinLock) +{ + if (SpinLock) + { + SpinLock->Acquire(); + } +} + +VOID +NTAPI +KeThreadAwareSpinLockRelease( + _Inout_ PKETHREADAWARESPINLOCK SpinLock) +{ + if (SpinLock) + { + SpinLock->Release(); + } +} + + + +