memprof 0.2.9 → 0.3.0

data/.gitignore

@@ -4,3 +4,7 @@
 Makefile
 *.dSYM
 *.log
+*.gem
+ext/dst/*
+ext/src/yajl-1.0.9/*
+

data/Rakefile

@@ -7,3 +7,33 @@ task :spec do
   sh "ruby spec/memprof_spec.rb"
 end
 task :default => :spec
+
+# Should be used like:
+# rake --trace ci[1.8.7,shared]
+
+task :ci, [:ruby_type, :lib_option] do |t, args|
+  ruby_type, lib_option = args[:ruby_type], args[:lib_option]
+  raise "#{ruby_type} is not a supported ruby version" unless ["1.8.6", "1.8.7", "ree"].include?(ruby_type)
+  raise "#{lib_option} is not a supported " unless ["shared", "static"].include?(lib_option)
+
+  lib_option = case lib_option
+  when "static"
+    "--disable-shared"
+  when "shared"
+    "--enable-shared"
+  end
+
+  sh "/usr/bin/env bash -c \"
+  source ~/.rvm/scripts/rvm &&
+  rvm install #{ruby_type} --reconfigure -C #{lib_option} &&
+  rvm #{ruby_type} --symlink memprof &&
+  memprof_gem install bacon\""
+
+  Dir.chdir('ext') do
+    sh '/usr/bin/env bash -c "make clean"' rescue nil
+    sh "~/.rvm/bin/memprof_ruby extconf.rb"
+    sh '/usr/bin/env bash -c "make"'
+  end
+  sh "~/.rvm/bin/memprof_ruby spec/memprof_spec.rb"
+end
+

data/bin/memprof

@@ -0,0 +1,158 @@
+#!/usr/bin/env ruby
+
+require 'rubygems'
+require 'optparse'
+require 'restclient'
+require 'term/ansicolor'
+
+MEMPROF_URL = "https://memprof.com/upload"
+
+MEMPROF_BANNER = <<-EOF
+Memprof Uploader
+http://www.memprof.com
+======================
+
+EOF
+
+class MemprofUploader
+  include Term::ANSIColor
+
+  def initialize(args=ARGV)
+    puts MEMPROF_BANNER
+
+    @parser = OptionParser.new do |opts|
+      opts.banner = "Usage:"
+      opts.on("-p", "--pid <pid>", Integer, "PID of the process to dump       (required)")      {|arg| @pid = arg }
+      opts.on("-n", "--name <name>",        "Name for your dump               (required)")      {|arg| @name = arg }
+      opts.on("-k", "--key <key>",          "Memprof.com API key              (required)")      {|arg| @key = arg }
+      opts.on("-d", "--[no-]delete",        "Delete dump file after uploading (default true)")  {|arg| @delete_dump = arg }
+      opts.on("-s", "--seconds <seconds>",  "Seconds to wait for the dump     (default 60)")     {|arg| @secs = arg}
+      opts.on("-t", "--[no-]test",          "Test run (don't actually upload) (default false)") {|arg| @test = arg }
+      opts.on("-f", "--file <path>",        "Upload specific json dump        (optional)")      {|arg| @file = arg }
+      opts.on("--put-my-data-on-the-internet",  "Confirm that you understand\n" +
+                                                "memprof.com will show all your            \n".rjust(80) +
+                                                "internal data on the internet    (required)".rjust(80)) {|arg| @confirmed = true}
+    end
+
+    begin
+      @parser.parse!
+    rescue Exception => e
+      if e.kind_of?(SystemExit)
+        raise e
+      else
+        fail_with(e.message)
+      end
+    end
+
+    # Make this default to true if the user didn't pass any preference.
+    @delete_dump = true if @delete_dump.nil?
+
+    # Make this default to 60 if the user didn't pass the number of seconds.
+    @secs = 60 if @secs.nil?
+
+    if @file
+      fail_with("File not found: #{@file}") unless File.exists?(@file)
+    end
+
+    if @pid.nil? and @file.nil?
+      fail_with("Missing PID! (-p PID)")
+    elsif @name.nil? || @name.empty?
+      fail_with("Missing name! (-n 'my application')")
+    elsif @key.nil? || @key.empty?
+      fail_with("Missing API key! (-k KEY)")
+    elsif !@confirmed
+      fail_with("\nERROR: You MUST CONFIRM that you understand ALL YOUR CODE "+
+                "WILL BE PUBLICLY \nAVAILABLE ON THE INTERNET by passing " +
+                "--put-my-data-on-the-internet", true)
+    end
+  end
+
+  def run!
+    dump_filename = @file ? compress_file(@file, "/tmp/#{File.basename @file}.gz") : compress_file(get_dump_filename(@pid))
+
+    begin
+      upload_dump(dump_filename, @name, @key)
+    ensure
+      if @delete_dump
+        File.delete(dump_filename)
+        puts "\nDeleted dump file."
+      end
+    end
+
+    puts "Finished!"
+  end
+
+  private
+
+  def compress_file(filename, output=nil)
+    puts "Compressing with gzip..."
+    output ||= "#{filename}.gz"
+    `gzip -c '#{filename}' > '#{output}'`
+    output
+  end
+
+  def get_dump_filename(pid)
+    # Get a listing of files before we signal for the new dump
+    # We'll compare against this later to see which is the new file.
+    old_files = Dir.glob("/tmp/memprof-#{pid}-*")
+    signal_process(pid)
+    puts "Waiting 5 seconds for process #{pid} to create a new dump..."
+
+    file = nil
+
+    5.times do |i|
+      sleep 1 unless file = (Dir.glob("/tmp/memprof-#{pid}-*") - old_files).first and break
+      fail_with("Timed out after waiting 5 seconds. Make sure you added require '#{File.expand_path('../../lib/memprof/signal', __FILE__)}' to your application.") if i+1 == 5
+    end
+
+    puts "\nFound file #{file}"
+
+    if file =~ /\.IN_PROGRESS/
+      file = file.sub(/\.IN_PROGRESS/, "")
+      puts "Dump in progress. Waiting #{@secs} seconds for it to complete..."
+      @secs.times do |i|
+        sleep 1 unless File.exist?(file) and break
+        fail_with("Timed out after waiting #{@secs} seconds") if i+1 == @secs
+      end
+    end
+
+    file
+  end
+
+  def upload_dump(filename, name, key)
+    return if @test
+
+    file = File.open(filename, "r")
+
+    puts "\nUploading to memprof.com..."
+    response = RestClient.post(MEMPROF_URL, :upload => file, :name => name, :key => key)
+
+    puts "Response from server:"
+    p response.to_s
+  end
+
+  def fail_with(str, yell=true)
+    puts @parser.to_s
+    puts
+    if yell
+      print red, bold, str, reset
+      puts
+    else
+      puts "\n" + str
+    end
+    puts
+    exit(1)
+  end
+
+  def signal_process(pid)
+    begin
+      Process.kill("URG", pid)
+    rescue Errno::ESRCH
+      fail_with("No such process #{pid}!")
+    end
+    puts "Signaled process #{pid} with SIGURG"
+  end
+
+end
+
+MemprofUploader.new(ARGV).run!

data/ext/bin_api.h

@@ -21,6 +21,8 @@ bin_init();
  * Given:
  *  - sym - a symbol name
  *  - size - optional out parameter
+ *  - search_libs - 0 to search only ruby/libruby, 1 to search other
+ *    libraries, too.
  *
  * This function will search for the symbol sym and return its address if
  * found, or NULL if the symbol could not be found.
@@ -29,7 +31,7 @@ bin_init();
  * size of the symbol.
  */
 void *
