watch for peer endpoint changes and preserve in machine state

internal/machine/network.go

@@ -26,9 +26,12 @@ const (
 )
 
 type networkController struct {
-	state        *State
-	store        *store.Store
-	wgnet        *network.WireGuardNetwork
+	state *State
+	store *store.Store
+
+	wgnet           *network.WireGuardNetwork
+	endpointChanges <-chan network.EndpointChangeEvent
+
 	server       *grpc.Server
 	corroService corroservice.Service
 
@@ -46,13 +49,15 @@ func newNetworkController(
 	if err != nil {
 		return nil, fmt.Errorf("create WireGuard network: %w", err)
 	}
+	endpointChanges := wgnet.WatchEndpoints()
 
 	return &networkController{
-		state:        state,
-		store:        store,
-		wgnet:        wgnet,
-		server:       server,
-		corroService: corroService,
+		state:           state,
+		store:           store,
+		wgnet:           wgnet,
+		endpointChanges: endpointChanges,
+		server:          server,
+		corroService:    corroService,
 	}, nil
 }
 
@@ -115,8 +120,38 @@ func (nc *networkController) Run(ctx context.Context) error {
 		},
 	)
 
-	// TODO: run another goroutine to watch WG network endpoint changes and update the state accordingly.
-	//  Network updates in the state should not occur outside of this controller.
+	// Watch for endpoint changes and update the machine state accordingly.
+	errGroup.Go(
+		func() error {
+			for {
+				select {
+				case e, ok := <-nc.endpointChanges:
+					if !ok {
+						// The channel was closed, stop watching for changes.
+						nc.endpointChanges = nil
+						return nil
+					}
+
+					nc.state.mu.Lock()
+					for i := range nc.state.Network.Peers {
+						if nc.state.Network.Peers[i].PublicKey.Equal(e.PublicKey) {
+							nc.state.Network.Peers[i].Endpoint = &e.Endpoint
+							break
+						}
+					}
+					if err = nc.state.Save(); err != nil {
+						slog.Error("Failed to save machine state.", "err", err)
+					}
+					nc.state.mu.Unlock()
+
+					slog.Debug("Preserved endpoint change in the machine state.",
+						"public_key", e.PublicKey, "endpoint", e.Endpoint)
+				case <-ctx.Done():
+					return nil
+				}
+			}
+		},
+	)
 
 	errGroup.Go(
 		func() error {

internal/machine/network/wireguard.go

@@ -3,6 +3,7 @@ package network
 import (
 	"fmt"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"net/netip"
 	"time"
 	"uncloud/internal/secret"
 )
@@ -14,6 +15,12 @@ const (
 	WireGuardKeepaliveInterval = 25 * time.Second
 )
 
+type EndpointChangeEvent struct {
+	PublicKey secret.Secret
+	// Endpoint is the new endpoint of the peer.
+	Endpoint netip.AddrPort
+}
+
 // NewMachineKeys generates a new WireGuard private and public key pair.
 func NewMachineKeys() (privKey, pubKey secret.Secret, err error) {
 	wgPrivKey, err := wgtypes.GeneratePrivateKey()

internal/machine/network/wireguard_darwin.go

@@ -20,3 +20,7 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 func (n *WireGuardNetwork) Run(ctx context.Context) error {
 	return errors.New("not implemented on darwin")
 }
+
+func (n *WireGuardNetwork) WatchEndpoints() <-chan EndpointChangeEvent {
+	return nil
+}

internal/machine/network/wireguard_linux.go

@@ -24,6 +24,8 @@ type WireGuardNetwork struct {
 	link netlink.Link
 	// peers is a map of peers indexed by their public key.
 	peers map[string]*peer
+	// watchers is a list of channels that are notified when the endpoints of the peers change.
+	watchers []chan EndpointChangeEvent
 	// mu synchronises concurrent network configuration changes.
 	mu sync.Mutex
 }
@@ -282,7 +284,7 @@ func (n *WireGuardNetwork) Run(ctx context.Context) error {
 		select {
 		case <-ticker.C:
 			n.mu.Lock()
-			if err = n.changeWireGuardEndpoints(); err != nil {
+			if err = n.changeWireGuardEndpoints(ctx); err != nil {
 				slog.Error("Failed to update peer endpoints on WireGuard interface.",
 					"name", n.link.Attrs().Name, "err", err)
 			}
@@ -291,17 +293,27 @@ func (n *WireGuardNetwork) Run(ctx context.Context) error {
 					"name", n.link.Attrs().Name, "err", err)
 			}
 			n.mu.Unlock()
-
-			// TODO: notify the controller through a channel to preserve the change in the machine state.
 		case <-ctx.Done():
+			for _, ch := range n.watchers {
+				close(ch)
+			}
 			return nil
 		}
 	}
 }
 
-// changeWireGuardEndpoints rotates the endpoints of the WireGuard peers that need to be changed.
-func (n *WireGuardNetwork) changeWireGuardEndpoints() error {
+// WatchEndpoints returns a channel that receives endpoint change events for the WireGuard peers.
+func (n *WireGuardNetwork) WatchEndpoints() <-chan EndpointChangeEvent {
+	ch := make(chan EndpointChangeEvent)
+	n.watchers = append(n.watchers, ch)
+	return ch
+}
+
+// changeWireGuardEndpoints rotates the endpoints of the WireGuard peers with 'down' status
+// in an attempt to find a working one.
+func (n *WireGuardNetwork) changeWireGuardEndpoints(ctx context.Context) error {
 	var wgPeerConfigs []wgtypes.PeerConfig
+	var events []EndpointChangeEvent
 	for _, p := range n.peers {
 		newEndpoint, ok := p.shouldChangeEndpoint()
 		if !ok {
@@ -323,6 +335,11 @@ func (n *WireGuardNetwork) changeWireGuardEndpoints() error {
 				Port: int(p.config.Endpoint.Port()),
 			},
 		})
+
+		events = append(events, EndpointChangeEvent{
+			PublicKey: p.config.PublicKey,
+			Endpoint:  *p.config.Endpoint,
+		})
 	}
 	if len(wgPeerConfigs) == 0 {
 		// No changes to the endpoints.
@@ -347,5 +364,22 @@ func (n *WireGuardNetwork) changeWireGuardEndpoints() error {
 		slog.Info("Changed peer endpoint on WireGuard interface.",
 			"name", n.link.Attrs().Name, "public_key", secret.Secret(pc.PublicKey[:]), "endpoint", pc.Endpoint)
 	}
+
+	// Notify the watchers about the endpoint changes.
+	for _, ch := range n.watchers {
+		for _, e := range events {
+			select {
+			case ch <- e:
+			// Use a timeout to avoid blocking the network control loop.
+			case <-time.After(1 * time.Second):
+				slog.Error("Timeout notifying watchers about a peer endpoint change.")
+				// As of October 2024, the machine is the only watcher to persist changed endpoints in the machine
+				// state which can tolerate missed events. Therefore, we can ignore the error here.
+				return nil
+			case <-ctx.Done():
+				return nil
+			}
+		}
+	}
 	return nil
 }