From 5db929f36933b2f309206cc8cfab1a4e5fc60071 Mon Sep 17 00:00:00 2001
From: Timo Kreuzer <timo.kreuzer@reactos.org>
Date: Mon, 18 Feb 2013 20:09:50 +0000
Subject: [PATCH] [FREELDR] Implement HeapVerify(), fix a buffer overrun.
 CORE-6893 #resolve

svn path=/trunk/; revision=58341
---
 reactos/boot/freeldr/freeldr/include/mm.h     |  4 +++
 .../boot/freeldr/freeldr/inifile/inifile.c    |  2 +-
 reactos/boot/freeldr/freeldr/mm/heap.c        | 34 +++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/reactos/boot/freeldr/freeldr/include/mm.h b/reactos/boot/freeldr/freeldr/include/mm.h
index 7b6ed61566a..bc7687a71ca 100644
--- a/reactos/boot/freeldr/freeldr/include/mm.h
+++ b/reactos/boot/freeldr/freeldr/include/mm.h
@@ -132,6 +132,10 @@ VOID
 HeapRelease(
     PVOID HeapHandle);
 
+VOID
+HeapVerify(
+    PVOID HeapHandle);
+
 VOID
 HeapCleanupAll(VOID);
 
diff --git a/reactos/boot/freeldr/freeldr/inifile/inifile.c b/reactos/boot/freeldr/freeldr/inifile/inifile.c
index 16a5b6e9ca2..7829bf38010 100644
--- a/reactos/boot/freeldr/freeldr/inifile/inifile.c
+++ b/reactos/boot/freeldr/freeldr/inifile/inifile.c
@@ -190,7 +190,7 @@ BOOLEAN IniAddSection(PCSTR SectionName, ULONG_PTR* SectionId)
 	RtlZeroMemory(Section, sizeof(INI_SECTION));
 
 	// Allocate the section name buffer
-	Section->SectionName = MmHeapAlloc(strlen(SectionName));
+	Section->SectionName = MmHeapAlloc(strlen(SectionName) + sizeof(CHAR));
 	if (!Section->SectionName)
 	{
 		MmHeapFree(Section);
diff --git a/reactos/boot/freeldr/freeldr/mm/heap.c b/reactos/boot/freeldr/freeldr/mm/heap.c
index 625754861f1..54c46d4b2e8 100644
--- a/reactos/boot/freeldr/freeldr/mm/heap.c
+++ b/reactos/boot/freeldr/freeldr/mm/heap.c
@@ -149,6 +149,32 @@ HeapDestroy(
                              LoaderFirmwareTemporary);
 }
 
+#ifdef FREELDR_HEAP_VERIFIER
+VOID
+HeapVerify(
+    PVOID HeapHandle)
+{
+    PHEAP Heap = HeapHandle;
+    PHEAP_BLOCK Block;
+
+    /* Loop all heap chunks */
+    for (Block = &Heap->Blocks;
+         Block->Size != 0;
+         Block = Block + 1 + Block->Size)
+    {
+        /* Continue, if its not free */
+        if (Block->Tag != 0)
+        {
+            /* Verify size and redzones */
+            ASSERT(*REDZONE_SIZE(Block) <= Block->Size * sizeof(HEAP_BLOCK));
+            ASSERT(*REDZONE_LOW(Block) == REDZONE_MARK);
+            ASSERT(*REDZONE_HI(Block) == REDZONE_MARK);
+            continue;
+        }
+    }
+}
+#endif /* FREELDR_HEAP_VERIFIER */
+
 VOID
 HeapRelease(
     PVOID HeapHandle)
@@ -296,6 +322,9 @@ HeapAllocate(
     ULONGLONG Time = __rdtsc();
 
 #ifdef FREELDR_HEAP_VERIFIER
+    /* Verify the heap */
+    HeapVerify(HeapHandle);
+
     /* Add space for a size field and 2 redzones */
     ByteSize += REDZONE_ALLOCATION;
 #endif
@@ -408,6 +437,11 @@ HeapFree(
     TRACE("HeapFree(%p, %p)\n", HeapHandle, Pointer);
     ASSERT(Tag != 'dnE#');
 
+#ifdef FREELDR_HEAP_VERIFIER
+    /* Verify the heap */
+    HeapVerify(HeapHandle);
+#endif
+
     /* Check if the block is really inside this heap */
     if ((Pointer < (PVOID)(Heap + 1)) ||
         (Pointer > (PVOID)((PUCHAR)Heap + Heap->MaximumSize)))