feat(l1): support snap capability message GetTrieNodes (#1031)

**Motivation**
Pass hive snap `GetTrieNodes` test
<!-- Why does this pull request exist? What are its goals? -->

**Description**

* Add trie method `get_node`
* Add snap capability messages `GetTrieNodes` and `TrieNodes`
* Support processing `GetTrieNodes` message

<!-- A clear and concise general description of the changes this PR
introduces -->

<!-- Link to issues: Resolves #111, Resolves #222 -->

Is part of #934

---------

Co-authored-by: Esteban Dimitroff Hodi <esteban.dimitroff@lambdaclass.com>
Co-authored-by: ElFantasma <estebandh@gmail.com>

.github/workflows/hive_and_assertoor.yaml

@@ -66,13 +66,17 @@ jobs:
             run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="discv4"
           - simulation: snap
             name: "Devp2p snap tests"
+<<<<<<< HEAD
+            run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="/AccountRange|StorageRanges|ByteCodes|TrieNodes"
+=======
             run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="/AccountRange|StorageRanges|ByteCodes"
           - simulation: eth
             name: "Devp2p eth tests"
             run_command: make run-hive SIMULATION=devp2p TEST_PATTERN="eth/getblockheaders"
           - simulation: eth
             name: "Devp2p eth tests"
             run_command: make run-hive SIMULATION=devp2p TEST_PATTERN="eth/getblockheaders|getblockbodies"
+>>>>>>> d15a89c9d7504b7f3eca3f20789dc38c449ee456
           - 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

@@ -11,6 +11,7 @@ use crate::{
     },
     snap::{
         process_account_range_request, process_byte_codes_request, process_storage_ranges_request,
+        process_trie_nodes_request,
     },
     MAX_DISC_PACKET_SIZE,
 };
@@ -194,6 +195,10 @@ impl<S: AsyncWrite + AsyncRead + std::marker::Unpin> RLPxConnection<S> {
                             let response = process_byte_codes_request(req, self.storage.clone())?;
                             self.send(Message::ByteCodes(response)).await?
                         }
+                        Message::GetTrieNodes(req) => {
+                            let response = process_trie_nodes_request(req, self.storage.clone())?;
+                            self.send(Message::TrieNodes(response)).await?
+                        }
                         // TODO: Add new message types and handlers as they are implemented
                         _ => return Err(RLPxError::MessageNotHandled()),
                     };

crates/networking/p2p/rlpx/error.rs

@@ -21,6 +21,8 @@ pub(crate) enum RLPxError {
     InvalidMessageLength(),
     #[error("Cannot handle message")]
     MessageNotHandled(),
+    #[error("Bad Request: {0}")]
+    BadRequest(String),
     #[error(transparent)]
     RLPDecodeError(#[from] RLPDecodeError),
     #[error(transparent)]

crates/networking/p2p/rlpx/message.rs

@@ -6,7 +6,8 @@ use super::eth::blocks::{BlockBodies, BlockHeaders, GetBlockBodies, GetBlockHead
 use super::eth::status::StatusMessage;
 use super::p2p::{DisconnectMessage, HelloMessage, PingMessage, PongMessage};
 use super::snap::{
-    AccountRange, ByteCodes, GetAccountRange, GetByteCodes, GetStorageRanges, StorageRanges,
+    AccountRange, ByteCodes, GetAccountRange, GetByteCodes, GetStorageRanges, GetTrieNodes,
+    StorageRanges, TrieNodes,
 };
 
 use ethereum_rust_rlp::encode::RLPEncode;
@@ -35,6 +36,8 @@ pub(crate) enum Message {
     StorageRanges(StorageRanges),
     GetByteCodes(GetByteCodes),
     ByteCodes(ByteCodes),
+    GetTrieNodes(GetTrieNodes),
+    TrieNodes(TrieNodes),
 }
 
 impl Message {
@@ -66,6 +69,8 @@ impl Message {
             0x24 => Ok(Message::StorageRanges(StorageRanges::decode(msg_data)?)),
             0x25 => Ok(Message::GetByteCodes(GetByteCodes::decode(msg_data)?)),
             0x26 => Ok(Message::ByteCodes(ByteCodes::decode(msg_data)?)),
+            0x27 => Ok(Message::GetTrieNodes(GetTrieNodes::decode(msg_data)?)),
+            0x28 => Ok(Message::TrieNodes(TrieNodes::decode(msg_data)?)),
             _ => Err(RLPDecodeError::MalformedData),
         }
     }
@@ -132,6 +137,14 @@ impl Message {
                 0x26_u8.encode(buf);
                 msg.encode(buf)
             }
+            Message::GetTrieNodes(msg) => {
+                0x27_u8.encode(buf);
+                msg.encode(buf)
+            }
+            Message::TrieNodes(msg) => {
+                0x28_u8.encode(buf);
+                msg.encode(buf)
+            }
         }
     }
 }
