scfi_instr.cpp

/*
 * Copyright (c) 2014-2015 - Zephyr Software LLC
 *
 * This file may be used and modified for non-commercial purposes as long as
 * all copyright, permission, and nonwarranty notices are preserved.
 * Redistribution is prohibited without prior written consent from Zephyr
 * Software.
 *
 * Please contact the authors for restrictions applying to commercial use.
 *
 * THIS SOURCE IS PROVIDED "AS IS" AND WITHOUT ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTIES OF
 * MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 *
 * Author: Zephyr Software
 * e-mail: jwd@zephyr-software.com
 * URL   : http://www.zephyr-software.com/
 *
 */


#include "utils.hpp"
#include "scfi_instr.hpp"
#include "Rewrite_Utility.hpp"
#include "color_map.hpp"
#include <stdlib.h>
#include <memory>
#include <math.h>
#include <exeio.h>
#include <elf.h>
#include "elfio/elfio.hpp"
#include "elfio/elfio_dump.hpp"




using namespace std;
using namespace libIRDB;

virtual_offset_t getAvailableAddress(FileIR_t *p_virp)
{

        static int counter = -16;
        counter += 16;
        return 0xf0020000 + counter;
}




static Instruction_t* registerCallbackHandler64(FileIR_t* firp, Instruction_t *p_orig, string p_callbackHandler, int p_numArgs)
{

        Instruction_t *instr;
        Instruction_t *first;
        char tmpbuf[1024];

        // save flags and 16 registers (136 bytes)
        // call pushes 8 bytes
        // Total: 8 * 18 = 144
        first = instr = addNewAssembly(firp,NULL, "push rsp");
        instr = addNewAssembly(firp,instr, "push rbp");
        instr = addNewAssembly(firp,instr, "push rdi");
        instr = addNewAssembly(firp,instr, "push rsi");
        instr = addNewAssembly(firp,instr, "push rdx");
        instr = addNewAssembly(firp,instr, "push rcx");
        instr = addNewAssembly(firp,instr, "push rbx");
        instr = addNewAssembly(firp,instr, "push rax");
        instr = addNewAssembly(firp,instr, "push r8");
        instr = addNewAssembly(firp,instr, "push r9");
        instr = addNewAssembly(firp,instr, "push r10");
        instr = addNewAssembly(firp,instr, "push r11");
        instr = addNewAssembly(firp,instr, "push r12");
        instr = addNewAssembly(firp,instr, "push r13");
        instr = addNewAssembly(firp,instr, "push r14");
        instr = addNewAssembly(firp,instr, "push r15");
        instr = addNewAssembly(firp,instr, "pushf");

        // handle the arguments (if any): rdi, rsi, rdx, rcx, r8, r9
        // first arg starts at byte +144
        instr = addNewAssembly(firp,instr, "mov rdi, rsp");

	if (p_numArgs >= 1)
                instr = addNewAssembly(firp,instr,  "mov rsi, [rsp+144]");
        if (p_numArgs >= 2)
                instr = addNewAssembly(firp,instr,  "mov rdx, [rsp+152]");
        if (p_numArgs >= 3)
                instr = addNewAssembly(firp,instr,  "mov rcx, [rsp+160]");
        if (p_numArgs >= 4)
                instr = addNewAssembly(firp,instr,  "mov r8, [rsp+168]");
        if (p_numArgs > 4)
                assert(0); // only handle up to 5 args

        // pin the instruction that follows the callback handler
        Instruction_t* postCallback = allocateNewInstruction(firp, BaseObj_t::NOT_IN_DATABASE, NULL);
        virtual_offset_t postCallbackReturn = getAvailableAddress(firp);
        postCallback->GetAddress()->SetVirtualOffset(postCallbackReturn);

        // push the address to return to once the callback handler is invoked
        sprintf(tmpbuf,"mov rax, 0x%x", (unsigned int)postCallbackReturn);
        instr = addNewAssembly(firp,instr, tmpbuf);

        instr = addNewAssembly(firp,instr, "push rax");

        // use a nop instruction for the actual callback
        instr = addNewAssembly(firp,instr, "nop");
        instr->SetComment(" -- callback: " + p_callbackHandler);
        instr->SetCallback(p_callbackHandler);
        instr->SetFallthrough(postCallback);


        // need to make sure the post callback address is pinned
        // (so that ILR and other transforms do not relocate it)
        AddressID_t *indTarg = new AddressID_t();
        firp->GetAddresses().insert(indTarg);
        indTarg->SetVirtualOffset(postCallback->GetAddress()->GetVirtualOffset());
        indTarg->SetFileID(BaseObj_t::NOT_IN_DATABASE); // SPRI global namespace
        postCallback->SetIndirectBranchTargetAddress(indTarg);

        // restore registers
        firp->RegisterAssembly(postCallback, "popf");


        instr = addNewAssembly(firp,postCallback, "pop r15");
        instr = addNewAssembly(firp,instr, "pop r14");
        instr = addNewAssembly(firp,instr, "pop r13");
        instr = addNewAssembly(firp,instr, "pop r12");
        instr = addNewAssembly(firp,instr, "pop r11");
        instr = addNewAssembly(firp,instr, "pop r10");
        instr = addNewAssembly(firp,instr, "pop r9");
        instr = addNewAssembly(firp,instr, "pop r8");
        instr = addNewAssembly(firp,instr, "pop rax");
        instr = addNewAssembly(firp,instr, "pop rbx");
        instr = addNewAssembly(firp,instr, "pop rcx");
        instr = addNewAssembly(firp,instr, "pop rdx");
        instr = addNewAssembly(firp,instr, "pop rsi");
        instr = addNewAssembly(firp,instr, "pop rdi");
        instr = addNewAssembly(firp,instr, "pop rbp");
        instr = addNewAssembly(firp,instr, "lea rsp, [rsp+8]");

        instr = addNewAssembly(firp,instr, "ret");

        // return first instruction in the callback handler chain
        return first;

}


