reactos/boot/environ/lib/mm/heapalloc.c

   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
  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
  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
  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
 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             EarlyPrint(L"Heap extension TODO\n");
 144             return STATUS_INSUFFICIENT_RESOURCES;
 145         }
 146     }
 147
 148     /* We do not -- allocate one */
 149     Status = MmPapAllocatePagesInRange((PULONG)&HeapBase,
 150                                        BlLoaderHeap,
 151                                        AlignedSize >> PAGE_SHIFT,
 152                                        HapAllocationAttributes,
 153                                        0,
 154                                        NULL,
 155                                        0);
 156     if (!NT_SUCCESS(Status))
 157     {
 158         return Status;
 159     }
 160
 161     /* Set the heap bottom, limit, and top */
 162     NewHeap = (PBL_HEAP_BOUNDARIES)HeapBase->Buffer;
 163     NewHeap->HeapBase = (ULONG_PTR)HeapBase;
 164     NewHeap->HeapLimit = (ULONG_PTR)HeapBase + AlignedSize;
 165     NewHeap->HeapStart = (PBL_BUSY_HEAP_ENTRY)(NewHeap + 1);
 166
 167     /* Set the buffer links */
 168     HeapBase->BufferPrevious.P = NULL;
 169     HeapBase->BufferNext.P = NewHeap->HeapStart;
 170
 171     /* Set the buffer at the top of the heap and mark it as being free */
 172     NewHeap->HeapStart->BufferPrevious.P = HeapBase;
 173     NewHeap->HeapStart->BufferNext.P =  NewHeap->HeapStart;
 174     NewHeap->HeapStart->BufferNext.BufferFree = 1;
 175     NewHeap->HeapStart->BufferNext.BufferOnHeap = 1;
 176
 177     /* Is this the first heap ever? */
 178     if (IsListEmpty(&MmHeapBoundaries))
 179     {
 180         /* We will host the free list at the top of the heap */
 181         MmFreeList = (PBL_FREE_HEAP_ENTRY*)((ULONG_PTR)NewHeap->HeapLimit - 8 * sizeof(PBL_FREE_HEAP_ENTRY));
 182         NewHeap->HeapLimit = (ULONG_PTR)MmFreeList;
 183         RtlZeroMemory(MmFreeList, 8 * sizeof(PBL_FREE_HEAP_ENTRY));
 184     }
 185
 186     /* Remove a page on top */
 187     HeapLimit = NewHeap->HeapLimit;
 188     NewHeap->HeapEnd = NewHeap->HeapLimit;
 189     NewHeap->HeapLimit -= PAGE_SIZE;
 190
 191     /* Add us into the heap list */
 192     InsertTailList(&MmHeapBoundaries, &NewHeap->ListEntry);
 193     return STATUS_SUCCESS;
 194 }
 195
 196 ULONG
 197 MmHapGetBucketId (
 198     _In_ ULONG Size
 199     )
 200 {
 201     ULONG BucketIndex = 0;
 202
 203     /* Use the last bucket if this is a large allocation */
 204     if (Size >= PAGE_SIZE)
 205     {
 206         return 7;
 207     }
 208
 209     /* Otherwise, use a higher index for each new power of two */
 210     while (Size >> BucketIndex)
 211     {
 212         BucketIndex++;
 213     }
 214
 215     /* Allocations are at least 16 bytes (2^4 = 5th index) */
 216     return BucketIndex - 5;
 217 }
 218
 219 VOID
 220 MmHapReportHeapCorruption (
 221     _In_ PBL_FREE_HEAP_ENTRY BufferEntry
 222     )
 223 {
 224 #if 0
 225     BOOLEAN DebuggerEnabled;
 226
 227     BlStatusPrint(L"Heap corruption in the links surrounding %p!\n", BufferEntry);
 228
 229     DebuggerEnabled = BlBdDebuggerEnabled();
 230     if (DebuggerEnabled)
 231     {
 232         BlStatusPrint(L"\n*** Fatal Error 0x%08x :\n                (0x%p, 0x%p, 0x%p, 0x%p)\n\n", 2, BufferEntry, NULL, NULL, NULL);
 233         __debugbreak();
 234     }
 235 #else
 236     EarlyPrint(L"Heap corruption in the links surrounding %p!\n", BufferEntry);
 237 #endif
 238 }
 239
 240 PVOID
 241 MmHapCheckFreeLinks (
 242     _In_ PVOID BufferEntry
 243     )
 244 {
 245     PBL_FREE_HEAP_ENTRY Prev, Next;
 246     PBL_FREE_HEAP_ENTRY Entry = BufferEntry;
 247
 248     /* Get the previous and next free pointers */
 249     Prev = MmHapDecodeLink(Entry->FreePrevious);
 250     Next = MmHapDecodeLink(Entry->FreeNext);
 251
 252     /* Make sure that both the previous and next entries point to this one */
 253     if (((Next) && (MmHapDecodeLink(Next->FreePrevious)) != Entry) ||
 254         ((Prev) && (MmHapDecodeLink(Prev->FreeNext)) != Entry))
 255     {
 256         /* They don't, so the free headers are corrupted */
 257         MmHapReportHeapCorruption(Entry);
 258         return NULL;
 259     }
 260
 261     /* They do, return the free entry as valid */
 262     return Entry;
 263 }
 264
 265 PVOID
 266 MmHapCheckBufferLinks (
 267     _In_ PVOID BufferEntry
 268     )
 269 {
 270     PBL_FREE_HEAP_ENTRY Prev, Next;
 271     PBL_FREE_HEAP_ENTRY Entry = BufferEntry;
 272
 273     /* Get the previous and next buffer pointers */
 274     Prev = MmHapDecodeLink(Entry->BufferPrevious);
 275     Next = MmHapDecodeLink(Entry->BufferNext);
 276
 277     /* Make sure that both the previous and next entries point to this one */
 278     if (((Next) && (MmHapDecodeLink(Next->BufferPrevious)) != Entry) ||
 279         ((Prev) && (MmHapDecodeLink(Prev->BufferNext)) != Entry))
 280     {
 281         /* They don't, so the heap headers are corrupted */
 282         MmHapReportHeapCorruption(Entry);
 283         return NULL;
 284     }
 285
 286     /* They, do the entry is valid */
 287     return Entry;
 288 }
 289
 290 PBL_FREE_HEAP_ENTRY
 291 MmHapRemoveBufferFromFreeList (
 292     _In_ PBL_FREE_HEAP_ENTRY FreeEntry
 293     )
 294 {
 295     PBL_FREE_HEAP_ENTRY Prev, Next;
 296
 297     /* Firest, make sure the free entry is valid */
 298     FreeEntry = MmHapCheckFreeLinks(FreeEntry);
 299     if (!FreeEntry)
 300     {
 301         return FreeEntry;
 302     }
 303
 304     /* Get the previous and next entry */
 305     Prev = MmHapDecodeLink(FreeEntry->FreePrevious);
 306     Next = MmHapDecodeLink(FreeEntry->FreeNext);
 307
 308     /* Update the next entry to point to our previous entry */
 309     if (Next)
 310     {
 311         Next->FreePrevious.P = Prev;
 312     }
 313
 314     /* Are we at the head? */
 315     if (Prev)
 316     {
 317         /* Nope, so update our previous entry to point to our next entry */
 318         Prev->FreeNext.P = Next;
 319     }
 320     else
 321     {
 322         /* Yep, so update the appropriate bucket listhead */
 323         MmFreeList[MmHapGetBucketId(MmHapBufferSize(FreeEntry))] = Prev;
 324     }
 325
 326     /* Return the (now removed) entry */
 327     return FreeEntry;
 328 }
 329
 330 PBL_FREE_HEAP_ENTRY
 331 MmHapCoalesceFreeBuffer (
 332     _In_ PBL_FREE_HEAP_ENTRY FreeEntry
 333     )
 334 {
 335     PBL_FREE_HEAP_ENTRY Prev, Next;
 336
 337     /* First make sure that this is a valid buffer entry */
 338     if (!MmHapCheckBufferLinks(FreeEntry))
 339     {
 340         return NULL;
 341     }
 342
 343     /* Get the next entry and check if it's free */
 344     Next = MmHapDecodeLink(FreeEntry->BufferNext);
 345     if (!(Next->BufferNext.BufferOnHeap) && (Next->BufferNext.BufferFree))
 346     {
 347         /* Remove the next buffer from the free list since we're coalescing */
 348         Next = MmHapRemoveBufferFromFreeList(Next);
 349         if (!Next)
 350         {
 351             return NULL;
 352         }
 353
 354         /* The forward link of the *new* free buffer should now point to us */
 355         MmHapDecodeLink(Next->BufferNext)->BufferPrevious.P = FreeEntry;
 356
 357         /* Our forward link should point to the *new* free buffer as well */
 358         FreeEntry->BufferNext.P = MmHapDecodeLink(Next->BufferNext);
 359
 360         /* Mark our buffer as free */
 361         FreeEntry->BufferNext.BufferFree = 1;
 362     }
 363
 364     /* Get the previous entry and check if it's free */
 365     Prev = MmHapDecodeLink(FreeEntry->BufferPrevious);
 366     if (!(Prev) || !(Prev->BufferNext.BufferFree))
 367     {
 368         return FreeEntry;
 369     }
 370
 371     /* It's free, so remove it */
 372     Prev = MmHapRemoveBufferFromFreeList(Prev);
 373     if (!Prev)
 374     {
 375         return NULL;
 376     }
 377
 378     /* The previous link of our next buffer should now point to our *previous* */
 379     MmHapDecodeLink(FreeEntry->BufferNext)->BufferPrevious.P = Prev;
 380
 381     /* Our previous link should point the next free buffer now */
 382     Prev->BufferNext.P = MmHapDecodeLink(FreeEntry->BufferNext);
 383
 384     /* Set the new freed buffer as the previous buffer, and mark it free */
 385     FreeEntry = Prev;
 386     FreeEntry->BufferNext.BufferFree = 1;
 387     return FreeEntry;
 388 }
 389
 390 PBL_FREE_HEAP_ENTRY
 391 MmHapAddToFreeList (
 392     _In_ PBL_BUSY_HEAP_ENTRY Entry,
 393     _In_ ULONG Flags
 394     )
 395 {
 396     PBL_FREE_HEAP_ENTRY FreeEntry, Head;
 397     ULONG BucketId;
 398     BL_LIBRARY_PARAMETERS LocalParameters;
 399
 400     /* First, check if the entry is valid */
 401     Entry = MmHapCheckBufferLinks(Entry);
 402     if (!Entry)
 403     {
 404         return NULL;
 405     }
 406
 407     /* Check if we should zero the entry */
 408     LocalParameters = BlpLibraryParameters;
 409     if ((LocalParameters.LibraryFlags & BL_LIBRARY_FLAG_ZERO_HEAP_ALLOCATIONS_ON_FREE) &&
 410         !(Flags))
 411     {
 412         /* Yep, zero it out */
 413         RtlZeroMemory(Entry->Buffer, MmHapUserBufferSize(Entry));
 414     }
 415
 416     /* Now mark the entry as free */
 417     Entry->BufferNext.BufferFree = 1;
 418
 419     /* Now that this buffer is free, try to coalesce it */
 420     FreeEntry = MmHapCoalesceFreeBuffer((PBL_FREE_HEAP_ENTRY)Entry);
 421     if (!FreeEntry)
 422     {
 423         return FreeEntry;
 424     }
 425
 426     /* Compute the bucket ID for the free list */
 427     BucketId = MmHapGetBucketId(MmHapBufferSize(Entry));
 428
 429     /* Get the current head for this bucket, if one exists */
 430     Head = MmFreeList ? MmFreeList[BucketId] : NULL;
 431
 432     /* Update the head's backlink to point to this newly freed entry */
 433     if (Head)
 434     {
 435         Head->FreePrevious.P = FreeEntry;
 436     }
 437
 438     /* Nobody behind us, the old head in front of us */
 439     FreeEntry->FreePrevious.P = NULL;
 440     FreeEntry->FreeNext.P = Head;
 441
 442     /* Put us at the head of list now, and return the entry */
 443     MmFreeList[BucketId] = FreeEntry;
 444     return FreeEntry;
 445 }
 446
 447 PBL_BUSY_HEAP_ENTRY
 448 MmHapFindBufferInFreeList (
 449     _In_ ULONG Size
 450     )
 451 {
 452     PBL_FREE_HEAP_ENTRY FreeEntry = NULL;
 453     PBL_BUSY_HEAP_ENTRY NextEntry;
 454     ULONG BucketId;
 455
 456     /* Get the appropriate bucket for our size */
 457     BucketId = MmHapGetBucketId(Size);
 458     if (BucketId >= 8)
 459     {
 460         return NULL;
 461     }
 462
 463     /* Keep going as long as we don't have a free entry */
 464     while (!FreeEntry)
 465     {
 466         /* Fet the first free entry in this list */
 467         FreeEntry = MmFreeList ? MmFreeList[BucketId] : NULL;
 468
 469         /* Loop as long as there's entries in the list */
 470         while (FreeEntry)
 471         {
 472             /* Can this free entry satisfy our needs? */
 473             if (MmHapBufferSize(FreeEntry) >= Size)
 474             {
 475                 /* All good */
 476                 break;
 477             }
 478
 479             /* It cannot, keep going to the next one */
 480             FreeEntry = MmHapDecodeLink(FreeEntry->FreeNext);
 481         }
 482
 483         /* Try the next list -- have we exhausted all the lists? */
 484         if (++BucketId >= 8)
 485         {
 486             /* Have we not found an entry yet? Fail if so... */
 487             if (!FreeEntry)
 488             {
 489                 return NULL;
 490             }
 491         }
 492     }
 493
 494     /* We should have an entry if we're here. Remove it from the free list */
 495     NT_ASSERT(FreeEntry != NULL);
 496     FreeEntry = MmHapRemoveBufferFromFreeList(FreeEntry);
 497     if (!FreeEntry)
 498     {
 499         return NULL;
 500     }
 501
 502     /* Make sure it's not corrupted */
 503     FreeEntry = MmHapCheckBufferLinks(FreeEntry);
 504     if (!FreeEntry)
 505     {
 506         return NULL;
 507     }
 508
 509     /* Do we have space for at least another buffer? */
 510     if ((MmHapBufferSize(FreeEntry) - Size) >= sizeof(BL_FREE_HEAP_ENTRY))
 511     {
 512         /* Go to where the new next buffer will  start */
 513         NextEntry = (PBL_BUSY_HEAP_ENTRY)((ULONG_PTR)FreeEntry + Size);
 514
 515         /* Make the new next buffer point to the next buffer */
 516         NextEntry->BufferNext.P = MmHapDecodeLink(FreeEntry->BufferNext);
 517
 518         /* Make the old next buffer point back to the new one */
 519         MmHapDecodeLink(FreeEntry->BufferNext)->BufferPrevious.P = NextEntry;
 520
 521         /* Point the new next buffer point back to us */
 522         NextEntry->BufferPrevious.P = FreeEntry;
 523
 524         /* Point us to the new next buffer */
 525         FreeEntry->BufferNext.P = NextEntry;
 526
 527         /* And insert the new next buffer into the free list */
 528         MmHapAddToFreeList(NextEntry, 1);
 529     }
 530
 531     /* Return the entry, which is now allocated */
 532     return (PBL_BUSY_HEAP_ENTRY)FreeEntry;
 533 }
 534
 535 NTSTATUS
 536 MmHaInitialize (
 537     _In_ ULONG HeapSize,
 538     _In_ ULONG HeapAttributes
 539     )
 540 {
 541     NTSTATUS Status;
 542
 543     /* No free list to begin with */
 544     MmFreeList = NULL;
 545
 546     /* Configure the minimum heap size and allocation attributes */
 547     HapMinimumHeapSize = ALIGN_UP_BY(HeapSize, PAGE_SIZE);
 548     HapAllocationAttributes = HeapAttributes & 0x20000;
 549
 550     /* Initialize the heap boundary list */
 551     InitializeListHead(&MmHeapBoundaries);
 552
 553     /* Initialize a heap big enough to handle a one pointer long allocation */
 554     Status = MmHapHeapAllocatorExtend(sizeof(PVOID));
 555     if (NT_SUCCESS(Status))
 556     {
 557         /* The heap is ready! */
 558         HapInitializationStatus = 1;
 559         Status = STATUS_SUCCESS;
 560     }
 561
 562     /* Return initialization status */
 563     return Status;
 564 }
 565
 566 PVOID
 567 BlMmAllocateHeap (
 568     _In_ ULONG Size
 569     )
 570 {
 571     ULONG BufferSize;
 572     PBL_HEAP_BOUNDARIES Heap;
 573     PBL_BUSY_HEAP_ENTRY BusyEntry, FreeEntry, NextEntry;
 574
 575     /* Ignore heap allocation if the heap allocator isn't ready yet */
 576     if (HapInitializationStatus != 1)
 577     {
 578         return NULL;
 579     }
 580
 581     /* Align the buffer size to the minimum size required */
 582     BufferSize = ALIGN_UP(Size + FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer),
 583                           FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer));
 584
 585     /* Watch out for overflow */
 586     if (BufferSize <= Size)
 587     {
 588         return NULL;
 589     }
 590
 591     /* Make sure it's at least big enough to hold a free entry later on */
 592     if (BufferSize < sizeof(BL_FREE_HEAP_ENTRY))
 593     {
 594         BufferSize = sizeof(BL_FREE_HEAP_ENTRY);
 595     }
 596
 597     /* Loop while we try to allocate memory */
 598     while (1)
 599     {
 600         /* Find a free buffer for this allocation */
 601         BusyEntry = MmHapFindBufferInFreeList(BufferSize);
 602         if (BusyEntry)
 603         {
 604             break;
 605         }
 606
 607         /* We couldn't find a free buffer. Do we have any heaps? */
 608         if (!IsListEmpty(&MmHeapBoundaries))
 609         {
 610             /* Get the current heap */
 611             Heap = CONTAINING_RECORD(MmHeapBoundaries.Flink,
 612                                      BL_HEAP_BOUNDARIES,
 613                                      ListEntry);
 614
 615             /* Check if we have space in the heap page for this allocation? */
 616             FreeEntry = Heap->HeapStart;
 617             NextEntry = (PBL_BUSY_HEAP_ENTRY)((ULONG_PTR)FreeEntry + BufferSize);
 618
 619             if ((NextEntry >= FreeEntry) &&
 620                 ((ULONG_PTR)NextEntry <=
 621                  Heap->HeapLimit - FIELD_OFFSET(BL_BUSY_HEAP_ENTRY, Buffer)))
 622             {
 623                 /* Update the heap top pointer past this allocation */
 624                 Heap->HeapStart = NextEntry;
 625
 626                 /* Make this allocation point to the slot */
 627                 FreeEntry->BufferNext.P = Heap->HeapStart;
 628
 629                 /* And make the free heap entry point back to us */
 630                 Heap->HeapStart->BufferPrevious.P = FreeEntry;
 631
 632                 /* Mark the heap entry as being free and on the heap */
 633                 Heap->HeapStart->BufferNext.BufferFree = 1;
 634                 Heap->HeapStart->BufferNext.BufferOnHeap = 1;
 635
 636                 /* The previously freed entry on the heap page is now ours */
 637                 BusyEntry = FreeEntry;
 638                 break;
 639             }
 640         }
 641
 642         /* We have no heaps or space on any heap -- extend the heap and retry */
 643         if (!NT_SUCCESS(MmHapHeapAllocatorExtend(BufferSize)))
 644         {
 645             EarlyPrint(L"Heap extension failed!\n");
 646             return NULL;
 647         }
 648
 649         EarlyPrint(L"Heap extended -- trying again\n");
 650     }
 651
 652     /* Clear all the bits, marking this entry as allocated */
 653     BusyEntry->BufferNext.P = MmHapDecodeLink(BusyEntry->BufferNext);
 654
 655     /* Return the entry's data buffer */
 656     EarlyPrint(L"Returning buffer at 0x%p\n", &BusyEntry->Buffer);
 657     return &BusyEntry->Buffer;
 658 }
 659
 660 NTSTATUS
 661 BlMmFreeHeap (
 662     _In_ PVOID Buffer
 663     )
 664 {
 665     PBL_BUSY_HEAP_ENTRY BusyEntry;
 666     PBL_HEAP_BOUNDARIES Heap;
 667     PLIST_ENTRY NextEntry;
 668
 669     /* If the heap is not initialized, fail */
 670     if (HapInitializationStatus != 1)
 671     {
 672         return STATUS_UNSUCCESSFUL;
 673     }
 674
 675     /* Get the heap header */
 676     BusyEntry = CONTAINING_RECORD(Buffer, BL_BUSY_HEAP_ENTRY, Buffer);
 677
 678     /* Loop all the heaps */
 679     NextEntry = MmHeapBoundaries.Flink;
 680     while (NextEntry != &MmHeapBoundaries)
 681     {
 682         /* Get the current heap in the list */
 683         Heap = CONTAINING_RECORD(NextEntry, BL_HEAP_BOUNDARIES, ListEntry);
 684
 685         /* Is this entry part of this heap? */
 686         if (((ULONG_PTR)Heap->HeapBase <= (ULONG_PTR)BusyEntry) &&
 687             ((ULONG_PTR)BusyEntry < (ULONG_PTR)Heap->HeapStart))
 688         {
 689             /* Ignore double-free */
 690             if (BusyEntry->BufferNext.BufferFree)
 691             {
 692                 return STATUS_INVALID_PARAMETER;
 693             }
 694
 695             /* It is -- add it to the free list */
 696             MmHapAddToFreeList(BusyEntry, 0);
 697             return STATUS_SUCCESS;
 698         }
 699
 700         /* It isn't, move to the next heap */
 701         NextEntry = NextEntry->Flink;
 702     }
 703
 704     /* The entry is not on any valid heap */
 705     return STATUS_INVALID_PARAMETER;
 706 }
 707