@@ -154,6 +167,8 @@ impl Display for Message {
             Message::StorageRanges(_) => "snap:StorageRanges".fmt(f),
             Message::GetByteCodes(_) => "snap:GetByteCodes".fmt(f),
             Message::ByteCodes(_) => "snap:ByteCodes".fmt(f),
+            Message::GetTrieNodes(_) => "snap:GetTrieNodes".fmt(f),
+            Message::TrieNodes(_) => "snap:TrieNodes".fmt(f),
         }
     }
 }

crates/networking/p2p/rlpx/snap.rs

@@ -64,6 +64,22 @@ pub(crate) struct ByteCodes {
     pub codes: Vec<Bytes>,
 }
 
+#[derive(Debug)]
+pub(crate) struct GetTrieNodes {
+    pub id: u64,
+    pub root_hash: H256,
+    // [[acc_path, slot_path_1, slot_path_2,...]...]
+    // The paths can be either full paths (hash) or only the partial path (compact-encoded nibbles)
+    pub paths: Vec<Vec<Bytes>>,
+    pub bytes: u64,
+}
+
+#[derive(Debug)]
+pub(crate) struct TrieNodes {
+    pub id: u64,
+    pub nodes: Vec<Bytes>,
+}
+
 impl RLPxMessage for GetAccountRange {
     fn encode(&self, buf: &mut dyn BufMut) -> Result<(), RLPEncodeError> {
         let mut encoded_data = vec![];
@@ -245,6 +261,63 @@ impl RLPxMessage for ByteCodes {
     }
 }
 
+impl RLPxMessage for GetTrieNodes {
+    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.paths)
+            .encode_field(&self.bytes)
+            .finish();
+
+        let msg_data = snappy_compress(encoded_data)?;
+        buf.put_slice(&msg_data);
+        Ok(())
+    }
+
+    fn decode(msg_data: &[u8]) -> Result<Self, RLPDecodeError> {
+        let decompressed_data = snappy_decompress(msg_data)?;
+        let decoder = Decoder::new(&decompressed_data)?;
+        let (id, decoder) = decoder.decode_field("request-id")?;
+        let (root_hash, decoder) = decoder.decode_field("root_hash")?;
+        let (paths, decoder) = decoder.decode_field("paths")?;
+        let (bytes, decoder) = decoder.decode_field("bytes")?;
+        decoder.finish()?;
+
+        Ok(Self {
+            id,
+            root_hash,
+            paths,
+            bytes,
+        })
+    }
+}
+
+impl RLPxMessage for TrieNodes {
+    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.nodes)
+            .finish();
+
+        let msg_data = snappy_compress(encoded_data)?;
+        buf.put_slice(&msg_data);
+        Ok(())
+    }
+
+    fn decode(msg_data: &[u8]) -> Result<Self, RLPDecodeError> {
+        let decompressed_data = snappy_decompress(msg_data)?;
+        let decoder = Decoder::new(&decompressed_data)?;
+        let (id, decoder) = decoder.decode_field("request-id")?;
+        let (nodes, decoder) = decoder.decode_field("nodes")?;
+        decoder.finish()?;
+
+        Ok(Self { id, nodes })
+    }
+}
+
 // Intermediate structures
 
 #[derive(Debug)]

crates/networking/p2p/snap.rs

@@ -2,9 +2,12 @@ use bytes::Bytes;
 use ethereum_rust_rlp::encode::RLPEncode;
 use ethereum_rust_storage::{error::StoreError, Store};
 
-use crate::rlpx::snap::{
-    AccountRange, AccountRangeUnit, AccountStateSlim, ByteCodes, GetAccountRange, GetByteCodes,
-    GetStorageRanges, StorageRanges, StorageSlot,
+use crate::rlpx::{
+    error::RLPxError,
+    snap::{
+        AccountRange, AccountRangeUnit, AccountStateSlim, ByteCodes, GetAccountRange, GetByteCodes,
+        GetStorageRanges, GetTrieNodes, StorageRanges, StorageSlot, TrieNodes,
+    },
 };
 
 pub fn process_account_range_request(
@@ -119,6 +122,37 @@ pub fn process_byte_codes_request(
     })
 }
 
