memprof.c

#define _GNU_SOURCE
#include <err.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <unistd.h>
#include <sysexits.h>
#include <sys/mman.h>
#include <mach-o/dyld.h>
#include <mach-o/getsect.h>
#include <mach-o/loader.h>
#include <mach-o/ldsyms.h>
#include <ruby.h>
#include <intern.h>
#include <errno.h>
#include <dlfcn.h>

struct tramp_tbl_entry {
  unsigned char mov[2];
  long long addr;
  unsigned char callq[2];
  unsigned char ret;
  unsigned char pad[3];
} __attribute__((__packed__));;

/*
  trampoline specific stuff
*/
static struct tramp_tbl_entry *tramp_table = NULL;
static size_t tramp_size = 0;

static void
error_tramp() {
  printf("WARNING: NO TRAMPOLINE SET.\n");
  return;
}

static VALUE
newobj_tramp() {
  printf("source = %s, line = %d\n", ruby_sourcefile, ruby_sourceline);
  return rb_newobj();
}

static void
create_tramp_table() {
  int i, j = 0;

  struct tramp_tbl_entry ent = {
    .mov = {'\x48', '\xbb'},
    .addr = (long long)&error_tramp,
    .callq = { '\xff', '\xd3' },
    .ret = '\xc3',
    .pad =  { '\x90', '\x90', '\x90'},
  };

  int pagesize = 4096;

  // walk up the lower 32 bits of the address space for a page
  for (; i < INT_MAX - pagesize; i++) {
    tramp_table = mmap((void*)(NULL + i), pagesize, PROT_WRITE|PROT_READ|PROT_EXEC, MAP_ANON|MAP_PRIVATE, -1, 0);
    if (tramp_table != MAP_FAILED) {
      for (; j < pagesize/sizeof(struct tramp_tbl_entry); j ++ ) {
        memcpy(tramp_table + j, &ent, sizeof(struct tramp_tbl_entry));
      }
      return;
    }
  }
  printf("failed to allocate a page for stage1\n");
  exit(1);
}

static void
update_callqs(int entry, void *trampee_addr, void* text_segment, unsigned long text_segment_len) {
  char *byte = text_segment;
  size_t count = 0;
  int fn_addr = 0;
  void *aligned_addr = NULL;

  for(; count < text_segment_len; count++) {
    if (*byte == '\xe8') {
      fn_addr = *(int *)(byte+1);
      if (((void *)trampee_addr - (void *)(byte+5)) == fn_addr) {
        aligned_addr = (void*)(((long)byte+1)&~(0xffff));
        mprotect(aligned_addr, (((void *)byte+1) - aligned_addr) + 10, PROT_READ|PROT_WRITE|PROT_EXEC);
        *(int  *)(byte+1) = (uint32_t)((void *)(tramp_table + entry) - (void *)(byte + 5));
        mprotect(aligned_addr, (((void *)byte+1) - aligned_addr) + 10, PROT_READ|PROT_EXEC);
      }
    }
    byte++;
  }
}

static void
update_dyld_stubs(int entry, void *trampee_addr, void* text_segment, unsigned long text_segment_len) {
  char *byte = text_segment;
  size_t count = 0;

  for(; count < text_segment_len; count++) {
    if (*byte == '\xff') {
      int off = *(int *)(byte+2);
      if (trampee_addr == (void*)(*(long long*)(byte + 6 + off))) {
        // do we need the stage1 trampoline address here? can we even use it?
        // this seems to work fine as is. could that change depending on some conditions?..
        *(long long*)(byte + 6 + off) = tramp_table[entry].addr;
      }
    }
    byte++;
  }
}

static void *
find_symbol(char *sym) {
  void *ptr = NULL;
  _dyld_lookup_and_bind((const char*)sym, (void*) &ptr, NULL);
  return ptr;
}

static void
update_image(int entry, void *trampee_addr)
{
  void *text_segment = NULL;
  unsigned long text_segment_len = 0;

  if (trampee_addr) {
    Dl_info info;
    // Find the mach object header where the function address resides
    if (dladdr(trampee_addr, &info)) {
      struct mach_header_64 *header = (struct mach_header_64*) info.dli_fbase;
      // Is the function within the executable? if so, grab the text segment
      // and update the callq instructions
      if (header == &_mh_execute_header) {
        text_segment = getsectdatafromheader_64(header, "__TEXT", "__text", (uint64_t*)&text_segment_len);
        if (!text_segment) {
          printf("failed to get the text segment\n");
          exit(1);
        }
        update_callqs(entry, trampee_addr, text_segment, text_segment_len);
        return;
      }
      // otherwise the function resides in a shared library
      // we need to get segment with the dyld_stub functions to update their jump offsets
      text_segment = getsectdatafromheader_64(header, "__TEXT", "__symbol_stub1", (uint64_t*)&text_segment_len);
      if (!text_segment) {
        printf("failed to get the stub segment\n");
        exit(1);
      }
      // find the index of this image and use it to get the vmaddr_slide and make it a real pointer per the man page
      unsigned long cur_idx, real_idx = 0;
      unsigned long count = _dyld_image_count();
      for (cur_idx = 0; cur_idx < count; cur_idx++) {
        if ((struct mach_header_64*)_dyld_get_image_header(cur_idx) == header) {
          real_idx = cur_idx;
          break;
        }
      }
      text_segment = text_segment + _dyld_get_image_vmaddr_slide(real_idx);
      // update the stubs
      update_dyld_stubs(entry, trampee_addr, text_segment, text_segment_len);
      return;
    }
  }
  printf("U FUCKING FAIL BADLY\n");
  exit(1);
}

static void
insert_tramp(char *trampee, void *tramp) {
  void *trampee_addr = find_symbol(trampee);
  int entry = tramp_size;
  tramp_table[tramp_size].addr = (long long)tramp;
  tramp_size++;
  update_image(entry, trampee_addr);
}

void Init_memprof()
{
  create_tramp_table();
  insert_tramp("_rb_newobj", newobj_tramp);
}