🚧 WIP: Elf tests

crates/mipsevm/src/patch.rs

@@ -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"))?;

crates/mipsevm/src/test_utils/evm.rs

@@ -214,9 +214,11 @@ 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 elf::{endian::AnyEndian, ElfBytes};
     use revm::primitives::ExecutionResult;
     use std::{
         cell::RefCell,
@@ -447,4 +449,93 @@ 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]
+    #[ignore]
+    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 mut instrumented = InstrumentedState::new(
+            state,
+            ClaimTestOracle::default(),
+            io::stdout(),
+            io::stderr(),
+        );
+
+        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");
+    }
 }

crates/mipsevm/src/test_utils/mod.rs

@@ -1,8 +1,10 @@
 //! Testing utilities.
 
-use crate::PreimageOracle;
-use alloy_primitives::B256;
-use preimage_oracle::{Keccak256Key, Key};
+use crate::{utils::concat_fixed, PreimageOracle};
+use alloy_primitives::{hex, keccak256, B256};
+use once_cell::sync::Lazy;
+use preimage_oracle::{Keccak256Key, Key, LocalIndexKey};
+use revm::primitives::HashMap;
 
 pub mod evm;
 
@@ -12,6 +14,7 @@ pub const BASE_ADDR_END: u32 = 0xBF_FF_FF_F0;
 /// Used as the return-address for tests
 pub const END_ADDR: u32 = 0xA7_EF_00_D0;
 
+#[derive(Default)]
 pub struct StaticOracle {
     preimage_data: Vec<u8>,
 }
@@ -34,3 +37,97 @@ impl PreimageOracle for StaticOracle {
         Ok(self.preimage_data.as_slice())
     }
 }
+
+pub struct ClaimTestOracle {
+    images: HashMap<B256, Vec<u8>>,
+}
+
+impl ClaimTestOracle {
+    const S: u64 = 1000;
+    const A: u64 = 3;
+    const B: u64 = 4;
+    const DIFF: Lazy<[u8; 64]> = Lazy::new(|| {
+        concat_fixed(
+            keccak256(Self::A.to_be_bytes()).into(),
+            keccak256(Self::B.to_be_bytes()).into(),
+        )
+    });
+    const PRE_HASH: Lazy<B256> = Lazy::new(|| keccak256(Self::S.to_be_bytes()));
+    const DIFF_HASH: Lazy<B256> = Lazy::new(|| keccak256(Self::DIFF.as_slice()));
+}
+
+impl Default for ClaimTestOracle {
+    fn default() -> Self {
+        let mut s = Self {
+            images: HashMap::new(),
+        };
+
+        s.images
+            .insert((0 as LocalIndexKey).preimage_key(), Self::PRE_HASH.to_vec());
+        s.images.insert(
+            (1 as LocalIndexKey).preimage_key(),
+            Self::DIFF_HASH.to_vec(),
+        );
+        s.images.insert(
+            (2 as LocalIndexKey).preimage_key(),
+            (Self::S * Self::A + Self::B).to_be_bytes().to_vec(),
+        );
+
+        s
+    }
+}
+
+impl PreimageOracle for ClaimTestOracle {
+    fn hint(&mut self, value: &[u8]) {
+        let s = String::from_utf8(value.to_vec()).unwrap();
+        let parts: Vec<&str> = s.split(" ").collect();
+
+        assert_eq!(parts.len(), 2);
+
+        let part = hex::decode(parts[1]).unwrap();
+        assert_eq!(part.len(), 32);
+        let hash = B256::from_slice(&part);
+
+        match parts[0] {
+            "fetch-state" => {
+                assert_eq!(
+                    hash,
+                    *Self::PRE_HASH,
+                    "Expecting request for pre-state preimage"
+                );
+
+                self.images.insert(
+                    (*Self::PRE_HASH as Keccak256Key).preimage_key(),
+                    Self::S.to_be_bytes().to_vec(),
+                );
+            }
+            "fetch-diff" => {
+                assert_eq!(
+                    hash,
+                    *Self::DIFF_HASH,
+                    "Expecting request for diff preimage"
+                );
+                self.images.insert(
+                    (*Self::DIFF_HASH as Keccak256Key).preimage_key(),
+                    Self::DIFF.to_vec(),
+                );
+                self.images.insert(
+                    (keccak256(Self::A.to_be_bytes()) as Keccak256Key).preimage_key(),
+                    Self::A.to_be_bytes().to_vec(),
+                );
+                self.images.insert(
+                    (keccak256(Self::B.to_be_bytes()) as Keccak256Key).preimage_key(),
+                    Self::B.to_be_bytes().to_vec(),
+                );
+            }
+            _ => panic!("Unexpected hint: {}", parts[0]),
+        }
+    }
+
+    fn get(&self, key: B256) -> anyhow::Result<&[u8]> {
+        Ok(self
+            .images
+            .get(&key)
+            .ok_or(anyhow::anyhow!("No image for key"))?)
+    }
+}