From ae7c5fdfdb359b3e9903c413ce1802576e679500 Mon Sep 17 00:00:00 2001
From: CTCaer <ctcaer@gmail.com>
Date: Thu, 25 Dec 2025 12:33:41 +0200
Subject: [PATCH] bdk: heap: add more safe guards

- Use magic for used regions
- Deduplicate calloc
Enable a print/abort in _heap_free() for debugging.
---
 bdk/mem/heap.c | 39 +++++++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 16 deletions(-)

diff --git a/bdk/mem/heap.c b/bdk/mem/heap.c
index ad707291..8187ffdd 100644
--- a/bdk/mem/heap.c
+++ b/bdk/mem/heap.c
@@ -1,6 +1,6 @@
 /*
  * Copyright (c) 2018 naehrwert
- * Copyright (c) 2018-2024 CTCaer
+ * Copyright (c) 2018-2025 CTCaer
  *
  * This program is free software; you can redistribute it and/or modify it
  * under the terms and conditions of the GNU General Public License,
@@ -19,6 +19,8 @@
 #include "heap.h"
 #include <gfx_utils.h>
 
+#define HEAP_USED_MAGIC 0x50414548 // "HEAP".
+
 heap_t _heap;
 
 static void _heap_create(void *start)
@@ -40,7 +42,7 @@ static void *_heap_alloc(u32 size)
 	if (!_heap.first)
 	{
 		node = (hnode_t *)_heap.start;
-		node->used = 1;
+		node->used = HEAP_USED_MAGIC;
 		node->size = size;
 		node->prev = NULL;
 		node->next = NULL;
@@ -85,7 +87,7 @@ static void *_heap_alloc(u32 size)
 				size += new_size;
 
 			node->size = size;
-			node->used = 1;
+			node->used = HEAP_USED_MAGIC;
 
 			return (void *)node + sizeof(hnode_t);
 		}
@@ -100,7 +102,7 @@ static void *_heap_alloc(u32 size)
 
 	// No unused node found, create a new one.
 	new_node = (hnode_t *)((void *)node + sizeof(hnode_t) + node->size);
-	new_node->used = 1;
+	new_node->used = HEAP_USED_MAGIC;
 	new_node->size = size;
 	new_node->prev = node;
 	new_node->next = NULL;
@@ -114,6 +116,14 @@ static void *_heap_alloc(u32 size)
 static void _heap_free(void *addr)
 {
 	hnode_t *node = (hnode_t *)(addr - sizeof(hnode_t));
+
+	// Check if heap owns this node.
+	if (addr < _heap.start || node->used != HEAP_USED_MAGIC)
+	{
+		//! BUGPRINTF("free error: addr %08p, used %08X!\n");
+		return;
+	}
+
 	node->used = 0;
 	node = _heap.first;
 
@@ -152,24 +162,21 @@ void *malloc(u32 size)
 	return _heap_alloc(size);
 }
 
-void *calloc(u32 num, u32 size)
-{
-	void *res = (void *)_heap_alloc(num * size);
-	memset(res, 0, ALIGN(num * size, sizeof(hnode_t))); // Clear the aligned size.
-	return res;
-}
-
 void *zalloc(u32 size)
 {
-	void *res = (void *)_heap_alloc(size);
-	memset(res, 0, ALIGN(size, sizeof(hnode_t))); // Clear the aligned size.
-	return res;
+	void *buf = (void *)_heap_alloc(size);
+	memset(buf, 0, ALIGN(size, sizeof(hnode_t))); // Clear the aligned size.
+	return buf;
+}
+
+void *calloc(u32 num, u32 size)
+{
+	return zalloc(num * size);
 }
 
 void free(void *buf)
 {
-	if (buf >= _heap.start)
-		_heap_free(buf);
+	_heap_free(buf);
 }
 
 void heap_monitor(heap_monitor_t *mon, bool print_node_stats)