gentrace.cpp

#include "pin.H"

#include <cassert>
#include <iostream>
#include <fstream>
#include <sstream>
#include <stack>
#include <vector>
#include <cstring>
#include <stdint.h>
#include <cstdlib>
#include <time.h>
//#include "pin_frame.h"
//#include "pin_trace.h"
#include "reg_mapping_pin.h"
//#include "cache.h"

/* The new trace container format */
#include <libtrace/trace.container.hpp>

#include "pivot.h"

#include "pin_taint.h"
#include "pin_misc.h"

//#define DEBUG
#ifdef DEBUG
  #define dbg_printf(...) printf(__VA_ARGS__)
#else
  #define dbg_printf(...)
#endif


/*
 * A useful set of macros that we can customize for different
 * architectures as simply as possible.
 */
#ifdef ARCH_64
  #define BFD_ARCH frame_arch_i386
  #define BFD_MACH frame_mach_x86_64
  #define STACK_OFFSET 8
  #define MAX_ADDRESS "0xffffffffffffffff"
  #define MEM_ACCESS qword
#elif defined(ARCH_32)
  #define BFD_ARCH frame_arch_i386
  #define BFD_MACH frame_mach_i386_i386
  #define STACK_OFFSET 4
  #define MAX_ADDRESS "0xffffffff"
  #define MEM_ACCESS dword
#endif

using namespace pintrace;
using namespace SerializedTrace;

const ADDRINT ehandler_fs_offset = 0;
const ADDRINT ehandler_nptr_offset = 0;
const ADDRINT ehandler_handler_offset = 4;
const ADDRINT ehandler_invalid_ptr = ~0x0;
const ADDRINT ehandler_size = 8;
/** The offset esp has from when the exception is initially handled to
    when the handler is called. */
const ADDRINT ehandler_esp_offset = 0xe0;

const int maxSehLength = 10;

#ifdef _WIN32

const char* const windowsDll = "kernel32.dll";
const char* const wsDll = "WS2_32.dll";

const int callbackNum = 5;

const unsigned int accessViolation = 0xc0000005;

namespace WINDOWS {
#include "Winsock2.h"
#include "Windows.h"
}
#endif

/* Environment variables on windows
 *
 * For a program that uses getenv, Windows does the following:
 *
 * 1. Call GetEnvironmentStringsW and set up an environment table.
 * 2. If using main (rather than wmain), WideCharToMultiByte is used
 * to convert to a multibyte environment table.
 *
 * WideCharToMultiByte is implemented using a conversion table; we
 * don't handle control-flow taint, and thus we cannot really handle
 * tainting of environment variables :-/
 */

/* Networking on windows
 *
 * Winsock appears to communicate with a windows subsystem using a
 * lightweight procedure calling interface, e.g., something we don't
 * want to parse.  So, we catch sockets() by instrumenting the socket
 * call itself.
 */

//
// CONFIGURATION
//

// Note: since we only flush the buffer once per basic block, the number
// of instructions per block should never exceed BUFFER_SIZE.
// TODO: See if there's some way to overcome this limitation, if
// necessary.
#define BUFFER_SIZE 10240

// Leave this much extra room in the frame buffer, for some unexpected
// frames.
#define FUDGE 5

// Add a keyframe every KEYFRAME_FREQ instructions.
#define KEYFRAME_FREQ 10240

// Use value caching.
#define USE_CACHING

// Use faster functions to append to the value buffer, where possible.
//#define USE_FASTPATH

#ifdef USE_FASTPATH
#define _FASTPATH true
#else
#define _FASTPATH false
#endif

/** Set to 1 to enable lock debug information */
#ifndef DEBUG_LOCK
#define DEBUG_LOCK 0
#endif

KNOB<string> KnobOut(KNOB_MODE_WRITEONCE, "pintool",
                     "o", "trace.frames",
                     "Trace file to output to.");

KNOB<int> KnobTrigAddr(KNOB_MODE_WRITEONCE, "pintool",
                       "trig_addr", "",
                       "Address of trigger point. No logging will occur until execution reaches this address.");

KNOB<string> KnobTrigModule(KNOB_MODE_WRITEONCE, "pintool",
                            "trig_mod", "",
                            "Module that trigger point is in.");

KNOB<int> KnobTrigCount(KNOB_MODE_WRITEONCE, "pintool",
                        "trig_count", "0",
                        "Number of times trigger will be executed before activating.");

//
// NOTE: This limit is not a hard limit; the generator only stops logging
// during buffer flushes, so the actual number of instructions logged
// might exceed log_limit, but at the most by BUFFER_SIZE.
// Also note that the limit is in terms of the number of _instructions_,
// not frames; things like keyframes, LoadModuleFrames, etc. are not
// included in the count.
//
KNOB<uint64_t> KnobLogLimit(KNOB_MODE_WRITEONCE, "pintool",
                            "log-limit", "0",
                            "Number of instructions to limit logging to.");
KNOB<bool> LogAllSyscalls(KNOB_MODE_WRITEONCE, "pintool",
                          "log-syscalls", "false",
                          "Log system calls (even those unrelated to taint)");

KNOB<bool> KnobTaintTracking(KNOB_MODE_WRITEONCE, "pintool",
                             "taint-track", "true",
                             "Enable taint tracking");

KNOB<bool> LogAllAfterTaint(KNOB_MODE_WRITEONCE, "pintool",
                            "logall-after", "false",
                            "Log all (even untainted) instructions after the first tainted instruction");

KNOB<bool> LogAllBeforeTaint(KNOB_MODE_WRITEONCE, "pintool",
                             "logall-before", "false",
                             "Log all (even untainted) instructions before and after the first tainted instruction");

// This option logs one instruction.  It then generates a fake
// standard frame to include operands after the instruction executed.
KNOB<bool> LogOneAfter(KNOB_MODE_WRITEONCE, "pintool",
                       "logone-after", "false",
                       "Log the first instruction outside of the log range (taint-start/end), and then exit.");

KNOB<bool> LogKeyFrames(KNOB_MODE_WRITEONCE, "pintool",
                        "log-key-frames", "false",
                        "Periodically output key frames containing important program values");

KNOB<string> TaintedFiles(KNOB_MODE_APPEND, "pintool",
                          "taint-files", "",
                          "Consider the given files as being tainted");

KNOB<bool> TaintedArgs(KNOB_MODE_WRITEONCE, "pintool",
                       "taint-args", "false",
                       "Command-line arguments will be considered tainted");

KNOB<bool> TaintedStdin(KNOB_MODE_WRITEONCE, "pintool",
                        "taint-stdin", "false",
                        "Everything read from stdin will be considered tainted");

KNOB<bool> TaintedNetwork(KNOB_MODE_WRITEONCE, "pintool",
                          "taint-net", "false",
                          "Everything read from network sockets will be considered tainted");

KNOB<bool> TaintedIndices(KNOB_MODE_WRITEONCE, "pintool",
                          "taint-indices", "false",
                          "Values loaded with tainted memory indices will be considered tainted");

// FIXME: we should be able to specify more refined tainted
// sources, e.g., that only the 5th argument should be considered
// tainted
KNOB<string> TaintedEnv(KNOB_MODE_APPEND, "pintool",
                        "taint-env", "",
                        "Environment variables to be considered tainted");

KNOB<ADDRINT> TaintStart(KNOB_MODE_WRITEONCE, "pintool",
                          "taint-start", "0x0",
                          "All logged instructions will have higher addresses");

KNOB<ADDRINT> TaintEnd(KNOB_MODE_WRITEONCE, "pintool",
                        "taint-end", MAX_ADDRESS,
                        "All logged instructions will have lower addresses");

KNOB<string> FollowProgs(KNOB_MODE_APPEND, "pintool",
                         "follow-progs", "",
                         "Follow the given program names if they are exec'd");

KNOB<string> PivotFile(KNOB_MODE_WRITEONCE, "pintool",
                       "pivots-file", "",
                       "Load file of pivot gadgets");

KNOB<bool> SEHMode(KNOB_MODE_WRITEONCE, "pintool",
                   "seh-mode", "false",
                   "Record an SEH exploits");

KNOB<int> CheckPointFreq(KNOB_MODE_WRITEONCE, "pintool",
                         "freq", "10000",
                         "Report value of eip every n instructions.");

KNOB<int> CacheLimit(KNOB_MODE_WRITEONCE, "pintool",
                     "cache-limit", "500000000",
                     "Code-cache size limit (bytes)");

KNOB<int> SkipTaints(KNOB_MODE_WRITEONCE, "pintool",
                     "skip-taints", "0",
                     "Skip this many taint introductions");

struct FrameBuf {
    ADDRINT addr;
    uint32_t tid;
    uint32_t insn_length;

    // The raw instruction bytes will be stored as 16 bytes placed over 4
    // integers. The conversion is equivalent to the casting of a char[16]
    // to a uint32_t[4].
    // NOTE: This assumes that MAX_INSN_BYTES == 16!!!
    uint32_t rawbytes0;
    uint32_t rawbytes1;
    uint32_t rawbytes2;
    uint32_t rawbytes3;

    uint32_t values_count;
    ValSpecRec valspecs[MAX_VALUES_COUNT];

};

/**
 * Temporary structure used during instrumentation.
 */
typedef struct TempOps_s {
    uint32_t reg;
    RegMem_t type;
    uint32_t taint;
} TempOps_t;

/**
 * Posible ways of passing a register to an analysis function
 */
enum RPassType { P_VALUE, P_REF, P_CONTEXT, P_FPX87 };

/**
 * Given a register, decide how to pass it.
 */
static RPassType howPass(REG r) {
    dbg_printf("howPass r=%s\n", pin_register_name(r).c_str());

    if(REG_is_fr_for_get_context(r))
      return P_CONTEXT;

    /* XMM and floating point registers can be passed by reference */
    if (REG_is_xmm(r) || REG_is_ymm(r) || REG_is_mm(r))
        return P_REF;

    if(REG_is_fr_or_x87(r))
        return P_FPX87;

    // For now, let's just use context
    return P_CONTEXT;
}

/**
 * Avoiding logging some addresses.
 */
static bool dontLog(ADDRINT addr) {
    dbg_printf("dontLog addr=0x%lx\n", addr);

    IMG i = IMG_FindByAddress(addr);
    if (IMG_Valid(i)) {

        char tempbuf[BUFSIZE];
        char *tok = NULL;
        char *lasttok = NULL;

        // Fill up the temporary buffer
        strncpy(tempbuf, IMG_Name(i).c_str(), BUFSIZE);

        // We don't need a lock, since this is an instrumentation function (strtok is not re-entrant)
        strtok(tempbuf, "\\");

        while ((tok = strtok(NULL, "\\")) != NULL) {
            // Just keep parsing...
            lasttok = tok;
        }

        if (lasttok) {
            if (lasttok == string("uxtheme.dll")) {
                return true;
            }
        }
    }

    return false;
}



