Add threshold to proxy lib to call system allocator

Add threshold to proxy lib to call system allocator
when a size is less than the given threshold.

Signed-off-by: Lukasz Dorau <lukasz.dorau@intel.com>

Author: ldorau

src/libumf.c

@@ -12,7 +12,9 @@
 #include "base_alloc_global.h"
 #include "memspace_internal.h"
 #include "provider_tracking.h"
+#include "utils_common.h"
 #include "utils_log.h"
+
 #if !defined(UMF_NO_HWLOC)
 #include "topology.h"
 #endif
@@ -39,6 +41,15 @@ void umfTearDown(void) {
         umfMemspaceLowestLatencyDestroy();
         umfDestroyTopology();
 #endif
+
+        if (utils_is_running_in_proxy_lib()) {
+            // We cannot destroy the TRACKER nor the base allocator
+            // when we are running in the proxy library,
+            // because it could lead to use-after-free in the program's unloader
+            // (for example _dl_fini() on Linux).
+            return;
+        }
+
         // make sure TRACKER is not used after being destroyed
         umf_memory_tracker_handle_t t = TRACKER;
         TRACKER = NULL;

src/provider/provider_tracking.c

@@ -103,16 +103,19 @@ umf_memory_pool_handle_t umfMemoryTrackerGetPool(const void *ptr) {
 
 umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
                                           umf_alloc_info_t *pAllocInfo) {
-    assert(ptr);
     assert(pAllocInfo);
 
+    if (ptr == NULL) {
+        return UMF_RESULT_ERROR_INVALID_ARGUMENT;
+    }
+
     if (TRACKER == NULL) {
-        LOG_ERR("tracker is not created");
+        LOG_ERR("tracker does not exist");
         return UMF_RESULT_ERROR_NOT_SUPPORTED;
     }
 
     if (TRACKER->map == NULL) {
-        LOG_ERR("tracker's map is not created");
+        LOG_ERR("tracker's map does not exist");
         return UMF_RESULT_ERROR_NOT_SUPPORTED;
     }
 
@@ -121,9 +124,8 @@ umf_result_t umfMemoryTrackerGetAllocInfo(const void *ptr,
     int found = critnib_find(TRACKER->map, (uintptr_t)ptr, FIND_LE,
                              (void *)&rkey, (void **)&rvalue);
     if (!found || (uintptr_t)ptr >= rkey + rvalue->size) {
-        LOG_WARN("pointer %p not found in the "
-                 "tracker, TRACKER=%p",
-                 ptr, (void *)TRACKER);
+        LOG_DEBUG("pointer %p not found in the tracker, TRACKER=%p", ptr,
+                  (void *)TRACKER);
         return UMF_RESULT_ERROR_INVALID_ARGUMENT;
     }
 

src/proxy_lib/proxy_lib.c

@@ -27,6 +27,12 @@
  * - _aligned_offset_recalloc()
  */
 
+#ifndef _WIN32
+#define _GNU_SOURCE
+#include <dlfcn.h>
+#undef _GNU_SOURCE
+#endif /* _WIN32 */
+
 #if (defined PROXY_LIB_USES_JEMALLOC_POOL)
 #include <umf/pools/pool_jemalloc.h>
 #define umfPoolManagerOps umfJemallocPoolOps
@@ -103,6 +109,21 @@ static umf_memory_pool_handle_t Proxy_pool = NULL;
 // it protects us from recursion in umfPool*()
 static __TLS int was_called_from_umfPool = 0;
 
+typedef void *(*system_aligned_alloc_t)(size_t alignment, size_t size);
+typedef void *(*system_calloc_t)(size_t nmemb, size_t size);
+typedef void (*system_free_t)(void *ptr);
+typedef void *(*system_malloc_t)(size_t size);
+typedef size_t (*system_malloc_usable_size_t)(void *ptr);
+typedef void *(*system_realloc_t)(void *ptr, size_t size);
+
+static system_aligned_alloc_t system_aligned_alloc;
+static system_calloc_t system_calloc;
+static system_free_t system_free;
+static system_malloc_t system_malloc;
+static system_malloc_usable_size_t system_malloc_usable_size;
+static system_realloc_t system_realloc;
+static size_t threshold_value = 0;
+
 /*****************************************************************************/
 /*** The constructor and destructor of the proxy library *********************/
 /*****************************************************************************/
@@ -149,16 +170,40 @@ void proxy_lib_create_common(void) {
         LOG_ERR("creating UMF pool manager failed");
         exit(-1);
     }
+
+    LOG_DEBUG("proxy pool created");
+
+#ifndef _WIN32
+    *((void **)(&system_aligned_alloc)) = dlsym(RTLD_NEXT, "aligned_alloc");
+    *((void **)(&system_calloc)) = dlsym(RTLD_NEXT, "calloc");
+    *((void **)(&system_free)) = dlsym(RTLD_NEXT, "free");
+    *((void **)(&system_malloc)) = dlsym(RTLD_NEXT, "malloc");
+    *((void **)(&system_malloc_usable_size)) =
+        dlsym(RTLD_NEXT, "malloc_usable_size");
+    *((void **)(&system_realloc)) = dlsym(RTLD_NEXT, "realloc");
+
+    if (system_aligned_alloc && system_calloc && system_free && system_malloc &&
+        system_malloc_usable_size && system_realloc) {
+        // TODO threshold_value will be read from the environment variable
+        threshold_value = 128;
+        LOG_DEBUG("all system hooks found, threshold_value = %zu",
+                  threshold_value);
+    } else {
+        LOG_WARN("system hooks NOT found");
+    }
+#endif
+
     // The UMF pool has just been created (Proxy_pool != NULL). Stop using
     // the linear allocator and start using the UMF pool allocator from now on.
+    LOG_DEBUG("proxy library initialized");
 }
 
 void proxy_lib_destroy_common(void) {
     if (utils_is_running_in_proxy_lib()) {
         // We cannot destroy 'Base_alloc_leak' nor 'Proxy_pool' nor 'OS_memory_provider',
         // because it could lead to use-after-free in the program's unloader
         // (for example _dl_fini() on Linux).
-        return;
+        goto fini_proxy_lib_destroy_common;
     }
 
     umf_memory_pool_handle_t pool = Proxy_pool;
@@ -168,6 +213,10 @@ void proxy_lib_destroy_common(void) {
     umf_memory_provider_handle_t provider = OS_memory_provider;
     OS_memory_provider = NULL;
     umfMemoryProviderDestroy(provider);
+    LOG_DEBUG("proxy library destroyed");
+
+fini_proxy_lib_destroy_common:
+    LOG_DEBUG("proxy library finalized");
 }
 
 /*****************************************************************************/
@@ -246,6 +295,10 @@ static inline size_t ba_leak_pool_contains_pointer(void *ptr) {
 /*****************************************************************************/
 
 void *malloc(size_t size) {
+    if (size < threshold_value) {
+        return system_malloc(size);
+    }
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolMalloc(Proxy_pool, size);
@@ -257,6 +310,10 @@ void *malloc(size_t size) {
 }
 
 void *calloc(size_t nmemb, size_t size) {
+    if ((nmemb * size) < threshold_value) {
+        return system_calloc(nmemb, size);
+    }
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolCalloc(Proxy_pool, nmemb, size);
@@ -276,15 +333,20 @@ void free(void *ptr) {
         return;
     }
 
-    if (Proxy_pool) {
+    if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
         if (umfPoolFree(Proxy_pool, ptr) != UMF_RESULT_SUCCESS) {
             LOG_ERR("umfPoolFree() failed");
-            assert(0);
         }
         return;
     }
 
-    assert(0);
+    if (threshold_value) {
+        system_free(ptr);
+        return;
+    }
+
+    LOG_ERR("free() failed: %p", ptr);
+
     return;
 }
 
@@ -303,18 +365,27 @@ void *realloc(void *ptr, size_t size) {
         return ba_leak_realloc(ptr, size, leak_pool_contains_pointer);
     }
 
-    if (Proxy_pool) {
+    if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
         was_called_from_umfPool = 1;
         void *new_ptr = umfPoolRealloc(Proxy_pool, ptr, size);
         was_called_from_umfPool = 0;
         return new_ptr;
     }
 
-    assert(0);
+    if (threshold_value) {
+        return system_realloc(ptr, size);
+    }
+
+    LOG_ERR("realloc() failed: %p", ptr);
+
     return NULL;
 }
 
 void *aligned_alloc(size_t alignment, size_t size) {
+    if (size < threshold_value) {
+        return system_aligned_alloc(alignment, size);
+    }
+
     if (!was_called_from_umfPool && Proxy_pool) {
         was_called_from_umfPool = 1;
         void *ptr = umfPoolAlignedMalloc(Proxy_pool, size, alignment);
@@ -330,19 +401,28 @@ size_t _msize(void *ptr) {
 #else
 size_t malloc_usable_size(void *ptr) {
 #endif
-
-    // a check to verify we are running the proxy library
+    // a check to verify if we are running the proxy library
     if (ptr == (void *)0x01) {
         return 0xDEADBEEF;
     }
 
-    if (!was_called_from_umfPool && Proxy_pool) {
+    if (ba_leak_pool_contains_pointer(ptr)) {
+        return 0; // unsupported in case of the ba_leak allocator
+    }
+
+    if (Proxy_pool && (umfPoolByPtr(ptr) == Proxy_pool)) {
         was_called_from_umfPool = 1;
         size_t size = umfPoolMallocUsableSize(Proxy_pool, ptr);
         was_called_from_umfPool = 0;
         return size;
     }
 
+    if (threshold_value) {
+        return system_malloc_usable_size(ptr);
+    }
+
+    LOG_ERR("malloc_usable_size() failed: %p", ptr);
+
     return 0; // unsupported in this case
 }