Implemented branch correction.

This commit is contained in:
Can Bölük 2020-06-15 16:07:13 +02:00
parent 50b0db03f9
commit 7fc98a17ee
11 changed files with 369 additions and 168 deletions

View file

@ -142,7 +142,7 @@ namespace vtil
// -- Control flow instructions
//
// JS Reg, Reg/Imm, Reg/Imm | Jumps to OP3 if OP1 != 0, else jumps to OP2, continues virtual execution
// JS Reg, Reg/Imm, Reg/Imm | Jumps to OP1 is set, jumps to OP2, otherwise OP3, continues virtual execution
// JMP Reg/Imm | Jumps to OP1, continues virtual execution
// VEXIT Reg/Imm | Jumps to OP1, continues real execution
// VXCALL Reg/Imm | Calls into OP1, pauses virtual execution until the call returns

View file

@ -48,7 +48,7 @@ namespace vtil
// Declare the lookup map for the cache mapping each variable to the
// result of the primitive traver.
//
cache_type cache;
mutable cache_type cache;
// Locks the cache.
//

View file

@ -9,7 +9,7 @@ add_library(${PROJECT_NAME} STATIC
passes/dead_code_elimination_pass.cpp
passes/mov_propagation_pass.cpp
passes/symbolic_rewrite_pass.cpp
passes/opaque_predicate_elimination_pass.cpp
passes/branch_correction_pass.cpp
passes/register_renaming_pass.cpp
)

View file

@ -111,7 +111,7 @@
<ClCompile Include="passes\dead_code_elimination_pass.cpp" />
<ClCompile Include="passes\istack_ref_substitution_pass.cpp" />
<ClCompile Include="passes\mov_propagation_pass.cpp" />
<ClCompile Include="passes\opaque_predicate_elimination_pass.cpp" />
<ClCompile Include="passes\branch_correction_pass.cpp" />
<ClCompile Include="passes\register_renaming_pass.cpp" />
<ClCompile Include="passes\stack_pinning_pass.cpp" />
<ClCompile Include="passes\stack_propagation_pass.cpp" />
@ -125,7 +125,7 @@
<ClInclude Include="passes\dead_code_elimination_pass.hpp" />
<ClInclude Include="passes\istack_ref_substitution_pass.hpp" />
<ClInclude Include="passes\mov_propagation_pass.hpp" />
<ClInclude Include="passes\opaque_predicate_elimination_pass.hpp" />
<ClInclude Include="passes\branch_correction_pass.hpp" />
<ClInclude Include="passes\register_renaming_pass.hpp" />
<ClInclude Include="passes\stack_pinning_pass.hpp" />
<ClInclude Include="passes\stack_propagation_pass.hpp" />

View file

@ -25,10 +25,10 @@
<ClCompile Include="passes\symbolic_rewrite_pass.cpp">
<Filter>Optimizer Passes</Filter>
</ClCompile>
<ClCompile Include="passes\opaque_predicate_elimination_pass.cpp">
<ClCompile Include="passes\register_renaming_pass.cpp">
<Filter>Optimizer Passes</Filter>
</ClCompile>
<ClCompile Include="passes\register_renaming_pass.cpp">
<ClCompile Include="passes\branch_correction_pass.cpp">
<Filter>Optimizer Passes</Filter>
</ClCompile>
</ItemGroup>
@ -74,10 +74,10 @@
<ClInclude Include="passes\symbolic_rewrite_pass.hpp">
<Filter>Optimizer Passes</Filter>
</ClInclude>
<ClInclude Include="passes\opaque_predicate_elimination_pass.hpp">
<ClInclude Include="passes\register_renaming_pass.hpp">
<Filter>Optimizer Passes</Filter>
</ClInclude>
<ClInclude Include="passes\register_renaming_pass.hpp">
<ClInclude Include="passes\branch_correction_pass.hpp">
<Filter>Optimizer Passes</Filter>
</ClInclude>
</ItemGroup>

View file

