[plaintext] Retain remaining read buffer after handshake.

protocols/plaintext/src/handshake.rs

@@ -22,7 +22,7 @@ use crate::PlainText2Config;
 use crate::error::PlainTextError;
 use crate::structs_proto::Exchange;
 
-use bytes::BytesMut;
+use bytes::{Bytes, BytesMut};
 use futures::prelude::*;
 use futures_codec::Framed;
 use libp2p_core::{PublicKey, PeerId};
@@ -111,7 +111,7 @@ impl HandshakeContext<Local> {
 }
 
 pub async fn handshake<S>(socket: S, config: PlainText2Config)
-    -> Result<(S, Remote), PlainTextError>
+    -> Result<(S, Remote, Bytes), PlainTextError>
 where
     S: AsyncRead + AsyncWrite + Send + Unpin,
 {
@@ -134,6 +134,12 @@ where
         }
     };
 
+    // The `Framed` wrapper may have buffered additional data that
+    // was already received but is no longer part of the plaintext
+    // handshake. We need to capture that data before dropping
+    // the `Framed` wrapper via `Framed::into_inner()`.
+    let read_buffer = framed_socket.read_buffer().clone().freeze();
+
     trace!("received exchange from remote; pubkey = {:?}", context.state.public_key);
-    Ok((framed_socket.into_inner(), context.state))
+    Ok((framed_socket.into_inner(), context.state, read_buffer))
 }

protocols/plaintext/src/lib.rs

@@ -20,6 +20,7 @@
 
 use crate::error::PlainTextError;
 
+use bytes::Bytes;
 use futures::future::{self, Ready};
 use futures::prelude::*;
 use futures::future::BoxFuture;
@@ -148,14 +149,15 @@ impl PlainText2Config {
         T: AsyncRead + AsyncWrite + Send + Unpin + 'static
     {
         debug!("Starting plaintext handshake.");
-        let (socket, remote) = handshake::handshake(socket, self).await?;
+        let (socket, remote, read_buffer) = handshake::handshake(socket, self).await?;
         debug!("Finished plaintext handshake.");
 
         Ok((
             remote.peer_id,
             PlainTextOutput {
                 socket,
                 remote_key: remote.public_key,
+                read_buffer,
             }
         ))
     }
@@ -170,12 +172,22 @@ where
     pub socket: S,
     /// The public key of the remote.
     pub remote_key: PublicKey,
+    /// Remaining bytes that have been already buffered
+    /// during the handshake but are not part of the
+    /// handshake. These must be consumed first by `poll_read`.
+    read_buffer: Bytes,
 }
 
 impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for PlainTextOutput<S> {
     fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut [u8])
         -> Poll<Result<usize, io::Error>>
     {
+        if !self.read_buffer.is_empty() {
+            let n = std::cmp::min(buf.len(), self.read_buffer.len());
+            let b = self.read_buffer.split_to(n);
+            buf[..n].copy_from_slice(&b[..]);
+            return Poll::Ready(Ok(n))
+        }
         AsyncRead::poll_read(Pin::new(&mut self.socket), cx, buf)
     }
 }