From 545ea51c99141d5093cb44e48055fef2c9c0bbc8 Mon Sep 17 00:00:00 2001
From: Augustin Cavalier <waddlesplash@gmail.com>
Date: Thu, 10 Oct 2024 15:04:47 -0400
Subject: [PATCH] bootloader: Actually release the heap.

This avoids leaking the bootloader heap memory into the kernel.
Ideally it'd be dropped automatically, but in seems in many cases
it isn't (even on EFI).

Adjust platform logic to always remove the physical allocated range,
and ignore the return code and reuse the memory anyway if we can.
---
 src/system/boot/loader/heap.cpp             | 6 ++++++
 src/system/boot/loader/main.cpp             | 3 +--
 src/system/boot/platform/amiga_m68k/mmu.cpp | 4 ++--
 src/system/boot/platform/atari_m68k/mmu.cpp | 4 ++--
 src/system/boot/platform/bios_ia32/mmu.cpp  | 5 ++---
 5 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/system/boot/loader/heap.cpp b/src/system/boot/loader/heap.cpp
index 153e8fc332..c7ec5002cd 100644
--- a/src/system/boot/loader/heap.cpp
+++ b/src/system/boot/loader/heap.cpp
@@ -385,6 +385,8 @@ FreeChunk::SetToAllocated(void* allocated)
 void
 heap_release(stage2_args* args)
 {
+	heap_print_statistics();
+
 	LargeAllocation* allocation = sLargeAllocations.Clear(true);
 	while (allocation != NULL) {
 		LargeAllocation* next = allocation->HashNext();
@@ -393,6 +395,10 @@ heap_release(stage2_args* args)
 	}
 
 	platform_free_heap_region(sHeapBase, (addr_t)sHeapEnd - (addr_t)sHeapBase);
+
+	sHeapBase = sHeapEnd = NULL;
+	memset((void*)&sFreeChunkTree, 0, sizeof(sFreeChunkTree));
+	memset((void*)&sLargeAllocations, 0, sizeof(sLargeAllocations));
 }
 
 
diff --git a/src/system/boot/loader/main.cpp b/src/system/boot/loader/main.cpp
index dea89496a8..d4f739da57 100644
--- a/src/system/boot/loader/main.cpp
+++ b/src/system/boot/loader/main.cpp
@@ -162,8 +162,7 @@ main(stage2_args *args)
 			gKernelArgs.boot_volume_size = gBootVolume.ContentSize();
 
 			platform_cleanup_devices();
-			// TODO: cleanup, heap_release() etc.
-			heap_print_statistics();
+			heap_release(args);
 			platform_start_kernel();
 		}
 	}
diff --git a/src/system/boot/platform/amiga_m68k/mmu.cpp b/src/system/boot/platform/amiga_m68k/mmu.cpp
index 2f75b0ed4f..0f12363fbe 100644
--- a/src/system/boot/platform/amiga_m68k/mmu.cpp
+++ b/src/system/boot/platform/amiga_m68k/mmu.cpp
@@ -682,8 +682,8 @@ void
 platform_free_heap_region(void *_base, size_t size)
 {
 	addr_t base = (addr_t)_base;
-	status_t status = remove_physical_allocated_range(base, size);
-	if (status == B_OK && sNextPhysicalAddress == (base + size))
+	remove_physical_allocated_range(base, size);
+	if (sNextPhysicalAddress == (base + size))
 		sNextPhysicalAddress -= size;
 
 	// Failures don't matter very much as regions should be freed automatically,
diff --git a/src/system/boot/platform/atari_m68k/mmu.cpp b/src/system/boot/platform/atari_m68k/mmu.cpp
index 2c6b5e2202..438c9cdb94 100644
--- a/src/system/boot/platform/atari_m68k/mmu.cpp
+++ b/src/system/boot/platform/atari_m68k/mmu.cpp
@@ -687,8 +687,8 @@ void
 platform_free_heap_region(void *_base, size_t size)
 {
 	addr_t base = (addr_t)_base;
-	status_t status = remove_physical_allocated_range(base, size);
-	if (status == B_OK && sNextPhysicalAddress == (base + size))
+	remove_physical_allocated_range(base, size);
+	if (sNextPhysicalAddress == (base + size))
 		sNextPhysicalAddress -= size;
 
 	// Failures don't matter very much as regions should be freed automatically,
diff --git a/src/system/boot/platform/bios_ia32/mmu.cpp b/src/system/boot/platform/bios_ia32/mmu.cpp
index ef02f56cdf..1a580f44cc 100644
--- a/src/system/boot/platform/bios_ia32/mmu.cpp
+++ b/src/system/boot/platform/bios_ia32/mmu.cpp
@@ -829,10 +829,9 @@ void
 platform_free_heap_region(void *_base, size_t size)
 {
 	addr_t base = (addr_t)_base;
-	if (sNextPhysicalAddress == (base + size)) {
+	remove_physical_allocated_range(base, size);
+	if (sNextPhysicalAddress == (base + size))
 		sNextPhysicalAddress -= size;
-		remove_physical_allocated_range(sNextPhysicalAddress, size);
-	}
 
 	// Failures don't matter very much as regions should be freed automatically,
 	// since they're in the identity map and not stored in the kernel's page tables.