// x86-64
// 20140421
static void ConvertCallToCallbackHandler64(FileIR_t* firp, Instruction_t *p_orig, string p_callbackHandler, int p_numArgs)
{
	static std::map<std::string, Instruction_t*> m_handlerMap;
        // nb: if first time, register and cache callback handler sequence
        if (m_handlerMap.count(p_callbackHandler) == 0)
        {
                m_handlerMap[p_callbackHandler] = registerCallbackHandler64(firp,p_orig, p_callbackHandler, p_numArgs);
        }

        if (p_orig)
                p_orig->SetTarget(m_handlerMap[p_callbackHandler]);
}


static Instruction_t* addCallbackHandlerSequence
	(
	  FileIR_t* firp,
  	  Instruction_t *p_orig, 
	  bool before,
	  std::string p_detector
	)
{

	if(before)
		insertAssemblyBefore(firp,p_orig,"lea rsp, [rsp-128]");
	else
		assert(0); // add handling  for inserting lea after given insn

        p_orig->SetComment("callback: " + p_detector);


        Instruction_t* call =insertAssemblyAfter(firp,p_orig,"call 0");

        ConvertCallToCallbackHandler64(firp, call, p_detector, 0); // no args for now

	insertAssemblyAfter(firp,call,"lea rsp, [rsp + 128 + 0]"); // no args for nwo 

        return p_orig;
}

Relocation_t* SCFI_Instrument::FindRelocation(Instruction_t* insn, string type)
{
        RelocationSet_t::iterator rit;
        for( rit=insn->GetRelocations().begin(); rit!=insn->GetRelocations().end(); ++rit)
        {
                Relocation_t& reloc=*(*rit);
                if(reloc.GetType()==type)
                {
                        return &reloc;
                }
        }
        return NULL;
}

bool SCFI_Instrument::isSafeFunction(Instruction_t* insn)
{
	return (insn && insn->GetFunction() && insn->GetFunction()->IsSafe());
}

bool SCFI_Instrument::isCallToSafeFunction(Instruction_t* insn)
{
	if (insn && insn->GetTarget() && insn->GetTarget()->GetFunction())
	{
		if(getenv("SCFI_VERBOSE")!=NULL)
		{
			if (insn->GetTarget()->GetFunction()->IsSafe())
			{
				cout << "Function " << insn->GetTarget()->GetFunction()->GetName() << " is deemed safe" << endl;
			}
		}

		return insn->GetTarget()->GetFunction()->IsSafe();
	}

	return false;
}

