Add support for return target set threshold

Former-commit-id: 3b75e01403bc75aaf545596fba0b34245d2f19f9

tools/simple_cdi/scdi_driver.cpp

@@ -35,7 +35,7 @@ using namespace libIRDB;
 
 void usage(char* name)
 {
-	cerr<<"Usage: "<<name<<" <variant_id>\n"; 
+	cerr<<"Usage: "<<name<<" <variant_id> --threshold <return_set_threshold> (default=1)\n"; 
 }
 
 int main(int argc, char **argv)
@@ -48,6 +48,9 @@ int main(int argc, char **argv)
 
         string programName(argv[0]);
         int variantID = atoi(argv[1]);
+	int threshold = 1;
+
+	// FIXME: implement getting threshold value from arguments
 
         VariantID_t *pidp=NULL;
 
@@ -74,8 +77,7 @@ int main(int argc, char **argv)
 
                 try
                 {
-			SimpleCDI_Instrument scdii(firp);
-
+			SimpleCDI_Instrument scdii(firp, threshold);
 
 			int success=scdii.execute();
 

tools/simple_cdi/scdi_instr.cpp

@@ -230,11 +230,7 @@ bool SimpleCDI_Instrument::add_scdi_instrumentation(Instruction_t* insn)
 		cout <<"["<<string(d.CompleteInstr)<<"] [" << string(d.Instruction.Mnemonic)<< "] IBTargets size: " << ibts->size() << " analysis_status: " << ibts->GetAnalysisStatus() << endl;
 	}
 
-	// only handle ret if complete && ib target size == 1
-
-	// @todo: for debugging, hlt
-	// @todo: can handle size == 2 with just one cmp
-	if (string(d.Instruction.Mnemonic) == string("ret ")) 
+	if (is_return(insn))
 	{
 		// instrumentation must be coordinated with needs_scdi_instrumentation()
 		if (ibts && ibts->IsComplete() && ibts->size() == 1)
@@ -266,6 +262,7 @@ bool SimpleCDI_Instrument::add_scdi_instrumentation(Instruction_t* insn)
 	}
 
 	assert(strstr("ret ", d.Instruction.Mnemonic)==NULL);
+	assert(strstr("retn ", d.Instruction.Mnemonic)==NULL);
 	
 	// pre-instrument
 	// push reg
@@ -306,17 +303,27 @@ bool SimpleCDI_Instrument::add_scdi_instrumentation(Instruction_t* insn)
 	return success;
 }
 
-
-bool SimpleCDI_Instrument::needs_scdi_instrumentation(Instruction_t* insn)
+bool SimpleCDI_Instrument::is_return(Instruction_t* insn)
 {
-	DISASM d;
-	insn->Disassemble(d);
-
-	bool isReturn = false;
+	if (insn) 
+	{
+		DISASM d;
+		insn->Disassemble(d);
+		return string(d.Instruction.Mnemonic) == string("ret "); 
+
+		// FIXME: handle retn immd, but this means the instrumentation should pop/lea immd
+	/*	return (string(d.Instruction.Mnemonic) == string("ret ") ||
+		    string(d.Instruction.Mnemonic) == string("retn "));
+	*/
+	}
 
-	if (string(d.Instruction.Mnemonic) == string("ret "))
-		isReturn = true;
+	return false;
+}
 
+// only complete returns need to be instrumented
+bool SimpleCDI_Instrument::needs_scdi_instrumentation(Instruction_t* insn, int target_size_threshold)
+{
+	const bool isReturn = is_return(insn);
 
 	if (isReturn)
 		num_returns++;
@@ -332,24 +339,19 @@ bool SimpleCDI_Instrument::needs_scdi_instrumentation(Instruction_t* insn)
 			num_complete_returns++;
 	}
 
-	if (string(d.Instruction.Mnemonic) == string("ret "))
+	if (isReturn)
 	{
-		if (ibts->IsComplete() && ibts->size() == 1)
-			return true;
+		if (ibts->IsComplete())
+		{
+			if (target_set_threshold < 0)
+				return true;
+			else 
+				return ibts->size() <= target_size_threshold;
+		}
 		else 
 			return false;
 	}
 
-/*
-	if (ibts->IsComplete() && ibts->size() <= 2)
-		return true;
-*/
-
-/*
-	if(ibts->IsComplete() && ibts->size() >0 )
-		return true;
-*/
-
 	return false;
 }
 
@@ -363,9 +365,8 @@ bool SimpleCDI_Instrument::convert_ibs()
                 ++it)
         {
 		Instruction_t* insn=*it;
-		if(needs_scdi_instrumentation(insn))
+		if(needs_scdi_instrumentation(insn, target_set_threshold))
 			success = success && add_scdi_instrumentation(insn);
-			
 	}
 
 	return success;
@@ -374,6 +375,7 @@ bool SimpleCDI_Instrument::convert_ibs()
 void SimpleCDI_Instrument::display_stats(std::ostream &out)
 {
 	float fraction = NAN;
+	out << "# ATTRIBUTE target_set_threshold=" << dec << target_set_threshold << endl;
 	out << "# ATTRIBUTE complete_ibts=" << dec << num_complete_ibts << endl;
 	out << "# ATTRIBUTE num_returns=" << num_returns << endl;
 	if (num_complete_returns>0)

tools/simple_cdi/scdi_instr.hpp

@@ -28,24 +28,27 @@
 class SimpleCDI_Instrument
 {
 	public:
-		SimpleCDI_Instrument(libIRDB::FileIR_t *the_firp) : firp(the_firp), 
+		SimpleCDI_Instrument(libIRDB::FileIR_t *the_firp, int p_target_set_threshold=1) : firp(the_firp), 
 			single_target_set_jumps(0),
 			single_target_set_returns(0),
 			num_complete_ibts(0),
 			num_returns(0),
-			num_complete_returns(0) {}
+			num_complete_returns(0),
+			target_set_threshold(p_target_set_threshold) {}
 		bool execute();
 
 	private:
 
-
+		bool is_return(libIRDB::Instruction_t* insn);
 		bool add_scdi_instrumentation(libIRDB::Instruction_t* insn);
-		bool needs_scdi_instrumentation(libIRDB::Instruction_t* insn);
+		bool needs_scdi_instrumentation(libIRDB::Instruction_t* insn, int p_target_set_threshold);
 		bool convert_ibs();
 		void display_stats(std::ostream &out);
 	
 		libIRDB::FileIR_t* firp;
 
+		int target_set_threshold;
+
 		int single_target_set_jumps;
 		int single_target_set_returns;
 		int num_complete_ibts;