1use crate::fastscan::{first_byte_in_raw, first_byte_in_slice};
20use crate::pattern::WildcardPattern;
21use crate::pe::{exec_sections, text_section_bounds};
22
23#[inline]
30unsafe fn matches_at(addr: usize, pattern: WildcardPattern<'_>) -> bool {
31 for (i, slot) in pattern.iter().enumerate() {
32 if let Some(want) = *slot {
33 let got = *((addr + i) as *const u8);
34 if got != want {
35 return false;
36 }
37 }
38 }
39 true
40}
41
42#[must_use]
53pub fn find_in_text(module_base: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
54 if module_base == 0 || pattern.is_empty() {
55 return None;
56 }
57 let (start, size) = text_section_bounds(module_base)?;
58 scan_range(start, size, pattern)
59}
60
61#[must_use]
69pub fn count_in_text(module_base: usize, pattern: WildcardPattern<'_>) -> usize {
70 if module_base == 0 || pattern.is_empty() {
71 return 0;
72 }
73 match text_section_bounds(module_base) {
74 Some((start, size)) => count_range(start, size, pattern),
75 None => 0,
76 }
77}
78
79#[must_use]
92pub fn find_in_exec_sections(module_base: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
93 if module_base == 0 || pattern.is_empty() {
94 return None;
95 }
96 for (start, size) in exec_sections(module_base)? {
97 if let Some(addr) = scan_range(start, size, pattern) {
98 return Some(addr);
99 }
100 }
101 None
102}
103
104#[must_use]
108pub fn count_in_exec_sections(module_base: usize, pattern: WildcardPattern<'_>) -> usize {
109 if module_base == 0 || pattern.is_empty() {
110 return 0;
111 }
112 let Some(sections) = exec_sections(module_base) else {
113 return 0;
114 };
115 let mut total = 0usize;
116 for (start, size) in sections {
117 total += count_range(start, size, pattern);
118 }
119 total
120}
121
122#[must_use]
133pub fn find_in_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> Option<usize> {
134 if pattern.is_empty() || haystack.len() < pattern.len() {
135 return None;
136 }
137 scan_slice(haystack, pattern).map(|off| haystack.as_ptr() as usize + off)
138}
139
140#[must_use]
144pub fn count_in_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> usize {
145 if pattern.is_empty() || haystack.len() < pattern.len() {
146 return 0;
147 }
148 count_slice(haystack, pattern)
149}
150
151#[derive(Debug, Clone)]
161pub struct SliceMatches<'a> {
162 haystack: &'a [u8],
163 pattern: WildcardPattern<'a>,
164 cursor: usize,
166}
167
168impl<'a> Iterator for SliceMatches<'a> {
169 type Item = usize;
170
171 fn next(&mut self) -> Option<usize> {
172 let pat_len = self.pattern.len();
173 if pat_len == 0 {
174 return None;
175 }
176 let off = scan_slice_from(self.haystack, self.cursor, self.pattern)?;
177 self.cursor = off + pat_len;
178 Some(self.haystack.as_ptr() as usize + off)
179 }
180}
181
182#[derive(Debug, Clone)]
193pub struct Matches<'a> {
194 sections: alloc::vec::Vec<(usize, usize)>,
197 section_idx: usize,
198 cursor: usize,
200 pattern: WildcardPattern<'a>,
201}
202
203impl<'a> Iterator for Matches<'a> {
204 type Item = usize;
205
206 fn next(&mut self) -> Option<usize> {
207 let pat_len = self.pattern.len();
208 if pat_len == 0 {
209 return None;
210 }
211 while self.section_idx < self.sections.len() {
212 let (start, size) = self.sections[self.section_idx];
213 if let Some(addr) = scan_range_from(start, size, self.cursor, self.pattern) {
214 let off = addr - start;
215 self.cursor = off + pat_len;
216 return Some(addr);
217 }
218 self.section_idx += 1;
220 self.cursor = 0;
221 }
222 None
223 }
224}
225
226#[must_use]
251pub fn iter_in_slice<'a>(haystack: &'a [u8], pattern: WildcardPattern<'a>) -> SliceMatches<'a> {
252 SliceMatches {
253 haystack,
254 pattern,
255 cursor: 0,
256 }
257}
258
259#[must_use]
270pub fn iter_in_text<'a>(module_base: usize, pattern: WildcardPattern<'a>) -> Matches<'a> {
271 let sections = if module_base == 0 || pattern.is_empty() {
272 alloc::vec::Vec::new()
273 } else {
274 match text_section_bounds(module_base) {
275 Some(range) => alloc::vec![range],
276 None => alloc::vec::Vec::new(),
277 }
278 };
279 Matches {
280 sections,
281 section_idx: 0,
282 cursor: 0,
283 pattern,
284 }
285}
286
287#[must_use]
295pub fn iter_in_exec_sections<'a>(module_base: usize, pattern: WildcardPattern<'a>) -> Matches<'a> {
296 let sections = if module_base == 0 || pattern.is_empty() {
297 alloc::vec::Vec::new()
298 } else {
299 exec_sections(module_base).unwrap_or_default()
300 };
301 Matches {
302 sections,
303 section_idx: 0,
304 cursor: 0,
305 pattern,
306 }
307}
308
309#[cfg(feature = "section-info")]
332#[must_use]
333pub fn find_in_section(
334 module_base: usize,
335 section_name: &[u8],
336 pattern: WildcardPattern<'_>,
337) -> Option<usize> {
338 if module_base == 0 || pattern.is_empty() {
339 return None;
340 }
341 let section = crate::pe::find_section(module_base, section_name)?;
342 scan_range(section.virtual_address, section.virtual_size, pattern)
343}
344
345#[cfg(feature = "section-info")]
352#[must_use]
353pub fn count_in_section(
354 module_base: usize,
355 section_name: &[u8],
356 pattern: WildcardPattern<'_>,
357) -> usize {
358 if module_base == 0 || pattern.is_empty() {
359 return 0;
360 }
361 let Some(section) = crate::pe::find_section(module_base, section_name) else {
362 return 0;
363 };
364 count_range(section.virtual_address, section.virtual_size, pattern)
365}
366
367#[cfg(feature = "section-info")]
374#[must_use]
375pub fn iter_in_section<'a>(
376 module_base: usize,
377 section_name: &[u8],
378 pattern: WildcardPattern<'a>,
379) -> Matches<'a> {
380 let sections = if module_base == 0 || pattern.is_empty() {
381 alloc::vec::Vec::new()
382 } else {
383 crate::pe::find_section(module_base, section_name)
384 .map(|s| alloc::vec![(s.virtual_address, s.virtual_size)])
385 .unwrap_or_default()
386 };
387 Matches {
388 sections,
389 section_idx: 0,
390 cursor: 0,
391 pattern,
392 }
393}
394
395#[inline]
399fn anchor(pattern: WildcardPattern<'_>) -> Option<(usize, u8)> {
400 pattern
401 .iter()
402 .enumerate()
403 .find_map(|(i, b)| b.map(|byte| (i, byte)))
404}
405
406#[inline]
425fn scan_slice_from(haystack: &[u8], from: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
426 let pat_len = pattern.len();
427 if haystack.len() < pat_len {
428 return None;
429 }
430 let upper = haystack.len() - pat_len;
431 if from > upper {
432 return None;
433 }
434 let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
435 return Some(from);
437 };
438
439 let mut i = from;
440 while i <= upper {
441 let search_from = i + anchor_off;
446 let Some(rel) = first_byte_in_slice(&haystack[search_from..], anchor_byte) else {
447 return None;
448 };
449 let candidate = search_from + rel - anchor_off;
450 if candidate > upper {
451 return None;
452 }
453 if unsafe { matches_at(haystack.as_ptr() as usize + candidate, pattern) } {
455 return Some(candidate);
456 }
457 i = candidate + 1;
458 }
459 None
460}
461
462#[inline]
465fn scan_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> Option<usize> {
466 scan_slice_from(haystack, 0, pattern)
467}
468
469fn count_slice(haystack: &[u8], pattern: WildcardPattern<'_>) -> usize {
471 let pat_len = pattern.len();
472 let upper = haystack.len() - pat_len;
473 let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
474 return haystack.len() / pat_len;
479 };
480
481 let mut count = 0usize;
482 let mut i = 0usize;
483 while i <= upper {
484 let search_from = i + anchor_off;
487 let Some(rel) = first_byte_in_slice(&haystack[search_from..], anchor_byte) else {
488 break;
489 };
490 let candidate = search_from + rel - anchor_off;
491 if candidate > upper {
492 break;
493 }
494 if unsafe { matches_at(haystack.as_ptr() as usize + candidate, pattern) } {
496 count += 1;
497 i = candidate + pat_len;
498 } else {
499 i = candidate + 1;
500 }
501 }
502 count
503}
504
505#[inline]
521fn scan_range_from(
522 start: usize,
523 size: usize,
524 from: usize,
525 pattern: WildcardPattern<'_>,
526) -> Option<usize> {
527 let pat_len = pattern.len();
528 if size < pat_len {
529 return None;
530 }
531 let upper = size - pat_len;
532 if from > upper {
533 return None;
534 }
535 let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
536 return Some(start + from);
537 };
538
539 let mut i = from;
540 while i <= upper {
541 let search_from = i + anchor_off;
544 let Some(rel) =
547 (unsafe { first_byte_in_raw(start + search_from, size - search_from, anchor_byte) })
548 else {
549 return None;
550 };
551 let candidate = search_from + rel - anchor_off;
552 if candidate > upper {
553 return None;
554 }
555 let addr = start + candidate;
556 if unsafe { matches_at(addr, pattern) } {
558 return Some(addr);
559 }
560 i = candidate + 1;
561 }
562 None
563}
564
565#[inline]
568fn scan_range(start: usize, size: usize, pattern: WildcardPattern<'_>) -> Option<usize> {
569 scan_range_from(start, size, 0, pattern)
570}
571
572fn count_range(start: usize, size: usize, pattern: WildcardPattern<'_>) -> usize {
580 let pat_len = pattern.len();
581 if size < pat_len {
582 return 0;
583 }
584 let upper = size - pat_len;
585 let Some((anchor_off, anchor_byte)) = anchor(pattern) else {
586 return size / pat_len;
589 };
590
591 let mut count = 0usize;
592 let mut i = 0usize;
593 while i <= upper {
594 let search_from = i + anchor_off;
597 let Some(rel) =
599 (unsafe { first_byte_in_raw(start + search_from, size - search_from, anchor_byte) })
600 else {
601 break;
602 };
603 let candidate = search_from + rel - anchor_off;
604 if candidate > upper {
605 break;
606 }
607 let addr = start + candidate;
608 if unsafe { matches_at(addr, pattern) } {
610 count += 1;
611 i = candidate + pat_len;
612 } else {
613 i = candidate + 1;
614 }
615 }
616 count
617}
618
619#[cfg(test)]
620mod tests {
621 use super::*;
622 use crate::pattern;
623 use crate::pe::IMAGE_SCN_MEM_EXECUTE;
624 use alloc::vec;
625 use alloc::vec::Vec;
626
627 fn synthetic_pe(sections: &[([u8; 8], u32, &[u8], u32)]) -> Vec<u8> {
631 let needed = sections
632 .iter()
633 .map(|(_, vaddr, bytes, _)| *vaddr as usize + bytes.len())
634 .max()
635 .unwrap_or(0)
636 .max(0x400);
637 let mut buf = vec![0u8; needed];
638 buf[0] = b'M';
639 buf[1] = b'Z';
640 let nt_offset: u32 = 0x80;
641 buf[0x3C..0x40].copy_from_slice(&nt_offset.to_le_bytes());
642 let nt = nt_offset as usize;
643 buf[nt..nt + 4].copy_from_slice(b"PE\0\0");
644 let num_sections: u16 = sections.len() as u16;
645 buf[nt + 4 + 2..nt + 4 + 4].copy_from_slice(&num_sections.to_le_bytes());
646 let opt_size: u16 = 0xF0;
647 buf[nt + 4 + 16..nt + 4 + 18].copy_from_slice(&opt_size.to_le_bytes());
648 let section_table = nt + 4 + 20 + opt_size as usize;
649 for (i, (name, vaddr, bytes, characteristics)) in sections.iter().enumerate() {
650 let sec = section_table + i * 40;
651 buf[sec..sec + 8].copy_from_slice(name);
652 let vsize: u32 = bytes.len() as u32;
653 buf[sec + 8..sec + 12].copy_from_slice(&vsize.to_le_bytes());
654 buf[sec + 12..sec + 16].copy_from_slice(&vaddr.to_le_bytes());
655 buf[sec + 36..sec + 40].copy_from_slice(&characteristics.to_le_bytes());
656 let v = *vaddr as usize;
657 buf[v..v + bytes.len()].copy_from_slice(bytes);
658 }
659 buf
660 }
661
662 #[test]
665 fn slice_find_basic() {
666 let haystack = [0x00, 0x11, 0x48, 0x8B, 0x05, 0x99, 0xAA];
667 let pat = pattern![0x48, 0x8B, 0x05];
668 let hit = find_in_slice(&haystack, pat).unwrap();
669 assert_eq!(hit, haystack.as_ptr() as usize + 2);
670 }
671
672 #[test]
673 fn slice_find_wildcard() {
674 let haystack = [0x00, 0x48, 0x77, 0x05, 0xFF];
675 let pat = pattern![0x48, _, 0x05];
676 let hit = find_in_slice(&haystack, pat).unwrap();
677 assert_eq!(hit, haystack.as_ptr() as usize + 1);
678 }
679
680 #[test]
681 fn slice_find_misses_returns_none() {
682 let haystack = [0x00, 0x11, 0x22];
683 let pat = pattern![0x48, 0x8B, 0x05];
684 assert!(find_in_slice(&haystack, pat).is_none());
685 }
686
687 #[test]
688 fn slice_find_empty_pattern_returns_none() {
689 let haystack = [0x00, 0x11];
690 let pat: &[Option<u8>] = &[];
691 assert!(find_in_slice(&haystack, pat).is_none());
692 }
693
694 #[test]
695 fn slice_find_pattern_longer_than_haystack() {
696 let haystack = [0x48];
697 let pat = pattern![0x48, 0x8B, 0x05];
698 assert!(find_in_slice(&haystack, pat).is_none());
699 }
700
701 #[test]
702 fn slice_find_all_wildcards_matches_first() {
703 let haystack = [0xAA, 0xBB, 0xCC];
706 let pat: &[Option<u8>] = &[None, None];
707 let hit = find_in_slice(&haystack, pat).unwrap();
708 assert_eq!(hit, haystack.as_ptr() as usize);
709 }
710
711 #[test]
712 fn slice_find_anchor_match_but_full_pattern_mismatch() {
713 let haystack = [0x48, 0xFF, 0x48, 0x8B];
724 let pat = pattern![0x48, 0x8B];
725 let hit = find_in_slice(&haystack, pat).unwrap();
726 assert_eq!(hit, haystack.as_ptr() as usize + 2);
727 }
728
729 #[test]
730 fn slice_count_anchor_match_but_full_pattern_mismatch() {
731 let haystack = [0x48, 0xFF, 0x48, 0x8B];
735 let pat = pattern![0x48, 0x8B];
736 assert_eq!(count_in_slice(&haystack, pat), 1);
737 }
738
739 #[test]
740 fn slice_find_anchor_in_middle() {
741 let haystack = [0xCC, 0x77, 0x99, 0xAA, 0x77, 0x99];
744 let pat: &[Option<u8>] = &[None, Some(0x77), Some(0x99)];
745 let hit = find_in_slice(&haystack, pat).unwrap();
746 assert_eq!(hit, haystack.as_ptr() as usize);
748 }
749
750 #[test]
751 fn slice_count_basic() {
752 let haystack = [0x48, 0x8B, 0x00, 0x48, 0x8B, 0x00, 0x48, 0x8B];
753 let pat = pattern![0x48, 0x8B];
754 assert_eq!(count_in_slice(&haystack, pat), 3);
755 }
756
757 #[test]
758 fn slice_count_no_overlap() {
759 let haystack = [0x42, 0x42, 0x42, 0x42];
762 let pat = pattern![0x42, 0x42];
763 assert_eq!(count_in_slice(&haystack, pat), 2);
764 }
765
766 #[test]
767 fn slice_count_zero_when_no_match() {
768 let haystack = [0x00, 0x11, 0x22];
769 let pat = pattern![0x48];
770 assert_eq!(count_in_slice(&haystack, pat), 0);
771 }
772
773 #[test]
774 fn slice_count_empty_pattern_returns_zero() {
775 let haystack = [0x00, 0x11];
776 let pat: &[Option<u8>] = &[];
777 assert_eq!(count_in_slice(&haystack, pat), 0);
778 }
779
780 #[test]
781 fn slice_count_pattern_longer_than_haystack() {
782 let haystack = [0x48];
783 let pat = pattern![0x48, 0x8B, 0x05];
784 assert_eq!(count_in_slice(&haystack, pat), 0);
785 }
786
787 #[test]
788 fn slice_count_all_wildcards() {
789 let haystack = [0xAA, 0xBB, 0xCC, 0xDD];
792 let pat: &[Option<u8>] = &[None, None];
793 assert_eq!(count_in_slice(&haystack, pat), 2);
794 }
795
796 #[test]
799 fn synthetic_pe_text_find_and_count() {
800 let text = [0x00u8, 0x11, 0x48, 0x8B, 0x05, 0xFF, 0x00, 0x48, 0x8B, 0x05];
801 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
802 let base = buf.as_ptr() as usize;
803 let pat = pattern![0x48, 0x8B, 0x05];
804
805 let hit = find_in_text(base, pat).unwrap();
807 assert_eq!(hit, base + 0x300 + 2);
808
809 assert_eq!(count_in_text(base, pat), 2);
811 }
812
813 #[test]
814 fn synthetic_pe_text_find_returns_none_when_no_match() {
815 let text = [0xAAu8, 0xBB, 0xCC];
816 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
817 let base = buf.as_ptr() as usize;
818 let pat = pattern![0x48, 0x8B];
819 assert!(find_in_text(base, pat).is_none());
820 assert_eq!(count_in_text(base, pat), 0);
821 }
822
823 #[test]
824 fn synthetic_pe_text_returns_none_when_no_text_section() {
825 let body = [0x48u8, 0x8B];
828 let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
829 let base = buf.as_ptr() as usize;
830 let pat = pattern![0x48, 0x8B];
831 assert!(find_in_text(base, pat).is_none());
832 assert_eq!(count_in_text(base, pat), 0);
833 }
834
835 #[test]
836 fn synthetic_pe_exec_sections_find_across_sections() {
837 let body_a = [0xAAu8, 0xBB];
840 let body_b = [0x90u8, 0x90, 0xC3];
841 let buf = synthetic_pe(&[
842 (*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
843 (*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
844 ]);
845 let base = buf.as_ptr() as usize;
846 let pat = pattern![0x90, 0x90, 0xC3];
847 let hit = find_in_exec_sections(base, pat).unwrap();
848 assert_eq!(hit, base + 0x310);
849 assert_eq!(count_in_exec_sections(base, pat), 1);
850 }
851
852 #[test]
853 fn synthetic_pe_exec_sections_count_sums_across_sections() {
854 let body = [0x90u8, 0x90];
857 let buf = synthetic_pe(&[
858 (*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
859 (*b".text$mn", 0x310, &body, IMAGE_SCN_MEM_EXECUTE),
860 ]);
861 let base = buf.as_ptr() as usize;
862 let pat = pattern![0x90, 0x90];
863 assert_eq!(count_in_exec_sections(base, pat), 2);
864 }
865
866 #[test]
867 fn synthetic_pe_exec_sections_returns_none_when_no_match() {
868 let body = [0xAAu8];
869 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
870 let base = buf.as_ptr() as usize;
871 let pat = pattern![0x48];
872 assert!(find_in_exec_sections(base, pat).is_none());
873 assert_eq!(count_in_exec_sections(base, pat), 0);
874 }
875
876 #[test]
879 fn null_module_returns_none_or_zero() {
880 let pat = pattern![0x48];
881 assert!(find_in_text(0, pat).is_none());
882 assert_eq!(count_in_text(0, pat), 0);
883 assert!(find_in_exec_sections(0, pat).is_none());
884 assert_eq!(count_in_exec_sections(0, pat), 0);
885 }
886
887 #[test]
888 fn empty_pattern_returns_none_or_zero() {
889 let body = [0x90u8];
890 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
891 let base = buf.as_ptr() as usize;
892 let pat: &[Option<u8>] = &[];
893 assert!(find_in_text(base, pat).is_none());
894 assert_eq!(count_in_text(base, pat), 0);
895 assert!(find_in_exec_sections(base, pat).is_none());
896 assert_eq!(count_in_exec_sections(base, pat), 0);
897 }
898
899 #[test]
900 fn malformed_module_returns_none_or_zero() {
901 let buf = vec![0u8; 0x400];
904 let base = buf.as_ptr() as usize;
905 let pat = pattern![0x48];
906 assert!(find_in_text(base, pat).is_none());
907 assert_eq!(count_in_text(base, pat), 0);
908 assert!(find_in_exec_sections(base, pat).is_none());
909 assert_eq!(count_in_exec_sections(base, pat), 0);
910 }
911
912 #[test]
917 fn synthetic_pe_text_all_wildcard_pattern() {
918 let text = [0xAAu8; 16];
919 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &text, IMAGE_SCN_MEM_EXECUTE)]);
920 let base = buf.as_ptr() as usize;
921 let pat: &[Option<u8>] = &[None, None, None, None];
922
923 let hit = find_in_text(base, pat).unwrap();
925 let (text_start, _) = crate::pe::text_section_bounds(base).unwrap();
926 assert_eq!(hit, text_start);
927
928 let count = count_in_text(base, pat);
932 assert!(count >= text.len() / 4);
933 }
934
935 #[test]
941 fn slice_anchor_at_tail_no_room_for_pattern() {
942 let mut buf = vec![0u8; 16];
945 buf[15] = 0x48;
946 let pat = pattern![0x48, 0x8B, 0x05, _];
947
948 assert!(crate::find_in_slice(&buf, pat).is_none());
949 assert_eq!(crate::count_in_slice(&buf, pat), 0);
950 }
951
952 #[test]
958 fn slice_find_loop_exhausts_when_last_candidate_fails() {
959 let haystack = [0x00, 0x00, 0x48, 0xFF];
962 let pat = pattern![0x48, 0x8B];
963 assert!(find_in_slice(&haystack, pat).is_none());
964 }
965
966 #[test]
973 fn synthetic_pe_text_anchor_match_but_full_pattern_mismatch() {
974 let body = [0x48, 0x00, 0x00, 0x00, 0x48];
979 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
980 let base = buf.as_ptr() as usize;
981 let pat = pattern![0x48, 0xAA, 0xBB, 0xCC];
982
983 assert!(find_in_text(base, pat).is_none());
984 assert_eq!(count_in_text(base, pat), 0);
985 }
986
987 #[test]
992 fn synthetic_pe_text_anchor_at_upper_then_loop_exhausts() {
993 let body = [0x48, 0x00, 0x00, 0x00];
996 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
997 let base = buf.as_ptr() as usize;
998 let pat = pattern![0x48, 0xAA, 0xBB, 0xCC];
999
1000 assert!(find_in_text(base, pat).is_none());
1001 assert_eq!(count_in_text(base, pat), 0);
1002 }
1003
1004 #[test]
1007 fn iter_slice_yields_all_non_overlapping_matches() {
1008 let bytes = [
1009 0x48, 0x8B, 0x05, 0x00, 0x48, 0x8B, 0x05, 0xFF, 0x48, 0x8B, 0x05,
1010 ];
1011 let pat = pattern![0x48, 0x8B, 0x05];
1012 let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
1013 assert_eq!(hits.len(), 3);
1014 assert_eq!(hits[0], bytes.as_ptr() as usize);
1015 assert_eq!(hits[1], bytes.as_ptr() as usize + 4);
1016 assert_eq!(hits[2], bytes.as_ptr() as usize + 8);
1017 }
1018
1019 #[test]
1020 fn iter_slice_count_matches_count_in_slice() {
1021 let bytes = [0x42, 0x42, 0x42, 0x42, 0x42];
1024 let pat = pattern![0x42, 0x42];
1025 assert_eq!(
1026 iter_in_slice(&bytes, pat).count(),
1027 count_in_slice(&bytes, pat)
1028 );
1029 }
1030
1031 #[test]
1032 fn iter_slice_empty_pattern_yields_nothing() {
1033 let bytes = [0x00, 0x11, 0x22];
1034 let pat: &[Option<u8>] = &[];
1035 assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
1036 }
1037
1038 #[test]
1039 fn iter_slice_no_match_yields_nothing() {
1040 let bytes = [0xAA, 0xBB, 0xCC];
1041 let pat = pattern![0x48, 0x8B];
1042 assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
1043 }
1044
1045 #[test]
1046 fn iter_slice_pattern_longer_than_haystack_yields_nothing() {
1047 let bytes = [0x48];
1048 let pat = pattern![0x48, 0x8B, 0x05];
1049 assert_eq!(iter_in_slice(&bytes, pat).count(), 0);
1050 }
1051
1052 #[test]
1053 fn iter_slice_with_wildcards() {
1054 let bytes = [0x48, 0x8B, 0x05, 0x48, 0x99, 0x05, 0x48, 0xAA, 0xFF];
1057 let pat = pattern![0x48, _, 0x05];
1058 let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
1059 assert_eq!(hits.len(), 2);
1060 assert_eq!(hits[0], bytes.as_ptr() as usize);
1061 assert_eq!(hits[1], bytes.as_ptr() as usize + 3);
1062 }
1063
1064 #[test]
1065 fn iter_slice_all_wildcard_pattern_strides_by_pat_len() {
1066 let bytes = [0xAA, 0xBB, 0xCC, 0xDD, 0xEE];
1070 let pat: &[Option<u8>] = &[None, None];
1071 let hits: Vec<usize> = iter_in_slice(&bytes, pat).collect();
1072 assert_eq!(hits.len(), 2);
1073 assert_eq!(hits[0], bytes.as_ptr() as usize);
1074 assert_eq!(hits[1], bytes.as_ptr() as usize + 2);
1075 }
1076
1077 #[test]
1078 fn iter_slice_clone_is_independent() {
1079 let bytes = [0x48, 0x8B, 0x00, 0x48, 0x8B];
1082 let pat = pattern![0x48, 0x8B];
1083 let it = iter_in_slice(&bytes, pat);
1084 let from_clone: Vec<usize> = it.clone().collect();
1085 let from_original: Vec<usize> = it.collect();
1086 assert_eq!(from_clone, from_original);
1087 assert_eq!(from_clone.len(), 2);
1088 }
1089
1090 #[test]
1093 fn iter_in_text_yields_all_matches() {
1094 let body = [0x48u8, 0x8B, 0x05, 0x00, 0x48, 0x8B, 0x05];
1095 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1096 let base = buf.as_ptr() as usize;
1097 let pat = pattern![0x48, 0x8B, 0x05];
1098 let hits: Vec<usize> = iter_in_text(base, pat).collect();
1099 assert_eq!(hits.len(), 2);
1100 assert_eq!(hits[0], base + 0x300);
1101 assert_eq!(hits[1], base + 0x300 + 4);
1102 }
1103
1104 #[test]
1105 fn iter_in_text_first_matches_find_in_text() {
1106 let body = [0x90u8, 0x90, 0x48, 0x8B, 0x05, 0xCC];
1109 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1110 let base = buf.as_ptr() as usize;
1111 let pat = pattern![0x48, 0x8B, 0x05];
1112 let from_iter = iter_in_text(base, pat).next();
1113 let from_find = find_in_text(base, pat);
1114 assert_eq!(from_iter, from_find);
1115 assert_eq!(from_iter, Some(base + 0x300 + 2));
1116 }
1117
1118 #[test]
1119 fn iter_in_text_count_matches_count_in_text() {
1120 let body = [0x48u8, 0x8B, 0x00, 0x48, 0x8B, 0x00, 0x48, 0x8B];
1122 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1123 let base = buf.as_ptr() as usize;
1124 let pat = pattern![0x48, 0x8B];
1125 assert_eq!(iter_in_text(base, pat).count(), count_in_text(base, pat));
1126 }
1127
1128 #[test]
1129 fn iter_in_text_null_module_yields_nothing() {
1130 let pat = pattern![0x48];
1131 assert_eq!(iter_in_text(0, pat).count(), 0);
1132 }
1133
1134 #[test]
1135 fn iter_in_text_empty_pattern_yields_nothing() {
1136 let body = [0x90u8];
1137 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1138 let base = buf.as_ptr() as usize;
1139 let pat: &[Option<u8>] = &[];
1140 assert_eq!(iter_in_text(base, pat).count(), 0);
1141 }
1142
1143 #[test]
1144 fn iter_in_text_missing_text_section_yields_nothing() {
1145 let body = [0x48u8, 0x8B];
1147 let buf = synthetic_pe(&[(*b".data\0\0\0", 0x300, &body, 0)]);
1148 let base = buf.as_ptr() as usize;
1149 let pat = pattern![0x48, 0x8B];
1150 assert_eq!(iter_in_text(base, pat).count(), 0);
1151 }
1152
1153 #[test]
1154 fn iter_in_text_malformed_module_yields_nothing() {
1155 let buf = vec![0u8; 0x400];
1157 let base = buf.as_ptr() as usize;
1158 let pat = pattern![0x48];
1159 assert_eq!(iter_in_text(base, pat).count(), 0);
1160 }
1161
1162 #[test]
1165 fn iter_in_exec_sections_yields_across_multiple_sections() {
1166 let body_a = [0x90u8, 0x90, 0xC3];
1167 let body_b = [0x48u8, 0x8B, 0x05, 0xCC, 0x48, 0x8B, 0x05];
1168 let buf = synthetic_pe(&[
1169 (*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
1170 (*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
1171 ]);
1172 let base = buf.as_ptr() as usize;
1173 let pat = pattern![0x48, 0x8B, 0x05];
1174 let hits: Vec<usize> = iter_in_exec_sections(base, pat).collect();
1175 assert_eq!(hits.len(), 2);
1177 assert_eq!(hits[0], base + 0x310);
1178 assert_eq!(hits[1], base + 0x310 + 4);
1179 }
1180
1181 #[test]
1182 fn iter_in_exec_sections_advances_to_next_section_after_exhaustion() {
1183 let body_a = [0xAAu8, 0xBB, 0xCC];
1186 let body_b = [0x90u8, 0x90, 0xC3];
1187 let buf = synthetic_pe(&[
1188 (*b".text\0\0\0", 0x300, &body_a, IMAGE_SCN_MEM_EXECUTE),
1189 (*b".text$mn", 0x310, &body_b, IMAGE_SCN_MEM_EXECUTE),
1190 ]);
1191 let base = buf.as_ptr() as usize;
1192 let pat = pattern![0x90, 0x90, 0xC3];
1193 let hits: Vec<usize> = iter_in_exec_sections(base, pat).collect();
1194 assert_eq!(hits, vec![base + 0x310]);
1195 }
1196
1197 #[test]
1198 fn iter_in_exec_sections_count_sums_across_sections() {
1199 let body = [0x90u8, 0x90];
1200 let buf = synthetic_pe(&[
1201 (*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE),
1202 (*b".text$mn", 0x310, &body, IMAGE_SCN_MEM_EXECUTE),
1203 ]);
1204 let base = buf.as_ptr() as usize;
1205 let pat = pattern![0x90, 0x90];
1206 assert_eq!(
1208 iter_in_exec_sections(base, pat).count(),
1209 count_in_exec_sections(base, pat)
1210 );
1211 assert_eq!(iter_in_exec_sections(base, pat).count(), 2);
1212 }
1213
1214 #[test]
1215 fn iter_in_exec_sections_null_module_yields_nothing() {
1216 let pat = pattern![0x48];
1217 assert_eq!(iter_in_exec_sections(0, pat).count(), 0);
1218 }
1219
1220 #[test]
1221 fn iter_in_exec_sections_empty_pattern_yields_nothing() {
1222 let body = [0x90u8];
1223 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1224 let base = buf.as_ptr() as usize;
1225 let pat: &[Option<u8>] = &[];
1226 assert_eq!(iter_in_exec_sections(base, pat).count(), 0);
1227 }
1228
1229 #[test]
1230 fn iter_in_exec_sections_malformed_module_yields_nothing() {
1231 let buf = vec![0u8; 0x400];
1232 let base = buf.as_ptr() as usize;
1233 let pat = pattern![0x48];
1234 assert_eq!(iter_in_exec_sections(base, pat).count(), 0);
1235 }
1236
1237 #[cfg(feature = "section-info")]
1238 mod section_info_tests {
1239 use super::synthetic_pe;
1240 use crate::pattern;
1241 use crate::pe::IMAGE_SCN_MEM_EXECUTE;
1242 use crate::scan::{count_in_section, find_in_section, iter_in_section};
1243 use alloc::vec;
1244 use alloc::vec::Vec;
1245
1246 fn multi_section_pe() -> Vec<u8> {
1257 let text_body = [0x90u8, 0xAA, 0xBB, 0xCC, 0xDD, 0xC3];
1258 let rdata_body = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x11, 0x22, 0xFF];
1259 synthetic_pe(&[
1260 (*b".text\0\0\0", 0x300, &text_body, IMAGE_SCN_MEM_EXECUTE),
1261 (*b".rdata\0\0", 0x400, &rdata_body, 0),
1262 ])
1263 }
1264
1265 #[test]
1268 fn returns_match_in_named_section() {
1269 let buf = multi_section_pe();
1270 let base = buf.as_ptr() as usize;
1271 let pat = pattern![0x11, 0x22, 0x33];
1272 assert_eq!(find_in_section(base, b".rdata", pat), Some(base + 0x401));
1273 }
1274
1275 #[test]
1276 fn does_not_cross_section_bounds() {
1277 let buf = multi_section_pe();
1281 let base = buf.as_ptr() as usize;
1282 let pat = pattern![0xAA, 0xBB, 0xCC, 0xDD];
1283 assert!(find_in_section(base, b".rdata", pat).is_none());
1284 assert_eq!(find_in_section(base, b".text", pat), Some(base + 0x301));
1285 }
1286
1287 #[test]
1288 fn matches_section_name_by_prefix() {
1289 let body = [0xDEu8, 0xAD, 0xBE, 0xEF];
1291 let buf = synthetic_pe(&[(*b".rdata$z", 0x300, &body, 0)]);
1292 let base = buf.as_ptr() as usize;
1293 let pat = pattern![0xDE, 0xAD, 0xBE, 0xEF];
1294 assert_eq!(find_in_section(base, b".rdata", pat), Some(base + 0x300));
1295 }
1296
1297 #[test]
1298 fn full_eight_byte_name_disambiguates() {
1299 let mn_body = [0x11u8, 0x22, 0x33];
1302 let text_body = [0xAAu8, 0xBB, 0xCC];
1303 let buf = synthetic_pe(&[
1304 (*b".text$mn", 0x300, &mn_body, IMAGE_SCN_MEM_EXECUTE),
1305 (*b".text\0\0\0", 0x400, &text_body, IMAGE_SCN_MEM_EXECUTE),
1306 ]);
1307 let base = buf.as_ptr() as usize;
1308 let pat = pattern![0xAA, 0xBB, 0xCC];
1310 assert_eq!(
1311 find_in_section(base, b".text\0\0\0", pat),
1312 Some(base + 0x400),
1313 );
1314 assert!(find_in_section(base, b".text$mn", pat).is_none());
1316 }
1317
1318 #[test]
1319 fn returns_none_when_section_missing() {
1320 let body = [0x90u8];
1321 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1322 let base = buf.as_ptr() as usize;
1323 let pat = pattern![0x90];
1324 assert!(find_in_section(base, b".rdata", pat).is_none());
1325 }
1326
1327 #[test]
1328 fn returns_none_when_pattern_absent() {
1329 let buf = multi_section_pe();
1330 let base = buf.as_ptr() as usize;
1331 let pat = pattern![0xFF, 0xFF, 0xFF, 0xFF];
1332 assert!(find_in_section(base, b".rdata", pat).is_none());
1333 }
1334
1335 #[test]
1336 fn null_module_returns_none() {
1337 let pat = pattern![0x90];
1338 assert!(find_in_section(0, b".rdata", pat).is_none());
1339 }
1340
1341 #[test]
1342 fn empty_pattern_returns_none() {
1343 let buf = multi_section_pe();
1344 let base = buf.as_ptr() as usize;
1345 let pat: &[Option<u8>] = &[];
1346 assert!(find_in_section(base, b".rdata", pat).is_none());
1347 }
1348
1349 #[test]
1350 fn malformed_module_returns_none() {
1351 let buf = vec![0u8; 0x400];
1352 let base = buf.as_ptr() as usize;
1353 let pat = pattern![0x90];
1354 assert!(find_in_section(base, b".rdata", pat).is_none());
1355 }
1356
1357 #[test]
1360 fn count_finds_all_matches_in_section() {
1361 let buf = multi_section_pe();
1362 let base = buf.as_ptr() as usize;
1363 let pat = pattern![0x11, 0x22];
1365 assert_eq!(count_in_section(base, b".rdata", pat), 2);
1366 }
1367
1368 #[test]
1369 fn count_does_not_include_other_sections() {
1370 let buf = multi_section_pe();
1371 let base = buf.as_ptr() as usize;
1372 let pat = pattern![0xAA, 0xBB];
1374 assert_eq!(count_in_section(base, b".rdata", pat), 0);
1375 assert_eq!(count_in_section(base, b".text", pat), 1);
1376 }
1377
1378 #[test]
1379 fn count_returns_zero_when_section_missing() {
1380 let body = [0x90u8];
1381 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1382 let base = buf.as_ptr() as usize;
1383 let pat = pattern![0x90];
1384 assert_eq!(count_in_section(base, b".rdata", pat), 0);
1385 }
1386
1387 #[test]
1388 fn count_null_module_returns_zero() {
1389 let pat = pattern![0x90];
1390 assert_eq!(count_in_section(0, b".rdata", pat), 0);
1391 }
1392
1393 #[test]
1394 fn count_empty_pattern_returns_zero() {
1395 let buf = multi_section_pe();
1396 let base = buf.as_ptr() as usize;
1397 let pat: &[Option<u8>] = &[];
1398 assert_eq!(count_in_section(base, b".rdata", pat), 0);
1399 }
1400
1401 #[test]
1402 fn count_malformed_module_returns_zero() {
1403 let buf = vec![0u8; 0x400];
1404 let base = buf.as_ptr() as usize;
1405 let pat = pattern![0x90];
1406 assert_eq!(count_in_section(base, b".rdata", pat), 0);
1407 }
1408
1409 #[test]
1412 fn iter_yields_all_matches_in_order() {
1413 let buf = multi_section_pe();
1414 let base = buf.as_ptr() as usize;
1415 let pat = pattern![0x11, 0x22];
1416 let hits: Vec<usize> = iter_in_section(base, b".rdata", pat).collect();
1417 assert_eq!(hits, vec![base + 0x401, base + 0x405]);
1418 }
1419
1420 #[test]
1421 fn iter_first_equals_find_in_section() {
1422 let buf = multi_section_pe();
1423 let base = buf.as_ptr() as usize;
1424 let pat = pattern![0x11, 0x22];
1425 assert_eq!(
1426 iter_in_section(base, b".rdata", pat).next(),
1427 find_in_section(base, b".rdata", pat),
1428 );
1429 }
1430
1431 #[test]
1432 fn iter_count_equals_count_in_section() {
1433 let buf = multi_section_pe();
1434 let base = buf.as_ptr() as usize;
1435 let pat = pattern![0x11, 0x22];
1436 assert_eq!(
1437 iter_in_section(base, b".rdata", pat).count(),
1438 count_in_section(base, b".rdata", pat),
1439 );
1440 }
1441
1442 #[test]
1443 fn iter_does_not_cross_section_bounds() {
1444 let buf = multi_section_pe();
1445 let base = buf.as_ptr() as usize;
1446 let pat = pattern![0xAA, 0xBB];
1447 assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
1448 }
1449
1450 #[test]
1451 fn iter_section_missing_yields_nothing() {
1452 let body = [0x90u8];
1453 let buf = synthetic_pe(&[(*b".text\0\0\0", 0x300, &body, IMAGE_SCN_MEM_EXECUTE)]);
1454 let base = buf.as_ptr() as usize;
1455 let pat = pattern![0x90];
1456 assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
1457 }
1458
1459 #[test]
1460 fn iter_null_module_yields_nothing() {
1461 let pat = pattern![0x90];
1462 assert_eq!(iter_in_section(0, b".rdata", pat).count(), 0);
1463 }
1464
1465 #[test]
1466 fn iter_empty_pattern_yields_nothing() {
1467 let buf = multi_section_pe();
1468 let base = buf.as_ptr() as usize;
1469 let pat: &[Option<u8>] = &[];
1470 assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
1471 }
1472
1473 #[test]
1474 fn iter_malformed_module_yields_nothing() {
1475 let buf = vec![0u8; 0x400];
1476 let base = buf.as_ptr() as usize;
1477 let pat = pattern![0x90];
1478 assert_eq!(iter_in_section(base, b".rdata", pat).count(), 0);
1479 }
1480 }
1481}