Ensure that handshake is sent back even in case of back-pressure

client/network/src/protocol/generic_proto/handler/group.rs

@@ -706,6 +706,16 @@ impl ProtocolsHandler for NotifsHandler {
 					}
 				}
 			}
+
+		} else {
+			// If `self.pending_legacy_handshake` is `None`, we don't want to accept events from
+			// the ingoing notifications substreams, but still want handshakes to go forward.
+			for (handler, _) in &mut self.in_handlers {
+				match handler.poll_process(cx) {
+					Poll::Pending => {},
+					Poll::Ready(v) => match v {}
+				}
+			}
 		}
 
 		for (handler_num, (handler, _)) in self.out_handlers.iter_mut().enumerate() {

client/network/src/protocol/generic_proto/handler/notif_in.rs

@@ -37,7 +37,7 @@ use libp2p::swarm::{
 	NegotiatedSubstream,
 };
 use log::{error, warn};
-use std::{borrow::Cow, collections::VecDeque, fmt, pin::Pin, task::{Context, Poll}};
+use std::{borrow::Cow, collections::VecDeque, convert::Infallible, fmt, pin::Pin, task::{Context, Poll}};
 
 /// Implements the `IntoProtocolsHandler` trait of libp2p.
 ///
@@ -139,6 +139,24 @@ impl NotifsInHandler {
 	pub fn protocol_name(&self) -> &[u8] {
 		self.in_protocol.protocol_name()
 	}
+
+	/// Equivalent to the `poll` method of `ProtocolsHandler`, except that it only drives the
+	/// handshake process.
+	///
+	/// Use this method in situations where it is not desirable to receive events but still
+	/// necessary to drive any potential incoming handshake or request.
+	pub fn poll_process(&mut self, cx: &mut Context) -> Poll<Infallible> {
+		match self.substream.as_mut().map(|s| NotificationsInSubstream::poll_process(Pin::new(s), cx)) {
+			None | Some(Poll::Pending) => {},
+			Some(Poll::Ready(Ok(v))) => match v {},
+			Some(Poll::Ready(Err(_))) => {
+				self.substream = None;
+				self.events_queue.push_back(ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::Closed));
+			},
+		}
+
+		Poll::Pending
+	}
 }
 
 impl ProtocolsHandler for NotifsInHandler {

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

@@ -39,7 +39,7 @@ use futures::prelude::*;
 use futures_codec::Framed;
 use libp2p::core::{UpgradeInfo, InboundUpgrade, OutboundUpgrade, upgrade};
 use log::error;
-use std::{borrow::Cow, io, iter, mem, pin::Pin, task::{Context, Poll}};
+use std::{borrow::Cow, convert::Infallible, io, iter, mem, pin::Pin, task::{Context, Poll}};
 use unsigned_varint::codec::UviBytes;
 
 /// Maximum allowed size of the two handshake messages, in bytes.
@@ -158,7 +158,7 @@ where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
 }
 
 impl<TSubstream> NotificationsInSubstream<TSubstream>
-where TSubstream: AsyncRead + AsyncWrite,
+where TSubstream: AsyncRead + AsyncWrite + Unpin,
 {
 	/// Sends the handshake in order to inform the remote that we accept the substream.
 	pub fn send_handshake(&mut self, message: impl Into<Vec<u8>>) {
@@ -169,6 +169,48 @@ where TSubstream: AsyncRead + AsyncWrite,
 
 		self.handshake = NotificationsInSubstreamHandshake::PendingSend(message.into());
 	}
+
+	/// Equivalent to `Stream::poll_next`, except that it only drives the handshake and is
+	/// guaranteed to not generate any notification.
+	pub fn poll_process(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<Infallible, io::Error>> {
+		let mut this = self.project();
+
+		loop {
+			match mem::replace(this.handshake, NotificationsInSubstreamHandshake::Sent) {
+				NotificationsInSubstreamHandshake::PendingSend(msg) =>
+					match Sink::poll_ready(this.socket.as_mut(), cx) {
+						Poll::Ready(_) => {
+							*this.handshake = NotificationsInSubstreamHandshake::Flush;
+							match Sink::start_send(this.socket.as_mut(), io::Cursor::new(msg)) {
+								Ok(()) => {},
+								Err(err) => return Poll::Ready(Err(err)),
+							}
+						},
+						Poll::Pending => {
+							*this.handshake = NotificationsInSubstreamHandshake::PendingSend(msg);
+							return Poll::Pending
+						}
+					},
+				NotificationsInSubstreamHandshake::Flush =>
+					match Sink::poll_flush(this.socket.as_mut(), cx)? {
+						Poll::Ready(()) =>
+							*this.handshake = NotificationsInSubstreamHandshake::Sent,
+						Poll::Pending => {
+							*this.handshake = NotificationsInSubstreamHandshake::Flush;
+							return Poll::Pending
+						}
+					},
+
+				st @ NotificationsInSubstreamHandshake::NotSent |
+				st @ NotificationsInSubstreamHandshake::Sent |
+				st @ NotificationsInSubstreamHandshake::ClosingInResponseToRemote |
+				st @ NotificationsInSubstreamHandshake::BothSidesClosed => {
+					*this.handshake = st;
+					return Poll::Pending;
+				}
+			}
+		}
+	}
 }
 
 impl<TSubstream> Stream for NotificationsInSubstream<TSubstream>