Wrap multistream-select streams under a Negotiated (#1001)

misc/multistream-select/src/dialer_select.rs

@@ -31,7 +31,7 @@ use crate::protocol::{
 use log::trace;
 use std::mem;
 use tokio_io::{AsyncRead, AsyncWrite};
-use crate::ProtocolChoiceError;
+use crate::{Negotiated, ProtocolChoiceError};
 
 /// Future, returned by `dialer_select_proto`, which selects a protocol and dialer
 /// either sequentially of by considering all protocols in parallel.
@@ -125,7 +125,7 @@ where
     I: Iterator,
     I::Item: AsRef<[u8]> + Clone
 {
-    type Item = (I::Item, R);
+    type Item = (I::Item, Negotiated<R>);
     type Error = ProtocolChoiceError;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
@@ -207,7 +207,7 @@ where
                         ListenerToDialerMessage::ProtocolAck { ref name }
                             if name.as_ref() == proto_name.as_ref() =>
                         {
-                            return Ok(Async::Ready((proto_name, r.into_inner())))
+                            return Ok(Async::Ready((proto_name, Negotiated(r.into_inner()))))
                         }
                         ListenerToDialerMessage::NotAvailable => {
                             let proto_name = protocols.next()
@@ -300,7 +300,7 @@ where
     I: Iterator,
     I::Item: AsRef<[u8]> + Clone
 {
-    type Item = (I::Item, R);
+    type Item = (I::Item, Negotiated<R>);
     type Error = ProtocolChoiceError;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
@@ -423,7 +423,7 @@ where
                         Some(ListenerToDialerMessage::ProtocolAck { ref name })
                             if name.as_ref() == proto_name.as_ref() =>
                         {
-                            return Ok(Async::Ready((proto_name, dialer.into_inner())))
+                            return Ok(Async::Ready((proto_name, Negotiated(dialer.into_inner()))))
                         }
                         _ => return Err(ProtocolChoiceError::UnexpectedMessage)
                     }

misc/multistream-select/src/lib.rs

@@ -74,6 +74,49 @@ mod tests;
 
 mod protocol;
 
+use futures::prelude::*;
+use std::io;
+
 pub use self::dialer_select::{dialer_select_proto, DialerSelectFuture};
 pub use self::error::ProtocolChoiceError;
 pub use self::listener_select::{listener_select_proto, ListenerSelectFuture};
+
+/// A stream after it has been negotiated.
+pub struct Negotiated<TInner>(pub(crate) TInner);
+
+impl<TInner> io::Read for Negotiated<TInner>
+where
+    TInner: io::Read
+{
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        self.0.read(buf)
+    }
+}
+
+impl<TInner> tokio_io::AsyncRead for Negotiated<TInner>
+where
+    TInner: tokio_io::AsyncRead
+{
+}
+
+impl<TInner> io::Write for Negotiated<TInner>
+where
+    TInner: io::Write
+{
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.0.write(buf)
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        self.0.flush()
+    }
+}
+
+impl<TInner> tokio_io::AsyncWrite for Negotiated<TInner>
+where
+    TInner: tokio_io::AsyncWrite
+{
+    fn shutdown(&mut self) -> Poll<(), io::Error> {
+        self.0.shutdown()
+    }
+}

misc/multistream-select/src/listener_select.rs

@@ -31,7 +31,7 @@ use crate::protocol::{
 use log::{debug, trace};
 use std::mem;
 use tokio_io::{AsyncRead, AsyncWrite};
-use crate::ProtocolChoiceError;
+use crate::{Negotiated, ProtocolChoiceError};
 
 /// Helps selecting a protocol amongst the ones supported.
 ///
@@ -99,7 +99,7 @@ where
     for<'a> &'a I: IntoIterator<Item = X>,
     X: AsRef<[u8]> + Clone
 {
-    type Item = (X, R, I);
+    type Item = (X, Negotiated<R>, I);
     type Error = ProtocolChoiceError;
 
     fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
@@ -171,7 +171,7 @@ where
                         }
                     };
                     if let Some(p) = outcome {
-                        return Ok(Async::Ready((p, listener.into_inner(), protocols)))
+                        return Ok(Async::Ready((p, Negotiated(listener.into_inner()), protocols)))
                     } else {
                         let stream = listener.into_future();
                         self.inner = ListenerSelectState::Incoming { stream, protocols }