<commit_before>/*
  This file is part of MAMBO, a low-overhead dynamic binary modification tool:
      https://github.com/beehive-lab/mambo

  Copyright 2013-2016 <NAME> <cosmin at linux-geek dot org>

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <assert.h>
#include <string.h>
#include <limits.h>
#include <signal.h>
#include <asm/unistd.h>
#include <linux/sched.h>
#include <pthread.h>

#include <sys/mman.h>
#include <unistd.h>

#include <libelf.h>

#include "dbm.h"
#include "common.h"
#include "scanner_common.h"

#include "elf_loader/elf_loader.h"

#include "pie/pie-thumb-decoder.h"
#include "pie/pie-thumb-encoder.h"

#include "pie/pie-arm-encoder.h"

#ifdef DEBUG
  #define debug(...) fprintf(stderr, __VA_ARGS__)
  #ifndef VERBOSE
    #define VERBOSE
  #endif
#else
  #define debug(...)
#endif

#ifdef VERBOSE
  #define info(...) fprintf(stderr, __VA_ARGS__)
#else
  #define info(...)
#endif

#ifdef CC_HUGETLB
  #define CC_PAGE_SIZE (2*1024*1024)
  #define CC_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB)
#else
  #define CC_PAGE_SIZE 4096
  #define CC_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS)
#endif

#ifdef METADATA_HUGETLB
  #define METADATA_PAGE_SIZE (2*1024*1024)
  #define METADATA_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB)
#else
  #define METADATA_PAGE_SIZE 4096
  #define METADATA_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS)
#endif

#define ROUND_UP(input, multiple_of) \
  (((input / multiple_of) * multiple_of) + ((input % multiple_of) ? multiple_of : 0))

#define CC_SZ_ROUND(input) ROUND_UP(input, CC_PAGE_SIZE)
#define METADATA_SZ_ROUND(input) ROUND_UP(input, CC_PAGE_SIZE)

dbm_global global_data;
__thread dbm_thread *current_thread;

void flush_code_cache(dbm_thread *thread_data) {
  thread_data->free_block = 2;
  hash_init(&thread_data->entry_address, CODE_CACHE_HASH_SIZE + CODE_CACHE_HASH_OVERP);
#ifdef DBM_TRACES
  thread_data->trace_cache_next = thread_data->code_cache->traces;
  thread_data->trace_id = CODE_CACHE_SIZE;

  hash_init(&thread_data->trace_entry_address, CODE_CACHE_HASH_SIZE + CODE_CACHE_HASH_OVERP);
#endif

  for (int i = 0; i < CODE_CACHE_SIZE; i++) {
    thread_data->code_cache_meta[i].exit_branch_type = unknown;
    thread_data->code_cache_meta[i].linked_from = NULL;
    thread_data->code_cache_meta[i].branch_cache_status = 0;
#ifdef DBM_TRACES
    thread_data->exec_count[i] = 0;
#endif
  }

  linked_list_init(thread_data->cc_links, MAX_CC_LINKS);
}

void *dbm_start_thread_pth(void *ptr) {
  dbm_thread *thread_data = (dbm_thread *)ptr;
  assert(thread_data->clone_args->child_stack);

  current_thread = thread_data;
  uint32_t addr = scan(thread_data, thread_data->clone_ret_addr, ALLOCATE_BB);
  uint32_t tid = syscall(__NR_gettid);

  if (thread_data->clone_args->flags & CLONE_PARENT_SETTID) {
    *thread_data->clone_args->ptid = tid;
  }
  if (thread_data->clone_args->flags & CLONE_CHILD_SETTID) {
    *thread_data->clone_args->ctid = tid;
  }
  if (thread_data->clone_args->flags & CLONE_CHILD_CLEARTID) {
		syscall(__NR_set_tid_address, thread_data->clone_args->ctid);
  }
  thread_data->tls = thread_data->clone_args->tls;

  // Copy the parent's saved register values to the child's stack
  uint32_t *child_stack = thread_data->clone_args->child_stack;
  child_stack -= 15; // reserve 15 words on the child's stack
  mambo_memcpy(child_stack, thread_data->clone_args, sizeof(uint32_t) * 14);
  child_stack[r0] = 0; // return 0
  child_stack[r8] = thread_data->scratch_regs[0];
  child_stack[r9] = thread_data->scratch_regs[1];
  child_stack[13] = thread_data->scratch_regs[2]; // R14
  child_stack[14] = addr; // pc

  // Release the lock
  __asm__ volatile("dmb");
  thread_data->tid = tid;

  th_enter(child_stack);
  return NULL;
}

dbm_thread *dbm_create_thread(dbm_thread *thread_data, void *next_inst, sys_clone_args *args) {
  pthread_t thread;
  dbm_thread *new_thread_data;

  if (!allocate_thread_data(&new_thread_data)) {
    fprintf(stderr, "Failed to allocate thread data\n");
    while(1);
  }
  init_thread(new_thread_data);
  new_thread_data->clone_ret_addr = next_inst;
  new_thread_data->tid = 0;
  new_thread_data->clone_args = args;
  for (int i = 0; i < 3; i++) {
    new_thread_data->scratch_regs[i] = thread_data->scratch_regs[i];
  }

  pthread_attr_t attr;
  pthread_attr_init(&attr);
  pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
  pthread_attr_setinheritsched(&attr, PTHREAD_INHERIT_SCHED);
  /* We're switching to the stack allocated by the application immediately, so make this
     as small as possible. Our glibc stores data here, so we can't unmap it.
     Also see man pthread_attr_setguardsize BUGS. */
  pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 4096);
  pthread_attr_setguardsize(&attr, 4096);
  pthread_create(&thread, &attr, dbm_start_thread_pth, new_thread_data);

  return new_thread_data;
}