@ -34,7 +34,7 @@
#include "../passes/dead_code_elimination_pass.hpp"
#include "../passes/mov_propagation_pass.hpp"
#include "../passes/symbolic_rewrite_pass.hpp"
#include "../passes/opaque_predicate_elimination_pass.hpp"
#include "../passes/branch_correction_pass.hpp"
#include "../passes/register_renaming_pass.hpp"
namespace vtil::optimizer
@ -47,7 +47,7 @@ namespace vtil::optimizer
using collective_routine_correction_pass = combine_pass<
stack_pinning_pass,
istack_ref_substitution_pass,
opaque_predicate_elimination_pass,
branch_correction_pass,
bblock_extension_pass
>;
@ -61,7 +61,7 @@ namespace vtil::optimizer
register_renaming_pass,
dead_code_elimination_pass,
conditional_pass<
opaque_predicate_elimination_pass,
branch_correction_pass,
conditional_pass<
bblock_extension_pass,
symbolic_rewrite_pass<true>

View file

@ -458,10 +458,6 @@ namespace vtil::optimizer::aux
{
fassert( var.is_register() );
// If cross block operation, lock routine mutex.
//
cnd_unique_lock _g{ it.container->owner->mutex, it.container != var.at.container };
// Drop const-qualifiers, this operation is not illegal since we're passed
// non-constant iterator, meaning we have access to the routine itself.
//
@ -487,24 +483,33 @@ namespace vtil::optimizer::aux
// Returns each possible branch destination of the given basic block in the format of:
// - [is_real, target] x N
//
std::vector<std::pair<bool, symbolic::expression>> discover_branches( const basic_block* blk, tracer* tracer, bool xblock )
branch_info analyze_branch( const basic_block* blk, tracer* tracer, bool xblock, bool pack )
{
// If block is not complete, return empty vector.
//
std::vector<std::pair<bool, symbolic::expression>> targets = {};
if ( !blk->is_complete() )
return targets;
return {};
// Declare tracer.
//
const auto trace = [ & ] ( symbolic::variable&& lookup )
{
auto exp = xblock ? tracer->rtrace( std::move( lookup ) ) : tracer->trace( std::move( lookup ) );
if ( pack )
symbolic::variable::pack_all( exp );
return exp;
};
// Declare operand->expression helper.
//
auto branch = std::prev( blk->end() );
auto discover = [ & ] ( const operand& op_dst, bool real )
auto discover = [ & ] ( const operand& op_dst, bool real, bool parse = true ) -> branch_info
{
// Determine the symbolic expression describing branch destination.
//
symbolic::expression destination = op_dst.is_immediate()
? symbolic::expression{ op_dst.imm().u64 }
: ( xblock ? tracer->rtrace_p( { branch, op_dst.reg() } ) : tracer->trace_p( { branch, op_dst.reg() } ) );
: trace( { branch, op_dst.reg() } );
// Remove any matches of REG_IMGBASE and pack.
//
@ -518,38 +523,95 @@ namespace vtil::optimizer::aux
}
} ).simplify( true );
// Match classic Jcc:
// If parsing requested:
//
using namespace symbolic::directive;
std::vector<symbol_table_t> results;
if ( fast_match( &results, A + ( __if( B, D ) + __if( C, E ) ), destination ) ||
fast_match( &results, A + ( __if( B, D ) | __if( C, E ) ), destination ) ||
fast_match( &results, __if( B, D ) + __if( C, E ), destination ) ||
fast_match( &results, __if( B, D ) | __if( C, E ), destination ) )
if ( parse )
{
auto& sym = results.front();
sym.add( A, symbolic::expression{ 0, 64 } );
if ( sym.translate( B )->equals( ~sym.translate( C ) ) )
// Match classic Jcc:
//
using namespace symbolic::directive;
std::vector<symbol_table_t> results;
if ( fast_match( &results, A + ( __if( B, D ) + __if( C, E ) ), destination ) ||
fast_match( &results, A + ( __if( B, D ) | __if( C, E ) ), destination ) ||
fast_match( &results, __if( B, D ) + __if( C, E ), destination ) ||
fast_match( &results, __if( B, D ) | __if( C, E ), destination ) )
{
auto reloc = sym.translate( A );
targets.emplace_back( real, reloc + sym.translate( D ) );
targets.emplace_back( real, reloc + sym.translate( E ) );
return;
}
}
// Pick the first result and translate conditions.
//
auto& sym = results.front();
auto cond1 = sym.translate( B );
auto cond2 = sym.translate( C );
// If not, push as is.
// If inverse conditionals:
//
if ( cond1->equals( ~cond2 ) )
{
// Translate destinations.
//
auto reloc = sym.translate( A ) ? *sym.translate( A ) : symbolic::expression{ 0, 64 };
auto dst1 = sym.translate( D );
auto dst2 = sym.translate( E );
// Make sure first condition is the simplest.
//
if ( cond1->complexity > cond2->complexity )
{
std::swap( cond1, cond2 );
std::swap( dst1, dst2 );
}
return {
.is_vm_exit = real,
.is_jcc = true,
.cc = *cond1,
.destinations = { reloc + dst1, reloc + dst2 }
};
}
}
// -- TODO: Handle jump tables.
//
}
// Otherwise assume direct jump.
//
targets.emplace_back( real, destination );
return {
.is_vm_exit = real,
.destinations = { destination }
};
};
// Discover all targets and return.
//
for ( int idx : branch->base->branch_operands_vip )
discover( branch->operands[ idx ], false );
for ( int idx : branch->base->branch_operands_rip )
discover( branch->operands[ idx ], true );
return targets;
if ( *branch->base == ins::jmp )
return discover( branch->operands[ 0 ], false );
if ( *branch->base == ins::vexit )
return discover( branch->operands[ 0 ], true );
if ( *branch->base == ins::vxcall )
return discover( branch->operands[ 0 ], true );
if ( *branch->base == ins::js )
{
// If condition can be resolved in compile time:
//
symbolic::expression cc = trace( { branch, branch->operands[ 0 ].reg() } );
if ( cc.is_constant() )
{
// Redirect to jmp resolver.
//
return discover( branch->operands[ *cc.get<bool>() ? 1 : 2 ], false, false );
}
// Resolve each individually and form jcc.
//
branch_info b1 = discover( branch->operands[ 1 ], false, false );
branch_info b2 = discover( branch->operands[ 2 ], false, false );
return {
.is_vm_exit = false,
.is_jcc = true,
.cc = cc,
.destinations = { b1.destinations[ 0 ], b2.destinations[ 0 ] }
};
}
unreachable();
}
};

