memtrace.cpp

#include <iostream>
#include <fstream>
#include <string>
#include "pin.H"

using std::endl;
using std::hex;
using std::dec;
using std::ofstream;
using std::string;

static UINT64 icount = 0; // total instructions executed
static UINT64 readCount  = 0;// memory reads logged
static UINT64 writeCount = 0;// memory writes logged

ADDRINT mainLow  = 0;
ADDRINT mainHigh = 0;

static UINT64 maxAccesses = 0;

ofstream InsFile;// instruction count output
ofstream MemFile;   // memory trace output

// KNOBS — these become command-line flags for the pintool
KNOB<string> KnobInsFile(
    KNOB_MODE_WRITEONCE, "pintool",
    "i", "icount.out",
    "Instruction count output file");

KNOB<string> KnobMemFile(
    KNOB_MODE_WRITEONCE, "pintool",
    "m", "memtrace.out",
    "Memory trace output file");

// NEW knob: put a ceiling on how many accesses we record
KNOB<UINT64> KnobMaxAccesses(
    KNOB_MODE_WRITEONCE, "pintool",
    "max", "0",
    "Max memory accesses to log (default 0 = no limit)");

// Analysis routines — these run at PROGRAM runtime for every access

VOID CountInstruction()
{
    icount++;
}

// Called before every memory READ instruction
VOID RecordMemRead(ADDRINT addr)
{
    // Only log addresses inside the main binary
    // if (addr < mainLow || addr > mainHigh)
    //     return;

    // Stop once we hit the cap
    if (maxAccesses > 0 && (readCount + writeCount) >= maxAccesses)
        return;

    MemFile << "R " << hex << addr << "\n";
    readCount++;
}

// Called before every memory WRITE instruction
VOID RecordMemWrite(ADDRINT addr)
{
    // if (addr < mainLow || addr > mainHigh)
    //     return;

    if (maxAccesses > 0 && (readCount + writeCount) >= maxAccesses)
        return;

    MemFile << "W " << hex << addr << "\n";
    writeCount++;
}

// Instrumentation callback
// PIN calls Instruction() once per instruction at JIT time.
// Here we DECIDE which analysis calls to inject.
// The actual RecordMem* calls happen at program runtime.
VOID Instruction(INS ins, VOID *v)
{
    // Guard: stop injecting once we've already hit the cap
    if (maxAccesses > 0 && (readCount + writeCount) >= maxAccesses)
        return;

    // Always count instructions
    INS_InsertCall(ins, IPOINT_BEFORE,
                   (AFUNPTR)CountInstruction,
                   IARG_END);

    //Memory READ
    // e.g.  MOV rax, [rbx]   → reads one memory operand
    if (INS_IsMemoryRead(ins))
    {
        INS_InsertCall(ins, IPOINT_BEFORE,
                       (AFUNPTR)RecordMemRead,
                       IARG_MEMORYREAD_EA,   // effective address of the read
                       IARG_END);
    }

    // Some instructions read from TWO memory locations (e.g. CMPS, MOVS).
    // INS_HasMemoryRead2 catches the second operand.
    if (INS_HasMemoryRead2(ins))
    {
        INS_InsertCall(ins, IPOINT_BEFORE,
                       (AFUNPTR)RecordMemRead,
                       IARG_MEMORYREAD2_EA,
                       IARG_END);
    }

    //Memory WRITE
    // e.g.  MOV [rbx], rax   → writes one memory operand
    if (INS_IsMemoryWrite(ins))
    {
        INS_InsertCall(ins, IPOINT_BEFORE,
                       (AFUNPTR)RecordMemWrite,
                       IARG_MEMORYWRITE_EA,  // effective address of the write
                       IARG_END);
    }
}

// Image Load callback
// Fires when a shared library or the main binary is loaded.
// We grab the address range of the main executable here.
VOID ImageLoad(IMG img, VOID *v)
{
    if (IMG_IsMainExecutable(img))
    {
        mainLow  = IMG_LowAddress(img);
        mainHigh = IMG_HighAddress(img);
        std::cerr << "[memtrace] Main image range: "
                  << hex << mainLow << " – " << mainHigh << "\n";
    }
}

// Fini — runs once when the instrumented program exits
VOID Fini(INT32 code, VOID *v)
{
    UINT64 totalMem = readCount + writeCount;

    // Write instruction count to its file
    InsFile << "InstructionCount " << dec << icount << "\n";
    InsFile.close();

    // Append summary to the bottom of the memory trace
    MemFile << "# ---\n";
    MemFile << "# TotalMemoryAccesses " << dec << totalMem   << "\n";
    MemFile << "# Reads               " << dec << readCount  << "\n";
    MemFile << "# Writes              " << dec << writeCount << "\n";
    MemFile.close();

    // Also print to stderr so you see it in the terminal
    std::cerr << "[memtrace] Done.\n";
    std::cerr << "  Instructions : " << dec << icount    << "\n";
    std::cerr << "  Mem accesses : " << dec << totalMem  << "\n";
    std::cerr << "    Reads      : " << dec << readCount  << "\n";
    std::cerr << "    Writes     : " << dec << writeCount << "\n";
}

// Usage — printed when wrong flags are passed
INT32 Usage()
{
    std::cerr << "CS204 Pintool: Instruction + Memory Tracer\n";
    std::cerr << KNOB_BASE::StringKnobSummary() << "\n";
    return -1;
}

// main
int main(int argc, char *argv[])
{
    // Initialize PIN (parses argc/argv for knobs)
    if (PIN_Init(argc, argv))
        return Usage();

    maxAccesses = KnobMaxAccesses.Value();

    // Open output files
    InsFile.open(KnobInsFile.Value().c_str());
    MemFile.open(KnobMemFile.Value().c_str());

    if (!InsFile.is_open() || !MemFile.is_open())
    {
        std::cerr << "ERROR: Cannot open output files.\n";
        return 1;
    }

    // Header lines in the trace file (lines starting with # are comments)
    MemFile << "# CS204 Memory Trace\n";
    MemFile << "# Format: [R|W] <hex_address>\n";
    MemFile << "# Only main-executable accesses are logged.\n";

    // Register our three callbacks
    IMG_AddInstrumentFunction(ImageLoad,   0);  // fires on image load
    INS_AddInstrumentFunction(Instruction, 0);  // fires per instruction (JIT)
    PIN_AddFiniFunction(Fini, 0);               // fires on program exit

    // Hand control to PIN — this never returns
    PIN_StartProgram();
    return 0;
}