Use new block requests protocol (#5760)

* Use new block requests protocol

* Tweak comment

substrate/client/network/src/protocol/block_requests.rs

@@ -27,7 +27,7 @@ use codec::{Encode, Decode};
 use crate::{
 	chain::Client,
 	config::ProtocolId,
-	protocol::{api, message::BlockAttributes}
+	protocol::{api, message::{self, BlockAttributes}}
 };
 use futures::{future::BoxFuture, prelude::*, stream::FuturesUnordered};
 use libp2p::{
@@ -36,13 +36,14 @@ use libp2p::{
 		Multiaddr,
 		PeerId,
 		connection::ConnectionId,
-		upgrade::{InboundUpgrade, ReadOneError, UpgradeInfo, Negotiated},
+		upgrade::{InboundUpgrade, OutboundUpgrade, ReadOneError, UpgradeInfo, Negotiated},
 		upgrade::{DeniedUpgrade, read_one, write_one}
 	},
 	swarm::{
 		NegotiatedSubstream,
 		NetworkBehaviour,
 		NetworkBehaviourAction,
+		NotifyHandler,
 		OneShotHandler,
 		OneShotHandlerConfig,
 		PollParameters,
@@ -53,8 +54,10 @@ use prost::Message;
 use sp_runtime::{generic::BlockId, traits::{Block, Header, One, Zero}};
 use std::{
 	cmp::min,
+	collections::VecDeque,
 	io,
 	iter,
+	marker::PhantomData,
 	sync::Arc,
 	time::Duration,
 	task::{Context, Poll}
@@ -64,11 +67,24 @@ use void::{Void, unreachable};
 // Type alias for convenience.
 pub type Error = Box<dyn std::error::Error + 'static>;
 
+/// Event generated by the block requests behaviour.
+#[derive(Debug)]
+pub enum Event<B: Block> {
+	/// A response to a block request has arrived.
+	Response {
+		peer: PeerId,
+		/// The original request passed to `send_request`.
+		original_request: message::BlockRequest<B>,
+		response: message::BlockResponse<B>,
+	},
+}
+
 /// Configuration options for `BlockRequests`.
 #[derive(Debug, Clone)]
 pub struct Config {
 	max_block_data_response: u32,
 	max_request_len: usize,
+	max_response_len: usize,
 	inactivity_timeout: Duration,
 	protocol: Bytes,
 }
@@ -78,11 +94,13 @@ impl Config {
 	///
 	/// - max. block data in response = 128
 	/// - max. request size = 1 MiB
+	/// - max. response size = 16 MiB
 	/// - inactivity timeout = 15s
 	pub fn new(id: &ProtocolId) -> Self {
 		let mut c = Config {
 			max_block_data_response: 128,
 			max_request_len: 1024 * 1024,
+			max_response_len: 16 * 1024 * 1024,
 			inactivity_timeout: Duration::from_secs(15),
 			protocol: Bytes::new(),
 		};
@@ -102,6 +120,12 @@ impl Config {
 		self
 	}
 
+	/// Limit the max. size of responses to our block requests.
+	pub fn set_max_response_len(&mut self, v: usize) -> &mut Self {
+		self.max_response_len = v;
+		self
+	}
+
 	/// Limit the max. duration the substream may remain inactive before closing it.
 	pub fn set_inactivity_timeout(&mut self, v: Duration) -> &mut Self {
 		self.inactivity_timeout = v;
@@ -127,6 +151,8 @@ pub struct BlockRequests<B: Block> {
 	chain: Arc<dyn Client<B>>,
 	/// Futures sending back the block request response.
 	outgoing: FuturesUnordered<BoxFuture<'static, ()>>,
+	/// Events to return as soon as possible from `poll`.
+	pending_events: VecDeque<NetworkBehaviourAction<OutboundProtocol<B>, Event<B>>>,
 }
 
 impl<B> BlockRequests<B>
@@ -138,9 +164,50 @@ where
 			config: cfg,
 			chain,
 			outgoing: FuturesUnordered::new(),
+			pending_events: VecDeque::new(),
 		}
 	}
 
+	/// Issue a new block request.
+	///
+	/// If the response doesn't arrive in time, or if the remote answers improperly, the target
+	/// will be disconnected.
+	pub fn send_request(&mut self, target: &PeerId, req: message::BlockRequest<B>) {
+		let protobuf_rq = api::v1::BlockRequest {
+			fields: u32::from_be_bytes([req.fields.bits(), 0, 0, 0]),
+			from_block: match req.from {
+				message::FromBlock::Hash(h) =>
+					Some(api::v1::block_request::FromBlock::Hash(h.encode())),
+				message::FromBlock::Number(n) =>
+					Some(api::v1::block_request::FromBlock::Number(n.encode())),
+			},
+			to_block: req.to.map(|h| h.encode()).unwrap_or_default(),
+			direction: match req.direction {
+				message::Direction::Ascending => api::v1::Direction::Ascending as i32,
+				message::Direction::Descending => api::v1::Direction::Descending as i32,
+			},
+			max_blocks: req.max.unwrap_or(0),
+		};
+
+		let mut buf = Vec::with_capacity(protobuf_rq.encoded_len());
+		if let Err(err) = protobuf_rq.encode(&mut buf) {
+			log::warn!("failed to encode block request {:?}: {:?}", protobuf_rq, err);
+			return;
+		}
+
+		log::trace!("enqueueing block request to {:?}: {:?}", target, protobuf_rq);
+		self.pending_events.push_back(NetworkBehaviourAction::NotifyHandler {
+			peer_id: target.clone(),
+			handler: NotifyHandler::Any,
+			event: OutboundProtocol {
+				request: buf,
+				original_request: req,
+				max_response_size: self.config.max_response_len,
+				protocol: self.config.protocol.clone(),
+			},
+		});
+	}
+
 	/// Callback, invoked when a new block request has been received from remote.
 	fn on_block_request
 		( &mut self
@@ -202,6 +269,12 @@ where
 			let number = header.number().clone();
 			let hash = header.hash();
 			let parent_hash = header.parent_hash().clone();
+			let justification = if get_justification {
+				self.chain.justification(&BlockId::Hash(hash))?
+			} else {
+				None
+			};
+			let is_empty_justification = justification.as_ref().map(|j| j.is_empty()).unwrap_or(false);
 
 			let block_data = api::v1::BlockData {
 				hash: hash.encode(),
@@ -221,11 +294,8 @@ where
 				},
 				receipt: Vec::new(),
 				message_queue: Vec::new(),
-				justification: if get_justification {
-					self.chain.justification(&BlockId::Hash(hash))?.unwrap_or(Vec::new())
-				} else {
-					Vec::new()
-				}
+				justification: justification.unwrap_or(Vec::new()),
+				is_empty_justification,
 			};
 
 			blocks.push(block_data);
@@ -251,13 +321,14 @@ impl<B> NetworkBehaviour for BlockRequests<B>
 where
 	B: Block
 {
-	type ProtocolsHandler = OneShotHandler<Protocol, DeniedUpgrade, Request<NegotiatedSubstream>>;
-	type OutEvent = Void;
+	type ProtocolsHandler = OneShotHandler<InboundProtocol<B>, OutboundProtocol<B>, NodeEvent<B, NegotiatedSubstream>>;
+	type OutEvent = Event<B>;
 
 	fn new_handler(&mut self) -> Self::ProtocolsHandler {
-		let p = Protocol {
+		let p = InboundProtocol {
 			max_request_len: self.config.max_request_len,
 			protocol: self.config.protocol.clone(),
+			marker: PhantomData,
 		};
 		let mut cfg = OneShotHandlerConfig::default();
 		cfg.inactive_timeout = self.config.inactivity_timeout;
@@ -278,44 +349,102 @@ where
 		&mut self,
 		peer: PeerId,
 		connection: ConnectionId,
-		Request(request, mut stream): Request<NegotiatedSubstream>
+		node_event: NodeEvent<B, NegotiatedSubstream>
 	) {
-		match self.on_block_request(&peer, &request) {
-			Ok(res) => {
-				log::trace!("enqueueing block response for peer {} with {} blocks", peer, res.blocks.len());
-				let mut data = Vec::with_capacity(res.encoded_len());
-				if let Err(e) = res.encode(&mut data) {
-					log::debug!("error encoding block response for peer {}: {}", peer, e)
-				} else {
-					let future = async move {
-						if let Err(e) = write_one(&mut stream, data).await {
-							log::debug!("error writing block response: {}", e)
+		match node_event {
+			NodeEvent::Request(request, mut stream) => {
+				match self.on_block_request(&peer, &request) {
+					Ok(res) => {
+						log::trace!("enqueueing block response for peer {} with {} blocks", peer, res.blocks.len());
+						let mut data = Vec::with_capacity(res.encoded_len());
+						if let Err(e) = res.encode(&mut data) {
+							log::debug!("error encoding block response for peer {}: {}", peer, e)
+						} else {
+							let future = async move {
+								if let Err(e) = write_one(&mut stream, data).await {
+									log::debug!("error writing block response: {}", e)
+								}
+							};
+							self.outgoing.push(future.boxed())
 						}
-					};
-					self.outgoing.push(future.boxed())
+					}
+					Err(e) => log::debug!("error handling block request from peer {}: {}", peer, e)
+				}
+			}
+			NodeEvent::Response(original_request, response) => {
+				log::trace!("received block response from peer {} with {} blocks", peer, response.blocks.len());
+				let blocks = response.blocks.into_iter().map(|block_data| {
+					Ok(message::BlockData::<B> {
+						hash: Decode::decode(&mut block_data.hash.as_ref())?,
+						header: if !block_data.header.is_empty() {
+							Some(Decode::decode(&mut block_data.header.as_ref())?)
+						} else {
+							None
+						},
+						body: if original_request.fields.contains(message::BlockAttributes::BODY) {
+							Some(block_data.body.iter().map(|body| {
+								Decode::decode(&mut body.as_ref())
+							}).collect::<Result<Vec<_>, _>>()?)
+						} else {
+							None
+						},
+						receipt: if !block_data.message_queue.is_empty() {
+							Some(block_data.receipt)
+						} else {
+							None
+						},
+						message_queue: if !block_data.message_queue.is_empty() {
+							Some(block_data.message_queue)
+						} else {
+							None
+						},
+						justification: if !block_data.justification.is_empty() {
+							Some(block_data.justification)
+						} else if block_data.is_empty_justification {
+							Some(Vec::new())
+						} else {
+							None
+						},
+					})
+				}).collect::<Result<Vec<_>, codec::Error>>();
+
+				match blocks {
+					Ok(blocks) => {
+						let id = original_request.id;
+						let ev = Event::Response {
+							peer,
+							original_request,
+							response: message::BlockResponse::<B> { id, blocks },
+						};
+						self.pending_events.push_back(NetworkBehaviourAction::GenerateEvent(ev));
+					}
+					Err(err) => {
+						log::debug!("failed to decode block response from peer {}: {}", peer, err);
+					}
 				}
 			}
-			Err(e) => log::debug!("error handling block request from peer {}: {}", peer, e)
 		}
 	}
 
-	fn poll(&mut self, cx: &mut Context, _: &mut impl PollParameters) -> Poll<NetworkBehaviourAction<DeniedUpgrade, Void>> {
+	fn poll(&mut self, cx: &mut Context, _: &mut impl PollParameters)
+		-> Poll<NetworkBehaviourAction<OutboundProtocol<B>, Event<B>>>
+	{
+		if let Some(ev) = self.pending_events.pop_front() {
+			return Poll::Ready(ev);
+		}
+
 		while let Poll::Ready(Some(_)) = self.outgoing.poll_next_unpin(cx) {}
 		Poll::Pending
 	}
 }
 
-/// The incoming block request.
-///
-/// Holds the protobuf value and the connection substream which made the
-/// request and over which to send the response.
+/// Output type of inbound and outbound substream upgrades.
 #[derive(Debug)]
-pub struct Request<T>(api::v1::BlockRequest, T);
-
-impl<T> From<Void> for Request<T> {
-	fn from(v: Void) -> Self {
-		unreachable(v)
-	}
+pub enum NodeEvent<B: Block, T> {
+	/// Incoming request from remote and substream to use for the response.
+	Request(api::v1::BlockRequest, T),
+	/// Incoming response from remote.
+	Response(message::BlockRequest<B>, api::v1::BlockResponse),
 }
 
 /// Substream upgrade protocol.
@@ -325,36 +454,39 @@ impl<T> From<Void> for Request<T> {
 /// will become visible via `inject_node_event` which then dispatches to the
 /// relevant callback to process the message and prepare a response.
 #[derive(Debug, Clone)]
-pub struct Protocol {
+pub struct InboundProtocol<B> {
 	/// The max. request length in bytes.
 	max_request_len: usize,
 	/// The protocol to use during upgrade negotiation.
 	protocol: Bytes,
+	/// Type of the block.
+	marker: PhantomData<B>,
 }
 
-impl UpgradeInfo for Protocol {
-    type Info = Bytes;
-    type InfoIter = iter::Once<Self::Info>;
+impl<B: Block> UpgradeInfo for InboundProtocol<B> {
+	type Info = Bytes;
+	type InfoIter = iter::Once<Self::Info>;
 
-    fn protocol_info(&self) -> Self::InfoIter {
-        iter::once(self.protocol.clone())
-    }
+	fn protocol_info(&self) -> Self::InfoIter {
+		iter::once(self.protocol.clone())
+	}
 }
 
-impl<T> InboundUpgrade<T> for Protocol
+impl<B, T> InboundUpgrade<T> for InboundProtocol<B>
 where
+	B: Block,
 	T: AsyncRead + AsyncWrite + Unpin + Send + 'static
 {
-    type Output = Request<T>;
-    type Error = ReadOneError;
-    type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
+	type Output = NodeEvent<B, T>;
+	type Error = ReadOneError;
+	type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
 
-    fn upgrade_inbound(self, mut s: T, _: Self::Info) -> Self::Future {
+	fn upgrade_inbound(self, mut s: T, _: Self::Info) -> Self::Future {
 		let future = async move {
 			let len = self.max_request_len;
 			let vec = read_one(&mut s, len).await?;
 			match api::v1::BlockRequest::decode(&vec[..]) {
-				Ok(r) => Ok(Request(r, s)),
+				Ok(r) => Ok(NodeEvent::Request(r, s)),
 				Err(e) => Err(ReadOneError::Io(io::Error::new(io::ErrorKind::Other, e)))
 			}
 		};
@@ -362,3 +494,49 @@ where
 	}
 }
 
+/// Substream upgrade protocol.
+///
+/// Sends a request to remote and awaits the response.
+#[derive(Debug, Clone)]
+pub struct OutboundProtocol<B: Block> {
+	/// The serialized protobuf request.
+	request: Vec<u8>,
+	/// The original request. Passed back through the API when the response comes back.
+	original_request: message::BlockRequest<B>,
+	/// The max. response length in bytes.
+	max_response_size: usize,
+	/// The protocol to use for upgrade negotiation.
+	protocol: Bytes,
+}
+
+impl<B: Block> UpgradeInfo for OutboundProtocol<B> {
+	type Info = Bytes;
+	type InfoIter = iter::Once<Self::Info>;
+
+	fn protocol_info(&self) -> Self::InfoIter {
+		iter::once(self.protocol.clone())
+	}
+}
+
+impl<B, T> OutboundUpgrade<T> for OutboundProtocol<B>
+where
+	B: Block,
+	T: AsyncRead + AsyncWrite + Unpin + Send + 'static
+{
+	type Output = NodeEvent<B, T>;
+	type Error = ReadOneError;
+	type Future = BoxFuture<'static, Result<Self::Output, Self::Error>>;
+
+	fn upgrade_outbound(self, mut s: T, _: Self::Info) -> Self::Future {
+		async move {
+			write_one(&mut s, &self.request).await?;
+			let vec = read_one(&mut s, self.max_response_size).await?;
+
+			api::v1::BlockResponse::decode(&vec[..])
+				.map(|r| NodeEvent::Response(self.original_request, r))
+				.map_err(|e| {
+					ReadOneError::Io(io::Error::new(io::ErrorKind::Other, e))
+				})
+		}.boxed()
+	}
+}