feat: fallback preimage recovery w/o debug_dbGet

Cargo.toml

@@ -37,6 +37,7 @@ futures = "0.3"
 url = "2.3"
 thiserror = "1.0.61"
 hex-literal = "0.4.1"
+rayon = "1.10.0"
 
 # workspace
 rsp-rpc-db = { path = "./crates/storage/rpc-db" }

README.md

@@ -22,9 +22,9 @@ and the command `rsp` will be installed.
 
 ### RPC Node Requirement
 
-RSP fetches block and state data from a JSON-RPC node. **But, you must use a RPC node that supports the `debug_dbGet` endpoint.**
+RSP fetches block and state data from a JSON-RPC node. It's recommended that you use a RPC node that supports the `debug_dbGet` endpoint.
 
-This is required because in some cases the host needs to recover the preimage of a [Merkle Patricia Trie](https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie/) node that's referenced by hash. To do this, the host utilizes the [`debug_dbGet` endpoint](https://geth.ethereum.org/docs/interacting-with-geth/rpc/ns-debug#debugdbget) of a Geth node running with options `--state.scheme=hash`, which is the default, and `--gcmode=archive`. An example command for running the node is:
+This is recommended because in some cases the host needs to recover the preimage of a [Merkle Patricia Trie](https://ethereum.org/en/developers/docs/data-structures-and-encoding/patricia-merkle-trie/) node that's referenced by hash. To do this, the host utilizes the [`debug_dbGet` endpoint](https://geth.ethereum.org/docs/interacting-with-geth/rpc/ns-debug#debugdbget) of a Geth node running with options `--state.scheme=hash`, which is the default, and `--gcmode=archive`. An example command for running the node is:
 
 ```bash
 geth \
@@ -33,7 +33,7 @@ geth \
   --http.api=eth,debug
 ```
 
-When running the host CLI or integration tests, **make sure to use an RPC URL pointing to a Geth node running with said options**, or errors will arise when preimage recovery is needed. You can reach out to the Succinct team to access an RPC URL that supports this endpoint.
+However, in the absence of the `debug_dbGet` method, the host is able to fall back to a less efficient process of recovering the preimages via the standard `eth_getProof`. The fallback works in most cases but not all, so if you encounter a preimage recovery failure, you can reach out to the Succinct team to access an RPC URL that supports `debug_dbGet`.
 
 ### Running the CLI
 

crates/mpt/src/lib.rs

@@ -12,7 +12,7 @@ use reth_trie::{
     EMPTY_ROOT_HASH,
 };
 use revm_primitives::{keccak256, HashMap};
-use rsp_primitives::storage::ExtDatabaseRef;
+use rsp_primitives::storage::{ExtDatabaseRef, PreimageContext};
 
 /// Computes the state root of a block's Merkle Patricia Trie given an [ExecutionOutcome] and a list
 /// of [EIP1186AccountProofResponse] storage proofs.
@@ -76,6 +76,7 @@ where
                     (storage_nibbles.clone(), encoded, storage_proof.proof.clone())
                 }),
                 db,
+                Some(address),
             )?
         };
         storage_roots.insert(hashed_address, root);
@@ -102,13 +103,15 @@ where
             (account_nibbles.clone(), encoded, proof.proof.clone())
         }),
         db,
+        None,
     )
 }
 
 /// Given a list of Merkle-Patricia proofs, compute the root of the trie.
 fn compute_root_from_proofs<DB>(
     items: impl IntoIterator<Item = (Nibbles, Option<Vec<u8>>, Vec<Bytes>)>,
     db: &DB,
+    root_context: Option<Address>,
 ) -> eyre::Result<B256>
 where
     DB: ExtDatabaseRef<Error: std::fmt::Debug>,
@@ -278,7 +281,12 @@ where
                     // technically have to modify this branch node, but the `alloy-trie` hash
                     // builder handles this automatically when supplying child nodes.
 
