reth_network_p2p/test_utils/
full_block.rs

1use crate::{
2    bodies::client::BodiesClient,
3    download::DownloadClient,
4    error::PeerRequestResult,
5    headers::client::{HeadersClient, HeadersRequest},
6    priority::Priority,
7    BlockClient,
8};
9use alloy_consensus::Header;
10use alloy_eips::{BlockHashOrNumber, BlockNumHash};
11use alloy_primitives::B256;
12use parking_lot::Mutex;
13use reth_eth_wire_types::HeadersDirection;
14use reth_ethereum_primitives::{Block, BlockBody};
15use reth_network_peers::{PeerId, WithPeerId};
16use reth_primitives_traits::{SealedBlock, SealedHeader};
17use std::{collections::HashMap, sync::Arc};
18
19/// A headers+bodies client that stores the headers and bodies in memory, with an artificial soft
20/// bodies response limit that is set to 20 by default.
21///
22/// This full block client can be [Clone]d and shared between multiple tasks.
23#[derive(Clone, Debug)]
24pub struct TestFullBlockClient {
25    headers: Arc<Mutex<HashMap<B256, Header>>>,
26    bodies: Arc<Mutex<HashMap<B256, BlockBody>>>,
27    // soft response limit, max number of bodies to respond with
28    soft_limit: usize,
29}
30
31impl Default for TestFullBlockClient {
32    fn default() -> Self {
33        Self {
34            headers: Arc::new(Mutex::new(HashMap::default())),
35            bodies: Arc::new(Mutex::new(HashMap::default())),
36            soft_limit: 20,
37        }
38    }
39}
40
41impl TestFullBlockClient {
42    /// Insert a header and body into the client maps.
43    pub fn insert(&self, header: SealedHeader, body: BlockBody) {
44        let hash = header.hash();
45        self.headers.lock().insert(hash, header.unseal());
46        self.bodies.lock().insert(hash, body);
47    }
48
49    /// Set the soft response limit.
50    pub const fn set_soft_limit(&mut self, limit: usize) {
51        self.soft_limit = limit;
52    }
53
54    /// Get the block with the highest block number.
55    pub fn highest_block(&self) -> Option<SealedBlock<Block>> {
56        self.headers.lock().iter().max_by_key(|(_, header)| header.number).and_then(
57            |(hash, header)| {
58                self.bodies.lock().get(hash).map(|body| {
59                    SealedBlock::from_parts_unchecked(header.clone(), body.clone(), *hash)
60                })
61            },
62        )
63    }
64}
65
66impl DownloadClient for TestFullBlockClient {
67    /// Reports a bad message from a specific peer.
68    fn report_bad_message(&self, _peer_id: PeerId) {}
69
70    /// Retrieves the number of connected peers.
71    ///
72    /// Returns the number of connected peers in the test scenario (1).
73    fn num_connected_peers(&self) -> usize {
74        1
75    }
76}
77
78/// Implements the `HeadersClient` trait for the `TestFullBlockClient` struct.
79impl HeadersClient for TestFullBlockClient {
80    type Header = Header;
81    /// Specifies the associated output type.
82    type Output = futures::future::Ready<PeerRequestResult<Vec<Header>>>;
83
84    /// Retrieves headers with a given priority level.
85    ///
86    /// # Arguments
87    ///
88    /// * `request` - A `HeadersRequest` indicating the headers to retrieve.
89    /// * `_priority` - A `Priority` level for the request.
90    ///
91    /// # Returns
92    ///
93    /// A `Ready` future containing a `PeerRequestResult` with a vector of retrieved headers.
94    fn get_headers_with_priority(
95        &self,
96        request: HeadersRequest,
97        _priority: Priority,
98    ) -> Self::Output {
99        let headers = self.headers.lock();
100
101        // Initializes the block hash or number.
102        let mut block: BlockHashOrNumber = match request.start {
103            BlockHashOrNumber::Hash(hash) => headers.get(&hash).cloned(),
104            BlockHashOrNumber::Number(num) => headers.values().find(|h| h.number == num).cloned(),
105        }
106        .map(|h| h.number.into())
107        .unwrap();
108
109        // Retrieves headers based on the provided limit and request direction.
110        let resp = (0..request.limit)
111            .filter_map(|_| {
112                headers.iter().find_map(|(hash, header)| {
113                    // Checks if the header matches the specified block or number.
114                    BlockNumHash::new(header.number, *hash).matches_block_or_num(&block).then(
115                        || {
116                            match request.direction {
117                                HeadersDirection::Falling => block = header.parent_hash.into(),
118                                HeadersDirection::Rising => block = (header.number + 1).into(),
119                            }
120                            header.clone()
121                        },
122                    )
123                })
124            })
125            .collect::<Vec<_>>();
126
127        // Returns a future containing the retrieved headers with a random peer ID.
128        futures::future::ready(Ok(WithPeerId::new(PeerId::random(), resp)))
129    }
130}
131
132/// Implements the `BodiesClient` trait for the `TestFullBlockClient` struct.
133impl BodiesClient for TestFullBlockClient {
134    type Body = BlockBody;
135    /// Defines the output type of the function.
136    type Output = futures::future::Ready<PeerRequestResult<Vec<BlockBody>>>;
137
138    /// Retrieves block bodies corresponding to provided hashes with a given priority.
139    ///
140    /// # Arguments
141    ///
142    /// * `hashes` - A vector of block hashes to retrieve bodies for.
143    /// * `_priority` - Priority level for block body retrieval (unused in this implementation).
144    ///
145    /// # Returns
146    ///
147    /// A future containing the result of the block body retrieval operation.
148    fn get_block_bodies_with_priority(
149        &self,
150        hashes: Vec<B256>,
151        _priority: Priority,
152    ) -> Self::Output {
153        // Acquire a lock on the bodies.
154        let bodies = self.bodies.lock();
155
156        // Create a future that immediately returns the result of the block body retrieval
157        // operation.
158        futures::future::ready(Ok(WithPeerId::new(
159            PeerId::random(),
160            hashes
161                .iter()
162                .filter_map(|hash| bodies.get(hash).cloned())
163                .take(self.soft_limit)
164                .collect(),
165        )))
166    }
167}
168
169impl BlockClient for TestFullBlockClient {
170    type Block = reth_ethereum_primitives::Block;
171}
reth_network_p2p/test_utils/full_block.rs

reth_network_p2p/test_utils/
full_block.rs
