Destroy global base allocator in the library destructor

The global base allocator should be destroyed in the library destructor
in order to avoid use-after-free.

Signed-off-by: Lukasz Dorau <[email protected]>

Author: ldorau

src/base_alloc/base_alloc.h

@@ -21,7 +21,9 @@ typedef struct umf_ba_pool_t umf_ba_pool_t;
 umf_ba_pool_t *umf_ba_create(size_t size);
 void *umf_ba_alloc(umf_ba_pool_t *pool);
 void umf_ba_free(umf_ba_pool_t *pool, void *ptr);
+void umf_ba_init(void);
 void umf_ba_destroy(umf_ba_pool_t *pool);
+void umf_ba_destroy_global(void);
 
 #ifdef __cplusplus
 }

src/base_alloc/base_alloc_global.c

@@ -17,17 +17,18 @@
 static umf_ba_pool_t *BA_pool = NULL;
 static UTIL_ONCE_FLAG ba_is_initialized = UTIL_ONCE_FLAG_INIT;
 
-static void umf_ba_destroy_global(void) {
-    assert(BA_pool);
-    umf_ba_destroy(BA_pool);
-    BA_pool = NULL;
-}
-
 static void umf_ba_create_global(void) {
+    umf_ba_init();
     assert(BA_pool == NULL);
     BA_pool = umf_ba_create(SIZE_BA_POOL_CHUNK);
+    assert(BA_pool);
+}
+
+void umf_ba_destroy_global(void) {
     if (BA_pool) {
-        atexit(umf_ba_destroy_global);
+        umf_ba_pool_t *pool = BA_pool;
+        BA_pool = NULL;
+        umf_ba_destroy(pool);
     }
 }
 

src/base_alloc/base_alloc_linux.c

@@ -13,6 +13,18 @@
 
 #include "base_alloc.h"
 
+// The highest possible priority (101) is used, because the constructor should be called
+// as the first one and the destructor as the last one in order to avoid use-after-free.
+void __attribute__((constructor(101))) umf_ba_constructor(void) {}
+
+void __attribute__((destructor(101))) umf_ba_destructor(void) {
+    umf_ba_destroy_global();
+}
+
+void umf_ba_init(void) {
+    // do nothing, additional initialization not needed
+}
+
 void *ba_os_alloc(size_t size) {
     return mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS,
                 -1, 0);

src/base_alloc/base_alloc_windows.c

@@ -5,8 +5,38 @@
  * SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 */
 
+#include "base_alloc.h"
 #include <windows.h>
 
+static void umf_ba_constructor(void) {}
+
+static void umf_ba_destructor(void) { umf_ba_destroy_global(); }
+
+#ifdef UMF_SHARED_LIBRARY
+
+// constructor and destructor are handled in libumf_windows.c
+
+void umf_ba_init(void) {
+    // do nothing, additional initialization not needed
+}
+
+#else /* UMF_SHARED_LIBRARY */
+
+INIT_ONCE BA_init_once_flag = INIT_ONCE_STATIC_INIT;
+
+BOOL CALLBACK BAinitOnceCb(PINIT_ONCE InitOnce, PVOID Parameter,
+                           PVOID *lpContext) {
+    umf_ba_constructor();
+    atexit(umf_ba_destructor);
+    return TRUE;
+}
+
+void umf_ba_init(void) {
+    InitOnceExecuteOnce(&BA_init_once_flag, BAinitOnceCb, NULL, NULL);
+}
+
+#endif /* UMF_SHARED_LIBRARY */
+
 void *ba_os_alloc(size_t size) {
     return VirtualAlloc(NULL, size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
 }

src/libumf_windows.c

@@ -17,7 +17,10 @@ umf_memory_tracker_handle_t TRACKER = NULL;
 
 static void umfCreate(void) { TRACKER = umfMemoryTrackerCreate(); }
 
-static void umfDestroy(void) { umfMemoryTrackerDestroy(TRACKER); }
+static void umfDestroy(void) {
+    umfMemoryTrackerDestroy(TRACKER);
+    umf_ba_destroy_global();
+}
 
 #if defined(UMF_SHARED_LIBRARY)
 BOOL APIENTRY DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved) {