Support instrumenting custom address range

arch/arm64/arm64_assembler.cpp

@@ -687,13 +687,6 @@ void Arm64Assembler::InstrumentCondJmp(
   const char *target_address1 = address + offset;
   const char *target_address2 = address + last_offset + branch_offset;
 
-  if (tinyinst_.GetModule((size_t)target_address2) != module) {
-    WARN("Relative jump to a differen module in bb at %p\n",
-         static_cast<const void *>(address));
-    tinyinst_.InvalidInstruction(module);
-    return;
-  }
-
   // preliminary encode cond branch instruction
   // offset will be changed later as we don't know
   // the size of edge instrumentation yet
@@ -720,6 +713,11 @@ void Arm64Assembler::InstrumentCondJmp(
   // fix conditional branch
   FixOffset(module, cond_branch_offset, label_offset);
 
+  if (tinyinst_.GetModule((size_t)target_address2) != module) {
+    tinyinst_.OutsideJump(module, (size_t)target_address2);
+    return;
+  }
+
   // instrument the 2nd edge
   tinyinst_.InstrumentEdge(module, module, (size_t)address,
                            (size_t)target_address2);
@@ -749,8 +747,7 @@ void Arm64Assembler::InstrumentJmp(
     const char *target_address = address + last_offset + branch_offset;
 
     if (tinyinst_.GetModule((size_t)target_address) != module) {
-      WARN("Relative jump to a differen module in bb at %p\n", (void *)address);
-      tinyinst_.InvalidInstruction(module);
+      tinyinst_.OutsideJump(module, (size_t)target_address);
       return;
     }
 
@@ -799,13 +796,6 @@ void Arm64Assembler::InstrumentCall(
     const char *return_address = address + offset;
     const char *call_address = address + last_offset + branch_offset;
 
-    if (tinyinst_.GetModule((size_t)call_address) != module) {
-      WARN("Relative jump to a differen module in bb at %p\n",
-           static_cast<const void *>(address));
-      tinyinst_.InvalidInstruction(module);
-      return;
-    }
-
     if (!tinyinst_.patch_return_addresses) {
       uint64_t addr = (uint64_t)module->instrumented_code_allocated +
                       (uint64_t)module->instrumented_code_local;
@@ -827,6 +817,11 @@ void Arm64Assembler::InstrumentCall(
           (uint32_t)(module->instrumented_code_allocated - 4), queue,
           offset_fixes);
 
+      if (tinyinst_.GetModule((size_t)call_address) != module) {
+        tinyinst_.OutsideJump(module, (size_t)call_address);
+        return;
+      }
+
       // jmp call_address
       tinyinst_.WriteCode(module, &branch_instr, sizeof(branch_instr));
 
@@ -838,6 +833,11 @@ void Arm64Assembler::InstrumentCall(
     } else {
       SetReturnAddress(module, (uint64_t)return_address);
 
+      if (tinyinst_.GetModule((size_t)call_address) != module) {
+        tinyinst_.OutsideJump(module, (size_t)call_address);
+        return;
+      }
+
       uint32_t branch_instr = b(0, 0);
       // jmp call_address
       tinyinst_.WriteCode(module, &branch_instr, sizeof(branch_instr));

litecov.cpp

@@ -57,6 +57,10 @@ void LiteCov::Init(int argc, char **argv) {
 void LiteCov::OnModuleInstrumented(ModuleInfo *module) {
   TinyInst::OnModuleInstrumented(module);
 
+  if(!module->client_data) {
+    module->client_data = new ModuleCovData();
+  }
+
   ModuleCovData *data = (ModuleCovData *)module->client_data;
 
   data->ClearInstrumentationData();

macOS/debugger.cpp

@@ -774,9 +774,15 @@ void Debugger::ExtractCodeRanges(void *base_address,
                                  size_t max_address,
                                  std::list<AddressRange> *executable_ranges,
                                  size_t *code_size) {
+
+  if(!base_address) {
+    ExtractSegmentCodeRanges(min_address, max_address, executable_ranges, code_size);
+    return;
+  }
+
   mach_header_64 mach_header;
   GetMachHeader(base_address, &mach_header);
-
+  
   void *load_commands_buffer = NULL;
   GetLoadCommandsBuffer(base_address, &mach_header, &load_commands_buffer);
 
@@ -1431,6 +1437,7 @@ void Debugger::HandleExceptionInternal(MachException *raised_mach_exception) {
       break;
 
     case EXC_BREAKPOINT:
+      WARN("Unhandled breakpoint\n");
       dbg_continue_status = KERN_FAILURE;
       break;
 

tinyinst.cpp

@@ -48,6 +48,7 @@ ModuleInfo::ModuleInfo() {
   instrumented_code_remote_previous = NULL;
   instrumented_code_size = 0;
   unwind_data = NULL;
+  client_data = NULL;
 }
 
 void ModuleInfo::ClearInstrumentation() {
@@ -306,7 +307,7 @@ bool TinyInst::HandleBreakpoint(void *address) {
   auto iter = module->outside_jumps.find((size_t)address);
   if (iter != module->outside_jumps.end()) {
 
-    WARN("Executing relative jump outside the current module");
+    // WARN("Executing relative jump outside the current module");
     SetRegister(ARCH_PC, iter->second);
 
     return true;
@@ -870,6 +871,26 @@ void TinyInst::ClearInstrumentation(ModuleInfo *module) {
   ClearCrossModuleLinks(module);
 }
 
+void TinyInst::InstrumentAddressRange(const char *name,
+                                      size_t min_address,
+                                      size_t max_address)
+{
+  ModuleInfo *module = GetModuleByName(name);
+  if(!module) {
+    module = new ModuleInfo();
+    module->module_name = name;
+    module->module_header = NULL;
+    instrumented_modules.push_back(module);
+  }
+
+  module->loaded = true;
+  module->min_address = min_address;
+  module->max_address = max_address;
+
+  InstrumentModule(module);
+}
+
+
 void TinyInst::InstrumentModule(ModuleInfo *module) {
   if (instrumentation_disabled) return;
 

tinyinst.h

@@ -161,6 +161,8 @@ class TinyInst : public Debugger {
   virtual void OnReturnAddress(ModuleInfo *module, size_t original_address, size_t translated_address);
 
   void RegisterHook(Hook *hook);
+
+  void InstrumentAddressRange(const char *name, size_t min_address, size_t max_address);
 
 private:
   bool HandleBreakpoint(void *address);