should handle tail calls using jumps and patching address in for bran…

…ches in traces

src/jit_internal.h

@@ -62,6 +62,12 @@ namespace redmagic {
     void *temp_disable(void *resume_pc);
     void *temp_enable(void *resume_pc);
 
+    void* is_traced_call();
+
+    void do_not_trace_method(void *addr);
+
+
+
     uint32_t get_thread_id();
 
     tracer_stack_state* push_tracer_stack();

src/manager.cc

@@ -62,6 +62,10 @@ extern "C" void* red_user_temp_disable(void *_, void *ret_addr) {
   //return NULL;
 }
 
+extern "C" void* red_user_is_traced(void *_, void *ret_addr) {
+  return manager->is_traced_call();
+}
+
 extern "C" void* red_user_temp_enable(void *_, void *ret_addr) {
   return manager->temp_enable(ret_addr);
   //assert(0);
@@ -84,16 +88,32 @@ extern "C" void redmagic_start() {
   // assert(!r);
 }
 
+extern "C" void redmagic_do_not_trace_function(void *function_pointer) {
+  manager->do_not_trace_method(function_pointer);
+}
+
 static const char *avoid_inlining_methods[] = {
   // inlining the allocator doesn't really help since it will have a lot of branching
   // in trying to find where there is open memory
   "malloc",
   "free",
+  "cfree", // python is somehow using this?
   "realloc",
   "calloc",
   "exit",
   "abort",
 
+  // we use the dl calls while debugging at least so don't inline them since there might be conflicts
+  "dlopen",
+  "dlclose",
+  "dlsym",
+  "dlmopen",
+  "dlvsym",
+  "dladdr",
+  "dladdr1",
+  "dlinfo",
+
+
   // we don't want to inline ourselves
   // so record the entry functions
   // TODO: don't have to use dladdr to resolve this since they are in the same binary
@@ -131,6 +151,10 @@ Manager::Manager() {
   get_tracer_head();
 }
 
+void Manager::do_not_trace_method(void *addr) {
+  no_trace_methods.insert(addr);
+}
+
 uint32_t Manager::get_thread_id() {
   if(this_thread_id == 0) {
     this_thread_id = ++thread_id_counter;
@@ -347,6 +371,14 @@ void* Manager::temp_enable(void *resume_pc) {
   // return NULL;
 }
 
+void* Manager::is_traced_call() {
+  auto head = get_tracer_head();
+  if(head->is_traced) {
+    return head->tracer->ReplaceIsTracedCall();
+  }
+  return NULL;
+}
+
 uint32_t Manager::tracer_stack_size() {
   return threadl_tracer_stack.size();
 }

src/redmagic.h

@@ -35,6 +35,13 @@ void redmagic_ensure_not_traced(void);
 void redmagic_temp_disable(void);
 void redmagic_temp_enable(void);
 
+
+// for setup to tell the jit to not trace some function, eg the call into a gc or some custom malloc
+void redmagic_do_not_trace_function(void* function_pointer);
+
+// return 0 if not traced, non zero otherwise
+unsigned long redmagic_is_traced(void);
+
 #ifdef __cplusplus
 }
 

src/simple_compiler.cc

@@ -322,8 +322,11 @@ CodeBuffer SimpleCompiler::MakeResumeTraceBlock(mem_loc_t resume_pc) {
 
 void SimpleCompiler::ResumeBlockJump(mem_loc_t resume_pc) {
   restore_registers();
+  auto label_top = newLabel();
   auto label = newLabel();
   bind(label);
+  jmp(label_top); // this should be exactly 5 bytes long with the destination address being the last part
+  bind(label_top);
   mov(x86::ptr(x86::rsp, -TRACE_STACK_OFFSET + 216), x86::r10);
   mov(x86::ptr(x86::rsp, -TRACE_STACK_OFFSET + 224), x86::r9);
   mov(x86::r10, imm_u(resume_pc));
@@ -336,6 +339,24 @@ void SimpleCompiler::ResumeBlockJump(mem_loc_t resume_pc) {
 
   // for identifying which instruction it jumped from
   mov(x86::r9, imm_u(0xfafafafafafafafa));
+
+  // No one is going to generate code after this since we have already used a jump so there is no point
+  // we also know that this will be generated at the bottom since there is no point of generating at the top
+  // thus we use this information to ensure that the address of the jump is aligned to a 4byte boundary which means that it
+  // can be atomically updated (hopefully)
+
+  auto loffset = getLabelOffset(label_top);
+
+  mem_loc_t laddr = (mem_loc_t)(buffer->buffer + buffer->size - buffer->trampolines_size - getCodeSize() + loffset);
+
+  for(int i = laddr & 0x3; i; i--) {
+    nop();
+  }
+
+  mem_loc_t laddr2 = (mem_loc_t)(buffer->buffer + buffer->size - buffer->trampolines_size - getCodeSize() + loffset);
+
+  assert((laddr2 & 0x3) == 0);
+  //assert(0);
 }
 
 

src/tracer.cc

@@ -90,10 +90,16 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
   // the dummy address
   assert(write_jump_address != 0xfafafafafafafafa);
 
+  // that this is a jump with a rel32 term to the next line and is aligned properly
+  assert(*(uint8_t*)write_jump_address == 0xE9);
+  assert(*(int32_t*)(write_jump_address + 1) == 0);
+  assert(((write_jump_address + 1) & 0x3) == 0);
+
+
   assert(regs_struct->rsp - TRACE_STACK_OFFSET == (register_t)regs_struct);
 
   void *ret = NULL;
-  void *patch = NULL;
+  //void *patch = NULL;
 
   auto head = manager->get_tracer_head();
 
@@ -112,11 +118,14 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
   assert(head->tracer == nullptr);
   head->tracer = l;
 
+  head->is_compiled = false;
+
   l->owning_thread = manager->get_thread_id();
 
   // calling Start will invalidate the stack
   ret = l->Start((void*)target_rip);
-  patch = l->get_start_location();
+  //patch = l->get_start_location();
+  l->set_where_to_patch((int32_t*)(write_jump_address + 1));
 
 
   // TODO: need to patch in the address only after the trace is done
@@ -274,6 +283,9 @@ void Tracer::Run(struct user_regs_struct *other_stack) {
     assert(current_location == last_location);
     //assert(generation_lock.owns_lock());
 
+    // TODO: manage cases where the buffer runs out of space
+    assert(buffer->getFree() > 100 * 1024);
+
   processes_instructions:
     while(ud_disassemble(&disassm)) {
 
@@ -353,6 +365,7 @@ void* Tracer::EndTraceFallThrough() {
   compiler.mov(asmjit::x86::rdi, asmjit::imm_u(last_call_ret_addr));
   compiler.jmp(asmjit::imm_ptr(&red_asm_end_trace));
   auto w = compiler.finalize();
+  finish_patch();
   tracing_from = 0;
   return (void*)w.getRawBuffer();
 }
@@ -370,6 +383,7 @@ void* Tracer::EndTraceLoop() {
   buffer->setOffset(last_call_generated_op);
   SimpleCompiler compiler(buffer.get());
   compiler.jmp(asmjit::imm_ptr(loop_location));
+  finish_patch();
   tracing_from = 0;
   return (void*)loop_location;
 }
@@ -415,6 +429,28 @@ void Tracer::JumpToNestedLoop(void *nested_trace_id) {
   written.replace_stump<uint64_t>(0xfafafafafafafafa, resume_addr);
 }
 
+void* Tracer::ReplaceIsTracedCall() {
+  assert(icount - last_call_instruction < 2);
+  buffer->setOffset(last_call_generated_op);
+  SimpleCompiler compiler(buffer.get());
+  compiler.mov(asmjit::x86::rax, asmjit::imm(1));
+  auto written = compiler.finalize();
+  write_interrupt_block();
+
+  mem_loc_t resume_addr = written.getRawBuffer() + written.getOffset();
+
+  return (void*)resume_addr;
+}
+
+void Tracer::finish_patch() {
+  if(finish_patch_addr != nullptr) {
+    mem_loc_t irip = ((mem_loc_t)finish_patch_addr) + 4;
+    int32_t d = trace_start_location - irip;
+    *finish_patch_addr = d;
+    finish_patch_addr = nullptr;
+  }
+}
+
 extern "C" void* red_asm_resume_eval_block(void*, void*);
 
 void Tracer::continue_program(mem_loc_t resume_loc) {
@@ -964,24 +1000,53 @@ void Tracer::evaluate_instruction() {
 
     process_jump_dest:
     assert(jump_dest != 0);
-    if(last_call_instruction + 1 == icount) {
-      // then the first operation in this method was a jump, which means that we were probably jumping through a redirect with the dynamic linker
-      if(!manager->should_trace_method((void*)jump_dest)) {
-        mem_loc_t cont_addr;
-        {
-          buffer->setOffset(last_call_generated_op);
-          pop_stack();
-          set_pc(last_call_ret_addr);
-          SimpleCompiler compiler(buffer.get());
-          compiler.call(asmjit::imm_ptr(jump_dest));
-          auto written = compiler.finalize();
-          write_interrupt_block();
-          cont_addr = written.getRawBuffer();
-        }
-        continue_program(cont_addr);
-        return;
+    if(!manager->should_trace_method((void*)jump_dest)) {
+      mem_loc_t cont_addr;
+      if(last_call_instruction + 1 == icount) {
+        // then the first operation in this method was a jump, which means that we were probably jumping through a redirect with the dynamic linker
+        buffer->setOffset(last_call_generated_op);
+        pop_stack();
+        set_pc(last_call_ret_addr);
+        SimpleCompiler compiler(buffer.get());
+        compiler.call(asmjit::imm_ptr(jump_dest));
+        auto written = compiler.finalize();
+        write_interrupt_block();
+        cont_addr = written.getRawBuffer();
+      } else {
+        // we are jumping to another method and isn't the first instruction, which means that this is behaving like a tail call optimization
+        register_t return_pc = peek_stack();
+        set_pc(return_pc);
+        SimpleCompiler compiler(buffer.get());
+        // pop the previous return address off the stack
+        compiler.add(asmjit::x86::rsp, asmjit::imm(8));
+        compiler.call(asmjit::imm_ptr(jump_dest));
+        auto written = compiler.finalize();
+        write_interrupt_block();
+        cont_addr = written.getRawBuffer();
       }
+      continue_program(cont_addr);
+      return;
+
     }
+
+    // if(last_call_instruction + 1 == icount) {
+    //   // then the first operation in this method was a jump, which means that we were probably jumping through a redirect with the dynamic linker
+    //   if(!manager->should_trace_method((void*)jump_dest)) {
+    //     mem_loc_t cont_addr;
+    //     {
+    //       buffer->setOffset(last_call_generated_op);
+    //       pop_stack();
+    //       set_pc(last_call_ret_addr);
+    //       SimpleCompiler compiler(buffer.get());
+    //       compiler.call(asmjit::imm_ptr(jump_dest));
+    //       auto written = compiler.finalize();
+    //       write_interrupt_block();
+    //       cont_addr = written.getRawBuffer();
+    //     }
+    //     continue_program(cont_addr);
+    //     return;
+    //   }
+    // }
     set_pc(jump_dest);
     return;
   }
@@ -1247,3 +1312,13 @@ void Tracer::abort() {
   }
   continue_program(current_location);
 }
+
+
+void Tracer::run_debugger() {
+  red_printf("\n----ABORT debugger---\n");
+  SimpleCompiler compiler(buffer.get());
+  compiler.int3();
+  //compiler.hlt();
+  compiler.finalize();
+  continue_program(current_location);
+}

src/tracer.h

@@ -38,8 +38,15 @@ namespace redmagic {
     void* TempDisableTrace();
     void TempEnableTrace(void *resume_pc) { set_pc((uint64_t)resume_pc); }
 
+    void* ReplaceIsTracedCall();
+
     inline void *get_start_location() { return (void*)trace_start_location; }
 
+    inline void set_where_to_patch(int32_t *addr) {
+      assert(finish_patch_addr == nullptr);
+      finish_patch_addr = addr;
+    }
+
     inline mem_loc_t get_origional_pc() { return udis_loc; }
 
   public:
@@ -48,7 +55,6 @@ namespace redmagic {
 
   private:
 
-
     void set_pc(uint64_t);
 
     struct jump_instruction_info decode_instruction();
@@ -62,12 +68,23 @@ namespace redmagic {
     // jump back to the normal execution of this program
     void abort();
 
+    // write int3 and switch the stack to that
+    void run_debugger();
+
+    // when done patch the address that needs to link to this trace
+    void finish_patch();
+
     inline register_t pop_stack() {
       register_t r = *((register_t*)((mem_loc_t)regs_struct->rsp + move_stack_by));
       move_stack_by += sizeof(register_t);
       return r;
     }
 
+    inline register_t peek_stack() {
+      register_t r = *((register_t*)((mem_loc_t)regs_struct->rsp + move_stack_by));
+      return r;
+    }
+
     inline void push_stack(register_t v) {
       move_stack_by -= sizeof(register_t);
       *((register_t*)((mem_loc_t)regs_struct->rsp + move_stack_by)) = v;
@@ -138,6 +155,7 @@ namespace redmagic {
     mem_loc_t trace_start_location; // where this current trace begins
     mem_loc_t loop_start_location; // where it should branch the loop back to
 
+    int32_t *finish_patch_addr = nullptr;
 
     // mem_loc_t stack;
 

src/user_interface.c

@@ -45,3 +45,4 @@ USER_INTERFACE(fellthrough_branch);
 USER_INTERFACE(ensure_not_traced);
 USER_INTERFACE(temp_disable);
 USER_INTERFACE(temp_enable);
+USER_INTERFACE(is_traced);