Fix memory issues causing EXC_BAD_ACCESS

Fix grabbed from PR facebook#84

Author: Pedro A Osuna

fishhook.c

@@ -35,6 +35,7 @@
 #include <mach-o/dyld.h>
 #include <mach-o/loader.h>
 #include <mach-o/nlist.h>
+#include <pthread.h>
 
 #ifdef __LP64__
 typedef struct mach_header_64 mach_header_t;
@@ -54,13 +55,18 @@ typedef struct nlist nlist_t;
 #define SEG_DATA_CONST  "__DATA_CONST"
 #endif
 
+#ifndef SEG_AUTH_CONST
+#define SEG_AUTH_CONST  "__AUTH_CONST"
+#endif
+
 struct rebindings_entry {
   struct rebinding *rebindings;
   size_t rebindings_nel;
   struct rebindings_entry *next;
 };
 
 static struct rebindings_entry *_rebindings_head;
+static pthread_mutex_t mutex = PTHREAD_MUTEX_INITIALIZER;
 
 static int prepend_rebindings(struct rebindings_entry **rebindings_head,
                               struct rebinding rebindings[],
@@ -90,7 +96,7 @@ static vm_prot_t get_protection(void *sectionStart) {
   mach_msg_type_number_t count = VM_REGION_BASIC_INFO_COUNT_64;
   vm_region_basic_info_data_64_t info;
   kern_return_t info_ret = vm_region_64(
-      task, &address, &size, VM_REGION_BASIC_INFO_64, (vm_region_info_64_t)&info, &count, &object);
+                                        task, &address, &size, VM_REGION_BASIC_INFO_64, (vm_region_info_64_t)&info, &count, &object);
 #else
   mach_msg_type_number_t count = VM_REGION_BASIC_INFO_COUNT;
   vm_region_basic_info_data_t info;
@@ -109,17 +115,23 @@ static void perform_rebinding_with_section(struct rebindings_entry *rebindings,
                                            char *strtab,
                                            uint32_t *indirect_symtab) {
   const bool isDataConst = strcmp(section->segname, SEG_DATA_CONST) == 0;
+  const bool isAuthConst = strcmp(section->segname, SEG_AUTH_CONST) == 0;
   uint32_t *indirect_symbol_indices = indirect_symtab + section->reserved1;
   void **indirect_symbol_bindings = (void **)((uintptr_t)slide + section->addr);
   vm_prot_t oldProtection = VM_PROT_READ;
-  if (isDataConst) {
-    oldProtection = get_protection(rebindings);
-    mprotect(indirect_symbol_bindings, section->size, PROT_READ | PROT_WRITE);
+  vm_size_t trunc_address = (vm_size_t)indirect_symbol_bindings;
+  vm_size_t trunc_size = 0;
+  if (isDataConst || isAuthConst) {
+    trunc_address = trunc_page((vm_size_t)indirect_symbol_bindings);
+    trunc_size =(vm_size_t)indirect_symbol_bindings -trunc_address;
+    pthread_mutex_lock(&mutex);
+    oldProtection = get_protection((void *)trunc_address);
+    mprotect((void *)trunc_address, section->size+trunc_size, PROT_READ | PROT_WRITE);
   }
   for (uint i = 0; i < section->size / sizeof(void *); i++) {
     uint32_t symtab_index = indirect_symbol_indices[i];
     if (symtab_index == INDIRECT_SYMBOL_ABS || symtab_index == INDIRECT_SYMBOL_LOCAL ||
-        symtab_index == (INDIRECT_SYMBOL_LOCAL   | INDIRECT_SYMBOL_ABS)) {
+        symtab_index == (INDIRECT_SYMBOL_LOCAL | INDIRECT_SYMBOL_ABS)) {
       continue;
     }
     uint32_t strtab_offset = symtab[symtab_index].n_un.n_strx;
@@ -142,7 +154,7 @@ static void perform_rebinding_with_section(struct rebindings_entry *rebindings,
     }
   symbol_loop:;
   }
-  if (isDataConst) {
+  if (isDataConst || isAuthConst) {
     int protection = 0;
     if (oldProtection & VM_PROT_READ) {
       protection |= PROT_READ;
@@ -153,7 +165,8 @@ static void perform_rebinding_with_section(struct rebindings_entry *rebindings,
     if (oldProtection & VM_PROT_EXECUTE) {
       protection |= PROT_EXEC;
     }
-    mprotect(indirect_symbol_bindings, section->size, protection);
+    mprotect((void *)trunc_address, section->size+trunc_size, protection);
+    pthread_mutex_unlock(&mutex);
   }
 }
 
@@ -164,12 +177,12 @@ static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
   if (dladdr(header, &info) == 0) {
     return;
   }
-
+  
   segment_command_t *cur_seg_cmd;
   segment_command_t *linkedit_segment = NULL;
   struct symtab_command* symtab_cmd = NULL;
   struct dysymtab_command* dysymtab_cmd = NULL;
-
+  
   uintptr_t cur = (uintptr_t)header + sizeof(mach_header_t);
   for (uint i = 0; i < header->ncmds; i++, cur += cur_seg_cmd->cmdsize) {
     cur_seg_cmd = (segment_command_t *)cur;
@@ -183,31 +196,32 @@ static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
       dysymtab_cmd = (struct dysymtab_command*)cur_seg_cmd;
     }
   }
-
+  
   if (!symtab_cmd || !dysymtab_cmd || !linkedit_segment ||
       !dysymtab_cmd->nindirectsyms) {
     return;
   }
-
+  
   // Find base symbol/string table addresses
   uintptr_t linkedit_base = (uintptr_t)slide + linkedit_segment->vmaddr - linkedit_segment->fileoff;
   nlist_t *symtab = (nlist_t *)(linkedit_base + symtab_cmd->symoff);
   char *strtab = (char *)(linkedit_base + symtab_cmd->stroff);
-
+  
   // Get indirect symbol table (array of uint32_t indices into symbol table)
   uint32_t *indirect_symtab = (uint32_t *)(linkedit_base + dysymtab_cmd->indirectsymoff);
-
+  
   cur = (uintptr_t)header + sizeof(mach_header_t);
   for (uint i = 0; i < header->ncmds; i++, cur += cur_seg_cmd->cmdsize) {
     cur_seg_cmd = (segment_command_t *)cur;
     if (cur_seg_cmd->cmd == LC_SEGMENT_ARCH_DEPENDENT) {
       if (strcmp(cur_seg_cmd->segname, SEG_DATA) != 0 &&
-          strcmp(cur_seg_cmd->segname, SEG_DATA_CONST) != 0) {
+          strcmp(cur_seg_cmd->segname, SEG_DATA_CONST) != 0 &&
+          strcmp(cur_seg_cmd->segname, SEG_AUTH_CONST) != 0) {
         continue;
       }
       for (uint j = 0; j < cur_seg_cmd->nsects; j++) {
         section_t *sect =
-          (section_t *)(cur + sizeof(segment_command_t)) + j;
+        (section_t *)(cur + sizeof(segment_command_t)) + j;
         if ((sect->flags & SECTION_TYPE) == S_LAZY_SYMBOL_POINTERS) {
           perform_rebinding_with_section(rebindings, sect, slide, symtab, strtab, indirect_symtab);
         }
@@ -221,21 +235,21 @@ static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
 
 static void _rebind_symbols_for_image(const struct mach_header *header,
                                       intptr_t slide) {
-    rebind_symbols_for_image(_rebindings_head, header, slide);
+  rebind_symbols_for_image(_rebindings_head, header, slide);
 }
 
 int rebind_symbols_image(void *header,
                          intptr_t slide,
                          struct rebinding rebindings[],
                          size_t rebindings_nel) {
-    struct rebindings_entry *rebindings_head = NULL;
-    int retval = prepend_rebindings(&rebindings_head, rebindings, rebindings_nel);
-    rebind_symbols_for_image(rebindings_head, (const struct mach_header *) header, slide);
-    if (rebindings_head) {
-      free(rebindings_head->rebindings);
-    }
-    free(rebindings_head);
-    return retval;
+  struct rebindings_entry *rebindings_head = NULL;
+  int retval = prepend_rebindings(&rebindings_head, rebindings, rebindings_nel);
+  rebind_symbols_for_image(rebindings_head, (const struct mach_header *) header, slide);
+  if (rebindings_head) {
+    free(rebindings_head->rebindings);
+  }
+  free(rebindings_head);
+  return retval;
 }
 
 int rebind_symbols(struct rebinding rebindings[], size_t rebindings_nel) {