pintool.cpp

/* INCLUDES */
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
#include "pin.H"
#pragma GCC diagnostic pop

#include <pin.H>
#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <algorithm>
#include <cctype>
#include <fstream>
#include <unordered_map>
#include <sstream>

// HACK for differing namespaces on Windows
#if _WIN32
#define unordered_set tr1::unordered_set
#define unordered_map tr1::unordered_map
#endif

#include "bashformats.h"
#include "outfile.h"
#include "syscalls.h"
#include "images.h"


/* GLOBAL VARIABLES */

// Command line options.
KNOB<std::string> KnobOutputFile(KNOB_MODE_WRITEONCE, "pintool", "o", "out", "Output file path");
KNOB<std::string> KnobTrackRoutines(KNOB_MODE_WRITEONCE, "pintool", "f", "", "Functions to track, separated by ';'");
KNOB<int> KnobSkip(KNOB_MODE_WRITEONCE, "pintool", "skip", "0", "Number of instructions to skip");
KNOB<bool> KnobPrintInstructions(KNOB_MODE_WRITEONCE, "pintool", "print_ins", "1", "Print executed instructions (0 / 1)");
KNOB<bool> KnobCountOpcodes(KNOB_MODE_WRITEONCE, "pintool", "count_opcodes", "0", "Count executed opcodes (0 / 1)");
KNOB<bool> KnobPrintMemoryWrites(KNOB_MODE_WRITEONCE, "pintool", "print_memwrites", "1", "Print memory writes (0 / 1)");
KNOB<bool> KnobPrintRegisterWrites(KNOB_MODE_WRITEONCE, "pintool", "print_regwrites", "1", "Print register writes (0 / 1)");
KNOB<bool> KnobPrintSyscalls(KNOB_MODE_WRITEONCE, "pintool", "print_syscalls", "0", "Print system calls (0 / 1)");
KNOB<bool> KnobPrintArguments(KNOB_MODE_WRITEONCE, "pintool", "print_args", "0", "Print function argument registers (0 / 1)");
KNOB<int> KnobIndentMode(KNOB_MODE_WRITEONCE, "pintool", "indent", "1", "Indent instructions by call level (0: no / 1: spaces / 2: tabs)");

// List of named functions.
std::vector<std::string> _functionNames;

// Mapping of instruction address to function names.
std::unordered_map<ADDRINT, std::string> _addressToFunctionNameMap;

// Current output indentation (call stack depth).
int _indent = 0;

// The current indent padding.
std::string _indentPadding = "";

// Instruction skip counter.
int _skip = 0;

// Stores the effective memory operand address of the current instruction.
UINT64 _currentMemoryOperandAddress;

// Buffer for temporarily storing accessed memory (must be copied by Pin, direct pointer dereferencing is not recommended).
UINT8* _memoryAccessBuffer = nullptr;

// Current size of the memory access buffer.
int _memoryAccessBufferSize = 0;

// Buffer for reading register values.
UINT8 _registerValueBuffer[64];

// List of tracked functions.
std::vector<std::string> _trackedRoutines;

// Mapping of instruction opcodes to counters.
std::unordered_map<OPCODE, int64_t> _instructionCounters;


/* CALLBACK PROTOTYPES */

VOID InstrumentTrace(TRACE trace, VOID* v);
VOID InstrumentImage(IMG img, VOID* v);
VOID InstrumentRoutine(RTN rtn, VOID* v);
VOID Fini(INT32 code, VOID *v);
EXCEPT_HANDLING_RESULT HandlePinToolException(THREADID tid, EXCEPTION_INFO* exceptionInfo, PHYSICAL_CONTEXT* physicalContext, VOID* v);


/* FUNCTIONS */

