<p class="MsoNormal"><span style="font-size: 10pt; font-family: "Courier New";">Hi
all,</span></p>
<p class="MsoNormal"><span style="font-size: 10pt; font-family: "Courier New";" lang="EN-US">I was trying to implement an obfuscation tool for C-code
on the basis of LLVM. <span style=""> </span>I got a prototype
of the simple obfuscation transformation which converting control flow graph to
something like a state machine. </span><span style="font-size: 10pt; font-family: "Courier New";">I am not sure I will have time to work on extending further
this tool with new transformations like opaque predicates and decided to put here
source code I have by now with hope that it can be helpful to somebody. </span><span style="font-size: 10pt; font-family: "Courier New";" lang="EN-US">I am newbie with LLVM and am sure there are a lot of bugs in current
implementation, but I have tested it with complex enough C-program and in my
case it worked.<span style=""> </span>Sorry for the poor
English and grammar.</span></p>
<p class="MsoNormal"><span style="font-size: 10pt; font-family: "Courier New";" lang="EN-US">Regards,<br>Evgeny Podlepaev</span></p>
--------------------------------------------- MakeDispatcherPass.h --------------------------------------------- <br><br>#include <map><br>#include <string><br><br>#include "llvm/Pass.h"<br>#include "llvm/Function.h"
<br>#include "llvm/BasicBlock.h"<br>#include "llvm/Instructions.h"<br><br>using namespace llvm;<br><br>class MakeDispatcherPass : public FunctionPass<br>{<br>public:<br> virtual bool runOnFunction( Function& currFunction );
<br><br>private:<br> typedef std::map< BasicBlock*, unsigned > BBindexBBMap;<br><br> static unsigned IndexSourceBasicBlocks( Function& function, BBindexBBMap& indexBBMap );<br><br> static BasicBlock* CreateNewEntryBlock( const std::string& name, BasicBlock *oldEntryBB );
<br><br> static void LinkBasicBlockToDispatcher( BasicBlock* basicBlock, BasicBlock* dispatcherBlock, Value* statePtr, BBindexBBMap& indexBBMap );<br><br> static void MovePHINodesToDispatcher( BBindexBBMap& indexBBMap, BasicBlock* dispatcherBB );
<br><br> static void ConvertSwitch( Function& function );<br><br> static BasicBlock* ProcessCase( SwitchInst* switchInst, int caseIdx, BasicBlock* prevBB, Value* testValuePtr );<br><br> static void ReduceTempVarsLifetime( Function& function );
<br><br> static bool IsUsedOutsideParentBlock( Instruction* value );<br><br> static const std::string PassName;<br>};<br><br>--------------------------------------------- MakeDispatcherPass.cpp ---------------------------------------------
<br><br>#include <map><br>#include <vector><br>#include <cassert><br>#include <iostream><br>#include <exception><br><br>#include "llvm/Type.h"<br>#include "llvm/Module.h"<br>
#include "llvm/Function.h"<br>#include "llvm/Constants.h"<br>#include "llvm/Instructions.h"<br>#include "llvm/Support/CFG.h"<br>#include "llvm/Support/Debug.h"<br>#include "llvm/Support/Casting.h"
<br>#include "llvm/Support/InstIterator.h"<br><br>#include "MakeDispatcherPass.h"<br><br>//<br>// Static variables.<br>//<br><br>const std::string MakeDispatcherPass::PassName = "MakeDispatcherPass";
<br><br>//<br>// Public methods.<br>//<br><br>bool MakeDispatcherPass::runOnFunction( Function& currFunction )<br>{ <br> DEBUG( std::cerr << PassName << ": Processing " << currFunction.getName
() << " ...\n" );<br> bool madeChange = true; <br><br> // Because of changes which we will make in control flow graph it is necessary<br> // to get rid of temporary variables used outside their parent basic block.
<br> ReduceTempVarsLifetime( currFunction );<br><br> // Replace all SwitchInst instructions with chained branch instructions<br> // in order to simplify the dispatcher embedding.<br> //<br> // NOTE: We can use LowerSwitch Pass instead of our implementation.
<br> ConvertSwitch( currFunction );<br><br> // Keep a reference to the old entry basic block.<br> BasicBlock* oldEntryBB = &currFunction.getEntryBlock();<br><br> // Collect basic blocks to be reordered and assign unique index for each.
<br> BBindexBBMap indexBBMap;<br> unsigned firstBBIndex = IndexSourceBasicBlocks( currFunction, indexBBMap );<br><br> // Create basic blocks for initialization and dispatching.<br> BasicBlock* disInitBB = CreateNewEntryBlock( "dispatcher_init", oldEntryBB );
<br> BasicBlock* disMainBB = new BasicBlock( "dispatcher_main", &currFunction, oldEntryBB );<br><br> // Add instructions that allocate and initialize the dispatcher's state variable<br> // to the initialization block.
<br> AllocaInst* statePtr = new AllocaInst( Type::UIntTy, 0, "state", disInitBB );<br> new StoreInst( ConstantUInt::get( Type::UIntTy, firstBBIndex ), statePtr, disInitBB );<br> new BranchInst( disMainBB, disInitBB ); // Connection to the dispatcher's main block.
<br> <br> // Add instructions to the dispatcher's main block that implement switching<br> // between the basic blocks depending on the state variable's value.<br> LoadInst* loadState = new LoadInst( statePtr, "next", disMainBB );
<br> SwitchInst* switchInst = new SwitchInst( loadState, oldEntryBB, unsigned( indexBBMap.size() ), disMainBB );<br> for( BBindexBBMap::iterator idxBBPair = indexBBMap.begin(); idxBBPair != indexBBMap.end(); idxBBPair++ )
<br> {<br> switchInst->addCase( ConstantUInt::get( Type::UIntTy, idxBBPair->second ), idxBBPair->first );<br> }<br><br> // Unlink all source basic blocks and relink them to the dispatcher's main block.
<br> for( BBindexBBMap::iterator idxBBPair = indexBBMap.begin(); idxBBPair != indexBBMap.end(); idxBBPair++ )<br> {<br> LinkBasicBlockToDispatcher( idxBBPair->first, disMainBB, statePtr, indexBBMap );<br> }
<br><br> // Since the predecessors of source basic blocks was replaced by the dispatcher's<br> // basic block it is necessary to update all PHI nodes in source blocks. PHI is a<br> // special artificial definition for supporting the Static Single Assignment form
<br> // used for data flow analysis (<a href="http://gcc.gnu.org/onlinedocs/gccint/SSA.html">http://gcc.gnu.org/onlinedocs/gccint/SSA.html</a>).<br> MovePHINodesToDispatcher( indexBBMap, disMainBB ); <br><br> return madeChange;
<br>}<br><br>//<br>// Private methods.<br>//<br><br>unsigned MakeDispatcherPass::IndexSourceBasicBlocks( Function& function, BBindexBBMap& indexBBMap )<br>{<br> unsigned startIdx = 0;<br><br> // Keep all source basic blocks, including entry and exit blocks.
<br> unsigned idx = startIdx;<br> for( Function::iterator i = function.begin(); i != function.end(); i++ )<br> {<br> // Add basic block with index.<br> indexBBMap[ &*i ] = idx++;<br> }<br><br>
// Return the index of the first basic block that is an entry to the function.<br> return startIdx;<br>}<br><br>BasicBlock* MakeDispatcherPass::CreateNewEntryBlock( const std::string& name, BasicBlock *oldEntryBB )
<br>{<br> // Create a new basic block and insert it before old entry block.<br> BasicBlock* newEntryBB = new BasicBlock( name, oldEntryBB->getParent(), oldEntryBB );<br> <br> // Move all allocation instructions for stack variables from the old entry block to the new one.
<br> // These instructions should be placed at the beginning of the function and should dominate<br> // all their uses.<br> for( BasicBlock::iterator i = oldEntryBB->begin(); i != oldEntryBB->end(); )<br> {
<br> if( AllocationInst* ai = dyn_cast< AllocationInst >( i ) )<br> { <br> oldEntryBB->getInstList().remove( i++ ); // Trick to preserve the valid iterator after removing an element.
<br> newEntryBB->getInstList().push_back( ai );<br> }<br> else<br> {<br> i++;<br> } <br> }<br><br> oldEntryBB->setName( "old_entry" );<br> return newEntryBB;
<br>}<br><br>void MakeDispatcherPass::LinkBasicBlockToDispatcher( BasicBlock* basicBlock, BasicBlock* dispatcherBlock, Value* statePtr, BBindexBBMap& indexBBMap )<br>{<br> // Get a reference to the basic block terminator. We should replace it
<br> // by jump to the dispatcher's basic block.<br> TerminatorInst* terminator = basicBlock->getTerminator();<br> assert( terminator && "Basic block is not well formed and has no terminator!" );
<br><br> if( isa< ReturnInst >( terminator ) )<br> {<br> // This block is exit from the function. We will not change it.<br> return;<br> }<br> else if( isa< BranchInst >( terminator ) )
<br> {<br> // This block is terminating by branch instruction. We will replace target<br> // of the branch by the dispatcher block and set current state variable to<br> // a proper basic block's index.
<br> BranchInst* branchInst = dynamic_cast< BranchInst* >( terminator );<br> <br> // Add instruction selecting a proper basic block's index.<br> Value* stateValue = 0;<br> if( !branchInst->isConditional() )
<br> { <br> stateValue = ConstantUInt::get(<br> Type::UIntTy,<br> indexBBMap[ branchInst->getSuccessor( 0 ) ]<br> );<br> }<br>
else<br> {<br> stateValue = new SelectInst(<br> branchInst->getCondition(),<br> ConstantUInt::get( Type::UIntTy, indexBBMap[ branchInst->getSuccessor( 0 ) ] ),
<br> ConstantUInt::get( Type::UIntTy, indexBBMap[ branchInst->getSuccessor( 1 ) ] ),<br> "select_state",<br> basicBlock<br> );
<br> } <br><br> // Add instructions for load of the state variable and branch to the dispatcher.<br> new StoreInst( stateValue, statePtr, basicBlock );<br> new BranchInst( dispatcherBlock, basicBlock );
<br> basicBlock->getInstList().erase( branchInst );<br> }<br> else<br> {<br> // TODO: Support other possible terminators.<br> throw std::exception( ( PassName + ": Unsupported basic block terminator: " + typeid( terminator ).name() ).c_str() );
<br> } <br>}<br><br>void MakeDispatcherPass::MovePHINodesToDispatcher( BBindexBBMap& indexBBMap, BasicBlock* dispatcherBB )<br>{<br> // Iterate through the list of indexed basic blocks and move all PHIs nodes
<br> // to the dispatcher's basic block (target basic block) that is new predecessor<br> // for all them.<br> for( BBindexBBMap::iterator idxBBPair = indexBBMap.begin(); idxBBPair != indexBBMap.end(); idxBBPair++ )
<br> {<br> // Iterate through the PHIs nodes that can be located only at the<br> // beginning of the basic block.<br> PHINode* phi = 0;<br> for( BasicBlock::iterator inst = idxBBPair->first->begin(); ( phi = dyn_cast< PHINode >( inst ) ) != 0; )
<br> {<br> // Move the PHI node to the target basic block.<br> phi->getParent()->getInstList().remove( inst++ );<br> dispatcherBB->getInstList().insert( dispatcherBB->getInstList().begin(), phi );
<br><br> // Since a PHI node must have an incoming value for every predecessor of<br> // the parent basic block we should add all predecessors of the target<br> // basic block that have no corresponding incoming value in the PHI node
<br> // as new incoming undefined value to the PHI node.<br> int wasIncomingCount = 0;<br> for( pred_iterator pred = pred_begin( dispatcherBB ); pred != pred_end( dispatcherBB ); pred++ )<br>
{ <br> int idx = phi->getBasicBlockIndex( *pred );<br> if( idx == -1 )<br> {<br> phi->addIncoming( UndefValue::get( phi->getType() ), *pred );
<br> }<br> else<br> {<br> wasIncomingCount++;<br> }<br> }<br> assert( wasIncomingCount && "No one of the predecessors was not incoming value to the PHI!" );
<br> }<br> }<br>}<br><br>void MakeDispatcherPass::ConvertSwitch( Function& function )<br>{<br> BasicBlock* entryBB = &function.getEntryBlock();<br><br> for( Function::iterator i = function.begin(); i !=
function.end(); i++ )<br> {<br> BasicBlock* basicBlock = &*i;<br><br> TerminatorInst* terminator = basicBlock->getTerminator();<br> assert( terminator && "Basic block is not well formed and has no terminator!" );
<br><br> if( isa< SwitchInst >( terminator ) )<br> {<br> SwitchInst* switchInst = dynamic_cast< SwitchInst * >( basicBlock->getTerminator() );<br><br> // Allocate a local variable and add an instruction saving the value
<br> // being tested by the switch instruction to this variable.<br> //<br> // Allocating of the new variable is necessary because initial order<br> // of the basic blocks will be changed by dispatcher embedding.
<br> AllocaInst* testValuePtr = new AllocaInst( switchInst->getCondition()->getType(), 0, "switch_test_val", entryBB->getTerminator() );<br> new StoreInst( switchInst->getCondition(), testValuePtr, false, switchInst );
<br><br> // Replace the switch instruction by basic blocks with conditional branches.<br> BasicBlock* caseBlock = ProcessCase( switchInst, 1, basicBlock, testValuePtr );<br> switchInst->eraseFromParent();
<br><br> new BranchInst( caseBlock, basicBlock );<br> }<br> }<br>}<br><br>BasicBlock* MakeDispatcherPass::ProcessCase( SwitchInst* switchInst, int caseIdx, BasicBlock* prevBB, Value* testValuePtr )<br>
{<br> // Create a new basic block for instructions checking the current case of the switch instruction.<br> BasicBlock* caseBlock = new BasicBlock( "case", prevBB->getParent(), prevBB->getNext() );<br>
<br> // If not all cases was handled ...<br> if( caseIdx != switchInst->getNumCases() )<br> { <br> // Instructions for checking the case.<br> LoadInst* testValue = new LoadInst( testValuePtr, "test_value", caseBlock );
<br> SetCondInst* setCondInst = new SetCondInst( Instruction::SetEQ, testValue, switchInst->getCaseValue( caseIdx ), "is_case", caseBlock );<br><br> // Get a reference to the basic block for the next case.
<br> BasicBlock* nextCaseBlock = ProcessCase( switchInst, caseIdx + 1, caseBlock, testValuePtr );<br><br> // Add an instruction for branching to the basic block<br> // corresponding to the condition of this case or to
<br> // the basic block checking the next case.<br> BranchInst* branchInst = new BranchInst( switchInst->getSuccessor( caseIdx ), nextCaseBlock, setCondInst, caseBlock );<br> }<br> else if( switchInst->getDefaultDest() != 0 )
<br> {<br> BranchInst* branchInst = new BranchInst( switchInst->getDefaultDest(), caseBlock );<br> }<br><br> return caseBlock;<br>}<br><br>void MakeDispatcherPass::ReduceTempVarsLifetime( Function& function )
<br>{<br> BasicBlock* entryBB = &function.getEntryBlock();<br><br> typedef std::vector< Instruction * > InstList;<br><br> // Collect all instructions (temporary variables) that have result used outside<br>
// their parent basic block. Skip the AllocInst instructions because it will be<br> // handled separately during creation of new entry basic block.<br> InstList insts;<br> for( inst_iterator i = inst_begin( function ); i != inst_end( function ); i++ )
<br> {<br> Instruction* inst = &*i;<br><br> if( IsUsedOutsideParentBlock( inst ) && !isa< AllocaInst >( inst ) )<br> {<br> insts.push_back( inst );<br> }<br> }
<br><br> // Walk through the list of temporary variables to reduce their lifetime<br> // to its parent basic block.<br> for( InstList::iterator i = insts.begin(); i != insts.end(); i++ )<br> {<br> Instruction* inst = *i;
<br><br> // Collect all instructions that are using the temporary variable.<br> InstList users;<br> for( Value::use_iterator j = inst->use_begin(); j != inst->use_end(); j++ )<br> {<br>
users.push_back( dynamic_cast< Instruction* >( *j ) ); <br> }<br><br> // Allocate local variable on the function's stack frame and save there<br> // the temporary variable (result of the instruction).
<br> AllocaInst* valuePtr = new AllocaInst( inst->getType(), 0, "reduced_var", entryBB->getTerminator() );<br> new StoreInst( inst, valuePtr, false, inst->getNext() );<br><br> // Patch users of the temporary variable to use the local function variable.
<br> for( InstList::iterator j = users.begin(); j != users.end(); j++ )<br> {<br> Instruction* user = dynamic_cast< Instruction* >( *j );<br><br> if( isa< PHINode >( user ) )
<br> {<br> // For the phi node we should look through every incoming value and,<br> // if it is necessary, add to the preceding basic block an instruction<br> // for loading the saved temporary variable's value.
<br> PHINode* phiNode = dynamic_cast< PHINode* >( user );<br> for( unsigned idx = 0; idx < phiNode->getNumIncomingValues(); idx++ )<br> {<br> if( phiNode->getIncomingValue( idx ) == inst )
<br> {<br> LoadInst* loadedValue = new LoadInst( valuePtr, "value_of_reduced_var", phiNode->getIncomingBlock( idx )->getTerminator() );<br> phiNode->setIncomingValue( idx, loadedValue );
<br> }<br> }<br> }<br> else<br> {<br> user->replaceUsesOfWith( inst, new LoadInst( valuePtr, "value_of_reduced_var", user ) );
<br> }<br> }<br> }<br>}<br><br>bool MakeDispatcherPass::IsUsedOutsideParentBlock( Instruction* value )<br>{<br> for( Value::use_iterator i = value->use_begin(); i != value->use_end(); i++ )<br>
{<br> Instruction* user = dynamic_cast< Instruction* >( *i );<br> if( user->getParent() != value->getParent() )<br> {<br> return true;<br> }<br> }<br> return false;
<br>}<br><br>--------------------------------------------- Main.cpp --------------------------------------------- <br><br>#include <iostream><br>#include <fstream><br><br>#include "llvm/Pass.h"<br>#include "llvm/Module.h"
<br>#include "llvm/PassManager.h"<br>#include "llvm/Bytecode/Reader.h"<br>#include "llvm/Analysis/Verifier.h"<br>#include "llvm/Support/CommandLine.h"<br>#include "llvm/Target/TargetMachine.h"
<br>#include "llvm/Bytecode/WriteBytecodePass.h"<br><br>#include "MakeDispatcherPass.h"<br><br>using namespace llvm;<br><br>static cl::opt< std::string > <br>InputFilename( cl::Positional, cl::desc( "<input .bc file>" ), cl::Required );
<br><br>static cl::opt< std::string > <br>OutputFilename( cl::Positional, cl::desc( "<output .bc file>" ), cl::Required );<br><br>static cl::opt< bool ><br>NoVerify( "disable-verify", cl::desc( "Do not verify result module" ), cl::Optional );
<br><br>int main( int argc, char* argv[] )<br>{<br> std::string programName = argv[ 0 ] = "obfuscator";<br><br> try<br> {<br> cl::ParseCommandLineOptions( argc, argv );<br> std::string errorMessage;
<br><br> // Load the input module.<br> std::auto_ptr< Module > M( ParseBytecodeFile( InputFilename, &errorMessage ) );<br> if( M.get() == 0 )<br> {<br> std::cerr << programName << ": " << errorMessage << "\n";
<br> return 1;<br> }<br><br> // Open the stream we are supposed to write to.<br> std::ostream *Out = new std::ofstream( OutputFilename.c_str(), std::ios::out | std::ios::trunc | std::ios::binary );
<br> if( !Out->good() )<br> {<br> std::cerr << programName << ": Error opening " << OutputFilename << "!\n";<br> return 1;<br> }<br>
<br> // Create a PassManager to hold and optimize the collection of passes we are about to build...<br> PassManager Passes;<br><br> // Add an appropriate TargetData instance for this module.<br>
Passes.add( new TargetData( "obfuscator", M.get() ) );<br><br> Passes.add( new MakeDispatcherPass() );<br><br> if( !NoVerify )<br> {<br> Passes.add( createVerifierPass() );<br> }
<br> Passes.add( new WriteBytecodePass( Out, true ) ); <br><br> // Now that we have all of the passes ready, run them.<br> Passes.run( *M.get() );<br> }<br> catch( std::exception e )<br> {
<br> std::cerr << programName << ": " << e.what() << "\n";<br> return 1;<br> }<br> catch( const std::string& msg )<br> {<br> std::cerr << programName << ": " << msg << "\n";
<br> return 1;<br> }<br> catch( ... )<br> {<br> std::cerr << programName << ": Unexpected exception occurred.\n";<br> return 1;<br> }<br><br> return 0;<br>}<br><br>