/**
 * This type preserves state between a system call entry and exit.
 */
typedef struct SyscallInfo_s {

    /** Frame for system call */
    frame sf;

    /** State shared between taintIntro and taintStart */
    uint32_t state;
} SyscallInfo_t;

/**
 * This type preserves state between a recv() call and return
 */
typedef struct RecvInfo_s {
    /** Fd */
    uint32_t fd;

    /** The address */
    void* addr;

    /** Bytes written ptr. */
    uint32_t *bytesOut;
} RecvInfo_t;

/**
 * Thread local information
 */

typedef struct ThreadInfo_s {
    // Stack keeping track of system calls
    // Needed because windows system calls can be nested!
    std::stack<SyscallInfo_t> scStack;
    std::stack<RecvInfo_t> recvStack;
    context delta;
} ThreadInfo_t;

int g_counter = 0;

//TraceWriter *g_tw;
TraceContainerWriter *g_twnew;

template <typename Iterator>
inline void add_frames(TraceContainerWriter* out,
                       Iterator first, Iterator last) {
    while (first != last) {
        out->add(*first++);
    }
}

// A taint tracker
TaintTracker * tracker;

FrameBuf g_buffer[BUFFER_SIZE];
uint32_t g_bufidx;

// Counter to keep track of when we should add a keyframe.
uint32_t g_kfcount;

// Caches.
//RegCache g_regcache;
//MemCache g_memcache;

// Profiling timer.
clock_t g_timer;

// True if logging is activated.
// Logging should be activated if it is possible for some instruction
// to be logged.  This could happen because 1) we are logging all
// instructions, or 2) taint is introduced, and so the instruction
// could be tainted.
bool g_active;

// Number of instructions logged so far.
uint64_t g_logcount;

// Number of instructions to limit logging to.
uint64_t g_loglimit;

// True if a trigger was specified.
bool g_usetrigger;

// Activate taint analysis
// bool t_active;

// Whether taint has been introduced
bool g_taint_introduced;

// True if the trigger address was resolved.
bool g_trig_resolved;

ADDRINT g_trig_addr;

// We use a signed integer because sometimes the countdown will be
// decremented past zero.
int g_trig_countdown;

// Name of our thread/process
char g_threadname[BUFFER_SIZE] = "s";

// A lock on any shared state
PIN_LOCK lock;

// An environment to keep all the values
ValSpecRec values[MAX_VALUES_COUNT];

// Address ranges
ADDRINT start_addr, end_addr;

// Pivot set
pivot_set ps;

// Exit after the next instruction
bool g_exit_next;

// Prototypes.
VOID Cleanup();

// Key for thread local system call stack
static TLS_KEY tl_key;


// Start of functions.

VOID ModLoad(IMG i, void*);

// Get Thread Info
ThreadInfo_t* GetThreadInfo(void) {
    ThreadInfo_t* ti;

    ti = static_cast<ThreadInfo_t*> (PIN_GetThreadData(tl_key, PIN_ThreadId()));
    assert(ti);
    return ti;
}

// Create a new thread information block for the current thread
ThreadInfo_t* NewThreadInfo(void) {
    ThreadInfo_t* ti = NULL;

    ti = new ThreadInfo_t;
    assert(ti);

    PIN_SetThreadData(tl_key, ti, PIN_ThreadId());

    return ti;
}

/** Given a REG, return the number of bits in the reg */
static uint32_t GetBitsOfReg(REG r) {
    dbg_printf("GetBitsOfReg r=%s\n", pin_register_name(r).c_str());
    if (REG_is_gr8(r)) return 8;
    if (REG_is_gr16(r)) return 16;
    if (REG_is_gr32(r)) return 32;
    if (REG_is_gr64(r)) return 64;

    /* REG_is_fr_or_x87 returns true on XMM registers and other
       non-x87 regs, so we can't use that. */
    if (REG_ST_BASE <= r && r <= REG_ST_LAST) return 80;

    string s = REG_StringShort(r);

    switch (r) {
    case REG_SEG_CS:
    case REG_SEG_DS:
    case REG_SEG_ES:
    case REG_SEG_FS:
    case REG_SEG_GS:
    case REG_SEG_SS:
        return 16;
        break;

    case REG_MXCSR:
        return 32;
        break;

    case REG_MM0:
    case REG_MM1:
    case REG_MM2:
    case REG_MM3:
    case REG_MM4:
    case REG_MM5:
    case REG_MM6:
    case REG_MM7:
        return 64;
        break;

    case REG_XMM0:
    case REG_XMM1:
    case REG_XMM2:
    case REG_XMM3:
    case REG_XMM4:
    case REG_XMM5:
    case REG_XMM6:
    case REG_XMM7:
        return 128;
        break;

    case REG_YMM0:
    case REG_YMM1:
    case REG_YMM2:
    case REG_YMM3:
    case REG_YMM4:
    case REG_YMM5:
    case REG_YMM6:
    case REG_YMM7:
        return 256;
        break;

/*
 * Handle any extra registers specific to an architecture
 * and defines any ambiguous registers
 * (E.g., REG_INST_PTR can be 64 or 32 based on the architecture)
 */
#if defined(ARCH_64)
    case REG_EIP:
    case REG_EFLAGS:
        return 32;
        break;

    case REG_INST_PTR:
    case REG_GFLAGS:
    case REG_SEG_GS_BASE:
    case REG_SEG_FS_BASE:
        return 64;
        break;

    case REG_XMM8:
    case REG_XMM9:
    case REG_XMM10:
    case REG_XMM11:
    case REG_XMM12:
    case REG_XMM13:
    case REG_XMM14:
    case REG_XMM15:
        return 128;
        break;

    case REG_YMM8:
    case REG_YMM9:
    case REG_YMM10:
    case REG_YMM11:
    case REG_YMM12:
    case REG_YMM13:
    case REG_YMM14:
    case REG_YMM15:
        return 256;
        break;

#elif defined(ARCH_32)
    case REG_INST_PTR:
    case REG_GFLAGS:
    case REG_SEG_GS_BASE:
    case REG_SEG_FS_BASE:
        return 32;
        break;
#endif

    default:
        break;
    }

    // Otherwise, exit because we don't know what's up
    cerr << "Warning: Unknown register size of register " << REG_StringShort(r) << endl;
    assert(false);
    return -1;
}

static size_t GetByteSize(RegMem_t vtype) {
    return (vtype.size / 8);
}

static size_t GetBitSize(RegMem_t type) {
    return type.size;
}

void LLOG(const char *str) {
#if DEBUG_LOCK
    LOG(str);
#else
    /* Disabled */
#endif
}

ADDRINT CheckTrigger()
{
    return --g_trig_countdown <= 0;
}

VOID Activate(CONTEXT *ctx)
{
    dbg_printf("Activate\n");
    cerr << "Activating logging" << endl;
    g_active = true;
    PIN_RemoveInstrumentation();
    PIN_ExecuteAt(ctx);
}

/** Activate taint analysis.

    Note: It's important to NOT hold locks when calling this function.
    PIN_RemoveInstrumentation obtains the VM lock, which is only possible
    when no analysis functions/etc are executing.  If one is waiting for
    one of our locks, this will cause a deadlock.
*/
VOID TActivate()
{
    dbg_printf("TActivate\n");
    cerr << "Activating taint analysis " << endl;
    g_active = true; /* Any instruction could be logged because taint is
                        introduced. */
    g_taint_introduced = true; /* Taint is definitely introduced now. */
    PIN_RemoveInstrumentation();
}

//
// Returns true if the buffer index with count added to it exceeds the
// maximum size of the buffer.
//
ADDRINT CheckBuffer(UINT32 count)
{
  dbg_printf("CheckBuffer count=%d\n", count);

  return (g_bufidx + count) >= BUFFER_SIZE - FUDGE;
}

ADDRINT CheckBufferEx(BOOL cond, UINT32 count, UINT32 count2)
{
  dbg_printf("CheckBufferEx cond=%d, count=%d, count2=%d\n", cond, count, count2);

  return cond && ((g_bufidx + count + count2) >= BUFFER_SIZE - FUDGE);
}

// Callers must ensure mutual exclusion
VOID FlushInstructions()
{
    dbg_printf("FlushInstructions\n");
    for(uint32_t i = 0; i < g_bufidx; i++) {

        frame fnew;
        fnew.mutable_std_frame()->set_address(g_buffer[i].addr);
        fnew.mutable_std_frame()->set_thread_id(g_buffer[i].tid);
        /* Ew. */
        fnew.mutable_std_frame()->set_rawbytes((void*)(&(g_buffer[i].rawbytes0)), g_buffer[i].insn_length);

        /* Add operands */

        // Go through each value and remove the ones that are cached.

        /* The operand_list is a required field, so we must access it
           even if there are no operands or protobuffers will complain to
           us. */
        fnew.mutable_std_frame()->mutable_operand_pre_list();

        for (uint32_t j = 0; j < g_buffer[i].values_count; j++) {

            ValSpecRec &v = g_buffer[i].valspecs[j];

            operand_info *o = fnew.mutable_std_frame()->mutable_operand_pre_list()->add_elem();
            o->set_bit_length(GetBitSize(v.type));
            o->mutable_operand_usage()->set_read(v.usage & RD);
            o->mutable_operand_usage()->set_written(v.usage & WR);
            /* XXX: Implement index and base */
            o->mutable_operand_usage()->set_index(false);
            o->mutable_operand_usage()->set_base(false);

            switch (v.taint) {
            case 0:
                o->mutable_taint_info()->set_no_taint(true);
                break;
            case -1:
                o->mutable_taint_info()->set_taint_multiple(true);
                break;
            default:
                o->mutable_taint_info()->set_taint_id(v.taint);
                break;
            }

            if (tracker->isMem(v.type)) {
                o->mutable_operand_info_specific()->mutable_mem_operand()->set_address(v.loc);

            } else {
                string t = pin_register_name((REG)v.loc);
                if (t == "Unknown") {
                    t = string("Unknown ") + REG_StringShort((REG)v.loc);
                }
                o->mutable_operand_info_specific()->mutable_reg_operand()->set_name(t);
            }

            o->set_value(&(v.value), GetByteSize(v.type));

            // We're in trouble if we don't know the type.
            if(v.type.type != REGISTER && v.type.type != MEM) {
                cerr << "v.type = " << v.type.type << endl;
                assert(false);
            }
        }

        g_twnew->add(fnew);
    }

    // Update counts.
    g_logcount += g_bufidx;
    g_kfcount += g_bufidx;

    g_bufidx = 0;

}

