#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
#[cfg(feature = "cbor")]
pub mod cbor;
mod codec;
mod handler;
#[cfg(feature = "json")]
pub mod json;
pub use codec::Codec;
pub use handler::ProtocolSupport;
use crate::handler::protocol::RequestProtocol;
use futures::channel::oneshot;
use handler::Handler;
use libp2p_core::{ConnectedPoint, Endpoint, Multiaddr};
use libp2p_identity::PeerId;
use libp2p_swarm::{
behaviour::{AddressChange, ConnectionClosed, DialFailure, FromSwarm},
dial_opts::DialOpts,
ConnectionDenied, ConnectionHandler, ConnectionId, NetworkBehaviour, NotifyHandler,
PollParameters, THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
};
use smallvec::SmallVec;
use std::{
collections::{HashMap, HashSet, VecDeque},
fmt,
sync::{atomic::AtomicU64, Arc},
task::{Context, Poll},
time::Duration,
};
#[derive(Debug)]
pub enum Message<TRequest, TResponse, TChannelResponse = TResponse> {
Request {
request_id: RequestId,
request: TRequest,
channel: ResponseChannel<TChannelResponse>,
},
Response {
request_id: RequestId,
response: TResponse,
},
}
#[derive(Debug)]
pub enum Event<TRequest, TResponse, TChannelResponse = TResponse> {
Message {
peer: PeerId,
message: Message<TRequest, TResponse, TChannelResponse>,
},
OutboundFailure {
peer: PeerId,
request_id: RequestId,
error: OutboundFailure,
},
InboundFailure {
peer: PeerId,
request_id: RequestId,
error: InboundFailure,
},
ResponseSent {
peer: PeerId,
request_id: RequestId,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OutboundFailure {
DialFailure,
Timeout,
ConnectionClosed,
UnsupportedProtocols,
}
impl fmt::Display for OutboundFailure {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
OutboundFailure::DialFailure => write!(f, "Failed to dial the requested peer"),
OutboundFailure::Timeout => write!(f, "Timeout while waiting for a response"),
OutboundFailure::ConnectionClosed => {
write!(f, "Connection was closed before a response was received")
}
OutboundFailure::UnsupportedProtocols => {
write!(f, "The remote supports none of the requested protocols")
}
}
}
}
impl std::error::Error for OutboundFailure {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InboundFailure {
Timeout,
ConnectionClosed,
UnsupportedProtocols,
ResponseOmission,
}
impl fmt::Display for InboundFailure {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
InboundFailure::Timeout => {
write!(f, "Timeout while receiving request or sending response")
}
InboundFailure::ConnectionClosed => {
write!(f, "Connection was closed before a response could be sent")
}
InboundFailure::UnsupportedProtocols => write!(
f,
"The local peer supports none of the protocols requested by the remote"
),
InboundFailure::ResponseOmission => write!(
f,
"The response channel was dropped without sending a response to the remote"
),
}
}
}
impl std::error::Error for InboundFailure {}
#[derive(Debug)]
pub struct ResponseChannel<TResponse> {
sender: oneshot::Sender<TResponse>,
}
impl<TResponse> ResponseChannel<TResponse> {
pub fn is_open(&self) -> bool {
!self.sender.is_canceled()
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct RequestId(u64);
impl fmt::Display for RequestId {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Debug, Clone)]
pub struct Config {
request_timeout: Duration,
connection_keep_alive: Duration,
}
impl Default for Config {
fn default() -> Self {
Self {
connection_keep_alive: Duration::from_secs(10),
request_timeout: Duration::from_secs(10),
}
}
}
impl Config {
#[deprecated(
note = "Set a global idle connection timeout via `SwarmBuilder::idle_connection_timeout` instead."
)]
pub fn set_connection_keep_alive(&mut self, v: Duration) -> &mut Self {
self.connection_keep_alive = v;
self
}
pub fn set_request_timeout(&mut self, v: Duration) -> &mut Self {
self.request_timeout = v;
self
}
}
pub struct Behaviour<TCodec>
where
TCodec: Codec + Clone + Send + 'static,
{
inbound_protocols: SmallVec<[TCodec::Protocol; 2]>,
outbound_protocols: SmallVec<[TCodec::Protocol; 2]>,
next_request_id: RequestId,
next_inbound_id: Arc<AtomicU64>,
config: Config,
codec: TCodec,
pending_events:
VecDeque<ToSwarm<Event<TCodec::Request, TCodec::Response>, RequestProtocol<TCodec>>>,
connected: HashMap<PeerId, SmallVec<[Connection; 2]>>,
addresses: HashMap<PeerId, HashSet<Multiaddr>>,
pending_outbound_requests: HashMap<PeerId, SmallVec<[RequestProtocol<TCodec>; 10]>>,
}
impl<TCodec> Behaviour<TCodec>
where
TCodec: Codec + Default + Clone + Send + 'static,
{
pub fn new<I>(protocols: I, cfg: Config) -> Self
where
I: IntoIterator<Item = (TCodec::Protocol, ProtocolSupport)>,
{
Self::with_codec(TCodec::default(), protocols, cfg)
}
}
impl<TCodec> Behaviour<TCodec>
where
TCodec: Codec + Clone + Send + 'static,
{
pub fn with_codec<I>(codec: TCodec, protocols: I, cfg: Config) -> Self
where
I: IntoIterator<Item = (TCodec::Protocol, ProtocolSupport)>,
{
let mut inbound_protocols = SmallVec::new();
let mut outbound_protocols = SmallVec::new();
for (p, s) in protocols {
if s.inbound() {
inbound_protocols.push(p.clone());
}
if s.outbound() {
outbound_protocols.push(p.clone());
}
}
Behaviour {
inbound_protocols,
outbound_protocols,
next_request_id: RequestId(1),
next_inbound_id: Arc::new(AtomicU64::new(1)),
config: cfg,
codec,
pending_events: VecDeque::new(),
connected: HashMap::new(),
pending_outbound_requests: HashMap::new(),
addresses: HashMap::new(),
}
}
pub fn send_request(&mut self, peer: &PeerId, request: TCodec::Request) -> RequestId {
let request_id = self.next_request_id();
let request = RequestProtocol {
request_id,
codec: self.codec.clone(),
protocols: self.outbound_protocols.clone(),
request,
};
if let Some(request) = self.try_send_request(peer, request) {
self.pending_events.push_back(ToSwarm::Dial {
opts: DialOpts::peer_id(*peer).build(),
});
self.pending_outbound_requests
.entry(*peer)
.or_default()
.push(request);
}
request_id
}
pub fn send_response(
&mut self,
ch: ResponseChannel<TCodec::Response>,
rs: TCodec::Response,
) -> Result<(), TCodec::Response> {
ch.sender.send(rs)
}
pub fn add_address(&mut self, peer: &PeerId, address: Multiaddr) {
self.addresses.entry(*peer).or_default().insert(address);
}
pub fn remove_address(&mut self, peer: &PeerId, address: &Multiaddr) {
let mut last = false;
if let Some(addresses) = self.addresses.get_mut(peer) {
addresses.retain(|a| a != address);
last = addresses.is_empty();
}
if last {
self.addresses.remove(peer);
}
}
pub fn is_connected(&self, peer: &PeerId) -> bool {
if let Some(connections) = self.connected.get(peer) {
!connections.is_empty()
} else {
false
}
}
pub fn is_pending_outbound(&self, peer: &PeerId, request_id: &RequestId) -> bool {
let est_conn = self
.connected
.get(peer)
.map(|cs| {
cs.iter()
.any(|c| c.pending_inbound_responses.contains(request_id))
})
.unwrap_or(false);
let pen_conn = self
.pending_outbound_requests
.get(peer)
.map(|rps| rps.iter().any(|rp| rp.request_id == *request_id))
.unwrap_or(false);
est_conn || pen_conn
}
pub fn is_pending_inbound(&self, peer: &PeerId, request_id: &RequestId) -> bool {
self.connected
.get(peer)
.map(|cs| {
cs.iter()
.any(|c| c.pending_outbound_responses.contains(request_id))
})
.unwrap_or(false)
}
fn next_request_id(&mut self) -> RequestId {
let request_id = self.next_request_id;
self.next_request_id.0 += 1;
request_id
}
fn try_send_request(
&mut self,
peer: &PeerId,
request: RequestProtocol<TCodec>,
) -> Option<RequestProtocol<TCodec>> {
if let Some(connections) = self.connected.get_mut(peer) {
if connections.is_empty() {
return Some(request);
}
let ix = (request.request_id.0 as usize) % connections.len();
let conn = &mut connections[ix];
conn.pending_inbound_responses.insert(request.request_id);
self.pending_events.push_back(ToSwarm::NotifyHandler {
peer_id: *peer,
handler: NotifyHandler::One(conn.id),
event: request,
});
None
} else {
Some(request)
}
}
fn remove_pending_outbound_response(
&mut self,
peer: &PeerId,
connection: ConnectionId,
request: RequestId,
) -> bool {
self.get_connection_mut(peer, connection)
.map(|c| c.pending_outbound_responses.remove(&request))
.unwrap_or(false)
}
fn remove_pending_inbound_response(
&mut self,
peer: &PeerId,
connection: ConnectionId,
request: &RequestId,
) -> bool {
self.get_connection_mut(peer, connection)
.map(|c| c.pending_inbound_responses.remove(request))
.unwrap_or(false)
}
fn get_connection_mut(
&mut self,
peer: &PeerId,
connection: ConnectionId,
) -> Option<&mut Connection> {
self.connected
.get_mut(peer)
.and_then(|connections| connections.iter_mut().find(|c| c.id == connection))
}
fn on_address_change(
&mut self,
AddressChange {
peer_id,
connection_id,
new,
..
}: AddressChange,
) {
let new_address = match new {
ConnectedPoint::Dialer { address, .. } => Some(address.clone()),
ConnectedPoint::Listener { .. } => None,
};
let connections = self
.connected
.get_mut(&peer_id)
.expect("Address change can only happen on an established connection.");
let connection = connections
.iter_mut()
.find(|c| c.id == connection_id)
.expect("Address change can only happen on an established connection.");
connection.remote_address = new_address;
}
fn on_connection_closed(
&mut self,
ConnectionClosed {
peer_id,
connection_id,
remaining_established,
..
}: ConnectionClosed<<Self as NetworkBehaviour>::ConnectionHandler>,
) {
let connections = self
.connected
.get_mut(&peer_id)
.expect("Expected some established connection to peer before closing.");
let connection = connections
.iter()
.position(|c| c.id == connection_id)
.map(|p: usize| connections.remove(p))
.expect("Expected connection to be established before closing.");
debug_assert_eq!(connections.is_empty(), remaining_established == 0);
if connections.is_empty() {
self.connected.remove(&peer_id);
}
for request_id in connection.pending_outbound_responses {
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::InboundFailure {
peer: peer_id,
request_id,
error: InboundFailure::ConnectionClosed,
}));
}
for request_id in connection.pending_inbound_responses {
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::OutboundFailure {
peer: peer_id,
request_id,
error: OutboundFailure::ConnectionClosed,
}));
}
}
fn on_dial_failure(&mut self, DialFailure { peer_id, .. }: DialFailure) {
if let Some(peer) = peer_id {
if let Some(pending) = self.pending_outbound_requests.remove(&peer) {
for request in pending {
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::OutboundFailure {
peer,
request_id: request.request_id,
error: OutboundFailure::DialFailure,
}));
}
}
}
}
fn preload_new_handler(
&mut self,
handler: &mut Handler<TCodec>,
peer: PeerId,
connection_id: ConnectionId,
remote_address: Option<Multiaddr>,
) {
let mut connection = Connection::new(connection_id, remote_address);
if let Some(pending_requests) = self.pending_outbound_requests.remove(&peer) {
for request in pending_requests {
connection
.pending_inbound_responses
.insert(request.request_id);
handler.on_behaviour_event(request);
}
}
self.connected.entry(peer).or_default().push(connection);
}
}
impl<TCodec> NetworkBehaviour for Behaviour<TCodec>
where
TCodec: Codec + Send + Clone + 'static,
{
type ConnectionHandler = Handler<TCodec>;
type ToSwarm = Event<TCodec::Request, TCodec::Response>;
fn handle_established_inbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
_: &Multiaddr,
_: &Multiaddr,
) -> Result<THandler<Self>, ConnectionDenied> {
let mut handler = Handler::new(
self.inbound_protocols.clone(),
self.codec.clone(),
self.config.request_timeout,
self.config.connection_keep_alive,
self.next_inbound_id.clone(),
);
self.preload_new_handler(&mut handler, peer, connection_id, None);
Ok(handler)
}
fn handle_pending_outbound_connection(
&mut self,
_connection_id: ConnectionId,
maybe_peer: Option<PeerId>,
_addresses: &[Multiaddr],
_effective_role: Endpoint,
) -> Result<Vec<Multiaddr>, ConnectionDenied> {
let peer = match maybe_peer {
None => return Ok(vec![]),
Some(peer) => peer,
};
let mut addresses = Vec::new();
if let Some(connections) = self.connected.get(&peer) {
addresses.extend(connections.iter().filter_map(|c| c.remote_address.clone()))
}
if let Some(more) = self.addresses.get(&peer) {
addresses.extend(more.iter().cloned());
}
Ok(addresses)
}
fn handle_established_outbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
remote_address: &Multiaddr,
_: Endpoint,
) -> Result<THandler<Self>, ConnectionDenied> {
let mut handler = Handler::new(
self.inbound_protocols.clone(),
self.codec.clone(),
self.config.request_timeout,
self.config.connection_keep_alive,
self.next_inbound_id.clone(),
);
self.preload_new_handler(
&mut handler,
peer,
connection_id,
Some(remote_address.clone()),
);
Ok(handler)
}
fn on_swarm_event(&mut self, event: FromSwarm<Self::ConnectionHandler>) {
match event {
FromSwarm::ConnectionEstablished(_) => {}
FromSwarm::ConnectionClosed(connection_closed) => {
self.on_connection_closed(connection_closed)
}
FromSwarm::AddressChange(address_change) => self.on_address_change(address_change),
FromSwarm::DialFailure(dial_failure) => self.on_dial_failure(dial_failure),
FromSwarm::ListenFailure(_) => {}
FromSwarm::NewListener(_) => {}
FromSwarm::NewListenAddr(_) => {}
FromSwarm::ExpiredListenAddr(_) => {}
FromSwarm::ListenerError(_) => {}
FromSwarm::ListenerClosed(_) => {}
FromSwarm::NewExternalAddrCandidate(_) => {}
FromSwarm::ExternalAddrExpired(_) => {}
FromSwarm::ExternalAddrConfirmed(_) => {}
}
}
fn on_connection_handler_event(
&mut self,
peer: PeerId,
connection: ConnectionId,
event: THandlerOutEvent<Self>,
) {
match event {
handler::Event::Response {
request_id,
response,
} => {
let removed = self.remove_pending_inbound_response(&peer, connection, &request_id);
debug_assert!(
removed,
"Expect request_id to be pending before receiving response.",
);
let message = Message::Response {
request_id,
response,
};
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::Message { peer, message }));
}
handler::Event::Request {
request_id,
request,
sender,
} => {
let channel = ResponseChannel { sender };
let message = Message::Request {
request_id,
request,
channel,
};
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::Message { peer, message }));
match self.get_connection_mut(&peer, connection) {
Some(connection) => {
let inserted = connection.pending_outbound_responses.insert(request_id);
debug_assert!(inserted, "Expect id of new request to be unknown.");
}
None => {
self.pending_events.push_back(ToSwarm::GenerateEvent(
Event::InboundFailure {
peer,
request_id,
error: InboundFailure::ConnectionClosed,
},
));
}
}
}
handler::Event::ResponseSent(request_id) => {
let removed = self.remove_pending_outbound_response(&peer, connection, request_id);
debug_assert!(
removed,
"Expect request_id to be pending before response is sent."
);
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::ResponseSent {
peer,
request_id,
}));
}
handler::Event::ResponseOmission(request_id) => {
let removed = self.remove_pending_outbound_response(&peer, connection, request_id);
debug_assert!(
removed,
"Expect request_id to be pending before response is omitted.",
);
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::InboundFailure {
peer,
request_id,
error: InboundFailure::ResponseOmission,
}));
}
handler::Event::OutboundTimeout(request_id) => {
let removed = self.remove_pending_inbound_response(&peer, connection, &request_id);
debug_assert!(
removed,
"Expect request_id to be pending before request times out."
);
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::OutboundFailure {
peer,
request_id,
error: OutboundFailure::Timeout,
}));
}
handler::Event::OutboundUnsupportedProtocols(request_id) => {
let removed = self.remove_pending_inbound_response(&peer, connection, &request_id);
debug_assert!(
removed,
"Expect request_id to be pending before failing to connect.",
);
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::OutboundFailure {
peer,
request_id,
error: OutboundFailure::UnsupportedProtocols,
}));
}
}
}
fn poll(
&mut self,
_: &mut Context<'_>,
_: &mut impl PollParameters,
) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
if let Some(ev) = self.pending_events.pop_front() {
return Poll::Ready(ev);
} else if self.pending_events.capacity() > EMPTY_QUEUE_SHRINK_THRESHOLD {
self.pending_events.shrink_to_fit();
}
Poll::Pending
}
}
const EMPTY_QUEUE_SHRINK_THRESHOLD: usize = 100;
struct Connection {
id: ConnectionId,
remote_address: Option<Multiaddr>,
pending_outbound_responses: HashSet<RequestId>,
pending_inbound_responses: HashSet<RequestId>,
}
impl Connection {
fn new(id: ConnectionId, remote_address: Option<Multiaddr>) -> Self {
Self {
id,
remote_address,
pending_outbound_responses: Default::default(),
pending_inbound_responses: Default::default(),
}
}
}