feat(l1): support snap capability message GetTrieNodes

.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|StorageRanges|ByteCodes"
+            run_command: make run-hive-on-latest SIMULATION=devp2p TEST_PATTERN="/AccountRange|StorageRanges|ByteCodes|TrieNodes"
           - 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

@@ -2,6 +2,7 @@ use crate::{
     rlpx::{eth::backend, handshake::encode_ack_message, message::Message, p2p, utils::id2pubkey},
     snap::{
         process_account_range_request, process_byte_codes_request, process_storage_ranges_request,
+        process_trie_nodes_request,
     },
     MAX_DISC_PACKET_SIZE,
 };
@@ -161,6 +162,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
                         message => return Err(RLPxError::UnexpectedMessage(message)),
                     };

crates/networking/p2p/rlpx/error.rs

@@ -13,4 +13,6 @@ pub(crate) enum RLPxError {
     UnexpectedMessage(Message),
     #[error(transparent)]
     Store(#[from] StoreError),
+    #[error("Bad Request: {0}")]
+    BadRequest(String),
 }

crates/networking/p2p/rlpx/message.rs

@@ -5,7 +5,8 @@ use std::fmt::Display;
 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;
@@ -29,6 +30,8 @@ pub(crate) enum Message {
     StorageRanges(StorageRanges),
     GetByteCodes(GetByteCodes),
     ByteCodes(ByteCodes),
+    GetTrieNodes(GetTrieNodes),
+    TrieNodes(TrieNodes),
 }
 
 impl Message {
@@ -47,6 +50,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),
         }
     }
@@ -82,6 +87,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)
+            }
         }
     }
 }
@@ -100,6 +113,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

@@ -63,6 +63,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![];
@@ -262,6 +278,69 @@ 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_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("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_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 (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 accoun'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/encoding.rs

@@ -0,0 +1,47 @@
+/// Converts a slice of compact-encoded nibbles into a byte slice
+/// If the nibble slice has odd-length (aka the last byte will be a half byte) returns true else false
+pub fn compact_nibbles_to_bytes(compact: &[u8]) -> (Vec<u8>, bool) {
+    // Convert compact nibbles to nibbles
+    let nibbles = compact_to_hex(compact);
+    // Convert nibbles to bytes, accouning for odd number of bytes
+    let mut last_is_half = false;
+    let bytes = nibbles
+        .chunks(2)
+        .map(|chunk| match chunk.len() {
+            1 => {
+                last_is_half = true;
+                chunk[0] << 4
+            }
+            // 2
+            _ => chunk[0] << 4 | chunk[1],
+        })
+        .collect::<Vec<_>>();
+    (bytes, last_is_half)
+}
+
+// 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
+}

crates/storage/trie/nibble.rs

@@ -322,7 +322,7 @@ impl NibbleVec {
         }
     }
 
-    #[cfg(test)]
+    //#[cfg(test)]
     /// Create a NibbleVec from an iterator of nibbles
     pub fn from_nibbles(
         data_iter: impl Iterator<Item = Nibble>,

crates/storage/trie/node.rs

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