feat(l1): support snap capability message GetStorageRanges

.github/workflows/hive_and_assertoor.yaml

@@ -66,7 +66,7 @@ jobs:
             run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="discv4"
           - simulation: snap
             name: "Devp2p snap tests"
-            run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="/AccountRange"
+            run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="/AccountRange|StorageRanges"
           - simulation: engine
             name: "Engine tests"
             run_command:  make run-hive-on-latest SIMULATION=ethereum/engine TEST_PATTERN="/Blob Transactions On Block 1, Cancun Genesis|Blob Transactions On Block 1, Shanghai Genesis|Blob Transaction Ordering, Single Account, Single Blob|Blob Transaction Ordering, Single Account, Dual Blob|Blob Transaction Ordering, Multiple Accounts|Replace Blob Transactions|Parallel Blob Transactions|ForkchoiceUpdatedV3 Modifies Payload ID on Different Beacon Root|NewPayloadV3 After Cancun|NewPayloadV3 Versioned Hashes|ForkchoiceUpdated Version on Payload Request"

crates/networking/p2p/rlpx/connection.rs

@@ -1,6 +1,6 @@
 use crate::{
     rlpx::{eth::backend, handshake::encode_ack_message, message::Message, p2p, utils::id2pubkey},
-    snap::process_account_range_request,
+    snap::{process_account_range_request, process_storage_ranges_request},
     MAX_DISC_PACKET_SIZE,
 };
 
@@ -150,6 +150,11 @@ impl<S: AsyncWrite + AsyncRead + std::marker::Unpin> RLPxConnection<S> {
                                 process_account_range_request(req, self.storage.clone())?;
                             self.send(Message::AccountRange(response)).await
                         }
+                        Message::GetStorageRanges(req) => {
+                            let response =
+                                process_storage_ranges_request(req, self.storage.clone())?;
+                            self.send(Message::StorageRanges(response)).await
+                        }
                         // TODO: Add new message types and handlers as they are implemented
                         message => return Err(RLPxError::UnexpectedMessage(message)),
                     };

crates/networking/p2p/rlpx/message.rs

@@ -4,7 +4,7 @@ use std::fmt::Display;
 
 use super::eth::status::StatusMessage;
 use super::p2p::{DisconnectMessage, HelloMessage, PingMessage, PongMessage};
-use super::snap::{AccountRange, GetAccountRange};
+use super::snap::{AccountRange, GetAccountRange, GetStorageRanges, StorageRanges};
 
 use ethereum_rust_rlp::encode::RLPEncode;
 
@@ -23,6 +23,8 @@ pub(crate) enum Message {
     // snap capability
     GetAccountRange(GetAccountRange),
     AccountRange(AccountRange),
+    GetStorageRanges(GetStorageRanges),
+    StorageRanges(StorageRanges),
 }
 
 impl Message {
@@ -35,6 +37,10 @@ impl Message {
             0x10 => Ok(Message::Status(StatusMessage::decode(msg_data)?)),
             0x21 => Ok(Message::GetAccountRange(GetAccountRange::decode(msg_data)?)),
             0x22 => Ok(Message::AccountRange(AccountRange::decode(msg_data)?)),
+            0x23 => Ok(Message::GetStorageRanges(GetStorageRanges::decode(
+                msg_data,
+            )?)),
+            0x24 => Ok(Message::StorageRanges(StorageRanges::decode(msg_data)?)),
             _ => Err(RLPDecodeError::MalformedData),
         }
     }
@@ -54,6 +60,14 @@ impl Message {
                 0x22_u8.encode(buf);
                 msg.encode(buf)
             }
+            Message::GetStorageRanges(msg) => {
+                0x23_u8.encode(buf);
+                msg.encode(buf)
+            }
+            Message::StorageRanges(msg) => {
+                0x24_u8.encode(buf);
+                msg.encode(buf)
+            }
         }
     }
 }
@@ -68,6 +82,8 @@ impl Display for Message {
             Message::Status(_) => "eth:Status".fmt(f),
             Message::GetAccountRange(_) => "snap:GetAccountRange".fmt(f),
             Message::AccountRange(_) => "snap:AccountRange".fmt(f),
+            Message::GetStorageRanges(_) => "snap:GetStorageRanges".fmt(f),
+            Message::StorageRanges(_) => "snap:StorageRanges".fmt(f),
         }
     }
 }

crates/networking/p2p/rlpx/snap.rs

