/*

    IDA trace: PIN tool to communicate with IDA's debugger

*/

#include "idadbg_local.h"
#include "idadbg.h"

//--------------------------------------------------------------------------
// forward declarations
static bool handle_packet(idapin_packet_t *res);
static ssize_t pin_recv(PIN_SOCKET fd, void *buf, size_t n, const char *from_where);
static ssize_t pin_send(PIN_SOCKET fd, const void *buf, size_t n, const char *from_where);
static bool handle_packets(int total, const string &until_packet = "");
inline static void mark_trace_flush_needed(bool pause = true);
static void add_instrumentation(trace_flags_t inst);

//--------------------------------------------------------------------------
// Command line switches
KNOB<int> knob_ida_port(KNOB_MODE_WRITEONCE, "pintool",
    "p", "23946", "Port where IDA Pro is listening for incoming PIN tool's connections");

KNOB<int> knob_debug_mode(KNOB_MODE_WRITEONCE, "pintool",
    "idadbg", "0", "Debug mode");

//--------------------------------------------------------------------------
// Lint suppressions
//
// The following suppressions are because of lint problems parsing PIN header files
// and the usage of the LOG macro.
//
//lint -emacro(10, MSG) Expecting a structure or union
//lint -emacro(40, MSG) Undeclared identifier 'MessageTypeLog'
//lint -emacro(746, MSG) call to function 'on()' not made in the presence of a prototype
//lint -emacro(1013, MSG) Symbol 'on' not a member of class ''
//lint -emacro(1055, MSG) Symbol 'on' undeclared, assumed to return int

//--------------------------------------------------------------------------
// macros
#define MSG(fmt, ...)                                     \
  do                                                      \
  {                                                       \
    if ( debug_tracer > 0 )                               \
    {                                                     \
      char buf[1024];                                     \
      pin_snprintf(buf, sizeof(buf), fmt, ##__VA_ARGS__); \
      fprintf(stderr, "%s", buf);                         \
      LOG(buf);                                           \
    }                                                     \
  }                                                       \
  while ( 0 )

#define DEBUG(fmt, ...)                                   \
  do                                                      \
  {                                                       \
    if ( debug_tracer > 1 )                               \
      MSG(fmt, ##__VA_ARGS__);                            \
  }                                                       \
  while ( 0 )

//--------------------------------------------------------------------------
// global variables

static const char *last_packet = "NONE";
static int debug_tracer = 0;

// Did we started the process? If not, don't report
// any event until the PROCESS_START packet is added
// to the events queue
static bool process_started = false;

// different trace modes
static bool tracing_instruction = true;
static bool tracing_bblock = false;
static bool tracing_routine = false;
static bool tracing_registers = false;

// only record new instructions?
static bool only_new_instructions = false;

// break at the very next instruction?
static bool break_at_next_inst = false;

// sockets
static PIN_SOCKET srv_socket, portno, cli_socket;

// limits to filter what to record
static ADDRINT min_address = BADADDR;
static ADDRINT max_address = BADADDR;

// name of the image to trace
static string image_name = "";

// is it exiting?
static bool process_exit = false;
// condition of the listener thread loop
static bool run_listener = false;
// listener thread identifier
static PIN_THREAD_UID listener_uid;

// 1 million or 100k
static unsigned int enqueue_limit = 1000000;
// remember only the last 1 million/100k instructions
static const unsigned int skip_limit = 1000000;

// semaphore used for pausing the whole application
static PIN_SEMAPHORE run_app_sem;
// semaphore used for pausing only in breakpoints
static PIN_SEMAPHORE bpt_sem;
// listener semaphore; the instrumentation code must wait, over all
// to the listener thread to start running before handling things
// or we will run into race conditions with breakpoints
static PIN_SEMAPHORE listener_sem;
// lock object for handling events
static PIN_LOCK events_lock;

// is the trace full?
static bool trace_flush_needed = false;

// Mistery: in wow64 the app_start_cb callback can be called twice?
static bool app_start_cb_called = false;

// list of active software breakpoints
typedef std::vector<ADDRINT> addrvec_t;
static addrvec_t soft_bpts;
// lock object for handling software breakpoints
static PIN_LOCK soft_locks;

// for displaying $PC when pausing from IDA
static addr_t   last_ea;
static pin_thid last_tid;

// last context saved for displaying register values
// in case tracing_registers == true when pausing from IDA
typedef std::map<THREADID, CONTEXT> threadctx_map_t;
static threadctx_map_t threads_context;

// already enabled instrumentations
static uchar instrumentations = 0;

//--------------------------------------------------------------------------
class janitor_for_pinlock_t
{
protected:
  PIN_LOCK *resource;
public:
  janitor_for_pinlock_t(PIN_LOCK *lock) : resource(lock)
  {
    GetLock(resource, PIN_GetTid());
  }
  ~janitor_for_pinlock_t(void)
  {
    ReleaseLock(resource);
    resource = NULL;
  }
};

//--------------------------------------------------------------------------
struct trc_element_t
{
  THREADID tid;
  ADDRINT  ea;
  pin_tev_type_t type;
  idapin_registers_t regs;

  trc_element_t(THREADID _tid, ADDRINT _ea, pin_tev_type_t _type) : tid(_tid), ea(_ea), type(_type)
  {
    regs.eip = ea;

    regs.eax =
    regs.ebx =
    regs.ecx =
    regs.edx =
    regs.esi =
    regs.edi =
    regs.ebp =
    regs.esp =
#if defined(PIN_64)
    regs.r8 =
    regs.r9 =
    regs.r10 =
    regs.r11 =
    regs.r12 =
    regs.r13 =
    regs.r14 =
    regs.r15 =
#endif
    regs.eflags =
    regs.cs =
    regs.ds =
    regs.es =
    regs.fs =
    regs.gs =
    regs.ss = BADADDR;
  };
};


// recorded instructions
typedef std::deque<trc_element_t> trc_deque_t;
static trc_deque_t trace_addrs;

// Already recorded instructions, those should be skipped if
// only_new_instructions is true.
// NOTE: as we have limited memory in the PIN tool we cannot let it grow
// without limit, we need to remember a maximum number of "skip_limit"
// element(s), or the PIN tool would die because it runs out of memory
typedef std::deque<ADDRINT> addr_deque_t;
static addr_deque_t all_addrs;

// list of all the events that will be passed to IDA
typedef std::deque<pin_debug_event_t> event_list_t;
static event_list_t events;

//--------------------------------------------------------------------------
static void dump_sizes(void)
{
#if 0
  MSG("Sizeof pin_module_info_t %d\n", sizeof(pin_module_info_t));
  MSG("Sizeof pin_e_breakpoint_t %d\n", sizeof(pin_e_breakpoint_t));
  MSG("Sizeof pin_e_exception_t %d\n", sizeof(pin_e_exception_t));
  MSG("Sizeof pin_debug_event_t %d\n", sizeof(pin_debug_event_t));
  MSG("Sizeof idapin_packet_t %d\n", sizeof(idapin_packet_t));
  MSG("Sizeof memimages_pkt_t %d\n", sizeof(memimages_pkt_t));
  MSG("Sizeof pin_memory_info_t %d\n", sizeof(pin_memory_info_t));
  MSG("Sizeof idamem_packet_t %d\n", sizeof(idamem_packet_t));
  MSG("Sizeof idamem_response_pkt_t %d\n", sizeof(idamem_response_pkt_t));
  MSG("Sizeof idapin_registers_t %d\n", sizeof(idapin_registers_t));
  MSG("Sizeof idatrace_data_t %d\n", sizeof(idatrace_data_t));
  MSG("Sizeof idatrace_events_t %d\n", sizeof(idatrace_events_t));
  MSG("Sizeof idabpt_packet_t %d\n", sizeof(idabpt_packet_t);
  MSG("Sizeof idalimits_packet_t %d\n", sizeof(idalimits_packet_t));
#endif
}

//--------------------------------------------------------------------------
const char *pin_basename(const char *path)
{
  if ( path != NULL )
  {
    const char *f1 = strrchr(path, '/');
    const char *f2 = strrchr(path, '\\');
    const char *file = max(f1, f2);
    if ( file != NULL )
      return file+1;
  }
  return path;
}

//--------------------------------------------------------------------------
inline       char *tail(      char *in_str) { return strchr(in_str, '\0'); }
inline const char *tail(const char *in_str) { return strchr(in_str, '\0'); }

//--------------------------------------------------------------------------
// wait for the process to be resumed only in case the listener thread is
// already spawned, if not, don't wait for anything as we can't receive
// any command from IDA side yet
static inline void wait_if_needed(PIN_SEMAPHORE *sem)
{
  if ( run_listener )
    PIN_SemaphoreWait(sem);
}

//--------------------------------------------------------------------------
bool check_address(ADDRINT addr)
{
  if ( break_at_next_inst )
    return true;

  if ( addr >= min_address && addr <= max_address )
  {
    if ( !only_new_instructions )
      return true;
    if ( only_new_instructions && std::find(all_addrs.begin(), all_addrs.end(), addr) == all_addrs.end() )
      return true;
  }
  return false;
}

//--------------------------------------------------------------------------
inline void get_context_regs(const CONTEXT *ctx, idapin_registers_t *regs)
{
  if ( ctx != NULL )
  {
    regs->eax = (ADDRINT)PIN_GetContextReg(ctx, REG_GAX);
    regs->ebx = (ADDRINT)PIN_GetContextReg(ctx, REG_GBX);
    regs->ecx = (ADDRINT)PIN_GetContextReg(ctx, REG_GCX);
    regs->edx = (ADDRINT)PIN_GetContextReg(ctx, REG_GDX);
    regs->esi = (ADDRINT)PIN_GetContextReg(ctx, REG_GSI);
    regs->edi = (ADDRINT)PIN_GetContextReg(ctx, REG_GDI);
    regs->ebp = (ADDRINT)PIN_GetContextReg(ctx, REG_GBP);
    regs->esp = (ADDRINT)PIN_GetContextReg(ctx, REG_STACK_PTR);
    regs->eip = (ADDRINT)PIN_GetContextReg(ctx, REG_INST_PTR);
#if defined(PIN_64)
    regs->r8  = (ADDRINT)PIN_GetContextReg(ctx, REG_R8);
    regs->r9  = (ADDRINT)PIN_GetContextReg(ctx, REG_R9);
    regs->r10 = (ADDRINT)PIN_GetContextReg(ctx, REG_R10);
    regs->r11 = (ADDRINT)PIN_GetContextReg(ctx, REG_R11);
    regs->r12 = (ADDRINT)PIN_GetContextReg(ctx, REG_R12);
    regs->r13 = (ADDRINT)PIN_GetContextReg(ctx, REG_R13);
    regs->r14 = (ADDRINT)PIN_GetContextReg(ctx, REG_R14);
    regs->r15 = (ADDRINT)PIN_GetContextReg(ctx, REG_R15);

    regs->eflags = (ADDRINT)PIN_GetContextReg(ctx, REG_RFLAGS);
#else
    regs->eflags = (ADDRINT)PIN_GetContextReg(ctx, REG_EFLAGS);
#endif
    regs->cs = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_CS);
    regs->ds = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_DS);
    regs->es = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_ES);
    regs->fs = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_FS);
    regs->gs = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_GS);
    regs->ss = (ADDRINT)PIN_GetContextReg(ctx, REG_SEG_SS);
  }
}

