td-loader: Add unit test case for PR 410

Fix: #416
Signed-off-by: Wei Liu <[email protected]>

Author: liuw1

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());
+        }
+    }
 }

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::<u32>(0x3c).unwrap() as usize;
+        let pe_region = &pe_image[coff_header_offset..];
+
+        let num_sections = pe_region.pread::<u16>(6).unwrap() as usize;
+        let optional_header_size = pe_region.pread::<u16>(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());
+    }
 }