How to hook malloc function on Linux?

I'm reading the Driving Compilers series of articles by Fabien Sanglard. In part 3 about The Compiler there is an example for hooking malloc function. First, an incorrect solution that falls in an infinite recursion is shown:

void* malloc(size_t sz) {
  void *(*libc_malloc)(size_t) = dlsym(RTLD_NEXT, "malloc");
  printf("malloced %zu bytes\n", sz);
  return libc_malloc(sz);
}

The reason for the recursion is that dlsym function internally calls malloc.

After that, a supposedly fixed solution is provided:

#include <stdio.h>
#include <dlfcn.h>

static void* (*real_malloc)(size_t) = nullptr;

void *malloc(size_t size) {
    if(!real_malloc)  {
      real_malloc = dlsym(RTLD_NEXT, "malloc");
    }

    printf("malloc(%d) = ", size);
    return real_malloc(size);
}

Except that the solution is not really fixed because it suffers from the same problem. If I rename the function from malloc to my_malloc for example then it will work, but it is no longer a hook because the other software uses malloc but not my_malloc. Is there any solution to the problem?

First, per 7.1.4 Use of library functions, paragraph 4 of the (draft) C11 standard:

> The functions in the standard library are not guaranteed to be reentrant and may modify objects with static or thread storage duration.

So you can't safely use any function from the C standard and be safe. You need to rely on platform-specific solutions.

First, you can find what functions your platform allows to be called from within signal handlers - these functions pretty much have to be reentrant.

For POSIX-based systems, which I'm assuming you're using because you use the POSIX function dlsym(), you can start with 2.4 Signal Concepts which has an extensive list.

Note that write() is on the list of async-signal-safe functions, but printf() is not.

So your code

static void* (*real_malloc)(size_t) = nullptr;

void *malloc(size_t size) {
    if(!real_malloc)  {
      real_malloc = dlsym(RTLD_NEXT, "malloc");
    }

    printf("malloc(%d) = ", size);
    return real_malloc(size);
}

can be replaced with

static void* (*real_malloc)(size_t) = nullptr;

void *malloc(size_t size) {
    if(!real_malloc)  {
      static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;

      pthread_mutex_lock( &mutex );
      if(!real_malloc)  {
        real_malloc = dlsym(RTLD_NEXT, "malloc");
      }
      pthread_mutex_unlock( &mutex );
    }

    write( STDOUT_FILENO, "malloc()", strlen( "malloc()" );
    return real_malloc(size);
}

Note that I omitted the address - that would need to be converted to a string. It's not hard to do, and examples on how to do that can easily be found, such as in the answers to How can I convert an int to a string in C?. Note that you can't use any answers there that use a standard library function that isn't async-signal-safe.

And if you're running on Solaris, s[n]printf() actually is async-signal-safe there.

I also added some protection for multithreaded use - there's a race condition in obtaining the pointer value to the actual malloc() that should be protected against, if only because if the pointer value gets corrupted you will likely never be able to reproduce whatever error(s) that causes.

EDIT

Per the comment from @ChrisDodd, fixed to address concerns about safety of dlsym():

static void* (*real_malloc)(size_t) = NULL;

__attribute__((constructor))
static void initValues(void) {
    real_malloc = dlsym(RTLD_NEXT, "malloc");
}

void *malloc(size_t size) {
    write( STDOUT_FILENO, "malloc()", strlen( "malloc()" );
    return real_malloc(size);
}

Note that the code inside the malloc() replacement is now much simpler - there's no race condition possible.

A little more investigation shows that what really calls malloc and leads to an infinite recursion is the printf function call, but not the dlsym as written in the article. This reduces the question to how to do the printing in such a way that this not happens. I came up with the following solution:

#include <stdio.h>
#include <dlfcn.h>

static void* (*real_malloc)(size_t) = NULL;

void* malloc(size_t size) {
  if(!real_malloc)  {
    real_malloc = dlsym(RTLD_NEXT, "malloc");
  }

  static char isPrintF = 0;
  if (isPrintF) {
    return real_malloc(size);
  }

  char* p = real_malloc(size);
  isPrintF = 1;
  printf("malloc(%zu) = %p\n", size, p);
  isPrintF = 0;
  return p;
}