Relocation_t* SCFI_Instrument::create_reloc(Instruction_t* insn)
{
        Relocation_t* reloc=new Relocation_t;
        insn->GetRelocations().insert(reloc);
        firp->GetRelocations().insert(reloc);

	return reloc;
}

bool SCFI_Instrument::add_scfi_instrumentation(Instruction_t* insn)
{
	bool success=true;

	if(getenv("SCFI_VERBOSE")!=NULL)
		;


	return success;
}

bool SCFI_Instrument::needs_scfi_instrumentation(Instruction_t* insn)
{
assert(0);
	return false;

}

unsigned int SCFI_Instrument::GetNonceOffset(Instruction_t* insn)
{
	if(color_map)
	{
		assert(insn->GetIBTargets());
		return (color_map->GetColorOfIB(insn).GetPosition()+1) * GetNonceSize(insn);
	}
	return GetNonceSize(insn);
}

NonceValueType_t SCFI_Instrument::GetNonce(Instruction_t* insn)
{
	/* in time we look up the nonce category for this insn */
	/* for now, it's just f4 as the nonce */
	if(color_map)
	{
		assert(insn->GetIBTargets());
		return color_map->GetColorOfIB(insn).GetNonceValue();
	}
	return 0xf4;
}

unsigned int SCFI_Instrument::GetNonceSize(Instruction_t* insn)
{
	/* in time we look up the nonce size for this insn */
	/* for now, it's just f4 as the nonce */
	return 1;
}

bool SCFI_Instrument::mark_targets() 
{
	int targets=0, ind_targets=0;
	for(InstructionSet_t::iterator it=firp->GetInstructions().begin();
		it!=firp->GetInstructions().end();
		++it)
	{
		targets++;
		Instruction_t* insn=*it;
		if(insn->GetIndirectBranchTargetAddress())
		{

			// make sure there are no fallthroughs to nonces.
			for(InstructionSet_t::iterator pred_it=preds[insn].begin(); pred_it!=preds[insn].end(); ++pred_it)
			{
				Instruction_t* the_pred=*pred_it;
				if(the_pred->GetFallthrough()==insn)
				{
					Instruction_t* jmp=addNewAssembly(firp,NULL, "jmp 0x0");
					the_pred->SetFallthrough(jmp);
					jmp->SetTarget(insn);
				}
			}

			ind_targets++;
			string type;
			if(do_coloring)
			{
				ColoredSlotValues_t v=color_map->GetColorsOfIBT(insn);
				int size=1;
				for(int i=0;i<v.size();i++)
				{
					if(!v[i].IsValid())
						continue;
					int position=v[i].GetPosition();

					// convert the colored "slot" into a position in the code.
					position++;
					position*=size;
					position = - position;

					// cfi_nonce=(pos=-1,nv=0x33,sz=1)
					NonceValueType_t noncevalue=v[i].GetNonceValue();
					type=string("cfi_nonce=(pos=") +  to_string(position) + ",nv="
						+ to_string(noncevalue) + ",sz="+ to_string(size)+ ")";
					Relocation_t* reloc=create_reloc(insn);
					reloc->SetOffset(-position*size);
					reloc->SetType(type);
					cout<<"Created reloc='"+type+"' for "<<std::dec<<insn->GetBaseID()<<":"<<insn->getDisassembly()<<endl;
				}
			}
			else
			{
				// cfi_nonce=f4.
				type="cfi_nonce=";
				type+=to_string(GetNonce(insn));

				Relocation_t* reloc=create_reloc(insn);
				reloc->SetOffset(-GetNonceOffset(insn));
				reloc->SetType(type);
				cout<<"Found nonce="+type+"  for "<<std::dec<<insn->GetBaseID()<<":"<<insn->getDisassembly()<<endl;
			}
		}
	}
	cout<<"# ATTRIBUTE ind_targets_found="<<std::dec<<ind_targets<<endl;
	cout<<"# ATTRIBUTE targets_found="<<std::dec<<targets<<endl;
	return true;
}

/*
 * targ_change_to_push - use the mode in the insnp to create a new instruction that is a push instruction.
 */
