From 8bd15bd272e426d0cf0c2ce9b3554a10cc23a1cc Mon Sep 17 00:00:00 2001
From: Luke Street <luke@street.dev>
Date: Mon, 27 Oct 2025 15:18:38 -0600
Subject: [PATCH] Fix heapapi ownership checks

`mi_heap_check_owned` does not work for allocations that span
multiple blocks; `mi_is_in_heap_region` is a less precise but
working way for us to check if a pointer is in our heap.

See https://github.com/microsoft/mimalloc/issues/298
---
 dll/kernel32/heapapi.cpp | 6 +++---
 test/test_heap.c         | 4 ++++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/dll/kernel32/heapapi.cpp b/dll/kernel32/heapapi.cpp
index 6a4f73b..b17415a 100644
--- a/dll/kernel32/heapapi.cpp
+++ b/dll/kernel32/heapapi.cpp
@@ -175,7 +175,7 @@ LPVOID WIN_FUNC HeapReAlloc(HANDLE hHeap, DWORD dwFlags, LPVOID lpMem, SIZE_T dw
 		VERBOSE_LOG("-> %p (alloc)\n", alloc);
 		return alloc;
 	}
-	if (!mi_heap_check_owned(record->heap, lpMem)) {
+	if (!mi_is_in_heap_region(lpMem)) {
 		VERBOSE_LOG("-> NULL (not owned)\n");
 		wibo::lastError = ERROR_INVALID_PARAMETER;
 		return nullptr;
@@ -244,7 +244,7 @@ SIZE_T WIN_FUNC HeapSize(HANDLE hHeap, DWORD dwFlags, LPCVOID lpMem) {
 		wibo::lastError = ERROR_INVALID_PARAMETER;
 		return static_cast<SIZE_T>(-1);
 	}
-	if (!mi_heap_check_owned(record->heap, lpMem)) {
+	if (!mi_is_in_heap_region(lpMem)) {
 		VERBOSE_LOG("-> ERROR_INVALID_PARAMETER (not owned)\n");
 		wibo::lastError = ERROR_INVALID_PARAMETER;
 		return static_cast<SIZE_T>(-1);
@@ -266,7 +266,7 @@ BOOL WIN_FUNC HeapFree(HANDLE hHeap, DWORD dwFlags, LPVOID lpMem) {
 		wibo::lastError = ERROR_INVALID_HANDLE;
 		return FALSE;
 	}
-	if (!mi_heap_check_owned(record->heap, lpMem)) {
+	if (!mi_is_in_heap_region(lpMem)) {
 		VERBOSE_LOG("-> ERROR_INVALID_PARAMETER (not owned)\n");
 		wibo::lastError = ERROR_INVALID_PARAMETER;
 		return FALSE;
diff --git a/test/test_heap.c b/test/test_heap.c
index 887e18f..c19c5a0 100644
--- a/test/test_heap.c
+++ b/test/test_heap.c
@@ -42,8 +42,12 @@ int main(void) {
     TEST_CHECK(privateBlock != NULL);
 
     SetLastError(0);
+    // Disabled temporarily; no good way to detect individual heap allocations
+    // in mimalloc currently. See https://github.com/microsoft/mimalloc/issues/298
+#if 0
     TEST_CHECK(!HeapFree(processHeap, 0, privateBlock));
     TEST_CHECK_EQ(ERROR_INVALID_PARAMETER, GetLastError());
+#endif
 
     TEST_CHECK(HeapFree(privateHeap, 0, privateBlock));
     TEST_CHECK(HeapDestroy(privateHeap));