fix issue with the merging happening at the wrong stack frame by chec…

…king the rsp pointer

still exists issue with some pointer getting currupted

make

@@ -22,6 +22,7 @@ C_FLAGS = (
     '-I ./deps/ '
     '-I ./deps/udis86 '
     '-I ./deps/asmjit/src '
+    '-I ./build/ '
 )
 CXX_FLAGS = (
     '-std=c++14 '
@@ -95,6 +96,13 @@ def link():
     after()
 
 def compile():
+    with open('build/build_version.h', 'w+') as f:
+        f.write('''
+#ifndef RED_BUILD_VERSION
+#define RED_BUILD_VERSION "{}"
+#endif
+        '''.format(GIT_VERSION))
+
     for f in glob.glob('src/*.cc'):
         Run('{CXX} {} -c {} -o {}'.format(
             C_FLAGS + CXX_FLAGS,

src/config.h

@@ -43,7 +43,7 @@
 // unable to determine where it should actually be performing this, so we are just makeing the end of the branchable frame
 // close out any traces that were created in this frame
 // ^^^ this might have been a bug with the intergration, but now it is using this auto closing as a "feature"
-#define CONF_ALLOW_UNCLOSED_TRACES
+//#define CONF_ALLOW_UNCLOSED_TRACES
 
 
 // using timers in addition to number of times it loops to determine what to trace

src/jit_internal.h

@@ -52,6 +52,13 @@ namespace redmagic {
 
   struct tracer_stack_state;
 
+  template<typename Key, typename Value>
+  using RealMallocMap = std::unordered_map<Key, Value, std::hash<Key>, std::equal_to<Key>, RealMallocAllocator<std::pair<const Key, Value>>>;
+
+  template<typename Value>
+  using RealMallocSet = std::unordered_set<Value, std::hash<Value>, std::equal_to<Value>, RealMallocAllocator<Value>>;
+
+
   class Manager {
   public:
     Manager();
@@ -119,31 +126,34 @@ namespace redmagic {
 #endif
     };
 
-    std::unordered_map<
-      void*,
-      branch_info,
-      // should be the same as normal
-      std::hash<void*>,
-      std::equal_to<void*>,
-      RealMallocAllocator<std::pair<const void*, branch_info>>
-      > branches;
+    // std::unordered_map<
+    //   void*,
+    //   branch_info,
+    //   // should be the same as normal
+    //   std::hash<void*>,
+    //   std::equal_to<void*>,
+    //   RealMallocAllocator<std::pair<const void*, branch_info>>
+    //   >
+    RealMallocMap<void*, branch_info> branches;
 
 #ifdef CONF_CHECK_MERGE_RIP
-    std::unordered_map<
-      mem_loc_t,
-      mem_loc_t,
-      std::hash<mem_loc_t>,
-      std::equal_to<mem_loc_t>,
-      RealMallocAllocator<std::pair<const mem_loc_t, mem_loc_t>>
-      > merge_rip;
+    // std::unordered_map<
+    //   mem_loc_t,
+    //   mem_loc_t,
+    //   std::hash<mem_loc_t>,
+    //   std::equal_to<mem_loc_t>,
+    //   RealMallocAllocator<std::pair<const mem_loc_t, mem_loc_t>>
+    //   >
+    RealMallocMap<mem_loc_t, RealMallocSet<mem_loc_t> > merge_rip;
 #endif
   private:
-    std::unordered_set<
-      void*,
-      std::hash<void*>,
-      std::equal_to<void*>,
-      RealMallocAllocator<void*>
-      > no_trace_methods;
+    // std::unordered_set<
+    //   void*,
+    //   std::hash<void*>,
+    //   std::equal_to<void*>,
+    //   RealMallocAllocator<void*>
+    //   >
+    RealMallocSet<void*> no_trace_methods;
 
     std::atomic<uint32_t> thread_id_counter;
 
@@ -179,18 +189,17 @@ namespace redmagic {
     mem_loc_t method_address;
     mem_loc_t return_stack_pointer;
 
+#ifdef CONF_MERGE_BACK_ON_RET
+    int corresponding_merge_block = 0;
+#endif
+
     tracer_method_stack_s(mem_loc_t a=0, mem_loc_t b=0):
       method_address(a), return_stack_pointer(b) {}
   };
 
   struct tracer_merge_block_stack_s {
     mem_loc_t merge_head = 0; // head of linked list for this merge point
 
-#ifdef CONF_CHECK_MERGE_RIP
-    mem_loc_t merge_rip_head = 0;
-#endif
-
-
     tracer_merge_block_stack_s() {}
   };
 
@@ -307,6 +316,7 @@ namespace redmagic {
         location++;
       }
       assert(did_find == 1);
+      assert(*ret == from);
       return ret;
     }
 

src/manager.cc

@@ -6,6 +6,8 @@
 
 #include <algorithm>
 
+#include "build_version.h"
+
 using namespace redmagic;
 using namespace std;
 
@@ -141,7 +143,7 @@ extern "C" void *__real_malloc(size_t);
 
 extern "C" void redmagic_start() {
   UnprotectMalloc upm;
-  red_printf("Using redmagic jit by Matthew Francis-Landau <[email protected]>\n");
+  red_printf("Using redmagic jit by Matthew Francis-Landau <[email protected]> version: " RED_BUILD_VERSION "\n");
   if(manager != nullptr) {
     red_printf("redmagic_start called twice");
     abort();
@@ -531,7 +533,7 @@ void* Manager::backwards_branch(void *id, void *ret_addr) {
       new_head = head;
 #ifdef CONF_ESTIMATE_INSTRUCTIONS
       head->num_backwards_loops++;
-      if(head->num_backwards_loops > (info->count / 8)) {
+      if(head->num_backwards_loops > (info->count / 8) || (head->num_backwards_loops > 2 && info->avg_observed_instructions == 0)) {
         uint64_t icnt = instruction_cnt();
         info->avg_observed_instructions = (icnt - head->instruction_cnt_at_start - head->sub_frame_num_instructions) / head->num_backwards_loops;
       }
@@ -565,10 +567,10 @@ void* Manager::backwards_branch(void *id, void *ret_addr) {
 #ifndef NDEBUG
         if(info->tracer->owning_thread == get_thread_id()) {
           // this is a recursive frame
-          for(auto t : threadl_tracer_stack) {
-            if(t.trace_id == id) {
+          for(int i = threadl_tracer_stack.size() - 2; i >= 0; i--) {
+            if(threadl_tracer_stack[i].trace_id == id) {
               // we found this that was tracing from
-              assert(t.frame_id != branchable_frame_id);
+              assert(threadl_tracer_stack[i].frame_id != branchable_frame_id);
               return start_addr;
             }
           }
@@ -875,10 +877,17 @@ void* Manager::end_branchable_frame(void *ret_addr, void **stack_ptr) {
       // this method will get call again with this current frame poped
       return ret;
     } else {
+      if(head->trace_id != nullptr) {
+        auto info = &branches[head->trace_id];
+        info->count_fellthrough++;
+      }
       pop_tracer_stack();
       head = get_tracer_head();
 
       if(head->resume_addr) {
+        // assert that this is a call instruction
+        assert(((uint8_t*)ret_addr)[-5] == 0xE8);
+
         if(head->tracer) {
           head->tracer->JumpFromNestedLoop((uint8_t*)ret_addr - 5); // backup the pc to the call instruction
         }
@@ -944,7 +953,7 @@ tracer_stack_state *Manager::get_tracer_head() {
     // tracer_stack_state e;
     // threadl_tracer_stack.push_back(e);
     push_tracer_stack();
-    stack_head = &threadl_tracer_stack[0];
+    //    stack_head = &threadl_tracer_stack[0];
   }
   return stack_head;
 }

src/tracer.cc

@@ -47,6 +47,11 @@ namespace redmagic {
     return false;
   }
 #endif
+
+#ifndef NDEBUG
+  // using a conditional break point is much slower then having the check internally
+  long global_debugger_int = -1;
+#endif
 }
 
 
@@ -160,6 +165,8 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
   head->is_compiled = false;
 
   l->owning_thread = manager->get_thread_id();
+  // vvv this might be "wrong" since the frame id might have changed within a given trace
+  // but this field is write only atm so might not be a problem
   l->owning_frame_id = branchable_frame_id;
 
   // calling Start will invalidate the previous stack
@@ -190,6 +197,8 @@ extern "C" void* red_end_trace(mem_loc_t normal_end_address) {
   auto head = manager->pop_tracer_stack();
   auto new_head = manager->get_tracer_head();
   assert(head.is_traced);
+  auto info = &manager->branches[head.trace_id];
+  info->count_fellthrough++;
   void *ret = (void*)normal_end_address;
   if(new_head->is_traced) {
     if(new_head->tracer) {
@@ -280,6 +289,9 @@ void* Tracer::Start(void *start_addr) {
   local_jump_min_addr = 0;
 #endif
 
+  current_not_traced_call_addr = 0;
+  current_not_traced_call_ret_loc = nullptr;
+
 
 #ifdef CONF_GLOBAL_ABORT
 #ifdef NDEBUG
@@ -340,6 +352,11 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
   red_printf("----->start %i\n", ++loop_n);
 #endif
 
+  assert(method_stack.size() == 1);
+  assert(method_stack.back().return_stack_pointer == -1);
+
+  run_starting_stack_pointer = regs_struct->rsp;
+
   while(true) {
     assert(before_stack == 0xdeadbeef);
     assert(after_stack == 0xdeadbeef);
@@ -361,6 +378,7 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
     assert(method_stack.back().return_stack_pointer >= regs_struct->rsp + move_stack_by);
 #ifdef CONF_MERGE_BACK_ON_RET
     assert(merge_block_stack.size() >= method_stack.size());
+    assert(merge_block_stack.size() >= method_stack.back().corresponding_merge_block);
 #endif
 
     // if we somehow have less then 1kb free then we might have overwritten something
@@ -412,6 +430,12 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
         red_printf("[%10lu %8i %#016lx] \t%-38s %-20s lib=%s\n", global_icount, icount, ins_loc, ud_insn_asm(&disassm), ud_insn_hex(&disassm), dlinfo.dli_fname);
 #endif
 
+#ifndef NDEBUG
+      if(global_debugger_int != -1 && global_debugger_int == global_icount) {
+        __asm__("int3\n");
+      }
+#endif
+
       jmp_info = decode_instruction();
       if(jmp_info.is_jump) {
         if(jmp_info.is_local_jump) {
@@ -800,6 +824,15 @@ void* Tracer::EndMergeBlock() {
 }
 
 mem_loc_t Tracer::merge_close_core() {
+#ifdef CONF_CHECK_MERGE_RIP
+  mem_loc_t check_rip;
+  if(current_not_traced_call_ret_loc != nullptr) {
+    check_rip = last_call_pc;
+  } else {
+    check_rip = udis_loc;
+  }
+#endif
+
   if(merge_block_stack.size() == 1) {
     assert(merge_resume != 0);
     mem_loc_t resume_a = merge_resume;
@@ -817,8 +850,11 @@ mem_loc_t Tracer::merge_close_core() {
       write_addr = next_addr;
     }
 #ifdef CONF_CHECK_MERGE_RIP
-    mem_loc_t merge_rip = manager->merge_rip[resume_a];
-    assert(merge_rip == udis_loc);
+    auto merge_f = manager->merge_rip.find(resume_a);
+    assert(merge_f != manager->merge_rip.end());
+    assert(merge_f->second.find(check_rip) != merge_f->second.end());
+    // mem_loc_t merge_rip = merge_f->second;
+    // assert(merge_rip == check_rip);
 #endif
 
     // the ending of this tracer instructions
@@ -853,7 +889,9 @@ mem_loc_t Tracer::merge_close_core() {
     Tracer *expected = nullptr;
     if(!free_tracer_list.compare_exchange_strong(expected, this)) {
       // gaaaa
+      protected_malloc = false;
       delete this;
+      protected_malloc = true;
     }
 
 #ifdef CONF_VERBOSE
@@ -864,15 +902,23 @@ mem_loc_t Tracer::merge_close_core() {
   } else {
     mem_loc_t merge_addr = buffer->getRawBuffer() + buffer->getOffset();
     mem_loc_t write_addr = merge_block_stack.back().merge_head;
+
+#ifdef CONF_CHECK_MERGE_RIP
+    // if(manager->merge_rip.find(merge_addr) == manager->merge_rip.end() || current_not_traced_call_ret_loc != nullptr)
+    //   manager->merge_rip[merge_addr] = check_rip;
+    if(write_addr != 0) {
+      auto merge_set = &manager->merge_rip[merge_addr];
+      merge_set->insert(check_rip);
+      assert(merge_set->size() < 3); // just trying to avoid filling the set with a lot of items since this is suppose to catch bugs
+    }
+#endif
+
     merge_block_stack.pop_back();
     while(write_addr != 0) {
       mem_loc_t next_addr = *(mem_loc_t*)write_addr;
       *(mem_loc_t*)write_addr = merge_addr;
       write_addr = next_addr;
     }
-#ifdef CONF_CHECK_MERGE_RIP
-    manager->merge_rip[merge_addr] = udis_loc;
-#endif
     return 0;
   }
   assert(0);
@@ -912,7 +958,9 @@ void Tracer::continue_program(mem_loc_t resume_loc) {
   regs_struct->rsp += move_stack_by;
   move_stack_by = 0;
   *((register_t*)(regs_struct->rsp - TRACE_RESUME_ADDRESS_OFFSET)) = resume_loc;
+  // assert((regs_struct->rax & ~0xff) != 0xfbfbfbfbfbfbfb00);
   regs_struct = (struct user_regs_struct*)red_asm_resume_eval_block(&resume_struct, regs_struct);
+  // assert((regs_struct->rax & ~0xff) != 0xfbfbfbfbfbfbfb00);
 }
 
 
@@ -1559,6 +1607,7 @@ void Tracer::evaluate_instruction() {
           if(merge_close) // this might have an issue with the tracer getting free and then using it when it trys to continue the program
             continue_program(merge_close);
         }
+        assert(merge_block_stack.size() < method_stack.back().corresponding_merge_block || method_stack.back().corresponding_merge_block == 0);
 #endif
         if(method_stack.back().return_stack_pointer == regs_struct->rsp + move_stack_by - sizeof(register_t)) {
           assert(!method_stack.empty());
@@ -1708,6 +1757,7 @@ void Tracer::evaluate_instruction() {
       method_stack.push_back(tracer_method_stack_s(call_pc, regs_struct->rsp + move_stack_by));
 #ifdef CONF_MERGE_BACK_ON_RET
       merge_block_stack.push_back(tracer_merge_block_stack_s());
+      method_stack.back().corresponding_merge_block = merge_block_stack.size();
 #endif
     }
 
@@ -1732,20 +1782,18 @@ void Tracer::evaluate_instruction() {
     // assert(disassm.pfx_repne == UD_NONE);
     mem_loc_t ret_addr = pop_stack();
     set_pc(ret_addr);
-// #ifdef CONF_CHECK_RET_ADDRESS
-
-// #else
     buffer->writeToEnd(cb_asm_pop_stack);
-    //#endif
 
     assert(method_stack.back().return_stack_pointer >= regs_struct->rsp + move_stack_by - sizeof(register_t));
 #ifdef CONF_MERGE_BACK_ON_RET
-    if(method_stack.back().return_stack_pointer == -1 ||
+    assert(merge_block_stack.size() >= method_stack.back().corresponding_merge_block);
+    if((method_stack.back().return_stack_pointer == -1 && regs_struct->rsp + move_stack_by - sizeof(register_t) >= run_starting_stack_pointer) ||
        method_stack.back().return_stack_pointer == regs_struct->rsp + move_stack_by - sizeof(register_t)) {
       mem_loc_t merge_close = merge_close_core();
       if(merge_close) // this might have an issue with the tracer getting free and then using it when it trys to continue the program
         continue_program(merge_close);
     }
+    //assert(merge_block_stack.size() >= method_stack.back().corresponding_merge_block);// || method_stack.back().corresponding_merge_block == 0);
 #endif
     if(method_stack.back().return_stack_pointer == regs_struct->rsp + move_stack_by - sizeof(register_t)) {
       assert(!method_stack.empty());