/* Add a PIN register to a value list. Helper function for FlushBuffer */
VOID AddRegister(tagged_value_list *tol, const CONTEXT *ctx, REG r, THREADID threadid) {
    dbg_printf("AddRegister r=%d, tid=%d\n", r, threadid);

    tol->mutable_value_source_tag()->set_thread_id(threadid);
    value_info *v = tol->mutable_value_list()->add_elem();
    v->mutable_operand_info_specific()->mutable_reg_operand()->set_name(REG_StringShort(r));
      size_t s_bytes = GetBitsOfReg(r) / 8;
      v->set_bit_length(s_bytes * 8);

    /* Make sure this register even fits in the context.  PIN would
       probably throw an error, but it's good to be paranoid. */
    assert (s_bytes <= sizeof(ADDRINT));
    ADDRINT regv = PIN_GetContextReg(ctx, r);
    v->set_value((void*)(&regv), s_bytes);
    //std::copy((uint8_t*) (&regv), ((uint8_t*) (&regv)) + s_bytes, v->
}

//
// Writes all instructions stored in the buffer to disk, and resets the
// buffer index to 0. Also checks to see if we need to insert a
// keyframe. If so, inserts the keyframe using the data in the supplied
// context.
//
VOID FlushBuffer(BOOL addKeyframe, const CONTEXT *ctx, THREADID threadid, BOOL needlock)
{
    dbg_printf("FlushBuffer\n");

    LLOG("Begin flushing buffer.\n");

    if (needlock) {
        PIN_GetLock(&lock, threadid+1);
    }

    FlushInstructions();


    // Check to see if we should insert a keyframe here.
    if (addKeyframe && (g_kfcount >= KEYFRAME_FREQ) && LogKeyFrames) {

        //LOG("Inserting keyframe:\n");
        //LOG("  addr: " + hexstr(PIN_GetContextReg(ctx, REG_EIP)) + "\n");

        assert(ctx);

        frame f;
        tagged_value_list *tol = f.mutable_key_frame()->mutable_tagged_value_lists()->add_elem();
        AddRegister(tol, ctx, LEVEL_BASE::REG_GAX, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GBX, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GCX, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GDX, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GSI, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GDI, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_STACK_PTR, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GBP, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_GFLAGS, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_CS, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_DS, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_SS, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_ES, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_FS, threadid);
        AddRegister(tol, ctx, LEVEL_BASE::REG_SEG_GS, threadid);

        g_twnew->add(f);

        g_kfcount = 0;

    }

    // See if we've gotten sufficient instructions.
    if ((g_loglimit != 0) && (g_logcount >= g_loglimit)) {
        LOG("Logged required number of instructions, quitting.\n");
        cerr << "Logged required number of instructions, quitting." << endl;
        Cleanup();
        //PIN_ReleaseLock(&lock);
        // Never release lock
        //PIN_Detach();
        exit(0);
    } else {
        if (needlock) {
            PIN_ReleaseLock(&lock);
        }
    }


    dbg_printf("End FlushBuffer\n");
    LLOG("End flushing buffer.\n");

}

#ifdef _WIN32

/** Wrapper for accept */
uint32_t AcceptWrapper(CONTEXT *ctx, AFUNPTR fp, THREADID tid, uint32_t s, void *addr, int *addrlen) {

    cerr << "AcceptWrapper" << endl;

    uint32_t ret;

    PIN_CallApplicationFunction(ctx, tid,
                                CALLINGSTD_STDCALL, fp,
                                PIN_PARG(uint32_t), &ret,
                                PIN_PARG(uint32_t), s,
                                PIN_PARG(void*), addr,
                                PIN_PARG(int*), addrlen,
                                PIN_PARG_END());

    PIN_GetLock(&lock, tid+1);
    tracker->acceptHelper(ret);
    PIN_ReleaseLock(&lock);

    return ret;

}

/** Wrapper for WSAConnect */
uint32_t WSAConnectWrapper(CONTEXT *ctx, AFUNPTR fp, THREADID tid, uint32_t s, void *arg2, void *arg3, void *arg4, void *arg5, void *arg6, void *arg7) {

    cerr << "WSAConnectWrapper" << endl;

    uint32_t ret;

    cerr << "Connect to socket " << s << endl;

    PIN_CallApplicationFunction(ctx, tid,
                                CALLINGSTD_STDCALL, fp,
                                PIN_PARG(uint32_t), &ret,
                                PIN_PARG(uint32_t), s,
                                PIN_PARG(void*), arg2,
                                PIN_PARG(void*), arg3,
                                PIN_PARG(void*), arg4,
                                PIN_PARG(void*), arg5,
                                PIN_PARG(void*), arg6,
                                PIN_PARG(void*), arg7,
                                PIN_PARG_END());

    PIN_GetLock(&lock, tid+1);
    if (ret != SOCKET_ERROR) {
        tracker->acceptHelper(s);
    } else {
        cerr << "WSAConnect error " << ret << endl;
    }
    PIN_ReleaseLock(&lock);

    return ret;

}

/** Wrapper for connect */
uint32_t ConnectWrapper(CONTEXT *ctx, AFUNPTR fp, THREADID tid, uint32_t s, void *arg2, void *arg3) {

    cerr << "ConnectWrapper" << endl;

    uint32_t ret;

    cerr << "Connect to socket " << s << endl;

    PIN_CallApplicationFunction(ctx, tid,
                                CALLINGSTD_STDCALL, fp,
                                PIN_PARG(uint32_t), &ret,
                                PIN_PARG(uint32_t), s,
                                PIN_PARG(void*), arg2,
                                PIN_PARG(void*), arg3,
                                PIN_PARG_END());

    PIN_GetLock(&lock, tid+1);
    //  if (ret != SOCKET_ERROR) {
    // Non-blocking sockets will return an "error".  However, we can't
    // call GetLastError to find out what the root problem is,
    // so... we'll just assume the connection was successful.
    tracker->acceptHelper(s);

    // } else {
    //    cerr << "connect error " << ret << endl;
    //  }
    PIN_ReleaseLock(&lock);

    return ret;

}

void BeforeRecv(THREADID tid, uint32_t s, char* buf) {

    RecvInfo_t r;

    r.fd = s;
    r.addr = buf;
    r.bytesOut = NULL;

    ThreadInfo_t *ti = GetThreadInfo();

    ti->recvStack.push(r);
}

void WSABeforeRecv(THREADID tid, uint32_t s, WINDOWS::LPWSABUF bufs, WINDOWS::LPDWORD bytesOut) {

    RecvInfo_t r;

    r.fd = s;
    r.addr = bufs[0].buf;
    r.bytesOut = (uint32_t*) bytesOut;

    ThreadInfo_t *ti = GetThreadInfo();

    ti->recvStack.push(r);
}

void AfterRecv(THREADID tid, int ret, char *f) {
    cerr << "afterrecv called by " << f << endl;
    ThreadInfo_t *ti = GetThreadInfo();
    uint32_t len = 0;

    if (ti->recvStack.empty()) {
        cerr << "WARNING: Stack empty in AfterRecv(). Thread " << tid << endl;
    } else {

        RecvInfo_t ri = ti->recvStack.top();
        ti->recvStack.pop();

        if (ret != SOCKET_ERROR) {
            PIN_GetLock(&lock, tid+1);
            //cerr << "fd: " << ri.fd << endl;

            uint32_t numbytes = 0;
            if (ri.bytesOut) {
                numbytes = *(ri.bytesOut);
            } else {
                numbytes = ret;
            }

            FrameOption_t fo = tracker->recvHelper(ri.fd, ri.addr, numbytes);
            PIN_ReleaseLock(&lock);

            if (fo.b) {

                if (!g_taint_introduced) {
                    TActivate();
                }

                PIN_GetLock(&lock, tid+1);
                g_twnew->add(fo.f);
                PIN_ReleaseLock(&lock);
            }
        } else {
            cerr << "recv() error " << endl;
        }
    }
}


/** Wrapper for calling GetEnvironmentStringsW() and tainting the output */
void* GetEnvWWrap(CONTEXT *ctx, AFUNPTR fp, THREADID tid) {
    void *ret = NULL;

    /*
       We must lock after the PIN_CallApplicationFunction call, since
       the called code is instrumented, and also tries to obtain the
       lock.

       This probably is not a big deal, but theoretically the
       instrumented code in another thread could change the memory as
       we're reading it.  This seems pretty unlikely.  If we ever feel
       like fixing it, we could obtain PIN's global vm lock.
    */

    PIN_CallApplicationFunction(ctx, tid,
                                CALLINGSTD_STDCALL, fp,
                                PIN_PARG(uint32_t), &ret,
                                PIN_PARG_END());

    LLOG("Getting lock in callback\n");

    PIN_GetLock(&lock, tid+1);
    LLOG("Got callback lock\n");

    std::vector<frame> frms = tracker->taintEnv(NULL, (wchar_t*) ret);
    add_frames(g_twnew, frms.begin(), frms.end());

    PIN_ReleaseLock(&lock);
    LLOG("Releasing callback lock\n");

    return ret;
}

/** Wrapper for calling GetEnvironmentStringsA() and tainting the output */
void* GetEnvAWrap(CONTEXT *ctx, AFUNPTR fp, THREADID tid) {
    void *ret = NULL;

    cerr << "In a wrap " << endl;

    /*
       We must lock after the PIN_CallApplicationFunction call, since
       the called code is instrumented, and also tries to obtain the
       lock.

       This probably is not a big deal, but theoretically the
       instrumented code in another thread could change the memory as
       we're reading it.  This seems pretty unlikely.  If we ever feel
       like fixing it, we could obtain PIN's global vm lock.
    */

    PIN_CallApplicationFunction(ctx, tid,
                                CALLINGSTD_STDCALL, fp,
                                PIN_PARG(uint32_t), &ret,
                                PIN_PARG_END());

    LLOG("Getting lock in callback\n");

    PIN_GetLock(&lock, tid+1);
    LLOG("Got callback lock\n");

    std::vector<frame> frms = tracker->taintEnv((char*) ret, NULL);
    add_frames(g_twnew, frms.begin(), frms.end());

    PIN_ReleaseLock(&lock);
    LLOG("Releasing callback lock\n");

    return ret;
}

#endif

/* This analysis function is called after the target instruction is
 * executed when using -logone-after. It transfer control back to the
 * same instruction to log its operands after execution. */
