1 files changed, 336 insertions, 0 deletions
diff --git a/src/swarm/peer/handshake.rs b/src/swarm/peer/handshake.rs
new file mode 100644
index 0000000..c46df9c
--- /dev/null
+++ b/src/swarm/peer/handshake.rs
@@ -0,0 +1,336 @@
+//! Discovering peers, the handshake (HELLO / HELLO_REPLY), and membership.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use anyhow::{Context, Result, anyhow, bail};
+use tokio::sync::Mutex;
+use tokio::time::Instant;
+use tracing::{debug, info, instrument, trace, warn};
+
+use crate::state::{AppState, HelloState, Liveness, PeerState, Standing};
+use crate::swarm::sync::SyncEvent;
+use crate::swarm::topology;
+use crate::transport::{self, ParsedMessage};
+
+pub(crate) async fn remember_peer(state: &AppState, peer: &str) {
+    if peer.is_empty() {
+        warn!("Ignoring empty peer public key");
+        return;
+    }
+
+    if peer == state.local_id() {
+        debug!("Ignoring own public key");
+        return;
+    }
+
+    let mut peer_states = state.peer_states.write().await;
+
+    if !peer_states.contains_key(peer) {
+        peer_states.insert(peer.to_owned(), state.new_peer_state(peer));
+        info!(
+            peer = %peer,
+            peer_count = peer_states.len(),
+            "Discovered peer"
+        );
+    } else {
+        trace!(peer = %peer, "Peer already registered");
+    }
+}
+
+/// Starts a new handshake with `peer` even though one is complete: its reply resets
+/// sync on both sides and brings its current peer list (`swarm::source`, after giving
+/// it up as the catch-up source).
+pub(crate) async fn rehandshake(state: &AppState, peer: &str) {
+    if let Some(peer_state) = state.get_peer_state(peer).await {
+        let mut peer_state = peer_state.lock().await;
+        if peer_state.hello_state == HelloState::Established {
+            peer_state.hello_state = HelloState::NotSent;
+        }
+    }
+    ensure_handshake(state, peer).await;
+}
+
+pub(crate) async fn send_hello_once(state: &AppState, peer: &str) -> Result<()> {
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        trace!(peer = %peer, "Peer was forgotten; not sending HELLO");
+        return Ok(());
+    };
+
+    let claimed = {
+        let mut peer_state = peer_state_mutex.lock().await;
+        if peer_state.hello_state == HelloState::NotSent {
+            // A new handshake: retries start over.
+            peer_state.hello_state = HelloState::Sending;
+            peer_state.hello_backoff.reset();
+            peer_state.hello_attempts = 0;
+            true
+        } else {
+            trace!(peer = %peer, hello_state = ?peer_state.hello_state, "not sending");
+            false
+        }
+    };
+
+    if !claimed {
+        return Ok(());
+    }
+
+    debug!(to = %peer, "Sending HELLO");
+    transmit_hello(state, peer).await
+}
+
+pub(super) async fn retry_hello_if_due(state: &AppState, peer: &str) -> Result<()> {
+    let Some(peer_state_mutex) = state.get_peer_state(peer).await else {
+        return Ok(());
+    };
+
+    {
+        let mut peer_state = peer_state_mutex.lock().await;
+        match peer_state.hello_state {
+            HelloState::AwaitingReply { next_attempt } if Instant::now() >= next_attempt => {
+                peer_state.hello_state = HelloState::Sending;
+            }
+            _ => return Ok(()),
+        }
+    }
+
+    debug!(peer = %peer, "No HELLO_REPLY yet; re-sending HELLO");
+    transmit_hello(state, peer).await
+}
+
+/// Sends a HELLO without holding any lock, then schedules the next retry.
+/// 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) -> 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 (epoch, epoch_seq) = state.epoch_stamp().await;
+    state.lock_rtt().hello_sent(peer, Instant::now());
+    let result = state
+        .transport
+        .send_message(
+            peer,
+            &ParsedMessage::Hello(transport::HelloRequest {
+                swarm_id,
+                epoch,
+                epoch_seq,
+                auth,
+                ready: state.is_ready().await,
+                run_id: state.run_id,
+            }),
+        )
+        .await;
+
+    if let Some(peer_state_mutex) = state.get_peer_state(peer).await {
+        let mut peer_state = peer_state_mutex.lock().await;
+
+        if peer_state.hello_state == HelloState::Sending {
+            let was_unreachable = peer_state.unreachable_since.is_some();
+            peer_state.hello_sent(Instant::now());
+            if !was_unreachable && peer_state.unreachable_since.is_some() {
+                state.topology_changed.notify_one();
+            }
+            let wait = match peer_state.liveness {
+                Liveness::Dormant { .. } => state.timing.dormant_probe,
+                Liveness::Active => peer_state.hello_backoff.next(),
+            };
+            peer_state.hello_state = HelloState::AwaitingReply {
+                next_attempt: Instant::now() + wait,
+            };
+        }
+    }
+
+    result.context("Sending HELLO")
+}
+
+pub(super) async fn ensure_handshake(state: &AppState, peer: &str) {
+    if let Err(error) = send_hello_once(state, peer).await {
+        warn!(peer = %peer, error = %format!("{error:#}"), "HELLO failed; will retry");
+    }
+}
+
+/// Answers a HELLO that carried `run_id`; `away` is how long before it we had last heard
+/// from the peer.
+pub(super) async fn send_hello_reply(
+    state: &AppState,
+    peer: &str,
+    run_id: transport::RunId,
+    away: Option<Duration>,
+) -> Result<()> {
+    let swarm_id = state
+        .swarm_id()
+        .await
+        .ok_or_else(|| anyhow!("Cannot send HELLO_REPLY without an established swarm ID"))?;
+
+    let candidates: Vec<(String, Arc<Mutex<PeerState>>)> = state
+        .peers_snapshot()
+        .await
+        .into_iter()
+        .filter(|(known, _)| known.as_str() != peer)
+        .collect();
+
+    let mut peers = Vec::new();
+    for (known, peer_state_mutex) in candidates {
+        let peer_state = peer_state_mutex.lock().await;
+        if peer_state.liveness == Liveness::Active
+            && peer_state.hello_state == HelloState::Established
+            && topology::is_neighbour(state, &known)
+        {
+            peers.push(known);
+        }
+    }
+
+    let auth = state
+        .swarm_key
+        .as_ref()
+        .map(|key| key.hello_reply_tag(&swarm_id, state.local_id(), peer))
+        .unwrap_or_default();
+    let ready = state.is_ready().await;
+    let (epoch, epoch_seq) = state.epoch_stamp().await;
+    let reply = transport::HelloReply {
+        swarm_id,
+        peers,
+        epoch,
+        epoch_seq,
+        auth,
+        ready,
+        run_id,
+        away_ms: away.map(|away| away.as_millis() as u64),
+        sender_run_id: state.run_id,
+    };
+
+    debug!(to = %peer, ready, "Sending HELLO_REPLY");
+
+    state
+        .transport
+        .send_message(peer, &ParsedMessage::HelloReply(reply))
+        .await
+        .context("Sending HELLO_REPLY")?;
+
+    trace!(to = %peer, "HELLO_REPLY sent");
+    Ok(())
+}
+
+pub(super) async fn discover_peers(state: &AppState, from_peer: &str, advertised_peers: &[String]) {
+    // TODO here we could add a trace t oshwo advertised_peers
+    for discovered_peer in advertised_peers {
+        if discovered_peer == from_peer {
+            continue;
+        }
+
+        remember_peer(state, discovered_peer).await;
+
+        if !topology::is_neighbour(state, discovered_peer) {
+            continue;
+        }
+        if let Err(error) = send_hello_once(state, discovered_peer).await {
+            warn!(
+                peer = %discovered_peer,
+                error = %format!("{error:#}"),
+                "Failed to initiate HELLO with discovered peer"
+            );
+        }
+    }
+}
+
+/// New neighbours get our HELLO (we may not have known them before: the member list is
+/// how a partial mesh discovers them) and a sync opening. The opening matters for a
+/// neighbour whose handshake is already complete, e.g. one that returned from dormancy
+/// while it wasn't our neighbour: it was only answered then, and with no HELLO to send,
+/// nothing else would open sync until the keepalive heal (`resync_after`). Former
+/// neighbours get no more HELLOs, keepalives or sync openings from us; their session is
+/// reset, so we don't keep pushing changes to them either. They are still answered if
+/// they contact us.
+pub(crate) async fn apply_neighbour_change(state: &AppState, change: topology::NeighbourChange) {
+    for peer in &change.added {
+        remember_peer(state, peer).await;
+        if let Err(error) = send_hello_once(state, peer).await {
+            warn!(peer = %peer, error = %format!("{error:#}"), "HELLO to new neighbour failed; will retry");
+        }
+        // Skipped there unless the handshake is complete; otherwise its reply restarts
+        // sync anyway.
+        let _ = state
+            .sync_events
+            .send(SyncEvent::PeerNeedsSync(peer.clone()));
+    }
+    for peer in &change.removed {
+        if let Some(peer_state) = state.get_peer_state(peer).await {
+            let mut peer_state = peer_state.lock().await;
+            peer_state.stop_hello_retries();
+            peer_state.reset_sync();
+        }
+    }
+}
+
+// --- 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.
+pub(super) async fn establish_peer(
+    state: &AppState,
+    peer: &str,
+    swarm_id: String,
+    from_hello: bool,
+    ready: bool,
+) -> Result<()> {
+    let expected_swarm = state.swarm_id().await;
+
+    if let Some(expected) = &expected_swarm {
+        if expected != &swarm_id {
+            bail!(
+                "Incompatible swarm at {peer}: expected {}, received {}",
+                expected,
+                swarm_id
+            );
+        }
+    }
+
+    let peer_state_mutex = state.get_or_create_peer_state(peer).await;
+    let mut peer_state = peer_state_mutex.lock().await;
+
+    if let Standing::Member { swarm_id: existing } = &peer_state.standing {
+        if existing != &swarm_id {
+            bail!(
+                "Peer {peer} changed swarm identity from {} to {}",
+                existing,
+                swarm_id
+            );
+        }
+    }
+
+    // Lock order: peer mutex -> lifecycle.
+    if expected_swarm.is_none() {
+        state.set_expected_swarm_id(swarm_id.clone()).await?;
+    }
+
+    if !matches!(peer_state.standing, Standing::Member { .. }) {
+        peer_state.standing = Standing::Member {
+            swarm_id: swarm_id.clone(),
+        };
+    }
+
+    // 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.
+    peer_state.reports_ready = Some(ready);
+    if from_hello || peer_state.hello_state != HelloState::Established {
+        peer_state.hello_state = HelloState::Established;
+        peer_state.reset_sync();
+        info!(peer = %peer, from_hello, "Handshake complete; sync state reset");
+    }
+
+    info!(peer = %peer, swarm_id = %swarm_id, "Peer membership established");
+    Ok(())
+}
+
+#[instrument(skip(state), fields(peer = %peer))]
+pub(crate) async fn peer_swarm(state: &AppState, peer: &str) -> Option<String> {
+    let peer_state_mutex = state.get_peer_state(peer).await?;
+    let peer_state = peer_state_mutex.lock().await;
+    peer_state.standing.swarm_id().map(str::to_owned)
+}