@@ -33,6 +33,23 @@ pub(crate) struct AccountRange {
     pub proof: Vec<Bytes>,
 }
 
+#[derive(Debug)]
+pub(crate) struct GetStorageRanges {
+    pub id: u64,
+    pub root_hash: H256,
+    pub account_hashes: Vec<H256>,
+    pub starting_hash: H256,
+    pub limit_hash: H256,
+    pub response_bytes: u64,
+}
+
+#[derive(Debug)]
+pub(crate) struct StorageRanges {
+    pub id: u64,
+    pub slots: Vec<Vec<StorageSlot>>,
+    pub proof: Vec<Bytes>,
+}
+
 impl RLPxMessage for GetAccountRange {
     fn encode(&self, buf: &mut dyn BufMut) -> Result<(), RLPEncodeError> {
         let mut encoded_data = vec![];
@@ -55,11 +72,12 @@ impl RLPxMessage for GetAccountRange {
             .decompress_vec(msg_data)
             .map_err(|e| RLPDecodeError::Custom(e.to_string()))?;
         let decoder = Decoder::new(&decompressed_data)?;
-        let (id, decoder): (u64, _) = decoder.decode_field("request-id")?;
-        let (root_hash, decoder): (H256, _) = decoder.decode_field("rootHash")?;
-        let (starting_hash, decoder): (H256, _) = decoder.decode_field("startingHash")?;
-        let (limit_hash, decoder): (H256, _) = decoder.decode_field("limitHash")?;
-        let (response_bytes, _): (u64, _) = decoder.decode_field("responseBytes")?;
+        let (id, decoder) = decoder.decode_field("request-id")?;
+        let (root_hash, decoder) = decoder.decode_field("rootHash")?;
+        let (starting_hash, decoder) = decoder.decode_field("startingHash")?;
+        let (limit_hash, decoder) = decoder.decode_field("limitHash")?;
+        let (response_bytes, decoder) = decoder.decode_field("responseBytes")?;
+        decoder.finish()?;
 
         Ok(Self {
             id,
@@ -104,6 +122,77 @@ impl RLPxMessage for AccountRange {
     }
 }
 
+impl RLPxMessage for GetStorageRanges {
+    fn encode(&self, buf: &mut dyn BufMut) -> Result<(), RLPEncodeError> {
+        let mut encoded_data = vec![];
+        Encoder::new(&mut encoded_data)
+            .encode_field(&self.id)
+            .encode_field(&self.root_hash)
+            .encode_field(&self.account_hashes)
+            .encode_field(&self.starting_hash)
+            .encode_field(&self.limit_hash)
+            .encode_field(&self.response_bytes)
+            .finish();
+
+        let msg_data = snappy_encode(encoded_data)?;
+        buf.put_slice(&msg_data);
+        Ok(())
+    }
+
+    fn decode(msg_data: &[u8]) -> Result<Self, RLPDecodeError> {
+        let mut snappy_decoder = SnappyDecoder::new();
+        let decompressed_data = snappy_decoder
+            .decompress_vec(msg_data)
+            .map_err(|e| RLPDecodeError::Custom(e.to_string()))?;
+        let decoder = Decoder::new(&decompressed_data)?;
+        let (id, decoder) = decoder.decode_field("request-id")?;
+        let (root_hash, decoder) = decoder.decode_field("rootHash")?;
+        let (account_hashes, decoder) = decoder.decode_field("accountHashes")?;
+        let (starting_hash, decoder) = decoder.decode_field("startingHash")?;
+        let (limit_hash, decoder) = decoder.decode_field("limitHash")?;
+        let (response_bytes, decoder) = decoder.decode_field("responseBytes")?;
+        decoder.finish()?;
+
+        Ok(Self {
+            id,
+            root_hash,
+            starting_hash,
+            account_hashes,
+            limit_hash,
+            response_bytes,
+        })
+    }
+}
+
+impl RLPxMessage for StorageRanges {
+    fn encode(&self, buf: &mut dyn BufMut) -> Result<(), RLPEncodeError> {
+        let mut encoded_data = vec![];
+        Encoder::new(&mut encoded_data)
+            .encode_field(&self.id)
+            .encode_field(&self.slots)
+            .encode_field(&self.proof)
+            .finish();
+
+        let msg_data = snappy_encode(encoded_data)?;
+        buf.put_slice(&msg_data);
+        Ok(())
+    }
+
+    fn decode(msg_data: &[u8]) -> Result<Self, RLPDecodeError> {
+        let mut snappy_decoder = SnappyDecoder::new();
+        let decompressed_data = snappy_decoder
+            .decompress_vec(msg_data)
+            .map_err(|e| RLPDecodeError::Custom(e.to_string()))?;
+        let decoder = Decoder::new(&decompressed_data)?;
+        let (id, decoder) = decoder.decode_field("request-id")?;
+        let (slots, decoder) = decoder.decode_field("slots")?;
+        let (proof, decoder) = decoder.decode_field("proof")?;
+        decoder.finish()?;
+
+        Ok(Self { id, slots, proof })
+    }
+}
+
 // Intermediate structures
 
 #[derive(Debug)]
@@ -120,6 +209,12 @@ pub struct AccountStateSlim {
     pub code_hash: Bytes,
 }
 
+#[derive(Debug)]
+pub struct StorageSlot {
+    pub hash: H256,
+    pub data: U256,
+}
+
 impl RLPEncode for AccountRangeUnit {
     fn encode(&self, buf: &mut dyn BufMut) {
         Encoder::new(buf)
@@ -209,3 +304,21 @@ impl From<AccountStateSlim> for AccountState {
         }
     }
 }
+
+impl RLPEncode for StorageSlot {
+    fn encode(&self, buf: &mut dyn BufMut) {
+        Encoder::new(buf)
+            .encode_field(&self.hash)
+            .encode_field(&self.data)
+            .finish();
+    }
+}
+
+impl RLPDecode for StorageSlot {
+    fn decode_unfinished(rlp: &[u8]) -> Result<(Self, &[u8]), RLPDecodeError> {
+        let decoder = Decoder::new(rlp)?;
+        let (hash, decoder) = decoder.decode_field("hash")?;
+        let (data, decoder) = decoder.decode_field("data")?;
+        Ok((Self { hash, data }, decoder.finish()?))
+    }
+}

crates/networking/p2p/snap.rs

@@ -2,7 +2,10 @@ use bytes::Bytes;
 use ethereum_rust_rlp::encode::RLPEncode;
 use ethereum_rust_storage::{error::StoreError, Store};
 
-use crate::rlpx::snap::{AccountRange, AccountRangeUnit, AccountStateSlim, GetAccountRange};
+use crate::rlpx::snap::{
+    AccountRange, AccountRangeUnit, AccountStateSlim, GetAccountRange, GetStorageRanges,
+    StorageRanges, StorageSlot,
+};
 
 pub fn process_account_range_request(
     request: GetAccountRange,
@@ -36,6 +39,65 @@ pub fn process_account_range_request(
     })
 }
 
+pub fn process_storage_ranges_request(
+    request: GetStorageRanges,
+    store: Store,
+) -> Result<StorageRanges, StoreError> {
+    let mut slots = vec![];
+    let mut proof = vec![];
+    let mut bytes_used = 0;
+
+    for hashed_address in request.account_hashes {
+        let mut account_slots = vec![];
+        let mut res_capped = false;
+
+        if let Some(storage_iter) = store.iter_storage(request.root_hash, hashed_address)? {
+            for (hash, data) in storage_iter {
+                if hash >= request.starting_hash {
+                    bytes_used += 64_u64; // slot size
+                    account_slots.push(StorageSlot { hash, data });
+                }
+                if hash >= request.limit_hash || bytes_used >= request.response_bytes {
+                    if bytes_used >= request.response_bytes {
+                        res_capped = true;
+                    }
+                    break;
+                }
+            }
+        }
+
+        // Generate proofs only if the response doesn't contain the full storage range for the account
+        // Aka if the starting hash is not zero or if the response was capped due to byte limit
+        if !request.starting_hash.is_zero() || res_capped && !account_slots.is_empty() {
+            proof.extend(
+                store
+                    .get_storage_range_proof(
+                        request.root_hash,
+                        hashed_address,
+                        request.starting_hash,
+                        account_slots.last().map(|acc| acc.hash),
+                    )?
+                    .unwrap_or_default()
+                    .iter()
+                    .map(|bytes| Bytes::copy_from_slice(bytes)),
+            );
+        }
+
+        if !account_slots.is_empty() {
+            slots.push(account_slots);
+        }
+
+        if bytes_used >= request.response_bytes {
+            break;
+        }
+    }
+    Ok(StorageRanges {
+        id: request.id,
+        slots,
+        proof,
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;

crates/storage/store/engines/api.rs

@@ -2,7 +2,7 @@ use bytes::Bytes;
 use ethereum_rust_core::types::{
     Block, BlockBody, BlockHash, BlockHeader, BlockNumber, ChainConfig, Index, Receipt, Transaction,
 };
-use ethereum_types::{Address, H256, U256};
+use ethereum_types::{H256, U256};
 use std::{fmt::Debug, panic::RefUnwindSafe};
 
 use crate::error::StoreError;
@@ -201,7 +201,7 @@ pub trait StoreEngine: Debug + Send + Sync + RefUnwindSafe {
     // Obtain a storage trie from the given address and storage_root
     // Doesn't check if the account is stored
     // Used for internal store operations
-    fn open_storage_trie(&self, address: Address, storage_root: H256) -> Trie;
+    fn open_storage_trie(&self, hashed_address: H256, storage_root: H256) -> Trie;
 
     // Obtain a state trie from the given state root
     // Doesn't check if the state root is valid

crates/storage/store/engines/in_memory.rs

@@ -4,7 +4,7 @@ use ethereum_rust_core::types::{
     Block, BlockBody, BlockHash, BlockHeader, BlockNumber, ChainConfig, Index, Receipt,
 };
 use ethereum_rust_trie::{InMemoryTrieDB, Trie};
-use ethereum_types::{Address, H256, U256};
+use ethereum_types::{H256, U256};
 use std::{
     collections::HashMap,
     fmt::Debug,
@@ -31,7 +31,8 @@ struct StoreInner {
     transaction_locations: HashMap<H256, Vec<(BlockNumber, BlockHash, Index)>>,
     receipts: HashMap<BlockHash, HashMap<Index, Receipt>>,
     state_trie_nodes: NodeMap,
-    storage_trie_nodes: HashMap<Address, NodeMap>,
+    // A storage trie for each hashed account address
+    storage_trie_nodes: HashMap<H256, NodeMap>,
     // TODO (#307): Remove TotalDifficulty.
     block_total_difficulties: HashMap<BlockHash, U256>,
     // Stores local blocks by payload id
@@ -295,9 +296,9 @@ impl StoreEngine for Store {
         Ok(self.inner().chain_data.pending_block_number)
     }
 
-    fn open_storage_trie(&self, address: Address, storage_root: H256) -> Trie {
+    fn open_storage_trie(&self, hashed_address: H256, storage_root: H256) -> Trie {
         let mut store = self.inner();
-        let trie_backend = store.storage_trie_nodes.entry(address).or_default();
+        let trie_backend = store.storage_trie_nodes.entry(hashed_address).or_default();
         let db = Box::new(InMemoryTrieDB::new(trie_backend.clone()));
         Trie::open(db, storage_root)
     }

crates/storage/store/engines/libmdbx.rs

@@ -12,7 +12,7 @@ use ethereum_rust_core::types::{
 use ethereum_rust_rlp::decode::RLPDecode;
 use ethereum_rust_rlp::encode::RLPEncode;
 use ethereum_rust_trie::{LibmdbxDupsortTrieDB, LibmdbxTrieDB, Trie};
-use ethereum_types::{Address, H256, U256};
+use ethereum_types::{H256, U256};
 use libmdbx::orm::{Decodable, Encodable, Table};
 use libmdbx::{
     dupsort,
@@ -324,10 +324,10 @@ impl StoreEngine for Store {
         }
     }
 
-    fn open_storage_trie(&self, address: Address, storage_root: H256) -> Trie {
-        let db = Box::new(LibmdbxDupsortTrieDB::<StorageTriesNodes, [u8; 20]>::new(
+    fn open_storage_trie(&self, hashed_address: H256, storage_root: H256) -> Trie {
+        let db = Box::new(LibmdbxDupsortTrieDB::<StorageTriesNodes, [u8; 32]>::new(
             self.db.clone(),
-            address.0,
+            hashed_address.0,
         ));
         Trie::open(db, storage_root)
     }
@@ -466,8 +466,8 @@ dupsort!(
 
 dupsort!(
     /// Table containing all storage trie's nodes
-    /// Each node is stored by address and node hash in order to keep different storage trie's nodes separate
-    ( StorageTriesNodes ) ([u8;20], [u8;33])[[u8;20]] => Vec<u8>
+    /// Each node is stored by hashed account address and node hash in order to keep different storage trie's nodes separate
+    ( StorageTriesNodes ) ([u8;32], [u8;33])[[u8;32]] => Vec<u8>
 );
 
 dupsort!(