update existing WG peers instead of replacing to not interrupt and lose track of handshakes

Makefile

@@ -14,6 +14,7 @@ uncloud-dev:
 reset-dev:
 	ssh spy@192.168.40.243 "sudo systemctl stop uncloud && sudo rm -rf /var/lib/uncloud"
 	ssh spy@192.168.40.176 "sudo systemctl stop uncloud && sudo rm -rf /var/lib/uncloud"
+	ssh ubuntu@152.67.101.197 "sudo systemctl stop uncloud && sudo rm -rf /var/lib/uncloud"
 
 .PHONY: proto
 proto:

internal/machine/network/config.go

@@ -41,7 +41,9 @@ func (c Config) IsConfigured() bool {
 		c.PrivateKey != nil && c.PublicKey != nil
 }
 
-func (c Config) toDeviceConfig() (wgtypes.Config, error) {
+// toDeviceConfig converts the configuration to a WireGuard device configuration. It updates the existing peers
+// without replacing them to not disrupt existing connections and to not lose track of the last handshake time.
+func (c Config) toDeviceConfig(currentPeers []wgtypes.Peer) (wgtypes.Config, error) {
 	privateKey, err := wgtypes.NewKey(c.PrivateKey)
 	if err != nil {
 		return wgtypes.Config{}, fmt.Errorf("parse private key: %w", err)
@@ -50,6 +52,8 @@ func (c Config) toDeviceConfig() (wgtypes.Config, error) {
 
 	persistentKeepalive := WireGuardKeepaliveInterval
 	wgPeerConfigs := make([]wgtypes.PeerConfig, len(c.Peers))
+	// A set of new peer public keys for checking which current peers should be removed.
+	newPeersSet := make(map[string]struct{}, len(c.Peers))
 	for i, peerConfig := range c.Peers {
 		peerPublicKey, kErr := wgtypes.NewKey(peerConfig.PublicKey)
 		if kErr != nil {
@@ -75,12 +79,24 @@ func (c Config) toDeviceConfig() (wgtypes.Config, error) {
 				Port: int(peerConfig.Endpoint.Port()),
 			}
 		}
+
+		newPeersSet[wgPeerConfigs[i].PublicKey.String()] = struct{}{}
+	}
+
+	// Remove peers that are not in the configuration.
+	for _, p := range currentPeers {
+		if _, ok := newPeersSet[p.PublicKey.String()]; !ok {
+			wgPeerConfigs = append(wgPeerConfigs, wgtypes.PeerConfig{
+				PublicKey: p.PublicKey,
+				Remove:    true,
+			})
+		}
 	}
 
 	return wgtypes.Config{
 		PrivateKey:   &privateKey,
 		ListenPort:   &listenPort,
-		ReplacePeers: true,
+		ReplacePeers: false,
 		Peers:        wgPeerConfigs,
 	}, nil
 }

internal/machine/network/wireguard_linux.go

@@ -97,7 +97,11 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	}
 	defer wg.Close()
 
-	wgConfig, err := config.toDeviceConfig()
+	dev, err := wg.Device(n.link.Attrs().Name)
+	if err != nil {
+		return fmt.Errorf("get WireGuard device %q: %w", n.link.Attrs().Name, err)
+	}
+	wgConfig, err := config.toDeviceConfig(dev.Peers)
 	if err != nil {
 		return err
 	}
@@ -159,11 +163,12 @@ func (n *WireGuardNetwork) updatePeersFromWireGuard() error {
 	}
 
 	for _, wgPeer := range dev.Peers {
-		if p, ok := n.peers[secret.Secret(wgPeer.PublicKey[:]).String()]; ok {
+		publicKey := secret.Secret(wgPeer.PublicKey[:])
+		if p, ok := n.peers[publicKey.String()]; ok {
 			p.updateFromWireGuard(wgPeer)
 		} else {
 			// Assume that WG peers are not updated out of band so they should always be in sync with the config.
-			slog.Warn("Found WireGuard peer that is not in the configuration.", "public_key", wgPeer.PublicKey)
+			slog.Warn("Found WireGuard peer that is not in the configuration.", "public_key", publicKey)
 		}
 	}
 	return nil
@@ -228,13 +233,16 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 				Scope:     netlink.SCOPE_LINK,
 				Dst:       &dst,
 			},
-		); err != nil && !errors.Is(err, unix.EEXIST) {
+		); err != nil {
-			return fmt.Errorf("add route to WireGuard link %q: %w", n.link.Attrs().Name, err)
+			if !errors.Is(err, unix.EEXIST) {
+				return fmt.Errorf("add route to WireGuard link %q: %w", n.link.Attrs().Name, err)
+			}
+		} else {
+			slog.Debug(
+				"Added route to peer(s) via WireGuard interface.",
+				"name", n.link.Attrs().Name, "dst", prefix,
+			)
 		}
-		slog.Debug(
-			"Added route to peer(s) via WireGuard interface.",
-			"name", n.link.Attrs().Name, "dst", prefix,
-		)
 	}
 
 	// Remove old routes to IP ranges that are no longer in the configuration.
@@ -338,7 +346,7 @@ func (n *WireGuardNetwork) changeWireGuardEndpoints() error {
 	}
 	for _, pc := range wgPeerConfigs {
 		slog.Info("Changed peer endpoint on WireGuard interface.",
-			"name", n.link.Attrs().Name, "public_key", pc.PublicKey.String(), "endpoint", pc.Endpoint)
+			"name", n.link.Attrs().Name, "public_key", secret.Secret(pc.PublicKey[:]), "endpoint", pc.Endpoint)
 	}
 	return nil
 }