use std::path::Path;
use rustc_hash::FxHashMap;
use crate::birthmarks::{Birthmark, BirthmarkType, Data, Elements, Kgram, Metadata};
use crate::program::{Function, Program};
use crate::{Error, Iterable, Result};
const HASH_PREFIX_LEN: usize = 16;
pub fn dest_file_name(program_path: &Path) -> Result<String> {
let file_name = program_path
.file_stem()
.ok_or_else(|| {
Error::Parse(format!(
"{}: cannot determine the file stem",
program_path.display()
))
})?
.to_string_lossy();
let hash = get_hash(program_path);
let new_filename = format!("{file_name}_{}.json", hash?);
Ok(new_filename)
}
fn get_hash(path: &Path) -> Result<String> {
use sha2::Digest;
use std::io::{Read, Seek};
let pbuf = path.to_path_buf();
let mut file = std::fs::File::open(path).map_err(|e| Error::Io(pbuf.clone(), e))?;
let len = file
.metadata()
.map_err(|e| Error::Io(pbuf.clone(), e))?
.len();
let mut hasher = sha2::Sha256::new();
hasher.update(len.to_le_bytes());
let mut head = vec![0; 4096.min(len as usize)];
file.read_exact(&mut head)
.map_err(|e| Error::Io(pbuf.clone(), e))?;
hasher.update(&head);
if len > 4096 {
let tail_len = 4096.min(len as usize - 4096);
let mut tail = vec![0; tail_len];
file.seek(std::io::SeekFrom::End(-(tail_len as i64)))
.map_err(|e| Error::Io(pbuf.clone(), e))?;
file.read_exact(&mut tail)
.map_err(|e| Error::Io(pbuf.clone(), e))?;
hasher.update(&tail);
}
let hash = hasher.finalize();
Ok(hash
.iter()
.take(HASH_PREFIX_LEN.div_ceil(2))
.map(|b| format!("{b:02x}"))
.collect::<String>()[..HASH_PREFIX_LEN]
.to_string())
}
pub struct Extractor {
bt: BirthmarkType,
}
impl Extractor {
pub fn new(bt: BirthmarkType) -> Self {
Self { bt }
}
pub fn extract<T: crate::Op>(&self, args: Vec<&Program<T>>) -> Result<Vec<Birthmark>> {
let result = args
.iter()
.map(|p| extract_birthmark_op(p, &self.bt))
.collect::<Vec<_>>();
Error::vec_result_to_result_vec(result)
}
pub fn extract_each<T: crate::Op>(&self, p: &Program<T>) -> Result<Birthmark> {
extract_birthmark_op(p, &self.bt)
}
pub fn extract_any(&self, p: &crate::program::AnyProgram) -> Result<Birthmark> {
use crate::program::AnyProgram;
match p {
AnyProgram::GhidraPcode(p) => self.extract_each(p),
AnyProgram::BinaryNinjaLlil(p) => self.extract_each(p),
AnyProgram::BinaryNinjaMlil(p) => self.extract_each(p),
AnyProgram::BinaryNinjaHlil(p) => self.extract_each(p),
AnyProgram::IdaMicrocode(p) => self.extract_each(p),
}
}
}
fn extract_birthmark_op<T: crate::Op>(p: &Program<T>, bt: &BirthmarkType) -> Result<Birthmark> {
let now = std::time::Instant::now();
if matches!(
bt,
BirthmarkType::FcSeq | BirthmarkType::FcSet | BirthmarkType::FcFreq
) && let Some(reason) = empty_family(p)
{
log::warn!("{}", reason.message(p.path(), p.ir()));
}
let elements = p
.iter()
.map(|f| {
let name = f.name().to_string();
let data = match bt {
BirthmarkType::FcFreq => Data::Freq(extract_function_calls_freq(f, p)),
BirthmarkType::FcSet => {
Data::Set(extract_function_calls_freq(f, p).into_keys().collect())
}
BirthmarkType::FcSeq => Data::Seq(extract_function_calls(f, p)),
BirthmarkType::OpFreq => Data::Freq(f.ops_freq()),
BirthmarkType::OpSet => Data::Set(f.ops_freq().into_keys().collect()),
BirthmarkType::OpSeq => Data::Seq(f.ops().map(|s| s.into()).collect()),
BirthmarkType::OpKgramSeq(k) => Data::KgramSeq(extract_op_kgram_seq(f, *k)),
BirthmarkType::OpKgramFreq(k) => Data::KgramFreq(extract_op_kgram_freq(f, *k)),
BirthmarkType::OpKgramSet(k) => {
Data::KgramSet(extract_op_kgram_seq(f, *k).into_iter().collect())
}
};
Elements { name, data }
})
.collect::<Vec<_>>();
let metadata = build_metadata(p, bt.clone(), now);
Ok(Birthmark {
metadata,
elements,
json_path: None,
})
}
fn build_metadata<T>(
p: &Program<T>,
bt: BirthmarkType,
start_time: std::time::Instant,
) -> Metadata {
let extracted_at = chrono::Utc::now();
let duration = start_time.elapsed();
let path = p.path().to_path_buf();
let name = path
.file_name()
.unwrap_or_default()
.to_string_lossy()
.to_string();
Metadata {
file_name: name,
path,
extracted_at,
duration,
birthmark_type: bt,
ir: p.ir(),
}
}
fn extract_op_kgram_seq<T: crate::Op>(f: &Function<T>, k: usize) -> Vec<Kgram> {
f.ops()
.map(|s| s.into())
.collect::<Vec<_>>()
.windows(k)
.map(|w| Kgram::new(w.to_vec()))
.collect()
}
fn extract_op_kgram_freq<T: crate::Op>(f: &Function<T>, k: usize) -> FxHashMap<Kgram, usize> {
seq_to_freq(extract_op_kgram_seq(f, k).into_iter())
}
pub(crate) fn seq_to_freq<T>(seq: impl Iterator<Item = T>) -> FxHashMap<T, usize>
where
T: std::hash::Hash + Eq,
{
seq.into_iter().fold(FxHashMap::default(), |mut acc, item| {
*acc.entry(item).or_insert(0) += 1;
acc
})
}
#[derive(Debug, PartialEq, Eq)]
enum EmptyFamily {
NoCalls,
NoneResolve(usize),
}
impl EmptyFamily {
fn message(&self, path: &std::path::Path, ir: crate::lift::Ir) -> String {
let path = path.display();
match self {
Self::NoCalls => format!(
"{path}: no operation is a call, so every fc-* birthmark of it is empty -- and \
two empty birthmarks score as a perfect match. Either the program really calls \
nothing, or oinkie's reader for {ir} does not recognise that representation's \
call operations"
),
Self::NoneResolve(calls) => format!(
"{path}: calls were found ({calls} in total), but none of them names anything in \
the symbol table, so every fc-* birthmark of it is empty -- and two empty \
birthmarks score as a perfect match. Either every call is indirect, or oinkie's \
reader for {ir} keys the symbol table differently from the way that \
representation renders a callee"
),
}
}
}
fn empty_family<T: crate::Op>(p: &Program<T>) -> Option<EmptyFamily> {
let mut calls = 0usize;
for function in p.iter() {
for op in function.iter() {
if !op.is_call() {
continue;
}
calls += 1;
if op.symbol_key().and_then(|key| p.symbol(&key)).is_some() {
return None;
}
}
}
Some(if calls == 0 {
EmptyFamily::NoCalls
} else {
EmptyFamily::NoneResolve(calls)
})
}
fn extract_function_calls<T: crate::Op>(f: &Function<T>, p: &Program<T>) -> Vec<String> {
f.iter()
.filter(|op| op.is_call())
.filter_map(|op| op.symbol_key())
.filter_map(|key| p.symbol(&key))
.map(|s| s.to_string())
.collect()
}
fn extract_function_calls_freq<T: crate::Op>(
f: &Function<T>,
p: &Program<T>,
) -> FxHashMap<String, usize> {
extract_function_calls(f, p)
.into_iter()
.fold(FxHashMap::default(), |mut acc, call| {
*acc.entry(call).or_insert(0) += 1;
acc
})
}
#[cfg(test)]
mod tests {
use super::*;
fn a_program_whose_calls_resolve_to_nothing() -> crate::program::AnyProgram {
let dir = tempdir().unwrap();
let path = dir.path().join("indirect.json");
std::fs::write(
&path,
r#"{"program":"indirect","path":"bin/indirect","ir":"ghidra-pcode",
"symbols":{"0x1000": "_printf"},
"functions":[{"name":"main","ops":[
{"op":"CALLIND","inputs":["(register, 0x20, 8)"]},
{"op":"CALLIND","inputs":["(register, 0x28, 8)"]}]}]}"#,
)
.unwrap();
crate::program::AnyProgram::load(&path).unwrap()
}
#[test]
fn test_calls_that_resolve_to_nothing_give_an_empty_birthmark() {
let p = a_program_whose_calls_resolve_to_nothing();
for bt in ["fc-set", "fc-seq", "fc-freq"] {
let b = Extractor::new(BirthmarkType::try_from(bt).unwrap())
.extract_any(&p)
.unwrap_or_else(|e| panic!("{bt} was refused: {e}"));
assert_eq!(b.len(), 1, "{bt}: one function was expected");
let calls: Vec<String> = b.iter().flat_map(|e| e.ops().map(String::from)).collect();
assert!(calls.is_empty(), "{bt}: expected nothing, got {calls:?}");
}
}
#[test]
fn test_a_program_that_calls_nothing_does_not_stop_the_rest() {
let dir = tempdir().unwrap();
let mut programs = Vec::new();
for (name, ops) in [
("calls", r#"{"op":"CALL","inputs":["(ram, 0x1000, 8)"]}"#),
(
"pure",
r#"{"op":"COPY","out":"(register, 0x0, 8)","inputs":["(const, 0x1, 8)"]}"#,
),
(
"calls_too",
r#"{"op":"CALL","inputs":["(ram, 0x1000, 8)"]}"#,
),
] {
let path = dir.path().join(format!("{name}.json"));
std::fs::write(
&path,
format!(
r#"{{"program":"{name}","path":"bin/{name}","ir":"ghidra-pcode",
"symbols":{{"0x1000":"_printf"}},
"functions":[{{"name":"main","ops":[{ops}]}}]}}"#
),
)
.unwrap();
programs.push(crate::program::AnyProgram::load(&path).unwrap());
}
let extractor = Extractor::new(BirthmarkType::try_from("fc-set").unwrap());
let sizes: Vec<usize> = programs
.iter()
.map(|p| {
extractor
.extract_any(p)
.unwrap_or_else(|e| panic!("{e}"))
.iter()
.flat_map(|e| e.ops())
.count()
})
.collect();
assert_eq!(
sizes,
vec![1, 0, 1],
"the middle one should be empty, not fatal"
);
}
#[test]
fn test_the_two_ways_of_emptying_a_family_are_told_apart() {
let unresolved = a_program_whose_calls_resolve_to_nothing();
assert!(matches!(
any_empty_family(&unresolved),
Some(EmptyFamily::NoneResolve(2))
));
let dir = tempdir().unwrap();
let path = dir.path().join("callless.json");
std::fs::write(
&path,
r#"{"program":"callless","path":"bin/callless","ir":"ghidra-pcode",
"symbols":{},
"functions":[{"name":"main","ops":[
{"op":"COPY","out":"(register, 0x0, 8)","inputs":["(const, 0x1, 8)"]}]}]}"#,
)
.unwrap();
let callless = crate::program::AnyProgram::load(&path).unwrap();
assert_eq!(any_empty_family(&callless), Some(EmptyFamily::NoCalls));
let fine = crate::program::AnyProgram::load(std::path::Path::new(
"testdata/lifted/pcodes/hello_clang.json",
))
.unwrap();
assert_eq!(any_empty_family(&fine), None);
}
#[test]
fn test_each_message_names_the_file_the_representation_and_the_cause() {
let path = std::path::Path::new("bin/sample");
let ir = crate::lift::Ir::GhidraPcode;
let none = EmptyFamily::NoCalls.message(path, ir);
assert!(
none.contains("bin/sample") && none.contains("ghidra-pcode"),
"{none}"
);
assert!(none.contains("no operation is a call"), "{none}");
let some = EmptyFamily::NoneResolve(1).message(path, ir);
assert!(some.contains("symbol table"), "{some}");
assert!(
!some.contains("1 operations"),
"reads as a plural for one: {some}"
);
assert!(some.contains("(1 in total)"), "{some}");
}
fn any_empty_family(p: &crate::program::AnyProgram) -> Option<EmptyFamily> {
match p {
crate::program::AnyProgram::GhidraPcode(p) => empty_family(p),
_ => unreachable!("the fixtures here are all P-Code"),
}
}
#[test]
fn test_the_op_families_are_not_refused_for_unresolvable_calls() {
let p = a_program_whose_calls_resolve_to_nothing();
let b = Extractor::new(BirthmarkType::try_from("op-set").unwrap())
.extract_any(&p)
.expect("op-set does not depend on symbols");
assert_eq!(b.len(), 1);
}
use std::io::Write;
use tempfile::tempdir;
#[test]
fn test_dest_file_name_small_file() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("small_test.txt");
let mut file = std::fs::File::create(&file_path).unwrap();
file.write_all(b"hello world").unwrap();
let name = dest_file_name(&file_path).unwrap();
assert!(name.starts_with("small_test_"));
assert!(name.ends_with(".json"));
}
#[test]
fn test_dest_file_name_large_file() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("large_test.txt");
let mut file = std::fs::File::create(&file_path).unwrap();
let data = vec![0u8; 10240];
file.write_all(&data).unwrap();
let name = dest_file_name(&file_path).unwrap();
assert!(name.starts_with("large_test_"));
assert!(name.ends_with(".json"));
}
#[test]
fn test_get_hash_reflects_tail_difference() {
let dir = tempdir().unwrap();
let path1 = dir.path().join("a.bin");
let path2 = dir.path().join("b.bin");
let mut data1 = vec![0u8; 5000];
let mut data2 = vec![0u8; 5000];
data1[4500] = 1;
data2[4500] = 2;
std::fs::write(&path1, &data1).unwrap();
std::fs::write(&path2, &data2).unwrap();
assert_ne!(get_hash(&path1).unwrap(), get_hash(&path2).unwrap());
}
#[test]
fn test_get_hash_of_a_known_input_is_exactly_this() {
let dir = tempdir().unwrap();
let path = dir.path().join("fixture.txt");
std::fs::write(&path, b"oinkie hash fixture\n").unwrap();
assert_eq!(get_hash(&path).unwrap(), "1a22831efd6c8d6e");
}
#[test]
fn test_get_hash_is_deterministic() {
let dir = tempdir().unwrap();
let path1 = dir.path().join("a.bin");
let path2 = dir.path().join("b.bin");
let data = vec![7u8; 10240];
std::fs::write(&path1, &data).unwrap();
std::fs::write(&path2, &data).unwrap();
assert_eq!(get_hash(&path1).unwrap(), get_hash(&path2).unwrap());
}
#[test]
fn test_get_hash_io_error() {
let path = Path::new("non_existent_file.xyz");
let result = get_hash(path);
assert!(result.is_err());
}
#[test]
fn test_extract_function_calls_resolves_the_symbol() {
let program: Program<crate::ghidra::Op> =
std::path::Path::new("testdata/lifted/pcodes/hello_clang.json")
.try_into()
.expect("failed to load the fixture");
let function = program.iter().next().expect("fixture has no function");
let calls = extract_function_calls(function, &program);
assert_eq!(
calls,
vec!["_printf".to_string()],
"the fixture calls _printf exactly once"
);
}
#[test]
fn test_fc_birthmarks_are_not_empty() {
for fixture in [
"testdata/lifted/pcodes/hello_clang.json",
"testdata/lifted/pcodes/hello_gcc.json",
] {
let program: Program<crate::ghidra::Op> = std::path::Path::new(fixture)
.try_into()
.unwrap_or_else(|e| panic!("{fixture}: {e}"));
let birthmark = Extractor::new(BirthmarkType::FcSet)
.extract_each(&program)
.unwrap_or_else(|e| panic!("{fixture}: {e}"));
assert!(
birthmark.elements.iter().any(|e| !e.is_empty()),
"{fixture}: every fc-set element is empty"
);
}
}
#[test]
fn test_fc_extraction_of_a_program_without_calls_is_empty_not_refused() {
let json = r#"{
"program": "callless",
"path": "bin/callless",
"ir": "ghidra-pcode",
"symbols": {"0x100000480": "_printf"},
"functions": [
{"name": "leaf", "ops": [
{"op": "COPY", "out": "(register, 0x4000, 8)", "inputs": ["(const, 0x0, 8)"]},
{"op": "RETURN", "inputs": ["(const, 0x0, 8)"]}
]}
]
}"#;
let program: Program<crate::ghidra::Op> = serde_json::from_str(json).unwrap();
for bt in [
BirthmarkType::FcSeq,
BirthmarkType::FcSet,
BirthmarkType::FcFreq,
] {
let b = Extractor::new(bt.clone())
.extract_each(&program)
.unwrap_or_else(|e| panic!("{bt} was refused: {e}"));
assert!(
b.iter().flat_map(|e| e.ops()).next().is_none(),
"{bt}: a leaf function calls nothing, so this should be empty"
);
}
}
#[test]
fn test_op_extraction_accepts_a_program_without_calls() {
let json = r#"{
"program": "callless",
"path": "bin/callless",
"symbols": {},
"functions": [
{"name": "leaf", "ops": [
{"op": "RETURN", "inputs": ["(const, 0x0, 8)"]}
]}
]
}"#;
let program: Program<crate::ghidra::Op> = serde_json::from_str(json).unwrap();
assert!(
Extractor::new(BirthmarkType::OpSeq)
.extract_each(&program)
.is_ok()
);
}
#[test]
fn test_extractor_extract_multiple() {
let extractor = Extractor::new(BirthmarkType::OpSeq);
let empty_args: Vec<&crate::program::Program<crate::ghidra::Op>> = vec![];
let result = extractor.extract(empty_args);
assert!(result.is_ok());
assert!(result.unwrap().is_empty());
}
#[test]
fn test_dest_file_name_hash_length() {
let dir = tempdir().unwrap();
let file_path = dir.path().join("sample.bin");
std::fs::write(&file_path, b"contents").unwrap();
let name = dest_file_name(&file_path).unwrap();
let hash = name
.strip_prefix("sample_")
.and_then(|s| s.strip_suffix(".json"))
.expect("unexpected file name layout");
assert_eq!(hash.len(), 16);
assert!(hash.chars().all(|c| c.is_ascii_hexdigit()));
}
}