// The main procedure of the tool.
int main(int argc, char* argv[])
{
	// Initialize PIN library
	if (PIN_Init(argc, argv))
	{
		// Print help message if -h(elp) is specified in the command line or the command line is invalid 
		std::cerr << KNOB_BASE::StringKnobSummary() << std::endl;
		return -1;
	}

	// Open output file stream
	OpenOutputFile(KnobOutputFile.Value());

	// Parse and save params
	_skip = KnobSkip.Value();

	// Parse tracked routines list
	std::istringstream knobTrackRoutinesValue(KnobTrackRoutines.Value());
	std::string currentRoutineName;
	while (std::getline(knobTrackRoutinesValue, currentRoutineName, ';'))
		_trackedRoutines.push_back(currentRoutineName);

	// Some initializations
	InitSyscalls();

	// Instrument instructions and routines
	IMG_AddInstrumentFunction(InstrumentImage, 0);
	TRACE_AddInstrumentFunction(InstrumentTrace, 0);
	RTN_AddInstrumentFunction(InstrumentRoutine, 0);

	// Register syscall instrumentation callbacks
	if (KnobPrintSyscalls.Value())
	{
		PIN_AddSyscallEntryFunction(HandleSyscallEntry, NULL);
		PIN_AddSyscallExitFunction(HandleSyscallExit, NULL);
	}

	// Register callback for program end
	PIN_AddFiniFunction(Fini, NULL);

	// Handle internal exceptions (for debugging)
	PIN_AddInternalExceptionHandler(HandlePinToolException, NULL);

	// Load symbols to access function name information
	PIN_InitSymbols();

	// Start the target program
	PIN_StartProgram();
	return 0;
}


/* CALLBACKS */

// [Callback] Prints executed instructions.
VOID PrintInstruction(CONTEXT* context, UINT64 rip, ImageData* image, UINT64 offset)
{
	// Skip?
	if (_skip > 0)
		return;

	// Print instruction info
	// Flush output file such that we get a full trace when the program crashes
	outputFile
		<< _indentPadding
		<< "I"
		<< " " << image->Id
		<< " " << std::hex << offset
		//	<< " " << "(" << std::hex << rip << ")"
		<< std::endl << std::flush;
}

// [Callback] Prints the value of the given register.
VOID PrintRegister(CONTEXT* context, REG reg, int width)
{
	if (_skip > 0)
		return;

	// Retrieve register value
	PIN_GetContextRegval(context, reg, _registerValueBuffer);

	// Print register value
	outputFile
		<< _indentPadding << "  R "
		<< REG_StringShort(reg) << " -> ";
	for (int i = 0; i < width; ++i)
		outputFile << std::setw(2) << std::hex << (int)_registerValueBuffer[i];

	outputFile << std::endl;
}

// [Callback] Prints the address and value of a memory write.
VOID PrintMemoryWrite(CONTEXT* context, ADDRINT rip, ImageData* image, UINT64 offset, UINT32 width)
{
	if (_skip > 0)
		return;

	outputFile << _indentPadding
		<< "  M"
		<< " " << _currentMemoryOperandAddress
		<< " -> ";

	// Resize memory access buffer, if necessary
	if ((int)width > _memoryAccessBufferSize)
	{
		if (_memoryAccessBuffer != nullptr)
			delete[] _memoryAccessBuffer;

		_memoryAccessBuffer = new UINT8[width];
		_memoryAccessBufferSize = width;
	}

	// Safely retrieve written data
	UINT8* data = _memoryAccessBuffer;
	size_t readBytes = PIN_SafeCopy(data, reinterpret_cast<UINT8*>(_currentMemoryOperandAddress), width);
	if (readBytes < width)
		outputFile << "<DATA READ ERROR>";
	else
	{
		// Dump data
		for (int i = 0; i < (int)width; ++i)
			outputFile << std::setw(2) << std::hex << (int)_memoryAccessBuffer[i];
	}

	// Flush output file such that we get a full trace when the program crashes
	outputFile << std::endl << std::flush;
}

// [Callback] Prints call targets.
VOID PrintFunctionName(ADDRINT branchTargetAddress)
{
	if (_skip > 0)
		return;

	// Try to find function name and print it
	auto search = _addressToFunctionNameMap.find(branchTargetAddress);
	if (search != _addressToFunctionNameMap.end())
	{
		outputFile
			<< _indentPadding
			<< "C "
			<< search->second << std::endl;
	}
	else
	{
		outputFile
			<< _indentPadding
			<< "C <?>"
			<< std::endl;
	}
}

