fix(optimizer): prevent unsafe thunk folding from deleting valid blocks ( #80 )

This commit is contained in:
dword64 2026-04-04 20:31:21 +02:00
parent 667eea2b1e
commit 81144e520a
2 changed files with 150 additions and 56 deletions

View file

@ -26,6 +26,7 @@
// POSSIBILITY OF SUCH DAMAGE.
//
#include "bblock_thunk_removal_pass.hpp"
#include <algorithm>
namespace vtil::optimizer
{
@ -64,70 +65,105 @@ namespace vtil::optimizer
fassert(blk->sp_offset == 0);
basic_block* next = blk->next[0];
basic_block* prev = blk->prev[0];
basic_block* prev = blk->prev[0];
bool can_fold = prev != next;
// Remove the block from the next hierarchy
//
for (auto& it : prev->next)
if (can_fold)
{
if (it->entry_vip == blk->entry_vip)
it = next;
}
auto prev_branch = prev->back();
fassert(prev_branch.base->is_branching_virt());
// Remove the block from the prev hierarchy
//
for (auto& it : next->prev)
{
if (it->entry_vip == blk->entry_vip)
it = prev;
}
fassert(prev->back().base->branch_operands_vip.size() != 0);
//Regular loop, because we need the operand index for make_mutable
//
for (size_t i = 0; i < prev->back().operands.size(); i++)
{
if (!prev->back().operands[i].is_immediate())
continue;
auto& ins = make_mutable(prev->back());
if (ins.operands[i].imm().ival == blk->entry_vip)
{
ins.operands[i].imm().ival = next->entry_vip;
}
}
// TODO: should we do this with another operands loop? We currently only have "js" that qualifies so a simple if does the job for now
//
auto branching_instruction = prev->back();
if (branching_instruction.base == &ins::js)
{
fassert(branching_instruction.operands.size() == 3);
// This pass should not interfer with blocks that aren't already touched by branch correction / our corrections above
// If any virtual branch target operand is register-based, we cannot
// safely retarget old_vip->new_vip for this predecessor.
//
if (branching_instruction.operands[1].is_immediate() && branching_instruction.operands[2].is_immediate())
for (int idx : prev_branch.base->branch_operands_vip)
{
if (branching_instruction.operands[1].imm().ival == branching_instruction.operands[2].imm().ival)
fassert(idx >= 0 && (size_t)idx < prev_branch.operands.size());
if (!prev_branch.operands[idx].is_immediate())
{
auto new_vip = branching_instruction.operands[1].imm();
auto ins = std::prev(prev->end());
(+ins)->base = &ins::jmp;
(+ins)->operands.resize(1);
(+ins)->operands[0] = { new_vip.ival, arch::bit_count };
prev->next.resize(1);
next->prev.erase(std::find(next->prev.begin(), next->prev.end(), prev));
can_fold = false;
break;
}
}
}
obsolete_blocks.emplace(blk);
counter++;
}
size_t rewritten_targets = 0;
if (can_fold)
{
// Retarget only branch target operands and ensure at least one
// successful rewrite happened before touching CFG links.
//
auto& mut_branch = make_mutable(prev->back());
for (int idx : mut_branch.base->branch_operands_vip)
{
fassert(idx >= 0 && (size_t)idx < mut_branch.operands.size());
auto& op = mut_branch.operands[idx];
if (op.is_immediate() && op.imm().ival == blk->entry_vip)
{
op.imm().ival = next->entry_vip;
rewritten_targets++;
}
}
can_fold = rewritten_targets != 0;
}
if (can_fold)
{
auto dedup_links = [] (auto& links)
{
for (auto it = links.begin(); it != links.end();)
{
if (std::find(links.begin(), it, *it) != it)
it = links.erase(it);
else
++it;
}
};
// Rewire predecessor/successor links and deduplicate adjacency.
//
prev->next.erase(std::remove(prev->next.begin(), prev->next.end(), blk), prev->next.end());
if (std::find(prev->next.begin(), prev->next.end(), next) == prev->next.end())
prev->next.emplace_back(next);
dedup_links(prev->next);
next->prev.erase(std::remove(next->prev.begin(), next->prev.end(), blk), next->prev.end());
if (std::find(next->prev.begin(), next->prev.end(), prev) == next->prev.end())
next->prev.emplace_back(prev);
dedup_links(next->prev);
// TODO: should we do this with another operands loop? We currently only have "js" that qualifies so a simple if does the job for now
//
auto branching_instruction = prev->back();
if (branching_instruction.base == &ins::js)
{
fassert(branching_instruction.operands.size() == 3);
// This pass should not interfer with blocks that aren't already touched by branch correction / our corrections above
//
if (branching_instruction.operands[1].is_immediate() && branching_instruction.operands[2].is_immediate())
{
if (branching_instruction.operands[1].imm().ival == branching_instruction.operands[2].imm().ival)
{
auto new_vip = branching_instruction.operands[1].imm();
auto ins = std::prev(prev->end());
(+ins)->base = &ins::jmp;
(+ins)->operands.resize(1);
(+ins)->operands[0] = { new_vip.ival, arch::bit_count };
prev->next.clear();
prev->next.emplace_back(next);
dedup_links(next->prev);
}
}
}
obsolete_blocks.emplace(blk);
counter++;
}
}
// Recurse into destinations:
@ -151,6 +187,9 @@ namespace vtil::optimizer
{
// Invoke recursive optimizer starting from entry point.
//
visited.clear();
obsolete_blocks.clear();
first_block = nullptr;
visited.reserve( rtn->num_blocks() );
return pass( rtn->entry_point, true );
}

