fluencelabs/rust-libp2p
1
0
Fork 0
mirror of https://github.com/fluencelabs/rust-libp2p synced 2025-06-12 17:41:22 +00:00
Code Issues Projects Releases Wiki Activity

Add a PeriodicPingHandler and a PingListenHandler (#574)

Browse Source 
* Add ProtocolsHandler trait

* Reexport symbols

* Add a note about shutting down

* Add a PeriodicPingHandler and a PingListenHandler

* Fix core doctest

* Add tolerating not supported

* Fix concerns
This commit is contained in:
Pierre Krieger
2018-11-04 09:47:15 +01:00
committed by GitHub
parent 631ceaea4b
commit 4405518045
6 changed files with 935 additions and 404 deletions
Download Patch File Download Diff File Expand all files Collapse all files

414
protocols/ping/src/lib.rs
View File

@ -1,4 +1,4 @@
// Copyright 2017 Parity Technologies (UK) Ltd.
// Copyright 2017-2018 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"),
@ -56,7 +56,7 @@
//! extern crate tokio;
//!
//! use futures::{Future, Stream};
//! use libp2p_ping::{Ping, PingOutput};
//! use libp2p_ping::protocol::{Ping, PingOutput};
//! use libp2p_core::Transport;
//! use tokio::runtime::current_thread::Runtime;
//!
@ -82,7 +82,9 @@
//! ```
//!
extern crate arrayvec;
extern crate bytes;
#[macro_use]
extern crate futures;
extern crate libp2p_core;
#[macro_use]
@ -92,407 +94,13 @@ extern crate parking_lot;
extern crate rand;
extern crate tokio_codec;
extern crate tokio_io;
extern crate tokio_timer;
extern crate void;
use bytes::{BufMut, Bytes, BytesMut};
use futures::{prelude::*, future::{FutureResult, IntoFuture}, task};
use libp2p_core::{ConnectionUpgrade, Endpoint};
use rand::{distributions::Standard, prelude::*, rngs::EntropyRng};
use std::collections::VecDeque;
use std::io::Error as IoError;
use std::{iter, marker::PhantomData, mem};
use tokio_codec::{Decoder, Encoder, Framed};
use tokio_io::{AsyncRead, AsyncWrite};
pub use self::dial_handler::PeriodicPingHandler;
pub use self::listen_handler::PingListenHandler;
/// Represents a prototype for an upgrade to handle the ping protocol.
///
/// According to the design of libp2p, this struct would normally contain the configuration options
/// for the protocol, but in the case of `Ping` no configuration is required.
#[derive(Debug, Copy, Clone)]
pub struct Ping<TUserData = ()>(PhantomData<TUserData>);
pub mod protocol;
impl<TUserData> Default for Ping<TUserData> {
#[inline]
fn default() -> Self {
Ping(PhantomData)
}
}
/// Output of a `Ping` upgrade.
pub enum PingOutput<TSocket, TUserData> {
/// We are on the dialing side.
Pinger(PingDialer<TSocket, TUserData>),
/// We are on the listening side.
Ponger(PingListener<TSocket>),
}
impl<TSocket, TUserData> ConnectionUpgrade<TSocket> for Ping<TUserData>
where
TSocket: AsyncRead + AsyncWrite,
{
type NamesIter = iter::Once<(Bytes, Self::UpgradeIdentifier)>;
type UpgradeIdentifier = ();
#[inline]
fn protocol_names(&self) -> Self::NamesIter {
iter::once(("/ipfs/ping/1.0.0".into(), ()))
}
type Output = PingOutput<TSocket, TUserData>;
type Future = FutureResult<Self::Output, IoError>;
#[inline]
fn upgrade(
self,
socket: TSocket,
_: Self::UpgradeIdentifier,
endpoint: Endpoint,
) -> Self::Future {
let out = match endpoint {
Endpoint::Dialer => upgrade_as_dialer(socket),
Endpoint::Listener => upgrade_as_listener(socket),
};
Ok(out).into_future()
}
}
/// Upgrades a connection from the dialer side.
fn upgrade_as_dialer<TSocket, TUserData>(socket: TSocket) -> PingOutput<TSocket, TUserData>
where TSocket: AsyncRead + AsyncWrite,
{
let dialer = PingDialer {
inner: Framed::new(socket, Codec),
need_writer_flush: false,
needs_close: false,
sent_pings: VecDeque::with_capacity(4),
rng: EntropyRng::default(),
pings_to_send: VecDeque::with_capacity(4),
};
PingOutput::Pinger(dialer)
}
/// Upgrades a connection from the listener side.
fn upgrade_as_listener<TSocket, TUserData>(socket: TSocket) -> PingOutput<TSocket, TUserData>
where TSocket: AsyncRead + AsyncWrite,
{
let listener = PingListener {
inner: Framed::new(socket, Codec),
state: PingListenerState::Listening,
};
PingOutput::Ponger(listener)
}
/// Sends pings and receives the pongs.
///
/// Implements `Stream`. The stream indicates when we receive a pong.
pub struct PingDialer<TSocket, TUserData> {
/// The underlying socket.
inner: Framed<TSocket, Codec>,
/// If true, need to flush the sink.
need_writer_flush: bool,
/// If true, need to close the sink.
needs_close: bool,
/// List of pings that have been sent to the remote and that are waiting for an answer.
sent_pings: VecDeque<(Bytes, TUserData)>,
/// Random number generator for the ping payload.
rng: EntropyRng,
/// List of pings to send to the remote.
pings_to_send: VecDeque<(Bytes, TUserData)>,
}
impl<TSocket, TUserData> PingDialer<TSocket, TUserData> {
/// Sends a ping to the remote.
///
/// The stream will produce an event containing the user data when we receive the pong.
pub fn ping(&mut self, user_data: TUserData) {
let payload: [u8; 32] = self.rng.sample(Standard);
debug!("Preparing for ping with payload {:?}", payload);
self.pings_to_send.push_back((Bytes::from(payload.to_vec()), user_data));
}
}
impl<TSocket, TUserData> Stream for PingDialer<TSocket, TUserData>
where TSocket: AsyncRead + AsyncWrite,
{
type Item = TUserData;
type Error = IoError;
fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
if self.needs_close {
match self.inner.close() {
Ok(Async::Ready(())) => return Ok(Async::Ready(None)),
Ok(Async::NotReady) => return Ok(Async::NotReady),
Err(err) => return Err(err),
}
}
while let Some((ping, user_data)) = self.pings_to_send.pop_front() {
match self.inner.start_send(ping.clone()) {
Ok(AsyncSink::Ready) => self.need_writer_flush = true,
Ok(AsyncSink::NotReady(_)) => {
self.pings_to_send.push_front((ping, user_data));
break;
},
Err(err) => return Err(err),
}
self.sent_pings.push_back((ping, user_data));
}
if self.need_writer_flush {
match self.inner.poll_complete() {
Ok(Async::Ready(())) => self.need_writer_flush = false,
Ok(Async::NotReady) => (),
Err(err) => return Err(err),
}
}
loop {
match self.inner.poll() {
Ok(Async::Ready(Some(pong))) => {
if let Some(pos) = self.sent_pings.iter().position(|&(ref p, _)| p == &pong) {
let (_, user_data) = self.sent_pings.remove(pos)
.expect("Grabbed a valid position just above");
return Ok(Async::Ready(Some(user_data)));
} else {
debug!("Received pong that doesn't match what we sent: {:?}", pong);
}
},
Ok(Async::NotReady) => break,
Ok(Async::Ready(None)) => {
// Notify the current task so that we poll again.
self.needs_close = true;
task::current().notify();
return Ok(Async::NotReady);
}
Err(err) => return Err(err),
}
}
Ok(Async::NotReady)
}
}
/// Listens to incoming pings and answers them.
///
/// Implements `Future`. The future terminates when the underlying socket closes.
pub struct PingListener<TSocket> {
/// The underlying socket.
inner: Framed<TSocket, Codec>,
/// State of the listener.
state: PingListenerState,
}
#[derive(Debug)]
enum PingListenerState {
/// We are waiting for the next ping on the socket.
Listening,
/// We are trying to send a pong.
Sending(Bytes),
/// We are flusing the underlying sink.
Flushing,
/// We are shutting down everything.
Closing,
/// A panic happened during the processing.
Poisoned,
}
impl<TSocket> Future for PingListener<TSocket>
where TSocket: AsyncRead + AsyncWrite
{
type Item = ();
type Error = IoError;
fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
loop {
match mem::replace(&mut self.state, PingListenerState::Poisoned) {
PingListenerState::Listening => {
match self.inner.poll() {
Ok(Async::Ready(Some(payload))) => {
debug!("Received ping (payload={:?}); sending back", payload);
self.state = PingListenerState::Sending(payload.freeze())
},
Ok(Async::Ready(None)) => self.state = PingListenerState::Closing,
Ok(Async::NotReady) => {
self.state = PingListenerState::Listening;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
}
},
PingListenerState::Sending(data) => {
match self.inner.start_send(data) {
Ok(AsyncSink::Ready) => self.state = PingListenerState::Flushing,
Ok(AsyncSink::NotReady(data)) => {
self.state = PingListenerState::Sending(data);
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
}
},
PingListenerState::Flushing => {
match self.inner.poll_complete() {
Ok(Async::Ready(())) => self.state = PingListenerState::Listening,
Ok(Async::NotReady) => {
self.state = PingListenerState::Flushing;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
}
},
PingListenerState::Closing => {
match self.inner.close() {
Ok(Async::Ready(())) => return Ok(Async::Ready(())),
Ok(Async::NotReady) => {
self.state = PingListenerState::Closing;
return Ok(Async::NotReady);
},
Err(err) => return Err(err),
}
},
PingListenerState::Poisoned => panic!("Poisoned or errored PingListener"),
}
}
}
}
// Implementation of the `Codec` trait of tokio-io. Splits frames into groups of 32 bytes.
#[derive(Copy, Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash)]
struct Codec;
impl Decoder for Codec {
type Item = BytesMut;
type Error = IoError;
#[inline]
fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<BytesMut>, IoError> {
if buf.len() >= 32 {
Ok(Some(buf.split_to(32)))
} else {
Ok(None)
}
}
}
impl Encoder for Codec {
type Item = Bytes;
type Error = IoError;
#[inline]
fn encode(&mut self, mut data: Bytes, buf: &mut BytesMut) -> Result<(), IoError> {
if !data.is_empty() {
let split = 32 * (1 + ((data.len() - 1) / 32));
buf.reserve(split);
buf.put(data.split_to(split));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
extern crate tokio;
extern crate tokio_tcp;
use self::tokio::runtime::current_thread::Runtime;
use self::tokio_tcp::TcpListener;
use self::tokio_tcp::TcpStream;
use super::{Ping, PingOutput};
use futures::{Future, Stream};
use libp2p_core::{ConnectionUpgrade, Endpoint};
// TODO: rewrite tests with the MemoryTransport
#[test]
fn ping_pong() {
let listener = TcpListener::bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
let listener_addr = listener.local_addr().unwrap();
let server = listener
.incoming()
.into_future()
.map_err(|(e, _)| e.into())
.and_then(|(c, _)| {
Ping::<()>::default().upgrade(
c.unwrap(),
(),
Endpoint::Listener,
)
})
.and_then(|out| match out {
PingOutput::Ponger(service) => service,
_ => unreachable!(),
});
let client = TcpStream::connect(&listener_addr)
.map_err(|e| e.into())
.and_then(|c| {
Ping::<()>::default().upgrade(
c,
(),
Endpoint::Dialer,
)
})
.and_then(|out| match out {
PingOutput::Pinger(mut pinger) => {
pinger.ping(());
pinger.into_future().map(|_| ()).map_err(|_| panic!())
},
_ => unreachable!(),
})
.map(|_| ());
let mut rt = Runtime::new().unwrap();
let _ = rt.block_on(server.select(client).map_err(|_| panic!())).unwrap();
}
#[test]
fn multipings() {
// Check that we can send multiple pings in a row and it will still work.
let listener = TcpListener::bind(&"127.0.0.1:0".parse().unwrap()).unwrap();
let listener_addr = listener.local_addr().unwrap();
let server = listener
.incoming()
.into_future()
.map_err(|(e, _)| e.into())
.and_then(|(c, _)| {
Ping::<u32>::default().upgrade(
c.unwrap(),
(),
Endpoint::Listener,
)
})
.and_then(|out| match out {
PingOutput::Ponger(service) => service,
_ => unreachable!(),
});
let client = TcpStream::connect(&listener_addr)
.map_err(|e| e.into())
.and_then(|c| {
Ping::<u32>::default().upgrade(
c,
(),
Endpoint::Dialer,
)
})
.and_then(|out| match out {
PingOutput::Pinger(mut pinger) => {
for n in 0..20 {
pinger.ping(n);
}
pinger
.take(20)
.collect()
.map(|val| { assert_eq!(val, (0..20).collect::<Vec<_>>()); })
.map_err(|_| panic!())
},
_ => unreachable!(),
});
let mut rt = Runtime::new().unwrap();
let _ = rt.block_on(server.select(client)).unwrap_or_else(|_| panic!());
}
}
mod dial_handler;
mod listen_handler;