resolve real_malloc earlier for C++ programs

cmake (as run by the Ruby test suite for RubyGems) uses libjson
and libtasn1, which respectively call malloc (via `new') and
free before our constructor can even fire.  Apparently, C++
variable initialization may call "new" outside of any functions;
and those run before any functions with the GCC constructor
attribute.

Disclaimer: I don't know C++

2 files changed, 40 insertions, 10 deletions

diff --git a/ext/mwrap/mwrap.c b/ext/mwrap/mwrap.c
index d08ebf0..73d5a80 100644
--- a/ext/mwrap/mwrap.c
+++ b/ext/mwrap/mwrap.c
@@ -42,12 +42,11 @@ void *__malloc(size_t);
 void __free(void *);
 static void *(*real_malloc)(size_t) = __malloc;
 static void (*real_free)(void *) = __free;
-static const int ready = 1;
 #else
-static int ready;
 static void *(*real_malloc)(size_t);
 static void (*real_free)(void *);
 #endif /* !FreeBSD */
+static int resolving_malloc;
 
 /*
  * we need to fake an OOM condition while dlsym is running,
@@ -55,7 +54,7 @@ static void (*real_free)(void *);
  * symbol for the jemalloc calloc, yet
  */
 #  define RETURN_IF_NOT_READY() do { \
-        if (!ready) { \
+        if (!real_malloc) { \
                 errno = ENOMEM; \
                 return NULL; \
         } \
@@ -93,14 +92,16 @@ __attribute__((constructor)) static void resolve_malloc(void)
         int err;
 
 #ifndef __FreeBSD__
-        real_malloc = dlsym(RTLD_NEXT, "malloc");
+        if (!real_malloc) {
+                resolving_malloc = 1;
+                real_malloc = dlsym(RTLD_NEXT, "malloc");
+        }
         real_free = dlsym(RTLD_NEXT, "free");
         if (!real_malloc || !real_free) {
                 fprintf(stderr, "missing malloc/aligned_alloc/free\n"
                         "\t%p %p\n", real_malloc, real_free);
                 _exit(1);
         }
-        ready = 1;
 #endif
         totals = lfht_new();
         if (!totals)
@@ -343,6 +344,8 @@ void free(void *p)
 {
         if (p) {
                 struct alloc_hdr *h = ptr2hdr(p);
+
+                if (!real_free) return; /* oh well, leak a little */
                 if (h->as.live.loc) {
                         h->size = 0;
                         mutex_lock(h->as.live.loc->mtx);
@@ -400,7 +403,7 @@ internal_memalign(void **pp, size_t alignment, size_t size, uintptr_t caller)
         size_t d = alignment / sizeof(void*);
         size_t r = alignment % sizeof(void*);
 
-        if (!ready) return ENOMEM;
+        if (!real_malloc) return ENOMEM;
 
         if (r != 0 || d == 0 || !is_power_of_two(d))
                 return EINVAL;
@@ -498,11 +501,19 @@ void *malloc(size_t size)
         size_t asize;
         void *p;
 
-        if (__builtin_add_overflow(size, sizeof(struct alloc_hdr), &asize)) {
-                errno = ENOMEM;
-                return 0;
+        if (__builtin_add_overflow(size, sizeof(struct alloc_hdr), &asize))
+                goto enomem;
+
+        /*
+         * Needed for C++ global declarations using "new",
+         * which happens before our constructor
+         */
+        if (!real_malloc) {
+                if (resolving_malloc) goto enomem;
+                resolving_malloc = 1;
+                real_malloc = dlsym(RTLD_NEXT, "malloc");
         }
-        RETURN_IF_NOT_READY();
+
         rcu_read_lock();
         l = update_stats_rcu(size, RETURN_ADDRESS(0));
         p = h = real_malloc(asize);
@@ -513,6 +524,9 @@ void *malloc(size_t size)
         rcu_read_unlock();
         if (caa_unlikely(!p)) errno = ENOMEM;
         return p;
+enomem:
+        errno = ENOMEM;
+        return 0;
 }
 
 void *calloc(size_t nmemb, size_t size)
diff --git a/test/test_mwrap.rb b/test/test_mwrap.rb
index d76e4da..d0af0f7 100644
--- a/test/test_mwrap.rb
+++ b/test/test_mwrap.rb
@@ -61,6 +61,22 @@ class TestMwrap < Test::Unit::TestCase
     end
   end
 
+  def test_cmake
+    begin
+      exp = `cmake -h`
+    rescue Errno::ENOENT
+      warn 'cmake missing'
+      return
+    end
+    assert_not_predicate exp.strip, :empty?
+    env = @@env.merge('MWRAP' => 'dump_fd:1')
+    out = IO.popen(env, %w(cmake -h), &:read)
+    assert out.start_with?(exp), 'original help exists'
+    assert_not_equal exp, out, 'includes dump output'
+    dump = out.delete_prefix(exp)
+    assert_match(/\b0x[a-f0-9]+\b/s, dump, 'dump output has addresses')
+  end
+
   def test_clear
     cmd = @@cmd + %w(
       -e ("0"*10000).clear