+pub fn process_trie_nodes_request(
+    request: GetTrieNodes,
+    store: Store,
+) -> Result<TrieNodes, RLPxError> {
+    let mut nodes = vec![];
+    let mut remaining_bytes = request.bytes;
+    for paths in request.paths {
+        if paths.is_empty() {
+            return Err(RLPxError::BadRequest(
+                "zero-item pathset requested".to_string(),
+            ));
+        }
+        let trie_nodes = store.get_trie_nodes(
+            request.root_hash,
+            paths.into_iter().map(|bytes| bytes.to_vec()).collect(),
+            remaining_bytes,
+        )?;
+        nodes.extend(trie_nodes.iter().map(|nodes| Bytes::copy_from_slice(nodes)));
+        remaining_bytes = remaining_bytes
+            .saturating_sub(trie_nodes.iter().fold(0, |acc, nodes| acc + nodes.len()) as u64);
+        if remaining_bytes == 0 {
+            break;
+        }
+    }
+
+    Ok(TrieNodes {
+        id: request.id,
+        nodes,
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use std::str::FromStr;

crates/storage/store/storage.rs

@@ -811,6 +811,57 @@ impl Store {
         Ok(Some(proof))
     }
 
+    /// Receives the root of the state trie and a list of paths where the first path will correspond to a path in the state trie
+    /// (aka a hashed account address) and the following paths will be paths in the account's storage trie (aka hashed storage keys)
+    /// If only one hash (account) is received, then the state trie node containing the account will be returned.
+    /// If more than one hash is received, then the storage trie nodes where each storage key is stored will be returned
+    /// For more information check out snap capability message [`GetTrieNodes`](https://github.com/ethereum/devp2p/blob/master/caps/snap.md#gettrienodes-0x06)
+    /// The paths can be either full paths (hash) or partial paths (compact-encoded nibbles), if a partial path is given for the account this method will not return storage nodes for it
+    pub fn get_trie_nodes(
+        &self,
+        state_root: H256,
+        paths: Vec<Vec<u8>>,
+        byte_limit: u64,
+    ) -> Result<Vec<Vec<u8>>, StoreError> {
+        let Some(account_path) = paths.first() else {
+            return Ok(vec![]);
+        };
+        let state_trie = self.engine.open_state_trie(state_root);
+        // State Trie Nodes Request
+        if paths.len() == 1 {
+            // Fetch state trie node
+            let node = state_trie.get_node(account_path)?;
+            return Ok(vec![node]);
+        }
+        // Storage Trie Nodes Request
+        let Some(account_state) = state_trie
+            .get(account_path)?
+            .map(|ref rlp| AccountState::decode(rlp))
+            .transpose()?
+        else {
+            return Ok(vec![]);
+        };
+        // We can't access the storage trie without the account's address hash
+        let Ok(hashed_address) = account_path.clone().try_into().map(H256) else {
+            return Ok(vec![]);
+        };
+        let storage_trie = self
+            .engine
+            .open_storage_trie(hashed_address, account_state.storage_root);
+        // Fetch storage trie nodes
+        let mut nodes = vec![];
+        let mut bytes_used = 0;
+        for path in paths.iter().skip(1) {
+            if bytes_used >= byte_limit {
+                break;
+            }
+            let node = storage_trie.get_node(path)?;
+            bytes_used += node.len() as u64;
+            nodes.push(node);
+        }
+        Ok(nodes)
+    }
+
     pub fn add_payload(&self, payload_id: u64, block: Block) -> Result<(), StoreError> {
         self.engine.add_payload(payload_id, block)
     }

crates/storage/trie/nibbles.rs

@@ -137,6 +137,11 @@ impl Nibbles {
         compact
     }
 
+    /// Encodes the nibbles in compact form
+    pub fn decode_compact(compact: &[u8]) -> Self {
+        Self::from_hex(compact_to_hex(compact))
+    }
+
     /// Returns true if the nibbles contain the leaf flag (16) at the end
     pub fn is_leaf(&self) -> bool {
         if self.is_empty() {
@@ -184,6 +189,33 @@ impl RLPDecode for Nibbles {
     }
 }
 
+// Code taken from https://github.com/ethereum/go-ethereum/blob/a1093d98eb3260f2abf340903c2d968b2b891c11/trie/encoding.go#L82
+fn compact_to_hex(compact: &[u8]) -> Vec<u8> {
+    if compact.is_empty() {
+        return vec![];
+    }
+    let mut base = keybytes_to_hex(compact);
+    // delete terminator flag
+    if base[0] < 2 {
+        base = base[..base.len() - 1].to_vec();
+    }
+    // apply odd flag
+    let chop = 2 - (base[0] & 1) as usize;
+    base[chop..].to_vec()
+}
+
+// Code taken from https://github.com/ethereum/go-ethereum/blob/a1093d98eb3260f2abf340903c2d968b2b891c11/trie/encoding.go#L96
+fn keybytes_to_hex(keybytes: &[u8]) -> Vec<u8> {
+    let l = keybytes.len() * 2 + 1;
+    let mut nibbles = vec![0; l];
+    for (i, b) in keybytes.iter().enumerate() {
+        nibbles[i * 2] = b / 16;
+        nibbles[i * 2 + 1] = b % 16;
+    }
+    nibbles[l - 1] = 16;
+    nibbles
+}
+
 #[cfg(test)]
 mod test {
     use super::*;

crates/storage/trie/node.rs

@@ -97,4 +97,12 @@ impl Node {
             Node::Leaf(n) => n.insert_self(state),
         }
     }
+
+    pub fn encode_raw(self) -> Vec<u8> {
+        match self {
+            Node::Branch(n) => n.encode_raw(),
+            Node::Extension(n) => n.encode_raw(),
+            Node::Leaf(n) => n.encode_raw(),
+        }
+    }
 }