static string change_to_push(Instruction_t *insn)
{
	string newbits=insn->GetDataBits();

	DISASM d; 
	insn->Disassemble(d);

	int opcode_offset=0;


	// FIXME: assumes REX is only prefix on jmp insn.  
	// does not assume rex exists.
	if(d.Prefix.REX.state == InUsePrefix)
		opcode_offset=1;

	unsigned char modregrm = (newbits[1+opcode_offset]);
	modregrm &= 0xc7;
	modregrm |= 0x30;
	newbits[0+opcode_offset] = 0xFF;
	newbits[1+opcode_offset] = modregrm;

        return newbits;
}


void mov_reloc(Instruction_t* from, Instruction_t* to, string type )
{
	for(
		/* start */
		RelocationSet_t::iterator it=from->GetRelocations().begin();

		/* continue */
		it!=from->GetRelocations().end();
		
		/* increment */
		/* empty */
	   )
	{
		Relocation_t* reloc=*it;

		if(reloc->GetType()==type)
		{
			to->GetRelocations().insert(reloc);	
	
			// odd standards-conforming way to delete object while iterating.
			from->GetRelocations().erase(it++);	
		}
		else
		{
			it++;
		}
	}
		
}


static void move_relocs(Instruction_t* from, Instruction_t* to)
{
	for(auto it=from->GetRelocations().begin(); it!=from->GetRelocations().end(); ) 
	{
		auto current=it++;
		Relocation_t* reloc=*current;
		if(reloc->GetType()=="fix_call_fallthrough")
		{
			// don't move it.
		}
		else
		{
			to->GetRelocations().insert(reloc);
			from->GetRelocations().erase(current);
		}
	}
}

void SCFI_Instrument::AddJumpCFI(Instruction_t* insn)
{
	assert(do_rets);
	ColoredSlotValue_t v2;
	if(insn->GetIBTargets() && color_map)
		v2=color_map->GetColorOfIB(insn);
	ColoredSlotValue_t *v=&v2;
	string reg="ecx";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		reg="rcx";	// 64-bit reg.

	string pushbits=change_to_push(insn);
	cout<<"Converting ' "<<insn->getDisassembly()<<"' to '";
	Instruction_t* after=insertDataBitsBefore(firp,insn,pushbits); 
	move_relocs(after,insn);
#ifdef CGC
	// insert the pop/checking code.
	cout<<insn->getDisassembly()<<"+jmp slowpath'"<<endl;

	string jmpBits=getJumpDataBits();
        after->SetDataBits(jmpBits);
        after->SetComment(insn->getDisassembly()+" ; scfi");
	assert(do_common_slow_path); /* fixme:  this defaults to the slow_cfi path.  need to color accordingly */
	createNewRelocation(firp,after,"slow_cfi_path",0);
	after->SetFallthrough(NULL);
	after->SetTarget(after);
	after->SetIBTargets(NULL); // lose information about ib targets.
	insn->SetIBTargets(NULL); // lose information about ib targets.
	return;
#else
	
	after->SetDataBits(getRetDataBits());
	cout <<insn->getDisassembly()<<" + ret' "<<endl ;

	// move any pc-rel relocation bits to the push, which will access memory now 
	mov_reloc(after,insn,"pcrel");
	after->SetIBTargets(insn->GetIBTargets());
	insn->SetIBTargets(NULL);

	AddReturnCFI(after,v);
	// cout<<"Warning, JUMPS not CFI's yet"<<endl;
	return;
#endif
}


void SCFI_Instrument::AddCallCFIWithExeNonce(Instruction_t* insn)
{
	// make a stub to call

	// the stub is:
	//	push [target]
	// 	jmp slow
	string pushbits=change_to_push(insn);
	Instruction_t* stub=addNewDatabits(firp,NULL,pushbits);
	stub->SetComment(insn->GetComment()+" cfi stub");
	

	string jmpBits=getJumpDataBits();
	Instruction_t* jmp=insertDataBitsAfter(firp, stub, jmpBits);

	assert(stub->GetFallthrough()==jmp);

	// create a reloc so the stub goes to the slow path, eventually
	createNewRelocation(firp,jmp,"slow_cfi_path",0);
	jmp->SetFallthrough(NULL);
	jmp->SetTarget(jmp);	// looks like infinite loop, but will go to slow apth

	// convert the indirct call to a direct call to the stub.
	string call_bits=insn->GetDataBits();
	call_bits.resize(5);
	call_bits[0]=0xe8;
	insn->SetTarget(stub);
	insn->SetDataBits(call_bits);
	insn->SetComment("Direct call to cfi stub");
	insn->SetIBTargets(NULL);	// lose info about branch targets.

}

