Merge of #3522

ZcashFoundation · Feb 15, 2022 · 016f4ad · 016f4ad
2 parents 1913df0 + dc6679b
commit 016f4ad
Showing 1 changed file with 36 additions and 8 deletions.
diff --git a/zebra-network/src/peer/handshake.rs b/zebra-network/src/peer/handshake.rs
@@ -588,13 +588,29 @@ where
     debug!(?our_version, "sending initial version message");
     peer_conn.send(our_version).await?;
 
-    let remote_msg = peer_conn
+    let mut remote_msg = peer_conn
         .next()
         .await
         .ok_or(HandshakeError::ConnectionClosed)??;
 
-    // Check that we got a Version and destructure its fields into the local scope.
-    debug!(?remote_msg, "got message from remote peer");
+    // Wait for next message if the one we got is not Version
+    loop {
+        match remote_msg {
+            Message::Version { .. } => {
+                debug!(?remote_msg, "got version message from remote peer");
+                break;
+            }
+            _ => {
+                remote_msg = peer_conn
+                    .next()
+                    .await
+                    .ok_or(HandshakeError::ConnectionClosed)??;
+                debug!(?remote_msg, "ignoring non-version message from remote peer");
+            }
+        }
+    }
+
+    // If we got a Version message, destructure its fields into the local scope.
     let (remote_nonce, remote_services, remote_version, remote_canonical_addr, user_agent) =
         if let Message::Version {
             version,
@@ -700,14 +716,26 @@ where
 
     peer_conn.send(Message::Verack).await?;
 
-    let remote_msg = peer_conn
+    let mut remote_msg = peer_conn
         .next()
         .await
         .ok_or(HandshakeError::ConnectionClosed)??;
-    if let Message::Verack = remote_msg {
-        debug!("got verack from remote peer");
-    } else {
-        Err(HandshakeError::UnexpectedMessage(Box::new(remote_msg)))?;
+
+    // Wait for next message if the one we got is not Verack
+    loop {
+        match remote_msg {
+            Message::Verack => {
+                debug!(?remote_msg, "got verack message from remote peer");
+                break;
+            }
+            _ => {
+                remote_msg = peer_conn
+                    .next()
+                    .await
+                    .ok_or(HandshakeError::ConnectionClosed)??;
+                debug!(?remote_msg, "ignoring non-verack message from remote peer");
+            }
+        }
     }
 
     Ok((remote_version, remote_services, remote_canonical_addr))