From 767804233fb3728c588bef90526c4c2b6710e4f4 Mon Sep 17 00:00:00 2001
From: mos <mos@cyclic.zone>
Date: Fri, 19 Apr 2024 11:23:48 +0200
Subject: [PATCH] add overflow check for vaddr

---
 src/elf.rs | 19 ++++++++++++-------
 src/lib.rs |  1 +
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/elf.rs b/src/elf.rs
index 724d77e..a3562ed 100644
--- a/src/elf.rs
+++ b/src/elf.rs
@@ -12,6 +12,7 @@ pub enum Error {
     BadMagic,
     BadType,
     BadMachine,
+    BadAddr,
 }
 
 struct Phdr {
@@ -101,15 +102,19 @@ pub fn load(elf: Vec<u8>, mem: &mut [u8]) -> Result<u32, Error> {
             continue;
         }
 
-        let vaddr = (phdr.vaddr - RAM_BASE) as usize;
-        let memsz = phdr.memsz as usize;
+        if let Some(vaddr) = phdr.vaddr.checked_sub(RAM_BASE) {
+            let vaddr = vaddr as usize;
+            let memsz = phdr.memsz as usize;
 
-        let pos = rdr.position();
-        rdr.set_position(phdr.offset as u64);
-        if rdr.read_exact(&mut mem[vaddr..vaddr + memsz]).is_err() {
-            return Err(Error::Eof);
+            let pos = rdr.position();
+            rdr.set_position(phdr.offset as u64);
+            if rdr.read_exact(&mut mem[vaddr..vaddr + memsz]).is_err() {
+                return Err(Error::Eof);
+            }
+            rdr.set_position(pos);
+        } else {
+            return Err(Error::BadAddr);
         }
-        rdr.set_position(pos);
     }
 
     Ok(entry)
diff --git a/src/lib.rs b/src/lib.rs
index 9b272da..64ac1fa 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -336,6 +336,7 @@ pub fn run(path: String) -> Result<()> {
             elf::Error::BadMagic => bail!("bad magic"),
             elf::Error::BadType => bail!("invalid executable"),
             elf::Error::BadMachine => bail!("foreign elf architecture"),
+            elf::Error::BadAddr => bail!("bad virt address offset"),
         },
     };