Skip to content

Commit

Permalink
Implement patching commands
Browse files Browse the repository at this point in the history 
These commands allow to replace removed instructions with NOPs
  • Loading branch information
@ergrelet
ergrelet committed Jul 17, 2024
1 parent 4b4413c commit eaa5376
Show file tree
Hide file tree
Showing 5 changed files with 133 additions and 56 deletions.
138 changes: 103 additions & 35 deletions src/commands.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,30 +13,71 @@ namespace triton_bn {

using namespace BinaryNinja;

static std::vector<MetaBasicBlock> SimplifyBasicBlockCommon(
BinaryNinja::BinaryView* p_view, bool padding);
static std::vector<MetaBasicBlock> SimplifyFunctionCommon(BinaryView* p_view,
bool padding);
static FlowGraph* GenerateFlowGraphFromMetaBasicBlocks(
std::vector<MetaBasicBlock> basic_blocks);

void SimplifyBasicBlockCommand(BinaryNinja::BinaryView* p_view) {
BinaryView& view = *p_view;
void SimplifyBasicBlockPreviewCommand(BinaryNinja::BinaryView* p_view) {
auto simplified_basic_blocks = SimplifyBasicBlockCommon(p_view, false);
if (simplified_basic_blocks.empty()) {
LogError("Failed to simplify basic block");
return;
}

// Construct result flow graph and display it
FlowGraph* flow_graph =
GenerateFlowGraphFromMetaBasicBlocks(std::move(simplified_basic_blocks));

const std::string report_title = fmt::format(
"Simplified basic block (0x{:x})", simplified_basic_blocks[0].GetStart());
p_view->ShowGraphReport(report_title, flow_graph);

LogInfo("Basic block has been simplified and preview rendered");
}

void SimplifyBasicBlockPatchCommand(BinaryNinja::BinaryView* p_view) {
auto simplified_basic_blocks = SimplifyBasicBlockCommon(p_view, true);
if (simplified_basic_blocks.empty()) {
LogError("Failed to simplify basic block");
return;
}

// Patch code
for (auto& basic_block : simplified_basic_blocks) {
for (auto& instruction : basic_block.triton_bb().getInstructions()) {
p_view->Write(instruction.getAddress(), instruction.getOpcode(),
instruction.getSize());
}
}
// Rerun analysis
p_view->UpdateAnalysis();

LogInfo("Basic block has been simplified and patches applied");
}

static std::vector<MetaBasicBlock> SimplifyBasicBlockCommon(
BinaryNinja::BinaryView* p_view, bool padding) {
// Get currently selected address in the view
const auto current_offset = view.GetCurrentOffset();
const auto current_offset = p_view->GetCurrentOffset();
LogDebug("Current offset=0x%p", (void*)current_offset);

// Find the function in which this address resides
const auto candidate_basic_blocks =
view.GetBasicBlocksForAddress(current_offset);
p_view->GetBasicBlocksForAddress(current_offset);
if (candidate_basic_blocks.empty()) {
LogError("Failed to find the currently selected basic block");
return;
return {};
}
// TODO: Alert the users if multiple candidate basic blocks exist?
const auto basic_block = candidate_basic_blocks[0];
LogDebug("Current basic block=0x%p", (void*)basic_block->GetStart());

// Determine the current platform/architecture
const std::string architecture_name =
view.GetDefaultArchitecture()->GetName();
p_view->GetDefaultArchitecture()->GetName();
LogDebug("Architecture is '%s'", architecture_name.c_str());

triton::Context triton{};
Expand All @@ -48,50 +89,87 @@ void SimplifyBasicBlockCommand(BinaryNinja::BinaryView* p_view) {
triton.setArchitecture(triton::arch::ARCH_AARCH64);
} else {
LogError("Unsupported architecture '%s'", architecture_name.c_str());
return;
return {};
}

auto meta_basic_blocks =
ExtractMetaBasicBlocksFromBasicBlock(view, basic_block, triton);
ExtractMetaBasicBlocksFromBasicBlock(*p_view, basic_block, triton);

// Simplify basic block
auto simplified_basic_blocks =
SimplifyMetaBasicBlocks(triton, std::move(meta_basic_blocks));
return SimplifyMetaBasicBlocks(triton, std::move(meta_basic_blocks), padding);
}

bool ValidateSimplifyBasicBlockCommand(BinaryView* p_view) {
return ValidateSimplifyFunctionCommand(p_view);
}

void SimplifyFunctionPreviewCommand(BinaryView* p_view) {
auto simplified_basic_blocks = SimplifyFunctionCommon(p_view, false);
if (simplified_basic_blocks.empty()) {
LogError("Failed to simplify function");
return;
}

// Construct result flow graph and display it
const auto candidate_functions =
p_view->GetAnalysisFunctionsContainingAddress(
simplified_basic_blocks[0].GetStart());
if (candidate_functions.empty()) {
LogError("Failed to simplify function");
return;
}

const std::string current_function_name =
candidate_functions[0]->GetSymbol()->GetFullName();
const std::string report_title =
fmt::format("Simplified function ({})", current_function_name);
FlowGraph* flow_graph =
GenerateFlowGraphFromMetaBasicBlocks(std::move(simplified_basic_blocks));
p_view->ShowGraphReport(report_title, flow_graph);

const std::string report_title =
fmt::format("Simplified basic block (0x{:x})", basic_block->GetStart());
view.ShowGraphReport(report_title, flow_graph);
LogInfo("Function has been simplified and preview rendered");
}

bool ValidateSimplifyBasicBlockCommand(BinaryView* p_view) {
return ValidateSimplifyFunctionCommand(p_view);
}
void SimplifyFunctionPatchCommand(BinaryView* p_view) {
auto simplified_basic_blocks = SimplifyFunctionCommon(p_view, true);
if (simplified_basic_blocks.empty()) {
LogError("Failed to simplify function");
return;
}

void SimplifyFunctionCommand(BinaryView* p_view) {
BinaryView& view = *p_view;
// Patch code
for (auto& basic_block : simplified_basic_blocks) {
for (auto& instruction : basic_block.triton_bb().getInstructions()) {
p_view->Write(instruction.getAddress(), instruction.getOpcode(),
instruction.getSize());
}
}
// Rerun analysis
p_view->UpdateAnalysis();

LogInfo("Function has been simplified and patches applied");
}

static std::vector<MetaBasicBlock> SimplifyFunctionCommon(BinaryView* p_view,
bool padding) {
// Get currently selected address in the view
const auto current_offset = view.GetCurrentOffset();
const auto current_offset = p_view->GetCurrentOffset();
LogDebug("Current offset=0x%p", (void*)current_offset);

// Find the function in which this address resides
const auto candidate_functions =
view.GetAnalysisFunctionsContainingAddress(current_offset);
p_view->GetAnalysisFunctionsContainingAddress(current_offset);
if (candidate_functions.empty()) {
LogError("Failed to find the currently selected function");
return;
return {};
}
// TODO: Alert the users if multiple candidate functions exist
const auto current_function = candidate_functions[0];
LogDebug("Current function=0x%p", (void*)current_function->GetStart());

// Determine the current architecture
const std::string architecture_name =
view.GetDefaultArchitecture()->GetName();
p_view->GetDefaultArchitecture()->GetName();
LogDebug("Architecture is '%s'", architecture_name.c_str());

// Intialize Triton's context
Expand All @@ -104,12 +182,12 @@ void SimplifyFunctionCommand(BinaryView* p_view) {
triton.setArchitecture(triton::arch::ARCH_AARCH64);
} else {
LogError("Unsupported architecture '%s'", architecture_name.c_str());
return;
return {};
}

// Create `MetaBasicBlock`s from the current Binja function's basic blocks
auto meta_basic_blocks =
ExtractMetaBasicBlocksFromFunction(view, current_function, triton);
ExtractMetaBasicBlocksFromFunction(*p_view, current_function, triton);
LogDebug("%zu meta basic block(s) extracted", meta_basic_blocks.size());

if (Settings::Instance()->Get<bool>("triton-bn.mergeBasicBlocks")) {
Expand All @@ -118,17 +196,7 @@ void SimplifyFunctionCommand(BinaryView* p_view) {
}

// Simplify basic blocks
auto simplified_basic_blocks =
SimplifyMetaBasicBlocks(triton, std::move(meta_basic_blocks));

// Construct result flow graph and display it
const std::string current_function_name =
current_function->GetSymbol()->GetFullName();
const std::string report_title =
fmt::format("Simplified function ({})", current_function_name);
FlowGraph* flow_graph =
GenerateFlowGraphFromMetaBasicBlocks(std::move(simplified_basic_blocks));
view.ShowGraphReport(report_title, flow_graph);
return SimplifyMetaBasicBlocks(triton, std::move(meta_basic_blocks), padding);
}

bool ValidateSimplifyFunctionCommand(BinaryView* p_view) {
Expand Down
6 changes: 4 additions & 2 deletions src/commands.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@

namespace triton_bn {

void SimplifyBasicBlockCommand(BinaryNinja::BinaryView* p_view);
void SimplifyBasicBlockPreviewCommand(BinaryNinja::BinaryView* p_view);
void SimplifyBasicBlockPatchCommand(BinaryNinja::BinaryView* p_view);
bool ValidateSimplifyBasicBlockCommand(BinaryNinja::BinaryView* p_view);

void SimplifyFunctionCommand(BinaryNinja::BinaryView* p_view);
void SimplifyFunctionPreviewCommand(BinaryNinja::BinaryView* p_view);
void SimplifyFunctionPatchCommand(BinaryNinja::BinaryView* p_view);
bool ValidateSimplifyFunctionCommand(BinaryNinja::BinaryView* p_view);

} // namespace triton_bn
19 changes: 14 additions & 5 deletions src/main.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,23 @@ BINARYNINJAPLUGIN bool CorePluginInit() {
"description" : "Automatically merge basic blocks linked with a single unconditional branch before running the simplification passes on functions."
})");

PluginCommand::Register("triton-bn\\Simplify basic block (DSE)",
// Preview commands
PluginCommand::Register("triton-bn\\Preview\\Simplify basic block (DSE)",
"Simplify basic block using Triton's DSE pass",
triton_bn::SimplifyBasicBlockCommand,
triton_bn::SimplifyBasicBlockPreviewCommand,
triton_bn::ValidateSimplifyBasicBlockCommand);

PluginCommand::Register("triton-bn\\Simplify function (DSE)",
PluginCommand::Register("triton-bn\\Preview\\Simplify function (DSE)",
"Simplify function using Triton's DSE pass",
triton_bn::SimplifyFunctionPreviewCommand,
triton_bn::ValidateSimplifyFunctionCommand);
// Patch commands
PluginCommand::Register("triton-bn\\Patch\\Simplify basic block (DSE)",
"Simplify basic block using Triton's DSE pass",
triton_bn::SimplifyBasicBlockPatchCommand,
triton_bn::ValidateSimplifyBasicBlockCommand);
PluginCommand::Register("triton-bn\\Patch\\Simplify function (DSE)",
"Simplify function using Triton's DSE pass",
triton_bn::SimplifyFunctionCommand,
triton_bn::SimplifyFunctionPatchCommand,
triton_bn::ValidateSimplifyFunctionCommand);

return true;
Expand Down
23 changes: 10 additions & 13 deletions src/meta_basic_block.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -235,16 +235,11 @@ static triton::arch::BasicBlock RemoveNopLikeInstructions(
triton::arch::BasicBlock in = triton_bb;
triton::arch::BasicBlock out;

const auto nop_instr = [&]() {
triton::Context tmp_ctx(triton.getArchitecture());
return tmp_ctx.getNopInstruction();
}();

triton::arch::Architecture arch;
arch.setArchitecture(triton.getArchitecture());
const auto nop_instr = arch.getNopInstruction();
const auto& pc_reg = arch.getProgramCounter();

auto last_addr = in.getFirstAddress();
for (auto& instr : in.getInstructions()) {
triton::Context tmp_ctx(triton.getArchitecture());
// Symbolize all registers
Expand Down Expand Up @@ -299,16 +294,16 @@ static triton::arch::BasicBlock RemoveNopLikeInstructions(
LogDebugF("NOP-like instruction removed: '{}'", instr.getDisassembly());
if (padding) {
// Replace with a nop padding of the appropriate size
while (instr.getAddress() > last_addr) {
size_t padding_size = 0;
while (instr.getSize() > padding_size) {
out.add(nop_instr);
last_addr += nop_instr.getSize();
padding_size += nop_instr.getSize();
}
}
} else {
// Instruction has side effects, keep it in the basic block
out.add(instr);
}
last_addr = instr.getNextAddress();
}

return out;
Expand All @@ -317,7 +312,8 @@ static triton::arch::BasicBlock RemoveNopLikeInstructions(
// Simplify the given `MetaBasicBlock`s with Triton's dead store elimination
// pass
std::vector<MetaBasicBlock> SimplifyMetaBasicBlocks(
const triton::Context& triton, std::vector<MetaBasicBlock> basic_blocks) {
const triton::Context& triton, std::vector<MetaBasicBlock> basic_blocks,
bool padding) {
// Simplify basic blocks
std::vector<MetaBasicBlock> simplified_basic_blocks(basic_blocks.size());
{
Expand All @@ -333,9 +329,10 @@ std::vector<MetaBasicBlock> SimplifyMetaBasicBlocks(

// Simplify basic blocks and disassemble the result
try {
auto simplified_triton_bb = triton.simplify(meta_bb.triton_bb());
simplified_triton_bb =
RemoveNopLikeInstructions(triton, simplified_triton_bb);
auto simplified_triton_bb =
triton.simplify(meta_bb.triton_bb(), padding);
simplified_triton_bb = RemoveNopLikeInstructions(
triton, simplified_triton_bb, padding);
triton.disassembly(simplified_triton_bb, meta_bb.GetStart());
meta_bb.set_triton_bb(simplified_triton_bb);
return std::move(meta_bb);
Expand Down
3 changes: 2 additions & 1 deletion src/meta_basic_block.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ std::vector<MetaBasicBlock> MergeMetaBasicBlocks(
std::vector<MetaBasicBlock> basic_blocks);

std::vector<MetaBasicBlock> SimplifyMetaBasicBlocks(
const triton::Context& triton, std::vector<MetaBasicBlock> basic_blocks);
const triton::Context& triton, std::vector<MetaBasicBlock> basic_blocks,
bool padding = false);

} // namespace triton_bn

0 comments on commit eaa5376

Please sign in to comment.