//--------------------------------------------------------------------------
// Although this function was primarily thought only for handling breakpoints,
// finally, it's going to be used also to really pause the application in case
// the user pressed the pause button.
static void PIN_FAST_ANALYSIS_CALL handle_bpt(ADDRINT addr, const CONTEXT *ctx)
{
  {
    janitor_for_pinlock_t plj(&events_lock);
    last_ea = addr;
    last_tid = PIN_GetTid();
  }

  // wait if the user paused the application
  wait_if_needed(&run_app_sem);
  // ...or if we're already handling some breakpoint
  wait_if_needed(&bpt_sem);
  bool bpt_exists = false;

  {
    janitor_for_pinlock_t plj(&soft_locks);
    if ( !soft_bpts.empty() )
      bpt_exists = std::binary_search(soft_bpts.begin(), soft_bpts.end(), addr);

    if ( !bpt_exists && break_at_next_inst )
    {
      bpt_exists = true;
      break_at_next_inst = false;
    }
  }

  if ( bpt_exists )
  {
    PIN_SemaphoreClear(&run_app_sem);
    PIN_SemaphoreClear(&bpt_sem);

    // suppressions for lint errors parsing PIN headers (4 warnings generated)
    //lint -e526 Symbol 'Message()' (line 330, file idadbg.cpp) not defined
    //lint -e628 no argument information provided for function 'Message()' (line 331, file idadbg.cpp)
    MSG("Breakpoint reached at %p\n", (void*)addr);

    // set the current thread's context, used for reading registers
    PIN_SaveContext(ctx, &threads_context[PIN_GetTid()]);

    pin_debug_event_t event;
    event.eid = BREAKPOINT;
    event.ea = addr;
    event.tid = PIN_GetTid();
    event.bpt.hea = BADADDR;
    event.bpt.kea = BADADDR;

    {
      janitor_for_pinlock_t plj(&events_lock);
      events.push_back(event);
    }

    PIN_SemaphoreWait(&bpt_sem);
  }
}

//--------------------------------------------------------------------------
inline void add_to_trace(const CONTEXT *ctx, ADDRINT ea, pin_tev_type_t tev_type)
{
  DEBUG("add_to_trace1: Adding instruction at %p\n", (void*)ea);
  // set the current thread's context, used for reading registers
  PIN_SaveContext(ctx, &threads_context[PIN_GetTid()]);

  bool trace_flush;
  {
    janitor_for_pinlock_t plj(&events_lock);
    last_ea = ea;
    last_tid = PIN_GetTid();
    all_addrs.push_front(ea);

    trc_element_t trc(last_tid, ea, tev_type);
    if ( tracing_registers && ctx != NULL )
      get_context_regs(ctx, &trc.regs);

    trace_addrs.push_back(trc);
    // wait if the user decided to pause the application or the trace is full
    trace_flush = trace_addrs.size() >= enqueue_limit;
  }
  if ( trace_flush )
    mark_trace_flush_needed(true);
}

//--------------------------------------------------------------------------
inline void add_to_trace(ADDRINT ea, pin_tev_type_t tev_type)
{
  DEBUG("add_to_trace2: Adding instruction at %p\n", (void*)ea);
  bool trace_flush;

  {
    janitor_for_pinlock_t plj(&events_lock);

    last_ea = ea;
    last_tid = PIN_GetTid();
    all_addrs.push_front(ea);

    trc_element_t trc(last_tid, ea, tev_type);
    trace_addrs.push_back(trc);
    // wait if the user decided to pause the application or the trace is full
    trace_flush = trace_addrs.size() >= enqueue_limit;
  }
  if ( trace_flush ) 
    mark_trace_flush_needed(true);
}

//--------------------------------------------------------------------------
static ADDRINT listener_not_running(VOID *)
{
  return !run_listener;
}

