handful of minor fixes

src/asm_snippets.S

@@ -136,11 +136,16 @@ red_asm_pop_all_regs_start:
   m_pop_all_regs
 red_asm_pop_all_regs_end:
 
-  .global red_asm_jump_rsi
-red_asm_jump_rsi:
+  .global red_asm_jump_rsi_temp_enable
+red_asm_jump_rsi_temp_enable:
+  jmp *%rsi
+
+  .global red_asm_jump_rsi_resume_trace
+red_asm_jump_rsi_resume_trace:
   // this is only used to resume a tracer at a different address
   // the address where we are expecting to resume will be loaded into rsi, so if the check fails then we will jump to that address
   jmp *%rsi
 
-// we don't need executable stack
+
+  // we don't need executable stack
   .section .note.GNU-stack,"",%progbits

src/config.h

@@ -29,7 +29,7 @@
 
 // makes it print all the instructions processed and extra info
 #ifdef CONF_DEBUG_BUILD
-# define CONF_VERBOSE
+//# define CONF_VERBOSE
 #endif
 
 // support aborting the system after some fixed number of instruction have been processed, see tools/bisect for debugging with this

src/jit_internal.h

@@ -69,7 +69,7 @@ namespace redmagic {
     void* jump_to_trace(void *id);
 
     void* backwards_branch(void *id, void *ret_addr, void **stack_ptr);
-    void* fellthrough_branch(void *id, void *ret_addr);
+    void* fellthrough_branch(void *id, void *ret_addr, void **stack_ptr);
 
     // void ensure_not_traced();
 

src/manager.cc

@@ -83,11 +83,10 @@ extern "C" void* red_user_backwards_branch(void *id, void *ret_addr, void **stac
 
 extern "C" void* red_user_fellthrough_branch(void *id, void *ret_addr, void **stack_ptr) {
   UnprotectMalloc upm;
-  return manager->fellthrough_branch(id, ret_addr);
+  return manager->fellthrough_branch(id, ret_addr, stack_ptr);
 }
 
 extern "C" void* red_user_ensure_not_traced(void *_, void *ret_addr, void **stack_ptr) {
-  // TODO:
   return manager->ensure_not_traced();
 }
 
@@ -669,7 +668,7 @@ void* Manager::backwards_branch(void *id, void *ret_addr, void **stack_ptr) {
 
 }
 
-void* Manager::fellthrough_branch(void *id, void *ret_addr) {
+void* Manager::fellthrough_branch(void *id, void *ret_addr, void **stack_ptr) {
   // ignore
   if(id == nullptr)
     return NULL;
@@ -705,6 +704,7 @@ void* Manager::fellthrough_branch(void *id, void *ret_addr) {
       // we have to pop this frame since we weren't being traced and there is nothing that will do it for us
       auto old_head = pop_tracer_stack();
       auto new_head = get_tracer_head();
+      ((mem_loc_t*)stack_ptr)[-1] = old_head.frame_stack_ptr - (mem_loc_t)stack_ptr;
       if(new_head->resume_addr) {
         assert(old_head.return_to_trace_when_done);
         if(new_head->tracer) {
@@ -737,22 +737,25 @@ void* Manager::fellthrough_branch(void *id, void *ret_addr) {
   assert(!head->is_compiled || head->frame_id != branchable_frame_id);
 
   return NULL;
-
 }
 
 void* Manager::temp_disable(void *ret_addr) {
   temp_disable_last_addr = ret_addr;
   auto head = get_tracer_head();
   assert(!head->is_temp_disabled);
-  head->is_temp_disabled = true;
+
   //head->d_ret = ret_addr;
   void *ret = NULL;
-  assert(!head->is_traced || head->tracer);
+
+  //assert(!head->is_traced || head->tracer);
+  // ^^^ due to the tracer sometimes
+
   if(head->tracer && !head->tracer->did_abort) {
     // this will push the stack
     ret = head->tracer->TempDisableTrace();
   } else {
     assert(!head->resume_addr);
+    head->is_temp_disabled = true;
     push_tracer_stack();
   }
   return ret;
@@ -769,7 +772,7 @@ void* Manager::temp_enable(void *ret_addr) {
   //head->is_temp_disabled = false;
   void *ret = NULL;
   //head->d_ret = nullptr;
-  if(head->tracer && !head->tracer->did_abort) {
+  if(old_head.return_to_trace_when_done && head->tracer && !head->tracer->did_abort) {
     head->tracer->TempEnableTrace(ret_addr);
   }
   if(head->resume_addr != nullptr) {
@@ -921,8 +924,10 @@ void* Manager::end_branchable_frame(void *ret_addr, void **stack_ptr) {
     }
   }
 #endif
+  auto head = get_tracer_head();
+  assert(head->frame_stack_ptr > (mem_loc_t)stack_ptr);
   branchable_frame_id--;
-  assert(get_tracer_head()->frame_id <= branchable_frame_id);
+  assert(head->frame_id <= branchable_frame_id);
   return NULL;
 }
 

src/tracer.cc

@@ -55,7 +55,8 @@ namespace redmagic {
 }
 
 
-extern "C" void red_asm_jump_rsi();
+extern "C" void red_asm_jump_rsi_temp_enable();
+extern "C" void red_asm_jump_rsi_resume_trace();
 
 
 static int udis_input_hook (ud_t *ud) {
@@ -364,7 +365,7 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
 
   run_starting_stack_pointer = regs_struct->rsp;
 
-  if(get_pc() == (mem_loc_t)&red_asm_jump_rsi) {
+  if(get_pc() == (mem_loc_t)&red_asm_jump_rsi_resume_trace) {
     // this is only when resuming a tracer after another one exited
     assert(regs_struct->rdi < 250); // this is the offset for the starting stack pointer (not a "real" assert, just some attempt at sanity)
     // we offset the starting stack pointer since that controls where we end up merging back to
@@ -739,7 +740,16 @@ void* Tracer::EndTraceEnsure() {
 
 
 void* Tracer::TempDisableTrace() {
-  assert(current_not_traced_call_addr == (mem_loc_t)&redmagic_temp_disable);
+  if(current_not_traced_call_addr != (mem_loc_t)&redmagic_temp_disable) {
+    // then this must be inside of some not traced call
+    // but we still have to behave the same
+    auto head = manager->get_tracer_head();
+    head->is_temp_disabled = true;
+    assert(head->resume_addr == nullptr);
+    manager->push_tracer_stack();
+    return NULL;
+  }
+  //assert(current_not_traced_call_addr == (mem_loc_t)&redmagic_temp_disable);
   assert(icount - last_call_instruction < 2);
   buffer->setOffset(last_call_generated_op);
   SimpleCompiler compiler(buffer);
@@ -762,15 +772,22 @@ void Tracer::TempEnableTrace(void *resume_pc) {
   set_pc((mem_loc_t)resume_pc);
   SimpleCompiler compiler(buffer);
   // the "normal" return address will be set to ris when this returns from the temp disabled region
-  auto wb = compiler.TestRegister((mem_loc_t)&red_asm_jump_rsi, RSI, (register_t)resume_pc, &merge_block_stack.back());
+  auto wb = compiler.TestRegister((mem_loc_t)&red_asm_jump_rsi_temp_enable, RSI, (register_t)resume_pc, &merge_block_stack.back());
   auto written = compiler.finalize();
   wb.replace_stump<uint64_t>(0xfafafafafafafafa, written.getRawBuffer());
   write_interrupt_block();
 }
 
 void Tracer::JumpFromNestedLoop(void *resume_pc) {
-  // same code where we check the rsi register for where we expect to resume the trace from
-  TempEnableTrace(resume_pc);
+  // this code is very similar to the above, todo:? make into function
+  // TempEnableTrace(resume_pc);
+  set_pc((mem_loc_t)resume_pc);
+  SimpleCompiler compiler(buffer);
+  // the "normal" return address will be set to ris when this returns from the temp disabled region
+  auto wb = compiler.TestRegister((mem_loc_t)&red_asm_jump_rsi_resume_trace, RSI, (register_t)resume_pc, &merge_block_stack.back());
+  auto written = compiler.finalize();
+  wb.replace_stump<uint64_t>(0xfafafafafafafafa, written.getRawBuffer());
+  write_interrupt_block();
 }
 
 extern "C" void red_asm_start_nested_trace();
@@ -795,7 +812,10 @@ void Tracer::JumpToNestedLoop(void *nested_trace_id) {
 }
 
 void* Tracer::ReplaceIsTracedCall() {
-  assert(current_not_traced_call_addr == (mem_loc_t)&redmagic_is_traced);
+  // if this wasn't the most recent call then don't delete
+  // also returning null will make this statement behave in a false way
+  if(current_not_traced_call_addr != (mem_loc_t)&redmagic_is_traced)
+    return NULL;
   assert(icount - last_call_instruction < 2);
   buffer->setOffset(last_call_generated_op);
   SimpleCompiler compiler(buffer);

src/user_interface.c

@@ -20,7 +20,9 @@ __asm__("jmp_rax: \n"
           "redmagic_" #method ": \n"                 \
           "movq 0(%rsp), %rsi \n"                    \
           "movq %rsp, %rdx \n"                       \
+          "pushq $0 \n"                              \
           "call red_user_" #method "@plt \n"         \
+          "popq %rdi \n"                             \
           "cmp $5, %rax \n"                          \
           "jg jmp_rax \n"                            \
           "ret \n"                                   \