From 397f368ecd8ab214858931e779b3bfa8f60942f5 Mon Sep 17 00:00:00 2001 From: Ganesh Vanahalli Date: Thu, 28 Mar 2024 11:38:50 -0400 Subject: [PATCH] Validate preimages in both JIT and Arbitrator [NIT-2377] --- arbitrator/Cargo.lock | 59 ++++++++++++++++++++++++++++++++++-- arbitrator/jit/Cargo.toml | 1 + arbitrator/jit/src/wavmio.rs | 17 +++++++++++ arbitrator/prover/Cargo.toml | 1 + arbitrator/prover/src/lib.rs | 48 +++++++++++++++++++++-------- 5 files changed, 110 insertions(+), 16 deletions(-) diff --git a/arbitrator/Cargo.lock b/arbitrator/Cargo.lock index 165fee89c3..6242986cbc 100644 --- a/arbitrator/Cargo.lock +++ b/arbitrator/Cargo.lock @@ -28,6 +28,18 @@ dependencies = [ "version_check", ] +[[package]] +name = "ahash" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91429305e9f0a25f6205c5b8e0d2db09e0708a7a6df0f42212bb56c32c8ac97a" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "0.7.19" @@ -43,6 +55,12 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "250f629c0161ad8107cf89319e990051fae62832fd343083bea452d93e2205fd" +[[package]] +name = "allocator-api2" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0942ffc6dcaadf03badf6e6a2d0228460359d5e34b57ccdc720b7382dfbd5ec5" + [[package]] name = "ansi_term" version = "0.11.0" @@ -645,7 +663,7 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" dependencies = [ - "ahash", + "ahash 0.7.6", ] [[package]] @@ -653,6 +671,10 @@ name = "hashbrown" version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2c6201b9ff9fd90a5a3bac2e56a830d0caa509576f0e503818ee82c181b3437a" +dependencies = [ + "ahash 0.8.6", + "allocator-api2", +] [[package]] name = "heck" @@ -787,6 +809,7 @@ dependencies = [ "parking_lot 0.12.1", "rand", "rand_pcg", + "sha2 0.9.9", "sha3 0.9.1", "structopt", "thiserror", @@ -885,6 +908,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "lru" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3262e75e648fce39813cb56ac41f3c3e3f65217ebf3844d818d1f9398cfb0dc" +dependencies = [ + "hashbrown 0.14.0", +] + [[package]] name = "mach" version = "0.3.2" @@ -1259,6 +1291,7 @@ dependencies = [ "hex", "lazy_static", "libc", + "lru", "nom", "nom-leb128", "num", @@ -1860,9 +1893,9 @@ checksum = "f1bddf1187be692e79c5ffeab891132dfb0f236ed36a43c7ed39f1165ee20191" [[package]] name = "version_check" -version = "0.9.3" +version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fecdca9a5291cc2b8dcf7dc02453fee791a280f3743cb0905f8822ae463b3fe" +checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" [[package]] name = "wasi" @@ -2397,6 +2430,26 @@ dependencies = [ "memchr", ] +[[package]] +name = "zerocopy" +version = "0.7.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74d4d3961e53fa4c9a25a8637fc2bfaf2595b3d3ae34875568a5cf64787716be" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ce1b18ccd8e73a9321186f97e46f9f04b778851177567b1975109d26a08d2a6" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.45", +] + [[package]] name = "zeroize" version = "1.7.0" diff --git a/arbitrator/jit/Cargo.toml b/arbitrator/jit/Cargo.toml index 75b3e3a74c..fd715eb47f 100644 --- a/arbitrator/jit/Cargo.toml +++ b/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"] diff --git a/arbitrator/jit/src/wavmio.rs b/arbitrator/jit/src/wavmio.rs index dfc7f21779..cebcdb26e9 100644 --- a/arbitrator/jit/src/wavmio.rs +++ b/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}"), diff --git a/arbitrator/prover/Cargo.toml b/arbitrator/prover/Cargo.toml index 51cbe84399..a596f0bf81 100644 --- a/arbitrator/prover/Cargo.toml +++ b/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" diff --git a/arbitrator/prover/src/lib.rs b/arbitrator/prover/src/lib.rs index c7610ab31f..1a8bf83b26 100644 --- a/arbitrator/prover/src/lib.rs +++ b/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> = RwLock::new(LruCache::new(NonZeroUsize::new(1024).unwrap())); +} + #[repr(C)] #[derive(Clone, Copy)] pub struct CByteArray { @@ -302,18 +308,34 @@ 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); + match CACHE_PREIMAGE_REHASH_CHECK.try_write() { + Ok(mut cache) => { + let _ = cache.pop(&cache_key); + cache.put(cache_key.clone(), true); + } + Err(_) => {} + }; } Some(data) },