//--------------------------------------------------------------------------
static void wait_for_listener(void)
{
  // Wait for the listener to start
#if defined(_WIN32) && !defined(TARGET_IA32)
  while ( !PIN_SemaphoreTimedWait(&listener_sem, 100) )
  {
    DEBUG("Waiting for the listener to start...\n");
    PIN_Yield();
  }
#else
  PIN_SemaphoreWait(&listener_sem);
#endif
}

//--------------------------------------------------------------------------
static ADDRINT ins_enabled(VOID *)
{
  return tracing_instruction;
}

//--------------------------------------------------------------------------
static ADDRINT trc_enabled(VOID *)
{
  return tracing_bblock;
}

//--------------------------------------------------------------------------
static ADDRINT rtn_enabled(VOID *)
{
  return tracing_routine;
}

//--------------------------------------------------------------------------
// This function is called before every instruction is executed
static VOID PIN_FAST_ANALYSIS_CALL ins_logic_cb(
        const CONTEXT *ctx,
        VOID *ip,
        pin_tev_type_t tev_type)
{
  add_to_trace(ctx, (ADDRINT)ip, tev_type);
}

//--------------------------------------------------------------------------
// This function is called before every routine
static VOID PIN_FAST_ANALYSIS_CALL rtn_logic_cb(VOID *ip)
{
  if ( check_address((ADDRINT)ip) )
    add_to_trace((ADDRINT)ip, tev_insn);
}

//--------------------------------------------------------------------------
inline static void mark_trace_flush_needed(bool pause)
{
  add_instrumentation(TF_TRACE_INSN);

  {
    janitor_for_pinlock_t plj(&events_lock);
    trace_flush_needed = true;
  }

  // pause the app until the whole trace is read
  if ( pause )
  {
    PIN_SemaphoreClear(&run_app_sem);
    PIN_SemaphoreWait(&run_app_sem);
  }
}

//--------------------------------------------------------------------------
// This function is called when the application exits
static VOID fini_cb(INT32, VOID *)
{
  DEBUG("Marking the trace as full\n");
  // mark the trace as full *and* pause the app
  mark_trace_flush_needed(true);

  // finally, add the process exit event
  pin_debug_event_t evt;
  evt.eid = PROCESS_EXIT;

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(evt);
  }

  MSG("Waiting for listener thread to exit...\n");
  if ( PIN_WaitForThreadTermination(listener_uid, 10000, NULL) )
  {
    MSG("Everything OK\n");
  }
  else
  {
    MSG("Timeout waiting for listener thread.\n");
  }
}

//--------------------------------------------------------------------------
//lint -e{1746} parameter 'img' could be made const reference
static VOID image_load_cb(IMG img, VOID *)
{
  wait_if_needed(&run_app_sem);

  pin_debug_event_t event;
  event.eid = LIBRARY_LOAD;
  event.ea = IMG_Entry(img);
  event.tid = PIN_GetTid();
  event.handled = true;
  pin_strncpy(event.modinfo.name, IMG_Name(img).c_str(), sizeof(event.modinfo.name));

  ADDRINT start_ea = IMG_LowAddress(img);
  ADDRINT end_ea = IMG_HighAddress(img);

  MSG("Loading library %s %p:%p\n", IMG_Name(img).c_str(), (void*)start_ea, (void*)end_ea);
  event.modinfo.base = start_ea;
  event.modinfo.size = (pin_size_t)(end_ea - start_ea);
  event.modinfo.rebase_to = BADADDR;

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(event);
  }

  if ( image_name != "*" && !image_name.empty() )
  {
    string base_head  = pin_basename(IMG_Name(img).c_str());
    string base_image = pin_basename(image_name.c_str());
    transform(base_head.begin(), base_head.end(), base_head.begin(), ::tolower);
    transform(base_image.begin(), base_image.end(), base_image.begin(), ::tolower);

    if ( base_head == base_image )
    {
      min_address = IMG_LowAddress(img);
      max_address = IMG_HighAddress(img);
      MSG("Setting default instrumentation limits: Min EA %p Max EA %p\n", (void*)start_ea, (void*)end_ea);
    }
  }
}

//--------------------------------------------------------------------------
// Pin calls this function every time a new img is unloaded
// You can't instrument an image that is about to be unloaded
//lint -e{1746} parameter 'img' could be made const reference
static VOID image_unload_cb(IMG img, VOID *)
{
  pin_debug_event_t ev;
  ev.eid     = LIBRARY_UNLOAD;
  ev.pid     = PIN_GetPid();
  ev.tid     = PIN_GetTid();
  ev.ea      = BADADDR;
  ev.handled = true;
  pin_strncpy(ev.info, IMG_Name(img).c_str(), sizeof(ev.info));

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(ev);
  }
  MSG("Unloading %s\n", IMG_Name(img).c_str());
}

//--------------------------------------------------------------------------
// This routine is executed every time a thread is created
//lint -e{818} Pointer parameter 'ctx' (line 560) could be declared as pointing to const
static VOID thread_start_cb(THREADID, CONTEXT *ctx, INT32, VOID *)
{
  PIN_SaveContext(ctx, &threads_context[PIN_GetTid()]);

  pin_debug_event_t ev;
  ev.eid = THREAD_START;
  ev.pid = PIN_GetPid();
  ev.tid = PIN_GetTid();
  if ( ctx != NULL )
    ev.ea = (ADDRINT)PIN_GetContextReg(ctx, REG_INST_PTR);
  else
    ev.ea = BADADDR;
  MSG("THREAD START: %d AT %p\n", ev.tid, (void*)ev.ea);
  ev.handled = true;

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(ev);
  }
}

//--------------------------------------------------------------------------
// This routine is executed every time a thread is destroyed.
static VOID thread_fini_cb(THREADID, const CONTEXT *ctx, INT32, VOID *)
{
  PIN_SaveContext(ctx, &threads_context[PIN_GetTid()]);

  pin_debug_event_t ev;
  ev.eid = THREAD_EXIT;
  ev.pid = PIN_GetPid();
  ev.tid = PIN_GetTid();
  ev.ea = ctx != NULL ? (ADDRINT)PIN_GetContextReg(ctx, REG_INST_PTR) : BADADDR;
  ev.handled = true;
  MSG("THREAD FINISH: %d AT %p\n", ev.tid, (void*)ev.ea);

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(ev);
  }

  mark_trace_flush_needed(true);
}

//--------------------------------------------------------------------------
// Pin calls this function every time a new instruction is encountered
//lint -e{1746} parameter 'ins' could be made const reference
static VOID instruction_cb(INS ins, VOID *)
{
  // resize it if needed
  if ( all_addrs.size() >= skip_limit )
    all_addrs.resize(skip_limit);

  // Insert a call to ins_logic_cb before every instruction
  ADDRINT addr = INS_Address(ins);
  if ( check_address(addr) )
  {
    if ( !run_listener )
    {
      // add the code for synching with the listener thread
      INS_InsertIfCall(ins, IPOINT_BEFORE, (AFUNPTR)listener_not_running, IARG_END);
      INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR)wait_for_listener, IARG_END);
    }

    INS_InsertIfCall(ins, IPOINT_BEFORE, (AFUNPTR)ins_enabled, IARG_INST_PTR, IARG_END);
    INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR)ins_logic_cb, IARG_FAST_ANALYSIS_CALL,
                IARG_CONST_CONTEXT, IARG_INST_PTR, IARG_UINT32, tev_insn, IARG_END);

    INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)handle_bpt, IARG_FAST_ANALYSIS_CALL,
                IARG_INST_PTR, IARG_CONST_CONTEXT, IARG_END);
  }
}

