Add a with_simple_protocol_upgrade alternative to with_upgrade

example/examples/echo-dialer.rs

@@ -26,15 +26,11 @@ extern crate libp2p_tcp_transport as tcp;
 extern crate tokio_core;
 extern crate tokio_io;
 
-use bytes::Bytes;
-use futures::future::{Future, FutureResult, IntoFuture};
-use futures::{Stream, Sink};
-use std::io::Error as IoError;
-use std::iter;
-use swarm::{Transport, ConnectionUpgrade};
+use bytes::BytesMut;
+use futures::{Stream, Sink, Future};
+use swarm::Transport;
 use tcp::TcpConfig;
 use tokio_core::reactor::Core;
-use tokio_io::{AsyncRead, AsyncWrite};
 use tokio_io::codec::length_delimited;
 
 fn main() {
@@ -51,11 +47,12 @@ fn main() {
             }
         });
 
-    let with_echo = with_secio.with_upgrade(Echo);
+    let with_echo = with_secio.with_simple_protocol_upgrade("/echo/1.0.0", |socket| {
+        Ok(length_delimited::Framed::<_, BytesMut>::new(socket))
+    });
 
     let dialer = with_echo.dial(swarm::multiaddr::Multiaddr::new("/ip4/127.0.0.1/tcp/10333").unwrap())
-        .map_err(|_| panic!())
-        .unwrap()
+        .unwrap_or_else(|_| panic!())
         .and_then(|f| {
             f.send("hello world".into())
         })
@@ -70,26 +67,3 @@ fn main() {
 
     core.run(dialer).unwrap();
 }
-
-// TODO: copy-pasted from echo-server
-#[derive(Debug, Copy, Clone)]
-pub struct Echo;
-impl<C> ConnectionUpgrade<C> for Echo
-    where C: AsyncRead + AsyncWrite
-{
-    type NamesIter = iter::Once<(Bytes, Self::UpgradeIdentifier)>;
-    type UpgradeIdentifier = ();
-
-    #[inline]
-    fn protocol_names(&self) -> Self::NamesIter {
-        iter::once(("/echo/1.0.0".into(), ()))
-    }
-
-    type Output = length_delimited::Framed<C>;
-    type Future = FutureResult<Self::Output, IoError>;
-
-    #[inline]
-    fn upgrade(self, socket: C, _: Self::UpgradeIdentifier) -> Self::Future {
-        Ok(length_delimited::Framed::new(socket)).into_future()
-    }
-}

example/examples/echo-server.rs

@@ -26,15 +26,14 @@ extern crate libp2p_tcp_transport as tcp;
 extern crate tokio_core;
 extern crate tokio_io;
 
-use bytes::Bytes;
-use futures::future::{Future, FutureResult, IntoFuture, loop_fn, Loop};
+use bytes::BytesMut;
+use futures::future::{Future, IntoFuture, loop_fn, Loop};
 use futures::{Stream, Sink};
 use std::io::Error as IoError;
 use std::iter;
 use swarm::{Transport, ConnectionUpgrade};
 use tcp::TcpConfig;
 use tokio_core::reactor::Core;
-use tokio_io::{AsyncRead, AsyncWrite};
 use tokio_io::codec::length_delimited;
 
 fn main() {
@@ -51,11 +50,12 @@ fn main() {
             }
         });
 
-    let with_echo = with_secio.with_upgrade(Echo);
+    let with_echo = with_secio.with_simple_protocol_upgrade("/echo/1.0.0", |socket| {
+        Ok(length_delimited::Framed::<_, BytesMut>::new(socket))
+    });
 
     let future = with_echo.listen_on(swarm::multiaddr::Multiaddr::new("/ip4/0.0.0.0/tcp/10333").unwrap())
