diff --git a/afl_transforms/tools/zax/test/test_context.c b/afl_transforms/tools/zax/test/test_context.c
new file mode 100644
index 0000000000000000000000000000000000000000..6e43960a11c49a0d468c728879f8bb2619e70ca2
--- /dev/null
+++ b/afl_transforms/tools/zax/test/test_context.c
@@ -0,0 +1,34 @@
+#include <stdio.h>
+volatile int bar(int a)
+{
+ if (a == 2)
+ return 5;
+ else
+ return 3;
+}
+
+volatile int foo(int a)
+{
+ if (a==3)
+ return bar(a);
+ else
+ return bar(2);
+
+}
+
+volatile int bob(int a)
+{
+ a = bar(7);
+ if (a==2)
+ return 1;
+ else return 2;
+
+}
+
+int main(int argc, char **argv)
+{
+ int x = foo(argc);
+ int y = bar(argc);
+ int z = bob(argc);
+ printf("out=%d\n", x+y+z);
+}
diff --git a/afl_transforms/tools/zax/test/test_context.sh b/afl_transforms/tools/zax/test/test_context.sh
new file mode 100755
index 0000000000000000000000000000000000000000..7505a722b9b7c82cbbcaadc975e007deb7bea232
--- /dev/null
+++ b/afl_transforms/tools/zax/test/test_context.sh
@@ -0,0 +1,106 @@
+cd $(dirname $(realpath $0) )
+
+PUT=test_context.exe
+ZAFL_PUT="$PUT.zafl $PUT.zafl.context_sensitive"
+MYARG="a"
+
+log_msg()
+{
+ echo "TEST PASS: $1"
+}
+
+log_error()
+{
+ echo "TEST FAIL: $1"
+ exit 1
+}
+
+check_afl()
+{
+ which afl-showmap >/dev/null 2>&1
+ if [ ! $? -eq 0 ]; then
+ log_error "AFL doesn't seem to be installed. Try: 'sudo apt install afl' before proceeding or download/build afl directly from source"
+ fi
+}
+
+build_one()
+{
+ local orig=$1
+ local zafl=$2
+ shift
+ shift
+ zafl.sh $orig $zafl $@
+ if [ $? -eq 0 ]; then
+ log_msg "build $zafl"
+ else
+ log_error "build $zafl"
+ fi
+}
+
+build_all()
+{
+ g++ test_context.c -o $PUT
+}
+
+zafl_all()
+{
+ for p in $*
+ do
+ build_one $p $p.zafl -v -t $p.analysis
+ build_one $p $p.zafl.context_sensitive --enable-context-sensitivity function -v -t $p.analysis.context_sensitive
+ done
+}
+
+clean_all()
+{
+ rm -fr ${PUT}*
+}
+
+verify_output()
+{
+ local arg=$1
+ shift
+ local orig_zafl=$1
+ shift
+ local all_configs=$*
+
+ ./$orig_zafl $arg > $orig_zafl.output.orig
+
+ for p in $all_configs
+ do
+ echo "Program under test: $p"
+ ./${p} $arg > $p.output
+ diff $orig_zafl.output.orig $p.output
+ if [ ! $? -eq 0 ]; then
+ log_error "output verification failure: $p.output"
+ fi
+ done
+
+ log_msg "output verified for $orig_zafl"
+}
+
+clean_all
+check_afl
+
+build_all
+
+zafl_all $PUT
+verify_output $MYARG $PUT $ZAFL_PUT
+
+zafl_map=${PUT}.zafl.map
+afl-showmap -o $zafl_map -- ./${PUT}.zafl
+count_zafl=$(wc -l ${zafl_map} | cut -d' ' -f1)
+
+zafl_cs_map=${PUT}.zafl.map
+afl-showmap -o $zafl_cs_map -- ./${PUT}.zafl.context_sensitive
+count_zafl_cs=$(wc -l ${zafl_cs_map} | cut -d' ' -f1)
+
+# difference must be exactly 2
+let diff=$count_zafl_cs-$count_zafl
+if [ $diff -eq 2 ]; then
+ log_msg "context sensitive map has +2 entries over baseline zafl map"
+else
+ log_error "context sensitive map does not have expected number of entries (should be +2): map_size(zafl):$count_zafl map_size(zafl_context_sensitive):$count_zafl_cs"
+fi
+
+clean_all
diff --git a/afl_transforms/tools/zax/zax.cpp b/afl_transforms/tools/zax/zax.cpp
index 334894d3c4f2dde5f10b304b705fac86ddfcdc14..351671a2d8a7437007bddbdd220dc723f0c61369 100644
--- a/afl_transforms/tools/zax/zax.cpp
+++ b/afl_transforms/tools/zax/zax.cpp
@@ -73,15 +73,18 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
char *reg_temp16 = NULL;
char *reg_trace_map = NULL;
char *reg_prev_id = NULL;
+ char *reg_context = NULL;
+ char *reg_context16 = NULL;
auto save_temp = true;
auto save_trace_map = true;
auto save_prev_id = true;
+ auto save_context = (getContextSensitivity() != ContextSensitivity_None) ? true : false;
auto block_record=BBRecord_t();
const auto trace_map_fixed_addr = getenv("ZAFL_TRACE_MAP_FIXED_ADDRESS");
const auto do_fixed_addr_optimization = (trace_map_fixed_addr!=nullptr);
- const auto num_free_regs_desired = 3;
+ const auto num_free_regs_desired = save_context ? 4 : 3;
auto instr = getInstructionToInstrument(p_bb, num_free_regs_desired);
if (!instr) throw;
@@ -97,8 +100,8 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
{
auto regset = getDeadRegs(instr);
live_flags = regset.find(IRDB_SDK::rn_EFLAGS)==regset.end();
- const auto allowed_regs = RegisterSet_t({rn_RAX, rn_RBX, rn_RCX, rn_RDX, rn_R8, rn_R9, rn_R10, rn_R11, rn_R12, rn_R13, rn_R14, rn_R15});
+ const auto allowed_regs = RegisterSet_t({rn_RAX, rn_RBX, rn_RCX, rn_RDX, rn_R8, rn_R9, rn_R10, rn_R11, rn_R12, rn_R13, rn_R14, rn_R15});
auto free_regs = getFreeRegs(regset, allowed_regs);
for (auto r : regset)
@@ -121,6 +124,13 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
save_prev_id = false;
free_regs.erase(r);
}
+ else if (r == rn_R8 && save_context)
+ {
+ reg_context=strdup("r8");
+ reg_context16=strdup("r8w");
+ save_context = false;
+ free_regs.erase(r);
+ }
}
// if we failed to do the assignment, check for any other register to fill the assignment.
@@ -151,6 +161,19 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
save_prev_id = false;
free_regs.erase(r);
}
+
+ if (getContextSensitivity() != ContextSensitivity_None)
+ {
+ if (free_regs.size() >= 1)
+ {
+ auto r = *free_regs.begin();
+ auto r16 = convertRegisterTo16bit(r);
+ reg_context = strdup(registerToString(r).c_str());
+ reg_context16 = strdup(registerToString(r16).c_str());
+ save_context = false;
+ free_regs.erase(r);
+ }
+ }
}
// In the event we couldn't find a free register, or we aren't using stars
@@ -161,13 +184,16 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
if (!reg_temp16) reg_temp16 = strdup("ax");
if (!reg_trace_map) reg_trace_map = strdup("rcx");
if (!reg_prev_id) reg_prev_id = strdup("rdx");
+ if (!reg_context) reg_context = strdup("r8");
+ if (!reg_context16) reg_context16 = strdup("r8w");
if (m_verbose)
{
cout << "save_temp: " << save_temp << " save_trace_map: " << save_trace_map << " save_prev_id: " << save_prev_id << " live_flags: " << live_flags << endl;
cout << "reg_temp: " << reg_temp << " " << reg_temp32 << " " << reg_temp16
<< " reg_trace_map: " << reg_trace_map
- << " reg_prev_id: " << reg_prev_id << endl;
+ << " reg_prev_id: " << reg_prev_id
+ << " reg_context: " << reg_context << endl;
}
// warning: first instrumentation must use insertAssemblyBefore
@@ -205,10 +231,11 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
// Omit saving any registers we don't need to save because we
// were able to locate a free register.
if (p_honorRedZone) do_insert("lea rsp, [rsp-128]");
- if (save_temp) do_insert("push rax");
- if (save_trace_map) do_insert("push rcx");
- if (save_prev_id) do_insert("push rdx");
- if (live_flags) do_insert("pushf");
+ if (save_temp) do_insert("push rax");
+ if (save_trace_map) do_insert("push rcx");
+ if (save_prev_id) do_insert("push rdx");
+ if (save_context) do_insert("push r8");
+ if (live_flags) do_insert("pushf");
const auto live_flags_str = live_flags ? "live" : "dead";
if (m_verbose) cout << " flags are " << live_flags_str << endl;
@@ -239,6 +266,17 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
do_insert(buf);
create_got_reloc(getFileIR(), m_prev_id, tmp);
+
+ if (getContextSensitivity() != ContextSensitivity_None)
+ {
+ sprintf(buf, "C%d: mov %s, QWORD [rel C%d]", labelid, reg_context, labelid);
+ do_insert(buf);
+ create_got_reloc(getFileIR(), m_context_id, tmp);
+
+ sprintf(buf,"mov %s, [%s]", reg_context, reg_context);
+ do_insert(buf);
+ }
+
// if we are using a variable address trace map, generate the address.
if(!do_fixed_addr_optimization)
{
@@ -248,6 +286,7 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
create_got_reloc(getFileIR(), m_trace_map, tmp);
}
+ // compute index into trace map
// do the calculation to has the previouus block ID with this block ID
// in the faster or slower fashion depending on the requested technique.
if (!p_collafl_optimization)
@@ -260,6 +299,14 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
sprintf(buf, "xor %s,0x%x", reg_temp16, blockid);
do_insert(buf);
+ // hash with calling context value
+ if (getContextSensitivity() != ContextSensitivity_None)
+ {
+ // xor ax, <context_id_register>
+ sprintf(buf, "xor %s,%s", reg_temp16, reg_context16);
+ do_insert(buf);
+ }
+
// 15: movzx eax,ax
sprintf(buf,"movzx %s,%s", reg_temp32, reg_temp16);
do_insert(buf);
@@ -298,16 +345,18 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
sprintf(buf, "mov %s, 0x%x", reg_temp32, blockid >> 1);
do_insert(buf);
+ // store prev_id
// 23: mov WORD PTR [rdx],ax
sprintf(buf, "mov WORD [%s], %s", reg_prev_id, reg_temp16);
do_insert(buf);
// finally, restore any flags/registers so that the program can execute.
- if (live_flags) do_insert("popf");
- if (save_prev_id) do_insert("pop rdx");
- if (save_trace_map) do_insert("pop rcx");
- if (save_temp) do_insert("pop rax");
- if (p_honorRedZone) do_insert("lea rsp, [rsp+128]");
+ if (live_flags) do_insert("popf");
+ if (save_context) do_insert("pop r8");
+ if (save_prev_id) do_insert("pop rdx");
+ if (save_trace_map) do_insert("pop rcx");
+ if (save_temp) do_insert("pop rax");
+ if (p_honorRedZone) do_insert("lea rsp, [rsp+128]");
m_modifiedBlocks[blockid] = block_record;
@@ -319,5 +368,7 @@ void Zax_t::instrumentBasicBlock(BasicBlock_t *p_bb, const bool p_honorRedZone,
free(reg_temp16);
free(reg_trace_map);
free(reg_prev_id);
+ free(reg_context);
+ free(reg_context16);
}
diff --git a/afl_transforms/tools/zax/zax_base.cpp b/afl_transforms/tools/zax/zax_base.cpp
index 90eb6f71ab83f022f8fe1f9f926b373f6957a975..bb723c66e1e915ab68ad709d549bba134b41e4b6 100644
--- a/afl_transforms/tools/zax/zax_base.cpp
+++ b/afl_transforms/tools/zax/zax_base.cpp
@@ -33,12 +33,14 @@
#include <irdb-deep>
#include "zax_base.hpp"
+#include "critical_edge_breaker.hpp"
using namespace std;
using namespace IRDB_SDK;
using namespace Zafl;
#define ALLOF(a) begin(a),end(a)
+#define FIRSTOF(a) (*(begin(a)))
void create_got_reloc(FileIR_t* fir, pair<DataScoop_t*,int> wrt, Instruction_t* i)
{
@@ -118,6 +120,7 @@ ZaxBase_t::ZaxBase_t(IRDB_SDK::pqxxDB_t &p_dbinterface, IRDB_SDK::FileIR_t *p_va
m_plt_zafl_initAflForkServer=ed->appendPltEntry("zafl_initAflForkServer");
m_trace_map = ed->appendGotEntry("zafl_trace_map");
m_prev_id = ed->appendGotEntry("zafl_prev_id");
+ m_context_id = ed->appendGotEntry("zafl_context");
// let's not instrument these functions ever
// see isBlacklisted() for other blacklisted functions
@@ -154,6 +157,10 @@ ZaxBase_t::ZaxBase_t(IRDB_SDK::pqxxDB_t &p_dbinterface, IRDB_SDK::FileIR_t *p_va
m_verbose = false;
m_bb_float_instrumentation = false;
+ setContextSensitivity(ContextSensitivity_None);
+
+ m_entry_point = nullptr;
+
m_labelid = 0;
m_blockid = 0;
@@ -171,6 +178,7 @@ ZaxBase_t::ZaxBase_t(IRDB_SDK::pqxxDB_t &p_dbinterface, IRDB_SDK::FileIR_t *p_va
m_num_exit_blocks_elided = 0;
m_num_entry_blocks_elided = 0;
m_num_single_block_function_elided = 0;
+ m_num_contexts = 0;
}
void ZaxBase_t::setVerbose(bool p_verbose)
@@ -219,6 +227,28 @@ bool ZaxBase_t::getBasicBlockFloatingInstrumentation() const
return m_bb_float_instrumentation;
}
+void ZaxBase_t::setContextSensitivity(ContextSensitivity_t p_context_style)
+{
+ m_context_sensitivity = p_context_style;
+ switch (m_context_sensitivity)
+ {
+ case ContextSensitivity_None:
+ cout << "disable context sensitivity" << endl;
+ break;
+ case ContextSensitivity_Function:
+ cout << "enable context sensitivity (style: function)" << endl;
+ break;
+ case ContextSensitivity_Callsite:
+ cout << "enable context sensitivity (style: callsite)" << endl;
+ break;
+ }
+}
+
+ContextSensitivity_t ZaxBase_t::getContextSensitivity() const
+{
+ return m_context_sensitivity;
+}
+
/*
* Only allow instrumentation in whitelisted functions/instructions
* Each line in file is either a function name or address
@@ -266,6 +296,24 @@ ZaflBlockId_t ZaxBase_t::getBlockId(const unsigned p_max)
return m_blockid;
}
+ZaflContextId_t ZaxBase_t::getContextId(const unsigned p_max)
+{
+ auto counter = 0;
+ auto contextid = 0;
+
+ // only try getting new context id 100 times
+ // avoid returning duplicate if we can help it
+ while (counter++ < 100) {
+ contextid = rand() % p_max;
+ if (m_used_contextid.find(contextid) == m_used_contextid.end())
+ {
+ m_used_contextid.insert(contextid);
+ return contextid;
+ }
+ }
+ return contextid;
+}
+
void ZaxBase_t::insertExitPoint(Instruction_t *p_inst)
{
assert(p_inst->getAddress()->getVirtualOffset());
@@ -304,6 +352,8 @@ void ZaxBase_t::insertForkServer(Instruction_t* p_entry)
cout << "Blacklisting entry point: " << ss.str() << endl;
m_blacklist.insert(ss.str());
+ m_entry_point = p_entry;
+
// insert the instrumentation
auto tmp=p_entry;
const auto regs = vector<string>({ "rdi", "rsi", "rbp", "rdx", "rcx", "rbx", "rax", "r8", "r9", "r10", "r11", "r12", "r13", "r14", "r15"});
@@ -474,7 +524,7 @@ bool ZaxBase_t::isBlacklisted(const Instruction_t *p_inst) const
{
stringstream ss;
ss << "0x" << hex << p_inst->getAddress()->getVirtualOffset();
- return (m_blacklist.count(ss.str()) > 0 || isBlacklisted(p_inst->getFunction()));
+ return (m_blacklist.count(ss.str()) > 0 || isBlacklisted(p_inst->getFunction()) || p_inst == m_entry_point);
}
bool ZaxBase_t::isWhitelisted(const Instruction_t *p_inst) const
@@ -768,6 +818,9 @@ Instruction_t* ZaxBase_t::getInstructionToInstrument(const BasicBlock_t *p_bb, c
for (auto i : p_bb->getInstructions())
{
+ if (isBlacklisted(i))
+ continue;
+
const auto dead_regs = getDeadRegs(i);
const auto num_free_regs = getFreeRegs(dead_regs, allowed_regs).size();
@@ -805,6 +858,202 @@ Instruction_t* ZaxBase_t::getInstructionToInstrument(const BasicBlock_t *p_bb, c
return best_i;
}
+void ZaxBase_t::dumpAttributes()
+{
+ cout << "#ATTRIBUTE num_bb=" << dec << m_num_bb << endl;
+ cout << "#ATTRIBUTE num_bb_instrumented=" << m_num_bb_instrumented << endl;
+ cout << "#ATTRIBUTE num_bb_skipped=" << m_num_bb_skipped << endl;
+ cout << "#ATTRIBUTE num_bb_skipped_pushjmp=" << m_num_bb_skipped_pushjmp << endl;
+ cout << "#ATTRIBUTE num_bb_skipped_nop_padding=" << m_num_bb_skipped_nop_padding << endl;
+ cout << "#ATTRIBUTE num_bb_float_instrumentation=" << m_num_bb_float_instrumentation << endl;
+ cout << "#ATTRIBUTE num_bb_float_register_saved=" << m_num_bb_float_regs_saved << endl;
+ cout << "#ATTRIBUTE graph_optimize=" << boolalpha << m_graph_optimize << endl;
+ cout << "#ATTRIBUTE num_bb_skipped_cond_branch=" << m_num_bb_skipped_cbranch << endl;
+ cout << "#ATTRIBUTE num_style_collafl=" << m_num_style_collafl << endl;
+ cout << "#ATTRIBUTE num_domgraph_blocks_elided=" << m_num_domgraph_blocks_elided << endl;
+ cout << "#ATTRIBUTE num_entry_blocks_elided=" << m_num_entry_blocks_elided << endl;
+ cout << "#ATTRIBUTE num_exit_blocks_elided=" << m_num_exit_blocks_elided << endl;
+ cout << "#ATTRIBUTE num_single_block_function_elided=" << m_num_single_block_function_elided << endl;
+ cout << "#ATTRIBUTE num_contexts=" << m_num_contexts << endl;
+}
+
+// file dump of modified basic block info
+void ZaxBase_t::dumpMap()
+{
+ getFileIR()->setBaseIDS(); // make sure instructions have IDs
+ getFileIR()->assembleRegistry(); // make sure to assemble all instructions
+
+ std::ofstream mapfile("zax.map");
+
+ mapfile << "# BLOCK_ID ID_EP:size ID_OLDEP:size (ID_INSTRUMENTATION:size)*" << endl;
+ for (auto &mb : m_modifiedBlocks)
+ {
+ const auto blockid = mb.first;
+ mapfile << dec << blockid << " ";
+ for (auto &entry : mb.second)
+ {
+ mapfile << hex << entry->getBaseID() << ":" << dec << entry->getDataBits().size() << " ";
+ }
+ mapfile << endl;
+ }
+}
+
+void ZaxBase_t::addContextSensitivity_Callsite(const ControlFlowGraph_t& cfg)
+{
+ assert(0);
+}
+
+// update calling context hash at entry point
+// revert calling context hash on exit
+void ZaxBase_t::addContextSensitivity_Function(const ControlFlowGraph_t& cfg)
+{
+ bool inserted_before = false;
+
+ // don't bother with single block functions
+ if (cfg.getBlocks().size() == 1)
+ return;
+
+ m_num_contexts++;
+
+ //
+ // entry_point
+ // context = prev_context % RANDOM_CONTEXT_ID
+ //
+ // exit point (returns)
+ // context = prev_context % RANDOM_CONTEXT_ID
+ //
+ const auto do_insert=[&](Instruction_t* instr, const string& insn_str) -> Instruction_t*
+ {
+ if (inserted_before)
+ {
+ instr = insertAssemblyAfter(instr, insn_str);
+ return instr;
+ }
+ else
+ {
+ insertAssemblyBefore(instr, insn_str);
+ inserted_before = true;
+ return instr;
+ }
+ };
+
+
+ auto compute_hash_chain = [&](ZaflContextId_t contextid, Instruction_t * instr, string reg_context, string reg_temp) -> Instruction_t*
+ {
+ auto labelid = getLabelId();
+ const auto hash_context_reloc = string("E") + to_string(labelid) + ": mov " + reg_context + ", QWORD [rel E" + to_string(labelid) + "]";
+ auto tmp = instr;
+ tmp = do_insert(tmp, hash_context_reloc);
+ create_got_reloc(getFileIR(), m_context_id, tmp);
+
+ const auto deref_context = string("mov ") + reg_temp + ", [" + reg_context + "]";
+ tmp = do_insert(tmp, deref_context);
+
+ const auto hash_chain = string("xor ") + reg_temp + ", " + to_string(contextid);
+ tmp = do_insert(tmp, hash_chain);
+
+ const auto store_context = string("mov [") + reg_context + "]" + "," + reg_temp;
+ tmp = do_insert(tmp, store_context);
+
+ // 5ae: 48 8b 05 23 0a 20 00 mov rax,QWORD PTR [rip+0x200a23] # 200fd8 <x>
+ // 5b5: c7 00 d2 04 00 00 mov DWORD PTR [rax],0x4d2
+ return tmp;
+ };
+
+ auto add_hash_context_instrumentation = [&](ZaflContextId_t contextid, Instruction_t* i)
+ {
+ inserted_before = false;
+
+ // look for instruction in entry block with at least 1 free reg
+ auto reg_context = string("r14");
+ auto reg_temp = string("r15");
+ bool save_context = false;
+ bool save_temp = false;
+
+ const auto allowed_regs = RegisterSet_t({rn_RAX, rn_RBX, rn_RCX, rn_RDX, rn_R8, rn_R9, rn_R10, rn_R11, rn_R12, rn_R13, rn_R14, rn_R15});
+ const auto dead_regs = getDeadRegs(i);
+ auto free_regs = getFreeRegs(dead_regs, allowed_regs);
+
+ // @todo: red zone?
+
+ if (free_regs.size() > 0)
+ {
+ reg_context = registerToString(FIRSTOF(free_regs));
+ free_regs.erase(FIRSTOF(free_regs));
+ }
+ else
+ {
+ save_context = true;
+ i = do_insert(i, "push " + reg_context);
+ }
+
+ if (free_regs.size() > 0)
+ {
+ reg_temp = registerToString(FIRSTOF(free_regs));
+ free_regs.erase(FIRSTOF(free_regs));
+ }
+ else
+ {
+ save_temp = true;
+ i = do_insert(i, "push " + reg_temp);
+ }
+
+ // compute new hash chain value
+ i = compute_hash_chain(contextid, i, reg_context, reg_temp);
+
+ if (save_temp)
+ i = do_insert(i, "pop " + reg_temp);
+
+ if (save_context)
+ i = do_insert(i, "pop " + reg_context);
+
+ };
+
+ auto contextid = getContextId();
+ auto entry_block = cfg.getEntry();
+ inserted_before = false;
+ add_hash_context_instrumentation(contextid, entry_block->getInstructions()[0]);
+
+ // find all exit blocks
+ auto find_exits = BasicBlockSet_t();
+ copy_if(ALLOF(cfg.getBlocks()), inserter(find_exits, find_exits.begin()), [entry_block](BasicBlock_t* bb) {
+ if (bb == entry_block) return false;
+ if (!bb->getIsExitBlock()) return false;
+ // make sure it's a ret
+ const auto last_instruction_index = bb->getInstructions().size() - 1;
+ return (bb->getInstructions()[last_instruction_index]->getDisassembly().find("ret")!=string::npos);
+ });
+
+ for (const auto &bb: find_exits)
+ {
+ const auto last_instruction_index = bb->getInstructions().size()-1;
+ inserted_before = false;
+ add_hash_context_instrumentation(contextid, bb->getInstructions()[last_instruction_index]);
+ }
+}
+
+void ZaxBase_t::addContextSensitivity(const ControlFlowGraph_t& cfg)
+{
+ if (m_entry_point && m_entry_point->getFunction())
+ {
+ cout << "cfg.func: " << cfg.getFunction()->getName() << " ep.func: " << m_entry_point->getFunction()->getName() << endl;
+ if (m_entry_point->getFunction()->getName() == cfg.getFunction()->getName())
+ {
+ cout << "Do not setup calling context in same function as entry point for fork server" << endl;
+ return;
+ }
+ }
+
+ if (getContextSensitivity() == ContextSensitivity_Callsite)
+ addContextSensitivity_Callsite(cfg);
+ else if (getContextSensitivity() == ContextSensitivity_Function)
+ addContextSensitivity_Function(cfg);
+ else if (getContextSensitivity() == ContextSensitivity_None)
+ return;
+ else
+ throw;
+}
+
/*
* Execute the transform.
*
@@ -814,6 +1063,15 @@ Instruction_t* ZaxBase_t::getInstructionToInstrument(const BasicBlock_t *p_bb, c
*/
int ZaxBase_t::execute()
{
+ if (m_breakupCriticalEdges)
+ {
+ CriticalEdgeBreaker_t ceb(getFileIR(), m_verbose);
+ cout << "#ATTRIBUTE num_bb_extra_blocks=" << ceb.getNumberExtraNodes() << endl;
+
+ getFileIR()->setBaseIDS();
+ getFileIR()->assembleRegistry();
+ }
+
setup();
// for all functions
@@ -931,6 +1189,16 @@ int ZaxBase_t::execute()
m_num_bb_instrumented += keepers.size();
m_num_bb_skipped += (num_blocks_in_func - keepers.size());
+ if (getContextSensitivity() != ContextSensitivity_None)
+ {
+ // this handles inserting the calling context sensitivity value
+ // at entry and exits of functions (todo: or at call sites a la Angora)
+ getFileIR()->assembleRegistry();
+ getFileIR()->setBaseIDS();
+ auto cs_cfg=ControlFlowGraph_t::factory(f);
+ addContextSensitivity(*cs_cfg);
+ }
+
if (m_verbose)
{
getFileIR()->assembleRegistry();
@@ -943,46 +1211,9 @@ int ZaxBase_t::execute()
cout << "Function " << f->getName() << ": " << dec << keepers.size() << "/" << num_blocks_in_func << " basic blocks instrumented." << endl;
};
+
teardown();
return 1;
}
-void ZaxBase_t::dumpAttributes()
-{
- cout << "#ATTRIBUTE num_bb=" << dec << m_num_bb << endl;
- cout << "#ATTRIBUTE num_bb_instrumented=" << m_num_bb_instrumented << endl;
- cout << "#ATTRIBUTE num_bb_skipped=" << m_num_bb_skipped << endl;
- cout << "#ATTRIBUTE num_bb_skipped_pushjmp=" << m_num_bb_skipped_pushjmp << endl;
- cout << "#ATTRIBUTE num_bb_skipped_nop_padding=" << m_num_bb_skipped_nop_padding << endl;
- cout << "#ATTRIBUTE num_bb_float_instrumentation=" << m_num_bb_float_instrumentation << endl;
- cout << "#ATTRIBUTE num_bb_float_register_saved=" << m_num_bb_float_regs_saved << endl;
- cout << "#ATTRIBUTE graph_optimize=" << boolalpha << m_graph_optimize << endl;
- cout << "#ATTRIBUTE num_bb_skipped_cond_branch=" << m_num_bb_skipped_cbranch << endl;
- cout << "#ATTRIBUTE num_style_collafl=" << m_num_style_collafl << endl;
- cout << "#ATTRIBUTE num_domgraph_blocks_elided=" << m_num_domgraph_blocks_elided << endl;
- cout << "#ATTRIBUTE num_entry_blocks_elided=" << m_num_entry_blocks_elided << endl;
- cout << "#ATTRIBUTE num_exit_blocks_elided=" << m_num_exit_blocks_elided << endl;
- cout << "#ATTRIBUTE num_single_block_function_elided=" << m_num_single_block_function_elided << endl;
-}
-
-// file dump of modified basic block info
-void ZaxBase_t::dumpMap()
-{
- getFileIR()->setBaseIDS(); // make sure instructions have IDs
- getFileIR()->assembleRegistry(); // make sure to assemble all instructions
-
- std::ofstream mapfile("zax.map");
-
- mapfile << "# BLOCK_ID ID_EP:size ID_OLDEP:size (ID_INSTRUMENTATION:size)*" << endl;
- for (auto &mb : m_modifiedBlocks)
- {
- const auto blockid = mb.first;
- mapfile << dec << blockid << " ";
- for (auto &entry : mb.second)
- {
- mapfile << hex << entry->getBaseID() << ":" << dec << entry->getDataBits().size() << " ";
- }
- mapfile << endl;
- }
-}
diff --git a/afl_transforms/tools/zax/zax_base.hpp b/afl_transforms/tools/zax/zax_base.hpp
index b2314dda630831cbba3fb3868709bb97416fc6f3..8f2a66f15dbbcf794fbfd968a923724d52425cca 100644
--- a/afl_transforms/tools/zax/zax_base.hpp
+++ b/afl_transforms/tools/zax/zax_base.hpp
@@ -18,8 +18,10 @@ namespace Zafl
using ZaflBlockId_t = uint32_t;
using ZaflLabelId_t = uint32_t;
+ using ZaflContextId_t = uint32_t;
using BBRecord_t = vector<Instruction_t*>;
using RegisterSet_t = IRDB_SDK::RegisterIDSet_t;
+ enum ContextSensitivity_t {ContextSensitivity_None, ContextSensitivity_Callsite, ContextSensitivity_Function};
/*
* Base class for afl-compatible instrumentation:
@@ -41,11 +43,13 @@ namespace Zafl
void setBasicBlockFloatingInstrumentation(bool);
void setEnableForkServer(bool);
void setBreakupCriticalEdges(bool);
+ void setContextSensitivity(ContextSensitivity_t);
void filterPaddingNOP(BasicBlockSet_t& p_in_out);
void filterBlocksByDomgraph(BasicBlockSet_t& in_out, const DominatorGraph_t * dg );
void filterConditionalBranches(BasicBlockSet_t& p_in_out);
void filterEntryBlock(BasicBlockSet_t& in_out, BasicBlock_t* p_entry);
void filterExitBlocks(BasicBlockSet_t& in_out);
+ void addContextSensitivity(const ControlFlowGraph_t&);
protected:
ZaxBase_t(pqxxDB_t &p_dbinterface, FileIR_t *p_variantIR, string p_entry, set<string> p_exits, bool p_use_stars=false, bool p_autozafl=false);
@@ -54,6 +58,7 @@ namespace Zafl
virtual ZaflBlockId_t getBlockId(const unsigned p_maxid=0xFFFF);
virtual ZaflLabelId_t getLabelId(const unsigned p_maxid=0xFFFF);
+ virtual ZaflContextId_t getContextId(const unsigned p_maxid=0xFFFF);
virtual BasicBlockSet_t getBlocksToInstrument (const ControlFlowGraph_t& cfg);
virtual Instruction_t* getInstructionToInstrument(const BasicBlock_t *, const unsigned p_num_free_regs_desired = 0);
virtual void setup();
@@ -73,10 +78,12 @@ namespace Zafl
bool BB_isPushJmp(const BasicBlock_t *) const;
bool BB_isPaddingNop(const BasicBlock_t *) const;
bool getBasicBlockFloatingInstrumentation() const;
+ ContextSensitivity_t getContextSensitivity() const;
bool hasLeafAnnotation(Function_t* fn) const;
RegisterSet_t getDeadRegs(Instruction_t* insn) const;
RegisterSet_t getFreeRegs(const RegisterSet_t& candidates, const RegisterSet_t& allowed) const;
-
+ void addContextSensitivity_Callsite(const ControlFlowGraph_t&);
+ void addContextSensitivity_Function(const ControlFlowGraph_t&);
protected:
pqxxDB_t& m_dbinterface;
@@ -91,13 +98,16 @@ namespace Zafl
bool m_breakupCriticalEdges;
bool m_bb_float_instrumentation; // skip basic blocks based on graph
bool m_verbose;
-
pair<DataScoop_t*,int> m_trace_map; // afl shared memory trace map
pair<DataScoop_t*,int> m_prev_id; // id of previous block
+ pair<DataScoop_t*,int> m_context_id; // calling context variable
Instruction_t* m_plt_zafl_initAflForkServer; // plt entry for afl fork server initialization routine
+ set<ZaflContextId_t> m_used_contextid; // internal bookkeeping to keep track of used block ids
map<ZaflBlockId_t, BBRecord_t> m_modifiedBlocks; // keep track of modified blocks
+ ContextSensitivity_t m_context_sensitivity; // none, @callsite, @function
+
// stats
size_t m_num_bb;
size_t m_num_bb_instrumented;
@@ -112,6 +122,7 @@ namespace Zafl
size_t m_num_entry_blocks_elided;
size_t m_num_exit_blocks_elided;
size_t m_num_single_block_function_elided;
+ size_t m_num_contexts;
private:
string m_fork_server_entry; // string to specify fork server entry point
@@ -120,6 +131,7 @@ namespace Zafl
set<string> m_blacklist; // blacklisted functions and/or instructions
ZaflLabelId_t m_labelid; // internal bookkeeping to generate labels
ZaflBlockId_t m_blockid; // internal bookkeeping to generate labels
+ Instruction_t* m_entry_point; // entry point where fork server was inserted
};
}
diff --git a/afl_transforms/tools/zax/zax_driver.cpp b/afl_transforms/tools/zax/zax_driver.cpp
index 40ff623696b475370316d08e1ecdcfac45d9c1bf..6b3493bbbb6352848ca762165e59c6130bd4eef9 100644
--- a/afl_transforms/tools/zax/zax_driver.cpp
+++ b/afl_transforms/tools/zax/zax_driver.cpp
@@ -49,7 +49,9 @@ static void usage(char* name)
cerr<<"\t[--enable-critical-edge-breakup|-c] Breakup critical edges"<<endl;
cerr<<"\t[--disable-critical-edge-breakup|-C] Do not breakup critical edges (default)"<<endl;
cerr<<"\t[--enable-floating-instrumentation|-i] Select best instrumentation within basic block"<<endl;
- cerr<<"\t[--disable-floating-instrumentation|-I] Instrument firt instruction in basic blocks"<<endl;
+ cerr<<"\t[--disable-floating-instrumentation|-I] Instrument first instruction in basic blocks"<<endl;
+ cerr<<"\t[--enable-context-sensitivity <style>] Use calling context sensitivity, style={callsite,function}"<<endl;
+ cerr<<"\t[--disable-context-sensitivity] Disable calling context sensitivity"<<endl;
}
int main(int argc, char **argv)
@@ -74,6 +76,7 @@ int main(int argc, char **argv)
auto untracer_mode=false;
auto breakup_critical_edges=false;
auto floating_instrumentation=false;
+ auto context_sensitivity=ContextSensitivity_None;
set<string> exitpoints;
srand(getpid()+time(NULL));
@@ -100,9 +103,10 @@ int main(int argc, char **argv)
{"disable-critical-edge-breakup", no_argument, 0, 'C'},
{"enable-floating-instrumentation", no_argument, 0, 'i'},
{"disable-floating-instrumentation", no_argument, 0, 'I'},
+ {"enable-context-sensitivity", required_argument, 0, 'z'},
{0,0,0,0}
};
- const char* short_opts="e:E:w:sv?hagGdDfFucCiI";
+ const char* short_opts="z:e:E:w:sv?hagGdDfFucCiI";
while(true)
{
@@ -172,6 +176,14 @@ int main(int argc, char **argv)
case 'I':
floating_instrumentation=false;
break;
+ case 'z':
+ if (optarg == string("callsite"))
+ context_sensitivity=ContextSensitivity_Callsite; // Angora fuzzer style
+ else if (optarg == string("function"))
+ context_sensitivity=ContextSensitivity_Function;
+ else
+ context_sensitivity=ContextSensitivity_None;
+ break;
default:
break;
}
@@ -231,6 +243,7 @@ int main(int argc, char **argv)
zax->setBasicBlockFloatingInstrumentation(floating_instrumentation);
zax->setEnableForkServer(forkserver_enabled);
zax->setBreakupCriticalEdges(breakup_critical_edges);
+ zax->setContextSensitivity(context_sensitivity);
int success=zax->execute();
diff --git a/afl_transforms/tools/zax/zuntracer.cpp b/afl_transforms/tools/zax/zuntracer.cpp
index 2018809ea64c450419b0f2ea95772860df79e6ba..82cb301283b458de1c339445b32ed3e5452ff7c8 100644
--- a/afl_transforms/tools/zax/zuntracer.cpp
+++ b/afl_transforms/tools/zax/zuntracer.cpp
@@ -1,5 +1,4 @@
#include "zuntracer.hpp"
-#include "critical_edge_breaker.hpp"
using namespace std;
using namespace IRDB_SDK;
@@ -226,22 +225,3 @@ set<BasicBlock_t*> ZUntracer_t::getBlocksToInstrument(ControlFlowGraph_t &cfg)
}
return keepers;
}
-
-//
-// breakup critical edges
-// use block-level instrumentation
-//
-int ZUntracer_t::execute()
-{
- if (m_breakupCriticalEdges)
- {
- CriticalEdgeBreaker_t ceb(getFileIR(), m_verbose);
- cout << "#ATTRIBUTE num_bb_extra_blocks=" << ceb.getNumberExtraNodes() << endl;
-
- getFileIR()->setBaseIDS();
- getFileIR()->assembleRegistry();
- }
-
- return ZaxBase_t::execute();
-}
-
diff --git a/afl_transforms/tools/zax/zuntracer.hpp b/afl_transforms/tools/zax/zuntracer.hpp
index b122f1b847d4d773bf78a566da370395fb6d2ba8..eba33a8874d490d4fdaefd6f393e0e8676f50ecb 100644
--- a/afl_transforms/tools/zax/zuntracer.hpp
+++ b/afl_transforms/tools/zax/zuntracer.hpp
@@ -16,7 +16,6 @@ namespace Zafl
ZUntracer_t(const ZUntracer_t&) = delete;
ZUntracer_t(IRDB_SDK::pqxxDB_t &p_dbinterface, IRDB_SDK::FileIR_t *p_variantIR, string p_entry, set<string> p_exits, bool p_use_stars=false, bool p_autozafl=false);
virtual ~ZUntracer_t() {};
- virtual int execute();
protected:
virtual void instrumentBasicBlock(BasicBlock_t*, const bool p_hasLeafAnnotation, const bool p_collafl_optimization=false);
diff --git a/bin/zafl.sh b/bin/zafl.sh
index d29ef894cc8eb08f4fa5c36c7ef658a64aea6138..8af80beecec7773afe94257545943eb0952b9d10 100755
--- a/bin/zafl.sh
+++ b/bin/zafl.sh
@@ -36,15 +36,17 @@ usage()
echo " -M, --disable-fixed-map Disable fixed address tracing map (default)"
echo " -i, --enable-floating-instrumentation Select best instrumentation point within basic block (default)"
echo " -I, --disable-floating-instrumentation Use first instruction for instrumentation in basic blocks"
+ echo " --enable-context-sensitivity <style> style={callsite,function} only function supported currently"
echo " -v Verbose mode"
echo
}
ida_or_rida_opt=" -c rida "
stars_opt=" -o zax:--stars "
-zax_opt=" "
+zax_opt=""
other_args=""
float_opt=" -o zax:--enable-floating-instrumentation "
+context_sensitivity_opt=""
me=$(whoami)
tmp_dir=/tmp/${me}/$$
@@ -180,13 +182,32 @@ parse_args()
shift
;;
-i | --enable-floating-instrumentation)
- float_opt= " -o zax:--enable-floating-instrumentation "
+ float_opt=" -o zax:--enable-floating-instrumentation "
shift
;;
-I | --disable-floating-instrumentation)
- float_opt= " -o zax:--disable-floating-instrumentation "
+ float_opt=" -o zax:--disable-floating-instrumentation "
+ shift
+ ;;
+ --enable-context-sensitivity)
shift
+ case $1 in
+ function)
+ context_sensitivity_opt=" -o zax:\"--enable-context-sensitivity function\" "
+ shift
+ ;;
+ callsite)
+ context_sensitivity_opt=" -o zax:--enable-context-sensitivity callsite "
+ echo "Error: context sensitivity <callsite> currently unsupported"
+ exit 1
+ ;;
+ *)
+ echo "Error: must specify function or callsite for context sensitivity"
+ exit 1
+ ;;
+ esac
;;
+
-*|--*=) # unsupported flags
echo "Error: Unsupported flag $1" >&2
exit 1
@@ -292,7 +313,7 @@ fi
#
log_msg "Transforming input binary $input_binary into $output_zafl_binary"
-zax_opt=" $zax_opt $float_opt "
+zax_opt=" $zax_opt $float_opt $context_sensitivity_opt "
cmd="$ZAFL_TM_ENV $PSZ $input_binary $output_zafl_binary $ida_or_rida_opt -c move_globals=on -c zax=on -o move_globals:--elftables-only $stars_opt $zax_opt $verbose_opt $options $other_args"
if [ ! -z "$ZAFL_TM_ENV" ]; then
diff --git a/libzafl b/libzafl
index 6f6f7350dc73d7305b58353b638be10148667015..97b561cb943dd9a15b3388cb19bd457eaf98fa37 160000
--- a/libzafl
+++ b/libzafl
@@ -1 +1 @@
-Subproject commit 6f6f7350dc73d7305b58353b638be10148667015
+Subproject commit 97b561cb943dd9a15b3388cb19bd457eaf98fa37