//--------------------------------------------------------------------------
// Pin calls this function every time a new basic block is encountered *BUT*
// we will use this callback for instrumenting routines instead of using the
// routine instrumentation API offered by the toolkit
static VOID routine_cb(TRACE trace, VOID *)
{
  // resize it if needed
  if ( all_addrs.size() >= skip_limit )
    all_addrs.resize(skip_limit);

  // Visit every basic block in the trace
  for ( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
  {
    for ( INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) )
    {
      // handle both calls and push + ret like in the following example:
      //
      // push offset some_func
      // retn
      //
      if ( INS_IsCall(ins) || INS_IsRet(ins) )
      {
        if ( !run_listener )
        {
          // add the code for synching with the listener thread
          INS_InsertIfCall(ins, IPOINT_BEFORE, (AFUNPTR)listener_not_running, IARG_END);
          INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR)wait_for_listener, IARG_END);
        }

        // add the real instruction instrumentation
        INS_InsertIfCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)rtn_enabled, IARG_INST_PTR, IARG_END);
        INS_InsertThenCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)rtn_logic_cb, IARG_FAST_ANALYSIS_CALL,
                IARG_BRANCH_TARGET_ADDR, IARG_END);
      }
    }
  }
}

//--------------------------------------------------------------------------
// Pin calls this function every time a new basic block is encountered.
static VOID trace_cb(TRACE trace, VOID *)
{
  // resize it if needed
  if ( all_addrs.size() >= skip_limit )
    all_addrs.resize(skip_limit);

  // Visit every basic block in the trace
  for ( BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl) )
  {
    bool first = true;
    for ( INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins) )
    {
      // Insert a call to ins_logic_cb for every instruction which is either
      // a call, branch, ret, syscall or invalid (i.e., UD2) and, also, to the
      // 1st instruction in the basic block
      ADDRINT addr = INS_Address(ins);
      if ( (first || INS_IsBranchOrCall(ins) || INS_IsRet(ins) || INS_IsSyscall(ins) || !ins.is_valid())
        && check_address(addr) )
      {
        if ( !run_listener )
        {
          // add the code for synching with the listener thread
          INS_InsertIfCall(ins, IPOINT_BEFORE, (AFUNPTR)listener_not_running, IARG_END);
          INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR)wait_for_listener, IARG_END);
        }

        first = false;
        pin_tev_type_t tev_type = tev_insn;
        if ( INS_IsCall(ins) )
          tev_type = tev_call;
        else if ( INS_IsRet(ins) )
          tev_type = tev_ret;

        // add the real instruction instrumentation
        INS_InsertIfCall(ins, IPOINT_BEFORE, (AFUNPTR)trc_enabled, IARG_INST_PTR, IARG_END);
        INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR)ins_logic_cb, IARG_FAST_ANALYSIS_CALL,
                IARG_CONST_CONTEXT, IARG_INST_PTR, IARG_UINT32, (uint32)tev_type, IARG_END);
      }
    }
  }
}

//--------------------------------------------------------------------------
static INT32 usage()
{
  fprintf(stderr, "Pin Tool to communicate with IDA's debugger\n");
  fprintf(stderr, "\n%s\n", KNOB_BASE::StringKnobSummary().c_str());
  return -1;
}

//--------------------------------------------------------------------------
inline static void error_msg(const char *msg)
{
  MSG("%s: %s\n", msg, strerror(errno));
}

//--------------------------------------------------------------------------
static bool sockets_startup(void)
{
  bool ret = true;
#ifdef _WIN32
  WINDOWS::WORD wVersionRequested = 0x0202;
  WINDOWS::WSADATA wsaData;
  int err = WINDOWS::WSAStartup(wVersionRequested, &wsaData);
  if ( err != 0 )
  {
    error_msg("WSAStartup");
    ret = false;
  }
#endif
  return ret;
}

//--------------------------------------------------------------------------
static bool init_socket(void)
{
  portno = knob_ida_port;

  sockets_startup();
  srv_socket = socket(AF_INET, SOCK_STREAM, 0);
  if ( srv_socket == (PIN_SOCKET)-1 )
  {
    error_msg("socket");
    return false;
  }

  int optval = 1;
  pin_setsockopt(srv_socket, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));

#ifdef _WIN32
  int timeout_ms = 30000;
  WINDOWS::setsockopt(srv_socket, SOL_SOCKET, SO_RCVTIMEO, (char*)&timeout_ms, sizeof(int));
#else
  struct timeval tv;
  tv.tv_sec = 30;
  setsockopt(srv_socket, SOL_SOCKET, SO_RCVTIMEO,(struct timeval *)&tv,sizeof(struct timeval));
#endif

  struct pin_sockaddr_in sa;
  memset(&sa, '\0', sizeof(sa));
  sa.sin_family = AF_INET;
  sa.sin_port   = pin_htons(portno);
  if ( pin_bind(srv_socket, (pin_sockaddr *)&sa, sizeof(sa)) != 0 )
  {
    error_msg("bind");
    return false;
  }

  if ( pin_listen(srv_socket, 1) == 0 )
  {
    MSG("Listening at port %d...\n", (int)portno);

    pin_socklen_t clilen = sizeof(sa);
    //lint -e565 tag 'sockaddr' not previously seen, assumed file-level scope
    cli_socket = pin_accept(srv_socket, ((struct sockaddr *)&sa), &clilen);
    if ( cli_socket > 0 )
      return true;
  }
  return false;
}

//--------------------------------------------------------------------------
static bool send_packet(
        const void *pkt,
        size_t size,
        void *answer,
        size_t ans_size,
        const char *from)
{
  ssize_t bytes = pin_send(cli_socket, pkt, size, from);
  if ( bytes > -1 )
  {
    if ( answer != NULL )
      bytes = pin_recv(cli_socket, answer, ans_size, from);
  }
  return bytes != -1;
}

//--------------------------------------------------------------------------
static VOID app_start_cb(VOID *)
{
  pin_debug_event_t event;
  event.pid = PIN_GetPid();
  event.tid = PIN_GetTid();
  event.eid = PROCESS_START;
  DEBUG("Setting process started to true\n");
  process_started = true;

  IMG img;
  img.invalidate();
  for( img = APP_ImgHead(); IMG_Valid(img); img = IMG_Next(img) )
  {
    if ( IMG_IsMainExecutable(img) )
      break;
  }

  if ( !img.is_valid() )
  {
    MSG("Cannot find the 1st instruction of the main executable!\n");
    abort();
  }
  event.ea = IMG_Entry(img);
  pin_strncpy(event.modinfo.name, IMG_Name(img).c_str(), sizeof(event.modinfo.name));

  // by default, we set the limits of the trace to the main binary
  ADDRINT start_ea = IMG_LowAddress(img);
  ADDRINT end_ea = IMG_HighAddress(img);

  event.modinfo.base = start_ea;
  event.modinfo.rebase_to = BADADDR;
  event.modinfo.size = (uint32)(end_ea - start_ea);

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(event);
  }

  if ( !app_start_cb_called )
  {
    app_start_cb_called = true;
    MSG("Head image: %s Start %p End %p\n", IMG_Name(img).c_str(), (void*)start_ea, (void*)end_ea);
    if ( image_name == "*" )
    {
      MSG("Image name set to '*', tracing everything!\n");
      min_address = 0;
      max_address = (unsigned)BADADDR;
    }
    else if ( min_address != start_ea && max_address != end_ea )
    {
      string base_head  = pin_basename(IMG_Name(img).c_str());
      string base_image = pin_basename(image_name.c_str());
      transform(base_head.begin(), base_head.end(), base_head.begin(), ::tolower);
      transform(base_image.begin(), base_image.end(), base_image.begin(), ::tolower);

      if ( image_name.empty() || base_head == base_image )
      {
        MSG("Setting default instrumentation limits: Min EA %p Max EA %p", (void*)start_ea, (void*)end_ea);
        min_address = start_ea;
        max_address = end_ea;
      }
    }
  }

  // Handle packets in the main thread until we receive the RESUME event.
  // We need this to add breakpoints before the application's code is
  // executed, otherwise, we will run into race conditions
  handle_packets(-1, "RESUME START");
  MSG("All breakpoints seems to be added, running the application...\n");

  if ( PIN_IsAttaching() )
    break_at_next_inst = true;
}

