secio: Defer nonce verification. (#667)

protocols/secio/src/handshake.rs

@@ -307,10 +307,8 @@ impl HandshakeContext<Ephemeral> {
 /// On success, returns an object that implements the `Sink` and `Stream` trait whose items are
 /// buffers of data, plus the public key of the remote, plus the ephemeral public key used during
 /// negotiation.
-pub fn handshake<'a, S: 'a>(
-    socket: S,
-    config: SecioConfig
-) -> Box<Future<Item = (FullCodec<S>, PublicKey, Vec<u8>), Error = SecioError> + Send + 'a>
+pub fn handshake<'a, S: 'a>(socket: S, config: SecioConfig)
+    -> impl Future<Item = (FullCodec<S>, PublicKey, Vec<u8>), Error = SecioError>
 where
     S: AsyncRead + AsyncWrite + Send,
 {
@@ -320,7 +318,7 @@ where
         .length_field_length(4)
         .new_framed(socket);
 
-    let future = future::ok::<_, SecioError>(HandshakeContext::new(config))
+    future::ok::<_, SecioError>(HandshakeContext::new(config))
         .and_then(|context| {
             // Generate our nonce.
             let context = context.with_local()?;
@@ -570,7 +568,14 @@ where
                 (cipher, hmac)
             };
 
-            let codec = full_codec(socket, encoding_cipher, encoding_hmac, decoding_cipher, decoding_hmac);
+            let codec = full_codec(
+                socket,
+                encoding_cipher,
+                encoding_hmac,
+                decoding_cipher,
+                decoding_hmac,
+                context.state.remote.local.nonce.to_vec()
+            );
             Ok((codec, context))
         })
         // We send back their nonce to check if the connection works.
@@ -578,32 +583,9 @@ where
             let remote_nonce = context.state.remote.nonce.clone();
             trace!("checking encryption by sending back remote's nonce");
             codec.send(BytesMut::from(remote_nonce))
-                .map(|s| (s, context))
+                .map(|s| (s, context.state.remote.public_key, context.state.local_tmp_pub_key))
                 .from_err()
         })
-        // Check that the received nonce is correct.
-        .and_then(|(codec, context)| {
-            codec.into_future()
-                .map_err(|(e, _)| e)
-                .and_then(move |(nonce, rest)| {
-                    match nonce {
-                        Some(ref n) if n == &context.state.remote.local.nonce => {
-                            trace!("secio handshake success");
-                            Ok((rest, context.state.remote.public_key, context.state.local_tmp_pub_key))
-                        },
-                        None => {
-                            debug!("unexpected eof during nonce check");
-                            Err(IoError::new(IoErrorKind::BrokenPipe, "unexpected eof").into())
-                        },
-                        _ => {
-                            debug!("failed nonce verification with remote");
-                            Err(SecioError::NonceVerificationFailed)
-                        }
-                    }
-                })
-        });
-
-    Box::new(future)
 }
 
 /// Custom algorithm translated from reference implementations. Needs to be the same algorithm
@@ -647,15 +629,16 @@ where ::hmac::Hmac<D>: Clone {
 mod tests {
     extern crate tokio;
     extern crate tokio_tcp;
+    use bytes::BytesMut;
     use self::tokio::runtime::current_thread::Runtime;
     use self::tokio_tcp::TcpListener;
     use self::tokio_tcp::TcpStream;
+    use crate::SecioError;
     use super::handshake;
     use super::stretch_key;
     use algo_support::Digest;
     use codec::Hmac;
-    use futures::Future;
-    use futures::Stream;
+    use futures::prelude::*;
     use {SecioConfig, SecioKeyPair};
 
     #[test]
@@ -707,11 +690,29 @@ mod tests {
             .incoming()
             .into_future()
             .map_err(|(e, _)| e.into())
-            .and_then(move |(connec, _)| handshake(connec.unwrap(), key1));
+            .and_then(move |(connec, _)| handshake(connec.unwrap(), key1))
+            .and_then(|(connec, _, _)| {
+                let (sink, stream) = connec.split();
+                stream
+                    .filter(|v| !v.is_empty())
+                    .forward(sink.with(|v| Ok::<_, SecioError>(BytesMut::from(v))))
+            });
 
         let client = TcpStream::connect(&listener_addr)
             .map_err(|e| e.into())
-            .and_then(move |stream| handshake(stream, key2));
+            .and_then(move |stream| handshake(stream, key2))
+            .and_then(|(connec, _, _)| {
+                connec.send("hello".into())
+                    .from_err()
+                    .and_then(|connec| {
+                        connec.filter(|v| !v.is_empty())
+                            .into_future()
+                            .map(|(v, _)| v)
+                            .map_err(|(e, _)| e)
+                    })
+                    .map(|v| assert_eq!(b"hello", &v.unwrap()[..]))
+            });
+
         let mut rt = Runtime::new().unwrap();
         let _ = rt.block_on(server.join(client)).unwrap();
     }