-                    let preimage = db.trie_node_ref(branch_hash).unwrap();
+                    let preimage = db
+                        .trie_node_ref(
+                            branch_hash,
+                            PreimageContext { address: &root_context, branch_path: &path },
+                        )
+                        .unwrap();
                     match TrieNode::decode(&mut &preimage[..]).unwrap() {
                         TrieNode::Branch(_) => {
                             // This node is a branch node that's referenced by hash. There's no need
@@ -365,7 +373,11 @@ mod tests {
     impl ExtDatabaseRef for TestTrieDb {
         type Error = std::convert::Infallible;
 
-        fn trie_node_ref(&self, hash: B256) -> std::result::Result<Bytes, Self::Error> {
+        fn trie_node_ref(
+            &self,
+            hash: B256,
+            _context: PreimageContext<'_>,
+        ) -> std::result::Result<Bytes, Self::Error> {
             for preimage in self.preimages.iter() {
                 if keccak256(preimage) == hash {
                     return std::result::Result::Ok(preimage.to_owned());
@@ -469,6 +481,7 @@ cb10a951f0e82cf2e461b98c4e5afb0348ccab5bb42180808080808080808080808080"
                 ],
             )],
             &TestTrieDb::new(),
+            None,
         )
         .unwrap();
 
@@ -587,6 +600,7 @@ f2e461b98c4e5afb0348ccab5bb421808080808080808080808080"
                 ),
             ],
             &TestTrieDb::new(),
+            None,
         )
         .unwrap();
 
@@ -653,6 +667,7 @@ f2e461b98c4e5afb0348ccab5bb421808080808080808080808080"
                 ),
             ],
             &TestTrieDb::new(),
+            None,
         )
         .unwrap();
 
@@ -713,6 +728,7 @@ f2e461b98c4e5afb0348ccab5bb421808080808080808080808080"
                 ),
             ],
             &TestTrieDb::new(),
+            None,
         )
         .unwrap();
 
@@ -757,6 +773,7 @@ f2e461b98c4e5afb0348ccab5bb421808080808080808080808080"
                 ],
             )],
             &TestTrieDb::new(),
+            None,
         )
         .unwrap();
 

crates/primitives/src/storage.rs

@@ -1,10 +1,21 @@
-use reth_primitives::{Bytes, B256};
+use reth_primitives::{Address, Bytes, B256};
+use reth_trie::Nibbles;
 
 /// Custom database access methods implemented by RSP storage backends.
 pub trait ExtDatabaseRef {
     /// The database error type.
     type Error;
 
     /// Gets the preimage of a trie node given its Keccak hash.
-    fn trie_node_ref(&self, hash: B256) -> Result<Bytes, Self::Error>;
+    fn trie_node_ref(&self, hash: B256, context: PreimageContext) -> Result<Bytes, Self::Error>;
+}
+
+/// Additional context for retrieving trie node preimages. These are useful when the JSON-RPC node
+/// does not serve the `debug_dbGet`.
+pub struct PreimageContext<'a> {
+    /// The account address if calculating a storage trie root; `None` if calculating the state
+    /// root.
+    pub address: &'a Option<Address>,
+    /// The trie key path of the branch child containing the hash whose preimage is being fetched.
+    pub branch_path: &'a Nibbles,
 }

crates/storage/rpc-db/Cargo.toml

@@ -12,6 +12,7 @@ tokio.workspace = true
 futures.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
+rayon.workspace = true
 
 # workspace
 rsp-witness-db.workspace = true
@@ -21,6 +22,7 @@ rsp-primitives.workspace = true
 reth-primitives.workspace = true
 reth-storage-errors.workspace = true
 reth-revm.workspace = true
+reth-trie.workspace = true
 
 # revm
 revm-primitives.workspace = true

crates/storage/rpc-db/src/lib.rs