VOID PostInstruction(ADDRINT addr, CONTEXT *ctx) {
    dbg_printf("PostInstruction addr=0x%lx\n", addr);

    g_exit_next = true;
    PIN_SetContextReg(ctx, REG_INST_PTR, addr);
    PIN_ExecuteAt(ctx);
}

VOID AppendBuffer(ADDRINT addr,
                  THREADID tid,
                  CONTEXT *ctx,
                  BOOL isBranch,
                  UINT32 insn_length,

                  UINT32 rawbytes0,
                  UINT32 rawbytes1,
                  UINT32 rawbytes2,
                  UINT32 rawbytes3,

                  /* Type contains the type of the operand. Location
                   * specifies the base address for memory operands.
                   * For registers, this holds the ID of the
                   * register.  Value is used to pass register values
                   * by reference.  For memory operands, it holds the
                   * byte offset into memory.  For instance, a 32-bit
                   * memory operand is broken into four 8-bit operands
                   * with the same address (specified in location),
                   * but with different offsets (0, 1, 2, 3) in
                   * value.  Usage specifies how the operand is used
                   * (read, write, etc.) */

                  UINT32 values_count,
                  ...
                  )
{
    dbg_printf("AppendBuffer addr=0x%lx tid=%d isBranch=%d, insn_length=%d "
               "rawbytes=%x %x %x %x\n", addr, tid, isBranch, insn_length,
               rawbytes0, rawbytes1, rawbytes2, rawbytes3);

    va_list va;
    va_start(va, values_count);

    static int firstTaint = true;
    static int firstLogged = true;
    REG r;

    //LOG("APPEND: " + hexstr(addr) + "\n");

    /* BUILD_VAL touches values, so we need the lock early. */

    /* Periodically report eip. */
    if ((g_counter++ % CheckPointFreq.Value()) == 0) {
        // PIN_LockClient();
        // IMG i = IMG_FindByAddress(addr);
        // PIN_UnlockClient();
        // cerr << "Checkpoint: Executing code at " << addr;
        // if (IMG_Valid(i)) {
        //   cerr << " (" << IMG_Name(i) << ")";
        // }
        cerr << " thread " << tid
             << "; " << g_counter << " instructions" << endl
             << "Code cache size is " << CODECACHE_CodeMemUsed() << endl
             << "Code cache limit is " << CODECACHE_CacheSizeLimit() << endl;
    }

    LLOG("big thing\n");

    PIN_GetLock(&lock, tid+1);

    LLOG("got big thing\n");

    for (unsigned int i = 0; i < values_count; i++) {
        values[i].type.type = (RegMemEnum_t)va_arg(va, uint32_t);
        assert(valid_regmem_type(values[i].type));

        values[i].type.size = va_arg(va, uint32_t);
        values[i].loc = va_arg(va, ADDRINT);
        values[i].value.MEM_ACCESS[0] = va_arg(va, ADDRINT);
        values[i].usage = va_arg(va, uint32_t);
        if (tracker->isMem(values[i].type)) {
            /* Add memory byte offset */
            values[i].loc += values[i].value.MEM_ACCESS[0];
        }
    }

    /* Perform taint propagation and checking */

    bool abort = false;
    bool log_addr =
        ((start_addr <= addr) && (addr <= end_addr)) || LogOneAfter.Value();
    bool log_all =
        ((LogAllAfterTaint.Value() && !firstTaint)
         || LogAllBeforeTaint.Value());
    uint32_t pretaint[MAX_VALUES_COUNT];
    LEVEL_VM::PIN_REGISTER *pr = NULL;
    ThreadInfo_t *ti = NULL;

    ti = GetThreadInfo();

    tracker->setCount(values_count);

    bool has_taint = tracker->hasTaint(ti->delta);

    if ((log_all || has_taint) && log_addr) {

        /* This instruction is tainted, or we're logging all
         * instructions */

        if (firstLogged) {
            cerr << "First logged instruction" << endl;
            firstLogged = false;
        }

        if (has_taint && firstTaint) {
            cerr << "First tainted instruction" << endl;
            LOG("First tainted instruction.\n");
            firstTaint = false;
        }

        // Mark everything as untainted
        for (uint32_t i = 0 ; i < values_count ; i++)
            values[i].taint = 0;

        // Set taint values from taint context
        tracker->setTaintContext(ti->delta);

        // Record pretaint (this goes in the log)
        for (uint32_t i = 0 ; i < values_count ; i++)
            pretaint[i] = values[i].taint;

        // Did this instruction propagate taint?
        //propagated_taint = tracker->propagatedTaint(isBranch);

        if (!isBranch)
            tracker->taintPropagation(ti->delta);

        // Taint checking
        abort = !tracker->taintChecking();

        //} FIXME: it there a case where the instruction contains taint
        //  but we do not need to log it?
        //if (log || (has_taint && propagated_taint)) {

        //cerr << "Logging instruction " << rawbytes0 << " " << rawbytes1 << endl;

        // Now, fill in the buffer with information

        assert (g_bufidx < BUFFER_SIZE);

        g_buffer[g_bufidx].addr = addr;
        g_buffer[g_bufidx].tid = tid;
        g_buffer[g_bufidx].insn_length = insn_length;

        g_buffer[g_bufidx].rawbytes0 = rawbytes0;
        g_buffer[g_bufidx].rawbytes1 = rawbytes1;
        g_buffer[g_bufidx].rawbytes2 = rawbytes2;
        g_buffer[g_bufidx].rawbytes3 = rawbytes3;

        // tracker->printRegs();

        g_buffer[g_bufidx].values_count = values_count;

        //g_buffer[g_bufidx].valspecs[i].taint = values[i].taint;

        // Values for floating point operations
        bool got_FP_state = false;
        FPSTATE fpState;
        void * fpValue;
        uint32_t s_i;

        // Store information to the buffer
        for (unsigned int i = 0; i < values_count; i++) {

            g_buffer[g_bufidx].valspecs[i].type = values[i].type;
            g_buffer[g_bufidx].valspecs[i].usage = values[i].usage;
            g_buffer[g_bufidx].valspecs[i].loc = values[i].loc;
            g_buffer[g_bufidx].valspecs[i].taint = pretaint[i];

            if(values[i].type.type == REGISTER) {
                       /*r = REG_FullRegName((REG) valspec##i##_loc);*/
                r = (REG)values[i].loc;

                /* Find how we should access the register value */
                switch(howPass(r)) {
                case P_CONTEXT:
                      g_buffer[g_bufidx].valspecs[i].value.MEM_ACCESS[0] =
                          PIN_GetContextReg(ctx, r);
                    break;

                case P_REF:
                    dbg_printf("in AppendBuffer: getting register value.\n");
                    pr = (LEVEL_VM::PIN_REGISTER*) values[i].value.MEM_ACCESS[0];
                    memcpy(&(g_buffer[g_bufidx].valspecs[i].value),
                           pr,
                           sizeof(LEVEL_VM::PIN_REGISTER));
                    dbg_printf("in AppendBuffer: got register value %#Lx.\n", pr);
                    break;

                case P_FPX87:
                    if(!got_FP_state) {
                        /* This is relativly expensive, so only do it once per
                           instruction */
                        PIN_GetContextFPState(ctx, &fpState);
                        got_FP_state = true;
                    }

                    // Figure out which st register we are using
                    s_i = r - REG_ST_BASE;

                    if (s_i > 7) {
                        cerr << "Unknown FP register " << r << " at addr "
                             << addr << endl;
                        assert(false);
                    }

                    fpValue = (void *)&(fpState.fxsave_legacy._sts[s_i]);

                    memcpy(&(g_buffer[g_bufidx].valspecs[i].value.flt[0]),
                           fpValue,
                           10); // FP are 80 bits = 10 bytes
                    break;

                default:
                    assert(false);
                    break;
                }
            } else if(values[i].type.type == MEM) {
                PIN_SafeCopy((VOID*) &(g_buffer[g_bufidx].valspecs[i].value),
                             (const VOID *)(g_buffer[g_bufidx].valspecs[i].loc),
                             GetByteSize(values[i].type));
            } else {
                cerr << "Unknown operand type at addr "
                     << addr << endl;
                assert(false);
            }

            //   cerr << "Building val specs now" << endl;

        }

        //   cerr << "... done" << endl;

        g_bufidx++;
    }


    /* For a non-SEH exploit, stop if taint checking fails.  In an SEH
       exploit, we may want an exception to trigger (e.g., from
       returning to a bad address). */
    if (abort && !SEHMode.Value()) {
        pivot_set::iterator i;
        cerr << "Stack smashing detected! @" << addr << endl;
        cerr << "Exiting...." << endl;
        LOG("Stack smashing detected\n");

        /* Let's assume this was a ret for now and increment esp by
         * four.  Of course, this isn't correct in general.  XXX: Fix
         * this so it works for any last instruction. */
        ADDRINT gsp = PIN_GetContextReg(ctx, REG_STACK_PTR);
        PIN_SetContextReg(ctx, REG_STACK_PTR, gsp + STACK_OFFSET);

        PIVOT_testpivot(ps, ctx, *tracker);

        FlushBuffer(true, ctx, tid, false);
        Cleanup();
        exit(0);
    }

    if (g_exit_next) {
        FlushBuffer(true, ctx, tid, false);
        Cleanup();
        exit(0);
    }

    PIN_ReleaseLock(&lock);
    LLOG("released big thing\n");

    va_end(va);

    return;

}

