use crate::fastscan::{first_byte_in_raw, first_byte_in_slice};
use crate::pattern::WildcardPattern;
use crate::pe::{exec_sections, text_section_bounds};
#[inline]
unsafe fn matches_at(addr: usize, pattern: WildcardPattern<'_>) -> bool {
for (i, slot) in pattern.iter().enumerate() {
if let Some(want) = *slot {
let got = *((addr + i) as *const u8);
if got != want {
return false;
}
}
}
true
}
#[must_use]
pub fn find_in_text(module_base: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
if module_base == 0 || pattern.is_empty() {
return None;
}
let (start, size) = text_section_bounds(module_base)?;
scan_range(start, size, pattern)
}
#[must_use]
pub fn count_in_text(module_base: usize, pattern: WildcardPattern<'_>) -> usize {
if module_base == 0 || pattern.is_empty() {
return 0;
}
match text_section_bounds(module_base) {
Some((start, size)) => count_range(start, size, pattern),
None => 0,
}
}
#[must_use]
pub fn find_in_exec_sections(module_base: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
if module_base == 0 || pattern.is_empty() {
return None;
}
for (start, size) in exec_sections(module_base)? {
if let Some(addr) = scan_range(start, size, pattern) {
return Some(addr);
}
}
None
}
#[must_use]
pub fn count_in_exec_sections(module_base: usize, pattern: WildcardPattern<'_>) -> usize {
if module_base == 0 || pattern.is_empty() {
return 0;
}
let Some(sections) = exec_sections(module_base) else {
return 0;
};
let mut total = 0usize;
for (start, size) in sections {
total += count_range(start, size, pattern);
}
total
}
#[must_use]
pub fn find_in_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> Option<usize> {
if pattern.is_empty() || haystack.len() < pattern.len() {
return None;
}
scan_slice(haystack, pattern).map(|off| haystack.as_ptr() as usize + off)
}
#[must_use]
pub fn count_in_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> usize {
if pattern.is_empty() || haystack.len() < pattern.len() {
return 0;
}
count_slice(haystack, pattern)
}
#[derive(Debug, Clone)]
pub struct SliceMatches<'a> {
haystack: &'a [u8],
pattern: WildcardPattern<'a>,
cursor: usize,
}
impl<'a> Iterator for SliceMatches<'a> {
type Item = usize;
fn next(&mut self) -> Option<usize> {
let pat_len = self.pattern.len();
if pat_len == 0 {
return None;
}
let off = scan_slice_from(self.haystack, self.cursor, self.pattern)?;
self.cursor = off + pat_len;
Some(self.haystack.as_ptr() as usize + off)
}
}
#[derive(Debug, Clone)]
pub struct Matches<'a> {
sections: alloc::vec::Vec<(usize, usize)>,
section_idx: usize,
cursor: usize,
pattern: WildcardPattern<'a>,
}
impl<'a> Iterator for Matches<'a> {
type Item = usize;
fn next(&mut self) -> Option<usize> {
let pat_len = self.pattern.len();
if pat_len == 0 {
return None;
}
while self.section_idx < self.sections.len() {
let (start, size) = self.sections[self.section_idx];
if let Some(addr) = scan_range_from(start, size, self.cursor, self.pattern) {
let off = addr - start;
self.cursor = off + pat_len;
return Some(addr);
}
self.section_idx += 1;
self.cursor = 0;
}
None
}
}
#[must_use]
pub fn iter_in_slice<'a>(haystack: &'a [u8], pattern: WildcardPattern<'a>) -> SliceMatches<'a> {
SliceMatches {
haystack,
pattern,
cursor: 0,
}
}
#[must_use]
pub fn iter_in_text<'a>(module_base: usize, pattern: WildcardPattern<'a>) -> Matches<'a> {
let sections = if module_base == 0 || pattern.is_empty() {
alloc::vec::Vec::new()
} else {
match text_section_bounds(module_base) {
Some(range) => alloc::vec![range],
None => alloc::vec::Vec::new(),
}
};
Matches {
sections,
section_idx: 0,
cursor: 0,
pattern,
}
}
#[must_use]
pub fn iter_in_exec_sections<'a>(module_base: usize, pattern: WildcardPattern<'a>) -> Matches<'a> {
let sections = if module_base == 0 || pattern.is_empty() {
alloc::vec::Vec::new()
} else {
exec_sections(module_base).unwrap_or_default()
};
Matches {
sections,
section_idx: 0,
cursor: 0,
pattern,
}
}
#[cfg(feature = "section-info")]
#[must_use]
pub fn find_in_section(
module_base: usize,
section_name: &[u8],
pattern: WildcardPattern<'_>,
) -> Option<usize> {
if module_base == 0 || pattern.is_empty() {
return None;
}
let section = crate::pe::find_section(module_base, section_name)?;
scan_range(section.virtual_address, section.virtual_size, pattern)
}
#[cfg(feature = "section-info")]
#[must_use]
pub fn count_in_section(
module_base: usize,
section_name: &[u8],
pattern: WildcardPattern<'_>,
) -> usize {
if module_base == 0 || pattern.is_empty() {
return 0;
}
let Some(section) = crate::pe::find_section(module_base, section_name) else {
return 0;
};
count_range(section.virtual_address, section.virtual_size, pattern)
}
#[cfg(feature = "section-info")]
#[must_use]
pub fn iter_in_section<'a>(
module_base: usize,
section_name: &[u8],
pattern: WildcardPattern<'a>,
) -> Matches<'a> {
let sections = if module_base == 0 || pattern.is_empty() {
alloc::vec::Vec::new()
} else {
crate::pe::find_section(module_base, section_name)
.map(|s| alloc::vec![(s.virtual_address, s.virtual_size)])
.unwrap_or_default()
};
Matches {
sections,
section_idx: 0,
cursor: 0,
pattern,
}
}
#[inline]
fn anchor(pattern: WildcardPattern<'_>) -> Option<(usize, u8)> {
pattern
.iter()
.enumerate()
.find_map(|(i, b)| b.map(|byte| (i, byte)))
}
#[inline]
fn scan_slice_from(haystack: &[u8], from: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
let pat_len = pattern.len();
if haystack.len() < pat_len {
return None;
}
let upper = haystack.len() - pat_len;
if from > upper {
return None;
}
let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
return Some(from);
};
let mut i = from;
while i <= upper {
let search_from = i + anchor_off;
let Some(rel) = first_byte_in_slice(&haystack[search_from..], anchor_byte) else {
return None;
};
let candidate = search_from + rel - anchor_off;
if candidate > upper {
return None;
}
if unsafe { matches_at(haystack.as_ptr() as usize + candidate, pattern) } {
return Some(candidate);
}
i = candidate + 1;
}
None
}
#[inline]
fn scan_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> Option<usize> {
scan_slice_from(haystack, 0, pattern)
}
fn count_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> usize {
let pat_len = pattern.len();
let upper = haystack.len() - pat_len;
let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
return haystack.len() / pat_len;
};
let mut count = 0usize;
let mut i = 0usize;
while i <= upper {
let search_from = i + anchor_off;
let Some(rel) = first_byte_in_slice(&haystack[search_from..], anchor_byte) else {
break;
};
let candidate = search_from + rel - anchor_off;
if candidate > upper {
break;
}
if unsafe { matches_at(haystack.as_ptr() as usize + candidate, pattern) } {
count += 1;
i = candidate + pat_len;
} else {
i = candidate + 1;
}
}
count
}
#[inline]
fn scan_range_from(
start: usize,
size: usize,
from: usize,
pattern: WildcardPattern<'_>,
) -> Option<usize> {
let pat_len = pattern.len();
if size < pat_len {
return None;
}
let upper = size - pat_len;
if from > upper {
return None;
}
let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
return Some(start + from);
};
let mut i = from;
while i <= upper {
let search_from = i + anchor_off;
let Some(rel) =
(unsafe { first_byte_in_raw(start + search_from, size - search_from, anchor_byte) })
else {
return None;
};
let candidate = search_from + rel - anchor_off;
if candidate > upper {
return None;
}
let addr = start + candidate;
if unsafe { matches_at(addr, pattern) } {
return Some(addr);
}
i = candidate + 1;
}
None
}
#[inline]
fn scan_range(start: usize, size: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
scan_range_from(start, size, 0, pattern)
}
fn count_range(start: usize, size: usize, pattern: WildcardPattern<'_>) -> usize {
let pat_len = pattern.len();
if size < pat_len {
return 0;
}
let upper = size - pat_len;
let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
return size / pat_len;
};
let mut count = 0usize;
let mut i = 0usize;
while i <= upper {
let search_from = i + anchor_off;
let Some(rel) =
(unsafe { first_byte_in_raw(start + search_from, size - search_from, anchor_byte) })
else {
break;
};
let candidate = search_from + rel - anchor_off;
if candidate > upper {
break;
}
let addr = start + candidate;
if unsafe { matches_at(addr, pattern) } {
count += 1;
i = candidate + pat_len;
} else {
i = candidate + 1;
}
}
count
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pattern;
use crate::pe::IMAGE_SCN_MEM_EXECUTE;
use alloc::vec;
use alloc::vec::Vec;
fn synthetic_pe(sections: &[([u8; 8], u32, &[u8], u32)]) -> Vec<u8> {
let needed = sections
.iter()
.map(|(_, vaddr, bytes, _)| *vaddr as usize + bytes.len())
.max()
.unwrap_or(0)
.max(0x400);
let mut buf = vec![0u8; needed];
buf[0] = b'M';
buf[1] = b'Z';
let nt_offset: u32 = 0x80;
buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
let nt = nt_offset as usize;
buf[nt..nt + 4].copy_from_slice(b"PE\0\0");
let num_sections: u16 = sections.len() as u16;
buf[nt + 4 + 2..nt + 4 + 4].copy_from_slice(&num_sections.to_le_bytes());
let opt_size: u16 = 0xF0;
buf[nt + 4 + 16..nt + 4 + 18].copy_from_slice(&opt_size.to_le_bytes());
let section_table = nt + 4 + 20 + opt_size as usize;
for (i, (name, vaddr, bytes, characteristics)) in sections.iter().enumerate() {
let sec = section_table + i * 40;
buf[sec..sec + 8].copy_from_slice(name);
let vsize: u32 = bytes.len() as u32;
buf[sec + 8..sec + 12].copy_from_slice(&vsize.to_le_bytes());
buf[sec + 12..sec + 16].copy_from_slice(&vaddr.to_le_bytes());
buf[sec + 36..sec + 40].copy_from_slice(&characteristics.to_le_bytes());
let v = *vaddr as usize;
buf[v..v + bytes.len()].copy_from_slice(bytes);
}
buf
}
#[test]
fn slice_find_basic() {
let haystack = [0x00, 0x11, 0x48, 0x8B, 0x05, 0x99, 0xAA];
let pat = pattern![0x48, 0x8B, 0x05];
let hit = find_in_slice(&haystack, pat).unwrap();
assert_eq!(hit, haystack.as_ptr() as usize + 2);
}
#[test]
fn slice_find_wildcard() {
let haystack = [0x00, 0x48, 0x77, 0x05, 0xFF];
let pat = pattern![0x48, _, 0x05];
let hit = find_in_slice(&haystack, pat).unwrap();
assert_eq!(hit, haystack.as_ptr() as usize + 1);
}
#[test]
fn slice_find_misses_returns_none() {
let haystack = [0x00, 0x11, 0x22];
let pat = pattern![0x48, 0x8B, 0x05];
assert!(find_in_slice(&haystack, pat).is_none());
}
#[test]
fn slice_find_empty_pattern_returns_none() {
let haystack = [0x00, 0x11];
let pat: &[Option<u8>] = &[];
assert!(find_in_slice(&haystack, pat).is_none());
}
#[test]
fn slice_find_pattern_longer_than_haystack() {
let haystack = [0x48];
let pat = pattern![0x48, 0x8B, 0x05];
assert!(find_in_slice(&haystack, pat).is_none());
}
#[test]
fn slice_find_all_wildcards_matches_first() {
let haystack = [0xAA, 0xBB, 0xCC];
let pat: &[Option<u8>] = &[None, None];
let hit = find_in_slice(&haystack, pat).unwrap();
assert_eq!(hit, haystack.as_ptr() as usize);
}
#[test]
fn slice_find_anchor_match_but_full_pattern_mismatch() {
let haystack = [0x48, 0xFF, 0x48, 0x8B];
let pat = pattern![0x48, 0x8B];
let hit = find_in_slice(&haystack, pat).unwrap();
assert_eq!(hit, haystack.as_ptr() as usize + 2);
}
#[test]
fn slice_count_anchor_match_but_full_pattern_mismatch() {
let haystack = [0x48, 0xFF, 0x48, 0x8B];
let pat = pattern![0x48, 0x8B];
assert_eq!(count_in_slice(&haystack, pat), 1);
}
#[test]
fn slice_find_anchor_in_middle() {
let haystack = [0xCC, 0x77, 0x99, 0xAA, 0x77, 0x99];
let pat: &[Option<u8>] = &[None, Some(0x77), Some(0x99)];
let hit = find_in_slice(&haystack, pat).unwrap();
assert_eq!(hit, haystack.as_ptr() as usize);
}
#[test]
fn slice_count_basic() {
let haystack = [0x48, 0x8B, 0x00, 0x48, 0x8B, 0x00, 0x48, 0x8B];
let pat = pattern![0x48, 0x8B];
assert_eq!(count_in_slice(&haystack, pat), 3);
}
#[test]
fn slice_count_no_overlap() {
let haystack = [0x42, 0x42, 0x42, 0x42];
let pat = pattern![0x42, 0x42];
assert_eq!(count_in_slice(&haystack, pat), 2);
}
#[test]
fn slice_count_zero_when_no_match() {
let haystack = [0x00, 0x11, 0x22];
let pat = pattern![0x48];
assert_eq!(count_in_slice(&haystack, pat), 0);
}
#[test]
fn slice_count_empty_pattern_returns_zero() {
let haystack = [0x00, 0x11];
let pat: &[Option<u8>] = &[];
assert_eq!(count_in_slice(&haystack, pat), 0);
}
#[test]
fn slice_count_pattern_longer_than_haystack() {
let haystack = [0x48];
let pat = pattern![0x48, 0x8B, 0x05];
assert_eq!(count_in_slice(&haystack, pat), 0);
}
#[test]
fn slice_count_all_wildcards() {
let haystack = [0xAA, 0xBB, 0xCC, 0xDD];
let pat: &[Option<u8>] = &[None, None];
assert_eq!(count_in_slice(&haystack, pat), 2);
}
#[test]
fn synthetic_pe_text_find_and_count() {
let text = [0x00u8, 0x11, 0x48, 0x8B, 0x05, 0xFF, 0x00, 0x48, 0x8B, 0x05];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B, 0x05];
let hit = find_in_text(base, pat).unwrap();
assert_eq!(hit, base + 0x300 + 2);
assert_eq!(count_in_text(base, pat), 2);
}
#[test]
fn synthetic_pe_text_find_returns_none_when_no_match() {
let text = [0xAAu8, 0xBB, 0xCC];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
}
#[test]
fn synthetic_pe_text_returns_none_when_no_text_section() {
let body = [0x48u8, 0x8B];
let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
}
#[test]
fn synthetic_pe_exec_sections_find_across_sections() {
let body_a = [0xAAu8, 0xBB];
let body_b = [0x90u8, 0x90, 0xC3];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90, 0x90, 0xC3];
let hit = find_in_exec_sections(base, pat).unwrap();
assert_eq!(hit, base + 0x310);
assert_eq!(count_in_exec_sections(base, pat), 1);
}
#[test]
fn synthetic_pe_exec_sections_count_sums_across_sections() {
let body = [0x90u8, 0x90];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90, 0x90];
assert_eq!(count_in_exec_sections(base, pat), 2);
}
#[test]
fn synthetic_pe_exec_sections_returns_none_when_no_match() {
let body = [0xAAu8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48];
assert!(find_in_exec_sections(base, pat).is_none());
assert_eq!(count_in_exec_sections(base, pat), 0);
}
#[test]
fn null_module_returns_none_or_zero() {
let pat = pattern![0x48];
assert!(find_in_text(0, pat).is_none());
assert_eq!(count_in_text(0, pat), 0);
assert!(find_in_exec_sections(0, pat).is_none());
assert_eq!(count_in_exec_sections(0, pat), 0);
}
#[test]
fn empty_pattern_returns_none_or_zero() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
assert!(find_in_exec_sections(base, pat).is_none());
assert_eq!(count_in_exec_sections(base, pat), 0);
}
#[test]
fn malformed_module_returns_none_or_zero() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x48];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
assert!(find_in_exec_sections(base, pat).is_none());
assert_eq!(count_in_exec_sections(base, pat), 0);
}
#[test]
fn synthetic_pe_text_all_wildcard_pattern() {
let text = [0xAAu8; 16];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[None, None, None, None];
let hit = find_in_text(base, pat).unwrap();
let (text_start, _) = crate::pe::text_section_bounds(base).unwrap();
assert_eq!(hit, text_start);
let count = count_in_text(base, pat);
assert!(count >= text.len() / 4);
}
#[test]
fn slice_anchor_at_tail_no_room_for_pattern() {
let mut buf = vec![0u8; 16];
buf[15] = 0x48;
let pat = pattern![0x48, 0x8B, 0x05, _];
assert!(crate::find_in_slice(&buf, pat).is_none());
assert_eq!(crate::count_in_slice(&buf, pat), 0);
}
#[test]
fn slice_find_loop_exhausts_when_last_candidate_fails() {
let haystack = [0x00, 0x00, 0x48, 0xFF];
let pat = pattern![0x48, 0x8B];
assert!(find_in_slice(&haystack, pat).is_none());
}
#[test]
fn synthetic_pe_text_anchor_match_but_full_pattern_mismatch() {
let body = [0x48, 0x00, 0x00, 0x00, 0x48];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0xAA, 0xBB, 0xCC];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
}
#[test]
fn synthetic_pe_text_anchor_at_upper_then_loop_exhausts() {
let body = [0x48, 0x00, 0x00, 0x00];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0xAA, 0xBB, 0xCC];
assert!(find_in_text(base, pat).is_none());
assert_eq!(count_in_text(base, pat), 0);
}
#[test]
fn iter_slice_yields_all_non_overlapping_matches() {
let bytes = [
0x48, 0x8B, 0x05, 0x00, 0x48, 0x8B, 0x05, 0xFF, 0x48, 0x8B, 0x05,
];
let pat = pattern![0x48, 0x8B, 0x05];
let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
assert_eq!(hits.len(), 3);
assert_eq!(hits[0], bytes.as_ptr() as usize);
assert_eq!(hits[1], bytes.as_ptr() as usize + 4);
assert_eq!(hits[2], bytes.as_ptr() as usize + 8);
}
#[test]
fn iter_slice_count_matches_count_in_slice() {
let bytes = [0x42, 0x42, 0x42, 0x42, 0x42];
let pat = pattern![0x42, 0x42];
assert_eq!(
iter_in_slice(&bytes, pat).count(),
count_in_slice(&bytes, pat)
);
}
#[test]
fn iter_slice_empty_pattern_yields_nothing() {
let bytes = [0x00, 0x11, 0x22];
let pat: &[Option<u8>] = &[];
assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
}
#[test]
fn iter_slice_no_match_yields_nothing() {
let bytes = [0xAA, 0xBB, 0xCC];
let pat = pattern![0x48, 0x8B];
assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
}
#[test]
fn iter_slice_pattern_longer_than_haystack_yields_nothing() {
let bytes = [0x48];
let pat = pattern![0x48, 0x8B, 0x05];
assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
}
#[test]
fn iter_slice_with_wildcards() {
let bytes = [0x48, 0x8B, 0x05, 0x48, 0x99, 0x05, 0x48, 0xAA, 0xFF];
let pat = pattern![0x48, _, 0x05];
let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0], bytes.as_ptr() as usize);
assert_eq!(hits[1], bytes.as_ptr() as usize + 3);
}
#[test]
fn iter_slice_all_wildcard_pattern_strides_by_pat_len() {
let bytes = [0xAA, 0xBB, 0xCC, 0xDD, 0xEE];
let pat: &[Option<u8>] = &[None, None];
let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0], bytes.as_ptr() as usize);
assert_eq!(hits[1], bytes.as_ptr() as usize + 2);
}
#[test]
fn iter_slice_clone_is_independent() {
let bytes = [0x48, 0x8B, 0x00, 0x48, 0x8B];
let pat = pattern![0x48, 0x8B];
let it = iter_in_slice(&bytes, pat);
let from_clone: Vec<usize> = it.clone().collect();
let from_original: Vec<usize> = it.collect();
assert_eq!(from_clone, from_original);
assert_eq!(from_clone.len(), 2);
}
#[test]
fn iter_in_text_yields_all_matches() {
let body = [0x48u8, 0x8B, 0x05, 0x00, 0x48, 0x8B, 0x05];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B, 0x05];
let hits: Vec<usize> = iter_in_text(base, pat).collect();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0], base + 0x300);
assert_eq!(hits[1], base + 0x300 + 4);
}
#[test]
fn iter_in_text_first_matches_find_in_text() {
let body = [0x90u8, 0x90, 0x48, 0x8B, 0x05, 0xCC];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B, 0x05];
let from_iter = iter_in_text(base, pat).next();
let from_find = find_in_text(base, pat);
assert_eq!(from_iter, from_find);
assert_eq!(from_iter, Some(base + 0x300 + 2));
}
#[test]
fn iter_in_text_count_matches_count_in_text() {
let body = [0x48u8, 0x8B, 0x00, 0x48, 0x8B, 0x00, 0x48, 0x8B];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B];
assert_eq!(iter_in_text(base, pat).count(), count_in_text(base, pat));
}
#[test]
fn iter_in_text_null_module_yields_nothing() {
let pat = pattern![0x48];
assert_eq!(iter_in_text(0, pat).count(), 0);
}
#[test]
fn iter_in_text_empty_pattern_yields_nothing() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert_eq!(iter_in_text(base, pat).count(), 0);
}
#[test]
fn iter_in_text_missing_text_section_yields_nothing() {
let body = [0x48u8, 0x8B];
let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B];
assert_eq!(iter_in_text(base, pat).count(), 0);
}
#[test]
fn iter_in_text_malformed_module_yields_nothing() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x48];
assert_eq!(iter_in_text(base, pat).count(), 0);
}
#[test]
fn iter_in_exec_sections_yields_across_multiple_sections() {
let body_a = [0x90u8, 0x90, 0xC3];
let body_b = [0x48u8, 0x8B, 0x05, 0xCC, 0x48, 0x8B, 0x05];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x48, 0x8B, 0x05];
let hits: Vec<usize> = iter_in_exec_sections(base, pat).collect();
assert_eq!(hits.len(), 2);
assert_eq!(hits[0], base + 0x310);
assert_eq!(hits[1], base + 0x310 + 4);
}
#[test]
fn iter_in_exec_sections_advances_to_next_section_after_exhaustion() {
let body_a = [0xAAu8, 0xBB, 0xCC];
let body_b = [0x90u8, 0x90, 0xC3];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90, 0x90, 0xC3];
let hits: Vec<usize> = iter_in_exec_sections(base, pat).collect();
assert_eq!(hits, vec![base + 0x310]);
}
#[test]
fn iter_in_exec_sections_count_sums_across_sections() {
let body = [0x90u8, 0x90];
let buf = synthetic_pe(&[
(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
(*b".text$mn", 0x310, &body, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90, 0x90];
assert_eq!(
iter_in_exec_sections(base, pat).count(),
count_in_exec_sections(base, pat)
);
assert_eq!(iter_in_exec_sections(base, pat).count(), 2);
}
#[test]
fn iter_in_exec_sections_null_module_yields_nothing() {
let pat = pattern![0x48];
assert_eq!(iter_in_exec_sections(0, pat).count(), 0);
}
#[test]
fn iter_in_exec_sections_empty_pattern_yields_nothing() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert_eq!(iter_in_exec_sections(base, pat).count(), 0);
}
#[test]
fn iter_in_exec_sections_malformed_module_yields_nothing() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x48];
assert_eq!(iter_in_exec_sections(base, pat).count(), 0);
}
#[cfg(feature = "section-info")]
mod section_info_tests {
use super::synthetic_pe;
use crate::pattern;
use crate::pe::IMAGE_SCN_MEM_EXECUTE;
use crate::scan::{count_in_section, find_in_section, iter_in_section};
use alloc::vec;
use alloc::vec::Vec;
fn multi_section_pe() -> Vec<u8> {
let text_body = [0x90u8, 0xAA, 0xBB, 0xCC, 0xDD, 0xC3];
let rdata_body = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x11, 0x22, 0xFF];
synthetic_pe(&[
(*b".text\0\0\0", 0x300, &text_body, IMAGE_SCN_MEM_EXECUTE),
(*b".rdata\0\0", 0x400, &rdata_body, 0),
])
}
#[test]
fn returns_match_in_named_section() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0x11, 0x22, 0x33];
assert_eq!(find_in_section(base, b".rdata", pat), Some(base + 0x401));
}
#[test]
fn does_not_cross_section_bounds() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0xAA, 0xBB, 0xCC, 0xDD];
assert!(find_in_section(base, b".rdata", pat).is_none());
assert_eq!(find_in_section(base, b".text", pat), Some(base + 0x301));
}
#[test]
fn matches_section_name_by_prefix() {
let body = [0xDEu8, 0xAD, 0xBE, 0xEF];
let buf = synthetic_pe(&[(*b".rdata$z", 0x300, &body, 0)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0xDE, 0xAD, 0xBE, 0xEF];
assert_eq!(find_in_section(base, b".rdata", pat), Some(base + 0x300));
}
#[test]
fn full_eight_byte_name_disambiguates() {
let mn_body = [0x11u8, 0x22, 0x33];
let text_body = [0xAAu8, 0xBB, 0xCC];
let buf = synthetic_pe(&[
(*b".text$mn", 0x300, &mn_body, IMAGE_SCN_MEM_EXECUTE),
(*b".text\0\0\0", 0x400, &text_body, IMAGE_SCN_MEM_EXECUTE),
]);
let base = buf.as_ptr() as usize;
let pat = pattern![0xAA, 0xBB, 0xCC];
assert_eq!(
find_in_section(base, b".text\0\0\0", pat),
Some(base + 0x400),
);
assert!(find_in_section(base, b".text$mn", pat).is_none());
}
#[test]
fn returns_none_when_section_missing() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert!(find_in_section(base, b".rdata", pat).is_none());
}
#[test]
fn returns_none_when_pattern_absent() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0xFF, 0xFF, 0xFF, 0xFF];
assert!(find_in_section(base, b".rdata", pat).is_none());
}
#[test]
fn null_module_returns_none() {
let pat = pattern![0x90];
assert!(find_in_section(0, b".rdata", pat).is_none());
}
#[test]
fn empty_pattern_returns_none() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert!(find_in_section(base, b".rdata", pat).is_none());
}
#[test]
fn malformed_module_returns_none() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert!(find_in_section(base, b".rdata", pat).is_none());
}
#[test]
fn count_finds_all_matches_in_section() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0x11, 0x22];
assert_eq!(count_in_section(base, b".rdata", pat), 2);
}
#[test]
fn count_does_not_include_other_sections() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0xAA, 0xBB];
assert_eq!(count_in_section(base, b".rdata", pat), 0);
assert_eq!(count_in_section(base, b".text", pat), 1);
}
#[test]
fn count_returns_zero_when_section_missing() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert_eq!(count_in_section(base, b".rdata", pat), 0);
}
#[test]
fn count_null_module_returns_zero() {
let pat = pattern![0x90];
assert_eq!(count_in_section(0, b".rdata", pat), 0);
}
#[test]
fn count_empty_pattern_returns_zero() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert_eq!(count_in_section(base, b".rdata", pat), 0);
}
#[test]
fn count_malformed_module_returns_zero() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert_eq!(count_in_section(base, b".rdata", pat), 0);
}
#[test]
fn iter_yields_all_matches_in_order() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0x11, 0x22];
let hits: Vec<usize> = iter_in_section(base, b".rdata", pat).collect();
assert_eq!(hits, vec![base + 0x401, base + 0x405]);
}
#[test]
fn iter_first_equals_find_in_section() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0x11, 0x22];
assert_eq!(
iter_in_section(base, b".rdata", pat).next(),
find_in_section(base, b".rdata", pat),
);
}
#[test]
fn iter_count_equals_count_in_section() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0x11, 0x22];
assert_eq!(
iter_in_section(base, b".rdata", pat).count(),
count_in_section(base, b".rdata", pat),
);
}
#[test]
fn iter_does_not_cross_section_bounds() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat = pattern![0xAA, 0xBB];
assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
}
#[test]
fn iter_section_missing_yields_nothing() {
let body = [0x90u8];
let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
}
#[test]
fn iter_null_module_yields_nothing() {
let pat = pattern![0x90];
assert_eq!(iter_in_section(0, b".rdata", pat).count(), 0);
}
#[test]
fn iter_empty_pattern_yields_nothing() {
let buf = multi_section_pe();
let base = buf.as_ptr() as usize;
let pat: &[Option<u8>] = &[];
assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
}
#[test]
fn iter_malformed_module_yields_nothing() {
let buf = vec![0u8; 0x400];
let base = buf.as_ptr() as usize;
let pat = pattern![0x90];
assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
}
}
}