1use crate::{
22 schema::v1::{KeyValueStateEntry, StateEntry, StateRequest, StateResponse},
23 LOG_TARGET,
24};
25
26use codec::{Decode, Encode};
27use futures::{channel::oneshot, stream::StreamExt};
28use log::{debug, trace};
29use prost::Message;
30use sc_network_types::PeerId;
31use schnellru::{ByLength, LruMap};
32
33use sc_client_api::{BlockBackend, ProofProvider};
34use sc_network::{
35 config::ProtocolId,
36 request_responses::{IncomingRequest, OutgoingResponse},
37 NetworkBackend, MAX_RESPONSE_SIZE,
38};
39use sp_runtime::traits::Block as BlockT;
40
41use std::{
42 hash::{Hash, Hasher},
43 sync::Arc,
44 time::{Duration, Instant},
45};
46
47const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; const MAX_NUMBER_OF_SAME_REQUESTS_PER_PEER: usize = 2;
49
50const SAME_REQUEST_WINDOW: Duration = Duration::from_secs(60);
53
54mod rep {
55 use sc_network::ReputationChange as Rep;
56
57 pub const SAME_REQUEST: Rep = Rep::new(-(1 << 12), "Same state request multiple times");
59}
60
61pub fn generate_protocol_config<
64 Hash: AsRef<[u8]>,
65 B: BlockT,
66 N: NetworkBackend<B, <B as BlockT>::Hash>,
67>(
68 protocol_id: &ProtocolId,
69 genesis_hash: Hash,
70 fork_id: Option<&str>,
71 inbound_queue: async_channel::Sender<IncomingRequest>,
72) -> N::RequestResponseProtocolConfig {
73 N::request_response_config(
74 generate_protocol_name(genesis_hash, fork_id).into(),
75 std::iter::once(generate_legacy_protocol_name(protocol_id).into()).collect(),
76 1024 * 1024,
77 MAX_RESPONSE_SIZE,
78 Duration::from_secs(40),
79 Some(inbound_queue),
80 )
81}
82
83fn generate_protocol_name<Hash: AsRef<[u8]>>(genesis_hash: Hash, fork_id: Option<&str>) -> String {
85 let genesis_hash = genesis_hash.as_ref();
86 if let Some(fork_id) = fork_id {
87 format!("/{}/{}/state/2", array_bytes::bytes2hex("", genesis_hash), fork_id)
88 } else {
89 format!("/{}/state/2", array_bytes::bytes2hex("", genesis_hash))
90 }
91}
92
93fn generate_legacy_protocol_name(protocol_id: &ProtocolId) -> String {
95 format!("/{}/state/2", protocol_id.as_ref())
96}
97
98#[derive(Eq, PartialEq, Clone)]
100struct SeenRequestsKey<B: BlockT> {
101 peer: PeerId,
102 block: B::Hash,
103 start: Vec<Vec<u8>>,
104}
105
106#[allow(clippy::derived_hash_with_manual_eq)]
107impl<B: BlockT> Hash for SeenRequestsKey<B> {
108 fn hash<H: Hasher>(&self, state: &mut H) {
109 self.peer.hash(state);
110 self.block.hash(state);
111 self.start.hash(state);
112 }
113}
114
115enum SeenRequestsValue {
117 First,
119 Fulfilled { requests: usize, since: Instant },
121}
122
123pub struct StateRequestHandler<B: BlockT, Client> {
125 client: Arc<Client>,
126 request_receiver: async_channel::Receiver<IncomingRequest>,
127 seen_requests: LruMap<SeenRequestsKey<B>, SeenRequestsValue>,
131}
132
133impl<B, Client> StateRequestHandler<B, Client>
134where
135 B: BlockT,
136 Client: BlockBackend<B> + ProofProvider<B> + Send + Sync + 'static,
137{
138 pub fn new<N: NetworkBackend<B, <B as BlockT>::Hash>>(
140 protocol_id: &ProtocolId,
141 fork_id: Option<&str>,
142 client: Arc<Client>,
143 num_peer_hint: usize,
144 ) -> (Self, N::RequestResponseProtocolConfig) {
145 let capacity = std::cmp::max(num_peer_hint, 1);
148 let (tx, request_receiver) = async_channel::bounded(capacity);
149
150 let protocol_config = generate_protocol_config::<_, B, N>(
151 protocol_id,
152 client
153 .block_hash(0u32.into())
154 .ok()
155 .flatten()
156 .expect("Genesis block exists; qed"),
157 fork_id,
158 tx,
159 );
160
161 let capacity = ByLength::new(num_peer_hint.max(1) as u32 * 2);
162 let seen_requests = LruMap::new(capacity);
163
164 (Self { client, request_receiver, seen_requests }, protocol_config)
165 }
166
167 pub async fn run(mut self) {
169 while let Some(request) = self.request_receiver.next().await {
170 let IncomingRequest { peer, payload, pending_response } = request;
171
172 match self.handle_request(payload, pending_response, &peer) {
173 Ok(()) => debug!(target: LOG_TARGET, "Handled block request from {}.", peer),
174 Err(e) => debug!(
175 target: LOG_TARGET,
176 "Failed to handle state request from {}: {}", peer, e,
177 ),
178 }
179 }
180 }
181
182 fn handle_request(
183 &mut self,
184 payload: Vec<u8>,
185 pending_response: oneshot::Sender<OutgoingResponse>,
186 peer: &PeerId,
187 ) -> Result<(), HandleRequestError> {
188 let request = StateRequest::decode(&payload[..])?;
189 let block: B::Hash = Decode::decode(&mut request.block.as_ref())?;
190
191 let key = SeenRequestsKey { peer: *peer, block, start: request.start.clone() };
192
193 let mut reputation_changes = Vec::new();
194
195 match self.seen_requests.get(&key) {
196 Some(SeenRequestsValue::First) => {},
197 Some(SeenRequestsValue::Fulfilled { requests, since })
198 if since.elapsed() <= SAME_REQUEST_WINDOW =>
199 {
200 *requests = requests.saturating_add(1);
201
202 if *requests > MAX_NUMBER_OF_SAME_REQUESTS_PER_PEER {
203 reputation_changes.push(rep::SAME_REQUEST);
204 }
205 },
206 Some(value @ SeenRequestsValue::Fulfilled { .. }) => {
207 *value = SeenRequestsValue::First;
208 },
209 None => {
210 self.seen_requests.insert(key.clone(), SeenRequestsValue::First);
211 },
212 }
213
214 trace!(
215 target: LOG_TARGET,
216 "Handling state request from {}: Block {:?}, Starting at {:x?}, no_proof={}",
217 peer,
218 request.block,
219 &request.start,
220 request.no_proof,
221 );
222
223 let result = if reputation_changes.is_empty() {
224 let mut response = StateResponse::default();
225
226 if !request.no_proof {
227 let (proof, _count) = self.client.read_proof_collection(
228 block,
229 request.start.as_slice(),
230 MAX_RESPONSE_BYTES,
231 )?;
232 response.proof = proof.encode();
233 } else {
234 let entries = self.client.storage_collection(
235 block,
236 request.start.as_slice(),
237 MAX_RESPONSE_BYTES,
238 )?;
239 response.entries = entries
240 .into_iter()
241 .map(|(state, complete)| KeyValueStateEntry {
242 state_root: state.state_root,
243 entries: state
244 .key_values
245 .into_iter()
246 .map(|(key, value)| StateEntry { key, value })
247 .collect(),
248 complete,
249 })
250 .collect();
251 }
252
253 trace!(
254 target: LOG_TARGET,
255 "StateResponse contains {} keys, {}, proof nodes, from {:?} to {:?}",
256 response.entries.len(),
257 response.proof.len(),
258 response.entries.get(0).and_then(|top| top
259 .entries
260 .first()
261 .map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key))),
262 response.entries.get(0).and_then(|top| top
263 .entries
264 .last()
265 .map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key))),
266 );
267 if let Some(value) = self.seen_requests.get(&key) {
268 if let SeenRequestsValue::First = value {
269 *value = SeenRequestsValue::Fulfilled { requests: 1, since: Instant::now() };
270 }
271 }
272
273 let mut data = Vec::with_capacity(response.encoded_len());
274 response.encode(&mut data)?;
275 Ok(data)
276 } else {
277 Err(())
278 };
279
280 pending_response
281 .send(OutgoingResponse { result, reputation_changes, sent_feedback: None })
282 .map_err(|_| HandleRequestError::SendResponse)
283 }
284}
285
286#[derive(Debug, thiserror::Error)]
287enum HandleRequestError {
288 #[error("Failed to decode request: {0}.")]
289 DecodeProto(#[from] prost::DecodeError),
290
291 #[error("Failed to encode response: {0}.")]
292 EncodeProto(#[from] prost::EncodeError),
293
294 #[error("Failed to decode block hash: {0}.")]
295 InvalidHash(#[from] codec::Error),
296
297 #[error(transparent)]
298 Client(#[from] sp_blockchain::Error),
299
300 #[error("Failed to send response.")]
301 SendResponse,
302}
303
304#[cfg(test)]
305mod tests {
306 use super::*;
307 use substrate_test_runtime_client::{
308 runtime::Block, DefaultTestClientBuilderExt, TestClient, TestClientBuilder,
309 TestClientBuilderExt,
310 };
311
312 fn test_handler() -> StateRequestHandler<Block, TestClient> {
313 let client = Arc::new(TestClientBuilder::new().build());
314 let (_tx, request_receiver) = async_channel::bounded(1);
315 StateRequestHandler {
316 client,
317 request_receiver,
318 seen_requests: LruMap::new(ByLength::new(16)),
319 }
320 }
321
322 fn send_request(
323 handler: &mut StateRequestHandler<Block, TestClient>,
324 peer: &PeerId,
325 no_proof: bool,
326 ) -> OutgoingResponse {
327 let request = StateRequest {
328 block: Encode::encode(&handler.client.chain_info().genesis_hash),
329 start: Vec::new(),
330 no_proof,
331 };
332 let (tx, mut rx) = oneshot::channel();
333 handler.handle_request(request.encode_to_vec(), tx, peer).unwrap();
334 rx.try_recv().unwrap().unwrap()
335 }
336
337 fn check_same_request_limit_resets_after_window(no_proof: bool) {
338 let mut handler = test_handler();
339 let peer = PeerId::random();
340
341 for _ in 0..MAX_NUMBER_OF_SAME_REQUESTS_PER_PEER {
342 let response = send_request(&mut handler, &peer, no_proof);
343 assert!(response.result.is_ok());
344 assert!(response.reputation_changes.is_empty());
345 }
346
347 let response = send_request(&mut handler, &peer, no_proof);
348 assert!(response.result.is_err());
349 assert_eq!(response.reputation_changes, vec![rep::SAME_REQUEST]);
350 assert!(rep::SAME_REQUEST.value > i32::MIN);
351
352 let key = SeenRequestsKey::<Block> {
354 peer,
355 block: handler.client.chain_info().genesis_hash,
356 start: Vec::new(),
357 };
358 match handler.seen_requests.get(&key) {
359 Some(SeenRequestsValue::Fulfilled { since, .. }) => {
360 *since = Instant::now() - SAME_REQUEST_WINDOW - Duration::from_secs(1)
361 },
362 _ => panic!("entry must be in the fulfilled state"),
363 }
364
365 for _ in 0..MAX_NUMBER_OF_SAME_REQUESTS_PER_PEER {
367 let response = send_request(&mut handler, &peer, no_proof);
368 assert!(response.result.is_ok());
369 assert!(response.reputation_changes.is_empty());
370 }
371 let response = send_request(&mut handler, &peer, no_proof);
372 assert!(response.result.is_err());
373 assert_eq!(response.reputation_changes, vec![rep::SAME_REQUEST]);
374 }
375
376 #[test]
377 fn same_request_limit_resets_after_window_with_proof() {
378 check_same_request_limit_resets_after_window(false);
379 }
380
381 #[test]
382 fn same_request_limit_resets_after_window_without_proof() {
383 check_same_request_limit_resets_after_window(true);
384 }
385}