Implement -instrument_transitive

googleprojectzero · Oct 31, 2024 · 7d2f643 · 7d2f643
1 parent e651645
commit 7d2f643
Show file tree

Hide file tree

Showing 9 changed files with 109 additions and 63 deletions.
diff --git a/Linux/debugger.cpp b/Linux/debugger.cpp
@@ -1499,7 +1499,8 @@ void Debugger::ExtractCodeRanges(void *module_base,
                                  size_t min_address,
                                  size_t max_address,
                                  std::list<AddressRange> *executable_ranges,
-                                 size_t *code_size)
+                                 size_t *code_size,
+                                 bool do_protect)
 {
   std::string elf_filename;
   if(module_base) {
@@ -1531,12 +1532,14 @@ void Debugger::ExtractCodeRanges(void *module_base,
         iter->addr_from = min_address;
       }
     }
-
-    int ret = RemoteMprotect((void *)iter->addr_from, 
-                             (iter->addr_to - iter->addr_from), 
-                             iter->permissions ^ PROT_EXEC);
-    if(ret) {
-      FATAL("Could not apply memory protection");
+
+    if(do_protect) {
+      int ret = RemoteMprotect((void *)iter->addr_from,
+                               (iter->addr_to - iter->addr_from),
+                               iter->permissions ^ PROT_EXEC);
+      if(ret) {
+        FATAL("Could not apply memory protection");
+      }
     }
 
     AddressRange range;

diff --git a/Linux/debugger.h b/Linux/debugger.h
@@ -162,7 +162,8 @@ class Debugger {
                          size_t min_address,
                          size_t max_address,
                          std::list<AddressRange> *executable_ranges,
-                         size_t *code_size);
+                         size_t *code_size,
+                         bool do_protect = true);
 
   void ProtectCodeRanges(std::list<AddressRange> *executable_ranges);
 

diff --git a/README.md b/README.md
@@ -184,6 +184,8 @@ In addition to the general-purpose API documented above, TinyInst also implement
 
 `-instrument_module [module name]` specifies which module to instrument, multiple `-instrument_module` options can be specified to instrument multiple modules.
 
+`-instrument_transitive [module name]` similar to `-instrument_module` except only code entered from other instrumented modules will run instrumented. Primarily used as optimization for calls like module1->module2->module1 where it's not important to instrument the entire module2 module, but module2->module1 entries are causing slowdowns.
+
 `-indirect_instrumentation [none|local|global|auto]` which instrumentation to use for indirect jump/calls
 
 `-patch_return_addresses` - replaces return address with the original value, causes returns to be instrumented using whatever `-indirect_instrumentation` method is specified

diff --git a/Windows/debugger.cpp b/Windows/debugger.cpp
@@ -553,7 +553,8 @@ void Debugger::ExtractCodeRanges(void *module_base,
                                  size_t min_address,
                                  size_t max_address,
                                  std::list<AddressRange> *executable_ranges,
-                                 size_t *code_size)
+                                 size_t *code_size,
+                                 bool do_protect)
 {
   LPCVOID end_address = (void *)max_address;
   LPCVOID cur_address = (void *)min_address;
@@ -593,17 +594,19 @@ void Debugger::ExtractCodeRanges(void *module_base,
         FATAL("Error in ReadProcessMemory");
       }
 
-      uint8_t low = meminfobuf.Protect & 0xFF;
-      low = low >> 4;
-      DWORD newProtect = (meminfobuf.Protect & 0xFFFFFF00) + low;
-      DWORD oldProtect;
-      if (!VirtualProtectEx(child_handle,
-        meminfobuf.BaseAddress,
-        meminfobuf.RegionSize,
-        newProtect,
-        &oldProtect))
-      {
-        FATAL("Error in VirtualProtectEx");
+      if(do_protect) {
+        uint8_t low = meminfobuf.Protect & 0xFF;
+        low = low >> 4;
+        DWORD newProtect = (meminfobuf.Protect & 0xFFFFFF00) + low;
+        DWORD oldProtect;
+        if (!VirtualProtectEx(child_handle,
+                              meminfobuf.BaseAddress,
+                              meminfobuf.RegionSize,
+                              newProtect,
+                              &oldProtect))
+        {
+          FATAL("Error in VirtualProtectEx");
+        }
       }
 
       newRange.from = (size_t)meminfobuf.BaseAddress;

diff --git a/Windows/debugger.h b/Windows/debugger.h
@@ -124,7 +124,8 @@ class Debugger {
                          size_t min_address,
                          size_t max_address,
                          std::list<AddressRange> *executable_ranges,
-                         size_t *code_size);
+                         size_t *code_size,
+                         bool do_protect = true);
 
   void ProtectCodeRanges(std::list<AddressRange> *executable_ranges);
 

diff --git a/macOS/debugger.cpp b/macOS/debugger.cpp
@@ -815,8 +815,9 @@ void Debugger::ExtractCodeRanges(void *base_address,
                                  size_t min_address,
                                  size_t max_address,
                                  std::list<AddressRange> *executable_ranges,
-                                 size_t *code_size) {
-
+                                 size_t *code_size,
+                                 bool do_protect)
+{
   if(!base_address) {
     ExtractSegmentCodeRanges(min_address, max_address, executable_ranges, code_size);
     return;
@@ -876,7 +877,7 @@ void Debugger::ExtractCodeRanges(void *base_address,
       mach_vm_address_t segment_start_addr = (mach_vm_address_t)segment_cmd->vmaddr + file_vm_slide;
       mach_vm_address_t segment_end_addr = (mach_vm_address_t)segment_cmd->vmaddr + file_vm_slide + segment_cmd->vmsize;
 
-      ExtractSegmentCodeRanges(segment_start_addr, segment_end_addr, executable_ranges, code_size);
+      ExtractSegmentCodeRanges(segment_start_addr, segment_end_addr, executable_ranges, code_size, do_protect);
 #endif
     }
 
@@ -889,7 +890,9 @@ void Debugger::ExtractCodeRanges(void *base_address,
 void Debugger::ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr,
                                         mach_vm_address_t segment_end_addr,
                                         std::list<AddressRange> *executable_ranges,
-                                        size_t *code_size) {
+                                        size_t *code_size,
+                                        bool do_protect)
+{
   mach_vm_address_t cur_address = segment_start_addr;
   while (cur_address < segment_end_addr) {
     mach_vm_size_t region_size = 0;
@@ -916,38 +919,42 @@ void Debugger::ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr,
       new_range.data = (char *)malloc(range_size);
       RemoteRead((void*)new_range.from, new_range.data, range_size);
 
-    retry_label:
-      RemoteProtect((void*)new_range.from, range_size, info.protection ^ VM_PROT_EXECUTE);
-      mach_vm_address_t region_addr = new_range.from;
-      mach_vm_size_t region_sz = range_size;
-      vm_region_submap_info_data_64_t region_info;
-      mach_target->GetRegionSubmapInfo(&region_addr, (mach_vm_size_t*)&region_sz, &region_info);
-      if (region_info.protection & VM_PROT_EXECUTE) {
-        if (retried) {
-          FATAL("Failed to mark the original code NON-EXECUTABLE\n");
-        }
-
-        kern_return_t krt;
-        krt = mach_vm_deallocate(mach_target->Task(),
-                                 (mach_vm_address_t)new_range.from,
-                                 range_size);
-
-        if (krt == KERN_SUCCESS) {
-          mach_vm_address_t alloc_address = new_range.from;
-          krt = mach_vm_allocate(mach_target->Task(),
-                                 (mach_vm_address_t*)&alloc_address,
-                                 range_size,
-                                 VM_FLAGS_FIXED);
-
-          if (krt == KERN_SUCCESS && alloc_address && new_range.from) {
-            RemoteWrite((void*)new_range.from, new_range.data, range_size);
-          } else {
-            FATAL("Unable to re-allocate memory after deallocate in ExtractSegmentCodeRanges\n");
+      if(do_protect) {
+      retry_label:
+        RemoteProtect((void*)new_range.from, range_size, info.protection ^ VM_PROT_EXECUTE);
+        mach_vm_address_t region_addr = new_range.from;
+        mach_vm_size_t region_sz = range_size;
+        vm_region_submap_info_data_64_t region_info;
+        mach_target->GetRegionSubmapInfo(&region_addr, (mach_vm_size_t*)&region_sz, &region_info);
+        if (region_info.protection & VM_PROT_EXECUTE) {
+          if (retried) {
+            FATAL("Failed to mark the original code NON-EXECUTABLE\n");
+          }
+
+          kern_return_t krt;
+          krt = mach_vm_deallocate(mach_target->Task(),
+                                   (mach_vm_address_t)new_range.from,
+                                   range_size);
+
+          if (krt == KERN_SUCCESS) {
+            mach_vm_address_t alloc_address = new_range.from;
+            krt = mach_vm_allocate(mach_target->Task(),
+                                   (mach_vm_address_t*)&alloc_address,
+                                   range_size,
+                                   VM_FLAGS_FIXED);
+
+            if (krt == KERN_SUCCESS && alloc_address && new_range.from) {
+              RemoteWrite((void*)new_range.from, new_range.data, range_size);
+            } else {
+              FATAL("Unable to re-allocate memory after deallocate in ExtractSegmentCodeRanges\n");
+            }
           }
+
+          retried = true;
+          goto retry_label;
         }
-
-        retried = true;
-        goto retry_label;
+      } else {
+        //WARN("skipping memory protection");
       }
 
       AddressRange *last_range = NULL;

diff --git a/macOS/debugger.h b/macOS/debugger.h
@@ -156,6 +156,9 @@ friend kern_return_t catch_mach_exception_raise_state_identity(
   DebuggerStatus Kill();
   DebuggerStatus Continue(uint32_t timeout);
   DebuggerStatus Attach(unsigned int pid, uint32_t timeout);
+  DebuggerStatus GetDebuggerStatus() {
+    return dbg_last_status;
+  }
 
   bool IsTargetAlive();
   bool IsTargetFunctionDefined() { return target_function_defined; }
@@ -254,7 +257,8 @@ friend kern_return_t catch_mach_exception_raise_state_identity(
                          size_t min_address,
                          size_t max_address,
                          std::list<AddressRange> *executable_ranges,
-                         size_t *code_size);
+                         size_t *code_size,
+                         bool do_protect = true);
 
   void ProtectCodeRanges(std::list<AddressRange> *executable_ranges);
 
@@ -398,7 +402,7 @@ friend kern_return_t catch_mach_exception_raise_state_identity(
   void ExtractSegmentCodeRanges(mach_vm_address_t segment_start_addr,
                                 mach_vm_address_t segment_end_addr,
                                 std::list<AddressRange> *executable_ranges,
-                                size_t *code_size);
+                                size_t *code_size, bool do_protect = true);
 
   void HandleDyld(void *module);
 

diff --git a/tinyinst.cpp b/tinyinst.cpp
@@ -49,6 +49,7 @@ ModuleInfo::ModuleInfo() {
   instrumented_code_size = 0;
   unwind_data = NULL;
   client_data = NULL;
+  do_protect = true;
 }
 
 void ModuleInfo::ClearInstrumentation() {
@@ -900,16 +901,19 @@ void TinyInst::InstrumentModule(ModuleInfo *module) {
     printf("Module %s already instrumented, "
            "reusing instrumentation data\n",
            module->module_name.c_str());
-    ProtectCodeRanges(&module->executable_ranges);
+    if (module->do_protect) {
+      ProtectCodeRanges(&module->executable_ranges);
+    }
     FixCrossModuleLinks(module);
     return;
   }
-
+  
   ExtractCodeRanges(module->module_header,
                     module->min_address,
                     module->max_address,
                     &module->executable_ranges,
-                    &module->code_size);
+                    &module->code_size,
+                    module->do_protect);
 
   // allocate buffer for instrumented code
   module->instrumented_code_size = module->code_size * CODE_SIZE_MULTIPLIER;
@@ -1213,6 +1217,18 @@ InstructionResult TinyInst::InstrumentInstruction(ModuleInfo *module,
   return INST_NOTHANDLED;
 }
 
+void TinyInst::AddInstrumentedModule(char* name, bool do_protect) {
+  std::string module_name = name;
+  for(auto iter=instrumented_modules.begin(); iter!=instrumented_modules.end(); iter++) {
+    if((*iter)->module_name == module_name) {
+      FATAL("Duplicate instrumented modules, module %s is already being instrumented", name);
+    }
+  }
+  ModuleInfo *new_module = new ModuleInfo();
+  new_module->module_name = module_name;
+  new_module->do_protect = do_protect;
+  instrumented_modules.push_back(new_module);
+}
 
 // initializes instrumentation from command line options
 void TinyInst::Init(int argc, char **argv) {
@@ -1298,9 +1314,15 @@ void TinyInst::Init(int argc, char **argv) {
 #endif
 
   for (const auto module_name: module_names) {
-    ModuleInfo *new_module = new ModuleInfo();
-    new_module->module_name = module_name;
-    instrumented_modules.push_back(new_module);
+    AddInstrumentedModule(module_name, true);
+    // SAY("--- %s\n", module_name);
+  }
+
+  std::list <char *> module_names_transitive;
+  GetOptionAll("-instrument_transitive", argc, argv, &module_names_transitive);
+
+  for (const auto module_name: module_names_transitive) {
+    AddInstrumentedModule(module_name, false);
     // SAY("--- %s\n", module_name);
   }
 

diff --git a/tinyinst.h b/tinyinst.h
@@ -165,6 +165,7 @@ class TinyInst : public Debugger {
   void InstrumentAddressRange(const char *name, size_t min_address, size_t max_address);
 
 private:
+  void AddInstrumentedModule(char* name, bool do_protect);
   bool HandleBreakpoint(void *address);
   void OnInstrumentModuleLoaded(void *module, ModuleInfo *target_module);
   ModuleInfo *IsInstrumentModule(char *module_name);
@@ -321,6 +322,8 @@ class ModuleInfo {
   std::unordered_set<size_t> entry_offsets;
 
   UnwindData *unwind_data;
+
+  bool do_protect;
 
   // clients can use this to store additional data
   // about the module