void SCFI_Instrument::AddExecutableNonce(Instruction_t* insn)
{
// this is now done by the nonce plugin's PlopDollopEntry routine.
//	insertDataBitsAfter(firp, insn, ExecutableNonceValue, NULL);
}


void SCFI_Instrument::AddReturnCFIForExeNonce(Instruction_t* insn, ColoredSlotValue_t *v)
{
	assert(!do_coloring);

	if(!do_rets)
		return;

	string reg="ecx";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		reg="r11";	// 64-bit reg.

	string rspreg="esp";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		rspreg="rsp";	// 64-bit reg.

	string worddec="dword";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		worddec="qword";	// 64-bit reg.

	
	DISASM d;
	insn->Disassemble(d);
	if(d.Argument1.ArgType!=NO_ARGUMENT)
	{
		unsigned int sp_adjust=d.Instruction.Immediat-firp->GetArchitectureBitWidth()/8;
		cout<<"Found relatively rare ret_with_pop insn: "<<d.CompleteInstr<<endl;
		char buf[30];
		sprintf(buf, "pop %s [%s+%d]", worddec.c_str(), rspreg.c_str(), sp_adjust);
		Instruction_t* newafter=insertAssemblyBefore(firp,insn,buf);

		if(sp_adjust>0)
		{
			sprintf(buf, "lea %s, [%s+%d]", rspreg.c_str(), rspreg.c_str(), sp_adjust);
		}

		// rewrite the "old" isntruction, as that's what insertAssemblyBefore returns
		insn=newafter;
	}
		
#ifdef CGC
#if 0
	ret -> jmp shared

	shared: 
		and [rsp], 0x7ffffffff
		mov reg <- [ rsp ] 
		cmp [reg], exe_nonce_value
		jne slow
		ret

#endif

	string jmpBits=getJumpDataBits();

#ifndef CGC
	assert(0); // not ported to non-cgc mode
#endif
	if(!ret_shared)
	{
        	string clamp_str="and "+worddec+"["+rspreg+"], 0x7fffffff";
		Instruction_t* tmp=NULL;
		ret_shared=
#ifdef CGC
        	tmp=addNewAssembly(firp,tmp,clamp_str); // FIXME???
		tmp=addNewAssembly(firp, tmp, "mov "+reg+", ["+rspreg+"]");
#else
		tmp=addNewAssembly(firp, "mov "+reg+", ["+rspreg+"]");
#endif
// fixme:  get value from ExecutableNonceString -- somewhat challening
		tmp=addNewAssembly(firp, tmp, "cmp byte ["+reg+"], 0x90");
		tmp=addNewAssembly(firp, tmp, "jne 0");
		createNewRelocation(firp,tmp,"slow_cfi_path",0);
		tmp->SetTarget(tmp);
		tmp=addNewAssembly(firp, tmp, "ret");

	}
	
        insn->SetDataBits(jmpBits);
	insn->SetTarget(ret_shared);

	return;
#else
	assert(0);
#endif
	
}

