[SHELL32][BOOTDATA] Implement Command Prompt here (#2029)
[reactos.git] / boot / freeldr / freeldr / lib / mm / heap.c
1 /*
2 * FreeLoader
3 * Copyright (C) 2011 Timo Kreuzer (timo.kreuzer@reactos.org)
4 *
5 * This program is free software; you can redistribute it and/or modify
6 * it under the terms of the GNU General Public License as published by
7 * the Free Software Foundation; either version 2 of the License, or
8 * (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13 * GNU General Public License for more details.
14 *
15 * You should have received a copy of the GNU General Public License along
16 * with this program; if not, write to the Free Software Foundation, Inc.,
17 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
18 */
19
20 #include <freeldr.h>
21
22 #include <debug.h>
23 DBG_DEFAULT_CHANNEL(HEAP);
24
25 #define FREELDR_HEAP_VERIFIER
26
27 #define REDZONE_MARK 0xCCCCCCCCCCCCCCCCULL
28 #define REDZONE_ALLOCATION 24
29 #define REDZONE_LOW_OFFSET 16
30 #define REDZONE_SIZE(Block) ((ULONG64*)Block->Data)
31 #define REDZONE_LOW(Block) ((ULONG64*)Block->Data + 1)
32 #define REDZONE_HI(Block) ((ULONG64*)((PUCHAR)Block->Data + 16 + *REDZONE_SIZE(Block)))
33
34 PVOID FrLdrDefaultHeap;
35 PVOID FrLdrTempHeap;
36
37 typedef struct _BLOCK_DATA
38 {
39 ULONG_PTR Flink:32;
40 ULONG_PTR Blink:32;
41 } BLOCK_DATA, *PBLOCK_DATA;
42
43 typedef struct _HEAP_BLOCK
44 {
45 USHORT Size;
46 USHORT PreviousSize;
47 ULONG Tag;
48 BLOCK_DATA Data[];
49 } HEAP_BLOCK, *PHEAP_BLOCK;
50
51 typedef struct _HEAP
52 {
53 SIZE_T MaximumSize;
54 SIZE_T CurrentAllocBytes;
55 SIZE_T MaxAllocBytes;
56 ULONG NumAllocs;
57 ULONG NumFrees;
58 SIZE_T LargestAllocation;
59 ULONGLONG AllocationTime;
60 ULONGLONG FreeTime;
61 ULONG_PTR TerminatingBlock;
62 HEAP_BLOCK Blocks;
63 } HEAP, *PHEAP;
64
65 PVOID
66 FrLdrHeapCreate(
67 SIZE_T MaximumSize,
68 TYPE_OF_MEMORY MemoryType)
69 {
70 PHEAP Heap;
71 PHEAP_BLOCK Block;
72 SIZE_T Remaining;
73 USHORT PreviousSize;
74 TRACE("HeapCreate(MemoryType=%ld)\n", MemoryType);
75
76 /* Allocate some memory for the heap */
77 MaximumSize = ALIGN_UP_BY(MaximumSize, MM_PAGE_SIZE);
78 Heap = MmAllocateMemoryWithType(MaximumSize, MemoryType);
79 if (!Heap)
80 {
81 ERR("HEAP: Failed to allocate heap of size 0x%lx, Type\n",
82 MaximumSize, MemoryType);
83 return NULL;
84 }
85
86 /* Initialize the heap header */
87 Heap->MaximumSize = MaximumSize;
88 Heap->CurrentAllocBytes = 0;
89 Heap->MaxAllocBytes = 0;
90 Heap->NumAllocs = 0;
91 Heap->NumFrees = 0;
92 Heap->LargestAllocation = 0;
93
94 /* Calculate what's left to process */
95 Remaining = (MaximumSize - sizeof(HEAP)) / sizeof(HEAP_BLOCK);
96 TRACE("Remaining = %ld\n", Remaining);
97
98 /* Substract 2 for the terminating entry (header + free entry) */
99 Remaining -= 2;
100
101 Block = &Heap->Blocks;
102 PreviousSize = 0;
103
104 /* Create free blocks */
105 while (Remaining > 1)
106 {
107 /* Initialize this free block */
108 Block->Size = (USHORT)min(MAXUSHORT, Remaining - 1);
109 Block->PreviousSize = PreviousSize;
110 Block->Tag = 0;
111 Block->Data[0].Flink = (Block - &Heap->Blocks) + Block->Size + 1;
112 Block->Data[0].Blink = (Block - &Heap->Blocks) - 1 - PreviousSize;
113
114 /* Substract current block size from remainder */
115 Remaining -= (Block->Size + 1);
116
117 /* Go to next block */
118 PreviousSize = Block->Size;
119 Block = Block + Block->Size + 1;
120
121 TRACE("Remaining = %ld\n", Remaining);
122 }
123
124 /* Now finish with a terminating block */
125 Heap->TerminatingBlock = Block - &Heap->Blocks;
126 Block->Size = 0;
127 Block->PreviousSize = PreviousSize;
128 Block->Tag = 'dnE#';
129 Block->Data[0].Flink = 0;
130 Block->Data[0].Blink = (Block - &Heap->Blocks) - 1 - PreviousSize;
131 Heap->Blocks.Data[0].Blink = Heap->TerminatingBlock;
132
133 return Heap;
134 }
135
136 VOID
137 FrLdrHeapDestroy(
138 PVOID HeapHandle)
139 {
140 PHEAP Heap = HeapHandle;
141
142 /* Mark all pages as firmware temporary, so they are free for the kernel */
143 MmMarkPagesInLookupTable(PageLookupTableAddress,
144 (ULONG_PTR)Heap / MM_PAGE_SIZE,
145 (PFN_COUNT)(Heap->MaximumSize / MM_PAGE_SIZE),
146 LoaderFirmwareTemporary);
147
148 #if DBG
149 /* Make sure everything is dead */
150 RtlFillMemory(Heap, Heap->MaximumSize, 0xCCCCCCCC);
151 #endif
152 }
153
154 #ifdef FREELDR_HEAP_VERIFIER
155 VOID
156 FrLdrHeapVerify(
157 PVOID HeapHandle)
158 {
159 PHEAP Heap = HeapHandle;
160 PHEAP_BLOCK Block;
161
162 /* Loop all heap chunks */
163 for (Block = &Heap->Blocks;
164 Block->Size != 0;
165 Block = Block + 1 + Block->Size)
166 {
167 /* Continue, if its not free */
168 if (Block->Tag != 0)
169 {
170 /* Verify size and redzones */
171 ASSERT(*REDZONE_SIZE(Block) <= Block->Size * sizeof(HEAP_BLOCK));
172 ASSERT(*REDZONE_LOW(Block) == REDZONE_MARK);
173 ASSERT(*REDZONE_HI(Block) == REDZONE_MARK);
174 continue;
175 }
176 }
177 }
178 #endif /* FREELDR_HEAP_VERIFIER */
179
180 VOID
181 FrLdrHeapRelease(
182 PVOID HeapHandle)
183 {
184 PHEAP Heap = HeapHandle;
185 PHEAP_BLOCK Block;
186 PUCHAR StartAddress, EndAddress;
187 PFN_COUNT FreePages, AllFreePages = 0;
188
189 TRACE("HeapRelease(%p)\n", HeapHandle);
190
191 /* Loop all heap chunks */
192 for (Block = &Heap->Blocks;
193 Block->Size != 0;
194 Block = Block + 1 + Block->Size)
195 {
196 /* Continue, if its not free */
197 if (Block->Tag != 0)
198 {
199 #ifdef FREELDR_HEAP_VERIFIER
200 /* Verify size and redzones */
201 ASSERT(*REDZONE_SIZE(Block) <= Block->Size * sizeof(HEAP_BLOCK));
202 ASSERT(*REDZONE_LOW(Block) == REDZONE_MARK);
203 ASSERT(*REDZONE_HI(Block) == REDZONE_MARK);
204 #endif
205 continue;
206 }
207
208 /* Calculate page aligned start address of the free region */
209 StartAddress = ALIGN_UP_POINTER_BY(Block->Data, PAGE_SIZE);
210
211 /* Walk over adjacent free blocks */
212 while (Block->Tag == 0) Block = Block + Block->Size + 1;
213
214 /* Check if this was the last block */
215 if (Block->Size == 0)
216 {
217 /* Align the end address up to cover the end of the heap */
218 EndAddress = ALIGN_UP_POINTER_BY(Block->Data, PAGE_SIZE);
219 }
220 else
221 {
222 /* Align the end address down to not cover any allocations */
223 EndAddress = ALIGN_DOWN_POINTER_BY(Block->Data, PAGE_SIZE);
224 }
225
226 /* Check if we have free pages */
227 if (EndAddress > StartAddress)
228 {
229 /* Calculate the size of the free region in pages */
230 FreePages = (PFN_COUNT)((EndAddress - StartAddress) / MM_PAGE_SIZE);
231 AllFreePages += FreePages;
232
233 /* Now mark the pages free */
234 MmMarkPagesInLookupTable(PageLookupTableAddress,
235 (ULONG_PTR)StartAddress / MM_PAGE_SIZE,
236 FreePages,
237 LoaderFree);
238 }
239
240 /* bail out, if it was the last block */
241 if (Block->Size == 0) break;
242 }
243
244 TRACE("HeapRelease() done, freed %lu of %lu pages\n", AllFreePages, Heap->MaximumSize / MM_PAGE_SIZE);
245 }
246
247 VOID
248 FrLdrHeapCleanupAll(VOID)
249 {
250 #if DBG
251 PHEAP Heap;
252
253 Heap = FrLdrDefaultHeap;
254 TRACE("Heap statistics for default heap:\n"
255 "CurrentAlloc=0x%lx, MaxAlloc=0x%lx, LargestAllocation=0x%lx\n"
256 "NumAllocs=%ld, NumFrees=%ld\n",
257 Heap->CurrentAllocBytes, Heap->MaxAllocBytes, Heap->LargestAllocation,
258 Heap->NumAllocs, Heap->NumFrees);
259 TRACE("AllocTime = %I64d, FreeTime = %I64d, sum = %I64d\n",
260 Heap->AllocationTime, Heap->FreeTime, Heap->AllocationTime + Heap->FreeTime);
261 #endif
262
263 /* Release free pages from the default heap */
264 FrLdrHeapRelease(FrLdrDefaultHeap);
265
266 #if DBG
267 Heap = FrLdrTempHeap;
268 TRACE("Heap statistics for temp heap:\n"
269 "CurrentAlloc=0x%lx, MaxAlloc=0x%lx, LargestAllocation=0x%lx\n"
270 "NumAllocs=%ld, NumFrees=%ld\n",
271 Heap->CurrentAllocBytes, Heap->MaxAllocBytes, Heap->LargestAllocation,
272 Heap->NumAllocs, Heap->NumFrees);
273 #endif
274
275 /* Destroy the temp heap */
276 FrLdrHeapDestroy(FrLdrTempHeap);
277 }
278
279 static VOID
280 FrLdrHeapRemoveFreeList(
281 PHEAP Heap,
282 PHEAP_BLOCK Block)
283 {
284 PHEAP_BLOCK Previous, Next;
285
286 Next = &Heap->Blocks + Block->Data[0].Flink;
287 Previous = &Heap->Blocks + Block->Data[0].Blink;
288 ASSERT((Next->Tag == 0) || (Next->Tag == 'dnE#'));
289 ASSERT(Next->Data[0].Blink == Block - &Heap->Blocks);
290 ASSERT((Previous->Tag == 0) || (Previous->Tag == 'dnE#'));
291 ASSERT(Previous->Data[0].Flink == Block - &Heap->Blocks);
292
293 Next->Data[0].Blink = Previous - &Heap->Blocks;
294 Previous->Data[0].Flink = Next - &Heap->Blocks;
295 }
296
297 static VOID
298 FrLdrHeapInsertFreeList(
299 PHEAP Heap,
300 PHEAP_BLOCK FreeBlock)
301 {
302 PHEAP_BLOCK ListHead, NextBlock;
303 ASSERT(FreeBlock->Tag == 0);
304
305 /* Terminating block serves as free list head */
306 ListHead = &Heap->Blocks + Heap->TerminatingBlock;
307
308 for (NextBlock = &Heap->Blocks + ListHead->Data[0].Flink;
309 NextBlock < FreeBlock;
310 NextBlock = &Heap->Blocks + NextBlock->Data[0].Flink);
311
312 FreeBlock->Data[0].Flink = NextBlock - &Heap->Blocks;
313 FreeBlock->Data[0].Blink = NextBlock->Data[0].Blink;
314 NextBlock->Data[0].Blink = FreeBlock - &Heap->Blocks;
315 NextBlock = &Heap->Blocks + FreeBlock->Data[0].Blink;
316 NextBlock->Data[0].Flink = FreeBlock - &Heap->Blocks;
317 }
318
319 PVOID
320 FrLdrHeapAllocateEx(
321 PVOID HeapHandle,
322 SIZE_T ByteSize,
323 ULONG Tag)
324 {
325 PHEAP Heap = HeapHandle;
326 PHEAP_BLOCK Block, NextBlock;
327 USHORT BlockSize, Remaining;
328 #if DBG && !defined(_M_ARM)
329 ULONGLONG Time = __rdtsc();
330 #endif
331
332 #ifdef FREELDR_HEAP_VERIFIER
333 /* Verify the heap */
334 FrLdrHeapVerify(HeapHandle);
335
336 /* Add space for a size field and 2 redzones */
337 ByteSize += REDZONE_ALLOCATION;
338 #endif
339
340 /* Check if the allocation is too large */
341 if ((ByteSize + sizeof(HEAP_BLOCK)) > MAXUSHORT * sizeof(HEAP_BLOCK))
342 {
343 ERR("HEAP: Allocation of 0x%lx bytes too large\n", ByteSize);
344 return NULL;
345 }
346
347 /* We need a proper tag */
348 if (Tag == 0) Tag = 'enoN';
349
350 /* Calculate alloc size */
351 BlockSize = (USHORT)((ByteSize + sizeof(HEAP_BLOCK) - 1) / sizeof(HEAP_BLOCK));
352
353 /* Walk the free block list */
354 Block = &Heap->Blocks + Heap->TerminatingBlock;
355 for (Block = &Heap->Blocks + Block->Data[0].Flink;
356 Block->Size != 0;
357 Block = &Heap->Blocks + Block->Data[0].Flink)
358 {
359 ASSERT(Block->Tag == 0);
360
361 /* Continue, if its too small */
362 if (Block->Size < BlockSize) continue;
363
364 /* This block is just fine, use it */
365 Block->Tag = Tag;
366
367 /* Remove this entry from the free list */
368 FrLdrHeapRemoveFreeList(Heap, Block);
369
370 /* Calculate the remaining size */
371 Remaining = Block->Size - BlockSize;
372
373 /* Check if the remaining space is large enough for a new block */
374 if (Remaining > 1)
375 {
376 /* Make the allocated block as large as necessary */
377 Block->Size = BlockSize;
378
379 /* Get pointer to the new block */
380 NextBlock = Block + 1 + BlockSize;
381
382 /* Make it a free block */
383 NextBlock->Tag = 0;
384 NextBlock->Size = Remaining - 1;
385 NextBlock->PreviousSize = BlockSize;
386 BlockSize = NextBlock->Size;
387 FrLdrHeapInsertFreeList(Heap, NextBlock);
388
389 /* Advance to the next block */
390 NextBlock = NextBlock + 1 + BlockSize;
391 }
392 else
393 {
394 /* Not enough left, use the full block */
395 BlockSize = Block->Size;
396
397 /* Get the next block */
398 NextBlock = Block + 1 + BlockSize;
399 }
400
401 /* Update the next blocks back link */
402 NextBlock->PreviousSize = BlockSize;
403
404 /* Update heap usage */
405 Heap->NumAllocs++;
406 Heap->CurrentAllocBytes += Block->Size * sizeof(HEAP_BLOCK);
407 Heap->MaxAllocBytes = max(Heap->MaxAllocBytes, Heap->CurrentAllocBytes);
408 Heap->LargestAllocation = max(Heap->LargestAllocation,
409 Block->Size * sizeof(HEAP_BLOCK));
410 #if DBG && !defined(_M_ARM)
411 Heap->AllocationTime += (__rdtsc() - Time);
412 #endif
413 TRACE("HeapAllocate(%p, %ld, %.4s) -> return %p\n",
414 HeapHandle, ByteSize, &Tag, Block->Data);
415
416 /* HACK: zero out the allocation */
417 RtlZeroMemory(Block->Data, Block->Size * sizeof(HEAP_BLOCK));
418
419 #ifdef FREELDR_HEAP_VERIFIER
420 /* Write size and redzones */
421 *REDZONE_SIZE(Block) = ByteSize - REDZONE_ALLOCATION;
422 *REDZONE_LOW(Block) = REDZONE_MARK;
423 *REDZONE_HI(Block) = REDZONE_MARK;
424
425 /* Allocation starts after size field and redzone */
426 return (PUCHAR)Block->Data + REDZONE_LOW_OFFSET;
427 #endif
428 /* Return pointer to the data */
429 return Block->Data;
430 }
431
432 /* We found nothing */
433 WARN("HEAP: nothing suitable found for 0x%lx bytes\n", ByteSize);
434 return NULL;
435 }
436
437 VOID
438 FrLdrHeapFreeEx(
439 PVOID HeapHandle,
440 PVOID Pointer,
441 ULONG Tag)
442 {
443 PHEAP Heap = HeapHandle;
444 PHEAP_BLOCK Block, PrevBlock, NextBlock;
445 #if DBG && !defined(_M_ARM)
446 ULONGLONG Time = __rdtsc();
447 #endif
448 TRACE("HeapFree(%p, %p)\n", HeapHandle, Pointer);
449 ASSERT(Tag != 'dnE#');
450
451 #ifdef FREELDR_HEAP_VERIFIER
452 /* Verify the heap */
453 FrLdrHeapVerify(HeapHandle);
454 #endif
455
456 /* Check if the block is really inside this heap */
457 if ((Pointer < (PVOID)(Heap + 1)) ||
458 (Pointer > (PVOID)((PUCHAR)Heap + Heap->MaximumSize)))
459 {
460 ERR("HEAP: trying to free %p outside of heap %p\n", Pointer, Heap);
461 ASSERT(FALSE);
462 }
463
464 Block = ((PHEAP_BLOCK)Pointer) - 1;
465 #ifdef FREELDR_HEAP_VERIFIER
466 Block = (PHEAP_BLOCK)((PUCHAR)Block - REDZONE_LOW_OFFSET);
467
468 /* Verify size and redzones */
469 ASSERT(*REDZONE_SIZE(Block) <= Block->Size * sizeof(HEAP_BLOCK));
470 ASSERT(*REDZONE_LOW(Block) == REDZONE_MARK);
471 ASSERT(*REDZONE_HI(Block) == REDZONE_MARK);
472 #endif
473
474 /* Check if the tag matches */
475 if ((Tag && (Block->Tag != Tag)) || (Block->Tag == 0))
476 {
477 ERR("HEAP: Bad tag! Pointer=%p: block tag '%.4s', requested '%.4s', size=0x%lx\n",
478 Pointer, &Block->Tag, &Tag, Block->Size);
479 ASSERT(FALSE);
480 }
481
482 /* Mark as free */
483 Block->Tag = 0;
484
485 #if DBG
486 /* Erase contents */
487 RtlFillMemory(Block->Data, Block->Size * sizeof(HEAP_BLOCK), 0xCCCCCCCC);
488 #endif
489
490 /* Update heap usage */
491 Heap->NumFrees++;
492 Heap->CurrentAllocBytes -= Block->Size * sizeof(HEAP_BLOCK);
493
494 /* Get pointers to the next and previous block */
495 PrevBlock = Block - Block->PreviousSize - 1;
496 NextBlock = Block + Block->Size + 1;
497
498 /* Check if next block is free */
499 if ((NextBlock->Tag == 0) &&
500 ((Block->Size + NextBlock->Size + 1) <= MAXUSHORT))
501 {
502 /* Merge next block into current */
503 Block->Size += NextBlock->Size + 1;
504 FrLdrHeapRemoveFreeList(Heap, NextBlock);
505
506 NextBlock = Block + Block->Size + 1;
507 }
508
509 /* Check if there is a block before and it's free */
510 if ((Block->PreviousSize != 0) && (PrevBlock->Tag == 0) &&
511 ((PrevBlock->Size + Block->Size + 1) <= MAXUSHORT))
512 {
513 /* Merge current block into previous */
514 PrevBlock->Size += Block->Size + 1;
515 Block = PrevBlock;
516 }
517 else
518 {
519 /* Insert the entry into the free list */
520 FrLdrHeapInsertFreeList(Heap, Block);
521 }
522
523 /* Update the next block's back link */
524 NextBlock->PreviousSize = Block->Size;
525 #if DBG && !defined(_M_ARM)
526 Heap->FreeTime += (__rdtsc() - Time);
527 #endif
528 }
529
530
531 /* Wrapper functions *********************************************************/
532
533 VOID
534 MmInitializeHeap(PVOID PageLookupTable)
535 {
536 TRACE("MmInitializeHeap()\n");
537
538 /* Create the default heap */
539 FrLdrDefaultHeap = FrLdrHeapCreate(DEFAULT_HEAP_SIZE, LoaderOsloaderHeap);
540 ASSERT(FrLdrDefaultHeap);
541
542 /* Create a temporary heap */
543 FrLdrTempHeap = FrLdrHeapCreate(TEMP_HEAP_SIZE, LoaderFirmwareTemporary);
544 ASSERT(FrLdrTempHeap);
545
546 TRACE("MmInitializeHeap() done, default heap %p, temp heap %p\n",
547 FrLdrDefaultHeap, FrLdrTempHeap);
548 }
549
550 PVOID
551 NTAPI
552 ExAllocatePoolWithTag(
553 IN POOL_TYPE PoolType,
554 IN SIZE_T NumberOfBytes,
555 IN ULONG Tag)
556 {
557 return FrLdrHeapAllocateEx(FrLdrDefaultHeap, NumberOfBytes, Tag);
558 }
559
560 PVOID
561 NTAPI
562 ExAllocatePool(
563 IN POOL_TYPE PoolType,
564 IN SIZE_T NumberOfBytes)
565 {
566 return FrLdrHeapAllocateEx(FrLdrDefaultHeap, NumberOfBytes, 0);
567 }
568
569 VOID
570 NTAPI
571 ExFreePool(
572 IN PVOID P)
573 {
574 FrLdrHeapFreeEx(FrLdrDefaultHeap, P, 0);
575 }
576
577 VOID
578 NTAPI
579 ExFreePoolWithTag(
580 IN PVOID P,
581 IN ULONG Tag)
582 {
583 FrLdrHeapFreeEx(FrLdrDefaultHeap, P, Tag);
584 }
585
586 PVOID
587 NTAPI
588 RtlAllocateHeap(
589 IN PVOID HeapHandle,
590 IN ULONG Flags,
591 IN SIZE_T Size)
592 {
593 PVOID ptr;
594
595 ptr = FrLdrHeapAllocateEx(FrLdrDefaultHeap, Size, ' ltR');
596 if (ptr && (Flags & HEAP_ZERO_MEMORY))
597 {
598 RtlZeroMemory(ptr, Size);
599 }
600
601 return ptr;
602 }
603
604 BOOLEAN
605 NTAPI
606 RtlFreeHeap(
607 IN PVOID HeapHandle,
608 IN ULONG Flags,
609 IN PVOID HeapBase)
610 {
611 FrLdrHeapFreeEx(FrLdrDefaultHeap, HeapBase, ' ltR');
612 return TRUE;
613 }
614