Make MemoryImage::new() failable + bound check (risc0#533)

* MemoryImage::new() returns a Result type and fixes Program input bug.

risc0/build/src/lib.rs

@@ -71,7 +71,7 @@ impl Risc0Method {
 
         let elf = fs::read(&self.elf_path).unwrap();
         let program = Program::load_elf(&elf, MEM_SIZE as u32).unwrap();
-        let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+        let image = MemoryImage::new(&program, PAGE_SIZE as u32).unwrap();
         image.get_root()
     }
 

risc0/tools/src/bin/make_id.rs

@@ -42,7 +42,7 @@ fn main() {
     let args = Args::parse();
     let elf_contents = fs::read(args.elf).unwrap();
     let program = Program::load_elf(&elf_contents, MEM_SIZE as u32).unwrap();
-    let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+    let image = MemoryImage::new(&program, PAGE_SIZE as u32).unwrap();
     let image_id = image.get_root();
     std::fs::write(args.out, image_id.as_bytes()).unwrap();
 }

risc0/zkvm/src/binfmt/image.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use anyhow::{Context, Result};
 use risc0_zkp::core::{
     digest::Digest,
     hash::sha::{Sha256, BLOCK_BYTES, SHA256_INIT},
@@ -119,23 +120,23 @@ impl MemoryImage {
     /// The result is a MemoryImage with the ELF of `program` loaded (but
     /// execution not yet begun), and with the page table Merkle tree
     /// constructed.
-    pub fn new(program: &Program, page_size: u32) -> Self {
+    pub fn new(program: &Program, page_size: u32) -> Result<Self> {
         let mut buf = vec![0_u8; MEM_SIZE];
 
         // Load the ELF into the memory image.
         for (addr, data) in program.image.iter() {
             let addr = *addr as usize;
             let bytes = data.to_le_bytes();
-            for i in 0..WORD_SIZE {
-                buf[addr + i] = bytes[i];
-            }
+            buf.get_mut(addr..(WORD_SIZE + addr))
+                .context("Invalid Elf Program, address outside MEM_SIZE")?
+                .copy_from_slice(&bytes[..WORD_SIZE]);
         }
 
         // Compute the page table hashes except for the very last root hash.
         let info = PageTableInfo::new(PAGE_TABLE.start() as u32, page_size);
         let mut img = Self { buf, info };
         img.hash_pages();
-        img
+        Ok(img)
     }
 
     /// Calculate and update the image merkle tree within this image.
@@ -157,7 +158,7 @@ impl MemoryImage {
     /// root and that the data from each page hashes to the expected page table
     /// entry.
     #[cfg(test)]
-    fn check(&self, addr: u32) -> anyhow::Result<()> {
+    fn check(&self, addr: u32) -> Result<()> {
         let mut page_idx = self.info.get_page_index(addr);
         while page_idx < self.info.root_idx {
             let page_addr = self.info.get_page_addr(page_idx);
@@ -212,7 +213,7 @@ fn hash_page(page: &[u8]) -> Digest {
 mod tests {
     use risc0_zkvm_methods::MULTI_TEST_ELF;
     use risc0_zkvm_platform::{
-        memory::{DATA, PAGE_TABLE, STACK, SYSTEM, TEXT},
+        memory::{DATA, MEM_SIZE, PAGE_TABLE, STACK, SYSTEM, TEXT},
         syscall::DIGEST_BYTES,
     };
     use test_log::test;
@@ -228,7 +229,7 @@ mod tests {
     fn check_integrity() {
         const PAGE_SIZE: u32 = 1024;
         let program = Program::load_elf(MULTI_TEST_ELF, TEXT.end() as u32).unwrap();
-        let image = MemoryImage::new(&program, PAGE_SIZE);
+        let image = MemoryImage::new(&program, PAGE_SIZE).unwrap();
         // This is useful in case one needs to manually inspect the memory image.
         // std::fs::write("/tmp/test.img", &image.image).unwrap();
         image.check(STACK.start() as u32).unwrap();
@@ -341,4 +342,13 @@ mod tests {
             (0x1A00 - 0x1800) / DIGEST_BYTES as u32 + 6
         );
     }
+
+    #[test]
+    #[should_panic(expected = "Invalid Elf Program, address outside MEM_SIZE")]
+    fn test_fuzzing_oob_idx_bug() {
+        let data = b"\x7f\x45\x4c\x46\x01\x01\x01\x01\x01\x01\xff\xff\x00\x00\x00\x00\x02\x00\xf3\x00\x00\x00\x00\x00\x00\x00\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x20\x00\x08\x00\x00\x00\x96\x96\x00\x94\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\xff\x00\x00\x94\x00\x00\x00\xff\xf6\x12\xa9\x00\x00\x00\x00\x00\x00\xfe\x00\x00\x00\x00\x00\x0a\x9a\x38\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x96\x4c\x46\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x0a\x9d\xd8\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x40\x1e\x00\x00\x46\x4c\x00\x00\x00\x00\x00\x02\x00\x40\x00\x01\x01\x01\x00\x04\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x05\x00\x00\x07\x78\xc1\x0a\x00\x00\xba\x00\x00\x00\x00\xe3\x04\x00\x00\x31\x35\x32\x37\x38\x31\x46\x01\x01\x01\x01\x01\x01\xff\xff\x00\x00\x00\x00\x02\x00\xe5\x00\x00\x00\x00\x96\x96\x00\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\xff\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x0b\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x06\x2e\xac\x00\x00\x00\x00\x00\x00\x0a\xce\x58\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\xff\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02\x40\x1e\x1e\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x40\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x05\x00\x00\x07\x00\xba\xe8\xad\x0a\x00\xe3\x04\x00\x00\x00\x00\x12\x00\x00\x00\x00\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x05\x00\x00\x05\x00\x00\x00\x01\x01\x01\x50\xcf\x0a\x00\x01\x01\x01\x01\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x04\x01\x01\x01\x01\x01\x01\x01\x00\x00\x31\x35\x31\x35\x32\x37\x38\x31\x30\x34\x02\x00\x00\x00\x00\x00\x00\x00\x00\x05\x00\x00\x05\x00\x00\x00\x01\x01\x01\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x07\x00\x00\x00\xff\xff\xff\xff\x00\x00\x00\x00\xff\x04\x92\x01\x01\x01\x01\x01\x01\xa2\xf8\x00\x20\x00\x00\x00\x00\xff\x00\x40\x00\x04\x00\x00\x00\x38\x00\x00\x00\x00\x00\x00\x00\x02\x00\x0a\x40\x40\x00\x1a\x00\x19\x00";
+        const PAGE_SIZE: u32 = 1024;
+        let prog = Program::load_elf(data, MEM_SIZE as u32).unwrap();
+        let _res = MemoryImage::new(&prog, PAGE_SIZE).unwrap();
+    }
 }

risc0/zkvm/src/exec/mod.rs

@@ -149,7 +149,7 @@ impl<'a> Executor<'a> {
     /// Construct a new [Executor] from an ELF binary.
     pub fn from_elf(env: ExecutorEnv<'a>, elf: &[u8]) -> Result<Self> {
         let program = Program::load_elf(&elf, MEM_SIZE as u32)?;
-        let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+        let image = MemoryImage::new(&program, PAGE_SIZE as u32)?;
         Ok(Self::new(env, image, program.entry))
     }
 

risc0/zkvm/src/exec/tests.rs

@@ -40,7 +40,7 @@ fn basic() {
         entry: 0x4000,
         image,
     };
-    let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+    let image = MemoryImage::new(&program, PAGE_SIZE as u32).unwrap();
     let pre_image_id = image.get_root();
 
     let mut exec = Executor::new(env, image, program.entry);
@@ -68,7 +68,7 @@ fn system_split() {
     image.insert(pc, 0x00000073); // ecall(halt)
 
     let program = Program { entry, image };
-    let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+    let image = MemoryImage::new(&program, PAGE_SIZE as u32).unwrap();
     let pre_image_id = image.get_root();
 
     let mut exec = Executor::new(env, image, program.entry);

risc0/zkvm/src/prove/tests.rs

@@ -233,7 +233,7 @@ mod riscv {
             entry.read_to_end(&mut elf).unwrap();
 
             let program = Program::load_elf(elf.as_slice(), MEM_SIZE as u32).unwrap();
-            let image = MemoryImage::new(&program, PAGE_SIZE as u32);
+            let image = MemoryImage::new(&program, PAGE_SIZE as u32).unwrap();
 
             let env = ExecutorEnv::default();
             let mut exec = Executor::new(env, image, program.entry);