User:ThFabba/Debug Heap

From ReactOS Wiki
Jump to: navigation, search

Here's my debug heap. In the form of a kernel pool in this case, but easily adjustable to other allocation functions.

debugpool.h

#pragma once

void* heap_alloc(size_t requested_size, const char* file, int line);
void heap_free(void* mem);

#define ExAllocatePoolWithTag(type,size,tag) heap_alloc(size, __FILE__, __LINE__)
#define ExAllocatePool(type,size,tag) heap_alloc(size, __FILE__, __LINE__)

#define ExFreePoolWithTag(p,tag) heap_free(p)
#define ExFreePool(p) heap_free(p)

static inline void RtlFreeUnicodeString_(PUNICODE_STRING String)
{
    if (String->Buffer)
    {
        heap_free(String->Buffer);
        String->Length = 0;
        String->MaximumLength = 0;
    }
}
#define RtlFreeUnicodeString(s) RtlFreeUnicodeString_(s)

debugpool.c

#include <debugpool.h>
#include <ntddk.h>
#include <debug.h>

static inline long round_up(long value, long align) {
	return (value + align - 1) & ~(align - 1);
}

typedef unsigned char byte;

static volatile LONG heap_end = 0;
__attribute__((aligned(32))) static byte heap[10 * 1024 * 1024];

#define HEAD_SIGNATURE_SIZE ((int)(32 - 4 * sizeof(int) - sizeof(void*)))
#define HEAD_SIGNATURE_VALUE 0xaa

#define TAIL_SIGNATURE_SIZE 8
#define TAIL_SIGNATURE_VALUE 0xbb

#define FREE_SIGNATURE_VALUE 0xcc

struct heap_entry {
	const char* file;
	int line;
	int size;
	int requested_size;
	int valid : 16;
	int freed : 16;
	byte signature[HEAD_SIGNATURE_SIZE];
};

static int assert_heap_entry_size[sizeof(struct heap_entry) == 32 ? 1 : -1];
static int assert_head_signature_size[HEAD_SIGNATURE_SIZE == 12 ? 1 : -1];

static void fill_head_signature(byte* signature) {
	int i;

	for (i = 0; i < HEAD_SIGNATURE_SIZE; i++)
		signature[i] = HEAD_SIGNATURE_VALUE;
}

static void check_head_signature(const byte* signature) {
	int i;

	for (i = 0; i < HEAD_SIGNATURE_SIZE; i++)
		if (signature[i] != HEAD_SIGNATURE_VALUE) {
			DPRINT1("signature[i]=%x\n", signature[i]);
			ASSERT(FALSE);
		}
}

static void fill_tail_signature(byte* signature) {
	int i;

	for (i = 0; i < TAIL_SIGNATURE_SIZE; i++)
		signature[i] = TAIL_SIGNATURE_VALUE;
}

static void check_tail_signature(const byte* signature) {
	int i;

	for (i = 0; i < TAIL_SIGNATURE_SIZE; i++)
		if (signature[i] != TAIL_SIGNATURE_VALUE) {
			DPRINT1("signature[i]=%x\n", signature[i]);
			ASSERT(FALSE);
		}
}

static void check_free_signature(const byte* signature, int size) {
	int i;

	for (i = 0; i < size; i++)
		if (signature[i] != FREE_SIGNATURE_VALUE) {
			DPRINT1("signature[i]=%x\n", signature[i]);
			ASSERT(FALSE);
		}
}

static void heap_validate(void) {
	struct heap_entry* previous_entry = NULL;
	struct heap_entry* entry;
	long end;
	byte* signature;
	byte* mem;

	/* has to be _somewhere_ ;p */
	(void)assert_heap_entry_size;
	(void)assert_head_signature_size;

	end = heap_end;
	entry = (struct heap_entry*)heap;
	while ((byte*)entry < heap + end) {
		if (entry->valid) {
			check_head_signature(entry->signature);
			mem = (void*)(entry + 1);
			signature = mem + entry->requested_size;
			check_tail_signature(signature);
			if (entry->freed)
				check_free_signature(mem, entry->requested_size);
		}
		previous_entry = entry;
		entry = (void*)((byte*)entry + entry->size);
	}
	if ((byte*)entry != heap + end) {
		DPRINT1("entry=%p, heap+end=%p\n", entry, heap + end);
		ASSERT(FALSE);
	}
}

