148 lines
4.9 KiB
C++
Executable File
148 lines
4.9 KiB
C++
Executable File
#include "initialize_stack.hpp"
|
|
#include <algorithm>
|
|
#include <fstream>
|
|
#include <math.h>
|
|
|
|
#define ALLOF(a) begin(a), end(a)
|
|
|
|
using namespace std;
|
|
using namespace IRDB_SDK;
|
|
using namespace InitStack;
|
|
|
|
//
|
|
// constructor
|
|
//
|
|
InitStack_t::InitStack_t(FileIR_t *p_variantIR,
|
|
const string &p_functions_filename, int p_init_value,
|
|
bool p_verbose)
|
|
: Transform_t(p_variantIR), // initialize the Transform class so things like
|
|
// insertAssembly and getFileIR() can be used
|
|
m_init_value(p_init_value), // member variable inits, these will vary
|
|
// depending on your transform's objectives
|
|
m_verbose(p_verbose), m_num_transformed(0) {
|
|
// check whether to read in a list of functions to transform
|
|
if (p_functions_filename == "") {
|
|
cout << "Auto-initialize all functions" << endl;
|
|
m_funcs_to_init =
|
|
getFileIR()->getFunctions(); // use all functions from the IR
|
|
} else {
|
|
cout << "Auto-initialize functions specified in: " << p_functions_filename
|
|
<< endl;
|
|
readFunctionsFromFile(p_functions_filename); // read functions from file
|
|
}
|
|
}
|
|
|
|
//
|
|
// read list of functions to auto-initialize
|
|
//
|
|
// post conditions: set of functions to auto-initialize
|
|
//
|
|
void InitStack_t::readFunctionsFromFile(const string &p_filename) {
|
|
// get all functions for readability of the rest of the code
|
|
const auto &all_funcs = getFileIR()->getFunctions();
|
|
|
|
// open input file and check for successful open
|
|
ifstream functionsFile(
|
|
p_filename); // can't use auto decl here because of lack of copy
|
|
// constructor in ifstream class
|
|
if (!functionsFile.is_open())
|
|
throw runtime_error("Cannot open " + p_filename);
|
|
|
|
// read each line of the input file.
|
|
auto line = string();
|
|
while (functionsFile >> line) {
|
|
// locate a function with the name read from the file.
|
|
const auto func_it = find_if(ALLOF(all_funcs), [&](const Function_t *f) {
|
|
return f->getName() == line;
|
|
});
|
|
|
|
// if found, log and insert it into the set to transform
|
|
if (func_it != end(all_funcs)) {
|
|
auto f = *func_it;
|
|
cout << "Adding " << f->getName() << " to function list" << endl;
|
|
m_funcs_to_init.insert(f);
|
|
}
|
|
}
|
|
}
|
|
|
|
//
|
|
// Execute the transform by transforming all to-transform functions
|
|
//
|
|
// preconditions: the FileIR is read as from the IRDB. valid file listing
|
|
// functions to auto-initialize postcondition: instructions added to
|
|
// auto-initialize stack for each specified function
|
|
//
|
|
//
|
|
bool InitStack_t::execute() {
|
|
// transform all functions
|
|
for (auto f : m_funcs_to_init)
|
|
initStack(f);
|
|
|
|
// #ATTRIBUTE is a convention used to help find useful information in log
|
|
// files
|
|
cout << "#ATTRIBUTE InitStack::num_transformed=" << m_num_transformed << endl;
|
|
|
|
return m_num_transformed > 0; // true means success
|
|
}
|
|
|
|
//
|
|
// preconditions : f is not NULL
|
|
// postconditions: stack auto-initialized if stack frame size > 0
|
|
//
|
|
void InitStack_t::initStack(Function_t *f) {
|
|
// preconditions
|
|
assert(f != nullptr);
|
|
|
|
// check that the frame size is in the area we care about
|
|
const auto frame_size = f->getStackFrameSize();
|
|
|
|
// nothing to init.
|
|
if (frame_size == 0)
|
|
return;
|
|
|
|
const auto num_locs = static_cast<uint64_t>(ceil(frame_size / 4.0));
|
|
|
|
// debug output
|
|
cout << "Function: " << f->getName()
|
|
<< " frame size: " << f->getStackFrameSize() << endl;
|
|
|
|
// not all functions have an entry point
|
|
const auto entry = f->getEntryPoint();
|
|
if (!entry)
|
|
return;
|
|
|
|
// log what we are doing
|
|
cout << "Function: " << f->getName() << " auto-initialize " << dec << num_locs
|
|
<< " stack memory locations (4 bytes at a time) with value = " << hex
|
|
<< m_init_value << endl;
|
|
|
|
// determine the registers to use on x86-32 or x86-64
|
|
const auto sp_reg =
|
|
getFileIR()->getArchitectureBitWidth() == 64 ? "rsp" : "esp";
|
|
const auto scratch_reg =
|
|
getFileIR()->getArchitectureBitWidth() == 64 ? "r11" : "ecx";
|
|
|
|
// Now, do the dirty work of inserting new assembly to initialize the stack.
|
|
// Insert these instructions at the start of the function (to initialize the
|
|
// stack frame before the function runs) Assume: flags are dead at function
|
|
// entry. Future work: Verify this is true using dead register list. Note: we
|
|
// spill a scratch register into the red zone at 120 bytes passed the end of
|
|
// the frame
|
|
|
|
const auto newInsns = insertAssemblyInstructionsBefore(
|
|
entry,
|
|
string() + " mov [%%1 + %%2], %%3\n"
|
|
" mov %%3, -%%4\n"
|
|
"L1: mov dword [%%1 + %%3 * 4 - 4], %%5\n"
|
|
" inc %%3\n"
|
|
" jnz 0\n"
|
|
" mov %%3, [%%1 + %%2]\n",
|
|
{sp_reg, to_string(-f->getStackFrameSize() - 120), scratch_reg,
|
|
to_string(num_locs), to_string(m_init_value)});
|
|
|
|
newInsns[4]->setTarget(newInsns[2]); // Link jnz to L1.
|
|
|
|
// bump stats
|
|
m_num_transformed++;
|
|
}
|