void SCFI_Instrument::AddReturnCFI(Instruction_t* insn, ColoredSlotValue_t *v)
{

	if(!do_rets)
		return;
	ColoredSlotValue_t v2; 
	if(v==NULL && color_map)
	{
		v2=color_map->GetColorOfIB(insn);
		v=&v2;
	}


	string reg="ecx";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		reg="r11";	// 64-bit reg.

	string rspreg="esp";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		rspreg="rsp";	// 64-bit reg.

	string worddec="dword";	// 32-bit reg 
	if(firp->GetArchitectureBitWidth()==64)
		worddec="qword";	// 64-bit reg.

	
	DISASM d;
	insn->Disassemble(d);
	if(d.Argument1.ArgType!=NO_ARGUMENT)
	{
		unsigned int sp_adjust=d.Instruction.Immediat-firp->GetArchitectureBitWidth()/8;
		cout<<"Found relatively rare ret_with_pop insn: "<<d.CompleteInstr<<endl;
		char buf[30];
		sprintf(buf, "pop %s [%s+%d]", worddec.c_str(), rspreg.c_str(), sp_adjust);
		Instruction_t* newafter=insertAssemblyBefore(firp,insn,buf);

		if(sp_adjust>0)
		{
			sprintf(buf, "lea %s, [%s+%d]", rspreg.c_str(), rspreg.c_str(), sp_adjust);
		}

		// rewrite the "old" isntruction, as that's what insertAssemblyBefore returns
		insn=newafter;
	}
		
	int size=1;
	int position=0;
	string slow_cfi_path_reloc_string;
	if(do_coloring && !do_common_slow_path)
	{
		slow_cfi_path_reloc_string="slow_cfi_path=(pos=-1,nv=244,sz=1)";
		if( v && v->IsValid())
		{
			slow_cfi_path_reloc_string="slow_cfi_path=(pos=-"+ to_string(v->GetPosition()+1) +",nv="
						  + to_string(v->GetNonceValue())+",sz="+ to_string(size) +")";
			size=v->GetPosition();
		}
	}
	else
	{
		slow_cfi_path_reloc_string="slow_cfi_path";
	}

	
	cout<<"Cal'd slow-path cfi reloc as: "<<slow_cfi_path_reloc_string<<endl;
// fixme:  would like to mark a slow path per nonce type using the variables calc'd above.
	
	

#ifdef CGC
	// insert the pop/checking code.
	Instruction_t* after=insn;

	string jmpBits=getJumpDataBits();
	
        after->SetDataBits(jmpBits);
        after->SetComment(insn->getDisassembly()+" ; scfi");
	createNewRelocation(firp,after,slow_cfi_path_reloc_string,0);
	after->SetFallthrough(NULL);
	after->SetTarget(after);
	return;
	
#else

	string decoration="";
	int nonce_size=GetNonceSize(insn);
	int nonce_offset=GetNonceOffset(insn);
	unsigned int nonce=GetNonce(insn);
	Instruction_t* jne=NULL, *tmp=NULL;
	

	// convert a return to:
	// 	pop ecx
	// 	cmp <nonce size> PTR [ecx-<nonce size>], Nonce
	// 	jne slow	; reloc such that strata/zipr can convert slow to new code
	//			; to handle places where nonce's can't be placed. 
	// 	jmp ecx

	switch(nonce_size)
	{
		case 1:
			decoration="byte ";
			break;
		case 2:	// handle later
		case 4: // handle later
		default:
			cerr<<"Cannot handle nonce of size "<<std::dec<<nonce_size<<endl;
			assert(0);
		
	}

	// insert the pop/checking code.
	insertAssemblyBefore(firp,insn,string("pop ")+reg);
	tmp=insertAssemblyAfter(firp,insn,string("cmp ")+decoration+
		" ["+reg+"-"+to_string(nonce_offset)+"], "+to_string(nonce));
    jne=tmp=insertAssemblyAfter(firp,tmp,"jne 0");

	// convert the ret instruction to a jmp ecx
	cout<<"Converting "<<dec<<tmp->GetFallthrough()->GetBaseID()<<":"<<tmp->GetFallthrough()->getDisassembly()<<"to jmp+reg"<<endl;
	setInstructionAssembly(firp,tmp->GetFallthrough(), string("jmp ")+reg, NULL,NULL);

	// set the jne's target to itself, and create a reloc that zipr/strata will have to resolve.
	jne->SetTarget(jne);	// needed so spri/spasm/irdb don't freak out about missing target for new insn.
	Relocation_t* reloc=create_reloc(jne);
	reloc->SetType(slow_cfi_path_reloc_string); 
	reloc->SetOffset(0);
	cout<<"Setting slow path for: "<<slow_cfi_path_reloc_string<<endl;

	return;
#endif
}

static void display_histogram(std::ostream& out, std::string attr_label, std::map<int,int> & p_map)
{
	if (p_map.size()) 
	{
		out<<"# ATTRIBUTE " << attr_label << "=";
		out<<"{ibt_size:count,";
		bool first_time=true;
		for (map<int,int>::iterator it = p_map.begin(); 
			it != p_map.end(); ++it)
		{
			if (!first_time)
				out << ",";
			out << it->first << ":" << it->second;
			first_time = false;
		}
		out<<"}"<<endl;
	}
}

