Skip to content

fix(heap): add CRITICAL_SECTION to protect tracking linked lists #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TorinKS merged 1 commit into from
Mar 22, 2026
Merged

fix(heap): add CRITICAL_SECTION to protect tracking linked lists #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions include/KernelHeap.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ typedef struct _GLOBAL_STATE {
HANDLE HeapHandle;
BOOL SuppressErrors; /* Control error message output */
BOOL TrackFreedMemory; /* Control whether to track freed memory for double-free detection */
CRITICAL_SECTION TrackingLock; /* Protects MemoryAllocations, FreedMemoryList, and counters */
} GLOBAL_STATE;

/* Function declarations - implementations in KernelHeap.c */
Expand Down
52 changes: 47 additions & 5 deletions src/KernelHeap.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,12 +48,14 @@ GLOBAL_STATE* GetGlobalState(void) {
temp->NextAllocationId = 1;
temp->SuppressErrors = FALSE;
temp->TrackFreedMemory = TRUE;
InitializeCriticalSection(&temp->TrackingLock);

/* Atomically publish: if another thread won the race, free our copy */
if (InterlockedCompareExchangePointer(
(PVOID*)&g_WinKernelLite_GlobalState, temp, NULL) != NULL) {
/* Another thread initialized first - free our allocation */
HEAP_TRACE("GetGlobalState: Lost race, freeing duplicate allocation at %p", temp);
DeleteCriticalSection(&temp->TrackingLock);
HeapFree(GetProcessHeap(), 0, temp);
} else {
HEAP_INFO("GetGlobalState: Global state initialized successfully");
Expand All @@ -68,6 +70,8 @@ BOOL InitHeap(void) {
return FALSE;
}

EnterCriticalSection(&state->TrackingLock);

/* Only reset if we have no active allocations */
if (IsListEmpty(&state->MemoryAllocations)) {
state->AllocationCount = 0;
Expand All @@ -82,6 +86,8 @@ BOOL InitHeap(void) {
}
}

LeaveCriticalSection(&state->TrackingLock);

return TRUE;
}

Expand All @@ -92,6 +98,8 @@ void CleanupHeap(void) {
PMEMORY_TRACKING_ENTRY entry;
PFREED_MEMORY_ENTRY freedEntry;

EnterCriticalSection(&state->TrackingLock);

/* Clean up allocation tracking entries */
current = state->MemoryAllocations.Flink;
while (current != &state->MemoryAllocations) {
Expand All @@ -100,6 +108,7 @@ void CleanupHeap(void) {
HeapFree(GetProcessHeap(), 0, entry);
current = next;
}
InitializeListHead(&state->MemoryAllocations);

/* Clean up freed memory tracking entries */
current = state->FreedMemoryList.Flink;
Expand All @@ -109,6 +118,10 @@ void CleanupHeap(void) {
HeapFree(GetProcessHeap(), 0, freedEntry);
current = next;
}
InitializeListHead(&state->FreedMemoryList);

LeaveCriticalSection(&state->TrackingLock);
DeleteCriticalSection(&state->TrackingLock);

HeapFree(GetProcessHeap(), 0, state);
g_WinKernelLite_GlobalState = NULL;
Expand Down Expand Up @@ -136,17 +149,21 @@ void TrackAllocation(PVOID Address, SIZE_T Size, const char* FileName, int LineN
entry->Size = Size;
entry->FileName = FileName;
entry->LineNumber = LineNumber;
entry->AllocationId = state->NextAllocationId++;

InsertHeadList(&state->MemoryAllocations, &entry->ListEntry);
EnterCriticalSection(&state->TrackingLock);

entry->AllocationId = state->NextAllocationId++;
InsertHeadList(&state->MemoryAllocations, &entry->ListEntry);
state->AllocationCount++;
state->TotalBytesAllocated += Size;
state->CurrentBytesAllocated += Size;
if (state->CurrentBytesAllocated > state->PeakBytesAllocated)
state->PeakBytesAllocated = state->CurrentBytesAllocated;

LeaveCriticalSection(&state->TrackingLock);
}

/* NOTE: Caller MUST hold state->TrackingLock */
void TrackFreedMemoryLocked(PVOID Address, SIZE_T Size, const char* AllocFileName, int AllocLineNumber, const char* FreeFileName, int FreeLineNumber, ULONGLONG AllocationId) {
GLOBAL_STATE* state = GetGlobalState();
PFREED_MEMORY_ENTRY entry;
Expand Down Expand Up @@ -253,6 +270,8 @@ BOOL UntrackAllocation(PVOID Address, const char* FreeFileName, int FreeLineNumb
return FALSE;
}

EnterCriticalSection(&state->TrackingLock);

current = state->MemoryAllocations.Flink;

while (current != &state->MemoryAllocations) {
Expand All @@ -275,10 +294,13 @@ BOOL UntrackAllocation(PVOID Address, const char* FreeFileName, int FreeLineNumb
current = current->Flink;
}

/* TrackFreedMemoryLocked expects lock held */
if (found && state->TrackFreedMemory) {
TrackFreedMemoryLocked(Address, allocSize, allocFileName, allocLineNumber, FreeFileName, FreeLineNumber, allocationId);
}

LeaveCriticalSection(&state->TrackingLock);

return found;
}

Expand Down Expand Up @@ -331,6 +353,8 @@ void _ExFreePoolWithTracking(PVOID pointer, const char* FileName, int LineNumber
state->HeapHandle, state->TrackFreedMemory);

if (state->TrackFreedMemory) {
EnterCriticalSection(&state->TrackingLock);

PLIST_ENTRY current = state->MemoryAllocations.Flink;
while (current != &state->MemoryAllocations) {
PMEMORY_TRACKING_ENTRY entry = CONTAINING_RECORD(current, MEMORY_TRACKING_ENTRY, ListEntry);
Expand All @@ -345,11 +369,14 @@ void _ExFreePoolWithTracking(PVOID pointer, const char* FileName, int LineNumber
HEAP_TRACE("_ExFreePoolWithTracking: Checking for double-free of %p (ID: %llu)", pointer, currentAllocationId);
isDoubleFree = CheckForDoubleFree(pointer, FileName, LineNumber, currentAllocationId);
if (isDoubleFree) {
LeaveCriticalSection(&state->TrackingLock);
HEAP_WARN("_ExFreePoolWithTracking: Double-free detected for %p, returning without freeing", pointer);
return;
}
HEAP_TRACE("_ExFreePoolWithTracking: No double-free detected for %p", pointer);
}

LeaveCriticalSection(&state->TrackingLock);
}

HEAP_TRACE("_ExFreePoolWithTracking: Calling UntrackAllocation for %p", pointer);
Expand Down Expand Up @@ -399,21 +426,27 @@ void _ExFreePoolWithTracking(PVOID pointer, const char* FileName, int LineNumber
BOOL IsValidHeapPointer(PVOID pointer) {
GLOBAL_STATE* state = GetGlobalState();
PLIST_ENTRY current;
BOOL result = FALSE;

if (!state || !pointer) return FALSE;

EnterCriticalSection(&state->TrackingLock);

current = state->MemoryAllocations.Flink;
while (current != &state->MemoryAllocations) {
PMEMORY_TRACKING_ENTRY entry = CONTAINING_RECORD(current, MEMORY_TRACKING_ENTRY, ListEntry);

if (entry->Address == pointer) {
return TRUE;
result = TRUE;
break;
}

current = current->Flink;
}

return FALSE;
LeaveCriticalSection(&state->TrackingLock);

return result;
}

void PrintMemoryLeaks(void) {
Expand All @@ -431,6 +464,8 @@ void PrintMemoryLeaks(void) {
leakCount = 0;
leakBytes = 0;

EnterCriticalSection(&state->TrackingLock);

HEAP_INFO("=== MEMORY LEAK REPORT ===");

current = state->MemoryAllocations.Flink;
Expand Down Expand Up @@ -468,6 +503,8 @@ void PrintMemoryLeaks(void) {
HEAP_INFO(" Double-free attempts: %d", (int)state->DoubleFreeCount);
HEAP_INFO(" Freed entries tracked: %d", (int)state->FreedEntryCount);
HEAP_INFO("===========================");

LeaveCriticalSection(&state->TrackingLock);
}

void PrintDoubleFreeReport(void) {
Expand All @@ -479,6 +516,8 @@ void PrintDoubleFreeReport(void) {
state = GetGlobalState();
if (!state) return;

EnterCriticalSection(&state->TrackingLock);

HEAP_INFO("=== FREED MEMORY REPORT ===");
HEAP_INFO("Total double-free attempts detected: %d", (int)state->DoubleFreeCount);
HEAP_INFO("Currently tracking %d freed allocations", (int)state->FreedEntryCount);
Expand Down Expand Up @@ -512,6 +551,8 @@ void PrintDoubleFreeReport(void) {
}

HEAP_INFO("==============================");

LeaveCriticalSection(&state->TrackingLock);
}

void SetErrorSuppression(BOOL suppress) {
Expand Down Expand Up @@ -547,11 +588,12 @@ BOOL GetFreedMemoryTracking(void) {
void SetMaxFreedEntries(SIZE_T maxEntries) {
GLOBAL_STATE* state = GetGlobalState();
if (state) {
EnterCriticalSection(&state->TrackingLock);
state->MaxFreedEntries = maxEntries;

if (state->FreedEntryCount > maxEntries) {
CleanupOldFreedEntries();
}
LeaveCriticalSection(&state->TrackingLock);
}
}

Expand Down
Loading