perform handshake in one round trip

1 files changed, 52 insertions, 62 deletions

diff --git a/src/swarm/peer.rs b/src/swarm/peer.rs
index 4c07d7d..bf552e9 100644
--- a/src/swarm/peer.rs
+++ b/src/swarm/peer.rs
@@ -253,17 +253,20 @@ async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> {
 /// Success and failure both end in AwaitingReply: Mycelium accepting the message
 /// says nothing about whether the peer received it.
 async fn transmit_hello(state: &AppState, peer: &str, delay: Duration) -> Result<()> {
+    let swarm_id = state.swarm_id().await;
+    let auth = state
+        .swarm_key
+        .as_ref()
+        .map(|key| key.hello_tag(swarm_id.as_deref().unwrap_or(""), state.local_id(), peer))
+        .unwrap_or_default();
     let result = state
         .transport
         .send_message(
             peer,
             &ParsedMessage::Hello(transport::HelloRequest {
+                swarm_id,
                 epoch: state.local_epoch(),
-                auth: state
-                    .swarm_key
-                    .as_ref()
-                    .map(|key| key.hello_tag(state.local_id(), peer))
-                    .unwrap_or_default(),
+                auth,
             }),
         )
         .await;
@@ -294,35 +297,6 @@ async fn ensure_handshake(state: &AppState, peer: &str) {
     }
 }
 
-/// Sends a HELLO_REPLY if one is owed and our swarm ID is known.
-/// Otherwise it stays owed and the peer worker tries again.
-async fn send_owed_hello_reply(state: &AppState, peer: &str) {
-    if state.swarm_id().await.is_none() {
-        return;
-    }
-
-    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
-        return;
-    };
-
-    let owed = std::mem::take(&mut peer_state_mutex.lock().await.reply_owed);
-    if !owed {
-        return;
-    }
-
-    if let Err(error) = send_hello_reply(state, peer).await {
-        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed; will retry");
-        peer_state_mutex.lock().await.reply_owed = true;
-    }
-}
-
-async fn flush_owed_hello_replies(state: &AppState) {
-    let peers: Vec<String> = state.peer_states.read().await.keys().cloned().collect();
-    for peer in &peers {
-        send_owed_hello_reply(state, peer).await;
-    }
-}
-
 async fn send_hello_reply(state: &AppState, peer: &str) -> Result<()> {
     let swarm_id = state
         .swarm_id()
@@ -460,8 +434,6 @@ pub(crate) async fn run_peer_worker(state: AppState) {
             if let Err(error) = retry_hello_if_due(&state, peer).await {
                 warn!(peer = %peer, error = %format!("{error:#}"), "HELLO retry failed");
             }
-
-            send_owed_hello_reply(&state, peer).await;
         }
 
         if last_forget_check.is_none_or(|t| t.elapsed() >= FORGET_CHECK_INTERVAL) {
@@ -473,7 +445,16 @@ pub(crate) async fn run_peer_worker(state: AppState) {
 
 // --- peer membership ---------------------------------------------------------
 
-async fn establish_peer(state: &AppState, peer: &str, swarm_id: String) -> Result<()> {
+/// Makes `peer` a member of `swarm_id`, after its valid HELLO (`cause: PeerHello`) or
+/// HELLO_REPLY (`HandshakeComplete`). 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.
+async fn establish_peer(
+    state: &AppState,
+    peer: &str,
+    swarm_id: String,
+    cause: SessionCause,
+) -> Result<()> {
     let expected_swarm = state.swarm_id().await;
 
     if let Some(expected) = &expected_swarm {
@@ -510,12 +491,12 @@ async fn establish_peer(state: &AppState, peer: &str, swarm_id: String) -> Resul
         };
     }
 
-    // First reply completes the handshake: stop retries and sync from a clean slate.
-    // Duplicate replies (caused by our own retries) don't reset again.
-    if peer_state.hello_state != HelloState::Established {
+    // 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.
+    if cause == SessionCause::PeerHello || peer_state.hello_state != HelloState::Established {
         peer_state.hello_state = HelloState::Established;
-        peer_state.reset_sync(SessionCause::HandshakeComplete);
-        info!(peer = %peer, "Handshake complete; sync state reset");
+        peer_state.reset_sync(cause);
+        info!(peer = %peer, ?cause, "Handshake complete; sync state reset");
     }
 
     info!(peer = %peer, swarm_id = %swarm_id, "Peer membership established");
@@ -604,7 +585,12 @@ async fn admit(state: &AppState, peer: &str, message: &ParsedMessage) -> bool {
     let local = state.local_id();
 
     let proof = match message {
-        ParsedMessage::Hello(hello) => Some(key.verify_hello(&hello.auth, peer, local)),
+        ParsedMessage::Hello(hello) => Some(key.verify_hello(
+            &hello.auth,
+            hello.swarm_id.as_deref().unwrap_or(""),
+            peer,
+            local,
+        )),
         ParsedMessage::HelloReply(reply) => {
             Some(key.verify_hello_reply(&reply.auth, &reply.swarm_id, peer, local))
         }
@@ -696,37 +682,41 @@ pub(crate) async fn handle_message(
     //debug!(from = %peer, message = ?message, "Received message");
 
     match message {
-        transport::ParsedMessage::Hello(_) => {
+        transport::ParsedMessage::Hello(hello) => {
             // A HELLO means the peer is (re)starting its handshake with us, e.g. after a
-            // restart. Whatever our sync state believed about it is stale.
-            {
-                let peer_state_mutex = state.get_or_create_peer_state(&peer).await;
-                let mut peer_state = peer_state_mutex.lock().await;
-                peer_state.reset_sync(SessionCause::PeerHello);
-                peer_state.reply_owed = true;
+            // restart. It becomes a member right away; our reply makes us one for it.
+            let Some(swarm_id) = state.swarm_id().await else {
+                // We are joining ourselves; the peer retries until we know our swarm.
+                debug!(peer = %peer, "HELLO before we know our swarm; not answering yet");
+                return Ok(());
+            };
+            if let Some(theirs) = &hello.swarm_id {
+                if theirs != &swarm_id {
+                    bail!("Incompatible swarm at {peer}: ours is {swarm_id}, it claims {theirs}");
+                }
             }
+            establish_peer(state, &peer, swarm_id, SessionCause::PeerHello).await?;
 
-            // Sends now if we know our swarm ID; otherwise stays owed.
-            send_owed_hello_reply(state, &peer).await;
-
-            // Our own HELLO goes out regardless of whether the reply could be sent.
-            if let Err(error) = send_hello_once(state, &peer).await {
-                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
+            if let Err(error) = send_hello_reply(state, &peer).await {
+                // The peer retries its HELLO.
+                warn!(peer = %peer, error = %format!("{error:#}"), "HELLO_REPLY failed");
             }
 
-            // Restart sync with the fresh state (no-op until the peer is a validated member).
             let _ = state
                 .sync_events
-                .send(SyncEvent::PeerNeedsSync(peer.clone()));
+                .send(SyncEvent::PeerValidated(peer.clone()));
         }
         transport::ParsedMessage::HelloReply(reply) => {
-            establish_peer(state, &peer, reply.swarm_id.clone()).await?;
+            establish_peer(
+                state,
+                &peer,
+                reply.swarm_id.clone(),
+                SessionCause::HandshakeComplete,
+            )
+            .await?;
 
             discover_peers(state, &peer, &reply.peers).await;
 
-            // Our swarm ID may have just become known: pay any replies we owe.
-            flush_owed_hello_replies(state).await;
-
             state
                 .sync_events
                 .send(SyncEvent::PeerValidated(peer.clone()))