bool SCFI_Instrument::is_plt_style_jmp(Instruction_t* insn) 
{
	DISASM d;
	insn->Disassemble(d);
	if((d.Argument1.ArgType&MEMORY_TYPE)==MEMORY_TYPE)
	{
		if(d.Argument1.Memory.BaseRegister == 0 && d.Argument1.Memory.IndexRegister == 0)  
			return true;
		return false;
	}
	return false;
}

bool SCFI_Instrument::is_jmp_a_fixed_call(Instruction_t* insn) 
{
	if(preds[insn].size()!=1)
		return false;

	Instruction_t* pred=*(preds[insn].begin());
	assert(pred);

	if(pred->GetDataBits()[0]==0x68)
		return true;
	return false;
}

bool SCFI_Instrument::instrument_jumps() 
{
	int cfi_checks=0;
	int cfi_branch_jmp_checks=0;
	int cfi_branch_jmp_complete=0;
	int cfi_branch_call_checks=0;
	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 cfi_safefn_call_skipped=0;
	int ibt_complete=0;
	double cfi_branch_jmp_complete_ratio = NAN;
	double cfi_branch_call_complete_ratio = NAN;
	double cfi_branch_ret_complete_ratio = NAN;

	std::map<int, int> calls;
	std::map<int, int> jmps;
	std::map<int, int> rets;

	// build histogram of target sizes

	// for each instruction
	for(InstructionSet_t::iterator it=firp->GetInstructions().begin();
		it!=firp->GetInstructions().end();
		++it)
	{
		Instruction_t* insn=*it;
		DISASM d;
		insn->Disassemble(d);


		// we always have to protect the zestcfi dispatcher, that we just added.
		if(zestcfi_function_entry==insn)
		{
			cout<<"Protecting zestcfi function for external entrances"<<endl;
			cfi_checks++;
			AddJumpCFI(insn);
			continue;
		}

		if(insn->GetBaseID()==BaseObj_t::NOT_IN_DATABASE)
			continue;

		if(string(d.Instruction.Mnemonic)==string("call ") && (protect_safefn && !do_exe_nonce_for_call))
		{
                	cerr<<"Fatal Error: Found call instruction!"<<endl;
                	cerr<<"FIX_CALLS_FIX_ALL_CALLS=1 should be set in the environment, or"<<endl;
                	cerr<<"--step-option fix_calls:--fix-all should be passed to ps_analyze."<<endl;
			exit(1);
		}

		// if marked safe
		if(FindRelocation(insn,"cf::safe"))
			continue;

		bool safefn = isSafeFunction(insn);
	
		
		switch(d.Instruction.BranchType)
		{
			case  JmpType:
			{
				if((d.Argument1.ArgType&CONSTANT_TYPE)!=CONSTANT_TYPE)
				{
					bool is_fixed_call=is_jmp_a_fixed_call(insn);
					bool is_plt_style=is_plt_style_jmp(insn);
					bool is_any_call_style = (is_fixed_call || is_plt_style);
					if(do_jumps && !is_any_call_style)
					{
						if (insn->GetIBTargets() && insn->GetIBTargets()->IsComplete())
						{
							cfi_branch_jmp_complete++;
							jmps[insn->GetIBTargets()->size()]++;
						}

						cfi_checks++;
						cfi_branch_jmp_checks++;

						AddJumpCFI(insn);
					}
					else if(do_calls && is_any_call_style)
					{
						if (insn->GetIBTargets() && insn->GetIBTargets()->IsComplete())
						{
							cfi_branch_call_complete++;
							calls[insn->GetIBTargets()->size()]++;
						}

						cfi_checks++;
						cfi_branch_call_checks++;
						
						AddJumpCFI(insn);
					}
					else 
					{	
						cout<<"Eliding protection for "<<insn->getDisassembly()<<std::boolalpha
							<<" is_fixed_call="<<is_fixed_call
							<<" is_plt_style="<<is_plt_style
							<<" is_any_call_style="<<is_any_call_style
							<<" do_jumps="<<do_jumps
							<<" do_calls="<<do_calls<<endl;
					}
				}
	
				break;
			}
			case  CallType:
			{

				// should only see calls if we are not CFI'ing safe functions
				// be sure to use with: --no-fix-safefn in fixcalls
				//    (1) --no-fix-safefn in fixcalls leaves call as call (instead of push/jmp)
				//    (2) and here, we don't plop down a nonce
				//    see (3) below where we don't instrument returns for safe functions
				if (!protect_safefn)
				{
					bool isDirectCall = (d.Argument1.ArgType&CONSTANT_TYPE)==CONSTANT_TYPE;

					if (safefn || (isDirectCall && isCallToSafeFunction(insn)))
					{
						cfi_safefn_call_skipped++;
						continue;
					}
				}

				AddExecutableNonce(insn);	// for all calls
				if((d.Argument1.ArgType&CONSTANT_TYPE)!=CONSTANT_TYPE)
				{
					// for indirect calls.
					AddCallCFIWithExeNonce(insn);
				}
				break;
			}
			case  RetType: 
			{
				if (insn->GetIBTargets() && insn->GetIBTargets()->IsComplete())
				{
					cfi_branch_ret_complete++;
					rets[insn->GetIBTargets()->size()]++;
				}

				// (3) and here, we don't instrument returns for safe function
				if (!protect_safefn && safefn)
				{
					cfi_safefn_ret_skipped++;
					continue;
				}

				cfi_checks++;
				cfi_branch_ret_checks++;

				if(do_exe_nonce_for_call)
					AddReturnCFIForExeNonce(insn);
				else
					AddReturnCFI(insn);
				break;

			}
			default:
			{
				break;
			}
		}
	}
	
	cout<<"# ATTRIBUTE cfi_jmp_checks="<<std::dec<<cfi_branch_jmp_checks<<endl;
	cout<<"# ATTRIBUTE cfi_jmp_complete="<<cfi_branch_jmp_complete<<endl;
	display_histogram(cout, "cfi_jmp_complete_histogram", jmps);

	cout<<"# ATTRIBUTE cfi_branch_call_checks="<<std::dec<<cfi_branch_call_checks<<endl;
	cout<<"# ATTRIBUTE cfi_branch_call_complete="<<std::dec<<cfi_branch_call_complete<<endl;
	display_histogram(cout, "cfi_call_complete_histogram", calls);

	cout<<"# ATTRIBUTE cfi_ret_checks="<<std::dec<<cfi_branch_ret_checks<<endl;
	cout<<"# ATTRIBUTE cfi_ret_complete="<<std::dec<<cfi_branch_ret_complete<<endl;
	display_histogram(cout, "cfi_ret_complete_histogram", rets);


	// 0 or 1 checks.
	cout<<"# ATTRIBUTE multimodule_checks="<< (unsigned int)(zestcfi_function_entry!=NULL) <<endl;

	cout<<"# ATTRIBUTE cfi_checks="<<std::dec<<cfi_checks<<endl;
	ibt_complete = cfi_branch_jmp_complete + cfi_branch_call_complete + cfi_branch_ret_complete;
	cout<<"# ATTRIBUTE ibt_complete="<<std::dec<<ibt_complete<<endl;

	if (cfi_branch_jmp_checks > 0) 
		cfi_branch_jmp_complete_ratio = (double)cfi_branch_jmp_complete / cfi_branch_jmp_checks;

	if (cfi_branch_call_checks > 0) 
		cfi_branch_call_complete_ratio = (double)cfi_branch_call_complete / cfi_branch_call_checks;

	if (cfi_branch_ret_checks > 0) 
		cfi_branch_ret_complete_ratio = (double)cfi_branch_ret_complete / cfi_branch_ret_checks;

	double cfi_branch_complete_ratio = NAN;
	if (ibt_complete > 0)
		cfi_branch_complete_ratio = (double) cfi_checks / ibt_complete;
	

	cout << "# ATTRIBUTE cfi_jmp_complete_ratio=" << cfi_branch_jmp_complete_ratio << endl;
	cout << "# ATTRIBUTE cfi_call_complete_ratio=" << cfi_branch_call_complete_ratio << endl;
	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;
	cout<<"# ATTRIBUTE cfi_safefn_call_skipped="<<cfi_safefn_call_skipped<<endl;

	return true;
}

// use this to determine whether a scoop has a given name.
static struct ScoopFinder : binary_function<const DataScoop_t*,const string,bool>
{
	// declare a simple scoop finder function that finds scoops by name
	bool operator()(const DataScoop_t* scoop, const string& name) const
	{
		return (scoop->GetName() == name);
	};