void mambo_deliver_callbacks(unsigned cb_id, dbm_thread *thread_data, inst_set inst_type,
                             cc_type fragment_type, int fragment_id, int inst, mambo_cond cond,
                             void *read_address, void *write_p, unsigned long *regs) {
#ifdef PLUGINS_NEW
  mambo_context ctx;

  assert(cb_id < CALLBACK_MAX_IDX);

  if (global_data.free_plugin > 0) {
    set_mambo_context(&ctx, thread_data, inst_type, fragment_type,
                      fragment_id, inst, cond, read_address, write_p, regs);
    for (int i = 0; i < global_data.free_plugin; i++) {
      if (global_data.plugins[i].cbs[cb_id] != NULL) {
        ctx.plugin_id = i;
        global_data.plugins[i].cbs[cb_id](&ctx);
      } // if
    } // for
  }
#endif
}

// return 0 to skip the syscall
int syscall_handler_pre(uint32_t syscall_no, uint32_t *args, uint16_t *next_inst, dbm_thread *thread_data) {
  struct sigaction *sig_action;
  sys_clone_args *clone_args;
  debug("syscall pre %d\n", syscall_no);

#ifdef PLUGINS_NEW
  mambo_context ctx;
  int cont;

  if (global_data.free_plugin > 0) {
    set_mambo_context(&ctx, thread_data, -1, -1, -1, -1, -1, NULL, NULL, (unsigned long *)args);
    for (int i = 0; i < global_data.free_plugin; i++) {
      if (global_data.plugins[i].cbs[PRE_SYSCALL_C] != NULL) {
        ctx.plugin_id = i;
        cont = global_data.plugins[i].cbs[PRE_SYSCALL_C](&ctx);
        if (!cont) return 0;
      } // if
    } // for
  }
#endif

  switch(syscall_no) {
    case SYSCALL_CLONE:
      clone_args = (sys_clone_args *)args;

      if (clone_args->flags & CLONE_VM) {
        if (!(clone_args->flags & CLONE_SETTLS)) {
          clone_args->tls = thread_data->tls;
        }
        thread_data->clone_vm = true;

        dbm_thread *child_data = dbm_create_thread(thread_data, next_inst, clone_args);
        while(child_data->tid == 0);
        args[0] = child_data->tid;

        return 0;
      } else {
        thread_data->child_tls = (clone_args->flags & CLONE_SETTLS) ? clone_args->tls : thread_data->tls;
        clone_args->flags &= ~CLONE_SETTLS;

        thread_data->clone_vm = false;
      }
      break;
    case SYSCALL_EXIT:
      debug("thread exit\n");
      mambo_deliver_callbacks(POST_THREAD_C, thread_data, -1, -1, -1, -1, -1, NULL, NULL, NULL);

      if (munmap(thread_data->code_cache, CC_SZ_ROUND(sizeof(dbm_code_cache))) != 0) {
        fprintf(stderr, "Error freeing code cache on exit()\n");
        while(1);
      }
      if (munmap(thread_data, METADATA_SZ_ROUND(sizeof(dbm_thread))) != 0) {
        fprintf(stderr, "Error freeing thread private structure on exit()\n");
        while(1);
      }
      pthread_exit(NULL); // this should never return
      while(1); 
      break;
    case SYSCALL_RT_SIGACTION:
      debug("sigaction %d\n", args[0]);
      debug("struct sigaction at 0x%x\n", args[1]);
      sig_action = (struct sigaction *)args[1];
      // If act is non-NULL, the new action for signal signum is installed from act. If oldact is non-NULL, the previous action is saved in oldact.
      debug("handler at %p\n", sig_action->sa_handler);
      if (sig_action
          && sig_action->sa_handler != SIG_IGN
          && sig_action->sa_handler != SIG_DFL) {
        sig_action->sa_handler = (void *)lookup_or_scan(thread_data, (uint32_t)sig_action->sa_handler, NULL);
      }
      break;
    case SYSCALL_EXIT_GROUP:
      dbm_exit(thread_data, args[0]);
      break;
    case SYSCALL_CLOSE:
      if (args[0] <= 2) { // stdin, stdout, stderr
        args[0] = 0;
        return 0;
      }
      break;
    case __ARM_NR_cacheflush:
      /* Returning to the calling BB is potentially unsafe because the remaining
         contents of the BB or other basic blocks it is linked against could be stale */
      flush_code_cache(thread_data);
      break;
    case __ARM_NR_set_tls:
      debug("set tls to %x\n", args[0]);
      thread_data->tls = args[0];
      args[0] = 0;
      return 0;
      break;
    /* Remove the execute permission from application mappings. At this point, this mostly acts
       as a safeguard in case a translation bug causes a branch to unmodified application code.
       Page permissions happen to be passed in the third argument both for mmap and mprotect. */
    case __NR_mmap2:
    case __NR_mprotect:
      /* Ensure that code pages are readable by the code scanner. */
      if (args[2] & PROT_EXEC) {
        assert(args[2] & PROT_READ);
      }
      args[2] &= ~PROT_EXEC;
      break;

    case __NR_munmap:
      flush_code_cache(thread_data);
      break;

    case __NR_vfork:
      assert(thread_data->is_vfork_child == false);
      thread_data->is_vfork_child = true;
      for (int i = 0; i < 3; i++) {
        thread_data->parent_scratch_regs[i] = thread_data->scratch_regs[i];
      }
      break;
  }
  
  return 1;
}

