Skip to content

Commit

Permalink
Remove relocation masking
Browse files Browse the repository at this point in the history 
Because of the extra constraints we can mask off entire instructions, this allows us to use signatures across units with differing relocations.
  • Loading branch information
@emesare
emesare committed Oct 23, 2024
1 parent 18a0434 commit db84122
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 103 deletions.
53 changes: 9 additions & 44 deletions warp_binja/src/cache.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,10 @@ use signaturebuild::function::constraints::FunctionConstraint;
use signaturebuild::prelude::*;

use crate::convert::from_bn_symbol;
use crate::{build_function, entry_basic_block_guid, function_guid, Relocations};
use crate::{build_function, function_guid};

pub static FUNCTION_CACHE: OnceLock<DashMap<ViewID, FunctionCache>> = OnceLock::new();
pub static GUID_CACHE: OnceLock<DashMap<ViewID, GUIDCache>> = OnceLock::new();
pub static RELOC_CACHE: OnceLock<DashMap<ViewID, Relocations>> = OnceLock::new();

pub fn cached_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
function: &BNFunction,
Expand All @@ -40,25 +39,6 @@ pub fn cached_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>
}
}

pub fn cached_function_entry_block<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
function: &BNFunction,
llil: &llil::Function<A, M, NonSSA<V>>,
) -> Option<BasicBlock> {
let view = function.view();
let view_id = ViewID::from(view.as_ref());
let reloc_cache = RELOC_CACHE.get_or_init(Default::default);
let guid = match reloc_cache.get(&view_id) {
Some(relocations) => entry_basic_block_guid(function, &relocations, llil),
None => {
let relocations = view.get_relocation_ranges();
let guid = entry_basic_block_guid(function, &relocations, llil);
reloc_cache.insert(view_id, relocations);
guid
}
};
Some(BasicBlock::new(guid?))
}

pub fn cached_call_site_constraints(function: &BNFunction) -> HashSet<FunctionConstraint> {
let view = function.view();
let view_id = ViewID::from(view);
Expand Down Expand Up @@ -225,30 +205,15 @@ impl GUIDCache {
function: &BNFunction,
llil: &llil::Function<A, M, NonSSA<V>>,
) -> Option<FunctionGUID> {
let function_guid_with_relocs = |relocations: &Relocations| {
let function_id = FunctionID::from(function);
match self.cache.try_get_mut(&function_id) {
TryResult::Present(function_guid) => function_guid.value().to_owned(),
TryResult::Absent => {
let function_guid = function_guid(function, relocations, llil);
self.cache.insert(function_id, function_guid);
function_guid
}
TryResult::Locked => function_guid(function, relocations, llil),
}
};

let view = function.view();
let view_id = ViewID::from(view.as_ref());
let reloc_cache = RELOC_CACHE.get_or_init(Default::default);
match reloc_cache.get(&view_id) {
Some(relocations) => function_guid_with_relocs(&relocations),
None => {
let relocations = view.get_relocation_ranges();
let guid = function_guid_with_relocs(&relocations);
reloc_cache.insert(view_id, relocations);
guid
let function_id = FunctionID::from(function);
match self.cache.try_get_mut(&function_id) {
TryResult::Present(function_guid) => function_guid.value().to_owned(),
TryResult::Absent => {
let function_guid = function_guid(function, llil);
self.cache.insert(function_id, function_guid);
function_guid
}
TryResult::Locked => function_guid(function, llil),
}
}
}
Expand Down
52 changes: 5 additions & 47 deletions warp_binja/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,9 @@ use binaryninja::llil;
use binaryninja::llil::{ExprInfo, FunctionMutability, NonSSA, NonSSAVariant, VisitorAction};
use binaryninja::rc::Ref as BNRef;
use signaturebuild::prelude::*;
use std::cmp::Ordering;
use std::ops::Range;

use crate::cache::{
cached_adjacency_constraints, cached_call_site_constraints, cached_function_entry_block,
cached_function_guid,
cached_adjacency_constraints, cached_call_site_constraints, cached_function_guid,
};
use crate::convert::{from_bn_symbol, from_bn_type};

Expand All @@ -21,8 +18,6 @@ mod matcher;
/// Only used when compiled for cdylib target.
mod plugin;

type Relocations = Vec<Range<u64>>;

pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
func: &BNFunction,
llil: &llil::Function<A, M, NonSSA<V>>,
Expand All @@ -42,18 +37,17 @@ pub fn build_function<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
caller_sites: Default::default(),
},
// TODO: We need more than one entry block.
entry: cached_function_entry_block(func, llil),
entry: entry_basic_block_guid(func, llil).map(BasicBlock::new),
})
}