View file

@ -31,6 +31,24 @@
namespace vtil::optimizer::aux
{
// Simple structure describing branch details.
//
struct branch_info
{
// If jump to real:
//
bool is_vm_exit = false;
// If jcc:
//
bool is_jcc = false;
symbolic::expression cc;
// Possible destination expressions:
//
std::vector<symbolic::expression> destinations;
};
// Helper to check if the expression given is block-local.
//
bool is_local( const symbolic::expression& ex );
@ -47,8 +65,7 @@ namespace vtil::optimizer::aux
//
register_desc revive_register( const symbolic::variable& var, const il_iterator& it );
// Returns each possible branch destination of the given basic block in the format of:
// - [is_real, target] x N
// Extracts the details of the branch taken at the end of the block where possible.
//
std::vector<std::pair<bool, symbolic::expression>> discover_branches( const basic_block* blk, tracer* tracer, bool xblock );
branch_info analyze_branch( const basic_block* blk, tracer* tracer, bool xblock, bool pack = true );
}

View file

@ -0,0 +1,235 @@
// Copyright (c) 2020 Can Boluk and contributors of the VTIL Project
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
// 3. Neither the name of mosquitto nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
#include "branch_correction_pass.hpp"
#include <vtil/symex>
#include <algorithm>
#include <future>
#include "../common/auxiliaries.hpp"
namespace vtil::optimizer
{
// Implement the pass.
//
size_t branch_correction_pass::pass( basic_block* blk, bool xblock )
{
// If block is not complete or not cross-block, skip.
//
if ( !blk->is_complete() || !xblock )
return 0;
size_t cnt = 0;
// Analyse the branching instruction.
//
auto branch_info = aux::analyze_branch( blk, &ctracer, true );
// If branching to real, assert single next block.
//
auto branch = std::prev( blk->end() );
if ( branch->base->is_branching_real() )
{
fassert( blk->next.size() <= 1 );
}
// If branching to virtual instruction:
//
else
{
fassert( branch->base->is_branching_virt() );
// For each destination block:
//
for ( auto it = blk->next.begin(); it != blk->next.end(); )
{
// Check if this destination is plausible or not.
//
vip_t target = ( *it )->entry_vip;
bool impossible = true;
for ( auto& branch : branch_info.destinations )
impossible &= ( branch != target ).get<bool>().value_or( false );
// If it is not:
//
if ( impossible )
{
// Delete prev and next links.
//
( *it )->prev.erase( std::remove( ( *it )->prev.begin(), ( *it )->prev.end(), blk ), ( *it )->prev.end() );
it = blk->next.erase( it );
// Increment counter and continue.
//
cnt++;
continue;
}
// Otherwise increment iterator and continue.
//
++it;
}
}
// If branch is jmp where it could be jcc:
//
if ( branch_info.is_jcc && *branch->base == ins::jmp )
{
// Demote conditional to local block:
//
cached_tracer local_tracer;
auto lbranch_info = aux::analyze_branch( blk, &local_tracer, false, false );
if ( lbranch_info.is_jcc )
{
// Attempts to revive an expression via cache.
//
const auto revive_via_cache = [ & ] ( const symbolic::expression& exp, cached_tracer* tr ) -> std::future<operand>
{
// If immediate return as is.
//
if ( exp.is_constant() )
return std::async( std::launch::deferred, [ op = operand{ *exp.get<uint64_t>(), exp.size() } ]() { return op; } );
// If expression is not a register:
//
symbolic::variable var_reg;
if ( !exp.is_variable() || !exp.uid.get<symbolic::variable>().is_register() )
{
// Iterate cache entries:
//
std::shared_lock _g{ tr->mtx };
for ( auto& [var, ex] : tr->cache )
{
// Skip if memory variable or has invalid iterator.
//
if ( var.is_memory() || !var.at.is_valid() )
continue;
// If expressions are not identical skip.
//
if ( !ex->is_identical( exp ) )
continue;
// Set var_reg and break.
//
var_reg = var;
break;
}
}
else
{
var_reg = exp.uid.get<symbolic::variable>();
}
// Fail if invalid.
//
if ( !var_reg.is_valid() )
return {};
// Check if alive, if not revive, else return as is.
//
if ( aux::is_alive( var_reg, branch, &ctracer ) )
return std::async( std::launch::deferred, [ op = operand{ var_reg.reg() } ]() { return op; } );
else
return std::async( std::launch::deferred, [ = ]() -> operand { return aux::revive_register( var_reg, branch ); } );
};
// Convert [cc] [d1] [d2] in order.
//
auto op_cc = revive_via_cache( lbranch_info.cc, &local_tracer );
if ( op_cc.valid() )
{
bool fail = false;
std::future<operand> dsts[ 2 ];
for ( auto [out, blocal, bglobal] : zip( dsts, lbranch_info.destinations, branch_info.destinations ) )
{
std::future<operand> op;
if ( blocal.complexity <= bglobal.complexity )
op = revive_via_cache( blocal, &local_tracer );
else
op = revive_via_cache( bglobal, &ctracer );
if ( !op.valid() )
{
fail = true;
break;
}
out = std::move(op);
}
// If we converted all succesfully:
//
if ( !fail )
{
branch->base = &ins::js;
branch->operands = {
op_cc.get(),
dsts[ 0 ].get(),
dsts[ 1 ].get()
};
cnt++;
}
}
else
{
logger::log( "opcc invalid!\n" );
}
}
}
// If branch is [j/c* reg] where it could be [j/c* reg] imm:
//
if ( branch_info.destinations.size() == 1 &&
branch_info.destinations[ 0 ].is_constant() &&
( *branch->base == ins::jmp || *branch->base == ins::vxcall || *branch->base == ins::vexit ) &&
branch->operands[ 0 ].is_register() )
{
branch->operands[ 0 ] = { *branch_info.destinations[ 0 ].get<vtil::vip_t>(), 64 };
cnt++;
}
return cnt;
}
size_t branch_correction_pass::xpass( routine* rtn )
{
// Invoke original method, if any removed:
//
if ( size_t cnt = pass_interface<>::xpass( rtn ) )
{
// Delete non-referenced blocks entirely.
//
for ( auto it = rtn->explored_blocks.begin(); it != rtn->explored_blocks.end(); )
{
if ( it->second->prev.size() == 0 )
it = rtn->explored_blocks.erase( it );
else
++it;
}
// Return counter as is.
//
return cnt;
}
return 0;
}
};