//--------------------------------------------------------------------------
static VOID context_change_cb(
        THREADID,
        CONTEXT_CHANGE_REASON reason,
        const CONTEXT *ctxt_from,
        CONTEXT *,
        INT32 sig,
        VOID *)
{
  pin_debug_event_t event;
  switch (reason)
  {
    case CONTEXT_CHANGE_REASON_FATALSIGNAL:
      event.exc.can_cont = false;
      snprintf(event.exc.info, sizeof(event.exc.info), "Fatal signal %d at 0x"HEX_FMT"", sig, event.ea);
      break;
    case CONTEXT_CHANGE_REASON_SIGNAL:
      snprintf(event.exc.info, sizeof(event.exc.info), "Signal %d at 0x"HEX_FMT"", sig, event.ea);
      break;
    case CONTEXT_CHANGE_REASON_SIGRETURN:
      snprintf(event.exc.info, sizeof(event.exc.info), "Signal return %d at 0x"HEX_FMT"", sig, event.ea);
      break;
    case CONTEXT_CHANGE_REASON_APC:
      snprintf(event.exc.info, sizeof(event.exc.info), "Asynchronous Process Call %d at 0x"HEX_FMT"", sig, event.ea);
      break;
    case CONTEXT_CHANGE_REASON_EXCEPTION:
      snprintf(event.exc.info, sizeof(event.exc.info), "Exception 0x%x at address 0x"HEX_FMT"", sig, event.ea);
      break;
    default:
      return;
  }

  event.eid = EXCEPTION;
  event.tid = PIN_GetTid();
  event.ea = ctxt_from != NULL ? (ADDRINT)PIN_GetContextReg(ctxt_from, REG_INST_PTR) : BADADDR;
  event.exc.code = sig;
  event.exc.ea = event.ea;

  {
    janitor_for_pinlock_t plj(&events_lock);
    events.push_back(event);
  }

  mark_trace_flush_needed(true);
}

//--------------------------------------------------------------------------
static VOID ida_pin_listener(VOID *)
{
  MSG("Listener started...\n");
  run_listener = true;

  while ( run_listener )
  {
    DEBUG("Handling events in ida_pin_listener\n");
    idapin_packet_t res;
    ssize_t bytes = pin_recv(cli_socket, &res, sizeof(res), "ida_pin_listener");
    if ( bytes == -1 )
    {
      error_msg("recv");
      continue;
    }

    if ( !handle_packet(&res) )
    {
      MSG("Error handling %s packet, exiting...\n", last_packet);
      PIN_ExitThread(0);
    }

    if ( PIN_IsProcessExiting() && events.empty() && process_exit )
    {
      MSG("Process is exiting...\n");
      break;
    }
  }
}

//--------------------------------------------------------------------------
static void add_instrumentation(trace_flags_t inst)
{
  if ( (instrumentations & inst) == 0 )
  {
    switch ( inst )
    {
      case TF_TRACE_INSN:
        // Register instruction_cb to be called to instrument instructions
        MSG("Adding instruction level instrumentation...\n");
        INS_AddInstrumentFunction(instruction_cb, 0);
        break;
      case TF_TRACE_BBLOCK:
        // Register trace_cb to be called to instrument basic blocks
        MSG("Adding trace level instrumentation...\n");
        TRACE_AddInstrumentFunction(trace_cb, 0);
        break;
      case TF_TRACE_ROUTINE:
        // Register routine_cb to be called to instrument routines
        MSG("Adding routine level instrumentation...\n");
        TRACE_AddInstrumentFunction(routine_cb, 0);
        break;
      default:
        MSG("Unknown instrumentation type %d!\n", inst);
        abort();
    }

    instrumentations |= inst;
  }
}

//--------------------------------------------------------------------------
static void handle_start_process(void)
{
  if ( !tracing_instruction && !tracing_bblock && !tracing_routine )
  {
    MSG("NOTICE: No tracing method selected, nothing will be recorded until some tracing method is selected.\n");
  }

  if ( tracing_instruction )
    add_instrumentation(TF_TRACE_INSN);

  if ( tracing_bblock )
    add_instrumentation(TF_TRACE_BBLOCK);

  if ( tracing_routine )
    add_instrumentation(TF_TRACE_ROUTINE);

  // initialized stuff
  soft_bpts.clear();

  // Initialize the semaphore used for the whole application pausing
  PIN_SemaphoreInit(&run_app_sem);
  // Initialize the semaphore used for breakpoints
  PIN_SemaphoreInit(&bpt_sem);
  // And immediately set it
  PIN_SemaphoreSet(&bpt_sem);
  // Initialize the listener semaphore
  PIN_SemaphoreInit(&listener_sem);
  // AFAIK, it isn't needed, but it doesn't harm anyway
  PIN_SemaphoreClear(&listener_sem);

  // Initialize the events list lock
  InitLock(&events_lock);
  // Initialize the breakpoints lock
  InitLock(&soft_locks);

  // Set the breakpoints semaphore as we don't set any breakpoint before
  // starting the application; we will set them after PIN_StartProgram if
  PIN_SemaphoreSet(&run_app_sem);

  // Register image_load_cb to be called when an image is loaded
  IMG_AddInstrumentFunction(image_load_cb, 0);

  // Register image_unload_cb to be called when an image is unloaded
  IMG_AddUnloadFunction(image_unload_cb, 0);

  // Register Analysis routines to be called when a thread begins/ends
  PIN_AddThreadStartFunction(thread_start_cb, 0);
  PIN_AddThreadFiniFunction(thread_fini_cb, 0);

  // Register fini_cb to be called when the application exits
  PIN_AddFiniUnlockedFunction(fini_cb, 0);

  // Register aplication start callback
  PIN_AddApplicationStartFunction(app_start_cb, 0);

  // Register context change function
  PIN_AddContextChangeFunction(context_change_cb, 0);

  // Create the thread for communicating with IDA
  THREADID thread_id = PIN_SpawnInternalThread(ida_pin_listener, NULL, 0, &listener_uid);
  if ( thread_id == INVALID_THREADID )
  {
    MSG("PIN_SpawnInternalThread(BufferProcessingThread) failed\n");
    exit(-1);
  }

  // Start the program, never returns
  PIN_StartProgram();
}

//--------------------------------------------------------------------------
inline static void pop_debug_event(pin_debug_event_t *evt)
{
  assert(!events.empty());
  {
    janitor_for_pinlock_t plj(&events_lock);
    *evt = events.front();
    events.pop_front();
  }
}

//--------------------------------------------------------------------------
static void add_segment(
        pin_meminfo_vec_t &miv,
        ADDRINT sec_ea,
        pin_memory_info_t &mi)
{
  bool added = false;
  pin_meminfo_vec_t::reverse_iterator p;
  for ( p = miv.rbegin(); p != miv.rend(); ++p )
  {
    const pin_memory_info_t &pmi = *p;

    // ignore duplicated segments
    if ( pmi.startEA == sec_ea )
      return;

    // if we found the correct position insert it
    if ( pmi.endEA <= sec_ea )
    {
      added = true;
      miv.insert(p.base(), mi);
      break;
    }
  }

  // no position found, put in the end of the list
  if ( !added )
    miv.push_back(mi);
}

#ifdef _WIN32

