From c37124800fe9ff9e4de274c232cd53b4d25ec640 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Mon, 9 Jan 2023 10:33:05 -0600
Subject: [PATCH] Don't use sbrk(0) to determine the initial heap size (#377)

* Don't use sbrk(0) to determine the initial heap size

This commit changes the `try_init_allocator` function as part of
dlmalloc to not use `sbrk(0)` to determine the initial heap size. The
purpose of this function is to use the extra memory at the end of linear
memory for the initial allocation heap before `memory.grow` is used to
allocate more memory. To learn the extent of this region the code
previously would use `sbrk(0)` to find the current size of linear
memory. This does not work, however, when other systems have called
`memory.grow` before this function is called. For example if another
allocator is used or if another component of a wasm binary grows memory
for its own purposes then that memory will be incorrectly claimed to be
owned by dlmalloc.

Instead this commit rounds up the `__heap_base` address to the nearest
page size, since that must be allocatable. Otherwise anything above this
rounded address is assumed to be used by something else, even if it's
addressable.

* Use `__heap_end` if defined

* Move mstate initialization earlier
---
 dlmalloc/src/malloc.c | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/dlmalloc/src/malloc.c b/dlmalloc/src/malloc.c
index 03da739e1..6e7a15c03 100644
--- a/dlmalloc/src/malloc.c
+++ b/dlmalloc/src/malloc.c
@@ -5214,26 +5214,43 @@ static void internal_inspect_all(mstate m,
 /* ------------------ Exported try_init_allocator -------------------- */
 
 /* Symbol marking the end of data, bss and explicit stack, provided by wasm-ld. */
-extern unsigned char __heap_base;
+extern char __heap_base;
+extern char __heap_end __attribute__((__weak__));
 
 /* Initialize the initial state of dlmalloc to be able to use free memory between __heap_base and initial. */
 static void try_init_allocator(void) {
   /* Check that it is a first-time initialization. */
   assert(!is_initialized(gm));
 
-  char *base = (char *)&__heap_base;
-  /* Calls sbrk(0) that returns the initial memory position. */
-  char *init = (char *)CALL_MORECORE(0);
-  int initial_heap_size = init - base;
+  /* Initialize mstate. */
+  ensure_initialization();
+
+  char *base = &__heap_base;
+  // Try to use the linker pseudo-symbol `__heap_end` for the initial size of
+  // the heap, but if that's not defined due to LLVM being too old perhaps then
+  // round up `base` to the nearest `PAGESIZE`. The initial size of linear
+  // memory will be at least the heap base to this page boundary, and it's then
+  // assumed that the initial linear memory image was truncated at that point.
+  // While this reflects the default behavior of `wasm-ld` it is also possible
+  // for users to craft larger linear memories by passing options to extend
+  // beyond this threshold. In this situation the memory will not be used for
+  // dlmalloc.
+  //
+  // Note that `sbrk(0)`, or in dlmalloc-ese `CALL_MORECORE(0)`, is specifically
+  // not used here. That captures the current size of the heap but is only
+  // correct if the we're the first to try to grow the heap. If the heap has
+  // grown elsewhere, such as a different allocator in place, then this would
+  // incorrectly claim such memroy as our own.
+  char *end = &__heap_end;
+  if (end == NULL)
+    end = (char*) page_align((size_t) base);
+  size_t initial_heap_size = end - base;
 
   /* Check that initial heap is long enough to serve a minimal allocation request. */
   if (initial_heap_size <= MIN_CHUNK_SIZE + TOP_FOOT_SIZE + MALLOC_ALIGNMENT) {
     return;
   }
 
-  /* Initialize mstate. */
-  ensure_initialization();
-
   /* Initialize the dlmalloc internal state. */
   gm->least_addr = base;
   gm->seg.base = base;