Make network-libp2p's Service generic over the message (#1708)

* Make network-libp2p's Service generic over the message

* Apply suggestions from code review

Co-Authored-By: tomaka <pierre.krieger1708@gmail.com>

* Fix warning

substrate/core/network/src/message.rs

@@ -125,6 +125,8 @@ pub struct RemoteReadResponse {
 
 /// Generic types.
 pub mod generic {
+	use parity_codec::{Encode, Decode};
+	use network_libp2p::CustomMessage;
 	use runtime_primitives::Justification;
 	use parity_codec_derive::{Encode, Decode};
 	use crate::config::Roles;
@@ -197,6 +199,18 @@ pub mod generic {
 		ChainSpecific(Vec<u8>),
 	}
 
+	impl<Header, Hash, Number, Extrinsic> CustomMessage for Message<Header, Hash, Number, Extrinsic>
+		where Self: Decode + Encode
+	{
+		fn into_bytes(self) -> Vec<u8> {
+			self.encode()
+		}
+
+		fn from_bytes(bytes: &[u8]) -> Result<Self, ()> {
+			Decode::decode(&mut &bytes[..]).ok_or(())
+		}
+	}
+
 	/// Status sent on connection.
 	#[derive(Debug, PartialEq, Eq, Clone, Encode, Decode)]
 	pub struct Status<Hash, Number> {

substrate/core/network/src/on_demand.rs

@@ -16,7 +16,6 @@
 
 //! On-demand requests service.
 
-use parity_codec::Encode;
 use std::collections::{HashMap, VecDeque};
 use std::sync::Arc;
 use std::time::{Instant, Duration};
@@ -79,7 +78,7 @@ pub trait OnDemandService<Block: BlockT>: Send + Sync {
 pub struct OnDemand<B: BlockT> {
 	core: Mutex<OnDemandCore<B>>,
 	checker: Arc<FetchChecker<B>>,
-	network_sender: Mutex<Option<NetworkChan>>,
+	network_sender: Mutex<Option<NetworkChan<B>>>,
 }
 
 /// On-demand remote call response.
@@ -150,11 +149,11 @@ impl<B: BlockT> OnDemand<B> where
 	}
 
 	/// Sets weak reference to network service.
-	pub fn set_network_sender(&self, network_sender: NetworkChan) {
+	pub fn set_network_sender(&self, network_sender: NetworkChan<B>) {
 		self.network_sender.lock().replace(network_sender);
 	}
 
-	fn send(&self, msg: NetworkMsg) {
+	fn send(&self, msg: NetworkMsg<B>) {
 		let _ = self.network_sender
 			.lock()
 			.as_ref()
@@ -459,7 +458,7 @@ impl<B> OnDemandCore<B> where
 			let mut request = self.pending_requests.pop_front().expect("checked in loop condition; qed");
 			request.timestamp = Instant::now();
 			trace!(target: "sync", "Dispatching remote request {} to peer {}", request.id, peer);
-			on_demand.send(NetworkMsg::Outgoing(peer, request.message().encode()));
+			on_demand.send(NetworkMsg::Outgoing(peer, request.message()));
 			self.active_peers.insert(peer, request);
 		}
 
@@ -604,7 +603,7 @@ pub mod tests {
 		}
 	}
 
-	fn assert_disconnected_peer(network_port: NetworkPort, expected_severity: Severity) {
+	fn assert_disconnected_peer(network_port: NetworkPort<Block>, expected_severity: Severity) {
 		let mut disconnect_count = 0;
 		while let Ok(msg) = network_port.receiver().try_recv() {
 			match msg {

substrate/core/network/src/protocol.rs

@@ -14,7 +14,6 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate.  If not, see <http://www.gnu.org/licenses/>.
 
-use parity_codec::Encode;
 use crossbeam_channel::{self as channel, Receiver, Sender, select};
 use network_libp2p::{NodeIndex, Severity};
 use primitives::storage::StorageKey;
@@ -56,7 +55,7 @@ const LIGHT_MAXIMAL_BLOCKS_DIFFERENCE: u64 = 8192;
 
 // Lock must always be taken in order declared here.
 pub struct Protocol<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> {
-	network_chan: NetworkChan,
+	network_chan: NetworkChan<B>,
 	port: Receiver<ProtocolMsg<B, S>>,
 	config: ProtocolConfig,
 	on_demand: Option<Arc<OnDemandService<B>>>,
@@ -133,14 +132,14 @@ pub trait Context<B: BlockT> {
 
 /// Protocol context.
 pub(crate) struct ProtocolContext<'a, B: 'a + BlockT, H: 'a + ExHashT> {
-	network_chan: &'a NetworkChan,
+	network_chan: &'a NetworkChan<B>,
 	context_data: &'a mut ContextData<B, H>,
 }
 
 impl<'a, B: BlockT + 'a, H: 'a + ExHashT> ProtocolContext<'a, B, H> {
 	pub(crate) fn new(
 		context_data: &'a mut ContextData<B, H>,
-		network_chan: &'a NetworkChan,
+		network_chan: &'a NetworkChan<B>,
 	) -> Self {
 		ProtocolContext {
 			network_chan,
@@ -276,7 +275,7 @@ pub enum ProtocolMsg<B: BlockT, S: NetworkSpecialization<B>,> {
 impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 	/// Create a new instance.
 	pub fn new<I: 'static + ImportQueue<B>>(
-		network_chan: NetworkChan,
+		network_chan: NetworkChan<B>,
 		config: ProtocolConfig,
 		chain: Arc<Client<B>>,
 		import_queue: Arc<I>,
@@ -1106,7 +1105,7 @@ impl<B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Protocol<B, S, H> {
 
 fn send_message<B: BlockT, H: ExHashT>(
 	peers: &mut HashMap<NodeIndex, Peer<B, H>>,
-	network_chan: &NetworkChan,
+	network_chan: &NetworkChan<B>,
 	who: NodeIndex,
 	mut message: Message<B>,
 ) {
@@ -1124,7 +1123,7 @@ fn send_message<B: BlockT, H: ExHashT>(
 		}
 		_ => (),
 	}
-	network_chan.send(NetworkMsg::Outgoing(who, message.encode()));
+	network_chan.send(NetworkMsg::Outgoing(who, message));
 }
 
 /// Construct a simple protocol that is composed of several sub protocols.

substrate/core/network/src/service.rs

@@ -25,8 +25,8 @@ use network_libp2p::{start_service, parse_str_addr, Service as NetworkService, S
 use network_libp2p::{Protocol as Libp2pProtocol, RegisteredProtocol};
 use consensus::import_queue::{ImportQueue, Link};
 use crate::consensus_gossip::ConsensusGossip;
+use crate::message::Message;
 use crate::protocol::{self, Context, Protocol, ProtocolMsg, ProtocolStatus, PeerInfo};
-use parity_codec::Decode;
 use crate::config::Params;
 use crossbeam_channel::{self as channel, Receiver, Sender, TryRecvError};
 use crate::error::Error;
@@ -72,7 +72,7 @@ pub struct NetworkLink<B: BlockT, S: NetworkSpecialization<B>> {
 	/// The protocol sender
 	pub(crate) protocol_sender: Sender<ProtocolMsg<B, S>>,
 	/// The network sender
-	pub(crate) network_sender: NetworkChan,
+	pub(crate) network_sender: NetworkChan<B>,
 }
 
 impl<B: BlockT, S: NetworkSpecialization<B>> Link<B> for NetworkLink<B, S> {
@@ -108,7 +108,7 @@ impl<B: BlockT, S: NetworkSpecialization<B>> Link<B> for NetworkLink<B, S> {
 /// Substrate network service. Handles network IO and manages connectivity.
 pub struct Service<B: BlockT + 'static, S: NetworkSpecialization<B>> {
 	/// Network service
-	network: Arc<Mutex<NetworkService>>,
+	network: Arc<Mutex<NetworkService<Message<B>>>>,
 	/// Protocol sender
 	protocol_sender: Sender<ProtocolMsg<B, S>>,
 	/// Sender for messages to the background service task, and handle for the background thread.
@@ -123,7 +123,7 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>> Service<B, S> {
 		params: Params<B, S, H>,
 		protocol_id: ProtocolId,
 		import_queue: Arc<I>,
-	) -> Result<(Arc<Service<B, S>>, NetworkChan), Error> {
+	) -> Result<(Arc<Service<B, S>>, NetworkChan<B>), Error> {
 		let (network_chan, network_port) = network_channel(protocol_id);
 		let protocol_sender = Protocol::new(
 			network_chan.clone(),
@@ -139,7 +139,6 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>> Service<B, S> {
 		let (thread, network) = start_thread(
 			protocol_sender.clone(),
 			network_port,
-			network_chan.clone(),
 			params.network_config,
 			registered,
 		)?;
@@ -332,7 +331,7 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>> ManageNetwork for Service
 
 
 /// Create a NetworkPort/Chan pair.
-pub fn network_channel(protocol_id: ProtocolId) -> (NetworkChan, NetworkPort) {
+pub fn network_channel<B: BlockT + 'static>(protocol_id: ProtocolId) -> (NetworkChan<B>, NetworkPort<B>) {
 	let (network_sender, network_receiver) = channel::unbounded();
 	let task_notify = Arc::new(AtomicTask::new());
 	let network_port = NetworkPort::new(network_receiver, protocol_id, task_notify.clone());
@@ -343,14 +342,14 @@ pub fn network_channel(protocol_id: ProtocolId) -> (NetworkChan, NetworkPort) {
 
 /// A sender of NetworkMsg that notifies a task when a message has been sent.
 #[derive(Clone)]
-pub struct NetworkChan {
-	sender: Sender<NetworkMsg>,
+pub struct NetworkChan<B: BlockT + 'static> {
+	sender: Sender<NetworkMsg<B>>,
 	task_notify: Arc<AtomicTask>,
 }
 
-impl NetworkChan {
+impl<B: BlockT + 'static> NetworkChan<B> {
 	/// Create a new network chan.
-	pub fn new(sender: Sender<NetworkMsg>, task_notify: Arc<AtomicTask>) -> Self {
+	pub fn new(sender: Sender<NetworkMsg<B>>, task_notify: Arc<AtomicTask>) -> Self {
 		 NetworkChan {
 			sender,
 			task_notify,
@@ -358,13 +357,13 @@ impl NetworkChan {
 	}
 
 	/// Send a messaging, to be handled on a stream. Notify the task handling the stream.
-	pub fn send(&self, msg: NetworkMsg) {
+	pub fn send(&self, msg: NetworkMsg<B>) {
 		let _ = self.sender.send(msg);
 		self.task_notify.notify();
 	}
 }
 
-impl Drop for NetworkChan {
+impl<B: BlockT + 'static> Drop for NetworkChan<B> {
 	/// Notifying the task when a sender is dropped(when all are dropped, the stream is finished).
 	fn drop(&mut self) {
 		self.task_notify.notify();
@@ -373,15 +372,15 @@ impl Drop for NetworkChan {
 
 
 /// A receiver of NetworkMsg that makes the protocol-id available with each message.
-pub struct NetworkPort {
-	receiver: Receiver<NetworkMsg>,
+pub struct NetworkPort<B: BlockT + 'static> {
+	receiver: Receiver<NetworkMsg<B>>,
 	protocol_id: ProtocolId,
 	task_notify: Arc<AtomicTask>,
 }
 
-impl NetworkPort {
+impl<B: BlockT + 'static> NetworkPort<B> {
 	/// Create a new network port for a given protocol-id.
-	pub fn new(receiver: Receiver<NetworkMsg>, protocol_id: ProtocolId, task_notify: Arc<AtomicTask>) -> Self {
+	pub fn new(receiver: Receiver<NetworkMsg<B>>, protocol_id: ProtocolId, task_notify: Arc<AtomicTask>) -> Self {
 		Self {
 			receiver,
 			protocol_id,
@@ -391,7 +390,7 @@ impl NetworkPort {
 
 	/// Receive a message, if any is currently-enqueued.
 	/// Register the current tokio task for notification when a new message is available.
-	pub fn take_one_message(&self) -> Result<Option<(ProtocolId, NetworkMsg)>, ()> {
+	pub fn take_one_message(&self) -> Result<Option<(ProtocolId, NetworkMsg<B>)>, ()> {
 		self.task_notify.register();
 		match self.receiver.try_recv() {
 			Ok(msg) => Ok(Some((self.protocol_id.clone(), msg))),
@@ -402,18 +401,18 @@ impl NetworkPort {
 
 	/// Get a reference to the underlying crossbeam receiver.
 	#[cfg(any(test, feature = "test-helpers"))]
-	pub fn receiver(&self) -> &Receiver<NetworkMsg> {
+	pub fn receiver(&self) -> &Receiver<NetworkMsg<B>> {
 		&self.receiver
 	}
 }
 
 /// Messages to be handled by NetworkService.
 #[derive(Debug)]
-pub enum NetworkMsg {
+pub enum NetworkMsg<B: BlockT + 'static> {
 	/// Ask network to convert a list of nodes, to a list of peers.
 	PeerIds(Vec<NodeIndex>, Sender<Vec<(NodeIndex, Option<PeerId>)>>),
 	/// Send an outgoing custom message.
-	Outgoing(NodeIndex, Vec<u8>),
+	Outgoing(NodeIndex, Message<B>),
 	/// Report a peer.
 	ReportPeer(NodeIndex, Severity),
 	/// Get a peer id.
@@ -423,11 +422,10 @@ pub enum NetworkMsg {
 /// Starts the background thread that handles the networking.
 fn start_thread<B: BlockT + 'static, S: NetworkSpecialization<B>>(
 	protocol_sender: Sender<ProtocolMsg<B, S>>,
-	network_port: NetworkPort,
-	network_sender: NetworkChan,
+	network_port: NetworkPort<B>,
 	config: NetworkConfiguration,
-	registered: RegisteredProtocol,
-) -> Result<((oneshot::Sender<()>, thread::JoinHandle<()>), Arc<Mutex<NetworkService>>), Error> {
+	registered: RegisteredProtocol<Message<B>>,
+) -> Result<((oneshot::Sender<()>, thread::JoinHandle<()>), Arc<Mutex<NetworkService<Message<B>>>>), Error> {
 	let protocol_id = registered.id();
 
 	// Start the main service.
@@ -447,7 +445,7 @@ fn start_thread<B: BlockT + 'static, S: NetworkSpecialization<B>>(
 	let service_clone = service.clone();
 	let mut runtime = Runtime::new()?;
 	let thread = thread::Builder::new().name("network".to_string()).spawn(move || {
-		let fut = run_thread(protocol_sender, service_clone, network_sender, network_port, protocol_id)
+		let fut = run_thread(protocol_sender, service_clone, network_port, protocol_id)
 			.select(close_rx.then(|_| Ok(())))
 			.map(|(val, _)| val)
 			.map_err(|(err,_ )| err);
@@ -466,9 +464,8 @@ fn start_thread<B: BlockT + 'static, S: NetworkSpecialization<B>>(
 /// Runs the background thread that handles the networking.
 fn run_thread<B: BlockT + 'static, S: NetworkSpecialization<B>>(
 	protocol_sender: Sender<ProtocolMsg<B, S>>,
-	network_service: Arc<Mutex<NetworkService>>,
-	network_sender: NetworkChan,
-	network_port: NetworkPort,
+	network_service: Arc<Mutex<NetworkService<Message<B>>>>,
+	network_port: NetworkPort<B>,
 	protocol_id: ProtocolId,
 ) -> impl Future<Item = (), Error = io::Error> {
 
@@ -538,28 +535,15 @@ fn run_thread<B: BlockT + 'static, S: NetworkSpecialization<B>>(
 			NetworkServiceEvent::ClosedCustomProtocol { node_index, debug_info, .. } => {
 				let _ = protocol_sender.send(ProtocolMsg::PeerDisconnected(node_index, debug_info));
 			}
-			NetworkServiceEvent::CustomMessage { node_index, data, .. } => {
-				if let Some(m) = Decode::decode(&mut (&data as &[u8])) {
-					let _ = protocol_sender.send(ProtocolMsg::CustomMessage(node_index, m));
-					return Ok(())
-				}
-				let _ = network_sender.send(
-					NetworkMsg::ReportPeer(
-						node_index,
-						Severity::Bad("Peer sent us a packet with invalid format".to_string())
-					)
-				);
+			NetworkServiceEvent::CustomMessage { node_index, message, .. } => {
+				let _ = protocol_sender.send(ProtocolMsg::CustomMessage(node_index, message));
+				return Ok(())
 			}
 			NetworkServiceEvent::Clogged { node_index, messages, .. } => {
 				debug!(target: "sync", "{} clogging messages:", messages.len());
-				for msg_bytes in messages.iter().take(5) {
-					if let Some(msg) = Decode::decode(&mut (&msg_bytes as &[u8])) {
-						debug!(target: "sync", "{:?}", msg);
-						let _ = protocol_sender.send(ProtocolMsg::PeerClogged(node_index, Some(msg)));
-					} else {
-						debug!(target: "sync", "{:?}", msg_bytes);
-						let _ = protocol_sender.send(ProtocolMsg::PeerClogged(node_index, None));
-					}
+				for msg in messages.into_iter().take(5) {
+					debug!(target: "sync", "{:?}", msg);
+					let _ = protocol_sender.send(ProtocolMsg::PeerClogged(node_index, Some(msg)));
 				}
 			}
 		};

substrate/core/network/src/test/mod.rs

@@ -28,7 +28,6 @@ use std::time::Duration;
 use log::trace;
 use client;
 use client::block_builder::BlockBuilder;
-use parity_codec::{Decode, Encode};
 use crate::config::ProtocolConfig;
 use consensus::import_queue::{import_many_blocks, ImportQueue, ImportQueueStatus, IncomingBlock};
 use consensus::import_queue::{Link, SharedBlockImport, SharedJustificationImport, Verifier};
@@ -39,7 +38,9 @@ use crossbeam_channel::{self as channel, Sender, select};
 use futures::Future;
 use futures::sync::{mpsc, oneshot};
 use keyring::Keyring;
-use network_libp2p::{NodeIndex, ProtocolId, Severity};
+use crate::message::Message;
+use network_libp2p::{NodeIndex, ProtocolId};
+use parity_codec::Encode;
 use parking_lot::Mutex;
 use primitives::{H256, Ed25519AuthorityId};
 use crate::protocol::{Context, Protocol, ProtocolMsg, ProtocolStatus};
@@ -236,9 +237,9 @@ pub type PeersClient = client::Client<test_client::Backend, test_client::Executo
 pub struct Peer<V: 'static + Verifier<Block>, D> {
 	client: Arc<PeersClient>,
 	pub protocol_sender: Sender<ProtocolMsg<Block, DummySpecialization>>,
-	network_port: Mutex<NetworkPort>,
+	network_port: Mutex<NetworkPort<Block>>,
 	import_queue: Arc<SyncImportQueue<Block, V>>,
-	network_sender: NetworkChan,
+	network_sender: NetworkChan<Block>,
 	pub data: D,
 }
 
@@ -247,8 +248,8 @@ impl<V: 'static + Verifier<Block>, D> Peer<V, D> {
 		client: Arc<PeersClient>,
 		import_queue: Arc<SyncImportQueue<Block, V>>,
 		protocol_sender: Sender<ProtocolMsg<Block, DummySpecialization>>,
-		network_sender: NetworkChan,
-		network_port: NetworkPort,
+		network_sender: NetworkChan<Block>,
+		network_port: NetworkPort<Block>,
 		data: D,
 	) -> Self {
 		let network_port = Mutex::new(network_port);
@@ -304,24 +305,14 @@ impl<V: 'static + Verifier<Block>, D> Peer<V, D> {
 	}
 
 	/// Receive a message from another peer. Return a set of peers to disconnect.
-	fn receive_message(&self, from: NodeIndex, msg: Vec<u8>) {
-		match Decode::decode(&mut (&msg as &[u8])) {
-			Some(m) => {
-				let _ = self
-					.protocol_sender
-					.send(ProtocolMsg::CustomMessage(from, m));
-			}
-			None => {
-				let _ = self.network_sender.send(NetworkMsg::ReportPeer(
-					from,
-					Severity::Bad("Peer sent us a packet with invalid format".to_string()),
-				));
-			}
-		}
+	fn receive_message(&self, from: NodeIndex, msg: Message<Block>) {
+		let _ = self
+			.protocol_sender
+			.send(ProtocolMsg::CustomMessage(from, msg));
 	}
 
 	/// Produce the next pending message to send to another peer.
-	fn pending_message(&self) -> Option<NetworkMsg> {
+	fn pending_message(&self) -> Option<NetworkMsg<Block>> {
 		select! {
 			recv(self.network_port.lock().receiver()) -> msg => return msg.ok(),
 			// If there are no messages ready, give protocol a change to send one.
@@ -330,7 +321,7 @@ impl<V: 'static + Verifier<Block>, D> Peer<V, D> {
 	}
 
 	/// Produce the next pending message to send to another peer, without waiting.
-	fn pending_message_fast(&self) -> Option<NetworkMsg> {
+	fn pending_message_fast(&self) -> Option<NetworkMsg<Block>> {
 		self.network_port.lock().receiver().try_recv().ok()
 	}