Validate preimages in both JIT and Arbitrator [NIT-2377]

arbitrator/jit/Cargo.toml

@@ -18,6 +18,7 @@ structopt = "0.3.26"
 sha3 = "0.9.1"
 libc = "0.2.132"
 ouroboros = "0.16.0"
+sha2 = "0.9.9"
 
 [features]
 llvm = ["dep:wasmer-compiler-llvm"]

arbitrator/jit/src/wavmio.rs

@@ -8,6 +8,8 @@ use crate::{
 };
 
 use arbutil::{Color, PreimageType};
+use sha2::Sha256;
+use sha3::{Digest, Keccak256};
 use std::{
     io,
     io::{BufReader, BufWriter, ErrorKind},
@@ -192,6 +194,21 @@ pub fn resolve_preimage_impl(
         error!("Missing requested preimage for preimage type {preimage_type:?} hash {hash_hex} in {name}");
     };
 
+    // Check if preimage rehashes to the provided hash. Exclude blob preimages
+    let calculated_hash: [u8; 32] = match preimage_type {
+        PreimageType::Keccak256 => Keccak256::digest(preimage).into(),
+        PreimageType::Sha2_256 => Sha256::digest(preimage).into(),
+        PreimageType::EthVersionedHash => *hash,
+    };
+    if calculated_hash != *hash {
+        error!(
+            "Calculated hash {} of preimage {} does not match provided hash {}",
+            hex::encode(calculated_hash),
+            hex::encode(preimage),
+            hex::encode(*hash)
+        );
+    }
+
     let offset = match u32::try_from(offset) {
         Ok(offset) if offset % 32 == 0 => offset as usize,
         _ => error!("bad offset {offset} in {name}"),

arbitrator/prover/Cargo.toml

@@ -30,6 +30,7 @@ smallvec = { version = "1.10.0", features = ["serde"] }
 arbutil = { path = "../arbutil/" }
 c-kzg = "0.4.0" # TODO: look into switching to rust-kzg (no crates.io release or hosted rustdoc yet)
 sha2 = "0.9.9"
+lru = "0.12.3"
 
 [lib]
 name = "prover"

arbitrator/prover/src/lib.rs

@@ -18,19 +18,25 @@ pub mod wavm;
 use crate::machine::{argument_data_to_inbox, Machine};
 use arbutil::PreimageType;
 use eyre::Result;
+use lru::LruCache;
 use machine::{get_empty_preimage_resolver, GlobalState, MachineStatus, PreimageResolver};
 use static_assertions::const_assert_eq;
 use std::{
     ffi::CStr,
+    num::NonZeroUsize,
     os::raw::{c_char, c_int},
     path::Path,
     sync::{
         atomic::{self, AtomicU8},
-        Arc,
+        Arc, RwLock,
     },
 };
 use utils::{Bytes32, CBytes};
 
+lazy_static::lazy_static! {
+    static ref CACHE_PREIMAGE_REHASH_CHECK: RwLock<LruCache<(Bytes32, PreimageType), bool>> = RwLock::new(LruCache::new(NonZeroUsize::new(1024).unwrap()));
+}
+
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct CByteArray {
@@ -302,18 +308,31 @@ pub unsafe extern "C" fn arbitrator_set_preimage_resolver(
                 return None;
             }
             let data = CBytes::from_raw_parts(res.ptr, res.len as usize);
-            #[cfg(debug_assertions)]
-            match crate::utils::hash_preimage(&data, ty) {
-                Ok(have_hash) if have_hash.as_slice() == *hash => {}
-                Ok(got_hash) => panic!(
-                    "Resolved incorrect data for hash {} (rehashed to {})",
-                    hash,
-                    Bytes32(got_hash),
-                ),
-                Err(err) => panic!(
-                    "Failed to hash preimage from resolver (expecting hash {}): {}",
-                    hash, err,
-                ),
+            // Check if preimage rehashes to the provided hash
+            let cache_key = (hash, ty);
+            let cache = CACHE_PREIMAGE_REHASH_CHECK.read().unwrap();
+            if !cache.contains(&cache_key) {
+                drop(cache);
+                match crate::utils::hash_preimage(&data, ty) {
+                    Ok(have_hash) if have_hash.as_slice() == *hash => {}
+                    Ok(got_hash) => panic!(
+                        "Resolved incorrect data for hash {} (rehashed to {})",
+                        hash,
+                        Bytes32(got_hash),
+                    ),
+                    Err(err) => panic!(
+                        "Failed to hash preimage from resolver (expecting hash {}): {}",
+                        hash, err,
+                    ),
+                }
+                let mut cache = CACHE_PREIMAGE_REHASH_CHECK.write().unwrap();
+                cache.put(cache_key, true);
+            } else {
+                drop(cache);
+                if let Ok(mut cache) = CACHE_PREIMAGE_REHASH_CHECK.try_write() {
+                    let _ = cache.pop(&cache_key);
+                    cache.put(cache_key, true);
+                };
             }
             Some(data)
         },