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, Mutex,
     },
 };
 use utils::{Bytes32, CBytes};
 
+lazy_static::lazy_static! {
+    static ref BLOBHASH_PREIMAGE_CACHE: Mutex<LruCache<Bytes32, CBytes>> = Mutex::new(LruCache::new(NonZeroUsize::new(12).unwrap()));
+}
+
 #[repr(C)]
 #[derive(Clone, Copy)]
 pub struct CByteArray {
@@ -290,32 +296,49 @@ pub struct ResolvedPreimage {
     pub len: isize, // negative if not found
 }
 
+macro_rules! handle_preimage_resolution {
+    ($context:expr, $ty:expr, $hash:expr, $resolver:expr) => {{
+        let res = $resolver($context, $ty.into(), $hash.as_ptr());
+        if res.len < 0 {
+            return None;
+        }
+        let data = CBytes::from_raw_parts(res.ptr, res.len as usize);
+        // Check if preimage rehashes to the provided hash
+        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,
+            ),
+        }
+        Some(data)
+    }};
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn arbitrator_set_preimage_resolver(
     mach: *mut Machine,
     resolver: unsafe extern "C" fn(u64, u8, *const u8) -> ResolvedPreimage,
 ) {
     (*mach).set_preimage_resolver(Arc::new(
         move |context: u64, ty: PreimageType, hash: Bytes32| -> Option<CBytes> {
-            let res = resolver(context, ty.into(), hash.as_ptr());
-            if res.len < 0 {
+            if let PreimageType::EthVersionedHash = ty {
+                let mut cache = BLOBHASH_PREIMAGE_CACHE.lock().unwrap();
+                if cache.contains(&hash) {
+                    return cache.get(&hash).cloned();
+                }
+                if let Some(data) = handle_preimage_resolution!(context, ty, hash, resolver) {
+                    cache.put(hash, data.clone());
+                    return Some(data);
+                }
                 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,
-                ),
-            }
-            Some(data)
+            handle_preimage_resolution!(context, ty, hash, resolver)
         },
     ) as PreimageResolver);
 }