mirror of
https://github.com/pezkuwichain/pezkuwi-subxt.git
synced 2026-06-14 06:21:11 +00:00
Open one substream for each notifications protocol (#4909)
* Open one substream for each notifications protocol * Fix WASM build * Apply suggestions from code review Co-Authored-By: Toralf Wittner <tw@dtex.org> * Address concerns * Use unsigned-varint to read the varint * Use unsigned-varint * Forgot Cargo.lock Co-authored-by: Toralf Wittner <tw@dtex.org>
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
// Copyright 2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
pub use self::group::{NotifsHandlerProto, NotifsHandler, NotifsHandlerIn, NotifsHandlerOut};
|
||||
|
||||
mod group;
|
||||
mod legacy;
|
||||
mod notif_in;
|
||||
mod notif_out;
|
||||
@@ -0,0 +1,523 @@
|
||||
// Copyright 2019-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Implementations of the `IntoProtocolsHandler` and `ProtocolsHandler` traits for both incoming
|
||||
//! and outgoing substreams for all gossiping protocols together.
|
||||
//!
|
||||
//! This is the main implementation of `ProtocolsHandler` in this crate, that handles all the
|
||||
//! protocols that are Substrate-related and outside of the scope of libp2p.
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! The handler can be in one of the following states: `Initial`, `Enabled`, `Disabled`.
|
||||
//!
|
||||
//! The `Initial` state is the state that the handler initially is in. It is a temporary state
|
||||
//! during which the user must either enable or disable the handler. After that, the handler stays
|
||||
//! either enabled or disabled.
|
||||
//!
|
||||
//! On the wire, we try to open the following substreams:
|
||||
//!
|
||||
//! - One substream for each notification protocol passed as parameter to the
|
||||
//! `NotifsHandlerProto::new` function.
|
||||
//! - One "legacy" substream used for anything non-related to gossiping, and used as a fallback
|
||||
//! in case the notification protocol can't be opened.
|
||||
//!
|
||||
//! When the handler is in the `Enabled` state, we immediately open and try to maintain all the
|
||||
//! aforementioned substreams. When the handler is in the `Disabled` state, we immediately close
|
||||
//! (or abort opening) all these substreams. It is intended that in the future we allow states in
|
||||
//! which some protocols are open and not others. Symmetrically, we allow incoming
|
||||
//! Substrate-related substreams if and only if we are in the `Enabled` state.
|
||||
//!
|
||||
//! The user has the choice between sending a message with `SendNotification`, to send a
|
||||
//! notification, and `SendLegacy`, to send any other kind of message.
|
||||
//!
|
||||
|
||||
use crate::protocol::generic_proto::{
|
||||
handler::legacy::{LegacyProtoHandler, LegacyProtoHandlerProto, LegacyProtoHandlerIn, LegacyProtoHandlerOut},
|
||||
handler::notif_in::{NotifsInHandlerProto, NotifsInHandler, NotifsInHandlerIn, NotifsInHandlerOut},
|
||||
handler::notif_out::{NotifsOutHandlerProto, NotifsOutHandler, NotifsOutHandlerIn, NotifsOutHandlerOut},
|
||||
upgrade::{NotificationsIn, NotificationsOut, NotificationsHandshakeError, RegisteredProtocol, UpgradeCollec},
|
||||
};
|
||||
use crate::protocol::message::generic::{Message as GenericMessage, ConsensusMessage};
|
||||
|
||||
use bytes::BytesMut;
|
||||
use codec::Encode as _;
|
||||
use libp2p::core::{either::{EitherError, EitherOutput}, ConnectedPoint, PeerId};
|
||||
use libp2p::core::upgrade::{EitherUpgrade, UpgradeError, SelectUpgrade, InboundUpgrade, OutboundUpgrade};
|
||||
use libp2p::swarm::{
|
||||
ProtocolsHandler, ProtocolsHandlerEvent,
|
||||
IntoProtocolsHandler,
|
||||
KeepAlive,
|
||||
ProtocolsHandlerUpgrErr,
|
||||
SubstreamProtocol,
|
||||
NegotiatedSubstream,
|
||||
};
|
||||
use log::error;
|
||||
use sp_runtime::ConsensusEngineId;
|
||||
use std::{borrow::Cow, error, io, task::{Context, Poll}};
|
||||
|
||||
/// Implements the `IntoProtocolsHandler` trait of libp2p.
|
||||
///
|
||||
/// Every time a connection with a remote starts, an instance of this struct is created and
|
||||
/// sent to a background task dedicated to this connection. Once the connection is established,
|
||||
/// it is turned into a [`NotifsHandler`].
|
||||
///
|
||||
/// See the documentation at the module level for more information.
|
||||
pub struct NotifsHandlerProto {
|
||||
/// Prototypes for handlers for inbound substreams.
|
||||
in_handlers: Vec<(NotifsInHandlerProto, ConsensusEngineId)>,
|
||||
|
||||
/// Prototypes for handlers for outbound substreams.
|
||||
out_handlers: Vec<(NotifsOutHandlerProto, ConsensusEngineId)>,
|
||||
|
||||
/// Prototype for handler for backwards-compatibility.
|
||||
legacy: LegacyProtoHandlerProto,
|
||||
}
|
||||
|
||||
/// The actual handler once the connection has been established.
|
||||
///
|
||||
/// See the documentation at the module level for more information.
|
||||
pub struct NotifsHandler {
|
||||
/// Handlers for inbound substreams.
|
||||
in_handlers: Vec<(NotifsInHandler, ConsensusEngineId)>,
|
||||
|
||||
/// Handlers for outbound substreams.
|
||||
out_handlers: Vec<(NotifsOutHandler, ConsensusEngineId)>,
|
||||
|
||||
/// Handler for backwards-compatibility.
|
||||
legacy: LegacyProtoHandler,
|
||||
|
||||
/// State of this handler.
|
||||
enabled: EnabledState,
|
||||
|
||||
/// If we receive inbound substream requests while in initialization mode,
|
||||
/// we push the corresponding index here and process them when the handler
|
||||
/// gets enabled/disabled.
|
||||
pending_in: Vec<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum EnabledState {
|
||||
Initial,
|
||||
Enabled,
|
||||
Disabled,
|
||||
}
|
||||
|
||||
impl IntoProtocolsHandler for NotifsHandlerProto {
|
||||
type Handler = NotifsHandler;
|
||||
|
||||
fn inbound_protocol(&self) -> SelectUpgrade<UpgradeCollec<NotificationsIn>, RegisteredProtocol> {
|
||||
let in_handlers = self.in_handlers.iter()
|
||||
.map(|(h, _)| h.inbound_protocol())
|
||||
.collect::<UpgradeCollec<_>>();
|
||||
|
||||
SelectUpgrade::new(in_handlers, self.legacy.inbound_protocol())
|
||||
}
|
||||
|
||||
fn into_handler(self, remote_peer_id: &PeerId, connected_point: &ConnectedPoint) -> Self::Handler {
|
||||
NotifsHandler {
|
||||
in_handlers: self.in_handlers
|
||||
.into_iter()
|
||||
.map(|(p, e)| (p.into_handler(remote_peer_id, connected_point), e))
|
||||
.collect(),
|
||||
out_handlers: self.out_handlers
|
||||
.into_iter()
|
||||
.map(|(p, e)| (p.into_handler(remote_peer_id, connected_point), e))
|
||||
.collect(),
|
||||
legacy: self.legacy.into_handler(remote_peer_id, connected_point),
|
||||
enabled: EnabledState::Initial,
|
||||
pending_in: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Event that can be received by a `NotifsHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsHandlerIn {
|
||||
/// The node should start using custom protocols.
|
||||
Enable,
|
||||
|
||||
/// The node should stop using custom protocols.
|
||||
Disable,
|
||||
|
||||
/// Sends a message through the custom protocol substream.
|
||||
///
|
||||
/// > **Note**: This must **not** be an encoded `ConsensusMessage` message.
|
||||
SendLegacy {
|
||||
/// The message to send.
|
||||
message: Vec<u8>,
|
||||
},
|
||||
|
||||
/// Sends a notifications message.
|
||||
SendNotification {
|
||||
/// Name of the protocol for the message.
|
||||
///
|
||||
/// Must match one of the registered protocols. For backwards-compatibility reasons, if
|
||||
/// the remote doesn't support this protocol, we use the legacy substream to send a
|
||||
/// `ConsensusMessage` message.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
|
||||
/// The engine ID to use, in case we need to send this message over the legacy substream.
|
||||
///
|
||||
/// > **Note**: Ideally this field wouldn't be necessary, and we would deduce the engine
|
||||
/// > ID from the existing handlers. However, it is possible (especially in test
|
||||
/// > situations) that we open connections before all the notification protocols
|
||||
/// > have been registered, in which case we always rely on the legacy substream.
|
||||
engine_id: ConsensusEngineId,
|
||||
|
||||
/// The message to send.
|
||||
message: Vec<u8>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Event that can be emitted by a `NotifsHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsHandlerOut {
|
||||
/// Opened the substreams with the remote.
|
||||
Open,
|
||||
|
||||
/// Closed the substreams with the remote.
|
||||
Closed {
|
||||
/// Reason why the substream closed, for diagnostic purposes.
|
||||
reason: Cow<'static, str>,
|
||||
},
|
||||
|
||||
/// Received a non-gossiping message on the legacy substream.
|
||||
CustomMessage {
|
||||
/// Message that has been received.
|
||||
///
|
||||
/// Keep in mind that this can be a `ConsensusMessage` message, which then contains a
|
||||
/// notification.
|
||||
message: BytesMut,
|
||||
},
|
||||
|
||||
/// Received a message on a custom protocol substream.
|
||||
Notification {
|
||||
/// Engine corresponding to the message.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
|
||||
/// For legacy reasons, the name to use if we had received the message from the legacy
|
||||
/// substream.
|
||||
engine_id: ConsensusEngineId,
|
||||
|
||||
/// Message that has been received.
|
||||
///
|
||||
/// If `protocol_name` is `None`, this decodes to a `Message`. If `protocol_name` is `Some`,
|
||||
/// this is directly a gossiping message.
|
||||
message: BytesMut,
|
||||
},
|
||||
|
||||
/// A substream to the remote is clogged. The send buffer is very large, and we should print
|
||||
/// a diagnostic message and/or avoid sending more data.
|
||||
Clogged {
|
||||
/// Copy of the messages that are within the buffer, for further diagnostic.
|
||||
messages: Vec<Vec<u8>>,
|
||||
},
|
||||
|
||||
/// An error has happened on the protocol level with this node.
|
||||
ProtocolError {
|
||||
/// If true the error is severe, such as a protocol violation.
|
||||
is_severe: bool,
|
||||
/// The error that happened.
|
||||
error: Box<dyn error::Error + Send + Sync>,
|
||||
},
|
||||
}
|
||||
|
||||
impl NotifsHandlerProto {
|
||||
/// Builds a new handler.
|
||||
pub fn new(legacy: RegisteredProtocol, list: impl Into<Vec<(Cow<'static, [u8]>, ConsensusEngineId, Vec<u8>)>>) -> Self {
|
||||
let list = list.into();
|
||||
|
||||
NotifsHandlerProto {
|
||||
in_handlers: list.clone().into_iter().map(|(p, e, _)| (NotifsInHandlerProto::new(p), e)).collect(),
|
||||
out_handlers: list.clone().into_iter().map(|(p, e, _)| (NotifsOutHandlerProto::new(p), e)).collect(),
|
||||
legacy: LegacyProtoHandlerProto::new(legacy),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProtocolsHandler for NotifsHandler {
|
||||
type InEvent = NotifsHandlerIn;
|
||||
type OutEvent = NotifsHandlerOut;
|
||||
type Error = EitherError<
|
||||
EitherError<
|
||||
<NotifsInHandler as ProtocolsHandler>::Error,
|
||||
<NotifsOutHandler as ProtocolsHandler>::Error,
|
||||
>,
|
||||
<LegacyProtoHandler as ProtocolsHandler>::Error,
|
||||
>;
|
||||
type InboundProtocol = SelectUpgrade<UpgradeCollec<NotificationsIn>, RegisteredProtocol>;
|
||||
type OutboundProtocol = EitherUpgrade<NotificationsOut, RegisteredProtocol>;
|
||||
// Index within the `out_handlers`; None for legacy
|
||||
type OutboundOpenInfo = Option<usize>;
|
||||
|
||||
fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
|
||||
let in_handlers = self.in_handlers.iter()
|
||||
.map(|h| h.0.listen_protocol().into_upgrade().1)
|
||||
.collect::<UpgradeCollec<_>>();
|
||||
|
||||
let proto = SelectUpgrade::new(in_handlers, self.legacy.listen_protocol().into_upgrade().1);
|
||||
SubstreamProtocol::new(proto)
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_inbound(
|
||||
&mut self,
|
||||
out: <Self::InboundProtocol as InboundUpgrade<NegotiatedSubstream>>::Output
|
||||
) {
|
||||
match out {
|
||||
EitherOutput::First((out, num)) =>
|
||||
self.in_handlers[num].0.inject_fully_negotiated_inbound(out),
|
||||
EitherOutput::Second(out) =>
|
||||
self.legacy.inject_fully_negotiated_inbound(out),
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_outbound(
|
||||
&mut self,
|
||||
out: <Self::OutboundProtocol as OutboundUpgrade<NegotiatedSubstream>>::Output,
|
||||
num: Self::OutboundOpenInfo
|
||||
) {
|
||||
match (out, num) {
|
||||
(EitherOutput::First(out), Some(num)) =>
|
||||
self.out_handlers[num].0.inject_fully_negotiated_outbound(out, ()),
|
||||
(EitherOutput::Second(out), None) =>
|
||||
self.legacy.inject_fully_negotiated_outbound(out, ()),
|
||||
_ => error!("inject_fully_negotiated_outbound called with wrong parameters"),
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_event(&mut self, message: NotifsHandlerIn) {
|
||||
match message {
|
||||
NotifsHandlerIn::Enable => {
|
||||
self.enabled = EnabledState::Enabled;
|
||||
self.legacy.inject_event(LegacyProtoHandlerIn::Enable);
|
||||
for (handler, _) in &mut self.out_handlers {
|
||||
handler.inject_event(NotifsOutHandlerIn::Enable {
|
||||
initial_message: vec![]
|
||||
});
|
||||
}
|
||||
for num in self.pending_in.drain(..) {
|
||||
self.in_handlers[num].0.inject_event(NotifsInHandlerIn::Accept(vec![]));
|
||||
}
|
||||
},
|
||||
NotifsHandlerIn::Disable => {
|
||||
self.legacy.inject_event(LegacyProtoHandlerIn::Disable);
|
||||
// The notifications protocols start in the disabled state. If we were in the
|
||||
// "Initial" state, then we shouldn't disable the notifications protocols again.
|
||||
if self.enabled != EnabledState::Initial {
|
||||
for (handler, _) in &mut self.out_handlers {
|
||||
handler.inject_event(NotifsOutHandlerIn::Disable);
|
||||
}
|
||||
}
|
||||
self.enabled = EnabledState::Disabled;
|
||||
for num in self.pending_in.drain(..) {
|
||||
self.in_handlers[num].0.inject_event(NotifsInHandlerIn::Refuse);
|
||||
}
|
||||
},
|
||||
NotifsHandlerIn::SendLegacy { message } =>
|
||||
self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage { message }),
|
||||
NotifsHandlerIn::SendNotification { message, engine_id, protocol_name } => {
|
||||
for (handler, ngn_id) in &mut self.out_handlers {
|
||||
if handler.protocol_name() != &protocol_name[..] {
|
||||
break;
|
||||
}
|
||||
|
||||
if handler.is_open() {
|
||||
handler.inject_event(NotifsOutHandlerIn::Send(message));
|
||||
return;
|
||||
} else {
|
||||
debug_assert_eq!(engine_id, *ngn_id);
|
||||
}
|
||||
}
|
||||
|
||||
let message = GenericMessage::<(), (), (), ()>::Consensus(ConsensusMessage {
|
||||
engine_id,
|
||||
data: message,
|
||||
});
|
||||
|
||||
self.legacy.inject_event(LegacyProtoHandlerIn::SendCustomMessage {
|
||||
message: message.encode()
|
||||
});
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_dial_upgrade_error(
|
||||
&mut self,
|
||||
num: Option<usize>,
|
||||
err: ProtocolsHandlerUpgrErr<EitherError<NotificationsHandshakeError, io::Error>>
|
||||
) {
|
||||
match (err, num) {
|
||||
(ProtocolsHandlerUpgrErr::Timeout, Some(num)) =>
|
||||
self.out_handlers[num].0.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Timeout
|
||||
),
|
||||
(ProtocolsHandlerUpgrErr::Timeout, None) =>
|
||||
self.legacy.inject_dial_upgrade_error((), ProtocolsHandlerUpgrErr::Timeout),
|
||||
(ProtocolsHandlerUpgrErr::Timer, Some(num)) =>
|
||||
self.out_handlers[num].0.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Timer
|
||||
),
|
||||
(ProtocolsHandlerUpgrErr::Timer, None) =>
|
||||
self.legacy.inject_dial_upgrade_error((), ProtocolsHandlerUpgrErr::Timer),
|
||||
(ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Select(err)), Some(num)) =>
|
||||
self.out_handlers[num].0.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Select(err))
|
||||
),
|
||||
(ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Select(err)), None) =>
|
||||
self.legacy.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Select(err))
|
||||
),
|
||||
(ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Apply(EitherError::A(err))), Some(num)) =>
|
||||
self.out_handlers[num].0.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Apply(err))
|
||||
),
|
||||
(ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Apply(EitherError::B(err))), None) =>
|
||||
self.legacy.inject_dial_upgrade_error(
|
||||
(),
|
||||
ProtocolsHandlerUpgrErr::Upgrade(UpgradeError::Apply(err))
|
||||
),
|
||||
_ => error!("inject_dial_upgrade_error called with bad parameters"),
|
||||
}
|
||||
}
|
||||
|
||||
fn connection_keep_alive(&self) -> KeepAlive {
|
||||
// Iterate over each handler and return the maximum value.
|
||||
|
||||
let mut ret = self.legacy.connection_keep_alive();
|
||||
if ret.is_yes() {
|
||||
return KeepAlive::Yes;
|
||||
}
|
||||
|
||||
for (handler, _) in &self.in_handlers {
|
||||
let val = handler.connection_keep_alive();
|
||||
if val.is_yes() {
|
||||
return KeepAlive::Yes;
|
||||
}
|
||||
if ret < val { ret = val; }
|
||||
}
|
||||
|
||||
for (handler, _) in &self.out_handlers {
|
||||
let val = handler.connection_keep_alive();
|
||||
if val.is_yes() {
|
||||
return KeepAlive::Yes;
|
||||
}
|
||||
if ret < val { ret = val; }
|
||||
}
|
||||
|
||||
ret
|
||||
}
|
||||
|
||||
fn poll(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
) -> Poll<
|
||||
ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>
|
||||
> {
|
||||
for (handler_num, (handler, engine_id)) in self.in_handlers.iter_mut().enumerate() {
|
||||
while let Poll::Ready(ev) = handler.poll(cx) {
|
||||
match ev {
|
||||
ProtocolsHandlerEvent::OutboundSubstreamRequest { .. } =>
|
||||
error!("Incoming substream handler tried to open a substream"),
|
||||
ProtocolsHandlerEvent::Close(err) => void::unreachable(err),
|
||||
ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::OpenRequest(_)) =>
|
||||
match self.enabled {
|
||||
EnabledState::Initial => self.pending_in.push(handler_num),
|
||||
EnabledState::Enabled =>
|
||||
handler.inject_event(NotifsInHandlerIn::Accept(vec![])),
|
||||
EnabledState::Disabled =>
|
||||
handler.inject_event(NotifsInHandlerIn::Refuse),
|
||||
},
|
||||
ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::Closed) => {},
|
||||
ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::Notif(message)) => {
|
||||
// Note that right now the legacy substream has precedence over
|
||||
// everything. If it is not open, then we consider that nothing is open.
|
||||
if self.legacy.is_open() {
|
||||
let msg = NotifsHandlerOut::Notification {
|
||||
message,
|
||||
engine_id: *engine_id,
|
||||
protocol_name: handler.protocol_name().to_owned().into(),
|
||||
};
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(msg));
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (handler_num, (handler, _)) in self.out_handlers.iter_mut().enumerate() {
|
||||
while let Poll::Ready(ev) = handler.poll(cx) {
|
||||
match ev {
|
||||
ProtocolsHandlerEvent::OutboundSubstreamRequest { protocol, info: () } =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: protocol.map_upgrade(EitherUpgrade::A),
|
||||
info: Some(handler_num),
|
||||
}),
|
||||
ProtocolsHandlerEvent::Close(err) => void::unreachable(err),
|
||||
|
||||
// At the moment we don't actually care whether any notifications protocol
|
||||
// opens or closes.
|
||||
// Whether our communications with the remote are open or closed entirely
|
||||
// depends on the legacy substream, because as long as we are open the user of
|
||||
// this struct might try to send legacy protocol messages which we need to
|
||||
// deliver for things to work properly.
|
||||
ProtocolsHandlerEvent::Custom(NotifsOutHandlerOut::Open { .. }) => {},
|
||||
ProtocolsHandlerEvent::Custom(NotifsOutHandlerOut::Closed) => {},
|
||||
ProtocolsHandlerEvent::Custom(NotifsOutHandlerOut::Refused) => {},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
while let Poll::Ready(ev) = self.legacy.poll(cx) {
|
||||
match ev {
|
||||
ProtocolsHandlerEvent::OutboundSubstreamRequest { protocol, info: () } =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: protocol.map_upgrade(EitherUpgrade::B),
|
||||
info: None,
|
||||
}),
|
||||
ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolOpen { .. }) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(
|
||||
NotifsHandlerOut::Open
|
||||
)),
|
||||
ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomProtocolClosed { reason }) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(
|
||||
NotifsHandlerOut::Closed { reason }
|
||||
)),
|
||||
ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::CustomMessage { message }) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(
|
||||
NotifsHandlerOut::CustomMessage { message }
|
||||
)),
|
||||
ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::Clogged { messages }) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(
|
||||
NotifsHandlerOut::Clogged { messages }
|
||||
)),
|
||||
ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::ProtocolError { is_severe, error }) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(
|
||||
NotifsHandlerOut::ProtocolError { is_severe, error }
|
||||
)),
|
||||
ProtocolsHandlerEvent::Close(err) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Close(EitherError::B(err))),
|
||||
}
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,639 @@
|
||||
// Copyright 2019-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use crate::protocol::generic_proto::upgrade::{RegisteredProtocol, RegisteredProtocolEvent, RegisteredProtocolSubstream};
|
||||
use bytes::BytesMut;
|
||||
use futures::prelude::*;
|
||||
use futures_timer::Delay;
|
||||
use libp2p::core::{ConnectedPoint, PeerId, Endpoint};
|
||||
use libp2p::core::upgrade::{InboundUpgrade, OutboundUpgrade};
|
||||
use libp2p::swarm::{
|
||||
ProtocolsHandler, ProtocolsHandlerEvent,
|
||||
IntoProtocolsHandler,
|
||||
KeepAlive,
|
||||
ProtocolsHandlerUpgrErr,
|
||||
SubstreamProtocol,
|
||||
NegotiatedSubstream,
|
||||
};
|
||||
use log::{debug, error};
|
||||
use smallvec::{smallvec, SmallVec};
|
||||
use std::{borrow::Cow, error, fmt, io, mem, time::Duration};
|
||||
use std::{pin::Pin, task::{Context, Poll}};
|
||||
|
||||
/// Implements the `IntoProtocolsHandler` trait of libp2p.
|
||||
///
|
||||
/// Every time a connection with a remote starts, an instance of this struct is created and
|
||||
/// sent to a background task dedicated to this connection. Once the connection is established,
|
||||
/// it is turned into a `LegacyProtoHandler`. It then handles all communications that are specific
|
||||
/// to Substrate on that single connection.
|
||||
///
|
||||
/// Note that there can be multiple instance of this struct simultaneously for same peer. However
|
||||
/// if that happens, only one main instance can communicate with the outer layers of the code. In
|
||||
/// other words, the outer layers of the code only ever see one handler.
|
||||
///
|
||||
/// ## State of the handler
|
||||
///
|
||||
/// There are six possible states for the handler:
|
||||
///
|
||||
/// - Enabled and open, which is a normal operation.
|
||||
/// - Enabled and closed, in which case it will try to open substreams.
|
||||
/// - Disabled and open, in which case it will try to close substreams.
|
||||
/// - Disabled and closed, in which case the handler is idle. The connection will be
|
||||
/// garbage-collected after a few seconds if nothing more happens.
|
||||
/// - Initializing and open.
|
||||
/// - Initializing and closed, which is the state the handler starts in.
|
||||
///
|
||||
/// The Init/Enabled/Disabled state is entirely controlled by the user by sending `Enable` or
|
||||
/// `Disable` messages to the handler. The handler itself never transitions automatically between
|
||||
/// these states. For example, if the handler reports a network misbehaviour, it will close the
|
||||
/// substreams but it is the role of the user to send a `Disabled` event if it wants the connection
|
||||
/// to close. Otherwise, the handler will try to reopen substreams.
|
||||
/// The handler starts in the "Initializing" state and must be transitionned to Enabled or Disabled
|
||||
/// as soon as possible.
|
||||
///
|
||||
/// The Open/Closed state is decided by the handler and is reported with the `CustomProtocolOpen`
|
||||
/// and `CustomProtocolClosed` events. The `CustomMessage` event can only be generated if the
|
||||
/// handler is open.
|
||||
///
|
||||
/// ## How it works
|
||||
///
|
||||
/// When the handler is created, it is initially in the `Init` state and waits for either a
|
||||
/// `Disable` or an `Enable` message from the outer layer. At any time, the outer layer is free to
|
||||
/// toggle the handler between the disabled and enabled states.
|
||||
///
|
||||
/// When the handler switches to "enabled", it opens a substream and negotiates the protocol named
|
||||
/// `/substrate/xxx`, where `xxx` is chosen by the user and depends on the chain.
|
||||
///
|
||||
/// For backwards compatibility reasons, when we switch to "enabled" for the first time (while we
|
||||
/// are still in "init" mode) and we are the connection listener, we don't open a substream.
|
||||
///
|
||||
/// In order the handle the situation where both the remote and us get enabled at the same time,
|
||||
/// we tolerate multiple substreams open at the same time. Messages are transmitted on an arbitrary
|
||||
/// substream. The endpoints don't try to agree on a single substream.
|
||||
///
|
||||
/// We consider that we are now "closed" if the remote closes all the existing substreams.
|
||||
/// Re-opening it can then be performed by closing all active substream and re-opening one.
|
||||
///
|
||||
pub struct LegacyProtoHandlerProto {
|
||||
/// Configuration for the protocol upgrade to negotiate.
|
||||
protocol: RegisteredProtocol,
|
||||
}
|
||||
|
||||
impl LegacyProtoHandlerProto {
|
||||
/// Builds a new `LegacyProtoHandlerProto`.
|
||||
pub fn new(protocol: RegisteredProtocol) -> Self {
|
||||
LegacyProtoHandlerProto {
|
||||
protocol,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoProtocolsHandler for LegacyProtoHandlerProto {
|
||||
type Handler = LegacyProtoHandler;
|
||||
|
||||
fn inbound_protocol(&self) -> RegisteredProtocol {
|
||||
self.protocol.clone()
|
||||
}
|
||||
|
||||
fn into_handler(self, remote_peer_id: &PeerId, connected_point: &ConnectedPoint) -> Self::Handler {
|
||||
LegacyProtoHandler {
|
||||
protocol: self.protocol,
|
||||
endpoint: connected_point.to_endpoint(),
|
||||
remote_peer_id: remote_peer_id.clone(),
|
||||
state: ProtocolState::Init {
|
||||
substreams: SmallVec::new(),
|
||||
init_deadline: Delay::new(Duration::from_secs(5))
|
||||
},
|
||||
events_queue: SmallVec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The actual handler once the connection has been established.
|
||||
pub struct LegacyProtoHandler {
|
||||
/// Configuration for the protocol upgrade to negotiate.
|
||||
protocol: RegisteredProtocol,
|
||||
|
||||
/// State of the communications with the remote.
|
||||
state: ProtocolState,
|
||||
|
||||
/// Identifier of the node we're talking to. Used only for logging purposes and shouldn't have
|
||||
/// any influence on the behaviour.
|
||||
remote_peer_id: PeerId,
|
||||
|
||||
/// Whether we are the connection dialer or listener. Used to determine who, between the local
|
||||
/// node and the remote node, has priority.
|
||||
endpoint: Endpoint,
|
||||
|
||||
/// Queue of events to send to the outside.
|
||||
///
|
||||
/// This queue must only ever be modified to insert elements at the back, or remove the first
|
||||
/// element.
|
||||
events_queue: SmallVec<[ProtocolsHandlerEvent<RegisteredProtocol, (), LegacyProtoHandlerOut, ConnectionKillError>; 16]>,
|
||||
}
|
||||
|
||||
/// State of the handler.
|
||||
enum ProtocolState {
|
||||
/// Waiting for the behaviour to tell the handler whether it is enabled or disabled.
|
||||
Init {
|
||||
/// List of substreams opened by the remote but that haven't been processed yet.
|
||||
substreams: SmallVec<[RegisteredProtocolSubstream<NegotiatedSubstream>; 6]>,
|
||||
/// Deadline after which the initialization is abnormally long.
|
||||
init_deadline: Delay,
|
||||
},
|
||||
|
||||
/// Handler is opening a substream in order to activate itself.
|
||||
/// If we are in this state, we haven't sent any `CustomProtocolOpen` yet.
|
||||
Opening {
|
||||
/// Deadline after which the opening is abnormally long.
|
||||
deadline: Delay,
|
||||
},
|
||||
|
||||
/// Normal operating mode. Contains the substreams that are open.
|
||||
/// If we are in this state, we have sent a `CustomProtocolOpen` message to the outside.
|
||||
Normal {
|
||||
/// The substreams where bidirectional communications happen.
|
||||
substreams: SmallVec<[RegisteredProtocolSubstream<NegotiatedSubstream>; 4]>,
|
||||
/// Contains substreams which are being shut down.
|
||||
shutdown: SmallVec<[RegisteredProtocolSubstream<NegotiatedSubstream>; 4]>,
|
||||
},
|
||||
|
||||
/// We are disabled. Contains substreams that are being closed.
|
||||
/// If we are in this state, either we have sent a `CustomProtocolClosed` message to the
|
||||
/// outside or we have never sent any `CustomProtocolOpen` in the first place.
|
||||
Disabled {
|
||||
/// List of substreams to shut down.
|
||||
shutdown: SmallVec<[RegisteredProtocolSubstream<NegotiatedSubstream>; 6]>,
|
||||
|
||||
/// If true, we should reactivate the handler after all the substreams in `shutdown` have
|
||||
/// been closed.
|
||||
///
|
||||
/// Since we don't want to mix old and new substreams, we wait for all old substreams to
|
||||
/// be closed before opening any new one.
|
||||
reenable: bool,
|
||||
},
|
||||
|
||||
/// In this state, we don't care about anything anymore and need to kill the connection as soon
|
||||
/// as possible.
|
||||
KillAsap,
|
||||
|
||||
/// We sometimes temporarily switch to this state during processing. If we are in this state
|
||||
/// at the beginning of a method, that means something bad happened in the source code.
|
||||
Poisoned,
|
||||
}
|
||||
|
||||
/// Event that can be received by a `LegacyProtoHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum LegacyProtoHandlerIn {
|
||||
/// The node should start using custom protocols.
|
||||
Enable,
|
||||
|
||||
/// The node should stop using custom protocols.
|
||||
Disable,
|
||||
|
||||
/// Sends a message through a custom protocol substream.
|
||||
SendCustomMessage {
|
||||
/// The message to send.
|
||||
message: Vec<u8>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Event that can be emitted by a `LegacyProtoHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum LegacyProtoHandlerOut {
|
||||
/// Opened a custom protocol with the remote.
|
||||
CustomProtocolOpen {
|
||||
/// Version of the protocol that has been opened.
|
||||
version: u8,
|
||||
},
|
||||
|
||||
/// Closed a custom protocol with the remote.
|
||||
CustomProtocolClosed {
|
||||
/// Reason why the substream closed, for diagnostic purposes.
|
||||
reason: Cow<'static, str>,
|
||||
},
|
||||
|
||||
/// Receives a message on a custom protocol substream.
|
||||
CustomMessage {
|
||||
/// Message that has been received.
|
||||
message: BytesMut,
|
||||
},
|
||||
|
||||
/// A substream to the remote is clogged. The send buffer is very large, and we should print
|
||||
/// a diagnostic message and/or avoid sending more data.
|
||||
Clogged {
|
||||
/// Copy of the messages that are within the buffer, for further diagnostic.
|
||||
messages: Vec<Vec<u8>>,
|
||||
},
|
||||
|
||||
/// An error has happened on the protocol level with this node.
|
||||
ProtocolError {
|
||||
/// If true the error is severe, such as a protocol violation.
|
||||
is_severe: bool,
|
||||
/// The error that happened.
|
||||
error: Box<dyn error::Error + Send + Sync>,
|
||||
},
|
||||
}
|
||||
|
||||
impl LegacyProtoHandler {
|
||||
/// Returns true if the legacy substream is currently open.
|
||||
pub fn is_open(&self) -> bool {
|
||||
match &self.state {
|
||||
ProtocolState::Init { substreams, .. } => !substreams.is_empty(),
|
||||
ProtocolState::Opening { .. } => false,
|
||||
ProtocolState::Normal { substreams, .. } => !substreams.is_empty(),
|
||||
ProtocolState::Disabled { .. } => false,
|
||||
ProtocolState::KillAsap => false,
|
||||
ProtocolState::Poisoned => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Enables the handler.
|
||||
fn enable(&mut self) {
|
||||
self.state = match mem::replace(&mut self.state, ProtocolState::Poisoned) {
|
||||
ProtocolState::Poisoned => {
|
||||
error!(target: "sub-libp2p", "Handler with {:?} is in poisoned state",
|
||||
self.remote_peer_id);
|
||||
ProtocolState::Poisoned
|
||||
}
|
||||
|
||||
ProtocolState::Init { substreams: incoming, .. } => {
|
||||
if incoming.is_empty() {
|
||||
if let Endpoint::Dialer = self.endpoint {
|
||||
self.events_queue.push(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: SubstreamProtocol::new(self.protocol.clone()),
|
||||
info: (),
|
||||
});
|
||||
}
|
||||
ProtocolState::Opening {
|
||||
deadline: Delay::new(Duration::from_secs(60))
|
||||
}
|
||||
|
||||
} else {
|
||||
let event = LegacyProtoHandlerOut::CustomProtocolOpen {
|
||||
version: incoming[0].protocol_version()
|
||||
};
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(event));
|
||||
ProtocolState::Normal {
|
||||
substreams: incoming.into_iter().collect(),
|
||||
shutdown: SmallVec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
st @ ProtocolState::KillAsap => st,
|
||||
st @ ProtocolState::Opening { .. } => st,
|
||||
st @ ProtocolState::Normal { .. } => st,
|
||||
ProtocolState::Disabled { shutdown, .. } => {
|
||||
ProtocolState::Disabled { shutdown, reenable: true }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Disables the handler.
|
||||
fn disable(&mut self) {
|
||||
self.state = match mem::replace(&mut self.state, ProtocolState::Poisoned) {
|
||||
ProtocolState::Poisoned => {
|
||||
error!(target: "sub-libp2p", "Handler with {:?} is in poisoned state",
|
||||
self.remote_peer_id);
|
||||
ProtocolState::Poisoned
|
||||
}
|
||||
|
||||
ProtocolState::Init { substreams: mut shutdown, .. } => {
|
||||
for s in &mut shutdown {
|
||||
s.shutdown();
|
||||
}
|
||||
ProtocolState::Disabled { shutdown, reenable: false }
|
||||
}
|
||||
|
||||
ProtocolState::Opening { .. } | ProtocolState::Normal { .. } =>
|
||||
// At the moment, if we get disabled while things were working, we kill the entire
|
||||
// connection in order to force a reset of the state.
|
||||
// This is obviously an extremely shameful way to do things, but at the time of
|
||||
// the writing of this comment, the networking works very poorly and a solution
|
||||
// needs to be found.
|
||||
ProtocolState::KillAsap,
|
||||
|
||||
ProtocolState::Disabled { shutdown, .. } =>
|
||||
ProtocolState::Disabled { shutdown, reenable: false },
|
||||
|
||||
ProtocolState::KillAsap => ProtocolState::KillAsap,
|
||||
};
|
||||
}
|
||||
|
||||
/// Polls the state for events. Optionally returns an event to produce.
|
||||
#[must_use]
|
||||
fn poll_state(&mut self, cx: &mut Context)
|
||||
-> Option<ProtocolsHandlerEvent<RegisteredProtocol, (), LegacyProtoHandlerOut, ConnectionKillError>> {
|
||||
match mem::replace(&mut self.state, ProtocolState::Poisoned) {
|
||||
ProtocolState::Poisoned => {
|
||||
error!(target: "sub-libp2p", "Handler with {:?} is in poisoned state",
|
||||
self.remote_peer_id);
|
||||
self.state = ProtocolState::Poisoned;
|
||||
None
|
||||
}
|
||||
|
||||
ProtocolState::Init { substreams, mut init_deadline } => {
|
||||
match Pin::new(&mut init_deadline).poll(cx) {
|
||||
Poll::Ready(()) => {
|
||||
init_deadline = Delay::new(Duration::from_secs(60));
|
||||
error!(target: "sub-libp2p", "Handler initialization process is too long \
|
||||
with {:?}", self.remote_peer_id)
|
||||
},
|
||||
Poll::Pending => {}
|
||||
}
|
||||
|
||||
self.state = ProtocolState::Init { substreams, init_deadline };
|
||||
None
|
||||
}
|
||||
|
||||
ProtocolState::Opening { mut deadline } => {
|
||||
match Pin::new(&mut deadline).poll(cx) {
|
||||
Poll::Ready(()) => {
|
||||
deadline = Delay::new(Duration::from_secs(60));
|
||||
let event = LegacyProtoHandlerOut::ProtocolError {
|
||||
is_severe: true,
|
||||
error: "Timeout when opening protocol".to_string().into(),
|
||||
};
|
||||
self.state = ProtocolState::Opening { deadline };
|
||||
Some(ProtocolsHandlerEvent::Custom(event))
|
||||
},
|
||||
Poll::Pending => {
|
||||
self.state = ProtocolState::Opening { deadline };
|
||||
None
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
ProtocolState::Normal { mut substreams, mut shutdown } => {
|
||||
for n in (0..substreams.len()).rev() {
|
||||
let mut substream = substreams.swap_remove(n);
|
||||
match Pin::new(&mut substream).poll_next(cx) {
|
||||
Poll::Pending => substreams.push(substream),
|
||||
Poll::Ready(Some(Ok(RegisteredProtocolEvent::Message(message)))) => {
|
||||
let event = LegacyProtoHandlerOut::CustomMessage {
|
||||
message
|
||||
};
|
||||
substreams.push(substream);
|
||||
self.state = ProtocolState::Normal { substreams, shutdown };
|
||||
return Some(ProtocolsHandlerEvent::Custom(event));
|
||||
},
|
||||
Poll::Ready(Some(Ok(RegisteredProtocolEvent::Clogged { messages }))) => {
|
||||
let event = LegacyProtoHandlerOut::Clogged {
|
||||
messages,
|
||||
};
|
||||
substreams.push(substream);
|
||||
self.state = ProtocolState::Normal { substreams, shutdown };
|
||||
return Some(ProtocolsHandlerEvent::Custom(event));
|
||||
}
|
||||
Poll::Ready(None) => {
|
||||
shutdown.push(substream);
|
||||
if substreams.is_empty() {
|
||||
let event = LegacyProtoHandlerOut::CustomProtocolClosed {
|
||||
reason: "All substreams have been closed by the remote".into(),
|
||||
};
|
||||
self.state = ProtocolState::Disabled {
|
||||
shutdown: shutdown.into_iter().collect(),
|
||||
reenable: true
|
||||
};
|
||||
return Some(ProtocolsHandlerEvent::Custom(event));
|
||||
}
|
||||
}
|
||||
Poll::Ready(Some(Err(err))) => {
|
||||
if substreams.is_empty() {
|
||||
let event = LegacyProtoHandlerOut::CustomProtocolClosed {
|
||||
reason: format!("Error on the last substream: {:?}", err).into(),
|
||||
};
|
||||
self.state = ProtocolState::Disabled {
|
||||
shutdown: shutdown.into_iter().collect(),
|
||||
reenable: true
|
||||
};
|
||||
return Some(ProtocolsHandlerEvent::Custom(event));
|
||||
} else {
|
||||
debug!(target: "sub-libp2p", "Error on extra substream: {:?}", err);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// This code is reached is none if and only if none of the substreams are in a ready state.
|
||||
self.state = ProtocolState::Normal { substreams, shutdown };
|
||||
None
|
||||
}
|
||||
|
||||
ProtocolState::Disabled { mut shutdown, reenable } => {
|
||||
shutdown_list(&mut shutdown, cx);
|
||||
// If `reenable` is `true`, that means we should open the substreams system again
|
||||
// after all the substreams are closed.
|
||||
if reenable && shutdown.is_empty() {
|
||||
self.state = ProtocolState::Opening {
|
||||
deadline: Delay::new(Duration::from_secs(60))
|
||||
};
|
||||
Some(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: SubstreamProtocol::new(self.protocol.clone()),
|
||||
info: (),
|
||||
})
|
||||
} else {
|
||||
self.state = ProtocolState::Disabled { shutdown, reenable };
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
ProtocolState::KillAsap => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Called by `inject_fully_negotiated_inbound` and `inject_fully_negotiated_outbound`.
|
||||
fn inject_fully_negotiated(
|
||||
&mut self,
|
||||
mut substream: RegisteredProtocolSubstream<NegotiatedSubstream>
|
||||
) {
|
||||
self.state = match mem::replace(&mut self.state, ProtocolState::Poisoned) {
|
||||
ProtocolState::Poisoned => {
|
||||
error!(target: "sub-libp2p", "Handler with {:?} is in poisoned state",
|
||||
self.remote_peer_id);
|
||||
ProtocolState::Poisoned
|
||||
}
|
||||
|
||||
ProtocolState::Init { mut substreams, init_deadline } => {
|
||||
if substream.endpoint() == Endpoint::Dialer {
|
||||
error!(target: "sub-libp2p", "Opened dialing substream with {:?} before \
|
||||
initialization", self.remote_peer_id);
|
||||
}
|
||||
substreams.push(substream);
|
||||
ProtocolState::Init { substreams, init_deadline }
|
||||
}
|
||||
|
||||
ProtocolState::Opening { .. } => {
|
||||
let event = LegacyProtoHandlerOut::CustomProtocolOpen {
|
||||
version: substream.protocol_version()
|
||||
};
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(event));
|
||||
ProtocolState::Normal {
|
||||
substreams: smallvec![substream],
|
||||
shutdown: SmallVec::new()
|
||||
}
|
||||
}
|
||||
|
||||
ProtocolState::Normal { substreams: mut existing, shutdown } => {
|
||||
existing.push(substream);
|
||||
ProtocolState::Normal { substreams: existing, shutdown }
|
||||
}
|
||||
|
||||
ProtocolState::Disabled { mut shutdown, .. } => {
|
||||
substream.shutdown();
|
||||
shutdown.push(substream);
|
||||
ProtocolState::Disabled { shutdown, reenable: false }
|
||||
}
|
||||
|
||||
ProtocolState::KillAsap => ProtocolState::KillAsap,
|
||||
};
|
||||
}
|
||||
|
||||
/// Sends a message to the remote.
|
||||
fn send_message(&mut self, message: Vec<u8>) {
|
||||
match self.state {
|
||||
ProtocolState::Normal { ref mut substreams, .. } =>
|
||||
substreams[0].send_message(message),
|
||||
|
||||
_ => debug!(target: "sub-libp2p", "Tried to send message over closed protocol \
|
||||
with {:?}", self.remote_peer_id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ProtocolsHandler for LegacyProtoHandler {
|
||||
type InEvent = LegacyProtoHandlerIn;
|
||||
type OutEvent = LegacyProtoHandlerOut;
|
||||
type Error = ConnectionKillError;
|
||||
type InboundProtocol = RegisteredProtocol;
|
||||
type OutboundProtocol = RegisteredProtocol;
|
||||
type OutboundOpenInfo = ();
|
||||
|
||||
fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
|
||||
SubstreamProtocol::new(self.protocol.clone())
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_inbound(
|
||||
&mut self,
|
||||
proto: <Self::InboundProtocol as InboundUpgrade<NegotiatedSubstream>>::Output
|
||||
) {
|
||||
self.inject_fully_negotiated(proto);
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_outbound(
|
||||
&mut self,
|
||||
proto: <Self::OutboundProtocol as OutboundUpgrade<NegotiatedSubstream>>::Output,
|
||||
_: Self::OutboundOpenInfo
|
||||
) {
|
||||
self.inject_fully_negotiated(proto);
|
||||
}
|
||||
|
||||
fn inject_event(&mut self, message: LegacyProtoHandlerIn) {
|
||||
match message {
|
||||
LegacyProtoHandlerIn::Disable => self.disable(),
|
||||
LegacyProtoHandlerIn::Enable => self.enable(),
|
||||
LegacyProtoHandlerIn::SendCustomMessage { message } =>
|
||||
self.send_message(message),
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn inject_dial_upgrade_error(&mut self, _: (), err: ProtocolsHandlerUpgrErr<io::Error>) {
|
||||
let is_severe = match err {
|
||||
ProtocolsHandlerUpgrErr::Upgrade(_) => true,
|
||||
_ => false,
|
||||
};
|
||||
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(LegacyProtoHandlerOut::ProtocolError {
|
||||
is_severe,
|
||||
error: Box::new(err),
|
||||
}));
|
||||
}
|
||||
|
||||
fn connection_keep_alive(&self) -> KeepAlive {
|
||||
match self.state {
|
||||
ProtocolState::Init { .. } | ProtocolState::Opening { .. } |
|
||||
ProtocolState::Normal { .. } => KeepAlive::Yes,
|
||||
ProtocolState::Disabled { .. } | ProtocolState::Poisoned |
|
||||
ProtocolState::KillAsap => KeepAlive::No,
|
||||
}
|
||||
}
|
||||
|
||||
fn poll(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
) -> Poll<
|
||||
ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>
|
||||
> {
|
||||
// Flush the events queue if necessary.
|
||||
if !self.events_queue.is_empty() {
|
||||
let event = self.events_queue.remove(0);
|
||||
return Poll::Ready(event)
|
||||
}
|
||||
|
||||
// Kill the connection if needed.
|
||||
if let ProtocolState::KillAsap = self.state {
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Close(ConnectionKillError));
|
||||
}
|
||||
|
||||
// Process all the substreams.
|
||||
if let Some(event) = self.poll_state(cx) {
|
||||
return Poll::Ready(event)
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for LegacyProtoHandler {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
|
||||
f.debug_struct("LegacyProtoHandler")
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Given a list of substreams, tries to shut them down. The substreams that have been successfully
|
||||
/// shut down are removed from the list.
|
||||
fn shutdown_list
|
||||
(list: &mut SmallVec<impl smallvec::Array<Item = RegisteredProtocolSubstream<NegotiatedSubstream>>>,
|
||||
cx: &mut Context)
|
||||
{
|
||||
'outer: for n in (0..list.len()).rev() {
|
||||
let mut substream = list.swap_remove(n);
|
||||
loop {
|
||||
match substream.poll_next_unpin(cx) {
|
||||
Poll::Ready(Some(Ok(_))) => {}
|
||||
Poll::Pending => break,
|
||||
Poll::Ready(Some(Err(_))) | Poll::Ready(None) => continue 'outer,
|
||||
}
|
||||
}
|
||||
list.push(substream);
|
||||
}
|
||||
}
|
||||
|
||||
/// Error returned when switching from normal to disabled.
|
||||
#[derive(Debug)]
|
||||
pub struct ConnectionKillError;
|
||||
|
||||
impl error::Error for ConnectionKillError {
|
||||
}
|
||||
|
||||
impl fmt::Display for ConnectionKillError {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
write!(f, "Connection kill when switching from normal to disabled")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
// Copyright 2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Implementations of the `IntoProtocolsHandler` and `ProtocolsHandler` traits for ingoing
|
||||
//! substreams for a single gossiping protocol.
|
||||
//!
|
||||
//! > **Note**: Each instance corresponds to a single protocol. In order to support multiple
|
||||
//! > protocols, you need to create multiple instances and group them.
|
||||
//!
|
||||
|
||||
use crate::protocol::generic_proto::upgrade::{NotificationsIn, NotificationsInSubstream};
|
||||
use bytes::BytesMut;
|
||||
use futures::prelude::*;
|
||||
use libp2p::core::{ConnectedPoint, PeerId};
|
||||
use libp2p::core::upgrade::{DeniedUpgrade, InboundUpgrade, OutboundUpgrade};
|
||||
use libp2p::swarm::{
|
||||
ProtocolsHandler, ProtocolsHandlerEvent,
|
||||
IntoProtocolsHandler,
|
||||
KeepAlive,
|
||||
ProtocolsHandlerUpgrErr,
|
||||
SubstreamProtocol,
|
||||
NegotiatedSubstream,
|
||||
};
|
||||
use log::{error, warn};
|
||||
use smallvec::SmallVec;
|
||||
use std::{borrow::Cow, fmt, pin::Pin, str, task::{Context, Poll}};
|
||||
|
||||
/// Implements the `IntoProtocolsHandler` trait of libp2p.
|
||||
///
|
||||
/// Every time a connection with a remote starts, an instance of this struct is created and
|
||||
/// sent to a background task dedicated to this connection. Once the connection is established,
|
||||
/// it is turned into a [`NotifsInHandler`].
|
||||
pub struct NotifsInHandlerProto {
|
||||
/// Configuration for the protocol upgrade to negotiate.
|
||||
in_protocol: NotificationsIn,
|
||||
}
|
||||
|
||||
/// The actual handler once the connection has been established.
|
||||
pub struct NotifsInHandler {
|
||||
/// Configuration for the protocol upgrade to negotiate for inbound substreams.
|
||||
in_protocol: NotificationsIn,
|
||||
|
||||
/// Substream that is open with the remote.
|
||||
substream: Option<NotificationsInSubstream<NegotiatedSubstream>>,
|
||||
|
||||
/// If the substream is opened and closed rapidly, we can emit several `OpenRequest` and
|
||||
/// `Closed` messages in a row without the handler having time to respond with `Accept` or
|
||||
/// `Refuse`.
|
||||
///
|
||||
/// In order to keep the state consistent, we increment this variable every time an
|
||||
/// `OpenRequest` is emitted and decrement it every time an `Accept` or `Refuse` is received.
|
||||
pending_accept_refuses: usize,
|
||||
|
||||
/// Queue of events to send to the outside.
|
||||
///
|
||||
/// This queue is only ever modified to insert elements at the back, or remove the first
|
||||
/// element.
|
||||
events_queue: SmallVec<[ProtocolsHandlerEvent<DeniedUpgrade, (), NotifsInHandlerOut, void::Void>; 16]>,
|
||||
}
|
||||
|
||||
/// Event that can be received by a `NotifsInHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsInHandlerIn {
|
||||
/// Can be sent back as a response to an `OpenRequest`. Contains the status message to send
|
||||
/// to the remote.
|
||||
///
|
||||
/// After sending this to the handler, the substream is now considered open and `Notif` events
|
||||
/// can be received.
|
||||
Accept(Vec<u8>),
|
||||
|
||||
/// Can be sent back as a response to an `OpenRequest`.
|
||||
Refuse,
|
||||
}
|
||||
|
||||
/// Event that can be emitted by a `NotifsInHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsInHandlerOut {
|
||||
/// The remote wants to open a substream. Contains the initial message sent by the remote
|
||||
/// when the substream has been opened.
|
||||
///
|
||||
/// Every time this event is emitted, a corresponding `Accepted` or `Refused` **must** be sent
|
||||
/// back even if a `Closed` is received.
|
||||
OpenRequest(Vec<u8>),
|
||||
|
||||
/// The notifications substream has been closed by the remote. In order to avoid race
|
||||
/// conditions, this does **not** cancel any previously-sent `OpenRequest`.
|
||||
Closed,
|
||||
|
||||
/// Received a message on the notifications substream.
|
||||
///
|
||||
/// Can only happen after an `Accept` and before a `Closed`.
|
||||
Notif(BytesMut),
|
||||
}
|
||||
|
||||
impl NotifsInHandlerProto {
|
||||
/// Builds a new `NotifsInHandlerProto`.
|
||||
pub fn new(
|
||||
protocol_name: impl Into<Cow<'static, [u8]>>
|
||||
) -> Self {
|
||||
NotifsInHandlerProto {
|
||||
in_protocol: NotificationsIn::new(protocol_name),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoProtocolsHandler for NotifsInHandlerProto {
|
||||
type Handler = NotifsInHandler;
|
||||
|
||||
fn inbound_protocol(&self) -> NotificationsIn {
|
||||
self.in_protocol.clone()
|
||||
}
|
||||
|
||||
fn into_handler(self, _: &PeerId, _: &ConnectedPoint) -> Self::Handler {
|
||||
NotifsInHandler {
|
||||
in_protocol: self.in_protocol,
|
||||
substream: None,
|
||||
pending_accept_refuses: 0,
|
||||
events_queue: SmallVec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NotifsInHandler {
|
||||
/// Returns the name of the protocol that we accept.
|
||||
pub fn protocol_name(&self) -> &[u8] {
|
||||
self.in_protocol.protocol_name()
|
||||
}
|
||||
}
|
||||
|
||||
impl ProtocolsHandler for NotifsInHandler {
|
||||
type InEvent = NotifsInHandlerIn;
|
||||
type OutEvent = NotifsInHandlerOut;
|
||||
type Error = void::Void;
|
||||
type InboundProtocol = NotificationsIn;
|
||||
type OutboundProtocol = DeniedUpgrade;
|
||||
type OutboundOpenInfo = ();
|
||||
|
||||
fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
|
||||
SubstreamProtocol::new(self.in_protocol.clone())
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_inbound(
|
||||
&mut self,
|
||||
(msg, proto): <Self::InboundProtocol as InboundUpgrade<NegotiatedSubstream>>::Output
|
||||
) {
|
||||
if self.substream.is_some() {
|
||||
warn!(
|
||||
target: "sub-libp2p",
|
||||
"Received duplicate inbound notifications substream for {:?}",
|
||||
str::from_utf8(self.in_protocol.protocol_name()),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
self.substream = Some(proto);
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::OpenRequest(msg)));
|
||||
self.pending_accept_refuses = self.pending_accept_refuses
|
||||
.checked_add(1)
|
||||
.unwrap_or_else(|| {
|
||||
error!(target: "sub-libp2p", "Overflow in pending_accept_refuses");
|
||||
usize::max_value()
|
||||
});
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_outbound(
|
||||
&mut self,
|
||||
out: <Self::OutboundProtocol as OutboundUpgrade<NegotiatedSubstream>>::Output,
|
||||
_: Self::OutboundOpenInfo
|
||||
) {
|
||||
// We never emit any outgoing substream.
|
||||
void::unreachable(out)
|
||||
}
|
||||
|
||||
fn inject_event(&mut self, message: NotifsInHandlerIn) {
|
||||
self.pending_accept_refuses = match self.pending_accept_refuses.checked_sub(1) {
|
||||
Some(v) => v,
|
||||
None => {
|
||||
error!(
|
||||
target: "sub-libp2p",
|
||||
"Inconsistent state: received Accept/Refuse when no pending request exists"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// If we send multiple `OpenRequest`s in a row, we will receive back multiple
|
||||
// `Accept`/`Refuse` messages. All of them are obsolete except the last one.
|
||||
if self.pending_accept_refuses != 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
match (message, self.substream.as_mut()) {
|
||||
(NotifsInHandlerIn::Accept(message), Some(sub)) => sub.send_handshake(message),
|
||||
(NotifsInHandlerIn::Accept(_), None) => {},
|
||||
(NotifsInHandlerIn::Refuse, _) => self.substream = None,
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_dial_upgrade_error(&mut self, _: (), _: ProtocolsHandlerUpgrErr<void::Void>) {
|
||||
error!(target: "sub-libp2p", "Received dial upgrade error in inbound-only handler");
|
||||
}
|
||||
|
||||
fn connection_keep_alive(&self) -> KeepAlive {
|
||||
if self.substream.is_some() {
|
||||
KeepAlive::Yes
|
||||
} else {
|
||||
KeepAlive::No
|
||||
}
|
||||
}
|
||||
|
||||
fn poll(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
) -> Poll<
|
||||
ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>
|
||||
> {
|
||||
// Flush the events queue if necessary.
|
||||
if !self.events_queue.is_empty() {
|
||||
let event = self.events_queue.remove(0);
|
||||
return Poll::Ready(event)
|
||||
}
|
||||
|
||||
match self.substream.as_mut().map(|s| Stream::poll_next(Pin::new(s), cx)) {
|
||||
None | Some(Poll::Pending) => {},
|
||||
Some(Poll::Ready(Some(Ok(msg)))) =>
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::Notif(msg))),
|
||||
Some(Poll::Ready(None)) | Some(Poll::Ready(Some(Err(_)))) => {
|
||||
self.substream = None;
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(NotifsInHandlerOut::Closed));
|
||||
},
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for NotifsInHandler {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
|
||||
f.debug_struct("NotifsInHandler")
|
||||
.field("substream_open", &self.substream.is_some())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,395 @@
|
||||
// Copyright 2019-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Implementations of the `IntoProtocolsHandler` and `ProtocolsHandler` traits for outgoing
|
||||
//! substreams of a single gossiping protocol.
|
||||
//!
|
||||
//! > **Note**: Each instance corresponds to a single protocol. In order to support multiple
|
||||
//! > protocols, you need to create multiple instances and group them.
|
||||
//!
|
||||
|
||||
use crate::protocol::generic_proto::upgrade::{NotificationsOut, NotificationsOutSubstream, NotificationsHandshakeError};
|
||||
use futures::prelude::*;
|
||||
use libp2p::core::{ConnectedPoint, PeerId};
|
||||
use libp2p::core::upgrade::{DeniedUpgrade, InboundUpgrade, OutboundUpgrade};
|
||||
use libp2p::swarm::{
|
||||
ProtocolsHandler, ProtocolsHandlerEvent,
|
||||
IntoProtocolsHandler,
|
||||
KeepAlive,
|
||||
ProtocolsHandlerUpgrErr,
|
||||
SubstreamProtocol,
|
||||
NegotiatedSubstream,
|
||||
};
|
||||
use log::error;
|
||||
use smallvec::SmallVec;
|
||||
use std::{borrow::Cow, fmt, mem, pin::Pin, task::{Context, Poll}, time::Duration};
|
||||
use wasm_timer::Instant;
|
||||
|
||||
/// Maximum duration to open a substream and receive the handshake message. After that, we
|
||||
/// consider that we failed to open the substream.
|
||||
const OPEN_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
/// After successfully establishing a connection with the remote, we keep the connection open for
|
||||
/// at least this amount of time in order to give the rest of the code the chance to notify us to
|
||||
/// open substreams.
|
||||
const INITIAL_KEEPALIVE_TIME: Duration = Duration::from_secs(5);
|
||||
|
||||
/// Implements the `IntoProtocolsHandler` trait of libp2p.
|
||||
///
|
||||
/// Every time a connection with a remote starts, an instance of this struct is created and
|
||||
/// sent to a background task dedicated to this connection. Once the connection is established,
|
||||
/// it is turned into a [`NotifsOutHandler`].
|
||||
///
|
||||
/// See the documentation of [`NotifsOutHandler`] for more information.
|
||||
pub struct NotifsOutHandlerProto {
|
||||
/// Name of the protocol to negotiate.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
}
|
||||
|
||||
impl NotifsOutHandlerProto {
|
||||
/// Builds a new [`NotifsOutHandlerProto`]. Will use the given protocol name for the
|
||||
/// notifications substream.
|
||||
pub fn new(protocol_name: impl Into<Cow<'static, [u8]>>) -> Self {
|
||||
NotifsOutHandlerProto {
|
||||
protocol_name: protocol_name.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl IntoProtocolsHandler for NotifsOutHandlerProto {
|
||||
type Handler = NotifsOutHandler;
|
||||
|
||||
fn inbound_protocol(&self) -> DeniedUpgrade {
|
||||
DeniedUpgrade
|
||||
}
|
||||
|
||||
fn into_handler(self, _: &PeerId, _: &ConnectedPoint) -> Self::Handler {
|
||||
NotifsOutHandler {
|
||||
protocol_name: self.protocol_name,
|
||||
when_connection_open: Instant::now(),
|
||||
state: State::Disabled,
|
||||
events_queue: SmallVec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handler for an outbound notification substream.
|
||||
///
|
||||
/// When a connection is established, this handler starts in the "disabled" state, meaning that
|
||||
/// no substream will be open.
|
||||
///
|
||||
/// One can try open a substream by sending an [`NotifsOutHandlerIn::Enable`] message to the
|
||||
/// handler. Once done, the handler will try to establish then maintain an outbound substream with
|
||||
/// the remote for the purpose of sending notifications to it.
|
||||
pub struct NotifsOutHandler {
|
||||
/// Name of the protocol to negotiate.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
|
||||
/// Relationship with the node we're connected to.
|
||||
state: State,
|
||||
|
||||
/// When the connection with the remote has been successfully established.
|
||||
when_connection_open: Instant,
|
||||
|
||||
/// Queue of events to send to the outside.
|
||||
///
|
||||
/// This queue must only ever be modified to insert elements at the back, or remove the first
|
||||
/// element.
|
||||
events_queue: SmallVec<[ProtocolsHandlerEvent<NotificationsOut, (), NotifsOutHandlerOut, void::Void>; 16]>,
|
||||
}
|
||||
|
||||
/// Our relationship with the node we're connected to.
|
||||
enum State {
|
||||
/// The handler is disabled and idle. No substream is open.
|
||||
Disabled,
|
||||
|
||||
/// The handler is disabled. A substream is still open and needs to be closed.
|
||||
///
|
||||
/// > **Important**: Having this state means that `poll_close` has been called at least once,
|
||||
/// > but the `Sink` API is unclear about whether or not the stream can then
|
||||
/// > be recovered. Because of that, we must never switch from the
|
||||
/// > `DisabledOpen` state to the `Open` state while keeping the same substream.
|
||||
DisabledOpen(NotificationsOutSubstream<NegotiatedSubstream>),
|
||||
|
||||
/// The handler is disabled but we are still trying to open a substream with the remote.
|
||||
///
|
||||
/// If the handler gets enabled again, we can immediately switch to `Opening`.
|
||||
DisabledOpening,
|
||||
|
||||
/// The handler is enabled and we are trying to open a substream with the remote.
|
||||
Opening {
|
||||
/// The initial message that we sent. Necessary if we need to re-open a substream.
|
||||
initial_message: Vec<u8>,
|
||||
},
|
||||
|
||||
/// The handler is enabled. We have tried opening a substream in the past but the remote
|
||||
/// refused it.
|
||||
Refused,
|
||||
|
||||
/// The handler is enabled and substream is open.
|
||||
Open {
|
||||
/// Substream that is currently open.
|
||||
substream: NotificationsOutSubstream<NegotiatedSubstream>,
|
||||
/// The initial message that we sent. Necessary if we need to re-open a substream.
|
||||
initial_message: Vec<u8>,
|
||||
},
|
||||
|
||||
/// Poisoned state. Shouldn't be found in the wild.
|
||||
Poisoned,
|
||||
}
|
||||
|
||||
/// Event that can be received by a `NotifsOutHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsOutHandlerIn {
|
||||
/// Enables the notifications substream for this node. The handler will try to maintain a
|
||||
/// substream with the remote.
|
||||
Enable {
|
||||
/// Initial message to send to remote nodes when we open substreams.
|
||||
initial_message: Vec<u8>,
|
||||
},
|
||||
|
||||
/// Disables the notifications substream for this node. This is the default state.
|
||||
Disable,
|
||||
|
||||
/// Sends a message on the notifications substream. Ignored if the substream isn't open.
|
||||
///
|
||||
/// It is only valid to send this if the notifications substream has been enabled.
|
||||
Send(Vec<u8>),
|
||||
}
|
||||
|
||||
/// Event that can be emitted by a `NotifsOutHandler`.
|
||||
#[derive(Debug)]
|
||||
pub enum NotifsOutHandlerOut {
|
||||
/// The notifications substream has been accepted by the remote.
|
||||
Open {
|
||||
/// Handshake message sent by the remote after we opened the substream.
|
||||
handshake: Vec<u8>,
|
||||
},
|
||||
|
||||
/// The notifications substream has been closed by the remote.
|
||||
Closed,
|
||||
|
||||
/// We tried to open a notifications substream, but the remote refused it.
|
||||
///
|
||||
/// Can only happen if we're in a closed state.
|
||||
Refused,
|
||||
}
|
||||
|
||||
impl NotifsOutHandler {
|
||||
/// Returns true if the substream is currently open.
|
||||
pub fn is_open(&self) -> bool {
|
||||
match &self.state {
|
||||
State::Disabled => false,
|
||||
State::DisabledOpening => false,
|
||||
State::DisabledOpen(_) => true,
|
||||
State::Opening { .. } => false,
|
||||
State::Refused => false,
|
||||
State::Open { .. } => true,
|
||||
State::Poisoned => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the name of the protocol that we negotiate.
|
||||
pub fn protocol_name(&self) -> &[u8] {
|
||||
&self.protocol_name
|
||||
}
|
||||
}
|
||||
|
||||
impl ProtocolsHandler for NotifsOutHandler {
|
||||
type InEvent = NotifsOutHandlerIn;
|
||||
type OutEvent = NotifsOutHandlerOut;
|
||||
type Error = void::Void;
|
||||
type InboundProtocol = DeniedUpgrade;
|
||||
type OutboundProtocol = NotificationsOut;
|
||||
type OutboundOpenInfo = ();
|
||||
|
||||
fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
|
||||
SubstreamProtocol::new(DeniedUpgrade)
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_inbound(
|
||||
&mut self,
|
||||
proto: <Self::InboundProtocol as InboundUpgrade<NegotiatedSubstream>>::Output
|
||||
) {
|
||||
// We should never reach here. `proto` is a `Void`.
|
||||
void::unreachable(proto)
|
||||
}
|
||||
|
||||
fn inject_fully_negotiated_outbound(
|
||||
&mut self,
|
||||
(handshake_msg, substream): <Self::OutboundProtocol as OutboundUpgrade<NegotiatedSubstream>>::Output,
|
||||
_: ()
|
||||
) {
|
||||
match mem::replace(&mut self.state, State::Poisoned) {
|
||||
State::Opening { initial_message } => {
|
||||
let ev = NotifsOutHandlerOut::Open { handshake: handshake_msg };
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(ev));
|
||||
self.state = State::Open { substream, initial_message };
|
||||
},
|
||||
// If the handler was disabled while we were negotiating the protocol, immediately
|
||||
// close it.
|
||||
State::DisabledOpening => self.state = State::DisabledOpen(substream),
|
||||
|
||||
// Any other situation should never happen.
|
||||
State::Disabled | State::Refused | State::Open { .. } | State::DisabledOpen(_) =>
|
||||
error!("State mismatch in notifications handler: substream already open"),
|
||||
State::Poisoned => error!("Notifications handler in a poisoned state"),
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_event(&mut self, message: NotifsOutHandlerIn) {
|
||||
match message {
|
||||
NotifsOutHandlerIn::Enable { initial_message } => {
|
||||
match mem::replace(&mut self.state, State::Poisoned) {
|
||||
State::Disabled => {
|
||||
let proto = NotificationsOut::new(self.protocol_name.clone(), initial_message.clone());
|
||||
self.events_queue.push(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: SubstreamProtocol::new(proto).with_timeout(OPEN_TIMEOUT),
|
||||
info: (),
|
||||
});
|
||||
self.state = State::Opening { initial_message };
|
||||
},
|
||||
State::DisabledOpening => self.state = State::Opening { initial_message },
|
||||
State::DisabledOpen(mut sub) => {
|
||||
// As documented above, in this state we have already called `poll_close`
|
||||
// once on the substream, and it is unclear whether the substream can then
|
||||
// be recovered. When in doubt, let's drop the existing substream and
|
||||
// open a new one.
|
||||
if sub.close().now_or_never().is_none() {
|
||||
log::warn!(
|
||||
target: "sub-libp2p",
|
||||
"Improperly closed outbound notifications substream"
|
||||
);
|
||||
}
|
||||
|
||||
let proto = NotificationsOut::new(self.protocol_name.clone(), initial_message.clone());
|
||||
self.events_queue.push(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: SubstreamProtocol::new(proto).with_timeout(OPEN_TIMEOUT),
|
||||
info: (),
|
||||
});
|
||||
self.state = State::Opening { initial_message };
|
||||
},
|
||||
State::Opening { .. } | State::Refused | State::Open { .. } =>
|
||||
error!("Tried to enable notifications handler that was already enabled"),
|
||||
State::Poisoned => error!("Notifications handler in a poisoned state"),
|
||||
}
|
||||
}
|
||||
|
||||
NotifsOutHandlerIn::Disable => {
|
||||
match mem::replace(&mut self.state, State::Poisoned) {
|
||||
State::Disabled | State::DisabledOpen(_) | State::DisabledOpening =>
|
||||
error!("Tried to disable notifications handler that was already disabled"),
|
||||
State::Opening { .. } => self.state = State::DisabledOpening,
|
||||
State::Refused => self.state = State::Disabled,
|
||||
State::Open { substream, .. } => self.state = State::DisabledOpen(substream),
|
||||
State::Poisoned => error!("Notifications handler in a poisoned state"),
|
||||
}
|
||||
}
|
||||
|
||||
NotifsOutHandlerIn::Send(msg) =>
|
||||
if let State::Open { substream, .. } = &mut self.state {
|
||||
if let Some(Ok(_)) = substream.send(msg).now_or_never() {
|
||||
} else {
|
||||
log::warn!(
|
||||
target: "sub-libp2p",
|
||||
"Failed to push message to queue, dropped it"
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// This is an API misuse.
|
||||
log::warn!(
|
||||
target: "sub-libp2p",
|
||||
"Tried to send a notification on a disabled handler"
|
||||
);
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn inject_dial_upgrade_error(&mut self, _: (), _: ProtocolsHandlerUpgrErr<NotificationsHandshakeError>) {
|
||||
match mem::replace(&mut self.state, State::Poisoned) {
|
||||
State::Disabled => {},
|
||||
State::DisabledOpen(_) | State::Refused | State::Open { .. } =>
|
||||
error!("State mismatch in NotificationsOut"),
|
||||
State::Opening { .. } => {
|
||||
self.state = State::Refused;
|
||||
let ev = NotifsOutHandlerOut::Refused;
|
||||
self.events_queue.push(ProtocolsHandlerEvent::Custom(ev));
|
||||
},
|
||||
State::DisabledOpening => self.state = State::Disabled,
|
||||
State::Poisoned => error!("Notifications handler in a poisoned state"),
|
||||
}
|
||||
}
|
||||
|
||||
fn connection_keep_alive(&self) -> KeepAlive {
|
||||
match self.state {
|
||||
// We have a small grace period of `INITIAL_KEEPALIVE_TIME` during which we keep the
|
||||
// connection open no matter what, in order to avoid closing and reopening
|
||||
// connections all the time.
|
||||
State::Disabled | State::DisabledOpen(_) | State::DisabledOpening =>
|
||||
KeepAlive::Until(self.when_connection_open + INITIAL_KEEPALIVE_TIME),
|
||||
State::Opening { .. } | State::Open { .. } => KeepAlive::Yes,
|
||||
State::Refused | State::Poisoned => KeepAlive::No,
|
||||
}
|
||||
}
|
||||
|
||||
fn poll(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
) -> Poll<ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>> {
|
||||
// Flush the events queue if necessary.
|
||||
if !self.events_queue.is_empty() {
|
||||
let event = self.events_queue.remove(0);
|
||||
return Poll::Ready(event);
|
||||
}
|
||||
|
||||
match &mut self.state {
|
||||
State::Open { substream, initial_message } =>
|
||||
match Sink::poll_flush(Pin::new(substream), cx) {
|
||||
Poll::Pending | Poll::Ready(Ok(())) => {},
|
||||
Poll::Ready(Err(_)) => {
|
||||
// We try to re-open a substream.
|
||||
let initial_message = mem::replace(initial_message, Vec::new());
|
||||
self.state = State::Opening { initial_message: initial_message.clone() };
|
||||
let proto = NotificationsOut::new(self.protocol_name.clone(), initial_message);
|
||||
self.events_queue.push(ProtocolsHandlerEvent::OutboundSubstreamRequest {
|
||||
protocol: SubstreamProtocol::new(proto).with_timeout(OPEN_TIMEOUT),
|
||||
info: (),
|
||||
});
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(NotifsOutHandlerOut::Closed));
|
||||
}
|
||||
},
|
||||
|
||||
State::DisabledOpen(sub) => match Sink::poll_close(Pin::new(sub), cx) {
|
||||
Poll::Pending => {},
|
||||
Poll::Ready(Ok(())) | Poll::Ready(Err(_)) => {
|
||||
self.state = State::Disabled;
|
||||
return Poll::Ready(ProtocolsHandlerEvent::Custom(NotifsOutHandlerOut::Closed));
|
||||
},
|
||||
},
|
||||
|
||||
_ => {}
|
||||
}
|
||||
|
||||
Poll::Pending
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for NotifsOutHandler {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
|
||||
f.debug_struct("NotifsOutHandler")
|
||||
.field("open", &self.is_open())
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,413 @@
|
||||
// Copyright 2019-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
#![cfg(test)]
|
||||
|
||||
use futures::{prelude::*, ready};
|
||||
use codec::{Encode, Decode};
|
||||
use libp2p::core::nodes::listeners::ListenerId;
|
||||
use libp2p::core::ConnectedPoint;
|
||||
use libp2p::swarm::{Swarm, ProtocolsHandler, IntoProtocolsHandler};
|
||||
use libp2p::swarm::{PollParameters, NetworkBehaviour, NetworkBehaviourAction};
|
||||
use libp2p::{PeerId, Multiaddr, Transport};
|
||||
use rand::seq::SliceRandom;
|
||||
use std::{error, io, task::Context, task::Poll, time::Duration};
|
||||
use crate::message::Message;
|
||||
use crate::protocol::generic_proto::{GenericProto, GenericProtoOut};
|
||||
use sp_test_primitives::Block;
|
||||
|
||||
/// Builds two nodes that have each other as bootstrap nodes.
|
||||
/// This is to be used only for testing, and a panic will happen if something goes wrong.
|
||||
fn build_nodes() -> (Swarm<CustomProtoWithAddr>, Swarm<CustomProtoWithAddr>) {
|
||||
let mut out = Vec::with_capacity(2);
|
||||
|
||||
let keypairs: Vec<_> = (0..2).map(|_| libp2p::identity::Keypair::generate_ed25519()).collect();
|
||||
let addrs: Vec<Multiaddr> = (0..2)
|
||||
.map(|_| format!("/memory/{}", rand::random::<u64>()).parse().unwrap())
|
||||
.collect();
|
||||
|
||||
for index in 0 .. 2 {
|
||||
let keypair = keypairs[index].clone();
|
||||
let transport = libp2p::core::transport::MemoryTransport
|
||||
.and_then(move |out, endpoint| {
|
||||
let secio = libp2p::secio::SecioConfig::new(keypair);
|
||||
libp2p::core::upgrade::apply(
|
||||
out,
|
||||
secio,
|
||||
endpoint,
|
||||
libp2p::core::upgrade::Version::V1
|
||||
)
|
||||
})
|
||||
.and_then(move |(peer_id, stream), endpoint| {
|
||||
libp2p::core::upgrade::apply(
|
||||
stream,
|
||||
libp2p::yamux::Config::default(),
|
||||
endpoint,
|
||||
libp2p::core::upgrade::Version::V1
|
||||
)
|
||||
.map_ok(|muxer| (peer_id, libp2p::core::muxing::StreamMuxerBox::new(muxer)))
|
||||
})
|
||||
.timeout(Duration::from_secs(20))
|
||||
.map_err(|err| io::Error::new(io::ErrorKind::Other, err))
|
||||
.boxed();
|
||||
|
||||
let (peerset, _) = sc_peerset::Peerset::from_config(sc_peerset::PeersetConfig {
|
||||
in_peers: 25,
|
||||
out_peers: 25,
|
||||
bootnodes: if index == 0 {
|
||||
keypairs
|
||||
.iter()
|
||||
.skip(1)
|
||||
.map(|keypair| keypair.public().into_peer_id())
|
||||
.collect()
|
||||
} else {
|
||||
vec![]
|
||||
},
|
||||
reserved_only: false,
|
||||
reserved_nodes: Vec::new(),
|
||||
});
|
||||
|
||||
let behaviour = CustomProtoWithAddr {
|
||||
inner: GenericProto::new(&b"test"[..], &[1], peerset),
|
||||
addrs: addrs
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter_map(|(n, a)| if n != index {
|
||||
Some((keypairs[n].public().into_peer_id(), a.clone()))
|
||||
} else {
|
||||
None
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
let mut swarm = Swarm::new(
|
||||
transport,
|
||||
behaviour,
|
||||
keypairs[index].public().into_peer_id()
|
||||
);
|
||||
Swarm::listen_on(&mut swarm, addrs[index].clone()).unwrap();
|
||||
out.push(swarm);
|
||||
}
|
||||
|
||||
// Final output
|
||||
let mut out_iter = out.into_iter();
|
||||
let first = out_iter.next().unwrap();
|
||||
let second = out_iter.next().unwrap();
|
||||
(first, second)
|
||||
}
|
||||
|
||||
/// Wraps around the `CustomBehaviour` network behaviour, and adds hardcoded node addresses to it.
|
||||
struct CustomProtoWithAddr {
|
||||
inner: GenericProto,
|
||||
addrs: Vec<(PeerId, Multiaddr)>,
|
||||
}
|
||||
|
||||
impl std::ops::Deref for CustomProtoWithAddr {
|
||||
type Target = GenericProto;
|
||||
|
||||
fn deref(&self) -> &Self::Target {
|
||||
&self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl std::ops::DerefMut for CustomProtoWithAddr {
|
||||
fn deref_mut(&mut self) -> &mut Self::Target {
|
||||
&mut self.inner
|
||||
}
|
||||
}
|
||||
|
||||
impl NetworkBehaviour for CustomProtoWithAddr {
|
||||
type ProtocolsHandler = <GenericProto as NetworkBehaviour>::ProtocolsHandler;
|
||||
type OutEvent = <GenericProto as NetworkBehaviour>::OutEvent;
|
||||
|
||||
fn new_handler(&mut self) -> Self::ProtocolsHandler {
|
||||
self.inner.new_handler()
|
||||
}
|
||||
|
||||
fn addresses_of_peer(&mut self, peer_id: &PeerId) -> Vec<Multiaddr> {
|
||||
let mut list = self.inner.addresses_of_peer(peer_id);
|
||||
for (p, a) in self.addrs.iter() {
|
||||
if p == peer_id {
|
||||
list.push(a.clone());
|
||||
}
|
||||
}
|
||||
list
|
||||
}
|
||||
|
||||
fn inject_connected(&mut self, peer_id: PeerId, endpoint: ConnectedPoint) {
|
||||
self.inner.inject_connected(peer_id, endpoint)
|
||||
}
|
||||
|
||||
fn inject_disconnected(&mut self, peer_id: &PeerId, endpoint: ConnectedPoint) {
|
||||
self.inner.inject_disconnected(peer_id, endpoint)
|
||||
}
|
||||
|
||||
fn inject_node_event(
|
||||
&mut self,
|
||||
peer_id: PeerId,
|
||||
event: <<Self::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::OutEvent
|
||||
) {
|
||||
self.inner.inject_node_event(peer_id, event)
|
||||
}
|
||||
|
||||
fn poll(
|
||||
&mut self,
|
||||
cx: &mut Context,
|
||||
params: &mut impl PollParameters
|
||||
) -> Poll<
|
||||
NetworkBehaviourAction<
|
||||
<<Self::ProtocolsHandler as IntoProtocolsHandler>::Handler as ProtocolsHandler>::InEvent,
|
||||
Self::OutEvent
|
||||
>
|
||||
> {
|
||||
self.inner.poll(cx, params)
|
||||
}
|
||||
|
||||
fn inject_replaced(&mut self, peer_id: PeerId, closed_endpoint: ConnectedPoint, new_endpoint: ConnectedPoint) {
|
||||
self.inner.inject_replaced(peer_id, closed_endpoint, new_endpoint)
|
||||
}
|
||||
|
||||
fn inject_addr_reach_failure(&mut self, peer_id: Option<&PeerId>, addr: &Multiaddr, error: &dyn std::error::Error) {
|
||||
self.inner.inject_addr_reach_failure(peer_id, addr, error)
|
||||
}
|
||||
|
||||
fn inject_dial_failure(&mut self, peer_id: &PeerId) {
|
||||
self.inner.inject_dial_failure(peer_id)
|
||||
}
|
||||
|
||||
fn inject_new_listen_addr(&mut self, addr: &Multiaddr) {
|
||||
self.inner.inject_new_listen_addr(addr)
|
||||
}
|
||||
|
||||
fn inject_expired_listen_addr(&mut self, addr: &Multiaddr) {
|
||||
self.inner.inject_expired_listen_addr(addr)
|
||||
}
|
||||
|
||||
fn inject_new_external_addr(&mut self, addr: &Multiaddr) {
|
||||
self.inner.inject_new_external_addr(addr)
|
||||
}
|
||||
|
||||
fn inject_listener_error(&mut self, id: ListenerId, err: &(dyn error::Error + 'static)) {
|
||||
self.inner.inject_listener_error(id, err);
|
||||
}
|
||||
|
||||
fn inject_listener_closed(&mut self, id: ListenerId) {
|
||||
self.inner.inject_listener_closed(id);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn two_nodes_transfer_lots_of_packets() {
|
||||
// We spawn two nodes, then make the first one send lots of packets to the second one. The test
|
||||
// ends when the second one has received all of them.
|
||||
|
||||
// Note that if we go too high, we will reach the limit to the number of simultaneous
|
||||
// substreams allowed by the multiplexer.
|
||||
const NUM_PACKETS: u32 = 5000;
|
||||
|
||||
let (mut service1, mut service2) = build_nodes();
|
||||
|
||||
let fut1 = future::poll_fn(move |cx| -> Poll<()> {
|
||||
loop {
|
||||
match ready!(service1.poll_next_unpin(cx)) {
|
||||
Some(GenericProtoOut::CustomProtocolOpen { peer_id, .. }) => {
|
||||
for n in 0 .. NUM_PACKETS {
|
||||
service1.send_packet(
|
||||
&peer_id,
|
||||
Message::<Block>::ChainSpecific(vec![(n % 256) as u8]).encode()
|
||||
);
|
||||
}
|
||||
},
|
||||
_ => panic!(),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut packet_counter = 0u32;
|
||||
let fut2 = future::poll_fn(move |cx| {
|
||||
loop {
|
||||
match ready!(service2.poll_next_unpin(cx)) {
|
||||
Some(GenericProtoOut::CustomProtocolOpen { .. }) => {},
|
||||
Some(GenericProtoOut::CustomMessage { message, .. }) => {
|
||||
match Message::<Block>::decode(&mut &message[..]).unwrap() {
|
||||
Message::<Block>::ChainSpecific(message) => {
|
||||
assert_eq!(message.len(), 1);
|
||||
packet_counter += 1;
|
||||
if packet_counter == NUM_PACKETS {
|
||||
return Poll::Ready(())
|
||||
}
|
||||
},
|
||||
_ => panic!(),
|
||||
}
|
||||
}
|
||||
_ => panic!(),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
futures::executor::block_on(async move {
|
||||
future::select(fut1, fut2).await;
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn basic_two_nodes_requests_in_parallel() {
|
||||
let (mut service1, mut service2) = build_nodes();
|
||||
|
||||
// Generate random messages with or without a request id.
|
||||
let mut to_send = {
|
||||
let mut to_send = Vec::new();
|
||||
for _ in 0..200 { // Note: don't make that number too high or the CPU usage will explode.
|
||||
let msg = (0..10).map(|_| rand::random::<u8>()).collect::<Vec<_>>();
|
||||
to_send.push(Message::<Block>::ChainSpecific(msg));
|
||||
}
|
||||
to_send
|
||||
};
|
||||
|
||||
// Clone `to_send` in `to_receive`. Below we will remove from `to_receive` the messages we
|
||||
// receive, until the list is empty.
|
||||
let mut to_receive = to_send.clone();
|
||||
to_send.shuffle(&mut rand::thread_rng());
|
||||
|
||||
let fut1 = future::poll_fn(move |cx| -> Poll<()> {
|
||||
loop {
|
||||
match ready!(service1.poll_next_unpin(cx)) {
|
||||
Some(GenericProtoOut::CustomProtocolOpen { peer_id, .. }) => {
|
||||
for msg in to_send.drain(..) {
|
||||
service1.send_packet(&peer_id, msg.encode());
|
||||
}
|
||||
},
|
||||
_ => panic!(),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let fut2 = future::poll_fn(move |cx| {
|
||||
loop {
|
||||
match ready!(service2.poll_next_unpin(cx)) {
|
||||
Some(GenericProtoOut::CustomProtocolOpen { .. }) => {},
|
||||
Some(GenericProtoOut::CustomMessage { message, .. }) => {
|
||||
let pos = to_receive.iter().position(|m| m.encode() == message).unwrap();
|
||||
to_receive.remove(pos);
|
||||
if to_receive.is_empty() {
|
||||
return Poll::Ready(())
|
||||
}
|
||||
}
|
||||
_ => panic!(),
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
futures::executor::block_on(async move {
|
||||
future::select(fut1, fut2).await;
|
||||
});
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconnect_after_disconnect() {
|
||||
// We connect two nodes together, then force a disconnect (through the API of the `Service`),
|
||||
// check that the disconnect worked, and finally check whether they successfully reconnect.
|
||||
|
||||
let (mut service1, mut service2) = build_nodes();
|
||||
|
||||
// For this test, the services can be in the following states.
|
||||
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
|
||||
enum ServiceState { NotConnected, FirstConnec, Disconnected, ConnectedAgain }
|
||||
let mut service1_state = ServiceState::NotConnected;
|
||||
let mut service2_state = ServiceState::NotConnected;
|
||||
|
||||
// Run the events loops.
|
||||
futures::executor::block_on(future::poll_fn(|cx| -> Poll<Result<_, io::Error>> {
|
||||
loop {
|
||||
let mut service1_not_ready = false;
|
||||
|
||||
match service1.poll_next_unpin(cx) {
|
||||
Poll::Ready(Some(GenericProtoOut::CustomProtocolOpen { .. })) => {
|
||||
match service1_state {
|
||||
ServiceState::NotConnected => {
|
||||
service1_state = ServiceState::FirstConnec;
|
||||
if service2_state == ServiceState::FirstConnec {
|
||||
service1.disconnect_peer(Swarm::local_peer_id(&service2));
|
||||
}
|
||||
},
|
||||
ServiceState::Disconnected => service1_state = ServiceState::ConnectedAgain,
|
||||
ServiceState::FirstConnec | ServiceState::ConnectedAgain => panic!(),
|
||||
}
|
||||
},
|
||||
Poll::Ready(Some(GenericProtoOut::CustomProtocolClosed { .. })) => {
|
||||
match service1_state {
|
||||
ServiceState::FirstConnec => service1_state = ServiceState::Disconnected,
|
||||
ServiceState::ConnectedAgain| ServiceState::NotConnected |
|
||||
ServiceState::Disconnected => panic!(),
|
||||
}
|
||||
},
|
||||
Poll::Pending => service1_not_ready = true,
|
||||
_ => panic!()
|
||||
}
|
||||
|
||||
match service2.poll_next_unpin(cx) {
|
||||
Poll::Ready(Some(GenericProtoOut::CustomProtocolOpen { .. })) => {
|
||||
match service2_state {
|
||||
ServiceState::NotConnected => {
|
||||
service2_state = ServiceState::FirstConnec;
|
||||
if service1_state == ServiceState::FirstConnec {
|
||||
service1.disconnect_peer(Swarm::local_peer_id(&service2));
|
||||
}
|
||||
},
|
||||
ServiceState::Disconnected => service2_state = ServiceState::ConnectedAgain,
|
||||
ServiceState::FirstConnec | ServiceState::ConnectedAgain => panic!(),
|
||||
}
|
||||
},
|
||||
Poll::Ready(Some(GenericProtoOut::CustomProtocolClosed { .. })) => {
|
||||
match service2_state {
|
||||
ServiceState::FirstConnec => service2_state = ServiceState::Disconnected,
|
||||
ServiceState::ConnectedAgain| ServiceState::NotConnected |
|
||||
ServiceState::Disconnected => panic!(),
|
||||
}
|
||||
},
|
||||
Poll::Pending if service1_not_ready => break,
|
||||
Poll::Pending => {}
|
||||
_ => panic!()
|
||||
}
|
||||
}
|
||||
|
||||
if service1_state == ServiceState::ConnectedAgain && service2_state == ServiceState::ConnectedAgain {
|
||||
Poll::Ready(Ok(()))
|
||||
} else {
|
||||
Poll::Pending
|
||||
}
|
||||
})).unwrap();
|
||||
|
||||
// Do a second 3-seconds run to make sure we don't get disconnected immediately again.
|
||||
let mut delay = futures_timer::Delay::new(Duration::from_secs(3));
|
||||
futures::executor::block_on(future::poll_fn(|cx| -> Poll<Result<_, io::Error>> {
|
||||
match service1.poll_next_unpin(cx) {
|
||||
Poll::Pending => {},
|
||||
_ => panic!()
|
||||
}
|
||||
|
||||
match service2.poll_next_unpin(cx) {
|
||||
Poll::Pending => {},
|
||||
_ => panic!()
|
||||
}
|
||||
|
||||
if let Poll::Ready(()) = delay.poll_unpin(cx) {
|
||||
Poll::Ready(Ok(()))
|
||||
} else {
|
||||
Poll::Pending
|
||||
}
|
||||
})).unwrap();
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
// Copyright 2018-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
pub use self::collec::UpgradeCollec;
|
||||
pub use self::legacy::{
|
||||
RegisteredProtocol,
|
||||
RegisteredProtocolEvent,
|
||||
RegisteredProtocolName,
|
||||
RegisteredProtocolSubstream
|
||||
};
|
||||
pub use self::notifications::{
|
||||
NotificationsIn,
|
||||
NotificationsInSubstream,
|
||||
NotificationsOut,
|
||||
NotificationsOutSubstream,
|
||||
NotificationsHandshakeError,
|
||||
NotificationsOutError,
|
||||
};
|
||||
|
||||
mod collec;
|
||||
mod legacy;
|
||||
mod notifications;
|
||||
@@ -0,0 +1,97 @@
|
||||
// Copyright 2018-2020 Parity Technologies (UK) Ltd.
|
||||
//
|
||||
// Permission is hereby granted, free of charge, to any person obtaining a
|
||||
// copy of this software and associated documentation files (the "Software"),
|
||||
// to deal in the Software without restriction, including without limitation
|
||||
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
|
||||
// and/or sell copies of the Software, and to permit persons to whom the
|
||||
// Software is furnished to do so, subject to the following conditions:
|
||||
//
|
||||
// The above copyright notice and this permission notice shall be included in
|
||||
// all copies or substantial portions of the Software.
|
||||
//
|
||||
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
|
||||
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
|
||||
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
|
||||
// DEALINGS IN THE SOFTWARE.
|
||||
|
||||
use futures::prelude::*;
|
||||
use libp2p::core::upgrade::{InboundUpgrade, ProtocolName, UpgradeInfo};
|
||||
use std::{iter::FromIterator, pin::Pin, task::{Context, Poll}, vec};
|
||||
|
||||
// TODO: move this to libp2p => https://github.com/libp2p/rust-libp2p/issues/1445
|
||||
|
||||
/// Upgrade that combines multiple upgrades of the same type into one. Supports all the protocols
|
||||
/// supported by either sub-upgrade.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct UpgradeCollec<T>(pub Vec<T>);
|
||||
|
||||
impl<T> From<Vec<T>> for UpgradeCollec<T> {
|
||||
fn from(list: Vec<T>) -> Self {
|
||||
UpgradeCollec(list)
|
||||
}
|
||||
}
|
||||
|
||||
impl<T> FromIterator<T> for UpgradeCollec<T> {
|
||||
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
|
||||
UpgradeCollec(iter.into_iter().collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: UpgradeInfo> UpgradeInfo for UpgradeCollec<T> {
|
||||
type Info = ProtoNameWithUsize<T::Info>;
|
||||
type InfoIter = vec::IntoIter<Self::Info>;
|
||||
|
||||
fn protocol_info(&self) -> Self::InfoIter {
|
||||
self.0.iter().enumerate()
|
||||
.flat_map(|(n, p)|
|
||||
p.protocol_info().into_iter().map(move |i| ProtoNameWithUsize(i, n)))
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
}
|
||||
}
|
||||
|
||||
impl<T, C> InboundUpgrade<C> for UpgradeCollec<T>
|
||||
where
|
||||
T: InboundUpgrade<C>,
|
||||
{
|
||||
type Output = (T::Output, usize);
|
||||
type Error = (T::Error, usize);
|
||||
type Future = FutWithUsize<T::Future>;
|
||||
|
||||
fn upgrade_inbound(mut self, sock: C, info: Self::Info) -> Self::Future {
|
||||
let fut = self.0.remove(info.1).upgrade_inbound(sock, info.0);
|
||||
FutWithUsize(fut, info.1)
|
||||
}
|
||||
}
|
||||
|
||||
/// Groups a `ProtocolName` with a `usize`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProtoNameWithUsize<T>(T, usize);
|
||||
|
||||
impl<T: ProtocolName> ProtocolName for ProtoNameWithUsize<T> {
|
||||
fn protocol_name(&self) -> &[u8] {
|
||||
self.0.protocol_name()
|
||||
}
|
||||
}
|
||||
|
||||
/// Equivalent to `fut.map_ok(|v| (v, num)).map_err(|e| (e, num))`, where `fut` and `num` are
|
||||
/// the two fields of this struct.
|
||||
#[pin_project::pin_project]
|
||||
pub struct FutWithUsize<T>(#[pin] T, usize);
|
||||
|
||||
impl<T: Future<Output = Result<O, E>>, O, E> Future for FutWithUsize<T> {
|
||||
type Output = Result<(O, usize), (E, usize)>;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
|
||||
let this = self.project();
|
||||
match Future::poll(this.0, cx) {
|
||||
Poll::Ready(Ok(v)) => Poll::Ready(Ok((v, *this.1))),
|
||||
Poll::Ready(Err(e)) => Poll::Ready(Err((e, *this.1))),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,298 @@
|
||||
// Copyright 2018-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use crate::config::ProtocolId;
|
||||
use bytes::BytesMut;
|
||||
use futures::prelude::*;
|
||||
use futures_codec::Framed;
|
||||
use libp2p::core::{Endpoint, UpgradeInfo, InboundUpgrade, OutboundUpgrade, upgrade::ProtocolName};
|
||||
use std::{collections::VecDeque, io, pin::Pin, vec::IntoIter as VecIntoIter};
|
||||
use std::task::{Context, Poll};
|
||||
use unsigned_varint::codec::UviBytes;
|
||||
|
||||
/// Connection upgrade for a single protocol.
|
||||
///
|
||||
/// Note that "a single protocol" here refers to `par` for example. However
|
||||
/// each protocol can have multiple different versions for networking purposes.
|
||||
pub struct RegisteredProtocol {
|
||||
/// Id of the protocol for API purposes.
|
||||
id: ProtocolId,
|
||||
/// Base name of the protocol as advertised on the network.
|
||||
/// Ends with `/` so that we can append a version number behind.
|
||||
base_name: Vec<u8>,
|
||||
/// List of protocol versions that we support.
|
||||
/// Ordered in descending order so that the best comes first.
|
||||
supported_versions: Vec<u8>,
|
||||
}
|
||||
|
||||
impl RegisteredProtocol {
|
||||
/// Creates a new `RegisteredProtocol`. The `custom_data` parameter will be
|
||||
/// passed inside the `RegisteredProtocolOutput`.
|
||||
pub fn new(protocol: impl Into<ProtocolId>, versions: &[u8])
|
||||
-> Self {
|
||||
let protocol = protocol.into();
|
||||
let mut base_name = b"/substrate/".to_vec();
|
||||
base_name.extend_from_slice(protocol.as_bytes());
|
||||
base_name.extend_from_slice(b"/");
|
||||
|
||||
RegisteredProtocol {
|
||||
base_name,
|
||||
id: protocol,
|
||||
supported_versions: {
|
||||
let mut tmp = versions.to_vec();
|
||||
tmp.sort_unstable_by(|a, b| b.cmp(&a));
|
||||
tmp
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for RegisteredProtocol {
|
||||
fn clone(&self) -> Self {
|
||||
RegisteredProtocol {
|
||||
id: self.id.clone(),
|
||||
base_name: self.base_name.clone(),
|
||||
supported_versions: self.supported_versions.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Output of a `RegisteredProtocol` upgrade.
|
||||
pub struct RegisteredProtocolSubstream<TSubstream> {
|
||||
/// If true, we are in the process of closing the sink.
|
||||
is_closing: bool,
|
||||
/// Whether the local node opened this substream (dialer), or we received this substream from
|
||||
/// the remote (listener).
|
||||
endpoint: Endpoint,
|
||||
/// Buffer of packets to send.
|
||||
send_queue: VecDeque<BytesMut>,
|
||||
/// If true, we should call `poll_complete` on the inner sink.
|
||||
requires_poll_flush: bool,
|
||||
/// The underlying substream.
|
||||
inner: stream::Fuse<Framed<TSubstream, UviBytes<BytesMut>>>,
|
||||
/// Version of the protocol that was negotiated.
|
||||
protocol_version: u8,
|
||||
/// If true, we have sent a "remote is clogged" event recently and shouldn't send another one
|
||||
/// unless the buffer empties then fills itself again.
|
||||
clogged_fuse: bool,
|
||||
}
|
||||
|
||||
impl<TSubstream> RegisteredProtocolSubstream<TSubstream> {
|
||||
/// Returns the version of the protocol that was negotiated.
|
||||
pub fn protocol_version(&self) -> u8 {
|
||||
self.protocol_version
|
||||
}
|
||||
|
||||
/// Returns whether the local node opened this substream (dialer), or we received this
|
||||
/// substream from the remote (listener).
|
||||
pub fn endpoint(&self) -> Endpoint {
|
||||
self.endpoint
|
||||
}
|
||||
|
||||
/// Starts a graceful shutdown process on this substream.
|
||||
///
|
||||
/// Note that "graceful" means that we sent a closing message. We don't wait for any
|
||||
/// confirmation from the remote.
|
||||
///
|
||||
/// After calling this, the stream is guaranteed to finish soon-ish.
|
||||
pub fn shutdown(&mut self) {
|
||||
self.is_closing = true;
|
||||
self.send_queue.clear();
|
||||
}
|
||||
|
||||
/// Sends a message to the substream.
|
||||
pub fn send_message(&mut self, data: Vec<u8>) {
|
||||
if self.is_closing {
|
||||
return
|
||||
}
|
||||
|
||||
self.send_queue.push_back(From::from(&data[..]));
|
||||
}
|
||||
}
|
||||
|
||||
/// Event produced by the `RegisteredProtocolSubstream`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum RegisteredProtocolEvent {
|
||||
/// Received a message from the remote.
|
||||
Message(BytesMut),
|
||||
|
||||
/// Diagnostic event indicating that the connection is clogged and we should avoid sending too
|
||||
/// many messages to it.
|
||||
Clogged {
|
||||
/// Copy of the messages that are within the buffer, for further diagnostic.
|
||||
messages: Vec<Vec<u8>>,
|
||||
},
|
||||
}
|
||||
|
||||
impl<TSubstream> Stream for RegisteredProtocolSubstream<TSubstream>
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin {
|
||||
type Item = Result<RegisteredProtocolEvent, io::Error>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||
// Flushing the local queue.
|
||||
while !self.send_queue.is_empty() {
|
||||
match Pin::new(&mut self.inner).poll_ready(cx) {
|
||||
Poll::Ready(Ok(())) => {},
|
||||
Poll::Ready(Err(err)) => return Poll::Ready(Some(Err(err))),
|
||||
Poll::Pending => break,
|
||||
}
|
||||
|
||||
if let Some(packet) = self.send_queue.pop_front() {
|
||||
Pin::new(&mut self.inner).start_send(packet)?;
|
||||
self.requires_poll_flush = true;
|
||||
}
|
||||
}
|
||||
|
||||
// If we are closing, close as soon as the Sink is closed.
|
||||
if self.is_closing {
|
||||
return match Pin::new(&mut self.inner).poll_close(cx) {
|
||||
Poll::Pending => Poll::Pending,
|
||||
Poll::Ready(Ok(_)) => Poll::Ready(None),
|
||||
Poll::Ready(Err(err)) => Poll::Ready(Some(Err(err))),
|
||||
}
|
||||
}
|
||||
|
||||
// Indicating that the remote is clogged if that's the case.
|
||||
if self.send_queue.len() >= 2048 {
|
||||
if !self.clogged_fuse {
|
||||
// Note: this fuse is important not just for preventing us from flooding the logs;
|
||||
// if you remove the fuse, then we will always return early from this function and
|
||||
// thus never read any message from the network.
|
||||
self.clogged_fuse = true;
|
||||
return Poll::Ready(Some(Ok(RegisteredProtocolEvent::Clogged {
|
||||
messages: self.send_queue.iter()
|
||||
.map(|m| m.clone().to_vec())
|
||||
.collect(),
|
||||
})))
|
||||
}
|
||||
} else {
|
||||
self.clogged_fuse = false;
|
||||
}
|
||||
|
||||
// Flushing if necessary.
|
||||
if self.requires_poll_flush {
|
||||
if let Poll::Ready(()) = Pin::new(&mut self.inner).poll_flush(cx)? {
|
||||
self.requires_poll_flush = false;
|
||||
}
|
||||
}
|
||||
|
||||
// Receiving incoming packets.
|
||||
// Note that `inner` is wrapped in a `Fuse`, therefore we can poll it forever.
|
||||
match Pin::new(&mut self.inner).poll_next(cx)? {
|
||||
Poll::Ready(Some(data)) => {
|
||||
Poll::Ready(Some(Ok(RegisteredProtocolEvent::Message(data))))
|
||||
}
|
||||
Poll::Ready(None) =>
|
||||
if !self.requires_poll_flush && self.send_queue.is_empty() {
|
||||
Poll::Ready(None)
|
||||
} else {
|
||||
Poll::Pending
|
||||
}
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UpgradeInfo for RegisteredProtocol {
|
||||
type Info = RegisteredProtocolName;
|
||||
type InfoIter = VecIntoIter<Self::Info>;
|
||||
|
||||
#[inline]
|
||||
fn protocol_info(&self) -> Self::InfoIter {
|
||||
// Report each version as an individual protocol.
|
||||
self.supported_versions.iter().map(|&version| {
|
||||
let num = version.to_string();
|
||||
|
||||
let mut name = self.base_name.clone();
|
||||
name.extend_from_slice(num.as_bytes());
|
||||
RegisteredProtocolName {
|
||||
name,
|
||||
version,
|
||||
}
|
||||
}).collect::<Vec<_>>().into_iter()
|
||||
}
|
||||
}
|
||||
|
||||
/// Implementation of `ProtocolName` for a custom protocol.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RegisteredProtocolName {
|
||||
/// Protocol name, as advertised on the wire.
|
||||
name: Vec<u8>,
|
||||
/// Version number. Stored in string form in `name`, but duplicated here for easier retrieval.
|
||||
version: u8,
|
||||
}
|
||||
|
||||
impl ProtocolName for RegisteredProtocolName {
|
||||
fn protocol_name(&self) -> &[u8] {
|
||||
&self.name
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> InboundUpgrade<TSubstream> for RegisteredProtocol
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
type Output = RegisteredProtocolSubstream<TSubstream>;
|
||||
type Future = future::Ready<Result<Self::Output, io::Error>>;
|
||||
type Error = io::Error;
|
||||
|
||||
fn upgrade_inbound(
|
||||
self,
|
||||
socket: TSubstream,
|
||||
info: Self::Info,
|
||||
) -> Self::Future {
|
||||
let framed = {
|
||||
let mut codec = UviBytes::default();
|
||||
codec.set_max_len(16 * 1024 * 1024); // 16 MiB hard limit for packets.
|
||||
Framed::new(socket, codec)
|
||||
};
|
||||
|
||||
future::ok(RegisteredProtocolSubstream {
|
||||
is_closing: false,
|
||||
endpoint: Endpoint::Listener,
|
||||
send_queue: VecDeque::new(),
|
||||
requires_poll_flush: false,
|
||||
inner: framed.fuse(),
|
||||
protocol_version: info.version,
|
||||
clogged_fuse: false,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> OutboundUpgrade<TSubstream> for RegisteredProtocol
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
type Output = <Self as InboundUpgrade<TSubstream>>::Output;
|
||||
type Future = <Self as InboundUpgrade<TSubstream>>::Future;
|
||||
type Error = <Self as InboundUpgrade<TSubstream>>::Error;
|
||||
|
||||
fn upgrade_outbound(
|
||||
self,
|
||||
socket: TSubstream,
|
||||
info: Self::Info,
|
||||
) -> Self::Future {
|
||||
let framed = Framed::new(socket, UviBytes::default());
|
||||
|
||||
future::ok(RegisteredProtocolSubstream {
|
||||
is_closing: false,
|
||||
endpoint: Endpoint::Dialer,
|
||||
send_queue: VecDeque::new(),
|
||||
requires_poll_flush: false,
|
||||
inner: framed.fuse(),
|
||||
protocol_version: info.version,
|
||||
clogged_fuse: false,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,622 @@
|
||||
// Copyright 2019-2020 Parity Technologies (UK) Ltd.
|
||||
// This file is part of Substrate.
|
||||
|
||||
// Substrate is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
|
||||
// Substrate is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU General Public License for more details.
|
||||
|
||||
// You should have received a copy of the GNU General Public License
|
||||
// along with Substrate. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
/// Notifications protocol.
|
||||
///
|
||||
/// The Substrate notifications protocol consists in the following:
|
||||
///
|
||||
/// - Node A opens a substream to node B and sends a message which contains some protocol-specific
|
||||
/// higher-level logic. This message is prefixed with a variable-length integer message length.
|
||||
/// This message can be empty, in which case `0` is sent.
|
||||
/// - If node B accepts the substream, it sends back a message with the same properties.
|
||||
/// Afterwards, the sending side of B is closed.
|
||||
/// - If instead B refuses the connection (which typically happens because no empty slot is
|
||||
/// available), then it immediately closes the substream without sending back anything.
|
||||
/// - Node A can then send notifications to B, prefixed with a variable-length integer indicating
|
||||
/// the length of the message.
|
||||
/// - Node A closes its writing side if it doesn't want the notifications substream anymore.
|
||||
///
|
||||
/// Notification substreams are unidirectional. If A opens a substream with B, then B is
|
||||
/// encouraged but not required to open a substream to A as well.
|
||||
///
|
||||
|
||||
use bytes::BytesMut;
|
||||
use futures::{prelude::*, ready};
|
||||
use futures_codec::Framed;
|
||||
use libp2p::core::{UpgradeInfo, InboundUpgrade, OutboundUpgrade, upgrade};
|
||||
use log::error;
|
||||
use std::{borrow::Cow, collections::VecDeque, io, iter, mem, pin::Pin, task::{Context, Poll}};
|
||||
use unsigned_varint::codec::UviBytes;
|
||||
|
||||
/// Maximum allowed size of the two handshake messages, in bytes.
|
||||
const MAX_HANDSHAKE_SIZE: usize = 1024;
|
||||
/// Maximum number of buffered messages before we consider the remote unresponsive and kill the
|
||||
/// substream.
|
||||
const MAX_PENDING_MESSAGES: usize = 256;
|
||||
|
||||
/// Upgrade that accepts a substream, sends back a status message, then becomes a unidirectional
|
||||
/// stream of messages.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NotificationsIn {
|
||||
/// Protocol name to use when negotiating the substream.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
}
|
||||
|
||||
/// Upgrade that opens a substream, waits for the remote to accept by sending back a status
|
||||
/// message, then becomes a unidirectional sink of data.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NotificationsOut {
|
||||
/// Protocol name to use when negotiating the substream.
|
||||
protocol_name: Cow<'static, [u8]>,
|
||||
/// Message to send when we start the handshake.
|
||||
initial_message: Vec<u8>,
|
||||
}
|
||||
|
||||
/// A substream for incoming notification messages.
|
||||
///
|
||||
/// When creating, this struct starts in a state in which we must first send back a handshake
|
||||
/// message to the remote. No message will come before this has been done.
|
||||
#[pin_project::pin_project]
|
||||
pub struct NotificationsInSubstream<TSubstream> {
|
||||
#[pin]
|
||||
socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
|
||||
handshake: NotificationsInSubstreamHandshake,
|
||||
}
|
||||
|
||||
/// State of the handshake sending back process.
|
||||
enum NotificationsInSubstreamHandshake {
|
||||
/// Waiting for the user to give us the handshake message.
|
||||
NotSent,
|
||||
/// User gave us the handshake message. Trying to push it in the socket.
|
||||
PendingSend(Vec<u8>),
|
||||
/// Handshake message was pushed in the socket. Still need to flush.
|
||||
Close,
|
||||
/// Handshake message successfully sent.
|
||||
Sent,
|
||||
}
|
||||
|
||||
/// A substream for outgoing notification messages.
|
||||
#[pin_project::pin_project]
|
||||
pub struct NotificationsOutSubstream<TSubstream> {
|
||||
/// Substream where to send messages.
|
||||
#[pin]
|
||||
socket: Framed<TSubstream, UviBytes<io::Cursor<Vec<u8>>>>,
|
||||
/// Queue of messages waiting to be sent.
|
||||
messages_queue: VecDeque<Vec<u8>>,
|
||||
/// If true, we need to flush `socket`.
|
||||
need_flush: bool,
|
||||
}
|
||||
|
||||
impl NotificationsIn {
|
||||
/// Builds a new potential upgrade.
|
||||
pub fn new(protocol_name: impl Into<Cow<'static, [u8]>>) -> Self {
|
||||
NotificationsIn {
|
||||
protocol_name: protocol_name.into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the name of the protocol that we accept.
|
||||
pub fn protocol_name(&self) -> &[u8] {
|
||||
&self.protocol_name
|
||||
}
|
||||
}
|
||||
|
||||
impl UpgradeInfo for NotificationsIn {
|
||||
type Info = Cow<'static, [u8]>;
|
||||
type InfoIter = iter::Once<Self::Info>;
|
||||
|
||||
fn protocol_info(&self) -> Self::InfoIter {
|
||||
iter::once(self.protocol_name.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> InboundUpgrade<TSubstream> for NotificationsIn
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
type Output = (Vec<u8>, NotificationsInSubstream<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,
|
||||
) -> 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 {
|
||||
return Err(NotificationsHandshakeError::TooLarge {
|
||||
requested: initial_message_len,
|
||||
max: MAX_HANDSHAKE_SIZE,
|
||||
});
|
||||
}
|
||||
|
||||
let mut initial_message = vec![0u8; initial_message_len];
|
||||
if !initial_message.is_empty() {
|
||||
socket.read(&mut initial_message).await?;
|
||||
}
|
||||
|
||||
let substream = NotificationsInSubstream {
|
||||
socket: Framed::new(socket, UviBytes::default()),
|
||||
handshake: NotificationsInSubstreamHandshake::NotSent,
|
||||
};
|
||||
|
||||
Ok((initial_message, substream))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> NotificationsInSubstream<TSubstream>
|
||||
where TSubstream: AsyncRead + AsyncWrite,
|
||||
{
|
||||
/// 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>>) {
|
||||
match self.handshake {
|
||||
NotificationsInSubstreamHandshake::NotSent => {}
|
||||
_ => {
|
||||
error!(target: "sub-libp2p", "Tried to send handshake twice");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
self.handshake = NotificationsInSubstreamHandshake::PendingSend(message.into());
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> Stream for NotificationsInSubstream<TSubstream>
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
type Item = Result<BytesMut, io::Error>;
|
||||
|
||||
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
|
||||
let mut this = self.project();
|
||||
|
||||
// This `Stream` implementation first tries to send back the handshake if necessary.
|
||||
loop {
|
||||
match mem::replace(this.handshake, NotificationsInSubstreamHandshake::Sent) {
|
||||
NotificationsInSubstreamHandshake::Sent =>
|
||||
return Stream::poll_next(this.socket.as_mut(), cx),
|
||||
NotificationsInSubstreamHandshake::NotSent =>
|
||||
return Poll::Pending,
|
||||
NotificationsInSubstreamHandshake::PendingSend(msg) =>
|
||||
match Sink::poll_ready(this.socket.as_mut(), cx) {
|
||||
Poll::Ready(_) => {
|
||||
*this.handshake = NotificationsInSubstreamHandshake::Close;
|
||||
match Sink::start_send(this.socket.as_mut(), io::Cursor::new(msg)) {
|
||||
Ok(()) => {},
|
||||
Err(err) => return Poll::Ready(Some(Err(err))),
|
||||
}
|
||||
},
|
||||
Poll::Pending =>
|
||||
*this.handshake = NotificationsInSubstreamHandshake::PendingSend(msg),
|
||||
},
|
||||
NotificationsInSubstreamHandshake::Close =>
|
||||
match Sink::poll_close(this.socket.as_mut(), cx)? {
|
||||
Poll::Ready(()) =>
|
||||
*this.handshake = NotificationsInSubstreamHandshake::Sent,
|
||||
Poll::Pending =>
|
||||
*this.handshake = NotificationsInSubstreamHandshake::Close,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl NotificationsOut {
|
||||
/// Builds a new potential upgrade.
|
||||
pub fn new(protocol_name: impl Into<Cow<'static, [u8]>>, initial_message: impl Into<Vec<u8>>) -> Self {
|
||||
let initial_message = initial_message.into();
|
||||
if initial_message.len() > MAX_HANDSHAKE_SIZE {
|
||||
error!(target: "sub-libp2p", "Outbound networking handshake is above allowed protocol limit");
|
||||
}
|
||||
|
||||
NotificationsOut {
|
||||
protocol_name: protocol_name.into(),
|
||||
initial_message,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UpgradeInfo for NotificationsOut {
|
||||
type Info = Cow<'static, [u8]>;
|
||||
type InfoIter = iter::Once<Self::Info>;
|
||||
|
||||
fn protocol_info(&self) -> Self::InfoIter {
|
||||
iter::once(self.protocol_name.clone())
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> OutboundUpgrade<TSubstream> for NotificationsOut
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin + Send + 'static,
|
||||
{
|
||||
type Output = (Vec<u8>, NotificationsOutSubstream<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,
|
||||
) -> Self::Future {
|
||||
Box::pin(async move {
|
||||
upgrade::write_with_len_prefix(&mut socket, &self.initial_message).await?;
|
||||
|
||||
// Reading handshake.
|
||||
let handshake_len = unsigned_varint::aio::read_usize(&mut socket).await?;
|
||||
if handshake_len > MAX_HANDSHAKE_SIZE {
|
||||
return Err(NotificationsHandshakeError::TooLarge {
|
||||
requested: handshake_len,
|
||||
max: MAX_HANDSHAKE_SIZE,
|
||||
});
|
||||
}
|
||||
|
||||
let mut handshake = vec![0u8; handshake_len];
|
||||
if !handshake.is_empty() {
|
||||
socket.read(&mut handshake).await?;
|
||||
}
|
||||
|
||||
Ok((handshake, NotificationsOutSubstream {
|
||||
socket: Framed::new(socket, UviBytes::default()),
|
||||
messages_queue: VecDeque::with_capacity(MAX_PENDING_MESSAGES),
|
||||
need_flush: false,
|
||||
}))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl<TSubstream> Sink<Vec<u8>> for NotificationsOutSubstream<TSubstream>
|
||||
where TSubstream: AsyncRead + AsyncWrite + Unpin,
|
||||
{
|
||||
type Error = NotificationsOutError;
|
||||
|
||||
fn poll_ready(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn start_send(mut self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
|
||||
if self.messages_queue.len() >= MAX_PENDING_MESSAGES {
|
||||
return Err(NotificationsOutError::Clogged);
|
||||
}
|
||||
|
||||
self.messages_queue.push_back(item);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
|
||||
let mut this = self.project();
|
||||
|
||||
while !this.messages_queue.is_empty() {
|
||||
match Sink::poll_ready(this.socket.as_mut(), cx) {
|
||||
Poll::Ready(Err(err)) => return Poll::Ready(Err(From::from(err))),
|
||||
Poll::Ready(Ok(())) => {
|
||||
let msg = this.messages_queue.pop_front()
|
||||
.expect("checked for !is_empty above; qed");
|
||||
Sink::start_send(this.socket.as_mut(), io::Cursor::new(msg))?;
|
||||
*this.need_flush = true;
|
||||
},
|
||||
Poll::Pending => return Poll::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
if *this.need_flush {
|
||||
match Sink::poll_flush(this.socket.as_mut(), cx) {
|
||||
Poll::Ready(Err(err)) => return Poll::Ready(Err(From::from(err))),
|
||||
Poll::Ready(Ok(())) => *this.need_flush = false,
|
||||
Poll::Pending => return Poll::Pending,
|
||||
}
|
||||
}
|
||||
|
||||
Poll::Ready(Ok(()))
|
||||
}
|
||||
|
||||
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
|
||||
ready!(Sink::poll_flush(self.as_mut(), cx))?;
|
||||
let this = self.project();
|
||||
match Sink::poll_close(this.socket, cx) {
|
||||
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
|
||||
Poll::Ready(Err(err)) => Poll::Ready(Err(From::from(err))),
|
||||
Poll::Pending => Poll::Pending,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Error generated by sending on a notifications out substream.
|
||||
#[derive(Debug, derive_more::From, derive_more::Display)]
|
||||
pub enum NotificationsHandshakeError {
|
||||
/// I/O error on the substream.
|
||||
Io(io::Error),
|
||||
|
||||
/// Initial message or handshake was too large.
|
||||
#[display(fmt = "Initial message or handshake was too large: {}", requested)]
|
||||
TooLarge {
|
||||
/// Size requested by the remote.
|
||||
requested: usize,
|
||||
/// Maximum allowed,
|
||||
max: usize,
|
||||
},
|
||||
|
||||
/// Error while decoding the variable-length integer.
|
||||
VarintDecode(unsigned_varint::decode::Error),
|
||||
}
|
||||
|
||||
impl From<unsigned_varint::io::ReadError> for NotificationsHandshakeError {
|
||||
fn from(err: unsigned_varint::io::ReadError) -> Self {
|
||||
match err {
|
||||
unsigned_varint::io::ReadError::Io(err) => NotificationsHandshakeError::Io(err),
|
||||
unsigned_varint::io::ReadError::Decode(err) => NotificationsHandshakeError::VarintDecode(err),
|
||||
_ => {
|
||||
log::warn!("Unrecognized varint decoding error");
|
||||
NotificationsHandshakeError::Io(From::from(io::ErrorKind::InvalidData))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Error generated by sending on a notifications out substream.
|
||||
#[derive(Debug, derive_more::From, derive_more::Display)]
|
||||
pub enum NotificationsOutError {
|
||||
/// I/O error on the substream.
|
||||
Io(io::Error),
|
||||
|
||||
/// Remote doesn't process our messages quickly enough.
|
||||
///
|
||||
/// > **Note**: This is not necessarily the remote's fault, and could also be caused by the
|
||||
/// > local node sending data too quickly. Properly doing back-pressure, however,
|
||||
/// > would require a deep refactoring effort in Substrate as a whole.
|
||||
Clogged,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{NotificationsIn, NotificationsOut};
|
||||
|
||||
use async_std::net::{TcpListener, TcpStream};
|
||||
use futures::{prelude::*, channel::oneshot};
|
||||
use libp2p::core::upgrade;
|
||||
use std::pin::Pin;
|
||||
|
||||
#[test]
|
||||
fn basic_works() {
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
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(
|
||||
socket,
|
||||
NotificationsOut::new(PROTO_NAME, &b"initial message"[..]),
|
||||
upgrade::Version::V1
|
||||
).await.unwrap();
|
||||
|
||||
assert_eq!(handshake, b"hello world");
|
||||
substream.send(b"test message".to_vec()).await.unwrap();
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (initial_message, mut substream) = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await.unwrap();
|
||||
|
||||
assert_eq!(initial_message, b"initial message");
|
||||
substream.send_handshake(&b"hello world"[..]);
|
||||
|
||||
let msg = substream.next().await.unwrap().unwrap();
|
||||
assert_eq!(msg.as_ref(), b"test message");
|
||||
});
|
||||
|
||||
async_std::task::block_on(client);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_handshake() {
|
||||
// Check that everything still works when the handshake messages are empty.
|
||||
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
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(
|
||||
socket,
|
||||
NotificationsOut::new(PROTO_NAME, vec![]),
|
||||
upgrade::Version::V1
|
||||
).await.unwrap();
|
||||
|
||||
assert!(handshake.is_empty());
|
||||
substream.send(Default::default()).await.unwrap();
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (initial_message, mut substream) = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await.unwrap();
|
||||
|
||||
assert!(initial_message.is_empty());
|
||||
substream.send_handshake(vec![]);
|
||||
|
||||
let msg = substream.next().await.unwrap().unwrap();
|
||||
assert!(msg.as_ref().is_empty());
|
||||
});
|
||||
|
||||
async_std::task::block_on(client);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refused() {
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
let client = async_std::task::spawn(async move {
|
||||
let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
|
||||
let outcome = upgrade::apply_outbound(
|
||||
socket,
|
||||
NotificationsOut::new(PROTO_NAME, &b"hello"[..]),
|
||||
upgrade::Version::V1
|
||||
).await;
|
||||
|
||||
// Despite the protocol negotiation being successfully conducted on the listener
|
||||
// side, we have to receive an error here because the listener didn't send the
|
||||
// handshake.
|
||||
assert!(outcome.is_err());
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (initial_msg, substream) = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await.unwrap();
|
||||
|
||||
assert_eq!(initial_msg, b"hello");
|
||||
|
||||
// We successfully upgrade to the protocol, but then close the substream.
|
||||
drop(substream);
|
||||
});
|
||||
|
||||
async_std::task::block_on(client);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_initial_message_refused() {
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
let client = async_std::task::spawn(async move {
|
||||
let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
|
||||
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<_>>()),
|
||||
upgrade::Version::V1
|
||||
).await;
|
||||
assert!(ret.is_err());
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let ret = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await;
|
||||
assert!(ret.is_err());
|
||||
});
|
||||
|
||||
async_std::task::block_on(client);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_handshake_refused() {
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
let client = async_std::task::spawn(async move {
|
||||
let socket = TcpStream::connect(listener_addr_rx.await.unwrap()).await.unwrap();
|
||||
let ret = upgrade::apply_outbound(
|
||||
socket,
|
||||
NotificationsOut::new(PROTO_NAME, &b"initial message"[..]),
|
||||
upgrade::Version::V1
|
||||
).await;
|
||||
assert!(ret.is_err());
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (initial_message, mut substream) = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await.unwrap();
|
||||
assert_eq!(initial_message, b"initial message");
|
||||
|
||||
// We check that a handshake that is too large gets refused.
|
||||
substream.send_handshake((0..32768).map(|_| 0).collect::<Vec<_>>());
|
||||
let _ = substream.next().await;
|
||||
});
|
||||
|
||||
async_std::task::block_on(client);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buffer_is_full_closes_connection() {
|
||||
const PROTO_NAME: &'static [u8] = b"/test/proto/1";
|
||||
let (listener_addr_tx, listener_addr_rx) = oneshot::channel();
|
||||
|
||||
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(
|
||||
socket,
|
||||
NotificationsOut::new(PROTO_NAME, vec![]),
|
||||
upgrade::Version::V1
|
||||
).await.unwrap();
|
||||
|
||||
assert!(handshake.is_empty());
|
||||
|
||||
// Push an item and flush so that the test works.
|
||||
substream.send(b"hello world".to_vec()).await.unwrap();
|
||||
|
||||
for _ in 0..32768 {
|
||||
// Push an item on the sink without flushing until an error happens because the
|
||||
// buffer is full.
|
||||
let message = b"hello world!".to_vec();
|
||||
if future::poll_fn(|cx| Sink::poll_ready(Pin::new(&mut substream), cx)).await.is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
if Sink::start_send(Pin::new(&mut substream), message).is_err() {
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
Err(())
|
||||
});
|
||||
|
||||
async_std::task::block_on(async move {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
listener_addr_tx.send(listener.local_addr().unwrap()).unwrap();
|
||||
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let (initial_message, mut substream) = upgrade::apply_inbound(
|
||||
socket,
|
||||
NotificationsIn::new(PROTO_NAME)
|
||||
).await.unwrap();
|
||||
|
||||
assert!(initial_message.is_empty());
|
||||
substream.send_handshake(vec![]);
|
||||
|
||||
// Process one message so that the handshake and all works.
|
||||
let _ = substream.next().await.unwrap().unwrap();
|
||||
|
||||
client.await.unwrap();
|
||||
});
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user