View file

@ -32,13 +32,15 @@
namespace vtil::optimizer
{
// Eliminates each prev-next link where the jump is not possible after
// the elimination of opaque predicates, removes the entire block if
// it was left without any references.
// This pass serves two purposes:
// 1) Eliminates each prev-next link where the jump is not possible after
// the elimination of opaque predicates, removes the entire block if
// it was left without any references.
// 2) Converts jmps to jccs where it can be inferred.
//
struct opaque_predicate_elimination_pass : pass_interface<>
struct branch_correction_pass : pass_interface<>
{
std::mutex mtx;
std::shared_mutex mutex;
cached_tracer ctracer = {};
size_t pass( basic_block* blk, bool xblock = false ) override;

View file

@ -1,115 +0,0 @@
// Copyright (c) 2020 Can Boluk and contributors of the VTIL Project
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
// 3. Neither the name of mosquitto nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.
//
#include "opaque_predicate_elimination_pass.hpp"
#include <vtil/symex>
#include <algorithm>
#include "../common/auxiliaries.hpp"
namespace vtil::optimizer
{
// Implement the pass.
//
size_t opaque_predicate_elimination_pass::pass( basic_block* blk, bool xblock )
{
// If block is not complete, skip.
//
if ( !blk->is_complete() )
return 0;
// If branching to real, skip.
//
auto branch = std::prev( blk->end() );
if ( !branch->base->is_branching_virt() )
return 0;
// Discover each possible branch target.
//
cached_tracer tmp = {};
auto branch_targets = aux::discover_branches( blk, xblock ? &ctracer : &tmp, xblock );
// For each destination block:
//
size_t cnt = 0;
for ( auto it = blk->next.begin(); it != blk->next.end(); )
{
// Check if this destination is plausible or not.
//
vip_t target = ( *it )->entry_vip;
bool impossible = true;
for ( auto& [real, dst] : branch_targets )
if ( !real )
impossible &= ( dst != target ).get<bool>().value_or( false );
// If it is not:
//
if ( impossible )
{
// Delete prev and next links.
//
( *it )->prev.erase( std::remove( ( *it )->prev.begin(), ( *it )->prev.end(), blk ), ( *it )->prev.end() );
it = blk->next.erase( it );
// Increment counter and continue.
//
cnt++;
continue;
}
// Otherwise increment iterator and continue.
//
++it;
}
return cnt;
}
size_t opaque_predicate_elimination_pass::xpass( routine* rtn )
{
// Flush cached tracer.
//
ctracer.flush();
// Invoke original method, if any removed:
//
if ( size_t cnt = pass_interface<>::xpass( rtn ) )
{
// Delete non-referenced blocks entirely.
//
for ( auto it = rtn->explored_blocks.begin(); it != rtn->explored_blocks.end(); )
{
if ( it->second->prev.size() == 0 )
it = rtn->explored_blocks.erase( it );
else
++it;
}
// Return counter as is.
//
return cnt;
}
return 0;
}
};