Merge pull request #15 from anton-rs/clabby/elf-tests

⚒️ ELF tests

crates/mipsevm/src/lib.rs

@@ -28,7 +28,7 @@ mod mips;
 pub use mips::InstrumentedState;
 
 mod patch;
-pub use patch::{load_elf, patch_go, patch_stack};
+pub use patch::{load_elf, patch_go, patch_stack, MultiReader};
 
 #[cfg(test)]
 mod test_utils;

crates/mipsevm/src/mips/instrumented.rs

@@ -105,19 +105,23 @@ where
 
 #[cfg(test)]
 mod test {
-    use crate::test_utils::StaticOracle;
+    use alloy_primitives::keccak256;
+
+    use crate::test_utils::{ClaimTestOracle, BASE_ADDR_END, END_ADDR};
+    use crate::witness::STATE_WITNESS_SIZE;
+    use crate::{load_elf, patch, StateWitnessHasher};
+    use crate::{test_utils::StaticOracle, Address, InstrumentedState, Memory, State};
+    use std::io::BufWriter;
+    use std::{
+        cell::RefCell,
+        fs,
+        io::{self, BufReader},
+        path::PathBuf,
+        rc::Rc,
+    };
 
     mod open_mips {
         use super::*;
-        use crate::test_utils::{BASE_ADDR_END, END_ADDR};
-        use crate::{Address, InstrumentedState, Memory, State};
-        use std::{
-            cell::RefCell,
-            fs,
-            io::{self, BufReader},
-            path::PathBuf,
-            rc::Rc,
-        };
 
         #[test]
         fn open_mips_tests() {
@@ -187,4 +191,114 @@ mod test {
             }
         }
     }
+
+    #[test]
+    fn state_hash() {
+        let cases = [
+            (false, 0),
+            (false, 1),
+            (false, 2),
+            (false, 3),
+            (true, 0),
+            (true, 1),
+            (true, 2),
+            (true, 3),
+        ];
+
+        for (exited, exit_code) in cases.into_iter() {
+            let mut state = State {
+                exited,
+                exit_code,
+                ..Default::default()
+            };
+
+            let actual_witness = state.encode_witness().unwrap();
+            let actual_state_hash = actual_witness.state_hash();
+            assert_eq!(actual_witness.len(), STATE_WITNESS_SIZE);
+
+            let mut expected_witness = [0u8; STATE_WITNESS_SIZE];
+            let mem_root = state.memory.borrow_mut().merkle_root().unwrap();
+            expected_witness[..32].copy_from_slice(mem_root.as_slice());
+            expected_witness[32 * 2 + 4 * 6] = exit_code;
+            expected_witness[32 * 2 + 4 * 6 + 1] = exited as u8;
+
+            assert_eq!(actual_witness, expected_witness, "Incorrect witness");
+
+            let mut expected_state_hash = keccak256(&expected_witness);
+            expected_state_hash[0] = State::vm_status(exited, exit_code) as u8;
+            assert_eq!(
+                actual_state_hash, expected_state_hash,
+                "Incorrect state hash"
+            );
+        }
+    }
+
+    #[test]
+    fn test_hello() {
+        let elf_bytes = include_bytes!("../../../../example/bin/hello.elf");
+        let mut state = load_elf(elf_bytes).unwrap();
+        patch::patch_go(elf_bytes, &state).unwrap();
+        patch::patch_stack(&mut state).unwrap();
+
+        let out = BufWriter::new(Vec::default());
+        let err = BufWriter::new(Vec::default());
+        let mut ins =
+            InstrumentedState::new(state, StaticOracle::new(b"hello world".to_vec()), out, err);
+
+        for _ in 0..400_000 {
+            if ins.state.exited {
+                break;
+            }
+            ins.step(false).unwrap();
+        }
+
+        assert!(ins.state.exited, "must exit");
+        assert_eq!(ins.state.exit_code, 0, "must exit with 0");
+
+        assert_eq!(
+            String::from_utf8(ins.std_out.buffer().to_vec()).unwrap(),
+            "hello world!\n"
+        );
+        assert_eq!(
+            String::from_utf8(ins.std_err.buffer().to_vec()).unwrap(),
+            ""
+        );
+    }
+
+    #[test]
+    fn test_claim() {
+        let elf_bytes = include_bytes!("../../../../example/bin/claim.elf");
+        let mut state = load_elf(elf_bytes).unwrap();
+        patch::patch_go(elf_bytes, &state).unwrap();
+        patch::patch_stack(&mut state).unwrap();
+
+        let out = BufWriter::new(Vec::default());
+        let err = BufWriter::new(Vec::default());
+        let mut ins = InstrumentedState::new(state, ClaimTestOracle::default(), out, err);
+
+        for _ in 0..2_000_000 {
+            if ins.state.exited {
+                break;
+            }
+            ins.step(false).unwrap();
+        }
+
+        assert!(ins.state.exited, "must exit");
+        assert_eq!(ins.state.exit_code, 0, "must exit with 0");
+
+        assert_eq!(
+            String::from_utf8(ins.std_out.buffer().to_vec()).unwrap(),
+            format!(
+                "computing {} * {} + {}\nclaim {} is good!\n",
+                ClaimTestOracle::S,
+                ClaimTestOracle::A,
+                ClaimTestOracle::B,
+                ClaimTestOracle::S * ClaimTestOracle::A + ClaimTestOracle::B
+            )
+        );
+        assert_eq!(
+            String::from_utf8(ins.std_err.buffer().to_vec()).unwrap(),
+            "started!"
+        );
+    }
 }

