Keep track of STARS SAFE_FUNC in IRDB

scfi_driver.cpp

@@ -39,9 +39,10 @@ void usage(char* name)
 "		[--color|--no-color]  \n"
 "		[--protect-jumps|--no-protect-jumps]  \n"
 "		[--protect-rets|--no-protect-rets] \n"
+"		[--protect-safefn|--no-protect-safefn]  \n"
 "		[ --common-slow-path | --no-common-slow-path ] \n"
 " \n"
-"default: --no-color --protect-jumps --protect-rets --common-slow-path\n"; 
+"default: --no-color --protect-jumps --protect-rets --no-protect-safefn --common-slow-path\n"; 
 }
 
 int main(int argc, char **argv)
@@ -62,6 +63,7 @@ int main(int argc, char **argv)
 	bool do_common_slow_path=true;
 	bool do_jumps=true;
 	bool do_rets=true;
+	bool do_safefn=false;
 	for(int  i=2;i<argc;i++)
 	{
 		if(string(argv[i])=="--color")
@@ -94,6 +96,16 @@ int main(int argc, char **argv)
 			cout<<"Not protecting returns..."<<endl;
 			do_rets=false;
 		}
+		else if(string(argv[i])=="--protect-safefn")
+		{
+			cout<<"protecting safe functions..."<<endl;
+			do_safefn=true;
+		}
+		else if(string(argv[i])=="--no-protect-safefn")
+		{
+			cout<<"Not protecting safe functions..."<<endl;
+			do_safefn=false;
+		}
 		else if(string(argv[i])=="--common-slow-path")
 		{
 			cout<<"Using common slow path..."<<endl;
@@ -140,7 +152,7 @@ int main(int argc, char **argv)
 
                 try
                 {
-			SCFI_Instrument scfii(firp, do_coloring, do_common_slow_path, do_jumps, do_rets);
+			SCFI_Instrument scfii(firp, do_coloring, do_common_slow_path, do_jumps, do_rets, do_safefn);
 
 
 			int success=scfii.execute();

scfi_instr.cpp

@@ -217,7 +217,10 @@ Relocation_t* SCFI_Instrument::FindRelocation(Instruction_t* insn, string type)
         return NULL;
 }
 
-
+bool SCFI_Instrument::isSafeFunction(Instruction_t* insn)
+{
+	return (insn && insn->GetFunction() && insn->GetFunction()->IsSafe());
+}
 
 
 Relocation_t* SCFI_Instrument::create_reloc(Instruction_t* insn)
@@ -602,6 +605,8 @@ bool SCFI_Instrument::instrument_jumps()
 	int cfi_branch_call_complete=0;
 	int cfi_branch_ret_checks=0;
 	int cfi_branch_ret_complete=0;
+	int	cfi_safefn_jmp_skipped=0;
+	int	cfi_safefn_ret_skipped=0;
 	int ibt_complete=0;
 	double cfi_branch_jmp_complete_ratio = NAN;
 	double cfi_branch_ret_complete_ratio = NAN;
@@ -625,6 +630,8 @@ bool SCFI_Instrument::instrument_jumps()
 		if(FindRelocation(insn,"cf::safe"))
 			continue;
 
+		bool safefn = isSafeFunction(insn);
+
 		DISASM d;
 		insn->Disassemble(d);
 
@@ -634,13 +641,20 @@ bool SCFI_Instrument::instrument_jumps()
 			case  JmpType:
 				if((d.Argument1.ArgType&MEMORY_TYPE)==MEMORY_TYPE)
 				{
-					cfi_checks++;
-					cfi_branch_jmp_checks++;
 					if (insn->GetIBTargets() && insn->GetIBTargets()->IsComplete())
 					{
 						cfi_branch_jmp_complete++;
 						jmps[insn->GetIBTargets()->size()]++;
 					}
+
+					if (!do_safefn && safefn)
+					{
+						cfi_safefn_jmp_skipped++;
+						continue;
+					}
+
+					cfi_checks++;
+					cfi_branch_jmp_checks++;
 					AddJumpCFI(insn);
 				}
 				break;
@@ -656,14 +670,22 @@ bool SCFI_Instrument::instrument_jumps()
 					cfi_checks++;
 				}
 				break;
+
 			case  RetType: 
-				cfi_branch_ret_checks++;
 				if (insn->GetIBTargets() && insn->GetIBTargets()->IsComplete())
 				{
 					cfi_branch_ret_complete++;
 					rets[insn->GetIBTargets()->size()]++;
 				}
+
+				if (!do_safefn && safefn)
+				{
+					cfi_safefn_ret_skipped++;
+					continue;
+				}
+
 				cfi_checks++;
+				cfi_branch_ret_checks++;
 				AddReturnCFI(insn);
 				break;
 
@@ -673,7 +695,7 @@ bool SCFI_Instrument::instrument_jumps()
 	}
 	
 	cout<<"# ATTRIBUTE cfi_jmp_checks="<<std::dec<<cfi_branch_jmp_checks<<endl;
-	cout<<"# ATTRIBUTE cfi_jmp_complete="<<std::dec<<cfi_branch_jmp_complete<<endl;
+	cout<<"# ATTRIBUTE cfi_jmp_complete="<<cfi_branch_jmp_complete<<endl;
 
 	display_histogram(cout, "cfi_jmp_complete_histogram", jmps);
 
@@ -705,6 +727,9 @@ bool SCFI_Instrument::instrument_jumps()
 	cout << "# ATTRIBUTE cfi_ret_complete_ratio=" << cfi_branch_ret_complete_ratio << endl;
 	cout << "# ATTRIBUTE cfi_complete_ratio=" << cfi_branch_ret_complete_ratio << endl;
 
+	cout<<"# ATTRIBUTE cfi_safefn_jmp_skipped="<<cfi_safefn_jmp_skipped<<endl;
+	cout<<"# ATTRIBUTE cfi_safefn_ret_skipped="<<cfi_safefn_ret_skipped<<endl;
+
 	return true;
 }
 

scfi_instr.hpp

@@ -33,12 +33,14 @@ class SCFI_Instrument
 				bool p_do_coloring=true,
 				bool p_do_common_slow_path=true,
 				bool p_do_jumps=true,
-				bool p_do_rets=true) 
+				bool p_do_rets=true,
+				bool p_do_safefn=true) 
 			: firp(the_firp), 
 			  do_coloring(p_do_coloring), 
 			  do_common_slow_path(p_do_common_slow_path), 
 			  do_jumps(p_do_jumps), 
 			  do_rets(p_do_rets), 
+			  do_safefn(p_do_safefn), 
 			  color_map(NULL) {}
 		bool execute();
 
@@ -52,6 +54,7 @@ class SCFI_Instrument
 		// helper
 		libIRDB::Relocation_t* create_reloc(libIRDB::Instruction_t* insn);
 		libIRDB::Relocation_t* FindRelocation(libIRDB::Instruction_t* insn, std::string type);
+		bool isSafeFunction(libIRDB::Instruction_t* insn);
 
 		// add instrumentation
 		bool add_scfi_instrumentation(libIRDB::Instruction_t* insn);
@@ -76,6 +79,7 @@ class SCFI_Instrument
 		bool do_common_slow_path;
 		bool do_jumps;
 		bool do_rets;
+		bool do_safefn;
 		ColoredInstructionNonces_t *color_map;