uint32_t syscall_handler_post(uint32_t syscall_no, uint32_t *args, uint16_t *next_inst, dbm_thread *thread_data) {
  dbm_thread *new_thread_data;
  uint32_t addr = 0;
  
  debug("syscall post %d\n", syscall_no);

  switch(syscall_no) {
    case SYSCALL_CLONE:
      debug("r0 (tid): %d\n", args[0]);
      if (args[0] == 0) { // the child
        if (thread_data->clone_vm) {
          debug("target: %p\n", next_inst);
          if (!allocate_thread_data(&new_thread_data)) {
            fprintf(stderr, "Failed to allocate thread data\n");
            while(1);
          }
          init_thread(new_thread_data);
          addr = scan(new_thread_data, next_inst, ALLOCATE_BB);
          new_thread_data->tls = thread_data->child_tls;
          /* There are a few race conditions in this implementation, which should be addressed.
             However, this code path is not used at the moment. We are using ptrace_create.
             TODO:
             * block the parent
             * copy all shared state to the child's private data (sr_regs, next_inst, th->child_tls)
               * args should be safe, they're pushed on the thread's stack
             * unblock the parent
          */
          assert(0);
        } else {
          /* Without CLONE_VM, the child runs in a separate memory space,
             no synchronisation is needed.*/
          thread_data->tls = thread_data->child_tls;
        }
      }
      break;

    case __NR_vfork:
      if (args[0] != 0) { // in the parent
        for (int i = 0; i < 3; i++) {
          thread_data->scratch_regs[i] = thread_data->parent_scratch_regs[i];
        }
        thread_data->is_vfork_child = false;
      }
      break;
  }

  mambo_deliver_callbacks(POST_SYSCALL_C, thread_data, -1, -1, -1, -1, -1, NULL, NULL, (unsigned long *)args);

  return addr;
}

