Two-level namespace

README.md

@@ -43,7 +43,10 @@ int my_open(const char *path, int oflag, ...) {
 int main(int argc, char * argv[])
 {
   @autoreleasepool {
-    rebind_symbols((struct rebinding[2]){{"close", my_close, (void *)&orig_close}, {"open", my_open, (void *)&orig_open}}, 2);
+    rebind_symbols((struct rebinding[2]){
+      {"close", "/usr/lib/libSystem.B.dylib", my_close, (void *)&orig_close},
+      {"open", "/usr/lib/libSystem.B.dylib", my_open, (void *)&orig_open}
+    }, 2);
 
     // Open our own binary and print out first 4 bytes (which is the same
     // for all Mach-O binaries on a given architecture)

fishhook.c

@@ -51,6 +51,7 @@ typedef struct nlist nlist_t;
 
 struct rebindings_entry {
   struct rebinding *rebindings;
+  uint8_t *dylib_ordinals;
   size_t rebindings_nel;
   struct rebindings_entry *next;
 };
@@ -69,6 +70,12 @@ static int prepend_rebindings(struct rebindings_entry **rebindings_head,
     free(new_entry);
     return -1;
   }
+  new_entry->dylib_ordinals = malloc(sizeof(uint8_t) * nel);
+  if (!new_entry->dylib_ordinals) {
+    free(new_entry->rebindings);
+    free(new_entry);
+    return -1;
+  }
   memcpy(new_entry->rebindings, rebindings, sizeof(struct rebinding) * nel);
   new_entry->rebindings_nel = nel;
   new_entry->next = *rebindings_head;
@@ -95,9 +102,14 @@ static void perform_rebinding_with_section(struct rebindings_entry *rebindings,
     if (strnlen(symbol_name, 2) < 2) {
       continue;
     }
+    int dylib_ordinal = GET_LIBRARY_ORDINAL(symtab[symtab_index].n_desc);
     struct rebindings_entry *cur = rebindings;
     while (cur) {
       for (uint j = 0; j < cur->rebindings_nel; j++) {
+        if (cur->dylib_ordinals[j] != 0 && dylib_ordinal != cur->dylib_ordinals[j] && dylib_ordinal != SELF_LIBRARY_ORDINAL) {
+          continue;
+        }
+
         if (strcmp(&symbol_name[1], cur->rebindings[j].name) == 0) {
           if (cur->rebindings[j].replaced != NULL &&
               indirect_symbol_bindings[i] != cur->rebindings[j].replacement) {
@@ -121,25 +133,78 @@ static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
     return;
   }
 
-  segment_command_t *cur_seg_cmd;
+  struct load_command *cur_load_cmd;
   segment_command_t *linkedit_segment = NULL;
   struct symtab_command* symtab_cmd = NULL;
   struct dysymtab_command* dysymtab_cmd = NULL;
 
+  for (struct rebindings_entry *entry = rebindings; entry != NULL; entry = entry->next) {
+    bzero(entry->dylib_ordinals, sizeof(uint8_t) * entry->rebindings_nel);
+  }
+
+  int cur_dylib_ordinal = 1;
+  bool image_has_dylibs_to_rebind = false;
+
   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;
-    if (cur_seg_cmd->cmd == LC_SEGMENT_ARCH_DEPENDENT) {
-      if (strcmp(cur_seg_cmd->segname, SEG_LINKEDIT) == 0) {
-        linkedit_segment = cur_seg_cmd;
+  for (uint i = 0; i < header->ncmds; i++, cur += cur_load_cmd->cmdsize) {
+    cur_load_cmd = (struct load_command *)cur;
+    switch (cur_load_cmd->cmd) {
+      case LC_LOAD_DYLIB:
+      case LC_LOAD_WEAK_DYLIB:
+      case LC_REEXPORT_DYLIB:
+      case LC_LOAD_UPWARD_DYLIB:
+      case LC_LAZY_LOAD_DYLIB: {
+        if ((header->flags & MH_TWOLEVEL) == 0) {
+          break;
+        }
+
+        struct dylib_command *dylib_cmd = (struct dylib_command *)cur_load_cmd;
+        char *dylib = (char *)(cur + dylib_cmd->dylib.name.offset);
+        for (struct rebindings_entry *entry = rebindings; entry != NULL; entry = entry->next) {
+          for (int el = 0; el < entry->rebindings_nel; ++el) {
+            struct rebinding *cur_rebinding = &entry->rebindings[el];
+            if (cur_rebinding->dylib == NULL) {
+              image_has_dylibs_to_rebind = true;
+            } else if (strcmp(cur_rebinding->dylib, dylib) == 0) {
+              entry->dylib_ordinals[el] = cur_dylib_ordinal;
+              image_has_dylibs_to_rebind = true;
+            }
+          }
+        }
+        ++cur_dylib_ordinal;
+        break;
+      }
+      case LC_SEGMENT_ARCH_DEPENDENT: {
+        segment_command_t *seg_cmd = (segment_command_t *)cur_load_cmd;
+        if (strcmp(seg_cmd->segname, SEG_LINKEDIT) == 0) {
+          linkedit_segment = seg_cmd;
+        }
+        break;
       }
-    } else if (cur_seg_cmd->cmd == LC_SYMTAB) {
-      symtab_cmd = (struct symtab_command*)cur_seg_cmd;
-    } else if (cur_seg_cmd->cmd == LC_DYSYMTAB) {
-      dysymtab_cmd = (struct dysymtab_command*)cur_seg_cmd;
+      case LC_SYMTAB:
+        symtab_cmd = (struct symtab_command*)cur_load_cmd;
+        break;
+      case LC_DYSYMTAB:
+        dysymtab_cmd = (struct dysymtab_command*)cur_load_cmd;
+        break;
+      default:
+        break;
     }
   }
 
+  // This optimization is enabled by two-level namespaces.
+  //
+  // There are two sets of dylibs:
+  //   1. The set of dylibs associated with the rebindings
+  //   2. The set of dylibs loaded by an image (i.e executable or library)
+  //
+  // When there is no overlap between these two sets of dylibs, then it becomes
+  // unnecessary to scan the symbol table of given image. However, if any
+  // rebinding has no target dylib (NULL), then all images are searched.
+  if (!image_has_dylibs_to_rebind) {
+    return;
+  }
+
   if (!symtab_cmd || !dysymtab_cmd || !linkedit_segment ||
       !dysymtab_cmd->nindirectsyms) {
     return;
@@ -154,14 +219,15 @@ static void rebind_symbols_for_image(struct rebindings_entry *rebindings,
   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) {
+  for (uint i = 0; i < header->ncmds; i++, cur += cur_load_cmd->cmdsize) {
+    cur_load_cmd = (struct load_command *)cur;
+    if (cur_load_cmd->cmd == LC_SEGMENT_ARCH_DEPENDENT) {
+      segment_command_t *seg_cmd = (segment_command_t *)cur_load_cmd;
+      if (strcmp(seg_cmd->segname, SEG_DATA) != 0 &&
+          strcmp(seg_cmd->segname, SEG_DATA_CONST) != 0) {
         continue;
       }
-      for (uint j = 0; j < cur_seg_cmd->nsects; j++) {
+      for (uint j = 0; j < seg_cmd->nsects; j++) {
         section_t *sect =
           (section_t *)(cur + sizeof(segment_command_t)) + j;
         if ((sect->flags & SECTION_TYPE) == S_LAZY_SYMBOL_POINTERS) {

fishhook.h

@@ -43,6 +43,7 @@ extern "C" {
  */
 struct rebinding {
   const char *name;
+  const char *dylib;
   void *replacement;
   void **replaced;
 };