@@ -1,19 +1,28 @@
-use std::{cell::RefCell, marker::PhantomData};
+use std::{cell::RefCell, iter::once, marker::PhantomData};
 
 use alloy_provider::Provider;
 use alloy_rpc_types::BlockId;
 use alloy_transport::Transport;
 use futures::future::join_all;
+use rayon::prelude::*;
 use reth_primitives::{
     revm_primitives::{AccountInfo, Bytecode},
     Address, Bytes, B256, U256,
 };
 use reth_revm::DatabaseRef;
 use reth_storage_errors::{db::DatabaseError, provider::ProviderError};
-use revm_primitives::{HashMap, HashSet};
-use rsp_primitives::{account_proof::AccountProofWithBytecode, storage::ExtDatabaseRef};
+use reth_trie::Nibbles;
+use revm_primitives::{keccak256, HashMap, HashSet};
+use rsp_primitives::{
+    account_proof::AccountProofWithBytecode,
+    storage::{ExtDatabaseRef, PreimageContext},
+};
 use rsp_witness_db::WitnessDb;
 
+/// The maximum number of addresses/slots to attempt for brute-forcing the key to be used for
+/// fetching trie node preimage via `eth_getProof`.
+const BRUTE_FORCE_LIMIT: u64 = 0xffffffff_u64;
+
 /// A database that fetches data from a [Provider] over a [Transport].
 #[derive(Debug, Clone)]
 pub struct RpcDb<T, P> {
@@ -42,6 +51,8 @@ pub enum RpcDbError {
     RpcError(String),
     #[error("failed to find block")]
     BlockNotFound,
+    #[error("failed to find trie node preimage")]
+    PreimageNotFound,
 }
 
 impl<T: Transport + Clone, P: Provider<T> + Clone> RpcDb<T, P> {
@@ -138,16 +149,30 @@ impl<T: Transport + Clone, P: Provider<T> + Clone> RpcDb<T, P> {
     }
 
     /// Fetch a trie node based on its Keccak hash using the `debug_dbGet` method.
-    pub async fn fetch_trie_node(&self, hash: B256) -> Result<Bytes, RpcDbError> {
+    pub async fn fetch_trie_node(
+        &self,
+        hash: B256,
+        context: PreimageContext<'_>,
+    ) -> Result<Bytes, RpcDbError> {
         tracing::info!("fetching trie node {}", hash);
 
         // Fetch the trie node value from a geth node with `state.scheme=hash`.
-        let value = self
-            .provider
-            .client()
-            .request::<_, Bytes>("debug_dbGet", (hash,))
-            .await
-            .map_err(|e| RpcDbError::RpcError(e.to_string()))?;
+        let value = match self.provider.client().request::<_, Bytes>("debug_dbGet", (hash,)).await {
+            Ok(value) => value,
+            Err(_) => {
+                // The `debug_dbGet` method failed for some reason. Fall back to brute-forcing the
+                // slot/address needed to recover the preimage via the `eth_getProof` method
+                // instead.
+                tracing::debug!(
+                    "failed to fetch preimage from debug_dbGet; \
+                    falling back to using eth_getProof: address={:?}, prefix={:?}",
+                    context.address,
+                    context.branch_path
+                );
+
+                self.fetch_trie_node_via_proof(hash, context).await?
+            }
+        };
 
         // Record the trie node value to the state.
         self.trie_nodes.borrow_mut().insert(hash, value.clone());
@@ -224,6 +249,68 @@ impl<T: Transport + Clone, P: Provider<T> + Clone> RpcDb<T, P> {
 
         account_proofs
     }
+
+    /// Fetches a trie node via `eth_getProof` with a hacky workaround when `debug_dbGet` is not
+    /// available.
+    async fn fetch_trie_node_via_proof(
+        &self,
+        hash: B256,
+        context: PreimageContext<'_>,
+    ) -> Result<Bytes, RpcDbError> {
+        let (address, storage_keys) = match context.address {
+            Some(address) => {
+                // Computing storage root. Brute force the slot.
+                let slot = Self::find_key_preimage::<32>(context.branch_path)
+                    .ok_or(RpcDbError::PreimageNotFound)?;
+
+                (address.to_owned(), vec![slot.into()])
+            }
+            None => {
+                // Computing state root. Brute force the address.
+                let address = Self::find_key_preimage::<20>(context.branch_path)
+                    .ok_or(RpcDbError::PreimageNotFound)?;
+
+                (address.into(), vec![])
+            }
+        };
+
+        let account_proof = self
+            .provider
+            .get_proof(address, storage_keys)
+            .block_id(self.block)
+            .await
+            .map_err(|e| RpcDbError::RpcError(e.to_string()))?;
+
+        for proof in account_proof
+            .storage_proof
+            .into_iter()
+            .map(|storage_proof| storage_proof.proof)
+            .chain(once(account_proof.account_proof))
+        {
+            // The preimage we're looking for is more likely to be at the end of the proof.
+            for node in proof.into_iter().rev() {
+                if hash == keccak256(&node) {
+                    return Ok(node)
+                }
+            }
+        }
+
+        Err(RpcDbError::PreimageNotFound)
+    }
+
+    /// Uses brute force to locate a key path preimage that contains a certain prefix.
+    fn find_key_preimage<const BYTES: usize>(prefix: &Nibbles) -> Option<[u8; BYTES]> {
+        (0..BRUTE_FORCE_LIMIT).into_par_iter().find_map_any(|nonce| {
+            let mut buffer = [0u8; BYTES];
+            buffer[(BYTES - 8)..].copy_from_slice(&nonce.to_be_bytes());
+
+            if Nibbles::unpack(keccak256(buffer)).starts_with(prefix) {
+                Some(buffer)
+            } else {
+                None
+            }
+        })
+    }
 }
 
 impl<T: Transport + Clone, P: Provider<T> + Clone> DatabaseRef for RpcDb<T, P> {
@@ -269,11 +356,12 @@ impl<T: Transport + Clone, P: Provider<T> + Clone> DatabaseRef for RpcDb<T, P> {
 impl<T: Transport + Clone, P: Provider<T> + Clone> ExtDatabaseRef for RpcDb<T, P> {
     type Error = ProviderError;
 
-    fn trie_node_ref(&self, hash: B256) -> Result<Bytes, Self::Error> {
+    fn trie_node_ref(&self, hash: B256, context: PreimageContext) -> Result<Bytes, Self::Error> {
         let handle = tokio::runtime::Handle::try_current().map_err(|_| {
             ProviderError::Database(DatabaseError::Other("no tokio runtime found".to_string()))
         })?;
-        let result = tokio::task::block_in_place(|| handle.block_on(self.fetch_trie_node(hash)));
+        let result =
+            tokio::task::block_in_place(|| handle.block_on(self.fetch_trie_node(hash, context)));
         let value =
             result.map_err(|e| ProviderError::Database(DatabaseError::Other(e.to_string())))?;
         Ok(value)

crates/storage/witness-db/src/lib.rs

@@ -4,7 +4,7 @@ use reth_primitives::{
 };
 use reth_storage_errors::provider::ProviderError;
 use revm_primitives::{Address, HashMap, U256};
-use rsp_primitives::storage::ExtDatabaseRef;
+use rsp_primitives::storage::{ExtDatabaseRef, PreimageContext};
 use serde::{Deserialize, Serialize};
 
 /// A database used to witness state inside the zkVM.
@@ -45,7 +45,7 @@ impl DatabaseRef for WitnessDb {
 impl ExtDatabaseRef for WitnessDb {
     type Error = ProviderError;
 
-    fn trie_node_ref(&self, hash: B256) -> Result<Bytes, Self::Error> {
+    fn trie_node_ref(&self, hash: B256, _context: PreimageContext) -> Result<Bytes, Self::Error> {
         // TODO: avoid cloning
         Ok(self.trie_nodes.get(&hash).unwrap().to_owned())
     }