diff --git a/tools/transforms/PNTransformDriver.cpp b/tools/transforms/PNTransformDriver.cpp index b10f5ce25acf0ce071f69f4f990d4855c894f3c0..6b10ce4744d8f45d02448255c275fd78822acbf3 100644 --- a/tools/transforms/PNTransformDriver.cpp +++ b/tools/transforms/PNTransformDriver.cpp @@ -115,7 +115,7 @@ PNTransformDriver::PNTransformDriver(VariantID_t *pidp,string BED_script, pqxxDB no_validation_level = -1; coverage_threshold = -1; do_shared_object_protection = false; - + m_mitigation_policy = P_CONTROLLED_EXIT; } PNTransformDriver::~PNTransformDriver() @@ -2388,7 +2388,7 @@ bool PNTransformDriver::Canary_Rewrite(PNStackLayout *orig_layout, Function_t *f //This could probably be done once, but having the original instruction //allows me to produce messages that indicate more precisely where //the overflow occurred. - Instruction_t *handler_code = getHandlerCode(virp,instr,P_CONTROLLED_EXIT); + Instruction_t *handler_code = getHandlerCode(virp,instr,GetMitigationPolicy()); //insert canary checks // @@ -2412,7 +2412,7 @@ bool PNTransformDriver::Canary_Rewrite(PNStackLayout *orig_layout, Function_t *f //This could probably be done once, but having the original instruction //allows me to produce messages that indicate more precisely where //the overflow occurred. - Instruction_t *handler_code = getHandlerCode(virp,instr,P_CONTROLLED_EXIT); + Instruction_t *handler_code = getHandlerCode(virp,instr,GetMitigationPolicy()); //insert canary checks // diff --git a/tools/transforms/PNTransformDriver.hpp b/tools/transforms/PNTransformDriver.hpp index 73fabcad05bb62ab3f43913700c7ee42ccb80e1e..d57d5b098c7e3ecb68b0934966528f8bafadb204 100644 --- a/tools/transforms/PNTransformDriver.hpp +++ b/tools/transforms/PNTransformDriver.hpp @@ -100,6 +100,8 @@ class PNTransformDriver // write stack objects to IRDB bool write_stack_ir_to_db; + mitigation_policy m_mitigation_policy; + // a way to map an instruction to it's set of predecessors. std::map< Instruction_t* , set<Instruction_t*> > preds; @@ -182,6 +184,9 @@ public: virtual void GenerateTransforms(); virtual void SetWriteStackIrToDb(bool setting) { write_stack_ir_to_db = setting; } + + inline virtual mitigation_policy GetMitigationPolicy() const { return m_mitigation_policy; } + virtual void SetMitigationPolicy(mitigation_policy policy) { m_mitigation_policy = policy; } }; #endif diff --git a/tools/transforms/Rewrite_Utility.cpp b/tools/transforms/Rewrite_Utility.cpp index 3a2a4e6821d52beacac7922dfc28ca09c8036320..957333fb091d42cdeb3ee72e5704c4f5101f7645 100644 --- a/tools/transforms/Rewrite_Utility.cpp +++ b/tools/transforms/Rewrite_Utility.cpp @@ -334,10 +334,21 @@ Instruction_t* getHandlerCode(FileIR_t* virp, Instruction_t* fallthrough, mitiga else { assert(virp->GetArchitectureBitWidth()==64); - handler_code= allocateNewInstruction(virp,fallthrough); - setInstructionAssembly(virp,handler_code,"hlt",NULL,NULL); - handler_code->SetComment("hlt ; Make this into a callback: jdh@getHandlerCode"); - handler_code->SetFallthrough(fallthrough); + if (policy == P_CONTROLLED_EXIT) + { + handler_code = allocateNewInstruction(virp,fallthrough); + setInstructionAssembly(virp,handler_code,"mov rdi, 189",NULL,NULL); + Instruction_t* syscall_num = insertAssemblyAfter(virp,handler_code,"mov rax, 60",NULL); + Instruction_t* syscall_i = insertAssemblyAfter(virp,syscall_num,"syscall",NULL); + syscall_i->SetFallthrough(fallthrough); + } + else + { + handler_code= allocateNewInstruction(virp,fallthrough); + setInstructionAssembly(virp,handler_code,"hlt",NULL,NULL); + handler_code->SetComment("hlt ; Make this into a callback: jdh@getHandlerCode"); + handler_code->SetFallthrough(fallthrough); + } } return handler_code;