shamalloc.c

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

#include "include/shamalloc.h"

void
break_alloc(int32_t until_broken) {
  shamalloc_time_until_broken = until_broken;
}

#define COUNTDOWN_OR_BREAK\
  if (shamalloc_time_until_broken == 0)\
    return NULL;\
  if (shamalloc_time_until_broken > 0)\
    shamalloc_time_until_broken -= 1;\

// Might be to little on some systems.
#define BUFFER_SIZE (1024)

void *
malloc(size_t size) {
  COUNTDOWN_OR_BREAK

  static void *(*fptr)(size_t) = NULL;
  if (fptr == NULL) {
    // To bootstrap malloc we need to allow dlsym to allocate
    // some memory. On my system only a single call is needed.
    static uint8_t is_recusive_call = 0;
    if (is_recusive_call++) {
      static uint8_t bootstrap_memory[BUFFER_SIZE];
      // Check that we don't bootstrap the same memory twice, or
      // that we try to allocate too much memory
      if (is_recusive_call > 1 || size > BUFFER_SIZE) abort();
      return bootstrap_memory;
    }
    fptr = (void *(*)(size_t)) dlsym(RTLD_NEXT, "malloc");
    is_recusive_call = 0;
  }

  if (fptr == NULL) abort();
  return (*fptr)(size);
}

void *
calloc(size_t nmemb, size_t size) {
  COUNTDOWN_OR_BREAK

  static void *(*fptr)(size_t, size_t) = NULL;
  if (fptr == NULL) {
    // To bootstrap calloc we need to allow dlsym to allocate some memmory.
    static uint8_t calloc_is_recusive_call = 0;
    if (calloc_is_recusive_call++) {
      static uint8_t bootstrap_memory[BUFFER_SIZE];
      static uint8_t * bootstrap_pointer = bootstrap_memory;
      // Check that we don't bootstrap too much memmory we try to allocate too
      // much memory
      if (size + (bootstrap_memory - bootstrap_pointer) > BUFFER_SIZE) 
        abort();
      void * ret = bootstrap_pointer;
      bootstrap_pointer += size + (size % 8 && 8 - size % 8);
      return ret;
    }
    fptr = ((void *(*)(size_t, size_t)) dlsym(RTLD_NEXT, "calloc"));
    calloc_is_recusive_call = 0;
  }

  if (fptr == NULL) abort();
  return (*fptr)(nmemb, size);
}

void *
realloc(void * ptr, size_t size) {
  COUNTDOWN_OR_BREAK

  static void *(*fptr)(void*, size_t) = NULL;
  fptr = fptr ? fptr
         : ((void *(*)(void*, size_t)) dlsym(RTLD_NEXT, "realloc"));

  if (fptr == NULL) abort();
  return (*fptr)(ptr, size);
}