VOID InstrBlock(BBL bbl)
{
    dbg_printf("InstrBlock\n");
    // Now we need to get the values.

    uint32_t valcount;
    uint32_t icount = BBL_NumIns(bbl);

    // Used to temporarily store the values we obtain from the operands,
    // to faciliate further analysis for fast paths.
    TempOps_t opndvals[MAX_VALUES_COUNT];


    // LOG("INS: BBL start.\n");

    if (g_active) {

        if (icount > BUFFER_SIZE) {
            LOG("WARNING: Basic block too many instructions: " +
                decstr(icount) + "\n");
            assert(false);
        }

        // Add instrumentation call to check if buffer needs to be flushed.
        BBL_InsertIfCall(bbl, IPOINT_BEFORE,
                         (AFUNPTR) CheckBuffer,
                         IARG_UINT32, icount,
                         IARG_END);

        BBL_InsertThenCall(bbl, IPOINT_BEFORE,
                           (AFUNPTR) FlushBuffer,
                           IARG_BOOL, true,
                           IARG_CONTEXT,
                           IARG_THREAD_ID,
                           IARG_BOOL, true,
                           IARG_END);

    }

    //LOG("INS: BBL ins start.\n");

    // Count of instructions that have yet to be inserted into the buffer,
    // at the point at which the current instruction will be executed.
    uint32_t insLeft = icount;

    for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins)) {

        if (!g_active && g_usetrigger) {
            // Logging has not been activated yet, so all we need to do now
            // is check for the trigger condition.

            if (INS_Address(ins) == g_trig_addr) {
                // Found the trigger address.

                INS_InsertIfCall(ins, IPOINT_BEFORE,
                                 (AFUNPTR) CheckTrigger,
                                 IARG_END);

                INS_InsertThenCall(ins, IPOINT_BEFORE,
                                   (AFUNPTR) Activate,
                                   IARG_CONTEXT,
                                   IARG_END);

            }

            // Skip the rest of the analysis and immediately go on to the
            // next instruction.
            continue;

        }

        // Skip instrumentation unless g_active is enabled
        if (!g_active) {
            continue;
        }

        // Add instrumentation call to insert instruction into buffer.

        if (INS_Category(ins) == XED_CATEGORY_X87_ALU) {

            // TODO: Handle floating point instructions.
            // LOG("Not logging FP instruction.\n");

            // cerr << "Not logging FP instruction @" << INS_Address(ins) << ": " << INS_Disassemble(ins) << endl;
            // continue;

        } else if (INS_Category(ins) == XED_CATEGORY_PREFETCH) {
            LOG("Not logging prefetch instruction.\n");
            cerr << "Not logging prefetch instruction @" << INS_Address(ins) << ": " << INS_Disassemble(ins) << endl;
            continue;
        } else if (INS_Category(ins) == XED_CATEGORY_MMX) {
            LOG("Not logging mmx instruction.\n");
            cerr << "Not logging mmx instruction @" << INS_Address(ins) << ": " << INS_Disassemble(ins) << endl;
            continue;
        } else if (INS_Category(ins) == XED_CATEGORY_FCMOV) {
            LOG("Not logging float move instruction.\n");
            cerr << "Not logging float move instruction @" << INS_Address(ins) << ": " << INS_Disassemble(ins) << endl;
            continue;
        }

        // Check if there's a REP prefix.
        if (INS_HasRealRep(ins)) {

            INS_InsertIfCall(ins, IPOINT_BEFORE,
                             (AFUNPTR) CheckBufferEx,
                             IARG_FIRST_REP_ITERATION,
                             IARG_REG_VALUE, INS_RepCountRegister(ins),
                             IARG_UINT32, insLeft,
                             IARG_END);

            INS_InsertThenCall(ins, IPOINT_BEFORE,
                               (AFUNPTR) FlushBuffer,
                               IARG_BOOL, false,
                               IARG_ADDRINT, 0,
                               IARG_THREAD_ID,
                               IARG_BOOL, true,
                               IARG_END);

        }

        // The argument list to be passed into the instruction analysis call.
        IARGLIST arglist = IARGLIST_Alloc();
        IARGLIST arglist_helper = IARGLIST_Alloc();
        valcount = 0;

        // The first few arguments to AppendBuffer.
        IARGLIST_AddArguments(arglist,
                              IARG_ADDRINT, INS_Address(ins),
                              IARG_THREAD_ID,
                              IARG_CONTEXT,
                              IARG_BOOL, INS_IsBranch(ins),
                              IARG_UINT32, INS_Size(ins),
                              IARG_END);

        // Now we need to gather the instruction bytes.

        // Wastes a copy, but required because the instruction may not be
        // 32-bit aligned, and we want to respect word alignment requirements.
        uint32_t rawbytes_i[4];
        // Is it an xor?
        bool is_xor = false;

        UINT sz = INS_Size(ins);
        assert(PIN_SafeCopy((void*)rawbytes_i, (const void*) INS_Address(ins), sz) == sz);

        IARGLIST_AddArguments(arglist,
                              IARG_UINT32, rawbytes_i[0],
                              IARG_UINT32, rawbytes_i[1],
                              IARG_UINT32, rawbytes_i[2],
                              IARG_UINT32, rawbytes_i[3],
                              IARG_END);

        for (uint32_t i = 0; i < MAX_VALUES_COUNT; i++) {
            opndvals[i].taint = 0;
            opndvals[i].type = INVALIDREGMEM;
        }

        // specializing xors
        if (INS_Mnemonic(ins) == string("XOR") ||
            INS_Mnemonic(ins) == string("PXOR")) {
            int opnum = -1;
            bool found = false;
            REG r = REG_INVALID();

            /* Find the source and destination operand. */
            for (uint32_t i = 0 ; i < INS_OperandCount(ins); i++) {
                if (INS_OperandReadAndWritten (ins, i) &&
                    INS_OperandIsReg(ins, i)) {
                } else {
                    found = true;
                    r = INS_OperandReg(ins, i);
                    opnum = -1;
                    break;
                }
            }

            /* Find the second operand, and ensure it's the same register
               as the first operand we found. */
            if (found) {
                for (uint32_t i = 0 ; i < INS_OperandCount(ins); i++) {
                    if (INS_OperandReadAndWritten (ins, i) &&
                        INS_OperandIsReg(ins, i) &&
                        r == INS_OperandReg(ins, i) &&
                        (unsigned)opnum != i) {
                        is_xor = true;
                        break;
                    }
                }
            }
        } /* end xor code */

        for(uint32_t i = 0; i < INS_OperandCount(ins); i++) {

            opndvals[valcount].taint = 0;
            if (INS_OperandRead(ins, i) && (!is_xor))
                opndvals[valcount].taint = RD;
            if (INS_OperandWritten(ins, i))
                opndvals[valcount].taint |= WR;

            /* Handle register operands */
            if (INS_OperandIsReg(ins, i)) {

                REG r = INS_OperandReg(ins, i);
                if(r == REG_INVALID()) {
                  cerr << "Warning: invalid register operand in " << INS_Disassemble(ins) << endl;
                  continue;
                }
                assert(r != REG_INVALID());
                opndvals[valcount].reg = r;
                opndvals[valcount].type.type = REGISTER;

                // This was causing problems with movd %eax, %xmm0,
                // because %xmm0's operand width is 32, but BAP needs
                // to know the full operand size, which is 128.
                // opndvals[valcount].type.size = INS_OperandWidth(ins, i);

                opndvals[valcount].type.size = GetBitsOfReg(r);

                REG fullr = REG_FullRegName(r);
                if (fullr != REG_INVALID() && fullr != r) {
                  /* We know the fuller register, so just use that! */
                    //        cerr << "partial " << REG_StringShort(r) << " full " << REG_StringShort(fullr) << endl;
                    opndvals[valcount].reg = fullr;
                    opndvals[valcount].type.type = REGISTER;
                    opndvals[valcount].type.size = GetBitsOfReg(fullr);
                }

                valcount++;

            } else if (INS_OperandIsMemory(ins, i) ||
                       INS_OperandIsAddressGenerator(ins, i)) {


                /* Note: Compiled code sometimes uses LEA instructions for
                 * arithmetic.  As such, we always want to treat reads of
                 * these base/index registers as tainted. */

                REG basereg = INS_OperandMemoryBaseReg(ins, i);
                if (basereg != REG_INVALID()) {

                    opndvals[valcount].reg = basereg;
                    opndvals[valcount].type.type = REGISTER;
                    opndvals[valcount].type.size = GetBitsOfReg(basereg);

                    if (TaintedIndices || INS_OperandIsAddressGenerator(ins, i))
                        opndvals[valcount].taint = RD;
                    else
                        opndvals[valcount].taint = 0;

                    valcount++;

                }

                REG idxreg = INS_OperandMemoryIndexReg(ins, i);
                if (idxreg != REG_INVALID()) {

                    opndvals[valcount].reg = idxreg;
                    opndvals[valcount].type.type = REGISTER;
                    opndvals[valcount].type.size = GetBitsOfReg(idxreg);

                    if (TaintedIndices || INS_OperandIsAddressGenerator(ins, i))
                        opndvals[valcount].taint = RD;
                    else
                        opndvals[valcount].taint = 0;

                    valcount++;

                }
            }
        }

        bool memRead = INS_IsMemoryRead(ins);
        bool memRead2 = INS_HasMemoryRead2(ins);
        bool memWrite = INS_IsMemoryWrite(ins);

        // Value type of memory read.
        RegMem_t memReadTy = {NONE , 0};
        if (memRead || memRead2) {
            memReadTy.size = (INS_MemoryReadSize(ins) * 8);
            memReadTy.type = MEM;
        }
        // Value type of memory write
        RegMem_t memWriteTy = {NONE , 0};
        if (memWrite) {
            memWriteTy.size = (INS_MemoryWriteSize(ins) * 8);
            memWriteTy.type = MEM;
        }

        // Insert the operand values we've previously identified into the arglist.
        for (unsigned int i = 0; i < valcount; i++) {

            // cerr << opndvals[i].type << " " << i << " " << valcount << endl;

            // LOG("Adding: " + REG_StringShort((REG)opndvals[i].reg) + "\n");

            /*
             * PIN has several ways of passing register values to analysis
             * functions.  Unfortunately, none of them works all the
             * time.  So, we need to decide how to pass the value, and set
             * the *_value arguments to AppendBuffer accordingly.
             */
            switch (howPass((REG) opndvals[i].reg)) {
            case P_FPX87:
            case P_CONTEXT:
                IARGLIST_AddArguments(arglist_helper,
                                      IARG_UINT32, (uint32_t)(opndvals[i].type.type),
                                      IARG_UINT32, opndvals[i].type.size,
                                      IARG_UINT32, opndvals[i].reg,
                                      /* We don't need the value
                                         argument for contexts */
                                      IARG_PTR, 0,
                                      IARG_UINT32, opndvals[i].taint,
                                      IARG_END);
                break;

            case P_REF:
                IARGLIST_AddArguments(arglist_helper,
                                      IARG_UINT32, (uint32_t)(opndvals[i].type.type),
                                      IARG_UINT32, opndvals[i].type.size,
                                      IARG_UINT32, opndvals[i].reg,
                                      /* Pass reference pointer */
                                      IARG_REG_CONST_REFERENCE, opndvals[i].reg,
                                      IARG_UINT32, opndvals[i].taint,
                                      IARG_END);
                break;

            default:
                cerr << "Unknown value passing method" << endl;
                assert(false);
            }
        }

        /* We break up memory operands into byte-wise operands.  This is
         * essential for taint analysis.  Code that utilizes taint
         * analysis assumes that a tainted value can be computed (e.g.,
         * symbolically executed) using the instructions in the trace.
         * However, if some of a memory operand are not tainted, then
         * they could have changed.  Thus, we must break up memory
         * operands to make this explicit. */

        if (memRead) {
            size_t bytes = GetByteSize(memReadTy);

            for (size_t offset = 0; offset < bytes; offset++) {
                IARGLIST_AddArguments(arglist_helper,
                                      IARG_UINT32, (uint32_t)MEM,
                                      IARG_UINT32, 8, // one byte
                                      IARG_MEMORYREAD_EA,
                                      //IARG_MEMORYREAD_SIZE,
                                      IARG_UINT32, offset,
                                      IARG_UINT32, RD,
                                      IARG_END);
                valcount++;
            }
        }

        if (memRead2) {
            size_t bytes = GetByteSize(memReadTy);

            for (size_t offset = 0; offset < bytes; offset++) {
                IARGLIST_AddArguments(arglist_helper,
                                      IARG_UINT32, (uint32_t)MEM,
                                      IARG_UINT32, 8, // one byte
                                      IARG_MEMORYREAD2_EA,
                                      //IARG_MEMORYREAD_SIZE,
                                      IARG_UINT32, offset,
                                      IARG_UINT32, RD,
                                      IARG_END);
                valcount++;
            }
        }

        if (memWrite) {
            size_t bytes = GetByteSize(memWriteTy);

            for (size_t offset = 0; offset < bytes; offset++) {

                IARGLIST_AddArguments(arglist_helper,
                                      IARG_UINT32, (uint32_t)MEM,
                                      IARG_UINT32, 8, // one byte
                                      IARG_MEMORYWRITE_EA,
                                      //IARG_MEMORYWRITE_SIZE,
                                      IARG_UINT32, offset,
                                      IARG_UINT32, WR,
                                      IARG_END);
                valcount++;
            }
        }

        if (INS_SegmentPrefix(ins)) {
            REG seg = INS_SegmentRegPrefix(ins);
            /* Pin only has base registers for FS and GS (probably since
               Linux uses GS, and Windows uses FS. So, we'll just output a
               base register if we see one of those for now, and hope we
               don't need ES/etc. */
            if (seg == REG_SEG_FS || seg == REG_SEG_GS) {
                REG addreg;

                /* Set the register to add to the buffer */
                switch(seg) {
                case REG_SEG_FS:
                    addreg = REG_SEG_FS_BASE;
                    break;

                case REG_SEG_GS:
                    addreg = REG_SEG_GS_BASE;
                    break;

                default:
                    assert(false);
                    break;
                }

                IARGLIST_AddArguments(arglist_helper,
                                      IARG_ADDRINT, (ADDRINT)REGISTER,
                                      IARG_UINT32, GetBitsOfReg(addreg), // Register size in bits
                                      IARG_UINT32, addreg,
                                      //IARG_MEMORYWRITE_SIZE,
                                      IARG_PTR, 0,
                                      IARG_UINT32, 0,
                                      IARG_END);
                valcount++;
            }
        }



        // TODO: Check if valcount has exceed the maximum number of
        // values. Also, figure out what to do if so.

        if (valcount >= MAX_VALUES_COUNT) {
            cerr << "Error: Too many values (" << valcount << "). Max: " << MAX_VALUES_COUNT << endl;
            cerr << "Instruction: " << INS_Disassemble(ins) << endl;
            cerr << "Category: " << CATEGORY_StringShort(INS_Category(ins)) << endl;
        }
        assert(valcount < MAX_VALUES_COUNT);


        IARGLIST_AddArguments(arglist,
                              IARG_UINT32, valcount,
                              IARG_END);

        /* Now, add the operands. */
        IARGLIST_AddArguments(arglist,
                              IARG_IARGLIST, arglist_helper,
                              IARG_END);

        // The argument list has been built, time to insert the call.

        INS_InsertCall(ins, IPOINT_BEFORE,
                       (AFUNPTR) AppendBuffer,
                       IARG_IARGLIST, arglist,
                       IARG_END);

        // If we are logging one instruction after exiting the recording
        // range, then arrange for the post instruction call to happen
        // if ins is outside of the range.
        if (LogOneAfter.Value() && !(INS_Address(ins) >= start_addr && INS_Address(ins) <= end_addr)) {
            cerr << "found the last one" << endl;
            if (INS_IsBranchOrCall(ins)) {
                INS_InsertCall(ins, IPOINT_TAKEN_BRANCH,
                               (AFUNPTR) PostInstruction,
                               IARG_ADDRINT, INS_Address(ins),
                               IARG_CONTEXT,
                               IARG_END);
            } else {
                INS_InsertCall(ins, IPOINT_AFTER,
                               (AFUNPTR) PostInstruction,
                               IARG_ADDRINT, INS_Address(ins),
                               IARG_CONTEXT,
                               IARG_END);
            }
        }

        insLeft--;

        // Free the memory.
        IARGLIST_Free(arglist);
        IARGLIST_Free(arglist_helper);

    }

    //LOG("INS: bbl ins end.\nINS: BBL end.\n");


}