-bin_find_symbol(const char *sym, size_t *size);
+bin_find_symbol(const char *sym, size_t *size, int search_libs);
 
 /*
  * bin_find_symbol_name - find a symbol's name
@@ -85,6 +87,8 @@ bin_type_member_offset(const char *type, const char *member);
  * Given:
  *  - trampee - the name of the symbol to hook
  *  - tramp - the stage 2 trampoline entry
+ *  - orig_func - out parameter storing the address of the function that was
+ *    originally being called.
  *
  * this function will update the binary image so that all calls to trampee will
  * be routed to tramp.
@@ -92,5 +96,5 @@ bin_type_member_offset(const char *type, const char *member);
  * Returns 0 on success.
  */
 int
-bin_update_image(const char *trampee, struct tramp_st2_entry *tramp);
+bin_update_image(const char *trampee, struct tramp_st2_entry *tramp, void **orig_func);
 #endif

data/ext/elf.c

@@ -90,6 +90,7 @@ typedef enum {
  * entry.
  */
 typedef linkmap_cb_status (*linkmap_cb)(struct link_map *, void *);
+typedef void (*linkmap_lib_cb)(struct elf_info *lib, void *data);
 
 static void open_elf(struct elf_info *info);
 static int dissect_elf(struct elf_info *info, int find_debug);
@@ -140,27 +141,36 @@ get_got_addr(struct plt_entry *plt)
 /*
  * overwrite_got - given the address of a PLT entry, overwrite the address
  * in the GOT that the PLT entry uses with the address in tramp.
+ *
+ * returns the original function address
  */
-static void
+static void *
 overwrite_got(void *plt, const void *tramp)
 {
   assert(plt != NULL);
   assert(tramp != NULL);
+  void *ret = NULL;
+
+  memcpy(&ret, get_got_addr(plt), sizeof(ret));
   memcpy(get_got_addr(plt), &tramp, sizeof(void *));
-  return;
+  return ret;
 }
 
 struct plt_hook_data {
   const char *sym;
-  const void *tramp;
+  void *addr;
+};
+
+struct dso_iter_data {
+  linkmap_lib_cb cb;
+  void *passthru;
 };
 
 static linkmap_cb_status
-hook_func_cb(struct link_map *map, void *data)
+for_each_dso_cb(struct link_map *map, void *data)
 {
-  struct plt_hook_data *hook_data = data;
+  struct dso_iter_data *iter_data = data;
   struct elf_info curr_lib;
-  void *trampee_addr = NULL;
 
   /* skip a few things we don't care about */
   if (strstr(map->l_name, "linux-vdso")) {
@@ -203,10 +213,7 @@ hook_func_cb(struct link_map *map, void *data)
   dbg_printf("dissected the elf file: %s, base: %lx\n",
       curr_lib.filename, (unsigned long)curr_lib.base_addr);
 
-  if ((trampee_addr = find_plt_addr(hook_data->sym, &curr_lib)) != NULL) {
-    dbg_printf("found: %s @ %p\n", hook_data->sym, trampee_addr);
-    overwrite_got(trampee_addr, hook_data->tramp);
-  }
+  iter_data->cb(&curr_lib, iter_data->passthru);
 
   elf_end(curr_lib.elf);
   close(curr_lib.fd);
@@ -214,12 +221,26 @@ hook_func_cb(struct link_map *map, void *data)
 }
 
 static void
-hook_required_objects(void *tramp)
+for_each_dso(linkmap_lib_cb cb, void *passthru)
 {
-  struct plt_hook_data data;
-  data.sym = "rb_newobj";
-  data.tramp = tramp;
-  walk_linkmap(hook_func_cb, &data);
+  struct dso_iter_data data;
+  data.cb = cb;
+  data.passthru = passthru;
+  walk_linkmap(for_each_dso_cb, &data);
+}
+
+static void
+hook_required_objects(struct elf_info *info, void *data)
+{
+  assert(info != NULL);
+  struct plt_hook_data *hook_data = data;
+  void *trampee_addr = NULL;
+
+  if ((trampee_addr = find_plt_addr(hook_data->sym, info)) != NULL) {
+    dbg_printf("found: %s @ %p\n", hook_data->sym, trampee_addr);
+    overwrite_got(trampee_addr, hook_data->addr);
+  }
+
   return;
 }
 
@@ -320,6 +341,8 @@ find_plt_addr(const char *symname, struct elf_info *info)
     info = ruby_info;
   }
 
+  assert(info != NULL);
+
   /* Search through each of the .rela.plt entries */
   for (i = 0; i < info->relplt_count; i++) {
     GElf_Rela rela;
@@ -365,38 +388,78 @@ find_plt_addr(const char *symname, struct elf_info *info)
 static void *
 do_bin_find_symbol(const char *sym, size_t *size, struct elf_info *elf)
 {
-  char *name = NULL;
+  const char *name = NULL;
 
   assert(sym != NULL);
   assert(elf != NULL);
 
-  assert(elf->symtab_data != NULL);
-  assert(elf->symtab_data->d_buf != NULL);
+  ElfW(Sym) *esym, *lastsym = NULL;
+  if (elf->symtab_data && elf->symtab_data->d_buf) {
+    dbg_printf("checking symtab....\n");
+    esym = (ElfW(Sym)*) elf->symtab_data->d_buf;
+    lastsym = (ElfW(Sym)*) ((char*) elf->symtab_data->d_buf + elf->symtab_data->d_size);
 
-  ElfW(Sym) *esym = (ElfW(Sym)*) elf->symtab_data->d_buf;
-  ElfW(Sym) *lastsym = (ElfW(Sym)*) ((char*) elf->symtab_data->d_buf + elf->symtab_data->d_size);
+    assert(esym <= lastsym);
 
-  assert(esym <= lastsym);
+    for (; esym < lastsym; esym++){
+      /* ignore numeric/empty symbols */
+      if ((esym->st_value == 0) ||
+          (ELF32_ST_BIND(esym->st_info)== STB_NUM))
+        continue;
 
-  for (; esym < lastsym; esym++){
-    /* ignore weak/numeric/empty symbols */
-    if ((esym->st_value == 0) ||
-        (ELF32_ST_BIND(esym->st_info)== STB_WEAK) ||
-        (ELF32_ST_BIND(esym->st_info)== STB_NUM))
-      continue;
+      name = elf_strptr(elf->elf, elf->symtab_shdr.sh_link, (size_t)esym->st_name);
 
-    name = elf_strptr(elf->elf, elf->symtab_shdr.sh_link, (size_t)esym->st_name);
+      if (name && strcmp(name, sym) == 0) {
+        if (size) {
+          *size = esym->st_size;
+        }
+        dbg_printf("Found symbol: %s in symtab\n", sym);
+        return elf->base_addr + (void *)esym->st_value;
+      }
+    }
+  }
+
+  if (elf->dynsym && elf->dynsym->d_buf) {
+    dbg_printf("checking dynsymtab....\n");
+    esym = (ElfW(Sym) *) elf->dynsym->d_buf;
+    lastsym = (ElfW(Sym) *) ((char *) elf->dynsym->d_buf + elf->dynsym->d_size);
+
+    for (; esym < lastsym; esym++){
+      /* ignore numeric/empty symbols */
+      if ((esym->st_value == 0) ||
+          (ELF32_ST_BIND(esym->st_info)== STB_NUM))
+        continue;
 
-    if (name && strcmp(name, sym) == 0) {
-      if (size) {
-        *size = esym->st_size;
+      name = elf->dynstr + esym->st_name;
+
+      if (name && strcmp(name, sym) == 0) {
+        if (size) {
+          *size = esym->st_size;
+        }
+        dbg_printf("Found symbol: %s in dynsym\n", sym);
+        return elf->base_addr + (void *)esym->st_value;
       }
-      return elf->base_addr + (void *)esym->st_value;
     }
   }
+
+  dbg_printf("Couldn't find symbol: %s in dynsym\n", sym);
   return NULL;
 }
 
+static void
+find_symbol_cb(struct elf_info *info, void *data)
+{
+  assert(info != NULL);
+  void *trampee_addr = NULL;
+  struct plt_hook_data *hook_data = data;
+  void *ret = do_bin_find_symbol(hook_data->sym, NULL, info);
+  if (ret) {
+    hook_data->addr = ret;
+    dbg_printf("found %s @ %p, fn addr: %p\n", hook_data->sym, trampee_addr,
+        hook_data->addr);
+  }
+}
+
 /*
  * bin_find_symbol - find the address of a given symbol and set its size if
  * desired.
@@ -404,9 +467,20 @@ do_bin_find_symbol(const char *sym, size_t *size, struct elf_info *elf)
  * This function is just a wrapper for the internal symbol lookup function.
  */
 void *
-bin_find_symbol(const char *sym, size_t *size)
+bin_find_symbol(const char *sym, size_t *size, int search_libs)
 {
-  return do_bin_find_symbol(sym, size, ruby_info);
+  void *ret = do_bin_find_symbol(sym, size, ruby_info);
+
+  if (!ret && search_libs) {
+    dbg_printf("Didn't find symbol: %s in ruby, searching other libs\n", sym);
+    struct plt_hook_data hd;
+    hd.sym = sym;
+    hd.addr = NULL;
+    for_each_dso(find_symbol_cb, &hd);
+    ret = hd.addr;
+  }
+
+  return ret;
 }
 
 /*
@@ -467,6 +541,8 @@ bin_find_symbol_name(void *sym) {
  * Given -
  *  trampee - the name of the symbol to hook
  *  tramp - the stage 2 trampoline entry
+ *  orig_func - out parameter storing the address of the function that was
+ *  originally called.
  *
  * This function will update the ruby binary image so that all calls to trampee
  * will be routed to tramp.
@@ -474,7 +550,7 @@ bin_find_symbol_name(void *sym) {
  * Returns 0 on success
  */
 int
-bin_update_image(const char *trampee, struct tramp_st2_entry *tramp)
+bin_update_image(const char *trampee, struct tramp_st2_entry *tramp, void **orig_func)
 {
   void *trampee_addr = NULL;
 
@@ -482,32 +558,56 @@ bin_update_image(const char *trampee, struct tramp_st2_entry *tramp)
   assert(tramp != NULL);
   assert(tramp->addr != NULL);
 
-  if (!libruby) {
+  /* first check if the symbol is in the PLT */
+  trampee_addr = find_plt_addr(trampee, NULL);
+
+  /* it isn't in the PLT, try to find it in the binary itself */
+  if (!trampee_addr) {
+    dbg_printf("Couldn't find %s in the PLT...\n", trampee);
     unsigned char *byte = ruby_info->text_segment;
-    trampee_addr = bin_find_symbol(trampee, NULL);
+    trampee_addr = bin_find_symbol(trampee, NULL, 0);
     size_t count = 0;
     int num = 0;
 
     assert(byte != NULL);
-    assert(trampee_addr != NULL);
+
+    if (!trampee_addr) {
+      dbg_printf("WARNING: Couldn't find: %s anywhere, so not tramping!\n", trampee);
+      return 0;
+    }
+
+    if (orig_func) {
+      *orig_func = trampee_addr;
+    }
 
     for(; count < ruby_info->text_segment_len; byte++, count++) {
       if (arch_insert_st1_tramp(byte, trampee_addr, tramp) == 0) {
         num++;
       }
     }
+    dbg_printf("Inserted %d tramps for: %s\n", num, trampee);
   } else {
-    trampee_addr = find_plt_addr(trampee, NULL);
-    assert(trampee_addr != NULL);
-    overwrite_got(trampee_addr, tramp->addr);
-  }
+    void *ret = NULL;
+    dbg_printf("Found %s in the PLT, inserting tramp...\n", trampee);
+    ret = overwrite_got(trampee_addr, tramp->addr);
 
-  if (strcmp("rb_newobj", trampee) == 0) {
-    dbg_printf("Trying to hook rb_newobj in other libraries...\n");
-    hook_required_objects(tramp->addr);
-    dbg_printf("Done searching other libraries for rb_newobj!\n");
+    assert(ret != NULL);
+
+    if (orig_func) {
+      *orig_func = ret;
+      dbg_printf("setting orig function: %p\n", *orig_func);
+    }
   }
 
+  dbg_printf("Trying to hook %s in other libraries...\n", trampee);
+
+  struct plt_hook_data data;
+  data.addr = tramp->addr;
+  data.sym = trampee;
+  for_each_dso(hook_required_objects, &data);
+
+  dbg_printf("Done searching other libraries for %s\n", trampee);
+
   return 0;
 }