-        .map_err(|_| panic!())
-        .unwrap().0
+        .unwrap_or_else(|_| panic!()).0
         .for_each(|(socket, _)| {
             loop_fn(socket, |socket| {
                 socket.into_future()
@@ -72,26 +72,3 @@ fn main() {
 
     core.run(future).unwrap();
 }
-
-// TODO: copy-pasted from echo-dialer
-#[derive(Debug, Copy, Clone)]
-pub struct Echo;
-impl<C> ConnectionUpgrade<C> for Echo
-    where C: AsyncRead + AsyncWrite
-{
-    type NamesIter = iter::Once<(Bytes, Self::UpgradeIdentifier)>;
-    type UpgradeIdentifier = ();
-
-    #[inline]
-    fn protocol_names(&self) -> Self::NamesIter {
-        iter::once(("/echo/1.0.0".into(), ()))
-    }
-
-    type Output = length_delimited::Framed<C>;
-    type Future = FutureResult<Self::Output, IoError>;
-
-    #[inline]
-    fn upgrade(self, socket: C, _: Self::UpgradeIdentifier) -> Self::Future {
-        Ok(length_delimited::Framed::new(socket)).into_future()
-    }
-}

libp2p-swarm/src/transport.rs

@@ -31,11 +31,12 @@
 
 use bytes::Bytes;
 use futures::{Stream, Poll, Async};
-use futures::future::{IntoFuture, Future, ok as future_ok, FutureResult};
+use futures::future::{IntoFuture, Future, ok as future_ok, FutureResult, FromErr};
 use multiaddr::Multiaddr;
 use multistream_select;
 use std::io::{Cursor, Error as IoError, Read, Write};
 use std::iter;
+use std::sync::Arc;
 use tokio_io::{AsyncRead, AsyncWrite};
 
 /// A transport is an object that can be used to produce connections by listening or dialing a
@@ -101,6 +102,25 @@ pub trait Transport {
 			upgrade: upgrade,
 		}
 	}
+
+	/// Wraps this transport inside an upgrade. Similar to `with_upgrade`, but more convenient to
+	/// use for small protocols.
+	#[inline]
+	fn with_simple_protocol_upgrade<N, F, O>(self, name: N, upgrade: F)
+											 -> UpgradedNode<Self, SimpleProtocolUpgrade<F>>
+		where Self: Sized,
+			  N: Into<Bytes>,
+			  F: Fn(Self::RawConn) -> O,
+			  O: IntoFuture<Error = IoError>,
+	{
+		UpgradedNode {
+			transports: self,
+			upgrade: SimpleProtocolUpgrade {
+				name: name.into(),
+				upgrade: Arc::new(upgrade),
+			},
+		}
+	}
 }
 
 /// Dummy implementation of `Transport` that just denies every single attempt.
@@ -164,6 +184,48 @@ impl<A, B> Transport for OrTransport<A, B>
 	}
 }
 
+/// Implementation of `ConnetionUpgrade`. See `with_simple_protocol_upgrade`.
+#[derive(Debug)]
+pub struct SimpleProtocolUpgrade<F> {
+	name: Bytes,
+	// Note: we put the closure `F` in an `Arc` because Rust closures aren't automatically clonable
+	// yet.
+	upgrade: Arc<F>,
+}
+
+impl<F> Clone for SimpleProtocolUpgrade<F> {
+	#[inline]
+	fn clone(&self) -> Self {
+		SimpleProtocolUpgrade {
+			name: self.name.clone(),
+			upgrade: self.upgrade.clone(),
+		}
+	}
+}
+
+impl<C, F, O> ConnectionUpgrade<C> for SimpleProtocolUpgrade<F>
+    where C: AsyncRead + AsyncWrite,
+		  F: Fn(C) -> O,
+		  O: IntoFuture<Error = IoError>,
+{
+    type NamesIter = iter::Once<(Bytes, ())>;
+    type UpgradeIdentifier = ();
+
+    #[inline]
+    fn protocol_names(&self) -> Self::NamesIter {
+        iter::once(("/echo/1.0.0".into(), ()))
+    }
+
+    type Output = O::Item;
+    type Future = FromErr<O::Future, IoError>;
+
+    #[inline]
+    fn upgrade(self, socket: C, _: ()) -> Self::Future {
+		let upgrade = &self.upgrade;
+		upgrade(socket).into_future().from_err()
+    }
+}
+
 /// Implements `Stream` and dispatches all method calls to either `First` or `Second`.
 #[derive(Debug, Copy, Clone)]
 pub enum EitherStream<A, B> {