number handshakes

This prevents late hellos from resetting the session.

2 files changed, 32 insertions, 7 deletions

diff --git a/src/swarm/peer/handshake.rs b/src/swarm/peer/handshake.rs
index 038dac6..f89e371 100644
--- a/src/swarm/peer/handshake.rs
+++ b/src/swarm/peer/handshake.rs
@@ -1,6 +1,7 @@
 //! Discovering peers, the handshake (HELLO / HELLO_REPLY), and membership.
 
 use std::sync::Arc;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 
 use anyhow::{Context, Result, anyhow, bail};
@@ -67,6 +68,9 @@ pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()>
         let mut peer_state = peer_state_mutex.lock().await;
         if peer_state.hello_state == HelloState::NotSent {
             // A new handshake: retries start over.
+            // Numbered across all peers, so a peer we forgot and meet again still
+            // sees a higher number than any it had from this run.
+            peer_state.handshake = state.handshakes.fetch_add(1, Ordering::Relaxed) + 1;
             peer_state.hello_state = HelloState::Sending;
             peer_state.hello_backoff.reset();
             peer_state.hello_attempts = 0;
@@ -115,6 +119,11 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
         .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
         .unwrap_or_default();
     let (epoch, epoch_seq) = state.epoch_stamp().await;
+    // Retries repeat the handshake's number.
+    let handshake = match state.get_peer_state(peer).await {
+        Some(peer_state) => peer_state.lock().await.handshake,
+        None => 0,
+    };
     state.lock_rtt().hello_sent(peer, Instant::now());
     let result = state
         .transport
@@ -127,6 +136,7 @@ async fn transmit_hello(state: &AppState, peer: &str) -> Result<()> {
                 auth,
                 ready: state.is_ready().await,
                 run_id: state.run_id,
+                handshake,
             }),
         )
         .await;
@@ -293,10 +303,11 @@ pub(crate) async fn apply_neighbour_change(state: &AppState, change: topology::N
 
 // --- peer membership ---------------------------------------------------------
 
-/// Makes `peer` a member of `swarm_id`, after its valid HELLO (`from_hello`) or
-/// HELLO_REPLY. Either completes the handshake: one round trip proves both sides to
-/// each other (the HELLO the sender's key, the reply the swarm ID and the responder's
-/// key), so neither side needs a HELLO of its own. `ready` is what the message said.
+/// Makes `peer` a member of `swarm_id`, after its valid HELLO or HELLO_REPLY. Either
+/// completes the handshake: one round trip proves both sides to each other (the HELLO
+/// the sender's key, the reply the swarm ID and the responder's key), so neither side
+/// needs a HELLO of its own. `from_hello`: a HELLO that starts a new handshake (a resent
+/// or late one doesn't). `ready` is what the message said.
 pub(super) async fn establish_peer(
     state: &AppState,
     peer: &str,
@@ -340,8 +351,9 @@ pub(super) async fn establish_peer(
         };
     }
 
-    // Stop our retries and sync from a clean slate. A HELLO always resets: the peer
-    // starts over with us. Duplicate replies (caused by our own retries) don't.
+    // Stop our retries and sync from a clean slate. A new handshake from the peer always
+    // resets: it starts over with us. Its resent HELLOs and duplicate replies (caused by
+    // our own retries) don't.
     peer_state.reports_ready = Some(ready);
     if from_hello || peer_state.hello_state != HelloState::Established {
         peer_state.hello_state = HelloState::Established;
diff --git a/src/swarm/peer/inbound.rs b/src/swarm/peer/inbound.rs
index 040cafb..684a358 100644
--- a/src/swarm/peer/inbound.rs
+++ b/src/swarm/peer/inbound.rs
@@ -84,7 +84,20 @@ pub(crate) async fn handle_message(
                     bail!("Incompatible swarm at {peer}: ours is {swarm_id}, it claims {theirs}");
                 }
             }
-            establish_peer(state, &peer, swarm_id, true, hello.ready).await?;
+            // A resent HELLO (our reply was lost) or a late copy of one repeats a
+            // handshake we answered already: it must not reset the session again.
+            let new_handshake = {
+                let peer_state = state.get_or_create_peer_state(&peer).await;
+                let mut peer_state = peer_state.lock().await;
+                let new = peer_state.last_hello.is_none_or(|(run, handshake)| {
+                    run != hello.run_id || hello.handshake > handshake
+                });
+                if new {
+                    peer_state.last_hello = Some((hello.run_id, hello.handshake));
+                }
+                new
+            };
+            establish_peer(state, &peer, swarm_id, new_handshake, hello.ready).await?;
 
             let away = previously_seen.map(|seen| seen.elapsed());
             if let Err(error) = send_hello_reply(state, &peer, hello.run_id, away).await {