#include "pin.H"
#include <sys/syscall.h>
#include <sys/mman.h>
#include <unistd.h>
#include <iostream>
#include <stack>


static void handle_indirect_call(ADDRINT inst_ptr, ADDRINT target_addr)
{
    std::cerr << "icall; inst_addr=" << std::hex << inst_ptr << " target_addr=" << target_addr << std::dec << std::endl;
}

/* shadow stack containing pairs of call instruction and expected return address. */
std::stack<std::pair<ADDRINT, ADDRINT>> shadow_stack{};

/* Register the expected return address for this call. */
static void handle_call(ADDRINT inst_ptr, ADDRINT expected_addr)
{
    shadow_stack.push({inst_ptr, expected_addr});
}

static void handle_ret(ADDRINT ret_addr)
{
    /* This assumes all rets are paired with a call. */
    auto const& [call_addr, expected_ret_addr] = shadow_stack.top();

    if (expected_ret_addr != ret_addr)
    {
        std::cerr << "unexpected_ret; call_addr=" << std::hex << call_addr << " expected_addr=" << expected_ret_addr << " actual_addr=" << ret_addr << std::dec << std::endl;
    }

    shadow_stack.pop();
}

static bool IsInPLT(ADDRINT addr)
{
    return (addr > 0x401020 && addr < 0x4012cf);
}

static ADDRINT handle_call_if_nonplt(ADDRINT target_addr)
{
    return !IsInPLT(target_addr);
}

/* PIN calls this while translating an instruction */
void Instrument(INS ins, void *)
{
#ifdef PRINT_ALL_INSTS
	printf("instrumenting@%p: %s\n", (void *)INS_Address(ins), INS_Disassemble(ins).c_str());
#endif

	if (INS_IsCall(ins))
    {
	    if (INS_IsDirectCall(ins))
        {
	        if (!IsInPLT(INS_DirectControlFlowTargetAddress(ins)))
            {
                INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) handle_call, IARG_INST_PTR, IARG_ADDRINT, INS_NextAddress(ins), IARG_END);
            }
        }
	    else
        {
	        INS_InsertIfCall(ins, IPOINT_BEFORE, AFUNPTR(handle_call_if_nonplt), IARG_BRANCH_TARGET_ADDR, IARG_END);
	        INS_InsertThenCall(ins, IPOINT_BEFORE, (AFUNPTR) handle_call, IARG_INST_PTR, IARG_ADDRINT, INS_NextAddress(ins), IARG_END);
        }


        if (!INS_IsDirectCall(ins))
        {
            // std::cerr << "Instrumenting call instruction: " << INS_Disassemble(ins) << std::endl;
            INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) handle_indirect_call, IARG_INST_PTR, IARG_BRANCH_TARGET_ADDR, IARG_END);
        }
    }
    else if (INS_IsBranch(ins))
    {
        // nothing for now
    }
    else if (INS_IsRet(ins))
    {
        INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR) handle_ret, IARG_BRANCH_TARGET_ADDR, IARG_END);
    }

}


void Trace(TRACE trace, void *) {
	/* Iterate through the basic blocks which PIN wants to instrument right now. */
	for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl)) {
		/*
		 * The system libraries use many more instructions than the BAMA binary does.
		 * We suggest that you exclude libraries from instrumentation (with the code below).
		 * and implement tainting by instrumenting the individual library calls instead.
		 */
		IMG img = IMG_FindByAddress(BBL_Address(bbl));
		if (!IMG_Valid(img) || !IMG_IsMainExecutable(img))
			continue;

#ifdef PRINT_BASIC_BLOCKS
		printf("Instrumenting basic block at %p\n", BBL_Address(bbl));
#endif

		/* Instrument every instruction in this basic block. */
		for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins)) {
			Instrument(ins, nullptr);
		}
	}
}

int main (int argc, char *argv[]) {
	/* This makes PIN handle IFuncs; RTN_FindByName gives us implementations in modern PIN, so we can pretend they don't exist. */
	PIN_InitSymbolsAlt(IFUNC_SYMBOLS);
	if (PIN_Init(argc, argv))
		return 1;

	/* We want to instrument instructions. */
	TRACE_AddInstrumentFunction(Trace, nullptr);

	/* Let's go! */
	PIN_StartProgram();
}