// [Callback] Prints call targets and the first argument registers.
VOID PrintFunctionNameWithArguments(ADDRINT branchTargetAddress, CONTEXT *context)
{
	if (_skip > 0)
		return;

	// Try to find function name and print it
	auto search = _addressToFunctionNameMap.find(branchTargetAddress);
	if (search != _addressToFunctionNameMap.end())
	{
		outputFile
			<< _indentPadding
			<< "C "
			<< search->second;
	}
	else
	{
		outputFile
			<< _indentPadding
			<< "C <?>";
	}

	outputFile
		<< " ("
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_RDI) << ", "
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_RSI) << ", "
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_RDX) << ", "
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_RCX) << ", "
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_R8) << ", "
		<< std::hex << PIN_GetContextReg(context, LEVEL_BASE::REG_R9)
		<< ")"
		<< std::endl;
}

// [Callback] Stores the memory operand address of the current instruction.
VOID StoreInstructionMemoryOperandAddress(UINT64 effectiveAddress)
{
	_currentMemoryOperandAddress = effectiveAddress;
}

// [Callback] Increments the counter for the given opcode.
VOID IncrementInstructionCounter(OPCODE opcode)
{
	++_instructionCounters[opcode];
}

// [Callback] Updates the skip counter.
VOID UpdateSkip()
{
	if (_skip > 0)
		--_skip;
}

// [Callback] Updates the indentation.
VOID UpdateIndent(int delta)
{
	_indent += delta;
	if (_indent < 0)
		_indent = 0;

	if (KnobIndentMode.Value() == 1)
		_indentPadding = std::string(3 * _indent, ' ');
	else if (KnobIndentMode.Value() == 2)
		_indentPadding = std::string(_indent, '\t');
}

// [Callback] Instruments basic blocks and instructions.
VOID InstrumentTrace(TRACE trace, VOID* v)
{
	for (BBL bbl = TRACE_BblHead(trace); BBL_Valid(bbl); bbl = BBL_Next(bbl))
	{
		for (INS ins = BBL_InsHead(bbl); INS_Valid(ins); ins = INS_Next(ins))
		{
			// Resolve instruction
			ADDRINT rip = INS_Address(ins);
			auto image = GetImageContainingAddress(rip);
			UINT64 offset = image->GetOffset(rip);

			if (KnobPrintInstructions.Value())
			{
				// Print executed instruction
				INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)PrintInstruction,
					IARG_CONST_CONTEXT,
					IARG_ADDRINT, rip,
					IARG_PTR, image,
					IARG_UINT64, offset,
					IARG_END);
			}

			if(KnobCountOpcodes.Value())
			{
				// Count instruction
				INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)IncrementInstructionCounter,
					IARG_UINT64, INS_Opcode(ins),
					IARG_END);
			}

			// Print function names
			if (INS_IsCall(ins))
			{
				if (KnobPrintArguments.Value())
				{
					INS_InsertCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)PrintFunctionNameWithArguments,
						IARG_BRANCH_TARGET_ADDR,
						IARG_CONST_CONTEXT,
						IARG_END);
				}
				else
				{
					INS_InsertCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)PrintFunctionName,
						IARG_BRANCH_TARGET_ADDR,
						IARG_END);
				}
			}

			// Print memory writes
			if (KnobPrintMemoryWrites.Value() && INS_IsMemoryWrite(ins) && !INS_IsControlFlow(ins))
			{
				INS_InsertPredicatedCall(ins, IPOINT_BEFORE, (AFUNPTR)StoreInstructionMemoryOperandAddress,
					IARG_MEMORYWRITE_EA,
					IARG_END);
				INS_InsertPredicatedCall(ins, IPOINT_AFTER, (AFUNPTR)PrintMemoryWrite,
					IARG_CONST_CONTEXT,
					IARG_ADDRINT, rip,
					IARG_PTR, image,
					IARG_UINT64, offset,
					IARG_MEMORYWRITE_SIZE,
					IARG_END);
			}

			// Print written registers
			if (KnobPrintRegisterWrites.Value())
			{
				for (UINT32 i = 0; i < INS_MaxNumWRegs(ins); ++i)
				{
					REG reg = INS_RegW(ins, i);
					if (REG_valid(reg) && !REG_is_any_flags_type(reg) && INS_IsValidForIpointAfter(ins))
					{
						// Track register
						INS_InsertPredicatedCall(ins, IPOINT_AFTER, (AFUNPTR)PrintRegister,
							IARG_CONST_CONTEXT,
							IARG_ADDRINT, reg,
							IARG_ADDRINT, REG_Size(reg),
							IARG_END);
					}
				}
			}

			// Update indentation
			if (KnobIndentMode.Value() != 0 && (INS_IsCall(ins) || INS_IsRet(ins)))
			{
				INS_InsertCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)UpdateIndent,
					IARG_UINT64, static_cast<UINT64>(INS_IsCall(ins) ? 1 : -1),
					IARG_END);
			}

			// Update skip counter
			if (INS_IsValidForIpointAfter(ins))
			{
				INS_InsertCall(ins, IPOINT_AFTER, (AFUNPTR)UpdateSkip,
					IARG_CALL_ORDER, CALL_ORDER_LAST,
					IARG_END);
			}
			else if (INS_IsValidForIpointTakenBranch(ins))
			{
				INS_InsertCall(ins, IPOINT_TAKEN_BRANCH, (AFUNPTR)UpdateSkip,
					IARG_CALL_ORDER, CALL_ORDER_LAST,
					IARG_END);
			}
			else
			{
				// Rare cases like syscall
				INS_InsertCall(ins, IPOINT_BEFORE, (AFUNPTR)UpdateSkip,
					IARG_CALL_ORDER, CALL_ORDER_LAST,
					IARG_END);
			}
		}
	}
}

