diff --git a/VTIL-Compiler/optimizer/bblock_thunk_removal_pass.cpp b/VTIL-Compiler/optimizer/bblock_thunk_removal_pass.cpp index 991d7ce..978c6ed 100644 --- a/VTIL-Compiler/optimizer/bblock_thunk_removal_pass.cpp +++ b/VTIL-Compiler/optimizer/bblock_thunk_removal_pass.cpp @@ -26,6 +26,7 @@ // POSSIBILITY OF SUCH DAMAGE. // #include "bblock_thunk_removal_pass.hpp" +#include 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 ); } diff --git a/VTIL-Tests/dummy.cpp b/VTIL-Tests/dummy.cpp index 4af75ff..251a94f 100644 --- a/VTIL-Tests/dummy.cpp +++ b/VTIL-Tests/dummy.cpp @@ -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__);