View file

@ -754,6 +754,61 @@ DOCTEST_TEST_CASE("Optimization bblock_thunk_removal_pass")
}
}
DOCTEST_TEST_CASE("Regression #80: avoid incorrect block removal")
{
vtil::logger::log("\n\n>> %s \n", __FUNCTION__);
auto block = test_helpers::begin_block(0x0ull);
auto [t0, t1, t2, t3] = block->tmp(64, 64, 1, 64);
auto rtn = block->owner;
block->mov(t0, vtil::REG_FLAGS);
block->bnot(t0);
block->ifs(t1, t0.select(1, 2), 0x1000ull);
block->mov(t2, t0.select(1, 2));
block->bnot(t2);
block->ifs(t3, t2, 0x2000ull);
block->add(t1, t3);
block->add(t1, vtil::REG_IMGBASE);
block->jmp(t1);
if (auto block_1000 = block->fork(0x1000ull))
{
block_1000->jmp(0x3000ull);
block_1000->fork(0x3000ull);
}
if (auto block_2000 = block->fork(0x2000ull))
{
block_2000->jmp(0x3000ull);
block_2000->fork(0x3000ull);
}
if (auto block_3000 = rtn->get_block(0x3000ull))
{
block_3000->vexit(uintptr_t(0xdeadc0de));
}
vtil::logger::log(":: Before:\n");
vtil::debug::dump(rtn);
CHECK(rtn->num_blocks() == 4);
CHECK(rtn->get_block(0x0ull) != nullptr);
CHECK(rtn->get_block(0x1000ull) != nullptr);
CHECK(rtn->get_block(0x2000ull) != nullptr);
CHECK(rtn->get_block(0x3000ull) != nullptr);
vtil::optimizer::bblock_thunk_removal_pass{}(rtn);
vtil::optimizer::branch_correction_pass{}(rtn);
vtil::logger::log(":: After:\n");
vtil::debug::dump(rtn);
CHECK(rtn->num_blocks() == 4);
CHECK(rtn->get_block(0x0ull) != nullptr);
CHECK(rtn->get_block(0x1000ull) != nullptr);
CHECK(rtn->get_block(0x2000ull) != nullptr);
CHECK(rtn->get_block(0x3000ull) != nullptr);
}
DOCTEST_TEST_CASE("Optimization branch_correction_pass")
{
vtil::logger::log("\n\n>> %s \n", __FUNCTION__);