pub fn entry_basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
func: &BNFunction,
relocations: &Relocations,
llil: &llil::Function<A, M, NonSSA<V>>,
) -> Option<BasicBlockGUID> {
// NOTE: This is not actually the entry point. This is the highest basic block.
let first_basic_block = sorted_basic_blocks(func).into_iter().next()?;
basic_block_guid(&first_basic_block, relocations, llil)
basic_block_guid(&first_basic_block, llil)
}

/// Basic blocks sorted from high to low.
Expand All @@ -69,21 +63,19 @@ pub fn sorted_basic_blocks(func: &BNFunction) -> Vec<BNRef<BNBasicBlock<NativeBl

pub fn function_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
func: &BNFunction,
relocations: &Relocations,
llil: &llil::Function<A, M, NonSSA<V>>,
) -> Option<FunctionGUID> {
// TODO: Sort the basic blocks.
let basic_blocks = sorted_basic_blocks(func);
let basic_block_guids = basic_blocks
.iter()
.filter_map(|bb| basic_block_guid(bb, relocations, llil))
.filter_map(|bb| basic_block_guid(bb, llil))
.collect::<Vec<_>>();
Some(FunctionGUID::from_basic_blocks(&basic_block_guids))
}

pub fn basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant>(
basic_block: &BNBasicBlock<NativeBlock>,
relocations: &Relocations,
llil: &llil::Function<A, M, NonSSA<V>>,
) -> Option<BasicBlockGUID> {
let func = basic_block.function();
Expand All @@ -97,47 +89,13 @@ pub fn basic_block_guid<A: Architecture, M: FunctionMutability, V: NonSSAVariant
// TODO: Could we keep the bytes and just zero mask them? At least then we don't completely get rid of them.

let basic_block_range = basic_block.raw_start()..basic_block.raw_end();
let basic_blocks_relocs = relocations
.iter()
.filter(|r| basic_block_range.contains(&r.start))
.collect::<Vec<_>>();

let mut basic_block_bytes = Vec::with_capacity(basic_block_range.count());
for instr_addr in basic_block.into_iter() {
let mut instr_bytes = view.read_vec(instr_addr, max_instr_len);
if let Some(instr_info) = arch.instruction_info(&instr_bytes, instr_addr) {
let instr_len = instr_info.len();
instr_bytes.truncate(instr_len);
let instr_range = instr_addr..(instr_addr + instr_len as u64);
// Check to see if instruction contains the start to a relocation
// TODO: What if it contains a partial relocation?
if let Some(instr_reloc) = basic_blocks_relocs
.iter()
.find(|r| instr_range.contains(&r.start))
{
// Found a relocatable instruction, mask off just the relocatable bytes
let mut reloc_start: usize = instr_reloc.start.saturating_sub(instr_addr) as usize;
let mut reloc_end: usize = instr_reloc.end.saturating_sub(instr_addr) as usize;
match reloc_end.cmp(&reloc_start) {
Ordering::Less => {
// If the relocation start and end are backwards (starts at 4, ends at 0), swap it.
std::mem::swap(&mut reloc_start, &mut reloc_end);
}
Ordering::Equal => {
// Relocation is not well-formed, mask off entire instruction.
log::warn!(
"Relocation not well-formed! Masking off entire instruction... 0x{:x}",
instr_addr
);
reloc_start = 0;
reloc_end = instr_len;
}
Ordering::Greater => {
// Relocation is well-formed, start < end
}
}
instr_bytes[reloc_start..reloc_end.min(instr_len)].fill(0);
} else if let Some(instr_llil) = llil.instruction_at(instr_addr) {
if let Some(instr_llil) = llil.instruction_at(instr_addr) {
if instr_llil.visit_tree(&mut |_expr, expr_info| match expr_info {
ExprInfo::ConstPtr(_) | ExprInfo::ExternPtr(_) => VisitorAction::Halt,
_ => VisitorAction::Descend,
Expand Down
23 changes: 17 additions & 6 deletions warp_binja/src/matcher.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::cmp::Ordering;
use dashmap::DashMap;
use fastbloom::BloomFilter;
use std::cmp::Ordering;
use std::collections::{HashMap, HashSet};
use std::hash::{DefaultHasher, Hasher};
use std::path::PathBuf;
Expand All @@ -18,14 +18,14 @@ use binaryninja::rc::Guard;
use binaryninja::rc::Ref as BNRef;

use signaturebuild::function::{Function, FunctionGUID};
use signaturebuild::prelude::BasicBlock;
use signaturebuild::Data;
use typebuild::guid::TypeGUID;
use typebuild::prelude::{Type, TypeClass};

use crate::cache::{
cached_call_site_constraints, cached_function_entry_block, cached_function_guid, FunctionID,
};
use crate::cache::{cached_call_site_constraints, cached_function_guid, FunctionID};
use crate::convert::to_bn_type;
use crate::entry_basic_block_guid;
use crate::plugin::on_matched_function;

pub const TRIVIAL_LLIL_THRESHOLD: usize = 8;
Expand Down Expand Up @@ -110,6 +110,11 @@ impl Matcher {
map
})
})
.map(|(guid, mut funcs)| {
funcs.sort_by_key(|f| f.symbol.name.to_owned());
funcs.dedup_by_key(|f| f.symbol.name.to_owned());
(guid, funcs)
})
.collect();

task.set_progress_text("Gathering matcher types...");
Expand Down Expand Up @@ -279,7 +284,7 @@ impl Matcher {
let is_function_trivial = { llil.instruction_count() < TRIVIAL_LLIL_THRESHOLD };

// Check to see if the functions entry block is even in the dataset
let entry_block = cached_function_entry_block(function, llil);
let entry_block = entry_basic_block_guid(function, llil).map(BasicBlock::new);
if self.basic_block_filter.contains(&entry_block) {
// Build the full function guid now
if let Some(warp_func_guid) = cached_function_guid(function, llil) {
Expand All @@ -295,6 +300,12 @@ impl Matcher {
function.start()
);
on_new_match(matched_function);
} else {
log::error!(
"Failed to find matching function `{}`... 0x{:x}",
matched.len(),
function.start()
);
}
}
}
Expand Down Expand Up @@ -334,7 +345,7 @@ impl Matcher {
match common_guid_count.cmp(&highest_guid_count) {
Ordering::Equal => {
// Multiple matches with same count, don't match on ONE of them.
matched_guid_func = None;
matched_guid_func = None;
}
Ordering::Greater => {
highest_guid_count = common_guid_count;
Expand Down
7 changes: 1 addition & 6 deletions warp_binja/src/plugin.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use binaryninja::rc::Ref;
use binaryninja::tags::TagType;

use crate::build_function;
use crate::cache::{ViewID, FUNCTION_CACHE, GUID_CACHE, RELOC_CACHE};
use crate::cache::{ViewID, FUNCTION_CACHE, GUID_CACHE};
use crate::convert::{to_bn_symbol_at_address, to_bn_type};
use crate::matcher::{PlatformID, PLAT_MATCHER_CACHE};

Expand Down Expand Up @@ -79,11 +79,6 @@ impl Command for DebugCache {
log::info!("View function guids: {}", cache.cache.len());
}

let relocation_cache = RELOC_CACHE.get_or_init(Default::default);
if let Some(cache) = relocation_cache.get(&view_id) {
log::info!("Relocations: {}", cache.len());
}

let plat_cache = PLAT_MATCHER_CACHE.get_or_init(Default::default);
if let Some(plat) = view.default_platform() {
let platform_id = PlatformID::from(plat);
Expand Down

0 comments on commit db84122

Please sign in to comment.