VOID InstrTrace(TRACE trace, VOID *v)
{
    dbg_printf("InstrTrace v=%p\n", v);

    /* Decide if we want to log this trace by examining the entrance address. */
    ADDRINT addr = TRACE_Address(trace);
    if (dontLog(addr)) {
    } else {
        for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl)) {
            InstrBlock(bbl);
        }
    }
}

VOID ThreadEnd(THREADID threadid, CONTEXT *ctx, INT32 code, VOID *v)
{
    dbg_printf("ThreadEnd tid=%d, code=%d, v=%p\n", threadid, code, v);

    ThreadInfo_t *ti = NULL;

    cerr << "Thread " << threadid << " ending" << endl;

    // Free thread-local data
    ti = GetThreadInfo();

    delete ti;
}

VOID ThreadStart(THREADID threadid, CONTEXT *ctx, INT32 flags, VOID *v)
{
    dbg_printf("ThreadStart tid=%d, flags=0x%x, v=%p\n", threadid, flags, v);

    // Get the command line arguments before _start is called
    // This only works with Linux conventions in mind
    static int firstthread = true;

    LLOG("new thread\n");

    NewThreadInfo();

    PIN_GetLock(&lock, threadid+1);

    LOG("New thread starting\n");
    cerr << "Thread " << threadid << " starting" << endl;

    if (firstthread) {
        firstthread = false;
#ifndef _WIN32 /* unix */
           int argc = *(int*)(PIN_GetContextReg(ctx, LEVEL_BASE::REG_STACK_PTR));
        char **argv = (char**) (PIN_GetContextReg(ctx, LEVEL_BASE::REG_STACK_PTR) + STACK_OFFSET);
        char **env = (char**) (PIN_GetContextReg(ctx, LEVEL_BASE::REG_STACK_PTR)+(argc+1) * STACK_OFFSET);
        std::vector<frame> frms = tracker->taintArgs(argc, argv);
        add_frames(g_twnew, frms.begin(), frms.end());
        frms = tracker->taintEnv(env);
        add_frames(g_twnew, frms.begin(), frms.end());
#else /* windows */
        /* On windows, we don't taint argc and argv, but rather taint the
           output of GetComamndLineA and GetCommandLineW.  On recent
           versions of Windows, these return a static pointer. */
        char *aptr = WINDOWS::GetCommandLineA();
        wchar_t *wptr = WINDOWS::GetCommandLineW();

        std::vector<frame> frms = tracker->taintArgs(aptr, wptr);
        add_frames(g_twnew, frms.begin(), frms.end());
#endif
    }

    PIN_ReleaseLock(&lock);

}

VOID ModLoad(IMG img, VOID *v)
{

    cerr << "This is modload()" << endl;

    const string &name = IMG_Name(img);

    frame f;
    f.mutable_modload_frame()->set_module_name(name);
    f.mutable_modload_frame()->set_low_address(IMG_LowAddress(img));
    f.mutable_modload_frame()->set_high_address(IMG_HighAddress(img));

    PIN_GetLock(&lock, 0);
    g_twnew->add(f);
    PIN_ReleaseLock(&lock);

#ifdef _WIN32
    // Try to find kernel32
    {
        char tempbuf[BUFSIZE];
        char *tok = NULL;
        char *lasttok = NULL;

        // Fill up the temporary buffer
        strncpy(tempbuf, name.c_str(), BUFSIZE);

        // We don't need a lock, since this is an instrumentation function (strtok is not re-entrant)
        strtok(tempbuf, "\\");

        while ((tok = strtok(NULL, "\\")) != NULL) {
            // Just keep parsing...
            lasttok = tok;
        }

        if (lasttok) {
            if (strncmp(windowsDll, lasttok, BUFSIZE) == 0) {

                /* GetEnvironmentStringsA uses a table-based conversion
                 * process that we can't analyze very well. So, these are
                 * disabled. */
#ifdef USE_GETENVSTRINGS
                RTN r;

                /** The prototype for GetEnvironmentStrings[WA] */
                PROTO proto = PROTO_Allocate( PIN_PARG(uint32_t), CALLINGSTD_STDCALL,
                                              "Windows API",
                                              PIN_PARG_END() );

                r = RTN_FindByName(img, "GetEnvironmentStringsW");
                if (r != RTN_Invalid()) {
                    RTN_ReplaceSignature(r, AFUNPTR(GetEnvWWrap),
                                         IARG_PROTOTYPE, proto,
                                         IARG_CONTEXT,
                                         IARG_ORIG_FUNCPTR,
                                         IARG_THREAD_ID,
                                         IARG_END);
                } else {
                    cerr << "Warning: Error instrumenting GetEnvironmentStringsW()" << endl;
                }

                r = RTN_FindByName(img, "GetEnvironmentStringsA");
                if (r != RTN_Invalid()) {
                    RTN_ReplaceSignature(r, AFUNPTR(GetEnvAWrap),
                                         IARG_PROTOTYPE, proto,
                                         IARG_CONTEXT,
                                         IARG_ORIG_FUNCPTR,
                                         IARG_THREAD_ID,
                                         IARG_END);
                } else {
                    cerr << "Warning: Error instrumenting GetEnvironmentStringsA()" << endl;
                }


                PROTO_Free(proto);
#endif

            } else if (strncmp(wsDll, lasttok, BUFSIZE) == 0) {
                /* Winsock */
                RTN r;

                cerr << "found winsock" << endl;

                r = RTN_FindByName(img, "accept");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(uint32_t), CALLINGSTD_STDCALL,
                                                 "accept",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(int*),
                                                 PIN_PARG_END());

                    RTN_ReplaceSignature(r, AFUNPTR(AcceptWrapper),
                                         IARG_PROTOTYPE, proto,
                                         IARG_CONTEXT,
                                         IARG_ORIG_FUNCPTR,
                                         IARG_THREAD_ID,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
                                         IARG_END);


                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find accept" << endl;
                }

                r = RTN_FindByName(img, "connect");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "connect",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG_END());

                    RTN_ReplaceSignature(r, AFUNPTR(ConnectWrapper),
                                         IARG_PROTOTYPE, proto,
                                         IARG_CONTEXT,
                                         IARG_ORIG_FUNCPTR,
                                         IARG_THREAD_ID,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
                                         IARG_END);


                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find connect" << endl;
                }

                r = RTN_FindByName(img, "WSAConnect");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "WSAConnect",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG_END());

                    RTN_ReplaceSignature(r, AFUNPTR(WSAConnectWrapper),
                                         IARG_PROTOTYPE, proto,
                                         IARG_CONTEXT,
                                         IARG_ORIG_FUNCPTR,
                                         IARG_THREAD_ID,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 4,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 5,
                                         IARG_FUNCARG_ENTRYPOINT_VALUE, 6,
                                         IARG_END);


                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find WSAConnect" << endl;
                }


                r = RTN_FindByName(img, "recv");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "recv",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(char*),
                                                 PIN_PARG(int),
                                                 PIN_PARG(int),
                                                 PIN_PARG_END());
                    RTN_Open(r);

                    RTN_InsertCall(r, IPOINT_BEFORE, (AFUNPTR)BeforeRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                   IARG_END);

                    RTN_InsertCall(r, IPOINT_AFTER, (AFUNPTR)AfterRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCRET_EXITPOINT_VALUE,
                                   IARG_PTR, "recv",
                                   IARG_END);

                    RTN_Close(r);
                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find recv" << endl;
                }


                r = RTN_FindByName(img, "recvfrom");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "recvfrom",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(char*),
                                                 PIN_PARG(int),
                                                 PIN_PARG(int),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(int*),
                                                 PIN_PARG_END());