//--------------------------------------------------------------------------
// In win32, if we try to enumerate the segments of the application we will
// read the PIN's segments; that's: we're going to mix the application's own
// segments with the conflicting real PIN segments (the areas can overlap) so,
// at least to be able to see the referenced memory in stack and heap and other
// areas we cannot list we create 2 big segments at start and at the end.
static void get_os_segments(
        pin_meminfo_vec_t &miv,
        ADDRINT img_min_ea,
        ADDRINT img_max_ea)
{
  WINDOWS::SYSTEM_INFO si;
  WINDOWS::GetSystemInfo(&si);

  pin_memory_info_t mi;
  mi.startEA = (ADDRINT)si.lpMinimumApplicationAddress;
  mi.endEA = img_min_ea;
  mi.bitness = BITNESS;
  mi.name_size = 0;
  memset(mi.name, '\0', sizeof(mi.name));
  mi.perm = SEGPERM_READ | SEGPERM_WRITE | SEGPERM_EXEC;
  miv.insert(miv.begin(), mi);

  pin_memory_info_t mi2;
  mi2.startEA = img_max_ea;
  mi2.endEA = (ADDRINT)si.lpMaximumApplicationAddress;
  mi2.bitness = BITNESS;
  mi2.name_size = 0;
  memset(mi2.name, '\0', sizeof(mi2.name));
  mi.perm = SEGPERM_READ | SEGPERM_WRITE | SEGPERM_EXEC;
  miv.push_back(mi2);
}
#else

#ifdef __linux__

//--------------------------------------------------------------------------
const char *skip_spaces(const char *ptr)
{
  if ( ptr != NULL )
  {
    while ( isspace(*ptr) )
      ptr++;
  }
  return ptr;
}

//--------------------------------------------------------------------------
struct mapfp_entry_t
{
  addr_t ea1;
  addr_t ea2;
  addr_t offset;
  uint64_t inode;
  char perm[8];
  char device[8];
  string fname;
};

//--------------------------------------------------------------------------
static bool read_mapping(FILE *mapfp, mapfp_entry_t *me)
{
  char line[2*MAXSTR];
  if ( !fgets(line, sizeof(line), mapfp) )
    return false;

  me->ea1 = BADADDR;

  uint32 len = 0;
  int code = sscanf(line, HEX_FMT"-"HEX_FMT" %s "HEX_FMT" %s " HEX64T_FMT "x%n",
                     &me->ea1,
                     &me->ea2,
                     me->perm,
                     &me->offset,
                     me->device,
                     &me->inode,
                     &len);
  if ( code == 6 && len < sizeof(line) )
  {
    char *ptr = &line[len];
    ptr = (char *)skip_spaces(ptr);
    // remove trailing spaces and eventual (deleted) suffix
    static const char delsuff[] = " (deleted)";
    const int suflen = sizeof(delsuff) - 1;
    char *end = (char*)tail(ptr);
    while ( end > ptr && isspace(end[-1]) )
      *--end = '\0';
    if ( end-ptr > suflen && strncmp(end-suflen, delsuff, suflen) == 0 )
      end[-suflen] = '\0';
    me->fname = ptr;
  }
  return (signed)me->ea1 != BADADDR;
}

//--------------------------------------------------------------------------
static void get_os_segments(pin_meminfo_vec_t &miv, ADDRINT, ADDRINT)
{
  FILE *mapfp = fopen("/proc/self/maps", "rb");
  mapfp_entry_t me;
  while ( read_mapping(mapfp, &me) )
  {
    // for some reason linux lists some areas twice
    // ignore them
    size_t i;
    for ( i=0; i < miv.size(); i++ )
      if ( miv[i].startEA == me.ea1 )
        break;
    if ( i != miv.size() )
      continue;

    // do we really want to hide the PIN's segments? I guess yes, but...
    if ( me.fname != "pinbin" )
    {
      pin_memory_info_t mi;
      mi.startEA = me.ea1;
      mi.endEA   = me.ea2;
      pin_strncpy(mi.name, me.fname.c_str(), sizeof(mi.name));
      mi.bitness = BITNESS;

      if ( strchr(me.perm, 'r') != NULL ) mi.perm |= SEGPERM_READ;
      if ( strchr(me.perm, 'w') != NULL ) mi.perm |= SEGPERM_WRITE;
      if ( strchr(me.perm, 'x') != NULL ) mi.perm |= SEGPERM_EXEC;

      add_segment(miv, me.ea1, mi);
    }
  }
}

#else
// MacOSX
static void get_os_segments(pin_meminfo_vec_t &, ADDRINT, ADDRINT){}
#endif

#endif

//--------------------------------------------------------------------------
static bool handle_memory_info(void)
{
  bool ret = false;

  // Visit every loaded image and fill miv vector
  pin_meminfo_vec_t miv;
  ADDRINT img_min_ea = 0;
  ADDRINT img_max_ea = 0;
  for( IMG img= APP_ImgHead(); IMG_Valid(img); img = IMG_Next(img) )
  {
    for( SEC sec= IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec) )
    {
      pin_memory_info_t mi;
      ea_t sec_ea = SEC_Address(sec);
      if ( sec_ea != 0 )
      {
        mi.startEA = sec_ea;
        mi.endEA   = sec_ea + SEC_Size(sec);

        if ( img_min_ea == 0 || img_min_ea > sec_ea )
          img_min_ea = sec_ea;

        if ( img_max_ea == 0 || img_max_ea < mi.endEA )
          img_max_ea = mi.endEA;

        string sec_name;
        sec_name = IMG_Name(img) + ":" + SEC_Name(sec);
        mi.name_size = (uint32)sec_name.size();
        memset(mi.name, '\0', sizeof(mi.name));
        pin_strncpy(mi.name, sec_name.c_str(), sizeof(mi.name));
        mi.bitness = BITNESS;

        mi.perm = 0;
        if ( SEC_IsReadable(sec) )
          mi.perm |= SEGPERM_READ;
        if ( SEC_IsWriteable(sec) )
          mi.perm |= SEGPERM_WRITE;
        if ( SEC_IsExecutable(sec) )
          mi.perm |= SEGPERM_EXEC;

        add_segment(miv, sec_ea, mi);
        //fprintf(stderr, "Segment %s Start 0x"HEX_FMT" End 0x"HEX_FMT"\n", mi.name, mi.startEA, mi.endEA);
      }
    }
  }

  get_os_segments(miv, img_min_ea, img_max_ea);

  memimages_pkt_t pkt(PTT_MEMORY_INFO, (uint32)miv.size());

  // send a packet with the number of segments
  if ( send_packet(&pkt, sizeof(pkt), NULL, 0, "handle_memory_info(1)") )
  {
    ret = true;

    // and, then, send the information for all images
    pin_meminfo_vec_t::iterator p;
    for ( p = miv.begin(); p != miv.end(); ++p )
    {
      pin_memory_info_t &mi = *p;
      if ( !send_packet(&mi, sizeof(mi), NULL, 0, "handle_memory_info(2)") )
      {
        ret = false;
        break;
      }
    }
  }

  return ret;
}

//--------------------------------------------------------------------------
static bool handle_read_memory(ADDRINT ea, pin_size_t size)
{
  DEBUG("Reading %d bytes at address %p\n", size, (void*)ea);

  idamem_response_pkt_t pkt;
  // read the data asked by IDA
  size_t copy_size = size < sizeof(pkt.buf) ? size : sizeof(pkt.buf);
  ssize_t read_bytes = PIN_SafeCopy(pkt.buf, (void*)ea, copy_size);
  pkt.size = (uint32)read_bytes;
  pkt.code = PTT_READ_MEMORY;

  ssize_t bytes = pin_send(cli_socket, &pkt, sizeof(pkt), __FUNCTION__);
  return bytes == sizeof(pkt);
}

