#include "pin.H"
#include <iostream>
#include <iomanip>
#include <algorithm>
#include <list>
#include <string.h>
#include <stdio.h>
#include <ostream>

list<ADDRINT> MallocAddrs;

VOID FreeBefore(ADDRINT target, ADDRINT inst)
{
   list<ADDRINT>::iterator p;
   p = find(MallocAddrs.begin(), MallocAddrs.end(), target);
   if ( (!(MallocAddrs.empty())) && (MallocAddrs.end() != p) ){
       p = MallocAddrs.erase(p); // Delete this from the allocated address list
   }else{   // We caught a Free of an un-allocated address
       cerr << "DOUBLE-FREE DETECTED: " << hex << target << " @" << inst << endl;
   }   // Using cerr is not a good practice, I do it oly for the sake of the example
}

VOID MallocAfter(ADDRINT ret, ADDRINT inst)
{
   // Save the address returned by malloc in our list
   if (ret != 0){
      list<ADDRINT>::iterator p;
      p = find(MallocAddrs.begin(), MallocAddrs.end(), ret);

      if (MallocAddrs.end() == p){ //not found
         MallocAddrs.push_back(ret);
         cerr << "Saving " << hex << ret << " in the address list @" << inst << endl;
      }else{
         // malloc address already in the list?!
         cerr << "already saved" << hex << " @" << inst << endl;
      }
   }else{
     cerr << "Malloc fail" << endl;
   }
}

// Instrument the malloc() and free() functions.
// note that there are malloc and free in the os loader and in libc
VOID Image(IMG img, VOID *v)
{
    cerr << "Hooking img: " << IMG_Name(img) << endl;

    // Find the malloc() function and add our function after it
    RTN mallocRtn = RTN_FindByName(img, "malloc");
    if (RTN_Valid(mallocRtn)){
		// print function name
        cerr << "Function name: " << RTN_Name(mallocRtn) << endl;
        RTN_Open(mallocRtn);
        RTN_InsertCall(mallocRtn, IPOINT_AFTER, (AFUNPTR)MallocAfter,
                           IARG_FUNCRET_EXITPOINT_VALUE, IARG_INST_PTR, 
                           IARG_END);
        // IARG_FUNCRET_EXITPOINT_VALUE function result, valid only at return instruction
        RTN_Close(mallocRtn);
    }

    // Find the free() function and add our function before it
    RTN freeRtn = RTN_FindByName(img, "free");
    if (RTN_Valid(freeRtn)) {
        // print function name
        cerr << "Function name: " << RTN_Name(freeRtn) << endl;
        RTN_Open(freeRtn);
        RTN_InsertCall(freeRtn, IPOINT_BEFORE, (AFUNPTR)FreeBefore,
                           IARG_FUNCARG_ENTRYPOINT_VALUE, 0, IARG_INST_PTR, 
                           IARG_END);
        // IARG_FUNCARG_ENTRYPOINT_VALUE int argument, valid only at the entry point of a routine
        RTN_Close(freeRtn);
    }
}

int main(int argc, char *argv[]){
    // Initialize pin & symbol manager
    PIN_InitSymbols();
    PIN_Init(argc,argv);

    // Register Image to be called to instrument functions.
    IMG_AddInstrumentFunction(Image, 0);

    PIN_StartProgram(); // Never returns
    
    return 0;
}
