Best effort support for relative jump outside the current module

googleprojectzero · Mar 20, 2024 · 2d9472e · 2d9472e
1 parent 9cdc11e
commit 2d9472e
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 15 deletions.
diff --git a/arch/x86/x86_assembler.cpp b/arch/x86/x86_assembler.cpp
@@ -785,12 +785,6 @@ void X86Assembler::HandleBasicBlockEnd(
     const char *target_address1 = address + offset;
     const char *target_address2 = address + offset + disp;
 
-    if (tinyinst_.GetModule((size_t)target_address2) != module) {
-      WARN("Relative jump to a differen module in bb at %p\n", address);
-      tinyinst_.InvalidInstruction(module);
-      return;
-    }
-
     // preliminary encode jump instruction
     // displacement might be changed later as we don't know
     // the size of edge instrumentation yet
@@ -843,6 +837,11 @@ void X86Assembler::HandleBasicBlockEnd(
                                   jump_size);
     }
 
+    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);
@@ -879,8 +878,7 @@ void X86Assembler::HandleBasicBlockEnd(
       const char *target_address = address + offset + disp;
 
       if (tinyinst_.GetModule((size_t)target_address) != module) {
-        WARN("Relative jump to a differen module in bb at %p\n", address);
-        tinyinst_.InvalidInstruction(module);
+        tinyinst_.OutsideJump(module, (size_t)target_address);
         return;
       }
 
@@ -937,12 +935,6 @@ void X86Assembler::HandleBasicBlockEnd(
       const char *return_address = address + offset;
       const char *call_address = address + offset + disp;
 
-      if (tinyinst_.GetModule((size_t)call_address) != module) {
-        WARN("Relative jump to a differen module in bb at %p\n", address);
-        tinyinst_.InvalidInstruction(module);
-        return;
-      }
-
       // fix the displacement and emit the call
       if (!tinyinst_.patch_return_addresses) {
         unsigned char encoded[15];
@@ -969,6 +961,11 @@ void X86Assembler::HandleBasicBlockEnd(
             (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, JMP, sizeof(JMP));
 
@@ -981,6 +978,11 @@ void X86Assembler::HandleBasicBlockEnd(
       } else {
         PushReturnAddress(module, (uint64_t)return_address);
 
+        if (tinyinst_.GetModule((size_t)call_address) != module) {
+          tinyinst_.OutsideJump(module, (size_t)call_address);
+          return;
+        }
+
         // jmp call_address
         tinyinst_.WriteCode(module, JMP, sizeof(JMP));
 

diff --git a/tinyinst.cpp b/tinyinst.cpp
@@ -78,6 +78,8 @@ void ModuleInfo::ClearInstrumentation() {
   jumptable_address_offset = 0;
 
   invalid_instructions.clear();
+  outside_jumps.clear();
+
   tracepoints.clear();
 }
 
@@ -301,6 +303,15 @@ bool TinyInst::HandleBreakpoint(void *address) {
     return true;
   }
 
+  auto iter = module->outside_jumps.find((size_t)address);
+  if (iter != module->outside_jumps.end()) {
+
+    WARN("Executing relative jump otside the current module");
+    SetRegister(ARCH_PC, iter->second);
+
+    return true;
+  }
+
   if(unwind_generator->HandleBreakpoint(module, address)) {
     return true;
   }
@@ -482,6 +493,13 @@ void TinyInst::InvalidInstruction(ModuleInfo *module) {
   assembler_->Crash(module);
 }
 
+void TinyInst::OutsideJump(ModuleInfo* module, size_t address) {
+  size_t breakpoint_address = (size_t)module->instrumented_code_remote +
+    module->instrumented_code_allocated;
+  assembler_->Breakpoint(module);
+  module->outside_jumps[breakpoint_address] = address;
+}
+
 void TinyInst::InstrumentIndirect(ModuleInfo *module,
                                   Instruction& inst,
                                   size_t instruction_address,

diff --git a/tinyinst.h b/tinyinst.h
@@ -181,7 +181,10 @@ class TinyInst : public Debugger {
                           uint32_t jmp_offset,
                           std::set<char *> *queue,
                           std::list<std::pair<uint32_t, uint32_t>> *offset_fixes);
-  void InvalidInstruction(ModuleInfo *module);
+  void InvalidInstruction(ModuleInfo* module);
+
+  // relative jump outside of current module
+  void OutsideJump(ModuleInfo* module, size_t address);
 
   // needed to support cross-module linking
   // on module unloads / reloads
@@ -310,6 +313,7 @@ class ModuleInfo {
   size_t jumptable_address_offset;
 
   std::unordered_set<size_t> invalid_instructions;
+  std::unordered_map<size_t, size_t> outside_jumps;
   std::unordered_map<size_t, size_t> tracepoints;
 
   std::unordered_set<size_t> entry_offsets;