automatically rotate WG peer endpoints if can't establish a connection

internal/machine/network/peer.go

@@ -3,6 +3,8 @@ package network
 import (
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"log/slog"
+	"net/netip"
+	"slices"
 	"time"
 )
 
@@ -117,6 +119,31 @@ func (p *peer) calculateStatus() {
 		p.status = PeerStatusUnknown
 	}
 	if p.status != lastStatus {
-		slog.Debug("Peer status changed.", "public_key", p.config.PublicKey, "status", p.status)
+		slog.Info("Peer status changed.", "public_key", p.config.PublicKey, "status", p.status)
 	}
 }
+
+// shouldChangeEndpoint returns the next endpoint to use and a boolean indicating if the endpoint should be changed.
+func (p *peer) shouldChangeEndpoint() (netip.AddrPort, bool) {
+	if p.config.Endpoint != nil && p.status != PeerStatusDown {
+		// Shouldn't change the endpoint if it's set and the status is 'up' or 'unknown'.
+		return netip.AddrPort{}, false
+	}
+	if len(p.config.AllEndpoints) == 0 {
+		// No endpoints to choose from.
+		return netip.AddrPort{}, false
+	}
+	if p.config.Endpoint == nil {
+		// No endpoint set, so choose the first one.
+		return p.config.AllEndpoints[0], true
+	}
+	if len(p.config.AllEndpoints) == 1 && p.config.Endpoint == &p.config.AllEndpoints[0] {
+		// Only one endpoint and it's the current one, can't rotate.
+		return netip.AddrPort{}, false
+	}
+
+	// The endpoint is set and the status is 'down', so rotate to the next one.
+	idx := slices.Index(p.config.AllEndpoints, *p.config.Endpoint)
+	endpoint := p.config.AllEndpoints[(idx+1)%len(p.config.AllEndpoints)]
+	return endpoint, true
+}

internal/machine/network/wireguard_linux.go

@@ -10,7 +10,9 @@ import (
 	"go4.org/netipx"
 	"golang.org/x/sys/unix"
 	"golang.zx2c4.com/wireguard/wgctrl"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"log/slog"
+	"net"
 	"net/netip"
 	"slices"
 	"sync"
@@ -261,22 +263,82 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 }
 
 func (n *WireGuardNetwork) Run(ctx context.Context) error {
+	wg, err := wgctrl.New()
+	if err != nil {
+		return fmt.Errorf("create WireGuard client: %w", err)
+	}
+	defer wg.Close()
+
 	ticker := time.NewTicker(1 * time.Second)
 	for {
 		select {
 		case <-ticker.C:
 			n.mu.Lock()
-			if err := n.updatePeersFromWireGuard(); err != nil {
+			if err = n.changeWireGuardEndpoints(); err != nil {
+				slog.Error("Failed to update peer endpoints on WireGuard interface.",
+					"name", n.link.Attrs().Name, "err", err)
+			}
+			if err = n.updatePeersFromWireGuard(); err != nil {
 				slog.Error("Failed to update peers status from WireGuard interface.",
 					"name", n.link.Attrs().Name, "err", err)
 			}
 			n.mu.Unlock()
 
-			// TODO: check if the endpoint should be changed for any peers. If so, change it and notify the controller
-			//  to preserve the change in the machine state.
+			// TODO: notify the controller through a channel to preserve the change in the machine state.
 		case <-ctx.Done():
 			return nil
 		}
 	}
 
 }
+
+// changeWireGuardEndpoints rotates the endpoints of the WireGuard peers that need to be changed.
+func (n *WireGuardNetwork) changeWireGuardEndpoints() error {
+	var wgPeerConfigs []wgtypes.PeerConfig
+	for _, p := range n.peers {
+		newEndpoint, ok := p.shouldChangeEndpoint()
+		if !ok {
+			continue
+		}
+		newConfig := p.config
+		newConfig.Endpoint = &newEndpoint
+		p.updateConfig(newConfig)
+
+		publicKey, err := wgtypes.NewKey(p.config.PublicKey)
+		if err != nil {
+			return fmt.Errorf("parse peer public key: %w", err)
+		}
+		wgPeerConfigs = append(wgPeerConfigs, wgtypes.PeerConfig{
+			PublicKey:  publicKey,
+			UpdateOnly: true,
+			Endpoint: &net.UDPAddr{
+				IP:   p.config.Endpoint.Addr().AsSlice(),
+				Port: int(p.config.Endpoint.Port()),
+			},
+		})
+	}
+	if len(wgPeerConfigs) == 0 {
+		// No changes to the endpoints.
+		return nil
+	}
+
+	wg, err := wgctrl.New()
+	if err != nil {
+		return fmt.Errorf("create WireGuard client: %w", err)
+	}
+	defer wg.Close()
+
+	wgConfigPatch := wgtypes.Config{
+		ReplacePeers: false,
+		Peers:        wgPeerConfigs,
+	}
+	// Apply the configuration patch to the WireGuard device.
+	if err = wg.ConfigureDevice(n.link.Attrs().Name, wgConfigPatch); err != nil {
+		return fmt.Errorf("configure WireGuard device %q with endpoint changes: %w", n.link.Attrs().Name, err)
+	}
+	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)
+	}
+	return nil
+}