crates/mipsevm/src/mips/mips_vm.rs

@@ -331,6 +331,7 @@ where
                                 break;
                             }
                         }
+                        v0 = a2;
                     }
                     Ok(Fd::PreimageWrite) => {
                         let effective_address = a1 & 0xFFFFFFFC;

crates/mipsevm/src/patch.rs

@@ -50,7 +50,7 @@ pub fn load_elf(raw: &[u8]) -> Result<State> {
             continue;
         }
 
-        let section_data = elf.segment_data(&header)?;
+        let section_data = &elf.segment_data(&header)?[..header.p_filesz as usize];
         let mut reader: Box<dyn Read> = Box::new(section_data);
 
         if header.p_filesz != header.p_memsz {
@@ -106,7 +106,8 @@ pub fn load_elf(raw: &[u8]) -> Result<State> {
 /// ### Returns
 /// - `Ok(())` if the patch was successful
 /// - `Err(_)` if the patch failed
-pub fn patch_go(elf: ElfBytes<AnyEndian>, state: &State) -> Result<()> {
+pub fn patch_go(raw: &[u8], state: &State) -> Result<()> {
+    let elf = ElfBytes::<AnyEndian>::minimal_parse(raw)?;
     let (parsing_table, string_table) = elf
         .symbol_table()?
         .ok_or(anyhow::anyhow!("Failed to load ELF symbol table"))?;
@@ -116,9 +117,12 @@ pub fn patch_go(elf: ElfBytes<AnyEndian>, state: &State) -> Result<()> {
         let name = string_table.get(symbol_idx as usize)?;
 
         if GO_SYMBOLS.contains(&name) {
+            // MIPS32 patch: ret (pseudo instruction)
+            // 03e00008 = jr $ra = ret (pseudo instruction)
+            // 00000000 = nop (executes with delay-slot, but does nothing)
             state.memory.borrow_mut().set_memory_range(
                 symbol.st_value as u32,
-                [0x03, 0xe0, 0x00, 0x08, 0, 0, 0, 0].as_slice(),
+                [0x03, 0xe0, 0x00, 0x08, 0x00, 0x00, 0x00, 0x00].as_slice(),
             )?;
         } else if name == "runtime.MemProfileRate" {
             // disable mem profiling, to avoid a lot of unnecessary floating point ops
@@ -146,7 +150,7 @@ pub fn patch_stack(state: &mut State) -> Result<()> {
     // Allocate 1 page for the initial stack data, and 16KB = 4 pages for the stack to grow.
     state.memory.borrow_mut().set_memory_range(
         ptr - 4 * page::PAGE_SIZE as u32,
-        [0; page::PAGE_SIZE * 5].as_slice(),
+        [0u8; page::PAGE_SIZE * 5].as_slice(),
     )?;
     state.registers[29] = ptr;
 
@@ -175,7 +179,7 @@ pub fn patch_stack(state: &mut State) -> Result<()> {
 }
 
 /// A multi reader is a reader that reads from the first reader until it returns 0, then reads from the second reader.
-struct MultiReader<R1: Read, R2: Read>(R1, R2);
+pub struct MultiReader<R1: Read, R2: Read>(R1, R2);
 
 impl<R1: Read, R2: Read> Read for MultiReader<R1, R2> {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {

crates/mipsevm/src/test_utils/evm.rs

@@ -214,14 +214,15 @@ impl MipsEVM<CacheDB<EmptyDB>> {
 mod test {
     use super::*;
     use crate::{
-        test_utils::{StaticOracle, BASE_ADDR_END, END_ADDR},
+        patch,
+        test_utils::{ClaimTestOracle, StaticOracle, BASE_ADDR_END, END_ADDR},
         Address, InstrumentedState, Memory, State,
     };
     use revm::primitives::ExecutionResult;
     use std::{
         cell::RefCell,
         fs,
-        io::{self, BufReader},
+        io::{self, BufReader, BufWriter},
         path::PathBuf,
         rc::Rc,
     };
@@ -447,4 +448,106 @@ mod test {
             assert!(mips_evm.step(step_witness).is_err());
         }
     }
+
+    #[test]
+    fn test_hello_evm() {
+        let mut mips_evm = MipsEVM::new();
+        mips_evm.try_init().unwrap();
+
+        let elf_bytes = include_bytes!("../../../../example/bin/hello.elf");
+        let mut state = patch::load_elf(elf_bytes).unwrap();
+        patch::patch_go(elf_bytes, &mut state).unwrap();
+        patch::patch_stack(&mut state).unwrap();
+
+        let mut instrumented =
+            InstrumentedState::new(state, StaticOracle::default(), io::stdout(), io::stderr());
+
+        for i in 0..400_000 {
+            if instrumented.state.exited {
+                break;
+            }
+
+            if i % 1000 == 0 {
+                let instruction = instrumented
+                    .state
+                    .memory
+                    .borrow_mut()
+                    .get_memory(instrumented.state.pc as Address)
+                    .unwrap();
+                println!(
+                    "step: {} pc: 0x{:08x} instruction: {:08x}",
+                    instrumented.state.step, instrumented.state.pc, instruction
+                );
+            }
+
+            let step_witness = instrumented.step(true).unwrap().unwrap();
+
+            let evm_post = mips_evm.step(step_witness).unwrap();
+            let rust_post = instrumented.state.encode_witness().unwrap();
+            assert_eq!(evm_post, rust_post);
+        }
+
+        assert!(instrumented.state.exited, "Must complete program");
+        assert_eq!(instrumented.state.exit_code, 0, "Must exit with 0");
+    }
+
+    #[test]
+    fn test_claim_evm() {
+        let mut mips_evm = MipsEVM::new();
+        mips_evm.try_init().unwrap();
+
+        let elf_bytes = include_bytes!("../../../../example/bin/claim.elf");
+        let mut state = patch::load_elf(elf_bytes).unwrap();
+        patch::patch_go(elf_bytes, &mut state).unwrap();
+        patch::patch_stack(&mut state).unwrap();
+
+        let out_buf = BufWriter::new(Vec::default());
+        let err_buf = BufWriter::new(Vec::default());
+
+        let mut instrumented =
+            InstrumentedState::new(state, ClaimTestOracle::default(), out_buf, err_buf);
+
+        for i in 0..2_000_000 {
+            if instrumented.state.exited {
+                break;
+            }
+
+            if i % 1000 == 0 {
+                let instruction = instrumented
+                    .state
+                    .memory
+                    .borrow_mut()
+                    .get_memory(instrumented.state.pc as Address)
+                    .unwrap();
+                println!(
+                    "step: {} pc: 0x{:08x} instruction: {:08x}",
+                    instrumented.state.step, instrumented.state.pc, instruction
+                );
+            }
+
+            let step_witness = instrumented.step(true).unwrap().unwrap();
+
+            let evm_post = mips_evm.step(step_witness).unwrap();
+            let rust_post = instrumented.state.encode_witness().unwrap();
+            assert_eq!(evm_post, rust_post);
+        }
+
+        assert!(instrumented.state.exited, "Must complete program");
+        assert_eq!(instrumented.state.exit_code, 0, "Must exit with 0");
+
+        assert_eq!(
+            String::from_utf8(instrumented.std_out.buffer().to_vec()).unwrap(),
+            format!(
+                "computing {} * {} + {}\nclaim {} is good!\n",
+                ClaimTestOracle::S,
+                ClaimTestOracle::A,
+                ClaimTestOracle::B,
+                ClaimTestOracle::S * ClaimTestOracle::A + ClaimTestOracle::B
+            )
+        );
+        assert_eq!(
+            String::from_utf8(instrumented.std_err.buffer().to_vec()).unwrap(),
+            "started!"
+        );
+    }
 }