//--------------------------------------------------------------------------
static bool handle_read_trace(void)
{
  idatrace_events_t trc_events;
  trc_events.code = PTT_ACK;
  trc_events.size = 0;

  {
    janitor_for_pinlock_t plj(&events_lock);

    do
    {
      if ( trace_addrs.empty() )
        break;

      trc_element_t trc = trace_addrs.front();
      trace_addrs.pop_front();
      trc_events.trace[trc_events.size].tid = trc.tid;
      trc_events.trace[trc_events.size].ea = trc.ea;
      trc_events.trace[trc_events.size].type = trc.type;
      trc_events.trace[trc_events.size].registers = trc.regs;
    } while ( ++trc_events.size < TRACE_EVENTS_SIZE );
  }

  ssize_t bytes = pin_send(cli_socket, &trc_events, sizeof(trc_events), __FUNCTION__);
  return bytes == sizeof(trc_events);
}

//--------------------------------------------------------------------------
static bool handle_read_regs(pin_thid tid)
{
  threadctx_map_t::iterator p;
  p = threads_context.find(tid);
  if ( p == threads_context.end() )
    return false;

  idapin_registers_t regs;
  const CONTEXT &ctx = p->second;
  get_context_regs(&ctx, &regs);
  ssize_t bytes = pin_send(cli_socket, &regs, sizeof(regs), __FUNCTION__);
  return bytes == sizeof(regs);
}

//--------------------------------------------------------------------------
static void update_instrumentation(uint32 trace_types)
{
  tracing_instruction = (trace_types & TF_TRACE_INSN) != 0;
  tracing_bblock = (trace_types & TF_TRACE_BBLOCK) != 0;
  tracing_routine = (trace_types & TF_TRACE_ROUTINE) != 0;
  tracing_registers = (trace_types & TF_REGISTERS) != 0;

  if ( tracing_instruction )
  {
    MSG("Enabling instruction tracing...\n");
    add_instrumentation(TF_TRACE_INSN);
  }
  else
  {
    MSG("Disabling instruction tracing...\n");
  }

  if ( tracing_bblock )
  {
    MSG("Enabling basic block tracing...\n");
    add_instrumentation(TF_TRACE_BBLOCK);
  }
  else
  {
    MSG("Disabling basic block tracing...\n");
  }

  if ( tracing_routine )
  {
    MSG("Enabling routine tracing...\n");
    add_instrumentation(TF_TRACE_ROUTINE);
  }
  else
  {
    MSG("Disabling routine tracing...\n");
  }

  MSG("%sabling register values tracing...\n", tracing_registers ? "En" : "Dis");
}

//--------------------------------------------------------------------------
static bool handle_limits(void)
{
  bool ret = false;
  idalimits_packet_t ans;
  ssize_t bytes = pin_recv(cli_socket, &ans, sizeof(ans), __FUNCTION__);
  if ( bytes == sizeof(ans) )
  {
    image_name = ans.image_name;
    enqueue_limit = ans.trace_limit;
    only_new_instructions = ans.only_new;

    MSG("Setting maximum enqueue limit to %d, tracing image '%s', new instructions only %d\n",
        enqueue_limit, image_name.c_str(), only_new_instructions);

#if _WIN32
    bool is_windows = true;
#else
    bool is_windows = false;
#endif

    // under windows, we do not support tracing basic blocks/instructions
    // when tracing everything since beginning
    idapin_packet_t res;
    if ( is_windows && image_name == "*" && (tracing_instruction || tracing_bblock) )
    {
      if ( tracing_instruction )
        MSG("Unsupported configuration: tracing everything and instruction level enabled\n");
      else
        MSG("Unsupported configuration: tracing everything and basic block level enabled\n");

      res.code = PTT_ERROR;
    }
    else
    {
      MSG("Correct configuration received\n");
      res.code = PTT_ACK;
    }

    // send the answer and terminate the application if the selected configuration
    // is not supported
    bytes = pin_send(cli_socket, &res, sizeof(res), __FUNCTION__);
    if ( res.code == PTT_ERROR || bytes != sizeof(res) )
    {
      MSG("Unsupported configuration or network error while setting limits, calling PIN_ExitApplication\n");
      PIN_ExitApplication(-1);
    }
    ret = true;
  }
  return ret;
}

//--------------------------------------------------------------------------
static bool do_resume(idapin_packet_t *ans)
{
  trace_flush_needed = false;
  PIN_SemaphoreSet(&run_app_sem);
  PIN_SemaphoreSet(&listener_sem);
  ans->code = PTT_ACK;
  last_ea = BADADDR;
  last_tid = 0;
  return send_packet(ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
}

//--------------------------------------------------------------------------
static const char *const packet_names[] =
{
  "ACK",           "ERROR",       "HELLO",       "EXIT PROCESS",
  "START PROCESS", "DEBUG EVENT", "READ EVENT",  "MEMORY INFO",
  "READ MEMORY",   "DETACH",      "COUNT TRACE", "READ TRACE",
  "CLEAR TRACE",   "PAUSE",       "RESUME",      "RESUME START",
  "ADD BPT",       "DEL BPT",     "RESUME BPT",  "CAN READ REGS",
  "READ REGS",     "SET TRACE",   "SET OPTIONS"
};

//--------------------------------------------------------------------------
static bool handle_packet(idapin_packet_t *res)
{
  bool ret = false;
  idapin_packet_t ans;
  ans.size = 0;
  ans.code = PTT_ERROR;

  if ( res->code > PTT_END )
  {
    MSG("Unknown packet type %d, exiting...\n", res->code);
    PIN_ExitProcess(0);
  }
  last_packet = packet_names[res->code];

  switch ( res->code )
  {
    case PTT_HELLO:
      // In the hello packet we send the size of ADDRINT to know
      // if we're using the correct IDA version (32 or 64 bits)
      // and in size the version of the protocol
      ans.size = PIN_PROTOCOL_VERSION;
      ans.data = sizeof(ADDRINT);
      ans.code = PTT_ACK;
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_START_PROCESS:
      // does not return
      handle_start_process();
      break;
    case PTT_EXIT_PROCESS:
      MSG("Received EXIT PROCESS, exiting from process...\n");
      run_listener = false;
      PIN_SemaphoreSet(&run_app_sem);
      PIN_SemaphoreSet(&bpt_sem);
      // does not return
      PIN_ExitProcess(0);
    case PTT_DEBUG_EVENT:
      ans.data = 0;
      if ( !events.empty() && process_started )
      {
        DEBUG("Total of %d events recorded\n", (uint32)events.size());
        ans.size = (uint32)events.size();
        ans.code = PTT_DEBUG_EVENT;
      }
      else
      {
        ans.size = 0;
        ans.code = PTT_ACK;
      }
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_READ_EVENT:
      {
        pin_debug_event_t evt;

        if ( trace_flush_needed )
          evt.eid = TRACE_FULL;
        else if ( !events.empty() )
          pop_debug_event(&evt);
        else
          evt.eid = NO_EVENT;

        if ( evt.eid == PROCESS_EXIT )
          process_exit = true;

        ret = send_packet(&evt, sizeof(evt), NULL, 0, __FUNCTION__);
      }
      break;
    case PTT_MEMORY_INFO:
      ret = handle_memory_info();
      break;
    case PTT_READ_MEMORY:
      ans.data = 0;
      ans.code = PTT_READ_MEMORY;
      ret = handle_read_memory(res->data, res->size);
      break;
    case PTT_DETACH:
      // this function is asynchronous
      PIN_Detach();

      ans.data = 0;
      ans.code = PTT_ACK;
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_PAUSE:
      PIN_SemaphoreClear(&run_app_sem);

      ans.code = PTT_ACK;
      ans.size = last_tid;
      ans.data = last_ea;
      mark_trace_flush_needed(false);

      DEBUG("Paused, thread %d, EA %p\n", last_tid, (void*)last_ea);
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_RESUME_BPT:
      DEBUG("Resuming from a breakpoint...\n");
      PIN_SemaphoreSet(&bpt_sem);
      ret = do_resume(&ans);
      break;
    case PTT_RESUME_START:
      DEBUG("Resuming from application start...\n");
      ret = do_resume(&ans);
      break;
    case PTT_RESUME:
      DEBUG("Resuming, thread %d, EA %p\n", last_tid, (void*)last_ea);
      ret = do_resume(&ans);
      break;
    case PTT_COUNT_TRACE:
      ans.code = PTT_ACK;
      ans.data = trace_addrs.size();
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_READ_TRACE:
      ret = handle_read_trace();
      {
        janitor_for_pinlock_t plj(&events_lock);
        trace_flush_needed = false;
      }
      break;
    case PTT_CLEAR_TRACE:
      {
        janitor_for_pinlock_t plj(&events_lock);
        trace_addrs.clear();
        trace_flush_needed = false;
      }
      ret = true;
      break;
    case PTT_ADD_BPT:
      MSG("Adding software breakpoint at %p\n", (void*)res->data);
      {
        janitor_for_pinlock_t plj(&soft_locks);
        addrvec_t::iterator p = std::lower_bound(soft_bpts.begin(), soft_bpts.end(), res->data);
        if ( p == soft_bpts.end() || *p != res->data )
        {
          soft_bpts.insert(p, res->data);
          // try to add instruction level instrumentation for handling breakpoints in case
          // it isn't already enabled
          add_instrumentation(TF_TRACE_INSN);
        }
      }

      ans.code = PTT_ACK;
      ans.data = trace_addrs.size();
      ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      break;
    case PTT_DEL_BPT:
      {
        {
          janitor_for_pinlock_t plj(&soft_locks);
          addrvec_t::iterator p = std::lower_bound(soft_bpts.begin(), soft_bpts.end(), res->data);
          if ( p != soft_bpts.end() )
            soft_bpts.erase(p);
        }

        ans.code = PTT_ACK;
        ans.data = trace_addrs.size();
        ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      }
      break;
    case PTT_CAN_READ_REGS:
      {
        pin_thid tid = (pin_thid)res->data;
        ans.code = (tracing_registers && threads_context.count(tid) > 0) ? PTT_ACK : PTT_ERROR;
        ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      }
      break;
    case PTT_READ_REGS:
      {
        pin_thid tid = (pin_thid)res->data;
        ret = handle_read_regs(tid);
      }
      break;
    case PTT_SET_TRACE:
      {
        uint32 trace_types = (uint32)res->data;
        MSG("Set trace to %d\n", trace_types);
        update_instrumentation(trace_types);
        ans.code = PTT_ACK;
        ret = send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__);
      }
      break;
    case PTT_SET_OPTIONS:
      ans.code = PTT_ACK;
      if ( send_packet(&ans, sizeof(idapin_packet_t), NULL, 0, __FUNCTION__) )
        ret = handle_limits();
      break;
    default:
      MSG("UNKNOWN PACKET RECEIVED WITH CODE %d\n", res->code);
      last_packet = "UNKNOWN " + res->code;
      PIN_ExitProcess(0);
      break;
  }
  DEBUG("LAST PACKET WAS %s\n", last_packet);
  return ret;
}