#if 0
                    RTN_Open(r);

                    RTN_InsertCall(r, IPOINT_BEFORE, (AFUNPTR)BeforeRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                   IARG_END);

                    RTN_InsertCall(r, IPOINT_AFTER, (AFUNPTR)AfterRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCRET_EXITPOINT_VALUE,
                                   IARG_PTR, "recvfrom",
                                   IARG_END);

                    RTN_Close(r);
#endif
                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find recvfrom" << endl;
                }

                r = RTN_FindByName(img, "WSARecv");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "WSARecv",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(uint32_t*),
                                                 PIN_PARG(uint32_t*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG_END());

                    RTN_Open(r);

                    RTN_InsertCall(r, IPOINT_BEFORE, (AFUNPTR)WSABeforeRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
                                   IARG_END);

                    RTN_InsertCall(r, IPOINT_AFTER, (AFUNPTR)AfterRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCRET_EXITPOINT_VALUE,
                                   IARG_PTR, "WSARecv",
                                   IARG_END);

                    RTN_Close(r);
                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find WSArecv" << endl;
                }

                r = RTN_FindByName(img, "WSARecvFrom");
                if (r != RTN_Invalid()) {

                    PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_STDCALL,
                                                 "WSARecvFrom",
                                                 PIN_PARG(uint32_t),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG(void*),
                                                 PIN_PARG_END());

#if 0
                    RTN_Open(r);

                    RTN_InsertCall(r, IPOINT_BEFORE, (AFUNPTR)WSABeforeRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
                                   IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
                                   IARG_END);

                    RTN_InsertCall(r, IPOINT_AFTER, (AFUNPTR)AfterRecv,
                                   IARG_PROTOTYPE, proto,
                                   IARG_THREAD_ID,
                                   IARG_FUNCRET_EXITPOINT_VALUE,
                                   IARG_PTR, "WSARecvFrom",
                                   IARG_END);

                    RTN_Close(r);
#endif

                    PROTO_Free(proto);

                } else {
                    cerr << "Couldn't find WSArecv" << endl;
                }

            } else {
                cerr << "Other img: " << lasttok << endl;
            }
        }
    }
#endif

    if (g_usetrigger && !g_trig_resolved) {
        // Check if this module can be used to resolve the trigger address.

        // If no trigger module is set, then we just use the first one we
        // find, i.e. the main module.
        if ((KnobTrigModule.Value() == "") ||
            name.find(KnobTrigModule.Value()) != string::npos) {
            // Found the module, resolve address.

            g_trig_addr = KnobTrigAddr.Value() + IMG_LoadOffset(img);
            g_trig_resolved = true;

        }

    }

}

VOID SyscallEntry(THREADID tid, CONTEXT *ctx, SYSCALL_STANDARD std, VOID *v)
{
    dbg_printf("SyscallEntry tid=%d, v=%p\n", tid, v);

    ThreadInfo_t *ti = NULL;
    SyscallInfo_t si;

    /*
     * Synchronization note: We assume there is only one system call per
     * thread, and thus the thread local syscall stack does not need any
     * locking.
     */

    // cerr << "syscall in " << PIN_GetSyscallNumber(ctx, std) << endl;

    ti = GetThreadInfo();
    //  cerr << "stack size " << ti->scStack.size() << endl;

    LLOG("syscall\n");

    // Ignore if not activated.
    //  if (!g_active) return;

    // Get the address from instruction pointer (should be EIP).
    si.sf.mutable_syscall_frame()->set_address(PIN_GetContextReg(ctx, REG_INST_PTR));

    si.sf.mutable_syscall_frame()->set_thread_id(tid);

    si.sf.mutable_syscall_frame()->set_number(PIN_GetSyscallNumber(ctx, std));

    for (int i = 0; i < MAX_SYSCALL_ARGS; i++)
        {
            if (i < PLAT_SYSCALL_ARGS) {
                dbg_printf("Syscall argument %d: 0x%lx\n", i, PIN_GetSyscallArgument(ctx, std, i));
                si.sf.mutable_syscall_frame()->mutable_argument_list()->add_elem(PIN_GetSyscallArgument(ctx, std, i));
            }
        }

    // XXX: This should really be above g_active probably, but it seems
    // unlikely that it would cause a problem.  It's here because
    // FlushBuffer obtains a lock of it's own.
    PIN_GetLock(&lock, tid+1);

    // First we need to flush the buffer, so we can directly add the
    // syscall frame after the frame for the instruction that led to the
    // syscall.
    FlushBuffer(true, ctx, tid, false);

    if (LogAllSyscalls.Value()) {
        g_twnew->add(si.sf);
    }

    if (tracker->taintPreSC(si.sf.mutable_syscall_frame()->number(), (const uint64_t *) (si.sf.syscall_frame().argument_list().elem().data()), si.state)) {
        // Do we need to do anything here? ...
    }
    ti->scStack.push(si);

    //e:

    LLOG("releasing sysenter\n");
    PIN_ReleaseLock(&lock);
    LLOG("really done with sysenter\n");
}

VOID SyscallExit(THREADID tid, CONTEXT *ctx, SYSCALL_STANDARD std, VOID *v)
{
    dbg_printf("SyscallExit tid=%d v=%p\n", tid, v);

    ThreadInfo_t *ti = NULL;
    SyscallInfo_t si;
    ADDRINT addr;
    uint32_t length;

    /*
     * Synchronization note: We assume there is only one system call per
     * thread, and thus the thread local syscall stack does not need any
     * locking.
     */

    // Ignore if not activated.
    // if (!g_active) return;

    LLOG("sysexit\n");

    ti = GetThreadInfo();

    si = ti->scStack.top();
    ti->scStack.pop();

    PIN_GetLock(&lock, tid+1);

    // Check to see if we need to introduce tainted bytes as a result of this
    // sytem call
    FrameOption_t fo = tracker->taintPostSC(PIN_GetSyscallReturn(ctx, std), (const uint64_t*) (si.sf.syscall_frame().argument_list().elem().data()), addr, length, si.state);

    if (fo.b) {
        if (!g_taint_introduced) {
            // Activate taint tracking
            TActivate();
        }
        g_twnew->add(fo.f);
    }

    //printf("syscall out %d\n", si.sf.callno);

    LLOG("releasing sysexit\n");
    PIN_ReleaseLock(&lock);

    // Untaint system call output registers (uses thread-local delta, so no lock needed)
    tracker->postSysCall(ti->delta);

    LLOG("really done with sysexit\n");
}


VOID FollowParent(THREADID threadid, const CONTEXT* ctxt, VOID * arg)
{
    dbg_printf("FollowParent tid=%d arg=%p\n", threadid, arg);

    int i;

    LLOG("fparent\n");

    PIN_GetLock(&lock, threadid+1);
    i = strlen(g_threadname);
    assert(i < BUFFER_SIZE);
    g_threadname[i++] = 'p';

    std::cerr << "Spawning parent: " << PIN_GetPid() << g_threadname << std::endl;

    PIN_ReleaseLock(&lock);
}

