Hook malloc

/*
 *  gcc -o hook main.c
 *  ./hook
 */
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <sys/mman.h>

typedef int (*foo_func)(int a);

foo_func ptr_foo = NULL;

int foo(int a)
{
    return a + 2;
}

int foo_hook(int a)
{
    printf("foo_hook - %d\n", a);
    a += 2;
    return ptr_foo(a);
}

#pragma pack(push, 1)
struct thunk
{
    char e9;
    unsigned int offset;
};
#pragma pack(pop)

struct thunk thunk_foo;

void test_malloc_hook();

int main()
{
    int ret;
    int foo_size;
    printf("hook!\n");
    printf("sizeof(struct thunk) =%ld\n", sizeof(struct thunk));
    printf("foo =%p\n", foo);
    printf("foo_hook =%p\n", foo_hook);

    foo_size = (int)((unsigned long)foo_hook - (unsigned long)foo);

    printf("foo_size =%d\n", foo_size);

    thunk_foo.e9 = 0xe9;
    thunk_foo.offset = foo_size - 5;

    printf("foo offset =%d\n", thunk_foo.offset);

    /*
     * make a copy of original foo
     */

    void *foo2 = mmap(NULL, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
                    MAP_ANON | MAP_PRIVATE, -1, 0);
    memcpy(foo2, foo, foo_size);
    ptr_foo = foo2;

    /*
     * test the copy foo2 works
     */
    printf("foo(10)=%d\n", foo(10));
    printf("ptr_foo(20)=%d\n", ptr_foo(20));

    /*
     * hook the original foo by thunk_foo which is a jump to foo_hook
     */
    void *page = (void *)((unsigned long)foo & ~0xfffUL);

    ret = mprotect(page, 4096, PROT_READ | PROT_WRITE | PROT_EXEC);
    if (ret < 0)
    {
        perror("mprotect failed!\n");
        exit(ret);
    }
    memcpy(foo, &thunk_foo, sizeof(struct thunk));
    printf("foo(10)=%d\n", foo(10));

    test_malloc_hook();

    return ret;
}

typedef int (*foo_func)(int a);

void af()
{
    printf("af\n");
    void *p = malloc(1024);
    free(p);
}

typedef void *(*malloc_func)(size_t size);

void malloc_hook_restore();
void malloc_hook_enable();

void *malloc_hook(size_t size)
{
    void *p;
    printf("malloc_hook- size = %ld\n", size);
    malloc_hook_restore();
    p = malloc(size);
    malloc_hook_enable();
    return p;
}

int malloc_size = 0;

void dump5(void *addr)
{
    unsigned char *p = addr;
    printf("%2.2X %2.2X %2.2X %2.2X %2.2X\n", p[0], p[1], p[2], p[3], p[4]);
}

unsigned char malloc5byte[5];
struct thunk thunk_malloc;
void malloc_hook_enable()
{
    /* 
     * save the first 5 bytes
     */

    memcpy(malloc5byte, malloc, 5);

    malloc_size = (int)((unsigned long)malloc_hook - (unsigned long)malloc);

    thunk_malloc.e9 = 0xe9;
    thunk_malloc.offset = malloc_size - 5;
    void *page = (void *)((unsigned long)malloc & ~0xfffUL);
    int ret = mprotect(page, 4096, PROT_READ | PROT_WRITE | PROT_EXEC);
    if (ret < 0)
    {
        perror("mprotect failed!\n");
        exit(ret);
    }
    memcpy(malloc, &thunk_malloc, 5);
}

void malloc_hook_restore()
{
    memcpy(malloc, malloc5byte, 5);
}

void test_malloc_hook()
{
    printf("test_malloc_hook\n");
    af();
    malloc_hook_enable();
    af();
    af();
    af();
}