void* heap_alloc(size_t requested_size, const char* file, int line) {
	void* mem;
	long start;
	long size;
	struct heap_entry* entry;
	byte* signature;

	heap_validate();

	/* this catches negative values in addition to too large ones */
	if (requested_size >= sizeof(heap)) {
		DPRINT1("requested_size=%lu\n", (ULONG)requested_size);
		ASSERT(FALSE);
	}

	size = requested_size;
	size += sizeof(struct heap_entry);
	size += TAIL_SIGNATURE_SIZE;
	size = round_up(size, 16);

	start = InterlockedExchangeAdd(&heap_end, size);
	if (start >= (int)sizeof(heap)) {
		DPRINT1("start=%d\n", start);
		ASSERT(FALSE);
	}

	entry = (void*)&heap[start];
	entry->file = file;
	entry->line = line;
	entry->size = size;
	entry->requested_size = requested_size;
	entry->freed = 0;
	fill_head_signature(entry->signature);

	mem = (void*)(entry + 1);

	signature = mem;
	signature += requested_size;
	fill_tail_signature(signature);

	if ((long)mem != round_up((long)mem, 16)) {
		DPRINT1("mem=%p\n", mem);
		ASSERT(FALSE);
	}

	entry->valid = 1;
	heap_validate();
	return mem;
}

void heap_free(void* mem) {
	struct heap_entry* entry;
	byte* signature;

	heap_validate();

	if ((long)mem != round_up((long)mem, 16)) {
		DPRINT1("mem=%p\n", mem);
		ASSERT(FALSE);
	}
	if ((byte*)mem <= heap) {
		DPRINT1("mem=%p, heap=%p\n", mem, heap);
		ASSERT(FALSE);
	}
	if ((byte*)mem >= heap + heap_end) {
		DPRINT1("mem=%p, heap-end=%p\n", mem, heap + heap_end);
		ASSERT(FALSE);
	}

	entry = (struct heap_entry*)mem - 1;
	if ((byte *)entry + entry->size > heap + heap_end) {
		DPRINT1("entry=%p, size=%lu, heap-end=%p\n", entry, (ULONG)entry->size, heap + heap_end);
		ASSERT(FALSE);
	}
	if (entry->freed != 0) {
		DPRINT1("entry=%p, freed=%d\n", entry, entry->freed);
		ASSERT(FALSE);
	}

	check_head_signature(entry->signature);

	signature = mem;
	signature += entry->requested_size;
	check_tail_signature(signature);

	memset(mem, FREE_SIGNATURE_VALUE, entry->requested_size);
	entry->freed = 1;

	heap_validate();
}


Win32k's heaps are particularly hard to debug since you cannot make too many changes to them (need to work in both kernel and user mode etc). The following is an example of a patch that tries to catch uses after free and buffer underruns/overruns in the user heap:

Index: usrheap.c
===================================================================
--- usrheap.c	(revision 74970)
+++ usrheap.c	(working copy)
@@ -334,4 +334,148 @@
     return STATUS_SUCCESS;
 }
 