void dispatcher(uint32_t target, uint32_t *next_addr, uint32_t source_index, dbm_thread *thread_data) {
  uint32_t block_address;
  uint32_t other_target;
  uint32_t pred_target;
  uint16_t *branch_addr; 
  uint32_t *branch_table;
  bool     cached;
  bool     other_target_in_cache;
  int      cache_index;
  uint8_t  *table;
  branch_type source_branch_type;
  uint8_t  sr[2];
  uint32_t reglist;
  uint32_t *ras;
  bool use_bx_lr;
  int unwind_len;
  bool is_taken;

/* It's essential to copy exit_branch_type before calling lookup_or_scan
     because when scanning a stub basic block the source block and its
     meta-information get overwritten */
  debug("Source block index: %d\n", source_index);
  source_branch_type = thread_data->code_cache_meta[source_index].exit_branch_type;

#ifdef DBM_TRACES
  // Handle trace exits separately
  if (source_index >= CODE_CACHE_SIZE && source_branch_type != tbb && source_branch_type != tbh) {
    return trace_dispatcher(target, next_addr, source_index, thread_data);
  }
#endif
  
  debug("Reached the dispatcher, target: 0x%x, ret: %p, src: %d thr: %p\n", target, next_addr, source_index, thread_data);
  block_address = lookup_or_scan(thread_data, target, &cached);
  if (cached) {
    debug("Found block from %d for 0x%x in cache at 0x%x\n", source_index, target, block_address);
  } else {
    debug("Scanned at 0x%x for 0x%x\n", block_address, target);
  }
  
  switch (source_branch_type) {
#ifdef DBM_TB_DIRECT
    case tbb:
    case tbh:
      /* the index is invalid only when the inline hash lookup is called for a new BB,
         no linking is required */
#ifdef FAST_BT
      if (thread_data->code_cache_meta[source_index].rn >= TB_CACHE_SIZE) {
        break;
      }
#else
      if (thread_data->code_cache_meta[source_index].rn >= MAX_TB_INDEX) {
        break;
      }
#endif
      //thread_data->code_cache_meta[source_index].count++;
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
#ifdef FAST_BT
      branch_table = (uint32_t *)(((uint32_t)branch_addr + 20 + 2) & 0xFFFFFFFC);
      branch_table[thread_data->code_cache_meta[source_index].rn] = block_address;
#else
      branch_addr += 7;
      table = (uint8_t *)branch_addr;
      if (thread_data->code_cache_meta[source_index].free_b == TB_CACHE_SIZE) {
        // if the list of linked blocks is full, link this index to the inline hash lookup
  #ifdef DBM_D_INLINE_HASH
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + TB_CACHE_SIZE * 2 + 1;
  #else
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + TB_CACHE_SIZE * 2;
  #endif
      } else {
        // allocate a branch slot and link it
        cache_index = thread_data->code_cache_meta[source_index].free_b++;
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + cache_index * 2;
        
        // insert the branch to the target BB
        branch_addr += MAX_TB_INDEX / 2 + cache_index * 2;
        thumb_cc_branch(thread_data, branch_addr, (uint32_t)block_address);
        __clear_cache(branch_addr, branch_addr + 5);
      }
  #endif
      
      // invalidate the saved rm value - required to detect calls from the inline hash lookup
      thread_data->code_cache_meta[source_index].rn = INT_MAX;

      break;
#endif
#ifdef DBM_LINK_UNCOND_IMM
    case uncond_imm_thumb:
    case uncond_b_to_bl_thumb:
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
      if (block_address & 0x1) {
        if (source_branch_type == uncond_b_to_bl_thumb) {
          thumb_b32_helper(branch_addr, (uint32_t)block_address);
        } else {
          thumb_cc_branch(thread_data, branch_addr, (uint32_t)block_address);
        }
        __clear_cache((char *)branch_addr-1, (char *)(branch_addr) + 8);
      } else {
        // The data word used for the address is word-aligned
        if (((uint32_t)branch_addr) & 2) {
          thumb_ldrl32(&branch_addr, pc, 4, 1);
          branch_addr += 3;
        } else {
          thumb_ldrl32(&branch_addr, pc, 0, 1);
          branch_addr += 2;
        }
        *(uint32_t *)branch_addr = block_address;
        __clear_cache((char *)branch_addr-7, (char *)branch_addr);
      }
      break;
#endif
#ifdef DBM_LINK_COND_IMM
    case cond_imm_arm:
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
      is_taken = target == thread_data->code_cache_meta[source_index].branch_taken_addr;
      if (is_taken) {
        other_target = hash_lookup(&thread_data->entry_address, thread_data->code_cache_meta[source_index].branch_skipped_addr);
        other_target_in_cache = (other_target != UINT_MAX);

        if (thread_data->code_cache_meta[source_index].branch_cache_status & 1) {
          branch_addr += 2;
        }

        arm_cc_branch(thread_data, (uint32_t *)branch_addr, (uint32_t)block_address,
                      thread_data->code_cache_meta[source_index].branch_condition);
      } else {
        other_target = hash_lookup(&thread_data->entry_address, thread_data->code_cache_meta[source_index].branch_taken_addr);
        other_target_in_cache = (other_target != UINT_MAX);

        if (thread_data->code_cache_meta[source_index].branch_cache_status & 2) {
          branch_addr += 2;
        }

        arm_cc_branch(thread_data, (uint32_t *)branch_addr, (uint32_t)block_address,
                       arm_inverse_cond_code[thread_data->code_cache_meta[source_index].branch_condition]);<commit_msg>ARM: add linking for unconditional direct branches<commit_after>/*
  This file is part of MAMBO, a low-overhead dynamic binary modification tool:
      https://github.com/beehive-lab/mambo

  Copyright 2013-2016 <NAME> <cosmin at linux-geek dot org>

  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <assert.h>
#include <string.h>
#include <limits.h>
#include <signal.h>
#include <asm/unistd.h>
#include <linux/sched.h>
#include <pthread.h>

#include <sys/mman.h>
#include <unistd.h>

#include <libelf.h>

#include "dbm.h"
#include "common.h"
#include "scanner_common.h"

#include "elf_loader/elf_loader.h"

#include "pie/pie-thumb-decoder.h"
#include "pie/pie-thumb-encoder.h"

#include "pie/pie-arm-encoder.h"

#ifdef DEBUG
  #define debug(...) fprintf(stderr, __VA_ARGS__)
  #ifndef VERBOSE
    #define VERBOSE
  #endif
#else
  #define debug(...)
#endif

#ifdef VERBOSE
  #define info(...) fprintf(stderr, __VA_ARGS__)
#else
  #define info(...)
#endif

#ifdef CC_HUGETLB
  #define CC_PAGE_SIZE (2*1024*1024)
  #define CC_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB)
#else
  #define CC_PAGE_SIZE 4096
  #define CC_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS)
#endif

#ifdef METADATA_HUGETLB
  #define METADATA_PAGE_SIZE (2*1024*1024)
  #define METADATA_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS|MAP_HUGETLB)
#else
  #define METADATA_PAGE_SIZE 4096
  #define METADATA_MMAP_OPTS (MAP_PRIVATE|MAP_ANONYMOUS)
#endif

#define ROUND_UP(input, multiple_of) \
  (((input / multiple_of) * multiple_of) + ((input % multiple_of) ? multiple_of : 0))

#define CC_SZ_ROUND(input) ROUND_UP(input, CC_PAGE_SIZE)
#define METADATA_SZ_ROUND(input) ROUND_UP(input, CC_PAGE_SIZE)

dbm_global global_data;
__thread dbm_thread *current_thread;

void flush_code_cache(dbm_thread *thread_data) {
  thread_data->free_block = 2;
  hash_init(&thread_data->entry_address, CODE_CACHE_HASH_SIZE + CODE_CACHE_HASH_OVERP);
#ifdef DBM_TRACES
  thread_data->trace_cache_next = thread_data->code_cache->traces;
  thread_data->trace_id = CODE_CACHE_SIZE;

  hash_init(&thread_data->trace_entry_address, CODE_CACHE_HASH_SIZE + CODE_CACHE_HASH_OVERP);
#endif

  for (int i = 0; i < CODE_CACHE_SIZE; i++) {
    thread_data->code_cache_meta[i].exit_branch_type = unknown;
    thread_data->code_cache_meta[i].linked_from = NULL;
    thread_data->code_cache_meta[i].branch_cache_status = 0;
#ifdef DBM_TRACES
    thread_data->exec_count[i] = 0;
#endif
  }

  linked_list_init(thread_data->cc_links, MAX_CC_LINKS);
}

void *dbm_start_thread_pth(void *ptr) {
  dbm_thread *thread_data = (dbm_thread *)ptr;
  assert(thread_data->clone_args->child_stack);

  current_thread = thread_data;
  uint32_t addr = scan(thread_data, thread_data->clone_ret_addr, ALLOCATE_BB);
  uint32_t tid = syscall(__NR_gettid);

  if (thread_data->clone_args->flags & CLONE_PARENT_SETTID) {
    *thread_data->clone_args->ptid = tid;
  }
  if (thread_data->clone_args->flags & CLONE_CHILD_SETTID) {
    *thread_data->clone_args->ctid = tid;
  }
  if (thread_data->clone_args->flags & CLONE_CHILD_CLEARTID) {
		syscall(__NR_set_tid_address, thread_data->clone_args->ctid);
  }
  thread_data->tls = thread_data->clone_args->tls;

  // Copy the parent's saved register values to the child's stack
  uint32_t *child_stack = thread_data->clone_args->child_stack;
  child_stack -= 15; // reserve 15 words on the child's stack
  mambo_memcpy(child_stack, thread_data->clone_args, sizeof(uint32_t) * 14);
  child_stack[r0] = 0; // return 0
  child_stack[r8] = thread_data->scratch_regs[0];
  child_stack[r9] = thread_data->scratch_regs[1];
  child_stack[13] = thread_data->scratch_regs[2]; // R14
  child_stack[14] = addr; // pc

  // Release the lock
  __asm__ volatile("dmb");
  thread_data->tid = tid;

  th_enter(child_stack);
  return NULL;
}

dbm_thread *dbm_create_thread(dbm_thread *thread_data, void *next_inst, sys_clone_args *args) {
  pthread_t thread;
  dbm_thread *new_thread_data;

  if (!allocate_thread_data(&new_thread_data)) {
    fprintf(stderr, "Failed to allocate thread data\n");
    while(1);
  }
  init_thread(new_thread_data);
  new_thread_data->clone_ret_addr = next_inst;
  new_thread_data->tid = 0;
  new_thread_data->clone_args = args;
  for (int i = 0; i < 3; i++) {
    new_thread_data->scratch_regs[i] = thread_data->scratch_regs[i];
  }

  pthread_attr_t attr;
  pthread_attr_init(&attr);
  pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_DETACHED);
  pthread_attr_setinheritsched(&attr, PTHREAD_INHERIT_SCHED);
  /* We're switching to the stack allocated by the application immediately, so make this
     as small as possible. Our glibc stores data here, so we can't unmap it.
     Also see man pthread_attr_setguardsize BUGS. */
  pthread_attr_setstacksize(&attr, PTHREAD_STACK_MIN + 4096);
  pthread_attr_setguardsize(&attr, 4096);
  pthread_create(&thread, &attr, dbm_start_thread_pth, new_thread_data);

  return new_thread_data;
}

void mambo_deliver_callbacks(unsigned cb_id, dbm_thread *thread_data, inst_set inst_type,
                             cc_type fragment_type, int fragment_id, int inst, mambo_cond cond,
                             void *read_address, void *write_p, unsigned long *regs) {
#ifdef PLUGINS_NEW
  mambo_context ctx;

  assert(cb_id < CALLBACK_MAX_IDX);

  if (global_data.free_plugin > 0) {
    set_mambo_context(&ctx, thread_data, inst_type, fragment_type,
                      fragment_id, inst, cond, read_address, write_p, regs);
    for (int i = 0; i < global_data.free_plugin; i++) {
      if (global_data.plugins[i].cbs[cb_id] != NULL) {
        ctx.plugin_id = i;
        global_data.plugins[i].cbs[cb_id](&ctx);
      } // if
    } // for
  }
#endif
}

// return 0 to skip the syscall
int syscall_handler_pre(uint32_t syscall_no, uint32_t *args, uint16_t *next_inst, dbm_thread *thread_data) {
  struct sigaction *sig_action;
  sys_clone_args *clone_args;
  debug("syscall pre %d\n", syscall_no);

#ifdef PLUGINS_NEW
  mambo_context ctx;
  int cont;

  if (global_data.free_plugin > 0) {
    set_mambo_context(&ctx, thread_data, -1, -1, -1, -1, -1, NULL, NULL, (unsigned long *)args);
    for (int i = 0; i < global_data.free_plugin; i++) {
      if (global_data.plugins[i].cbs[PRE_SYSCALL_C] != NULL) {
        ctx.plugin_id = i;
        cont = global_data.plugins[i].cbs[PRE_SYSCALL_C](&ctx);
        if (!cont) return 0;
      } // if
    } // for
  }
#endif

  switch(syscall_no) {
    case SYSCALL_CLONE:
      clone_args = (sys_clone_args *)args;

      if (clone_args->flags & CLONE_VM) {
        if (!(clone_args->flags & CLONE_SETTLS)) {
          clone_args->tls = thread_data->tls;
        }
        thread_data->clone_vm = true;

        dbm_thread *child_data = dbm_create_thread(thread_data, next_inst, clone_args);
        while(child_data->tid == 0);
        args[0] = child_data->tid;

        return 0;
      } else {
        thread_data->child_tls = (clone_args->flags & CLONE_SETTLS) ? clone_args->tls : thread_data->tls;
        clone_args->flags &= ~CLONE_SETTLS;

        thread_data->clone_vm = false;
      }
      break;
    case SYSCALL_EXIT:
      debug("thread exit\n");
      mambo_deliver_callbacks(POST_THREAD_C, thread_data, -1, -1, -1, -1, -1, NULL, NULL, NULL);

      if (munmap(thread_data->code_cache, CC_SZ_ROUND(sizeof(dbm_code_cache))) != 0) {
        fprintf(stderr, "Error freeing code cache on exit()\n");
        while(1);
      }
      if (munmap(thread_data, METADATA_SZ_ROUND(sizeof(dbm_thread))) != 0) {
        fprintf(stderr, "Error freeing thread private structure on exit()\n");
        while(1);
      }
      pthread_exit(NULL); // this should never return
      while(1); 
      break;
    case SYSCALL_RT_SIGACTION:
      debug("sigaction %d\n", args[0]);
      debug("struct sigaction at 0x%x\n", args[1]);
      sig_action = (struct sigaction *)args[1];
      // If act is non-NULL, the new action for signal signum is installed from act. If oldact is non-NULL, the previous action is saved in oldact.
      debug("handler at %p\n", sig_action->sa_handler);
      if (sig_action
          && sig_action->sa_handler != SIG_IGN
          && sig_action->sa_handler != SIG_DFL) {
        sig_action->sa_handler = (void *)lookup_or_scan(thread_data, (uint32_t)sig_action->sa_handler, NULL);
      }
      break;
    case SYSCALL_EXIT_GROUP:
      dbm_exit(thread_data, args[0]);
      break;
    case SYSCALL_CLOSE:
      if (args[0] <= 2) { // stdin, stdout, stderr
        args[0] = 0;
        return 0;
      }
      break;
    case __ARM_NR_cacheflush:
      /* Returning to the calling BB is potentially unsafe because the remaining
         contents of the BB or other basic blocks it is linked against could be stale */
      flush_code_cache(thread_data);
      break;
    case __ARM_NR_set_tls:
      debug("set tls to %x\n", args[0]);
      thread_data->tls = args[0];
      args[0] = 0;
      return 0;
      break;
    /* Remove the execute permission from application mappings. At this point, this mostly acts
       as a safeguard in case a translation bug causes a branch to unmodified application code.
       Page permissions happen to be passed in the third argument both for mmap and mprotect. */
    case __NR_mmap2:
    case __NR_mprotect:
      /* Ensure that code pages are readable by the code scanner. */
      if (args[2] & PROT_EXEC) {
        assert(args[2] & PROT_READ);
      }
      args[2] &= ~PROT_EXEC;
      break;

    case __NR_munmap:
      flush_code_cache(thread_data);
      break;

    case __NR_vfork:
      assert(thread_data->is_vfork_child == false);
      thread_data->is_vfork_child = true;
      for (int i = 0; i < 3; i++) {
        thread_data->parent_scratch_regs[i] = thread_data->scratch_regs[i];
      }
      break;
  }
  
  return 1;
}

uint32_t syscall_handler_post(uint32_t syscall_no, uint32_t *args, uint16_t *next_inst, dbm_thread *thread_data) {
  dbm_thread *new_thread_data;
  uint32_t addr = 0;
  
  debug("syscall post %d\n", syscall_no);

  switch(syscall_no) {
    case SYSCALL_CLONE:
      debug("r0 (tid): %d\n", args[0]);
      if (args[0] == 0) { // the child
        if (thread_data->clone_vm) {
          debug("target: %p\n", next_inst);
          if (!allocate_thread_data(&new_thread_data)) {
            fprintf(stderr, "Failed to allocate thread data\n");
            while(1);
          }
          init_thread(new_thread_data);
          addr = scan(new_thread_data, next_inst, ALLOCATE_BB);
          new_thread_data->tls = thread_data->child_tls;
          /* There are a few race conditions in this implementation, which should be addressed.
             However, this code path is not used at the moment. We are using ptrace_create.
             TODO:
             * block the parent
             * copy all shared state to the child's private data (sr_regs, next_inst, th->child_tls)
               * args should be safe, they're pushed on the thread's stack
             * unblock the parent
          */
          assert(0);
        } else {
          /* Without CLONE_VM, the child runs in a separate memory space,
             no synchronisation is needed.*/
          thread_data->tls = thread_data->child_tls;
        }
      }
      break;

    case __NR_vfork:
      if (args[0] != 0) { // in the parent
        for (int i = 0; i < 3; i++) {
          thread_data->scratch_regs[i] = thread_data->parent_scratch_regs[i];
        }
        thread_data->is_vfork_child = false;
      }
      break;
  }

  mambo_deliver_callbacks(POST_SYSCALL_C, thread_data, -1, -1, -1, -1, -1, NULL, NULL, (unsigned long *)args);

  return addr;
}

void dispatcher(uint32_t target, uint32_t *next_addr, uint32_t source_index, dbm_thread *thread_data) {
  uint32_t block_address;
  uint32_t other_target;
  uint32_t pred_target;
  uint16_t *branch_addr; 
  uint32_t *branch_table;
  bool     cached;
  bool     other_target_in_cache;
  int      cache_index;
  uint8_t  *table;
  branch_type source_branch_type;
  uint8_t  sr[2];
  uint32_t reglist;
  uint32_t *ras;
  bool use_bx_lr;
  int unwind_len;
  bool is_taken;

/* It's essential to copy exit_branch_type before calling lookup_or_scan
     because when scanning a stub basic block the source block and its
     meta-information get overwritten */
  debug("Source block index: %d\n", source_index);
  source_branch_type = thread_data->code_cache_meta[source_index].exit_branch_type;

#ifdef DBM_TRACES
  // Handle trace exits separately
  if (source_index >= CODE_CACHE_SIZE && source_branch_type != tbb && source_branch_type != tbh) {
    return trace_dispatcher(target, next_addr, source_index, thread_data);
  }
#endif
  
  debug("Reached the dispatcher, target: 0x%x, ret: %p, src: %d thr: %p\n", target, next_addr, source_index, thread_data);
  block_address = lookup_or_scan(thread_data, target, &cached);
  if (cached) {
    debug("Found block from %d for 0x%x in cache at 0x%x\n", source_index, target, block_address);
  } else {
    debug("Scanned at 0x%x for 0x%x\n", block_address, target);
  }
  
  switch (source_branch_type) {
#ifdef DBM_TB_DIRECT
    case tbb:
    case tbh:
      /* the index is invalid only when the inline hash lookup is called for a new BB,
         no linking is required */
#ifdef FAST_BT
      if (thread_data->code_cache_meta[source_index].rn >= TB_CACHE_SIZE) {
        break;
      }
#else
      if (thread_data->code_cache_meta[source_index].rn >= MAX_TB_INDEX) {
        break;
      }
#endif
      //thread_data->code_cache_meta[source_index].count++;
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
#ifdef FAST_BT
      branch_table = (uint32_t *)(((uint32_t)branch_addr + 20 + 2) & 0xFFFFFFFC);
      branch_table[thread_data->code_cache_meta[source_index].rn] = block_address;
#else
      branch_addr += 7;
      table = (uint8_t *)branch_addr;
      if (thread_data->code_cache_meta[source_index].free_b == TB_CACHE_SIZE) {
        // if the list of linked blocks is full, link this index to the inline hash lookup
  #ifdef DBM_D_INLINE_HASH
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + TB_CACHE_SIZE * 2 + 1;
  #else
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + TB_CACHE_SIZE * 2;
  #endif
      } else {
        // allocate a branch slot and link it
        cache_index = thread_data->code_cache_meta[source_index].free_b++;
        table[thread_data->code_cache_meta[source_index].rn] = MAX_TB_INDEX / 2 + cache_index * 2;
        
        // insert the branch to the target BB
        branch_addr += MAX_TB_INDEX / 2 + cache_index * 2;
        thumb_cc_branch(thread_data, branch_addr, (uint32_t)block_address);
        __clear_cache(branch_addr, branch_addr + 5);
      }
  #endif
      
      // invalidate the saved rm value - required to detect calls from the inline hash lookup
      thread_data->code_cache_meta[source_index].rn = INT_MAX;

      break;
#endif
#ifdef DBM_LINK_UNCOND_IMM
    case uncond_imm_thumb:
    case uncond_b_to_bl_thumb:
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
      if (block_address & 0x1) {
        if (source_branch_type == uncond_b_to_bl_thumb) {
          thumb_b32_helper(branch_addr, (uint32_t)block_address);
        } else {
          thumb_cc_branch(thread_data, branch_addr, (uint32_t)block_address);
        }
        __clear_cache((char *)branch_addr-1, (char *)(branch_addr) + 8);
      } else {
        // The data word used for the address is word-aligned
        if (((uint32_t)branch_addr) & 2) {
          thumb_ldrl32(&branch_addr, pc, 4, 1);
          branch_addr += 3;
        } else {
          thumb_ldrl32(&branch_addr, pc, 0, 1);
          branch_addr += 2;
        }
        *(uint32_t *)branch_addr = block_address;
        __clear_cache((char *)branch_addr-7, (char *)branch_addr);
      }
      break;

    case uncond_imm_arm:
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
      arm_b32_helper((uint32_t *)branch_addr, (uint32_t)block_address, AL);
      __clear_cache(branch_addr, (char *)branch_addr+5);
      break;
#endif
#ifdef DBM_LINK_COND_IMM
    case cond_imm_arm:
      branch_addr = thread_data->code_cache_meta[source_index].exit_branch_addr;
      is_taken = target == thread_data->code_cache_meta[source_index].branch_taken_addr;
      if (is_taken) {
        other_target = hash_lookup(&thread_data->entry_address, thread_data->code_cache_meta[source_index].branch_skipped_addr);
        other_target_in_cache = (other_target != UINT_MAX);

        if (thread_data->code_cache_meta[source_index].branch_cache_status & 1) {
          branch_addr += 2;
        }

        arm_cc_branch(thread_data, (uint32_t *)branch_addr, (uint32_t)block_address,
                      thread_data->code_cache_meta[source_index].branch_condition);
      } else {
        other_target = hash_lookup(&thread_data->entry_address, thread_data->code_cache_meta[source_index].branch_taken_addr);
        other_target_in_cache = (other_target != UINT_MAX);

        if (thread_data->code_cache_meta[source_index].branch_cache_status & 2) {
          branch_addr += 2;
        }

        arm_cc_branch(thread_data, (uint32_t *)branch_addr, (uint32_t)block_address,
                       arm_inverse_cond_code[thread_data->code_cache_meta[source_index].branch_condition]);
{'fraction_non_alphanumeric': 0.10918505177847816, 'fraction_numerical': 0.014323502926609635, 'mean_word_length': 3.3427838201148723, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 240, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1800288', 'n_tokens_mistral': 12843, 'n_tokens_neox': 12257, 'n_words': 2779}