From bd5e3b9fb6b2e7bb0032ec9dd819d98d96837f85 Mon Sep 17 00:00:00 2001
From: riccardom <riccardomanfrin@gmail.com>
Date: Tue, 20 Oct 2026 12:26:35 +0200
Subject: [PATCH] [client] Run per-map bookkeeping once instead of on every
 pass

Offline peers, SSH host keys, the SSH client config and the SSH server auth
ran on every bounded pass over the whole map. Each pass rewrote the SSH
client config file, and recording host keys for peers not added yet logged
a warning per missing peer. Offline peers now apply with the wholesale
sections on the first pass, and the SSH steps on the converging pass, once
every peer is in the store. The per-pass peer diff stays: it is how each
pass finds the next batch.
---
 client/internal/engine.go      | 30 +++++++++++++--------------
 client/internal/engine_test.go | 38 ++++++++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/client/internal/engine.go b/client/internal/engine.go
index d73beb400..d7a8de4d9 100644
--- a/client/internal/engine.go
+++ b/client/internal/engine.go
@@ -1709,14 +1709,14 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap, maxBatch int,
 	// remaining peer batches they are skipped.
 	if firstPass {
 		e.applyWholesale(networkMap, serial)
+
+		doneOffline := e.phase("offline_peers")
+		e.updateOfflinePeers(networkMap.GetOfflinePeers())
+		doneOffline()
 	}
 
 	log.Debugf("got peers update from Management Service, total peers to connect to = %d", len(networkMap.GetRemotePeers()))
 
-	doneOffline := e.phase("offline_peers")
-	e.updateOfflinePeers(networkMap.GetOfflinePeers())
-	doneOffline()
-
 	// Filter out own peer from the remote peers list
 	localPubKey := e.config.WgPrivateKey.PublicKey().String()
 	remotePeers := make([]*mgmProto.RemotePeerConfig, 0, len(networkMap.GetRemotePeers()))
@@ -1756,18 +1756,18 @@ func (e *Engine) updateNetworkMap(networkMap *mgmProto.NetworkMap, maxBatch int,
 
 	e.statusRecorder.FinishPeerListModifications()
 
-	e.updatePeerSSHHostKeys(remotePeers)
-
-	if err := e.updateSSHClientConfig(remotePeers); err != nil {
-		log.Warnf("failed to update SSH client config: %v", err)
-	}
-
-	e.updateSSHServerAuth(networkMap.GetSshAuth())
-
-	// Set the exclude list only once peers have fully converged (this pass added
-	// the last batch). It needs all target peers present in the store, and
-	// ExcludePeer has replace-semantics — a partial set mid-convergence would be wrong.
+	// The SSH and lazy-exclude steps walk the whole map and need every target peer
+	// in the store, so they run once, on the pass that added the last batch. The
+	// exclude list also has replace-semantics: a partial set would be wrong.
 	if !needMore {
+		e.updatePeerSSHHostKeys(remotePeers)
+
+		if err := e.updateSSHClientConfig(remotePeers); err != nil {
+			log.Warnf("failed to update SSH client config: %v", err)
+		}
+
+		e.updateSSHServerAuth(networkMap.GetSshAuth())
+
 		doneLazy := e.phase("lazy_exclude")
 		excludedLazyPeers := e.toExcludedLazyPeers(remotePeers)
 		e.connMgr.SetExcludeList(e.ctx, excludedLazyPeers)
diff --git a/client/internal/engine_test.go b/client/internal/engine_test.go
index 7394de106..ac2db189b 100644
--- a/client/internal/engine_test.go
+++ b/client/internal/engine_test.go
@@ -790,6 +790,44 @@ func TestEngine_UpdateNetworkMap(t *testing.T) {
 		require.False(t, more, "pass 2 should converge")
 		assert.Equal(t, before+1, countMapEvents(), "the converging pass should publish exactly one event")
 	})
+
+	// The SSH steps walk the whole map and need every peer in the store, so the
+	// draining passes skip them and the converging pass records every host key.
+	t.Run("chunked apply records SSH host keys once peers converge", func(t *testing.T) {
+		sshPeer := func(key, ip string) *mgmtProto.RemotePeerConfig {
+			return &mgmtProto.RemotePeerConfig{
+				WgPubKey:   key,
+				AllowedIps: []string{ip},
+				SshConfig:  &mgmtProto.SSHConfig{SshPubKey: []byte("ssh-ed25519 " + key)},
+			}
+		}
+		peers := []*mgmtProto.RemotePeerConfig{
+			sshPeer("SSHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=", "100.64.0.20/32"),
+			sshPeer("TSHf3Ma6z6mdLbriAJbqhX7+nM/B71lgw2+91q3LfhU=", "100.64.0.21/32"),
+		}
+		hostKey := func(p *mgmtProto.RemotePeerConfig) []byte {
+			state, err := engine.statusRecorder.GetPeer(p.GetWgPubKey())
+			if err != nil {
+				return nil
+			}
+			return state.SSHHostKey
+		}
+		nm := &mgmtProto.NetworkMap{Serial: 9, RemotePeers: peers}
+
+		more, err := engine.updateNetworkMap(nm, 1, true)
+		require.NoError(t, err)
+		require.True(t, more, "the first pass should leave batches to drain")
+		assert.Nil(t, hostKey(peers[0]), "a draining pass should not record SSH host keys")
+
+		for i := 0; more && i < 10; i++ {
+			more, err = engine.updateNetworkMap(nm, 1, false)
+			require.NoError(t, err)
+		}
+		require.False(t, more, "the map should converge")
+		for _, p := range peers {
+			assert.Equal(t, p.GetSshConfig().GetSshPubKey(), hostKey(p), "the converging pass should record the host key of %s", p.GetWgPubKey())
+		}
+	})
 }
 
 func TestEngine_UpdateNetworkMapWithRoutes(t *testing.T) {