+#define DBG_SIG1 0xaefea6c1
+#define DBG_SIG2 0x54e9861e
+typedef struct _DBG_HEADER
+{
+    ULONG Sig1;
+    PCSTR File;
+    INT Line;
+    ULONG Sig2;
+} DBG_HEADER, *PDBG_HEADER;
+
+static PDBG_HEADER DbgHolding[1024];
+static ULONG DbgHoldingIndex;
+static PDBG_HEADER MinAlloc = (PVOID)(LONG_PTR)-1, MaxAlloc;
+
+PVOID
+UserHeapAlloc_(SIZE_T Bytes,
+               PCSTR File,
+               INT Line)
+{
+    PDBG_HEADER Header;
+
+    Bytes = ALIGN_UP_BY(Bytes + sizeof(DBG_HEADER) + 4 * sizeof(ULONG), 2 * sizeof(PVOID));
+
+    Header = RtlAllocateHeap(GlobalUserHeap,
+                             HEAP_NO_SERIALIZE,
+                             Bytes);
+    if (!Header) return NULL;
+    Header->Sig1 = DBG_SIG1;
+    Header->File = File;
+    Header->Line = Line;
+    Header->Sig2 = DBG_SIG2;
+    if (Header < MinAlloc)
+        MinAlloc = Header;
+    if (Header > MaxAlloc)
+        MaxAlloc = Header;
+    RtlFillMemory(Header + 1, Bytes - sizeof(DBG_HEADER) - 4 * sizeof(ULONG), 0xe5);
+    RtlFillMemory((PUCHAR)(Header + 1) + Bytes - sizeof(DBG_HEADER) - 4 * sizeof(ULONG), 4 * sizeof(ULONG), 0xdc);
+    return Header + 1;
+}
+
+BOOL
+UserHeapFree(PVOID lpMem)
+{
+    PDBG_HEADER Header;
+    ULONG i, Size;
+    PULONG Data;
+    ULONG Sig;
+
+    Header = lpMem;
+    Header--;
+    if (Header < MinAlloc || Header > MaxAlloc ||
+        (ULONG_PTR)Header % (2 * sizeof(PVOID)) != 0)
+    {
+        DPRINT1("Invalid free. lpMem=%p, Min=%p, Max=%p\n",
+                lpMem, MinAlloc, MaxAlloc);
+        __debugbreak();
+        return FALSE;
+    }
+    if (DbgHolding[DbgHoldingIndex])
+    {
+        if (DbgHolding[DbgHoldingIndex]->Sig1 != DBG_SIG1 ||
+            DbgHolding[DbgHoldingIndex]->Sig2 != DBG_SIG2)
+        {
+            DPRINT1("Header sig mismatch for %p\n", DbgHolding[DbgHoldingIndex]);
+            DPRINT1("Allocation is from %s:%d\n", DbgHolding[DbgHoldingIndex]->File, DbgHolding[DbgHoldingIndex]->Line);
+            DPRINT1("Sigs are 0x%08x, 0x%08x\n", DbgHolding[DbgHoldingIndex]->Sig1, DbgHolding[DbgHoldingIndex]->Sig2);
+            __debugbreak();
+        }
+        Data = (PULONG)(DbgHolding[DbgHoldingIndex] + 1);
+        Size = (RtlSizeHeap(GlobalUserHeap, HEAP_NO_SERIALIZE, DbgHolding[DbgHoldingIndex]) - sizeof(DBG_HEADER)) / sizeof(ULONG);
+        Sig = (ULONG_PTR)DbgHolding[DbgHoldingIndex]->File;
+        Sig ^= DbgHolding[DbgHoldingIndex]->Line;
+        Sig ^= DBG_SIG2;
+        for (i = 0; i < Size; i++)
+        {
+            if (Data[i] != Sig + i)
+            {
+                DPRINT1("Content sig mismatch for %p\n", DbgHolding[DbgHoldingIndex]);
+                DPRINT1("Allocation is from %s:%d\n", DbgHolding[DbgHoldingIndex]->File, DbgHolding[DbgHoldingIndex]->Line);
+                DPRINT1("Sig at offset %lu is 0x%08x, expected 0x%08x\n", i, Data[i], Sig + i);
+                __debugbreak();
+            }
+        }
+        RtlFreeHeap(GlobalUserHeap,
+                    HEAP_NO_SERIALIZE,
+                    DbgHolding[DbgHoldingIndex]);
+    }
+    Data = (PULONG)(Header + 1);
+    Size = (RtlSizeHeap(GlobalUserHeap, HEAP_NO_SERIALIZE, Header) - sizeof(DBG_HEADER)) / sizeof(ULONG);
+    Sig = (ULONG_PTR)Header->File;
+    Sig ^= Header->Line;
+    Sig ^= DBG_SIG2;
+    for (i = Size - 4; i < Size; i++)
+    {
+        if (Data[i] != 0xdcdcdcdc)
+        {
+            DPRINT1("Post buffer sig mismatch for %p\n", Header);
+            DPRINT1("Allocation is from %s:%d\n", Header->File, Header->Line);
+            DPRINT1("Sig at offset %lu is 0x%08x\n", i, Data[i]);
+            __debugbreak();
+        }
+    }
+    for (i = 0; i < Size; i++)
+    {
+        Data[i] = Sig + i;
+    }
+    DbgHolding[DbgHoldingIndex] = Header;
+    DbgHoldingIndex = (DbgHoldingIndex + 1) % _countof(DbgHolding);
+    return TRUE;
+}
+
+PVOID
+UserHeapReAlloc_(PVOID lpMem,
+                 SIZE_T BytesAsked,
+                 PCSTR File,
+                 INT Line)
+{
+    PDBG_HEADER Header;
+    SIZE_T PrevSize;
+    PVOID pNew;
+
+    Header = lpMem;
+    Header--;
+
+    PrevSize = RtlSizeHeap(GlobalUserHeap,
+                           HEAP_NO_SERIALIZE,
+                           lpMem);
+
+    pNew = UserHeapAlloc_(BytesAsked, File, Line);
+    if (pNew != NULL)
+    {
+        if (PrevSize < BytesAsked)
+            BytesAsked = PrevSize;
+
+        RtlCopyMemory(pNew,
+                      lpMem,
+                      BytesAsked);
+
+        UserHeapFree(lpMem);
+    }
+
+    return pNew;
+}
+
 /* EOF */
Index: usrheap.h
===================================================================
--- usrheap.h	(revision 74970)
+++ usrheap.h	(working copy)
@@ -30,6 +30,23 @@
                   OUT PVOID* KernelMapping,
                   OUT PVOID* UserMapping);
 
+#if 1
+#define UserHeapAlloc(b) UserHeapAlloc_(b, __FILE__, __LINE__)
+PVOID
+UserHeapAlloc_(SIZE_T Bytes,
+               PCSTR File,
+               INT Line);
+
+BOOL
+UserHeapFree(PVOID lpMem);
+
+#define UserHeapReAlloc(m, b) UserHeapReAlloc(m, b, __FILE__, __LINE__)
+PVOID
+UserHeapReAlloc_(PVOID lpMem,
+                 SIZE_T Bytes,
+                 PCSTR File,
+                 INT Line);
+#else
 static __inline PVOID
 UserHeapAlloc(SIZE_T Bytes)
 {
@@ -87,6 +104,7 @@
     return pNew;
 #endif
 }
+#endif
 
 static __inline PVOID
 UserHeapAddressToUser(PVOID lpMem)