xref: /reactos/boot/environ/lib/mm/heapalloc.c (revision 3be40816)
1 /*
2  * COPYRIGHT:       See COPYING.ARM in the top level directory
3  * PROJECT:         ReactOS UEFI Boot Library
4  * FILE:            boot/environ/lib/mm/heapalloc.c
5  * PURPOSE:         Boot Library Memory Manager Heap Allocator
6  * PROGRAMMER:      Alex Ionescu (alex.ionescu@reactos.org)
7  */
8 
9 /* INCLUDES ******************************************************************/
10 
11 #include "bl.h"
12 
13 /* DATA VARIABLES ************************************************************/
14 
15 #define BL_HEAP_POINTER_FLAG_BITS   3
16 
17 typedef struct _BL_HEAP_POINTER
18 {
19     union
20     {
21         struct
22         {
23             ULONG_PTR BufferFree : 1;
24             ULONG_PTR BufferOnHeap : 1;
25             ULONG_PTR NotUsed : 1;
26             ULONG_PTR BufferPointer : ((8 * sizeof(ULONG_PTR)) - BL_HEAP_POINTER_FLAG_BITS);
27         };
28         PVOID P;
29     };
30 } BL_HEAP_POINTER, *PBL_HEAP_POINTER;
31 
32 typedef struct _BL_FREE_HEAP_ENTRY
33 {
34     BL_HEAP_POINTER BufferNext;
35     BL_HEAP_POINTER BufferPrevious;
36     BL_HEAP_POINTER FreeNext;
37     BL_HEAP_POINTER FreePrevious;
38 } BL_FREE_HEAP_ENTRY, *PBL_FREE_HEAP_ENTRY;
39 
40 typedef struct _BL_BUSY_HEAP_ENTRY
41 {
42     BL_HEAP_POINTER BufferNext;
43     BL_HEAP_POINTER BufferPrevious;
44     UCHAR Buffer[ANYSIZE_ARRAY];
45 } BL_BUSY_HEAP_ENTRY, *PBL_BUSY_HEAP_ENTRY;
46 
47 typedef struct _BL_HEAP_BOUNDARIES
48 {
49     LIST_ENTRY ListEntry;
50     ULONG_PTR HeapEnd;
51     ULONG_PTR HeapLimit;
52     ULONG_PTR HeapBase;
53     PBL_BUSY_HEAP_ENTRY HeapStart;
54 } BL_HEAP_BOUNDARIES, *PBL_HEAP_BOUNDARIES;
55 
56 ULONG HapInitializationStatus;
57 LIST_ENTRY MmHeapBoundaries;
58 ULONG HapMinimumHeapSize;
59 ULONG HapAllocationAttributes;
60 PBL_FREE_HEAP_ENTRY* MmFreeList;
61 
62 /* INLINES *******************************************************************/
63 
64 FORCEINLINE
65 PBL_FREE_HEAP_ENTRY
MmHapDecodeLink(_In_ BL_HEAP_POINTER Link)66 MmHapDecodeLink (
67     _In_ BL_HEAP_POINTER Link
68     )
69 {
70     /* Decode the buffer pointer by ignoring the flags */
71     return (PBL_FREE_HEAP_ENTRY)(Link.BufferPointer << BL_HEAP_POINTER_FLAG_BITS);
72 }
73 
74 FORCEINLINE
75 ULONG
MmHapBufferSize(_In_ PVOID FreeEntry)76 MmHapBufferSize (
77     _In_ PVOID FreeEntry
78     )
79 {
80     PBL_FREE_HEAP_ENTRY Entry = FreeEntry;
81 
82     /* The space between the next buffer header and this one is the size */
83     return (ULONG_PTR)MmHapDecodeLink(Entry->BufferNext) - (ULONG_PTR)Entry;
84 }
85 
86 FORCEINLINE
87 ULONG
MmHapUserBufferSize(_In_ PVOID FreeEntry)88 MmHapUserBufferSize (
89     _In_ PVOID FreeEntry
90     )
91 {
92     PBL_FREE_HEAP_ENTRY Entry = FreeEntry;
93 
94     /* Get the size of the buffer as the user sees it */
95     return MmHapBufferSize(Entry) - FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer);
96 }
97 
98 
99 /* FUNCTIONS *****************************************************************/
100 
101 NTSTATUS
MmHapHeapAllocatorExtend(_In_ ULONG ExtendSize)102 MmHapHeapAllocatorExtend (
103     _In_ ULONG ExtendSize
104     )
105 {
106     ULONG HeapSize, AlignedSize, HeapLimit;
107     PBL_HEAP_BOUNDARIES Heap, NewHeap;
108     NTSTATUS Status;
109     PBL_BUSY_HEAP_ENTRY HeapBase = NULL;
110 
111     /* Compute a new heap, and add 2 more pages for the free list */
112     HeapSize = ExtendSize + (2 * PAGE_SIZE);
113     if (HeapSize < ExtendSize)
114     {
115         return STATUS_INTEGER_OVERFLOW;
116     }
117 
118     /* Make sure the new heap is at least the minimum configured size */
119     if (HapMinimumHeapSize > HeapSize)
120     {
121         HeapSize = HapMinimumHeapSize;
122     }
123 
124     /* Align it on a page boundary */
125     AlignedSize = ALIGN_UP_BY(HeapSize, PAGE_SIZE);
126     if (!AlignedSize)
127     {
128         return STATUS_INTEGER_OVERFLOW;
129     }
130 
131     /* Check if we already have a heap */
132     if (!IsListEmpty(&MmHeapBoundaries))
133     {
134         /* Find the first heap*/
135         Heap = CONTAINING_RECORD(MmHeapBoundaries.Flink,
136                                  BL_HEAP_BOUNDARIES,
137                                  ListEntry);
138 
139         /* Check if we have a page free above the heap */
140         HeapLimit = Heap->HeapLimit + PAGE_SIZE;
141         if (HeapLimit <= Heap->HeapEnd)
142         {
143             EfiPrintf(L"Heap extension TODO\r\n");
144             return STATUS_INSUFFICIENT_RESOURCES;
145         }
146     }
147 
148     /* We do not -- allocate one */
149     Status = MmPapAllocatePagesInRange((PVOID*)&HeapBase,
150                                        BlLoaderHeap,
151                                        AlignedSize >> PAGE_SHIFT,
152                                        HapAllocationAttributes,
153                                        0,
154                                        NULL,
155                                        0);
156     if (!NT_SUCCESS(Status))
157     {
158         EfiPrintf(L"HEAP ALLOCATION FAILED\r\n");
159         EfiStall(1000000);
160         return Status;
161     }
162 
163     /* Set the heap bottom, limit, and top */
164     NewHeap = (PBL_HEAP_BOUNDARIES)HeapBase->Buffer;
165     NewHeap->HeapBase = (ULONG_PTR)HeapBase;
166     NewHeap->HeapLimit = (ULONG_PTR)HeapBase + AlignedSize;
167     NewHeap->HeapStart = (PBL_BUSY_HEAP_ENTRY)(NewHeap + 1);
168 
169     /* Set the buffer links */
170     HeapBase->BufferPrevious.P = NULL;
171     HeapBase->BufferNext.P = NewHeap->HeapStart;
172 
173     /* Set the buffer at the top of the heap and mark it as being free */
174     NewHeap->HeapStart->BufferPrevious.P = HeapBase;
175     NewHeap->HeapStart->BufferNext.P =  NewHeap->HeapStart;
176     NewHeap->HeapStart->BufferNext.BufferFree = 1;
177     NewHeap->HeapStart->BufferNext.BufferOnHeap = 1;
178 
179     /* Is this the first heap ever? */
180     if (IsListEmpty(&MmHeapBoundaries))
181     {
182         /* We will host the free list at the top of the heap */
183         MmFreeList = (PBL_FREE_HEAP_ENTRY*)((ULONG_PTR)NewHeap->HeapLimit - 8 * sizeof(PBL_FREE_HEAP_ENTRY));
184         NewHeap->HeapLimit = (ULONG_PTR)MmFreeList;
185         RtlZeroMemory(MmFreeList, 8 * sizeof(PBL_FREE_HEAP_ENTRY));
186     }
187 
188     /* Remove a page on top */
189     HeapLimit = NewHeap->HeapLimit;
190     NewHeap->HeapEnd = NewHeap->HeapLimit;
191     NewHeap->HeapLimit -= PAGE_SIZE;
192 
193     /* Add us into the heap list */
194     InsertTailList(&MmHeapBoundaries, &NewHeap->ListEntry);
195     return STATUS_SUCCESS;
196 }
197 
198 ULONG
MmHapGetBucketId(_In_ ULONG Size)199 MmHapGetBucketId (
200     _In_ ULONG Size
201     )
202 {
203     ULONG BucketIndex = 0;
204 
205     /* Use the last bucket if this is a large allocation */
206     if (Size >= PAGE_SIZE)
207     {
208         return 7;
209     }
210 
211     /* Otherwise, use a higher index for each new power of two */
212     while (Size >> BucketIndex)
213     {
214         BucketIndex++;
215     }
216 
217     /* Allocations are at least 16 bytes (2^4 = 5th index) */
218     return BucketIndex - 5;
219 }
220 
221 VOID
MmHapReportHeapCorruption(_In_ PBL_FREE_HEAP_ENTRY BufferEntry)222 MmHapReportHeapCorruption (
223     _In_ PBL_FREE_HEAP_ENTRY BufferEntry
224     )
225 {
226 #if 0
227     BOOLEAN DebuggerEnabled;
228 
229     BlStatusPrint(L"Heap corruption in the links surrounding %p!\r\n", BufferEntry);
230 
231     DebuggerEnabled = BlBdDebuggerEnabled();
232     if (DebuggerEnabled)
233     {
234         BlStatusPrint(L"\n*** Fatal Error 0x%08x :\n                (0x%p, 0x%p, 0x%p, 0x%p)\n\r\n", 2, BufferEntry, NULL, NULL, NULL);
235         __debugbreak();
236     }
237 #else
238     EfiPrintf(L"Heap corruption in the links surrounding %p!\r\n", BufferEntry);
239 #endif
240 }
241 
242 PVOID
MmHapCheckFreeLinks(_In_ PVOID BufferEntry)243 MmHapCheckFreeLinks (
244     _In_ PVOID BufferEntry
245     )
246 {
247     PBL_FREE_HEAP_ENTRY Prev, Next;
248     PBL_FREE_HEAP_ENTRY Entry = BufferEntry;
249 
250     /* Get the previous and next free pointers */
251     Prev = MmHapDecodeLink(Entry->FreePrevious);
252     Next = MmHapDecodeLink(Entry->FreeNext);
253 
254     /* Make sure that both the previous and next entries point to this one */
255     if (((Next) && (MmHapDecodeLink(Next->FreePrevious)) != Entry) ||
256         ((Prev) && (MmHapDecodeLink(Prev->FreeNext)) != Entry))
257     {
258         /* They don't, so the free headers are corrupted */
259         MmHapReportHeapCorruption(Entry);
260         return NULL;
261     }
262 
263     /* They do, return the free entry as valid */
264     return Entry;
265 }
266 
267 PVOID
MmHapCheckBufferLinks(_In_ PVOID BufferEntry)268 MmHapCheckBufferLinks (
269     _In_ PVOID BufferEntry
270     )
271 {
272     PBL_FREE_HEAP_ENTRY Prev, Next;
273     PBL_FREE_HEAP_ENTRY Entry = BufferEntry;
274 
275     /* Get the previous and next buffer pointers */
276     Prev = MmHapDecodeLink(Entry->BufferPrevious);
277     Next = MmHapDecodeLink(Entry->BufferNext);
278 
279     /* Make sure that both the previous and next entries point to this one */
280     if (((Next) && (MmHapDecodeLink(Next->BufferPrevious)) != Entry) ||
281         ((Prev) && (MmHapDecodeLink(Prev->BufferNext)) != Entry))
282     {
283         /* They don't, so the heap headers are corrupted */
284         MmHapReportHeapCorruption(Entry);
285         return NULL;
286     }
287 
288     /* They, do the entry is valid */
289     return Entry;
290 }
291 
292 PBL_FREE_HEAP_ENTRY
MmHapRemoveBufferFromFreeList(_In_ PBL_FREE_HEAP_ENTRY FreeEntry)293 MmHapRemoveBufferFromFreeList (
294     _In_ PBL_FREE_HEAP_ENTRY FreeEntry
295     )
296 {
297     PBL_FREE_HEAP_ENTRY Prev, Next;
298 
299     /* Firest, make sure the free entry is valid */
300     FreeEntry = MmHapCheckFreeLinks(FreeEntry);
301     if (!FreeEntry)
302     {
303         return FreeEntry;
304     }
305 
306     /* Get the previous and next entry */
307     Prev = MmHapDecodeLink(FreeEntry->FreePrevious);
308     Next = MmHapDecodeLink(FreeEntry->FreeNext);
309 
310     /* Update the next entry to point to our previous entry */
311     if (Next)
312     {
313         Next->FreePrevious.P = Prev;
314     }
315 
316     /* Are we at the head? */
317     if (Prev)
318     {
319         /* Nope, so update our previous entry to point to our next entry */
320         Prev->FreeNext.P = Next;
321     }
322     else
323     {
324         /* Yep, so update the appropriate bucket listhead */
325         MmFreeList[MmHapGetBucketId(MmHapBufferSize(FreeEntry))] = Prev;
326     }
327 
328     /* Return the (now removed) entry */
329     return FreeEntry;
330 }
331 
332 PBL_FREE_HEAP_ENTRY
MmHapCoalesceFreeBuffer(_In_ PBL_FREE_HEAP_ENTRY FreeEntry)333 MmHapCoalesceFreeBuffer (
334     _In_ PBL_FREE_HEAP_ENTRY FreeEntry
335     )
336 {
337     PBL_FREE_HEAP_ENTRY Prev, Next;
338 
339     /* First make sure that this is a valid buffer entry */
340     if (!MmHapCheckBufferLinks(FreeEntry))
341     {
342         return NULL;
343     }
344 
345     /* Get the next entry and check if it's free */
346     Next = MmHapDecodeLink(FreeEntry->BufferNext);
347     if (!(Next->BufferNext.BufferOnHeap) && (Next->BufferNext.BufferFree))
348     {
349         /* Remove the next buffer from the free list since we're coalescing */
350         Next = MmHapRemoveBufferFromFreeList(Next);
351         if (!Next)
352         {
353             return NULL;
354         }
355 
356         /* The forward link of the *new* free buffer should now point to us */
357         MmHapDecodeLink(Next->BufferNext)->BufferPrevious.P = FreeEntry;
358 
359         /* Our forward link should point to the *new* free buffer as well */
360         FreeEntry->BufferNext.P = MmHapDecodeLink(Next->BufferNext);
361 
362         /* Mark our buffer as free */
363         FreeEntry->BufferNext.BufferFree = 1;
364     }
365 
366     /* Get the previous entry and check if it's free */
367     Prev = MmHapDecodeLink(FreeEntry->BufferPrevious);
368     if (!(Prev) || !(Prev->BufferNext.BufferFree))
369     {
370         return FreeEntry;
371     }
372 
373     /* It's free, so remove it */
374     Prev = MmHapRemoveBufferFromFreeList(Prev);
375     if (!Prev)
376     {
377         return NULL;
378     }
379 
380     /* The previous link of our next buffer should now point to our *previous* */
381     MmHapDecodeLink(FreeEntry->BufferNext)->BufferPrevious.P = Prev;
382 
383     /* Our previous link should point the next free buffer now */
384     Prev->BufferNext.P = MmHapDecodeLink(FreeEntry->BufferNext);
385 
386     /* Set the new freed buffer as the previous buffer, and mark it free */
387     FreeEntry = Prev;
388     FreeEntry->BufferNext.BufferFree = 1;
389     return FreeEntry;
390 }
391 
392 PBL_FREE_HEAP_ENTRY
MmHapAddToFreeList(_In_ PBL_BUSY_HEAP_ENTRY Entry,_In_ ULONG Flags)393 MmHapAddToFreeList (
394     _In_ PBL_BUSY_HEAP_ENTRY Entry,
395     _In_ ULONG Flags
396     )
397 {
398     PBL_FREE_HEAP_ENTRY FreeEntry, Head;
399     ULONG BucketId;
400     BL_LIBRARY_PARAMETERS LocalParameters;
401 
402     /* First, check if the entry is valid */
403     Entry = MmHapCheckBufferLinks(Entry);
404     if (!Entry)
405     {
406         return NULL;
407     }
408 
409     /* Check if we should zero the entry */
410     LocalParameters = BlpLibraryParameters;
411     if ((LocalParameters.LibraryFlags & BL_LIBRARY_FLAG_ZERO_HEAP_ALLOCATIONS_ON_FREE) &&
412         !(Flags))
413     {
414         /* Yep, zero it out */
415         RtlZeroMemory(Entry->Buffer, MmHapUserBufferSize(Entry));
416     }
417 
418     /* Now mark the entry as free */
419     Entry->BufferNext.BufferFree = 1;
420 
421     /* Now that this buffer is free, try to coalesce it */
422     FreeEntry = MmHapCoalesceFreeBuffer((PBL_FREE_HEAP_ENTRY)Entry);
423     if (!FreeEntry)
424     {
425         return FreeEntry;
426     }
427 
428     /* Compute the bucket ID for the free list */
429     BucketId = MmHapGetBucketId(MmHapBufferSize(Entry));
430 
431     /* Get the current head for this bucket, if one exists */
432     Head = MmFreeList ? MmFreeList[BucketId] : NULL;
433 
434     /* Update the head's backlink to point to this newly freed entry */
435     if (Head)
436     {
437         Head->FreePrevious.P = FreeEntry;
438     }
439 
440     /* Nobody behind us, the old head in front of us */
441     FreeEntry->FreePrevious.P = NULL;
442     FreeEntry->FreeNext.P = Head;
443 
444     /* Put us at the head of list now, and return the entry */
445     MmFreeList[BucketId] = FreeEntry;
446     return FreeEntry;
447 }
448 
449 PBL_BUSY_HEAP_ENTRY
MmHapFindBufferInFreeList(_In_ ULONG Size)450 MmHapFindBufferInFreeList (
451     _In_ ULONG Size
452     )
453 {
454     PBL_FREE_HEAP_ENTRY FreeEntry = NULL;
455     PBL_BUSY_HEAP_ENTRY NextEntry;
456     ULONG BucketId;
457 
458     /* Get the appropriate bucket for our size */
459     BucketId = MmHapGetBucketId(Size);
460     if (BucketId >= 8)
461     {
462         return NULL;
463     }
464 
465     /* Keep going as long as we don't have a free entry */
466     while (!FreeEntry)
467     {
468         /* Fet the first free entry in this list */
469         FreeEntry = MmFreeList ? MmFreeList[BucketId] : NULL;
470 
471         /* Loop as long as there's entries in the list */
472         while (FreeEntry)
473         {
474             /* Can this free entry satisfy our needs? */
475             if (MmHapBufferSize(FreeEntry) >= Size)
476             {
477                 /* All good */
478                 break;
479             }
480 
481             /* It cannot, keep going to the next one */
482             FreeEntry = MmHapDecodeLink(FreeEntry->FreeNext);
483         }
484 
485         /* Try the next list -- have we exhausted all the lists? */
486         if (++BucketId >= 8)
487         {
488             /* Have we not found an entry yet? Fail if so... */
489             if (!FreeEntry)
490             {
491                 return NULL;
492             }
493         }
494     }
495 
496     /* We should have an entry if we're here. Remove it from the free list */
497     NT_ASSERT(FreeEntry != NULL);
498     FreeEntry = MmHapRemoveBufferFromFreeList(FreeEntry);
499     if (!FreeEntry)
500     {
501         return NULL;
502     }
503 
504     /* Make sure it's not corrupted */
505     FreeEntry = MmHapCheckBufferLinks(FreeEntry);
506     if (!FreeEntry)
507     {
508         return NULL;
509     }
510 
511     /* Do we have space for at least another buffer? */
512     if ((MmHapBufferSize(FreeEntry) - Size) >= sizeof(BL_FREE_HEAP_ENTRY))
513     {
514         /* Go to where the new next buffer will  start */
515         NextEntry = (PBL_BUSY_HEAP_ENTRY)((ULONG_PTR)FreeEntry + Size);
516 
517         /* Make the new next buffer point to the next buffer */
518         NextEntry->BufferNext.P = MmHapDecodeLink(FreeEntry->BufferNext);
519 
520         /* Make the old next buffer point back to the new one */
521         MmHapDecodeLink(FreeEntry->BufferNext)->BufferPrevious.P = NextEntry;
522 
523         /* Point the new next buffer point back to us */
524         NextEntry->BufferPrevious.P = FreeEntry;
525 
526         /* Point us to the new next buffer */
527         FreeEntry->BufferNext.P = NextEntry;
528 
529         /* And insert the new next buffer into the free list */
530         MmHapAddToFreeList(NextEntry, 1);
531     }
532 
533     /* Return the entry, which is now allocated */
534     return (PBL_BUSY_HEAP_ENTRY)FreeEntry;
535 }
536 
537 NTSTATUS
MmHaInitialize(_In_ ULONG HeapSize,_In_ ULONG HeapAttributes)538 MmHaInitialize (
539     _In_ ULONG HeapSize,
540     _In_ ULONG HeapAttributes
541     )
542 {
543     NTSTATUS Status;
544 
545     /* No free list to begin with */
546     MmFreeList = NULL;
547 
548     /* Configure the minimum heap size and allocation attributes */
549     HapMinimumHeapSize = ALIGN_UP_BY(HeapSize, PAGE_SIZE);
550     HapAllocationAttributes = HeapAttributes & 0x20000;
551 
552     /* Initialize the heap boundary list */
553     InitializeListHead(&MmHeapBoundaries);
554 
555     /* Initialize a heap big enough to handle a one pointer long allocation */
556     Status = MmHapHeapAllocatorExtend(sizeof(PVOID));
557     if (NT_SUCCESS(Status))
558     {
559         /* The heap is ready! */
560         HapInitializationStatus = 1;
561         Status = STATUS_SUCCESS;
562     }
563 
564     /* Return initialization status */
565     return Status;
566 }
567 
568 PVOID
BlMmAllocateHeap(_In_ SIZE_T Size)569 BlMmAllocateHeap (
570     _In_ SIZE_T Size
571     )
572 {
573     ULONG BufferSize;
574     PBL_HEAP_BOUNDARIES Heap;
575     PBL_BUSY_HEAP_ENTRY BusyEntry, FreeEntry, NextEntry;
576 
577     /* Ignore heap allocation if the heap allocator isn't ready yet */
578     if (HapInitializationStatus != 1)
579     {
580         return NULL;
581     }
582 
583     /* Align the buffer size to the minimum size required */
584     BufferSize = ALIGN_UP_BY(Size + FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer),
585                              FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer));
586 
587     /* Watch out for overflow */
588     if (BufferSize <= Size)
589     {
590         return NULL;
591     }
592 
593     /* Make sure it's at least big enough to hold a free entry later on */
594     if (BufferSize < sizeof(BL_FREE_HEAP_ENTRY))
595     {
596         BufferSize = sizeof(BL_FREE_HEAP_ENTRY);
597     }
598 
599     /* Loop while we try to allocate memory */
600     while (1)
601     {
602         /* Find a free buffer for this allocation */
603         BusyEntry = MmHapFindBufferInFreeList(BufferSize);
604         if (BusyEntry)
605         {
606             break;
607         }
608 
609         /* We couldn't find a free buffer. Do we have any heaps? */
610         if (!IsListEmpty(&MmHeapBoundaries))
611         {
612             /* Get the current heap */
613             Heap = CONTAINING_RECORD(MmHeapBoundaries.Flink,
614                                      BL_HEAP_BOUNDARIES,
615                                      ListEntry);
616 
617             /* Check if we have space in the heap page for this allocation? */
618             FreeEntry = Heap->HeapStart;
619             NextEntry = (PBL_BUSY_HEAP_ENTRY)((ULONG_PTR)FreeEntry + BufferSize);
620 
621             if ((NextEntry >= FreeEntry) &&
622                 ((ULONG_PTR)NextEntry <=
623                  Heap->HeapLimit - FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer)))
624             {
625                 /* Update the heap top pointer past this allocation */
626                 Heap->HeapStart = NextEntry;
627 
628                 /* Make this allocation point to the slot */
629                 FreeEntry->BufferNext.P = Heap->HeapStart;
630 
631                 /* And make the free heap entry point back to us */
632                 Heap->HeapStart->BufferPrevious.P = FreeEntry;
633 
634                 /* Mark the heap entry as being free and on the heap */
635                 Heap->HeapStart->BufferNext.BufferFree = 1;
636                 Heap->HeapStart->BufferNext.BufferOnHeap = 1;
637 
638                 /* The previously freed entry on the heap page is now ours */
639                 BusyEntry = FreeEntry;
640                 break;
641             }
642         }
643 
644         /* We have no heaps or space on any heap -- extend the heap and retry */
645         if (!NT_SUCCESS(MmHapHeapAllocatorExtend(BufferSize)))
646         {
647             EfiPrintf(L"Heap extension failed!\r\n");
648             return NULL;
649         }
650 
651         EfiPrintf(L"Heap extended -- trying again\r\n");
652     }
653 
654     /* Clear all the bits, marking this entry as allocated */
655     BusyEntry->BufferNext.P = MmHapDecodeLink(BusyEntry->BufferNext);
656 
657     /* Return the entry's data buffer */
658     //EfiPrintf(L"Returning buffer at 0x%p\r\n", &BusyEntry->Buffer);
659     return &BusyEntry->Buffer;
660 }
661 
662 NTSTATUS
BlMmFreeHeap(_In_ PVOID Buffer)663 BlMmFreeHeap (
664     _In_ PVOID Buffer
665     )
666 {
667     PBL_BUSY_HEAP_ENTRY BusyEntry;
668     PBL_HEAP_BOUNDARIES Heap;
669     PLIST_ENTRY NextEntry;
670 
671     /* If the heap is not initialized, fail */
672     if (HapInitializationStatus != 1)
673     {
674         return STATUS_UNSUCCESSFUL;
675     }
676 
677     /* Get the heap header */
678     //EfiPrintf(L"Freeing entry at: %p\r\n", Buffer);
679     if (Buffer)
680     {
681         /* Don't free heap until we discover the corruption */
682         return STATUS_SUCCESS;
683     }
684 
685     BusyEntry = CONTAINING_RECORD(Buffer, BL_BUSY_HEAP_ENTRY, Buffer);
686 
687     /* Loop all the heaps */
688     NextEntry = MmHeapBoundaries.Flink;
689     while (NextEntry != &MmHeapBoundaries)
690     {
691         /* Get the current heap in the list */
692         Heap = CONTAINING_RECORD(NextEntry, BL_HEAP_BOUNDARIES, ListEntry);
693 
694         /* Is this entry part of this heap? */
695         if (((ULONG_PTR)Heap->HeapBase <= (ULONG_PTR)BusyEntry) &&
696             ((ULONG_PTR)BusyEntry < (ULONG_PTR)Heap->HeapStart))
697         {
698             /* Ignore double-free */
699             if (BusyEntry->BufferNext.BufferFree)
700             {
701                 return STATUS_INVALID_PARAMETER;
702             }
703 
704             /* It is -- add it to the free list */
705             MmHapAddToFreeList(BusyEntry, 0);
706             return STATUS_SUCCESS;
707         }
708 
709         /* It isn't, move to the next heap */
710         NextEntry = NextEntry->Flink;
711     }
712 
713     /* The entry is not on any valid heap */
714     return STATUS_INVALID_PARAMETER;
715 }
716 
717