//--------------------------------------------------------------------------
static void check_network_error(ssize_t ret, const char *from_where)
{
  if ( ret == -1 )
  {
#ifdef _WIN32
    int err = WINDOWS::GetLastError();
    bool timeout = err == WSAETIMEDOUT;
#else
    bool timeout = errno == EAGAIN;
#endif
    if ( !timeout )
    {
      MSG("A network error happened, exiting from application...\n");
      PIN_ExitProcess(-1);
    }
    MSG("Timeout, called from %s\n", from_where);
  }
}

//--------------------------------------------------------------------------
static ssize_t pin_recv(PIN_SOCKET fd, void *buf, size_t n, const char *from_where)
{
  //debug("NET: RECEIVING "  + intstr((int)n) + " BYTES\n");
#ifdef _WIN32
  ssize_t ret = WINDOWS::recv(fd, (char *)buf, (int)n, MSG_WAITALL);
#else
  ssize_t ret = recv(fd, buf, n, MSG_WAITALL);
#endif
  check_network_error(ret, from_where);
  return ret;
}

//--------------------------------------------------------------------------
static ssize_t pin_send(PIN_SOCKET fd, const void *buf, size_t n, const char *from_where)
{
  //debug("NET: SENDING " + intstr((int)n) + " BYTES\n");
#ifdef _WIN32
  ssize_t ret = WINDOWS::send(fd, (const char *)buf, (int)n, 0);
#else
  ssize_t ret = send(fd, buf, n, 0);
#endif
  check_network_error(ret, from_where);
  return ret;
}

//--------------------------------------------------------------------------
static bool handle_packets(int total, const string &until_packet)
{
  int packets = 0;
  while ( (total != -1 && packets++ < total) || (!until_packet.empty() && last_packet != until_packet) )
  {
    DEBUG("Receiving 1st packets, expected %d bytes...\n", (uint32)sizeof(idapin_packet_t));
    idapin_packet_t res;
    ssize_t bytes = pin_recv(cli_socket, &res, sizeof(res), __FUNCTION__);
    if ( bytes != sizeof(res) )
    {
      error_msg("pin_recv");
      return false;
    }

    DEBUG("Handling packet ... \n");
    if ( !handle_packet(&res) )
    {
      MSG("Error handling %s packet, exiting...\n", last_packet);
      return false;
    }
  }

  if ( total == packets )
    DEBUG("Maximum number of packets reached, exiting from handle_packets...\n");
  else
    DEBUG("Expected packet '%s' received, exiting from handle_packets...\n", until_packet.c_str());

  return true;
}

//--------------------------------------------------------------------------
static bool listen_to_ida(void)
{
  // initialize the socket and connect to ida
  if ( !init_socket() )
  {
    MSG("listen_to_ida: init_socket() failed!\n");
    return false;
  }

  MSG("CONNECTED TO IDA\n");

  // Handle the 1st packets. Then, we will handle the next (variable
  // number) of packets to add breakpoints in the application start
  // callback and, finally, handle all the rest of packets send to the
  // PIN tool in the handle_packet function
  bool ret = handle_packets(5);
  MSG("Exiting from listen_to_ida\n");

  return ret;
}

//--------------------------------------------------------------------------
static void open_console(void)
{
#ifdef _WIN32
  if ( WINDOWS::AllocConsole() )
  {
    freopen("CONIN$" , "rb", stdin);
    freopen("CONOUT$", "wb", stdout);
    freopen("CONOUT$", "wb", stderr);
    std::ios::sync_with_stdio();
  }
#endif
}

//--------------------------------------------------------------------------
int main(int argc, char * argv[])
{
  // Initialize pin command line
  if ( PIN_Init(argc, argv) )
  {
    MSG("PIN_Init call failed!\n");
    return usage();
  }

  int value = knob_debug_mode.Value();
  if ( getenv("IDAPIN_DEBUG") != NULL || value > 0 )
  {
    debug_tracer = (value > 0) ? value : atoi(getenv("IDAPIN_DEBUG"));
    open_console();
    MSG("IDA PIN Tool version $Revision: #62 $\nInitializing PIN tool...\n\n");
  }

  dump_sizes();
  // Connect to IDA's debugger; it only returns in case of error
  if ( !listen_to_ida() )
    PIN_ExitApplication(-1);

  return 0;
}
