Notes

#include "BPatch.h"
#include "BPatch_addressSpace.h"
#include "BPatch_binaryEdit.h"
#include "BPatch_function.h"
#include "BPatch_point.h"
#include <string>

typedef set<ADDRINT> VTableSet; // Related vtables
typedef map<DynInst::Address, VTableSet> CallSiteMapping;

// Contains valid vtables for a vcall site
CallSiteMapping callsite_mapping;
size_t data_size = 0;
uint8_t addr_size = 0;

bool parse_callsite_mapping(const char *input_file) {
callsite_mapping.clear();

ifstream file(input_file);
string line;
string app_name;

bool first = true;
while(std::getline(file, line)) {
std::istringstream parser(line);
DynInst::Address callsite_address = 0;
std::string vtable_entry;

// Read application name.
if(first) {
first = false;
parser >> app_name;
global_app_name = app_name;
continue;
}

parser >> hex >> callsite_address;
if(parser.fail()) {
return false;
}
data_size += 2; // Reserve 2 bytes for hierarchy id, if new callsite

while(parser >> vtable_entry) {
if(parser.fail()) {
return false;
}
DynInst::Address vtable_addr = 0;
std::string vtable_addr_str, module_name;

std::istringstream parser_entry(vtable_entry);
if(parser_entry.fail()) {
return false;
}
getline(parser_entry, module_name, ':');
getline(parser_entry, vtable_addr_str, ':');

// Only support vtables residing in applications for the moment.
if(module_name != app_name) {
std::cerr << "At the moment only vtables in "
<< "the same module are supported" << std::endl;
return false;
}

std::istringstream parser_vtable_addr(vtable_addr_str);
if(parser_vtable_addr.fail()) {
return false;
}
parser_vtable_addr >> hex >> vtable_addr;

// TODO: QUICK FIX FOR MYSQL PROFILING
if(vtable_addr == 1) {
continue;
}

callsite_mapping[callsite_address].insert(vtable_addr);
data_size += addr_size;
}

data_size += 8; // Reserve 8 Bytes to mark end of set (null byte)
}

return true;
}

bool insertInstrumentation(const BPatch_snippet &snippet) {
for(const auto &it : callsite_mapping) {
std::vector<BPatch_point *> points;
app->findPoints(it.first, &points); // BPatch_object
app->insertSnippet(snippet, points, BPatch_callBefore);
}
}

int main(int argc, char **argv) {
if (argc < 2) {
std::cout << "Usage: " << argv[0] << " binary" << std::endl;
return 1;
}

std::string progName = argv[1];
auto app = bpatch.openBinary(progName);
addr_size = app->getModules()[0]->getAddressWidth();

// Do instrumentation
/**
* Class dependency:
* 1) Read vcalls from progName + ".vcalls"
* 2) Insert data section containing allowed sets of vtables
* 3) At each vcall allow hierarchy that includes all vtables found there
*/
// 1) Read vcalls
if(!parse_callsite_mapping(progName + ".vcalls")) {
std::cerr << "Could not parse the callsite information file."
<< std::endl;
return -1;
}

// 2) Insert data section containing allowed sets of vtables
// Maybe use BPatch_snippet::BPatch_constExpr(const void *value)
/*auto alloc_expr = app->createVariable(
addr,
type,
var_name,
in_module
);*/
uint8_t *data = malloc(data_size);
size_t *offset = 0;
uint16_t hier_id;
for(const auto &it : callsite_mapping) {
data[offset] = hier_id;
++hier_id;
offset += 2;

for(const auto &vtbl_it : it.second) {
data[offset] = vtbl_it;
offset += 8;
}

data[offset] = 0;
offset += 8; // TODO: Use BPatch_module::getAddressWidth()
}
auto alloc_expr = app->malloc(data_size);
alloc_expr->writeValue(data, data_size);

// 3) Create instrumentation snippet
// Get objptr
std::vector<BPatch_register> regs;
if(!app->getRegisters(&regs)) {
cerr << "At the moment only vtables in "
<< "the same module are supported" << endl;
return -1;
}
BPatch_register *thisptr;
if(addr_size == 8) {
for(const auto &reg : regs) {
if(reg.name() == "rdi") {
thisptr = &reg;
}
}
} else if(addr_size == 4) {
for(const auto &reg : regs) {
if(reg.name() == "ecx") {
thisptr = &reg;
}
}
}
auto objptr = BPatch_registerExpr(*thisptr);

// Get vtblptr by dereferencing the objptr
//objptr->readValue(); //TODO: How to write it into a register?


/**
* Class hierarchy:
* 1) Read vcalls from progName + ".vcalls"
* 2) At each vcall allow hierarchy that includes the most derived vtable
* found there and all of its base classes
*/



/**
* Readonly memory:
* 1) Read vcalls from progName + ".vcalls"
* 2) Add new section containing bitfield indicating writable memory
* sections
* 3) Add hook to the binary when new shared libraries are loaded into
* address space
* 3.1) Make section writable
* 3.2) Set bits according to memory pages of mapped library
* 3.3) Make section non-writable
* 4) At each vcall check if vtable address is in writable memory
*/



// Insert instrumentation and write binary to disk
insertInstrumentation(snippet);
app->writeFile(std::string(progName) + ".instrumented");

return 0;
}