diff --git a/internal/machine/network.go b/internal/machine/network.go index 8078fd81..626284d6 100644 --- a/internal/machine/network.go +++ b/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 { diff --git a/internal/machine/network/wireguard.go b/internal/machine/network/wireguard.go index c44a09e7..648fb5ef 100644 --- a/internal/machine/network/wireguard.go +++ b/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() diff --git a/internal/machine/network/wireguard_darwin.go b/internal/machine/network/wireguard_darwin.go index 8b60a512..99b15faf 100644 --- a/internal/machine/network/wireguard_darwin.go +++ b/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 +} diff --git a/internal/machine/network/wireguard_linux.go b/internal/machine/network/wireguard_linux.go index 5e08fc99..b3c21f94 100644 --- a/internal/machine/network/wireguard_linux.go +++ b/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 }