VOID ExceptionHandler(THREADID threadid, CONTEXT_CHANGE_REASON reason, const CONTEXT *from, CONTEXT *to, INT32 info, VOID *v) {

    dbg_printf("ExceptionHandler tid=%d info=0x%x v=%p\n", threadid, info, v);
    /*
      CONTEXT_CHANGE_REASON_FATALSIGNAL          Receipt of fatal Unix signal.
      CONTEXT_CHANGE_REASON_SIGNAL       Receipt of handled Unix signal.
      CONTEXT_CHANGE_REASON_SIGRETURN    Return from Unix signal handler.
      CONTEXT_CHANGE_REASON_APC          Receipt of Windows APC.
      CONTEXT_CHANGE_REASON_EXCEPTION    Receipt of Windows exception.
      CONTEXT_CHANGE_REASON_CALLBACK     Receipt of Windows call-back.
    */

    /*
      If there is a fatal exception, we should halt the trace as soon
      as possible, so we can exit.

      Also, FlushInstructions() needs mutual exclusivity.
    */

    // SWHITMANXXX Get information and put it into exception frame here
    // XXX Put this into frame buffer

    frame f;
    f.mutable_exception_frame()->set_exception_number(info);
    f.mutable_exception_frame()->set_thread_id(threadid);
    if (from) {
        f.mutable_exception_frame()->set_from_addr(PIN_GetContextReg(from, REG_INST_PTR));
    }
    if (to) {
        f.mutable_exception_frame()->set_to_addr(PIN_GetContextReg(to, REG_INST_PTR));
    }

    PIN_GetLock(&lock, threadid+1);
    LLOG("got except lock!\n");

    // If we want the exception to be the last thing in the trace when
    // we crash, then we need to flush.
    FlushBuffer(false, from, threadid, false);
    g_twnew->add(f);

    if (reason == CONTEXT_CHANGE_REASON_FATALSIGNAL) {
        std::cerr << "Received fatal signal " << info << endl;
        FlushBuffer(false, from, threadid, false);
        Cleanup();
        exit(1);
    } else if (reason == CONTEXT_CHANGE_REASON_EXCEPTION) {

#ifdef _WIN32
        ADDRINT pc = PIN_GetContextReg(from, REG_INST_PTR);
        cerr << "Received windows exception @" << pc << " " << info << " in thread " << threadid << endl;

        if (info == accessViolation && SEHMode.Value() && g_taint_introduced) {
            cerr << "SEH mode activated!" << endl;
            ADDRINT old_esp = PIN_GetContextReg(from, REG_STACK_PTR);
            ADDRINT new_esp = PIN_GetContextReg(to, REG_STACK_PTR);

            cerr << "old esp: " << old_esp << " new esp: " << new_esp
                 << " old eip: " << PIN_GetContextReg(from, REG_INST_PTR) << " new eip: " << PIN_GetContextReg(to, REG_INST_PTR) << endl;

            /* The windows exception handler just pushed a bunch of crap
               onto the stack.  Although some of this is user controllable,
               we can untaint it for now, since we are mainly concerned with
               accessing our buffer. */
            assert (new_esp < old_esp);
            for (ADDRINT ptr = new_esp; ptr < old_esp; ptr++) {
                tracker->untaintMem(ptr);
            }

            /* Try to find a tainted exception handler. */
            ADDRINT eptr = PIN_GetContextReg(to, REG_SEG_FS_BASE)
                + ehandler_fs_offset;

            /* eptr points to the &(head of SEH). */
            assert(PIN_SafeCopy(&eptr, (void*)eptr, sizeof(ADDRINT)) == sizeof(ADDRINT));

            /* eptr points to head of SEH. */

            for (int i = 0; i < maxSehLength; i++) {
                // while (true) {
                struct {
                    ADDRINT nptr;
                    ADDRINT handler;
                } buf;
                /* We supposedly have a pointer to an exception handler
                   structure.  Let's make sure it's mapped. */
                size_t b = PIN_SafeCopy((void*)&buf, (void*)eptr, ehandler_size);
                if (b == ehandler_size) {

                    /* Okay, we have an exception handler.  Let's see if the
                       pointer handler is tainted.  So, check if M[eptr+4] is
                       tainted. */
                    ADDRINT hptr = eptr + ehandler_handler_offset;

                    /* hptr holds the address of the handler. */
                    cerr << "SEH handler M[" << hptr
                         << "] = " << buf.handler
                         << " (" << tracker->getMemTaint(hptr, pintrace::INVALIDREGMEM)
                         << ")"
                         << endl;

                    eptr = buf.nptr;

                } else {
                    cerr << "Unable to read from " << eptr << endl;
                    break;
                }

            }

            ADDRINT esp = PIN_GetContextReg(to, REG_STACK_PTR);

            /* The exception handling stuff will push a lot of data to the
               stack, so take account for that here. */
            PIN_SetContextReg(to, REG_STACK_PTR, esp-ehandler_esp_offset);

            PIVOT_testpivot(ps, to, *tracker);

            FlushBuffer(false, from, threadid, false);
            Cleanup();
            exit(1);

        } else {
            cerr << "Ignoring exception!" << endl;
        }
#endif
    } else if (reason == CONTEXT_CHANGE_REASON_CALLBACK) {
#if 0
        cerr << "Received windows callback" << endl;
#endif

    } else {
        std::cerr << "Received other exception " << reason << endl;
    }

    LLOG("done handling exception\n");
    PIN_ReleaseLock(&lock);
}

VOID FollowChild(THREADID threadid, const CONTEXT* ctxt, VOID * arg)
{
    dbg_printf("FollowChild tid=%d arg=%p\n", threadid, arg);

    int i;

    LLOG("follow child\n");

    PIN_GetLock(&lock, threadid+1);
    i = strlen(g_threadname);
    assert(i < BUFFER_SIZE);
    g_threadname[i++] = 'c';

    g_twnew = new TraceContainerWriter((g_threadname + KnobOut.Value()).c_str(), BFD_ARCH, BFD_MACH, default_frames_per_toc_entry);

    g_bufidx = 0;
    g_kfcount = 0;

    g_logcount = 0;
    g_loglimit = KnobLogLimit.Value();

    g_timer = clock();
    std::cerr << "Spawning child: " << PIN_GetPid() << g_threadname << std::endl;
    PIN_ReleaseLock(&lock);

}

bool FollowExec(CHILD_PROCESS cp, VOID *v) {
    bool follow = false;
    int argc;
    const char * const * argv;

    CHILD_PROCESS_GetCommandLine(cp, &argc, &argv);
    assert (argc >= 0);
    cerr << "Exec: ";
    for (int i = 0; i < argc; i++) {
        cerr << argv[i] << " ";
    }
    cerr << endl;

    /* See if we should follow this */
    for (unsigned int i = 0; i < FollowProgs.NumberOfValues(); i++) {
        if (FollowProgs.Value(i) == argv[0]) {
            follow = true;
        }
    }


    if (follow)
        cerr << "Following" << endl;
    else
        cerr << "Not following" << endl;

#ifndef _WIN32
    /* If we're on Linux, this means we're about to call execv(), and
     * we're going to disappear! We had better write out our trace! */

    cerr << "Flushing buffer before execv()" << endl;
    FlushBuffer(false, NULL, PIN_ThreadId(), true);
    Cleanup();
#endif

    return follow;
}

VOID Fini(INT32 code, VOID *v)
{
    LOG("In Fini");
    Cleanup();
}

// Caller responsible for mutual exclusion
VOID Cleanup()
{
    g_twnew->finish();

    LOG("done.\n");

    clock_t endtime = clock();

    LOG("Time taken: " + decstr((UINT64) (endtime - g_timer)));

    exit(0);

}

INT32 Usage()
{
    cerr << endl << KNOB_BASE::StringKnobSummary() << endl;
    return -1;
}

int main(int argc, char *argv[])
{
    dbg_printf("Entered main.\n");
    stringstream ss;

    cerr << hex;

    // A sanity check for AppendBuffer
    //assert(sizeof(RPassType) == sizeof(ADDRINT));

    PIN_InitSymbols();

    if (PIN_Init(argc,argv))
        return Usage();

    PIN_InitLock(&lock);

    // Check if a trigger was specified.
    if (KnobTrigAddr.Value() != 0) {
        g_usetrigger = true;
        g_trig_resolved = false;

        // Set trigger countdown to initial value.
        g_trig_countdown = KnobTrigCount.Value();

    } else {
        g_usetrigger = false;
    }
    // Check if taint tracking is on
    if (KnobTaintTracking.Value()) {
        tracker = new TaintTracker(values);
        for (uint32_t i = 0 ; i < TaintedFiles.NumberOfValues() ; i++) {
            if (TaintedFiles.Value(i) != "") {
                tracker->trackFile(TaintedFiles.Value(i));
            }
        }

        tracker->setTaintArgs(TaintedArgs);
        if (TaintedStdin)
            tracker->setTaintStdin();
        if (TaintedNetwork)
            tracker->setTaintNetwork();
    
        for (uint32_t i = 0 ; i < TaintedEnv.NumberOfValues() ; i++) {
            if (TaintedEnv.Value(i) != "") {
                tracker->setTaintEnv(TaintedEnv.Value(i));
            }
        }
    }

    /* Get a key for thread info */
    tl_key = PIN_CreateThreadDataKey(NULL);
    assert(tl_key != -1);

    // We must activate taint tracking early if we have tainted args
    // or envs
    if ((TaintedEnv.Value() != "") || TaintedArgs.Value()) {
        g_taint_introduced = true;
    } else {
        g_taint_introduced = false;
    }

    // Determine whether logging is enabled.
    // If a trigger is specified, logging is never enabled, because
    // nothing is logged until the trigger point.  Otherwise, logging
    // is enabled if taint is introduced, or if logging before taint is enabled.
    if (g_usetrigger) {
        g_active = false;
    } else {
        if (g_taint_introduced || LogAllBeforeTaint.Value()) {
            g_active = true;
        } else {
            g_active = false;
        }
    }

    cerr << "Logging initially enabled: " << g_active << endl;

    /** Read pivot gadgets */
    if (PivotFile.Value() != "") {
        fstream f;
        pivot_set::iterator i;

        f.open(PivotFile.Value().c_str());
        if (!f.is_open()) {
            cerr << "Could not open pivot gadget file: " << PivotFile.Value() << endl;
            exit(1);
        }

        ps = PIVOT_parseinput(f);
        cerr << "Read " << ps.size() << " pivots" << endl;

        f.close();
    }

    IMG_AddInstrumentFunction(ModLoad, 0);
    TRACE_AddInstrumentFunction(InstrTrace, 0);
    PIN_AddThreadStartFunction(ThreadStart, 0);
    PIN_AddThreadFiniFunction((THREAD_FINI_CALLBACK)ThreadEnd, 0);

    PIN_AddContextChangeFunction(ExceptionHandler, 0);

#ifndef _WIN32
    PIN_AddForkFunction(FPOINT_AFTER_IN_CHILD, FollowChild, 0);
    PIN_AddForkFunction(FPOINT_AFTER_IN_PARENT, FollowParent, 0);
#endif
    PIN_AddFollowChildProcessFunction(FollowExec, 0);

    PIN_AddSyscallEntryFunction(SyscallEntry, 0);
    PIN_AddSyscallExitFunction(SyscallExit, 0);

    PIN_AddFiniFunction(Fini, 0);

    ss << PIN_GetPid() << "-" << KnobOut.Value();

    g_twnew = new TraceContainerWriter(ss.str().c_str(), BFD_ARCH, BFD_MACH, default_frames_per_toc_entry);

    g_bufidx = 0;
    g_kfcount = 0;

    g_logcount = 0;
    g_loglimit = KnobLogLimit.Value();

    g_skipTaints = SkipTaints.Value();

    g_timer = clock();

    g_exit_next = false;

    start_addr = TaintStart.Value();
    end_addr = TaintEnd.Value();

    cerr << "Code cache limit is " << CODECACHE_CacheSizeLimit() << endl;
    assert(CODECACHE_ChangeCacheLimit(CacheLimit.Value()));

    LOG("Starting program\n");
    cerr << "Starting program" << endl;

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

    return 0;

}