// [Callback] Outputs a function call.
VOID LogFunctionCall(UINT32 id)
{
	std::cerr
		<< _colorDefault << _functionNames[id]
		<< "(...)"
		<< std::endl;

	outputFile
		<< _indentPadding
		<< "F "
		<< _functionNames[id]
		<< std::endl;
}

// Tries to locate a function with the given name and to register a
// logging callback when this function is called.
void TryRegisterRoutine(IMG& img, std::string name)
{
	auto rtn = RTN_FindByName(img, name.c_str());
	if (RTN_Valid(rtn))
	{
		int index = _functionNames.size();
		_functionNames.push_back(name);

		RTN_Open(rtn);
		RTN_InsertCall(rtn, IPOINT_BEFORE, AFUNPTR(LogFunctionCall),
			IARG_UINT32, index,
			IARG_END);

		RTN_Close(rtn);

		std::cerr << "Instrumented function " << name << std::endl;
	}
}

// [Callback] Stores image info and instruments specific functions.
VOID InstrumentImage(IMG img, VOID* v)
{
	// Record and print image load
	HandleImageLoad(img);

	// Track interesting routines
	for (auto& name : _trackedRoutines)
		TryRegisterRoutine(img, name);
}

// [Callback] Records information for fast address to symbol lookup.
VOID InstrumentRoutine(RTN rtn, VOID* v)
{
	std::string rtnName = RTN_Name(rtn);
	ADDRINT rtnAddress = RTN_Address(rtn);

	_addressToFunctionNameMap[rtnAddress] = rtnName;
}

// [Callback] Outputs some information at program end.
VOID Fini(INT32 code, VOID *v)
{
	if(KnobCountOpcodes.Value())
	{
		// Sort counts
		std::vector<std::pair<OPCODE, int64_t>> countsSorted(_instructionCounters.begin(), _instructionCounters.end());
		sort(countsSorted.begin(), countsSorted.end(), [](const std::pair<OPCODE, int64_t>& a, const std::pair<OPCODE, int64_t>& b) { return a.second < b.second; });

		std::cerr << "Opcode counts:" << std::endl;
		for(auto opcodeCount : countsSorted)
		{
			std::cerr << "  "
				<< std::dec << OPCODE_StringShort(opcodeCount.first)
				<< ": " << opcodeCount.second
				<< std::endl;
		}
	}
}

// [Callback] Handles an internal exception of this trace tool.
EXCEPT_HANDLING_RESULT HandlePinToolException(THREADID tid, EXCEPTION_INFO* exceptionInfo, PHYSICAL_CONTEXT* physicalContext, VOID* v)
{
	// Output exception data
	std::cerr << "Internal exception: " << PIN_ExceptionToString(exceptionInfo) << std::endl;
	return EHR_UNHANDLED;
}