not crashing the python repl

src/jit_internal.h

@@ -117,6 +117,7 @@ namespace redmagic {
   //extern thread_local void *trace_id; // id of current executing trace
   //extern thread_local bool is_traced;
   extern Manager *manager;
+  extern thread_local bool protected_malloc;
 
 
   class CodeBuffer final {

src/manager.cc

@@ -168,6 +168,8 @@ uint32_t Manager::get_thread_id() {
 
 void* Manager::begin_trace(void *id, void *ret_addr) {
 
+  protected_malloc = false;
+
   branch_info *info = &branches[id];
   if(info->disabled)
     return NULL; // do not trace this loop
@@ -211,6 +213,7 @@ void* Manager::begin_trace(void *id, void *ret_addr) {
     ret = l->Start(trace_pc);
     new_head->is_traced = true;
     info->starting_point = l->get_start_location();
+    protected_malloc = true;
   }
   //return NULL;
   return ret;
@@ -225,7 +228,7 @@ void* Manager::end_trace(void *id) {
   assert(head->trace_id == id);
   assert(head->tracer == info->tracer);
   l = head->tracer;
-  head->tracer = nullptr;
+
 
   // ret is going to be the address of the normal execution
   ret = l->EndTraceFallthrough();
@@ -238,11 +241,13 @@ void* Manager::end_trace(void *id) {
   //is_traced = false;
   l->tracing_from.store(0);
 
+  info->tracer = head->tracer = nullptr;
 
   Tracer *expected = nullptr;
   if(!free_tracer_list.compare_exchange_strong(expected, l)) {
     // failled to save the tracer to the free list head
-    delete l;
+    // TODO: the tracer holds the only refernece to the code buff which is in a shared pointer, so can't delete.....
+    //delete l;
   }
 
 
@@ -257,20 +262,26 @@ void* Manager::jump_to_trace(void *id) {
 }
 
 void* Manager::backwards_branch(void *id, void *ret_addr) {
+  // ignore
+  if(id == nullptr)
+    return NULL;
   //return NULL;
   auto head = get_tracer_head();
   if(head->is_traced) {
     if(id == head->trace_id) {
+      auto info = &branches[id];
       assert(!head->is_compiled);
+      assert(head->tracer == info->tracer);
       void *ret = head->tracer->EndTraceLoop();
       head->is_compiled = true;
       Tracer *l = head->tracer;
-      head->tracer = nullptr;
+      info->tracer = head->tracer = nullptr;
 
       Tracer *expected = nullptr;
       if(!free_tracer_list.compare_exchange_strong(expected, l)) {
         // failled to save the tracer to the free list head
-        delete l;
+        // TODO: the tracer holds the only refernece to the code buff which is in a shared pointer, so can't delete.....
+        //delete l;
       }
 
       // We are continuing the loop so there is no need to check the parent stack frame
@@ -321,6 +332,9 @@ void* Manager::backwards_branch(void *id, void *ret_addr) {
 }
 
 void* Manager::fellthrough_branch(void *id) {
+  // ignore
+  if(id == nullptr)
+    return NULL;
   auto head = get_tracer_head();
   if(head->trace_id == id && head->is_traced) {
     return end_trace(id);
@@ -348,22 +362,27 @@ void* Manager::temp_disable(void *resume_pc) {
   assert(!head->is_temp_disabled);
   head->is_temp_disabled = true;
   void *ret = NULL;
-  if(head->is_traced) {
+  assert(!head->is_traced || head->tracer);
+  if(head->tracer) {
+    // this will push the stack
     ret = head->tracer->TempDisableTrace();
+    protected_malloc = false;
+  } else {
+    push_tracer_stack();
   }
-  push_tracer_stack();
   return ret;
 }
 
 void* Manager::temp_enable(void *resume_pc) {
   auto old_head = pop_tracer_stack();
-  assert(!old_head.is_traced);
+  assert(!old_head.is_temp_disabled);
   auto head = get_tracer_head();
   assert(head->is_temp_disabled);
   head->is_temp_disabled = false;
   void *ret = NULL;
-  if(head->is_traced) {
+  if(head->tracer) {
     head->tracer->TempEnableTrace(resume_pc);
+    protected_malloc = true;
   }
   if(head->resume_addr != nullptr) {
     ret = head->resume_addr;
@@ -416,12 +435,14 @@ tracer_stack_state* Manager::push_tracer_stack() {
 tracer_stack_state Manager::pop_tracer_stack() {
   auto r = threadl_tracer_stack.back();
   threadl_tracer_stack.pop_back();
+  assert(!threadl_tracer_stack.empty());
   stack_head = &threadl_tracer_stack.back();
   return r;
 }
 
 tracer_stack_state *Manager::get_tracer_head() {
   if(stack_head == nullptr) {
+    assert(threadl_tracer_stack.capacity() == 0);
     threadl_tracer_stack.reserve(50);
     tracer_stack_state e;
     threadl_tracer_stack.push_back(e);

src/simple_compiler.cc

@@ -41,7 +41,8 @@ CodeBuffer SimpleCompiler::finalize_bottom() {
   //void *start = make();
   void *gen_loc = buffer->buffer + buffer->size - buffer->trampolines_size - size;
   size_t gen_c_size = relocCode(gen_loc);
-  assert(gen_c_size == minSize);
+  assert(minSize == getOffset());
+  assert(gen_c_size == size); // return the full size
 
   buffer->trampolines_size += minSize;
   assert(gen_loc == (void*)(buffer->buffer + buffer->size - buffer->trampolines_size));
@@ -518,10 +519,6 @@ size_t SimpleCompiler::_relocCode(void* _dst, asmjit::Ptr baseAddress) const noe
     }
   }
 
-
-
-
-
   // if (arch == kArchX64)
   //   return (size_t)(dst_end - dst);
   // else

src/tracer.cc

@@ -102,6 +102,8 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
   //void *patch = NULL;
 
   auto head = manager->get_tracer_head();
+  assert(head->trace_id != nullptr);
+  auto info = &manager->branches[head->trace_id];
 
   // remove the free tracer
   // TODO: make this into an actual linked list so that we can keep these tracers around and reuse?
@@ -116,7 +118,8 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
   assert(old_tracing_from == 0);
 
   assert(head->tracer == nullptr);
-  head->tracer = l;
+  assert(info->tracer == nullptr);
+  info->tracer = head->tracer = l;
 
   head->is_compiled = false;
 
@@ -140,6 +143,7 @@ extern "C" void red_resume_trace(mem_loc_t target_rip, mem_loc_t write_jump_addr
 
   assert(ret != NULL);
 
+  protected_malloc = true;
 
   *((register_t*)(regs_struct->rsp - TRACE_RESUME_ADDRESS_OFFSET)) = (register_t)ret;
   //return ret;
@@ -153,6 +157,7 @@ extern "C" void red_set_temp_resume(void *resume_addr) {
   assert(head->resume_addr == nullptr);
   head->resume_addr = resume_addr;
   head->is_temp_disabled = true;
+  manager->push_tracer_stack();
   //temp_disable_resume = resume_addr;
 }
 
@@ -162,8 +167,19 @@ 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);
-  assert(!new_head->is_traced); // TODO: resuming a previously interrupted trace
-  return (void*)normal_end_address;
+  //assert(!new_head->is_traced); // TODO: resuming a previously interrupted trace
+  void *ret = (void*)normal_end_address;
+  if(new_head->is_traced) {
+    if(new_head->tracer) {
+      new_head->tracer->JumpFromNestedLoop((void*)normal_end_address);
+    }
+    assert(new_head->resume_addr);
+    ret = new_head->resume_addr;
+    new_head->resume_addr = nullptr;
+  } else {
+    protected_malloc = false;
+  }
+  return ret;
 }
 
 extern "C" void* red_branch_to_sub_trace(void *resume_addr, void *sub_trace_id, void* target_rip) {
@@ -182,6 +198,10 @@ extern "C" void* red_branch_to_sub_trace(void *resume_addr, void *sub_trace_id,
     // maybe treat this as a temp disabled inner loop
   }
   assert(info->starting_point != nullptr);
+  auto new_head = manager->push_tracer_stack();
+  new_head->is_traced = true;
+  new_head->trace_id = sub_trace_id;
+
   return info->starting_point;
 }
 
@@ -484,7 +504,9 @@ void* Tracer::ReplaceIsTracedCall() {
 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;
+    int64_t dl = trace_start_location - irip;
+    int32_t d = dl;
+    assert(d == dl);
     *finish_patch_addr = d;
     finish_patch_addr = nullptr;
   }
@@ -1381,6 +1403,7 @@ void Tracer::abort() {
     // auto r_cb = compiler.MakeResumeTraceBlock(current_location);
     // compiler.jmp(asmjit::imm_ptr(r_cb.getRawBuffer()));
   }
+  protected_malloc = false;
   continue_program(current_location);
 }
 

src/wrap_malloc.cc

@@ -11,7 +11,7 @@
 // we can't direclty use allocators when we are tracing since it might screw with their internal states
 
 // use mmap and mprotect to wrap all alloc options
-#define MPROTECTED_ALLOC (40*1024)
+//#define MPROTECTED_ALLOC (40*1024)
 
 #define NBUFFERS 6
 

tools/gdb-helper.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+
+# load using source tools/gdb-helper.py
+
+import gdb
+
+
+class TraceJumps(gdb.Command):
+
+    def __init__(self):
+        super().__init__("trace-jumps", gdb.COMMAND_USER)
+
+    def invoke(self, args, from_tty):
+        argv = gdb.string_to_argv(args)
+
+        redmagic_info = gdb.execute('info shared redmagic', to_string=True).split()
+        redmagic_start = int(redmagic_info[7], 16)
+        redmagic_end = int(redmagic_info[8], 16)
+
+        branches_taken = []
+
+        def get_rip():
+            return int(gdb.parse_and_eval('$rip'))
+
+        current_rip = get_rip()
+
+        while True:
+            last_rip = current_rip
+            if redmagic_start < last_rip < redmagic_end:
+                gdb.execute('n', to_string=True)
+                current_rip = get_rip()
+            else:
+                gdb.execute('si', to_string=True)
+                current_rip = get_rip()
+                if not (0 < current_rip - last_rip < 15):
+                    # then we probably have taken a branch or something
+                    li = gdb.execute('x/i {}'.format(last_rip), to_string=True)
+                    if '__tls_get_addr' not in li:
+                        gdb.write(li)
+
+
+TraceJumps()