diff --git a/td-loader/src/elf64.rs b/td-loader/src/elf64.rs index 06d1c58a..6d5d22db 100644 --- a/td-loader/src/elf64.rs +++ b/td-loader/src/elf64.rs @@ -924,4 +924,47 @@ mod test_elf_loader { hdr.sh_size = 0x1; hdr.vm_range(); } + + #[test] + fn test_programheaders_iterator() { + const ENTRY_SIZE: usize = 56; + let pe_image = &include_bytes!("../../data/blobs/td-payload.elf")[..]; + let elf = crate::elf64::Elf::parse(pe_image).unwrap(); + let mut p_hdrs = elf.program_headers().unwrap(); + + p_hdrs.index = (usize::MAX as usize / ENTRY_SIZE) + 1; + assert!(p_hdrs.next().is_none()); + } + + #[test] + fn test_sectionheaders_iterator() { + const ENTRY_SIZE: usize = 64; + let pe_image = &include_bytes!("../../data/blobs/td-payload.elf")[..]; + let elf = crate::elf64::Elf::parse(pe_image).unwrap(); + let mut s_hdrs = elf.section_headers().unwrap(); + + s_hdrs.index = (usize::MAX as usize / ENTRY_SIZE) + 1; + assert!(s_hdrs.next().is_none()); + } + + #[test] + fn test_dyns_iterator() { + const ENTRY_SIZE: usize = 16; + let pe_image = &include_bytes!("../../data/blobs/td-payload.elf")[..]; + let elf = crate::elf64::Elf::parse(pe_image).unwrap(); + let elf_bin = elf.bytes; + + for header in elf.program_headers().unwrap() { + println!("header: {:?}\n", header); + + let mut dyns = Dyns::parse( + &elf_bin[header.p_offset as usize..], + header.p_filesz as usize, + ) + .unwrap(); + + dyns.index = (usize::MAX as usize / ENTRY_SIZE) + 1; + assert!(dyns.next().is_none()); + } + } } diff --git a/td-loader/src/pe.rs b/td-loader/src/pe.rs index 6fb1fd58..a6b30588 100644 --- a/td-loader/src/pe.rs +++ b/td-loader/src/pe.rs @@ -598,4 +598,22 @@ mod test { } assert_eq!(entries, 1); } + + #[test] + fn test_sections_iterator() { + let pe_image = &include_bytes!("../../data/blobs/td-payload.efi")[..]; + let coff_header_offset = pe_image.pread::(0x3c).unwrap() as usize; + let pe_region = &pe_image[coff_header_offset..]; + + let num_sections = pe_region.pread::(6).unwrap() as usize; + let optional_header_size = pe_region.pread::(20).unwrap() as usize; + let optional_region = &pe_image[24 + coff_header_offset..]; + + let mut num_section = 0; + let sections_buffer = &pe_image[(24 + coff_header_offset + optional_header_size)..]; + let mut sections = Sections::parse(sections_buffer, num_sections as usize).unwrap(); + + sections.index = (usize::MAX as usize / COFF_SECTION_SIZE) + 1; + assert!(sections.next().is_none()); + } }