Allow fallback names for protocols (#8682)

* Allow fallback names for protocols

* Apply suggestions from code review

Co-authored-by: Roman Proskuryakov <humbug@deeptown.org>

* Fix some issues

* Fix compilation after merging master

Co-authored-by: Roman Proskuryakov <humbug@deeptown.org>

substrate/client/network/src/protocol/notifications/upgrade/notifications.rs

@@ -41,7 +41,7 @@ use futures::prelude::*;
 use asynchronous_codec::Framed;
 use libp2p::core::{UpgradeInfo, InboundUpgrade, OutboundUpgrade, upgrade};
 use log::error;
-use std::{borrow::Cow, convert::{Infallible, TryFrom as _}, io, iter, mem, pin::Pin, task::{Context, Poll}};
+use std::{borrow::Cow, convert::{Infallible, TryFrom as _}, io, mem, pin::Pin, task::{Context, Poll}, vec};
 use unsigned_varint::codec::UviBytes;
 
 /// Maximum allowed size of the two handshake messages, in bytes.
@@ -52,7 +52,8 @@ const MAX_HANDSHAKE_SIZE: usize = 1024;
 #[derive(Debug, Clone)]
 pub struct NotificationsIn {
 	/// Protocol name to use when negotiating the substream.
-	protocol_name: Cow<'static, str>,
+	/// The first one is the main name, while the other ones are fall backs.
+	protocol_names: Vec<Cow<'static, str>>,
 	/// Maximum allowed size for a single notification.
 	max_notification_size: u64,
 }
@@ -62,7 +63,8 @@ pub struct NotificationsIn {
 #[derive(Debug, Clone)]
 pub struct NotificationsOut {
 	/// Protocol name to use when negotiating the substream.
-	protocol_name: Cow<'static, str>,
+	/// The first one is the main name, while the other ones are fall backs.
+	protocol_names: Vec<Cow<'static, str>>,
 	/// Message to send when we start the handshake.
 	initial_message: Vec<u8>,
 	/// Maximum allowed size for a single notification.
@@ -106,51 +108,54 @@ pub struct NotificationsOutSubstream<TSubstream> {
 
 impl NotificationsIn {
 	/// Builds a new potential upgrade.
-	pub fn new(protocol_name: impl Into<Cow<'static, str>>, max_notification_size: u64) -> Self {
+	pub fn new(
+		main_protocol_name: impl Into<Cow<'static, str>>,
+		fallback_names: Vec<Cow<'static, str>>,
+		max_notification_size: u64
+	) -> Self {
+		let mut protocol_names = fallback_names;
+		protocol_names.insert(0, main_protocol_name.into());
+
 		NotificationsIn {
-			protocol_name: protocol_name.into(),
+			protocol_names,
 			max_notification_size,
 		}
 	}
 }
 
 impl UpgradeInfo for NotificationsIn {
-	type Info = Cow<'static, [u8]>;
-	type InfoIter = iter::Once<Self::Info>;
+	type Info = StringProtocolName;
+	type InfoIter = vec::IntoIter<Self::Info>;
 
 	fn protocol_info(&self) -> Self::InfoIter {
-		let bytes: Cow<'static, [u8]> = match &self.protocol_name {
-			Cow::Borrowed(s) => Cow::Borrowed(s.as_bytes()),
-			Cow::Owned(s) => Cow::Owned(s.as_bytes().to_vec())
-		};
-		iter::once(bytes)
+		self.protocol_names.iter().cloned().map(StringProtocolName).collect::<Vec<_>>().into_iter()
 	}
 }
 
 impl<TSubstream> InboundUpgrade<TSubstream> for NotificationsIn
 where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 {
-	type Output = (Vec<u8>, NotificationsInSubstream<TSubstream>);
+	type Output = NotificationsInOpen<TSubstream>;
 	type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
 	type Error = NotificationsHandshakeError;
 
 	fn upgrade_inbound(
 		self,
 		mut socket: TSubstream,
-		_: Self::Info,
+		negotiated_name: Self::Info,
 	) -> Self::Future {
 		Box::pin(async move {
-			let initial_message_len = unsigned_varint::aio::read_usize(&mut socket).await?;
-			if initial_message_len > MAX_HANDSHAKE_SIZE {
+			let handshake_len = unsigned_varint::aio::read_usize(&mut socket).await?;
+			if handshake_len > MAX_HANDSHAKE_SIZE {
 				return Err(NotificationsHandshakeError::TooLarge {
-					requested: initial_message_len,
+					requested: handshake_len,
 					max: MAX_HANDSHAKE_SIZE,
 				});
 			}
 
-			let mut initial_message = vec![0u8; initial_message_len];
-			if !initial_message.is_empty() {
-				socket.read_exact(&mut initial_message).await?;
+			let mut handshake = vec![0u8; handshake_len];
+			if !handshake.is_empty() {
+				socket.read_exact(&mut handshake).await?;
 			}
 
 			let mut codec = UviBytes::default();
@@ -161,11 +166,30 @@ where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 				handshake: NotificationsInSubstreamHandshake::NotSent,
 			};
 
-			Ok((initial_message, substream))
+			Ok(NotificationsInOpen {
+				handshake,
+				negotiated_fallback: if negotiated_name.0 == self.protocol_names[0] {
+					None
+				} else {
+					Some(negotiated_name.0)
+				},
+				substream,
+			})
 		})
 	}
 }
 
+/// Yielded by the [`NotificationsIn`] after a successfuly upgrade.
+pub struct NotificationsInOpen<TSubstream> {
+	/// Handshake sent by the remote.
+	pub handshake: Vec<u8>,
+	/// If the negotiated name is not the "main" protocol name but a fallback, contains the
+	/// name of the negotiated fallback.
+	pub negotiated_fallback: Option<Cow<'static, str>>,
+	/// Implementation of `Stream` that allows receives messages from the substream.
+	pub substream: NotificationsInSubstream<TSubstream>,
+}
+
 impl<TSubstream> NotificationsInSubstream<TSubstream>
 where TSubstream: AsyncRead + AsyncWrite + Unpin,
 {
@@ -296,7 +320,8 @@ where TSubstream: AsyncRead + AsyncWrite + Unpin,
 impl NotificationsOut {
 	/// Builds a new potential upgrade.
 	pub fn new(
-		protocol_name: impl Into<Cow<'static, str>>,
+		main_protocol_name: impl Into<Cow<'static, str>>,
+		fallback_names: Vec<Cow<'static, str>>,
 		initial_message: impl Into<Vec<u8>>,
 		max_notification_size: u64,
 	) -> Self {
@@ -305,38 +330,47 @@ impl NotificationsOut {
 			error!(target: "sub-libp2p", "Outbound networking handshake is above allowed protocol limit");
 		}
 
+		let mut protocol_names = fallback_names;
+		protocol_names.insert(0, main_protocol_name.into());
+
 		NotificationsOut {
-			protocol_name: protocol_name.into(),
+			protocol_names,
 			initial_message,
 			max_notification_size,
 		}
 	}
 }
 
+/// Implementation of the `ProtocolName` trait, where the protocol name is a string.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct StringProtocolName(Cow<'static, str>);
+
+impl upgrade::ProtocolName for StringProtocolName {
+	fn protocol_name(&self) -> &[u8] {
+		self.0.as_bytes()
+	}
+}
+
 impl UpgradeInfo for NotificationsOut {
-	type Info = Cow<'static, [u8]>;
-	type InfoIter = iter::Once<Self::Info>;
+	type Info = StringProtocolName;
+	type InfoIter = vec::IntoIter<Self::Info>;
 
 	fn protocol_info(&self) -> Self::InfoIter {
-		let bytes: Cow<'static, [u8]> = match &self.protocol_name {
-			Cow::Borrowed(s) => Cow::Borrowed(s.as_bytes()),
-			Cow::Owned(s) => Cow::Owned(s.as_bytes().to_vec())
-		};
-		iter::once(bytes)
+		self.protocol_names.iter().cloned().map(StringProtocolName).collect::<Vec<_>>().into_iter()
 	}
 }
 
 impl<TSubstream> OutboundUpgrade<TSubstream> for NotificationsOut
 where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 {
-	type Output = (Vec<u8>, NotificationsOutSubstream<TSubstream>);
+	type Output = NotificationsOutOpen<TSubstream>;
 	type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
 	type Error = NotificationsHandshakeError;
 
 	fn upgrade_outbound(
 		self,
 		mut socket: TSubstream,
-		_: Self::Info,
+		negotiated_name: Self::Info,
 	) -> Self::Future {
 		Box::pin(async move {
 			upgrade::write_with_len_prefix(&mut socket, &self.initial_message).await?;
@@ -358,13 +392,32 @@ where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 			let mut codec = UviBytes::default();
 			codec.set_max_len(usize::try_from(self.max_notification_size).unwrap_or(usize::max_value()));
 
-			Ok((handshake, NotificationsOutSubstream {
-				socket: Framed::new(socket, codec),
-			}))
+			Ok(NotificationsOutOpen {
+				handshake,
+				negotiated_fallback: if negotiated_name.0 == self.protocol_names[0] {
+					None
+				} else {
+					Some(negotiated_name.0)
+				},
+				substream: NotificationsOutSubstream {
+					socket: Framed::new(socket, codec),
+				}
+			})
 		})
 	}
 }
 
+/// Yielded by the [`NotificationsOut`] after a successfuly upgrade.
+pub struct NotificationsOutOpen<TSubstream> {
+	/// Handshake returned by the remote.
+	pub handshake: Vec<u8>,
+	/// If the negotiated name is not the "main" protocol name but a fallback, contains the
+	/// name of the negotiated fallback.
+	pub negotiated_fallback: Option<Cow<'static, str>>,
+	/// Implementation of `Sink` that allows sending messages on the substream.
+	pub substream: NotificationsOutSubstream<TSubstream>,
+}
+
 impl<TSubstream> Sink<Vec<u8>> for NotificationsOutSubstream<TSubstream>
 	where TSubstream: AsyncRead + AsyncWrite + Unpin,
 {
@@ -436,7 +489,7 @@ pub enum NotificationsOutError {
 
 #[cfg(test)]
 mod tests {
-	use super::{NotificationsIn, NotificationsOut};
+	use super::{NotificationsIn, NotificationsInOpen, NotificationsOut, NotificationsOutOpen};
 
 	use async_std::net::{TcpListener, TcpStream};
 	use futures::{prelude::*, channel::oneshot};
@@ -450,9 +503,9 @@ mod tests {
 
 		let client = async_std::task::spawn(async move {
 			let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
-			let (handshake, mut substream) = upgrade::apply_outbound(
+			let NotificationsOutOpen { handshake, mut substream, .. } = upgrade::apply_outbound(
 				socket,
-				NotificationsOut::new(PROTO_NAME, &b"initial message"[..], 1024 * 1024),
+				NotificationsOut::new(PROTO_NAME, Vec::new(), &b"initial message"[..], 1024 * 1024),
 				upgrade::Version::V1
 			).await.unwrap();
 
@@ -465,12 +518,12 @@ mod tests {
 			listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
 
 			let (socket, _) = listener.accept().await.unwrap();
-			let (initial_message, mut substream) = upgrade::apply_inbound(
+			let NotificationsInOpen { handshake, mut substream, .. } = upgrade::apply_inbound(
 				socket,
-				NotificationsIn::new(PROTO_NAME, 1024 * 1024)
+				NotificationsIn::new(PROTO_NAME, Vec::new(), 1024 * 1024)
 			).await.unwrap();
 
-			assert_eq!(initial_message, b"initial message");
+			assert_eq!(handshake, b"initial message");
 			substream.send_handshake(&b"hello world"[..]);
 
 			let msg = substream.next().await.unwrap().unwrap();
@@ -489,9 +542,9 @@ mod tests {
 
 		let client = async_std::task::spawn(async move {
 			let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
-			let (handshake, mut substream) = upgrade::apply_outbound(
+			let NotificationsOutOpen { handshake, mut substream, .. } = upgrade::apply_outbound(
 				socket,
-				NotificationsOut::new(PROTO_NAME, vec![], 1024 * 1024),
+				NotificationsOut::new(PROTO_NAME, Vec::new(), vec![], 1024 * 1024),
 				upgrade::Version::V1
 			).await.unwrap();
 
@@ -504,12 +557,12 @@ mod tests {
 			listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
 
 			let (socket, _) = listener.accept().await.unwrap();
-			let (initial_message, mut substream) = upgrade::apply_inbound(
+			let NotificationsInOpen { handshake, mut substream, .. } = upgrade::apply_inbound(
 				socket,
-				NotificationsIn::new(PROTO_NAME, 1024 * 1024)
+				NotificationsIn::new(PROTO_NAME, Vec::new(), 1024 * 1024)
 			).await.unwrap();
 
-			assert!(initial_message.is_empty());
+			assert!(handshake.is_empty());
 			substream.send_handshake(vec![]);
 
 			let msg = substream.next().await.unwrap().unwrap();
@@ -528,7 +581,7 @@ mod tests {
 			let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
 			let outcome = upgrade::apply_outbound(
 				socket,
-				NotificationsOut::new(PROTO_NAME, &b"hello"[..], 1024 * 1024),
+				NotificationsOut::new(PROTO_NAME, Vec::new(), &b"hello"[..], 1024 * 1024),
 				upgrade::Version::V1
 			).await;
 
@@ -543,12 +596,12 @@ mod tests {
 			listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
 
 			let (socket, _) = listener.accept().await.unwrap();
-			let (initial_msg, substream) = upgrade::apply_inbound(
+			let NotificationsInOpen { handshake, substream, .. } = upgrade::apply_inbound(
 				socket,
-				NotificationsIn::new(PROTO_NAME, 1024 * 1024)
+				NotificationsIn::new(PROTO_NAME, Vec::new(), 1024 * 1024)
 			).await.unwrap();
 
-			assert_eq!(initial_msg, b"hello");
+			assert_eq!(handshake, b"hello");
 
 			// We successfully upgrade to the protocol, but then close the substream.
 			drop(substream);
@@ -567,7 +620,7 @@ mod tests {
 			let ret = upgrade::apply_outbound(
 				socket,
 				// We check that an initial message that is too large gets refused.
-				NotificationsOut::new(PROTO_NAME, (0..32768).map(|_| 0).collect::<Vec<_>>(), 1024 * 1024),
+				NotificationsOut::new(PROTO_NAME, Vec::new(), (0..32768).map(|_| 0).collect::<Vec<_>>(), 1024 * 1024),
 				upgrade::Version::V1
 			).await;
 			assert!(ret.is_err());
@@ -580,7 +633,7 @@ mod tests {
 			let (socket, _) = listener.accept().await.unwrap();
 			let ret = upgrade::apply_inbound(
 				socket,
-				NotificationsIn::new(PROTO_NAME, 1024 * 1024)
+				NotificationsIn::new(PROTO_NAME, Vec::new(), 1024 * 1024)
 			).await;
 			assert!(ret.is_err());
 		});
@@ -597,7 +650,7 @@ mod tests {
 			let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
 			let ret = upgrade::apply_outbound(
 				socket,
-				NotificationsOut::new(PROTO_NAME, &b"initial message"[..], 1024 * 1024),
+				NotificationsOut::new(PROTO_NAME, Vec::new(), &b"initial message"[..], 1024 * 1024),
 				upgrade::Version::V1
 			).await;
 			assert!(ret.is_err());
@@ -608,11 +661,11 @@ mod tests {
 			listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
 
 			let (socket, _) = listener.accept().await.unwrap();
-			let (initial_message, mut substream) = upgrade::apply_inbound(
+			let NotificationsInOpen { handshake, mut substream, .. } = upgrade::apply_inbound(
 				socket,
-				NotificationsIn::new(PROTO_NAME, 1024 * 1024)
+				NotificationsIn::new(PROTO_NAME, Vec::new(), 1024 * 1024)
 			).await.unwrap();
-			assert_eq!(initial_message, b"initial message");
+			assert_eq!(handshake, b"initial message");
 
 			// We check that a handshake that is too large gets refused.
 			substream.send_handshake((0..32768).map(|_| 0).collect::<Vec<_>>());