td-loader: Fix potential panic issue

Fix: 439 Signed-off-by: Wei Liu <[email protected]>
confidential-containers · Nov 4, 2022 · e5961c3 · e5961c3
1 parent 0c18c17
commit e5961c3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/td-loader/fuzz/fuzz_targets/fuzzlib.rs b/td-loader/fuzz/fuzz_targets/fuzzlib.rs
@@ -19,6 +19,9 @@ pub fn fuzz_elf_loader(data: &[u8]) {
     if let Some(elf) = Elf::parse(data) {
         log::info!("{:?}\n", elf.header);
 
+        if elf.program_headers().is_none() {
+            return;
+        }
         if let Some(hd) = elf.program_headers().unwrap().next() {
             let status = hd.is_executable();
             log::info!("executable status: {}", status);

diff --git a/td-loader/src/elf.rs b/td-loader/src/elf.rs
@@ -38,7 +38,7 @@ pub fn relocate_elf_with_per_program_header(
     let mut bottom: u64 = 0xFFFFFFFFu64;
     let mut top: u64 = 0u64;
 
-    for ph in elf.program_headers().unwrap() {
+    for ph in elf.program_headers()? {
         if bottom > ph.p_vaddr {
             bottom = ph.p_vaddr;
         }
@@ -54,7 +54,7 @@ pub fn relocate_elf_with_per_program_header(
     bottom = align_value(bottom, SIZE_4KB, true);
     top = align_value(top, SIZE_4KB, false);
     // load per program header
-    for ph in elf.program_headers().unwrap() {
+    for ph in elf.program_headers()? {
         if ph.p_memsz != 0 {
             if ph.p_offset.checked_add(ph.p_filesz)? > image.len() as u64
                 || ph.p_vaddr.checked_add(ph.p_filesz)? > loaded_buffer.len() as u64