Add fixed map option to block-level instrumentation

Add tests

afl_transforms/tools/zax/zuntracer.cpp

@@ -21,6 +21,33 @@ void ZUntracer_t::afl_instrument_bb(Instruction_t *p_inst, const bool p_redZoneH
 {
 	assert(p_inst);
 
+	const auto trace_map_fixed_addr       = getenv("ZAFL_TRACE_MAP_FIXED_ADDRESS");
+	const auto do_fixed_addr_optimization = (trace_map_fixed_addr!=nullptr);
+
+	if (do_fixed_addr_optimization)
+		_afl_instrument_bb_fixed(p_inst, trace_map_fixed_addr);
+	else
+		_afl_instrument_bb(p_inst, p_redZoneHint);
+}
+
+void ZUntracer_t::_afl_instrument_bb_fixed(Instruction_t *p_inst, char* p_tracemap_addr)
+{
+	// 1st instruction in block record is the new entry point of the block (p_inst)
+	// 2nd instruction in block record is where the instruction at the old entry point is now at (orig)
+	const auto blockid = get_blockid();
+	BBRecord_t block_record;
+	block_record.push_back(p_inst);
+
+	// e.g.: mov BYTE [ 0x10000 + blockid ], 0x1
+	const auto s = string("mov BYTE [") + p_tracemap_addr + "+" + to_string(blockid) + "], 0x1";
+	const auto orig = insertAssemblyBefore(p_inst, s);
+	block_record.push_back(orig);
+
+	m_modifiedBlocks[blockid] = block_record;
+}
+
+void ZUntracer_t::_afl_instrument_bb(Instruction_t *p_inst, const bool p_redZoneHint)
+{
 	/*
 	Original afl instrumentation:
 	        block_id = <random>;

afl_transforms/tools/zax/zuntracer.hpp

@@ -11,20 +11,24 @@ namespace Zafl
 // Block-level instrumentation for Untracer
 class ZUntracer_t : public Zax_t
 {
-public:
-	ZUntracer_t() = delete;
-	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, bool p_verbose=false);
-	virtual ~ZUntracer_t() {};
-	virtual int execute();
-
-protected:
-	virtual zafl_blockid_t get_blockid(const unsigned p_maxid = 0xFFFF);
-	virtual void afl_instrument_bb(Instruction_t *p_inst, const bool p_hasLeafAnnotation, const bool p_collafl_optimization=false);
-	virtual set<libIRDB::BasicBlock_t*> getBlocksToInstrument(libIRDB::ControlFlowGraph_t &cfg);
-
-private:
-	zafl_blockid_t   m_blockid;
+	public:
+		ZUntracer_t() = delete;
+		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, bool p_verbose=false);
+		virtual ~ZUntracer_t() {};
+		virtual int execute();
+
+	protected:
+		virtual zafl_blockid_t get_blockid(const unsigned p_maxid = 0xFFFF);
+		virtual void afl_instrument_bb(Instruction_t *p_inst, const bool p_hasLeafAnnotation, const bool p_collafl_optimization=false);
+		virtual set<libIRDB::BasicBlock_t*> getBlocksToInstrument(libIRDB::ControlFlowGraph_t &cfg);
+
+	private:
+		void _afl_instrument_bb_fixed(Instruction_t *p_inst, char* p_tracemap_addr);
+		void _afl_instrument_bb(Instruction_t *p_inst, const bool p_redZoneHint);
+
+	private:
+		zafl_blockid_t   m_blockid;
 };
 
 } 

bin/zafl.sh

@@ -30,6 +30,8 @@ usage()
 	echo "     -c, --enable-breakup-critical-edges    Breakup critical edges"
 	echo "     -C, --disable-breakup-critical-edges   Do not breakup critical edges"
 	echo "     -f, --fork-server-only                 Fork server only"
+	echo "     -m, --enable-fixed-map                 Use fixed address for tracing map"
+	echo "     -M, --disable-fixed-map                Disable fixed address tracing map (default)"
 	echo 
 }
 
@@ -42,6 +44,10 @@ me=$(whoami)
 tmp_dir=/tmp/${me}/$$
 mkdir -p $tmp_dir
 
+# by default, use fixed address for map
+trace_map_address="0x10000"
+ZAFL_TM_ENV=""
+
 cleanup()
 {
 	if [ ! -z "$tmp_dir" ]; then
@@ -143,6 +149,20 @@ parse_args()
 				log_warning "Fork Server Only mode: no block-level instrumentation will be performed"
 				shift
 				;;
+			-m | --enable-fixed-map)
+				shift
+				case $1 in
+					0x*)
+						trace_map_address="$1"
+						shift
+					;;
+				esac
+				ZAFL_TM_ENV="ZAFL_TRACE_MAP_FIXED_ADDRESS=$trace_map_address"
+				;;
+			-M | --disable-fixed-map)
+				ZAFL_TM_ENV=""
+				shift
+				;;
 			-*|--*=) # unsupported flags
 				echo "Error: Unsupported flag $1" >&2
 				exit 1
@@ -249,7 +269,16 @@ fi
 log_msg "Transforming input binary $input_binary into $output_zafl_binary"
 
 #cmd="$PSZ $input_binary $output_zafl_binary $ida_or_rida_opt -c move_globals=on -c zax=on -o move_globals:--elftables-only -o zipr:--traceplacement:on $stars_opt $zax_opt $verbose_opt $options $other_args"
-cmd="$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"
+
+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
+	log_msg "Trace map will be expected at fixed address"
+	if [ -z $ZAFL_TRACE_MAP_FIXED_ADDRESS ]; then
+		log_warning "When running afl-fuzz, make sure that the environment variable ZAFL_TRACE_MAP_FIXED_ADDRESS is exported and set properly to (otherwise the instrumented binary will crash):"
+		log_warning "export $ZAFL_TM_ENV"
+	fi
+fi
 
 log_msg "Issuing command: $cmd"
 eval $cmd

test/ls.zuntracer/test_ls.sh

-AFL_TIMEOUT=30
+AFL_TIMEOUT=15
 session=/tmp/tmp.ls.$$
 TMP_FILE_1="${session}/ls.tmp.$$"
 
@@ -34,7 +34,7 @@ build_zuntracer()
 {
 	ls_zafl=$1
 	shift
-	zafl.sh `which ls` $ls_zafl --untracer --tempdir analysis.${ls_zafl}
+	zafl.sh `which ls` $ls_zafl --untracer --tempdir analysis.${ls_zafl} $*
 	if [ ! $? -eq 0 ]; then
 		log_error "$ls_zafl: unable to generate zafl version"	
 	else
@@ -102,18 +102,25 @@ build_zuntracer ls.untracer.critical_edge -c
 test_zuntracer ./ls.untracer.critical_edge -lt
 
 # zuntracer, do break critical edges, optimize graph
-build_zuntracer ls.untracer.critical_edge.graph -c -g
+build_zuntracer ls.untracer.critical_edge.graph -c -g -M
 test_zuntracer ./ls.untracer.critical_edge.graph -lt
 
 log_message "Fuzz zuntracer (basic block coverage) for $AFL_TIMEOUT seconds"
 fuzz_with_zafl ./ls.untracer.no_critical_edge -lt
 
 log_message "Fuzz zuntracer (break critical edges) for $AFL_TIMEOUT seconds"
-fuzz_with_zafl ./ls.untracer.critical_edge -lt
+fuzz_with_zafl ./ls.untracer.critical_edge -lt 
 
 log_message "Fuzz zuntracer (break critical edges + graph optimization) for $AFL_TIMEOUT seconds"
 fuzz_with_zafl ./ls.untracer.critical_edge.graph -lt
 
+build_zuntracer ls.untracer.fixed -m 
+test_zuntracer ./ls.untracer.fixed -lt
+build_zuntracer ls.untracer.fixed.0x10000 -m 0x10000
+test_zuntracer ./ls.untracer.fixed.0x10000 -lt
+
+fuzz_with_zafl ./ls.untracer.fixed.0x10000 -lt
+
 log_success "all tests passed: zafl/zuntracer instrumentation operational on ls"
 
 cleanup