calculate WG peer connection status

internal/machine/network/config.go

@@ -84,3 +84,15 @@ func (c Config) toDeviceConfig() (wgtypes.Config, error) {
 		Peers:        wgPeerConfigs,
 	}, nil
 }
+
+func (p *PeerConfig) prefixes() ([]netip.Prefix, error) {
+	managePrefix, err := addrToSingleIPPrefix(p.ManagementIP)
+	if err != nil {
+		return nil, fmt.Errorf("parse management IP: %w", err)
+	}
+	prefixes := []netip.Prefix{managePrefix}
+	if p.Subnet != nil {
+		prefixes = append(prefixes, *p.Subnet)
+	}
+	return prefixes, nil
+}

internal/machine/network/peer.go

@@ -0,0 +1,117 @@
+package network
+
+import (
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"time"
+)
+
+const (
+	PeerStatusUnknown = "unknown"
+	PeerStatusUp      = "up"
+	PeerStatusDown    = "down"
+)
+
+type peer struct {
+	config                 PeerConfig
+	lastEndpointChangeTime time.Time
+	lastHandshakeTime      time.Time
+	receiveBytes           int64
+	transmitBytes          int64
+	status                 string
+}
+
+func newPeer(config PeerConfig) peer {
+	p := peer{
+		config: config,
+		status: PeerStatusUnknown,
+	}
+	if p.config.Endpoint != nil {
+		p.lastEndpointChangeTime = time.Now()
+	}
+	return p
+}
+
+func (p *peer) updateConfig(config PeerConfig) {
+	if p.config.Endpoint != config.Endpoint {
+		p.lastEndpointChangeTime = time.Now()
+		p.status = PeerStatusUnknown
+	}
+	p.config = config
+}
+
+func (p *peer) updateFromWireGuard(wgPeer wgtypes.Peer) {
+	p.lastHandshakeTime = wgPeer.LastHandshakeTime
+	p.receiveBytes = wgPeer.ReceiveBytes
+	p.transmitBytes = wgPeer.TransmitBytes
+	p.calculateStatus()
+}
+
+// Peer status calculation is based on Talos Kubespan implementation:
+// https://github.com/siderolabs/talos/blob/v1.8.0/internal/app/machined/pkg/adapters/kubespan/peer_status.go
+
+// endpointConnectionTimeout is time to wait for initial handshake when the endpoint is just set.
+const endpointConnectionTimeout = 15 * time.Second
+
+// peerDownInterval is the time since last handshake when established peer is considered to be down.
+//
+// WG whitepaper defines a downed peer as being:
+// Handshake Timeout (180s) + Rekey Timeout (5s) + Rekey Attempt Timeout (90s)
+//
+// This interval is applied when the link is already established.
+const peerDownInterval = (180 + 5 + 90) * time.Second
+
+// calculateStatus updates the peer's connection status based on other field values.
+//
+// Goal: endpoint is ultimately down if we haven't seen handshake for more than peerDownInterval,
+// but as the endpoints get updated we want faster feedback, so we start checking more aggressively
+// that the handshake happened within endpointConnectionTimeout since last endpoint change.
+//
+// Timeline:
+//
+// ---------------------------------------------------------------------->
+// ^            ^                                   ^
+// |            |                                   |
+// T0           T0+endpointConnectionTimeout        T0+peerDownInterval
+//
+// Where T0 = lastEndpointChangeTime
+//
+// The question is where is LastHandshakeTimeout vs. those points above:
+//
+//   - if we're past (T0+peerDownInterval), simply check that time since last handshake < peerDownInterval
+//   - if we're between (T0+endpointConnectionTimeout) and (T0+peerDownInterval), and there's no handshake
+//     after the endpoint change, assume that the endpoint is down
+//   - if we're between (T0) and (T0+endpointConnectionTimeout), and there's no handshake since the endpoint change,
+//     consider the state to be unknown
+func (p *peer) calculateStatus() {
+	sinceLastHandshake := time.Since(p.lastHandshakeTime)
+	sinceEndpointChange := time.Since(p.lastEndpointChangeTime)
+
+	switch {
+	case sinceEndpointChange > peerDownInterval: // past T0+peerDownInterval
+		// If we got handshake in the last peerDownInterval, endpoint is up.
+		if sinceLastHandshake < peerDownInterval {
+			p.status = PeerStatusUp
+		} else {
+			p.status = PeerStatusDown
+		}
+	case sinceEndpointChange < endpointConnectionTimeout: // between (T0) and (T0+endpointConnectionTimeout)
+		// Endpoint got recently updated, consider no handshake as 'unknown'.
+		if p.lastHandshakeTime.After(p.lastEndpointChangeTime) {
+			p.status = PeerStatusUp
+		} else {
+			p.status = PeerStatusUnknown
+		}
+	default: // otherwise, we're between (T0+endpointConnectionTimeout) and (T0+peerDownInterval)
+		// If we haven't had the handshake yet, consider the endpoint to be down.
+		if p.lastHandshakeTime.After(p.lastEndpointChangeTime) {
+			p.status = PeerStatusUp
+		} else {
+			p.status = PeerStatusDown
+		}
+	}
+
+	if p.status == PeerStatusDown && p.config.Endpoint == nil {
+		// No endpoint, so unknown.
+		p.status = PeerStatusUnknown
+	}
+}

internal/machine/network/wireguard_linux.go

@@ -14,28 +14,25 @@ import (
 	"net/netip"
 	"slices"
 	"sync"
-	"time"
+	"uncloud/internal/secret"
 )
 
 type WireGuardNetwork struct {
-	link  netlink.Link
+	link netlink.Link
-	peers []peer
+	// peers is a map of peers indexed by their public key.
+	peers map[string]peer
 	// mu synchronises concurrent network configuration changes.
 	mu sync.Mutex
 }
 
-type peer struct {
-	config                 PeerConfig
-	lastEndpointChangeTime time.Time
-}
-
 func NewWireGuardNetwork() (*WireGuardNetwork, error) {
 	link, err := createOrGetLink(WireGuardInterfaceName)
 	if err != nil {
 		return nil, fmt.Errorf("create or get WireGuard link %q: %v", WireGuardInterfaceName, err)
 	}
 	return &WireGuardNetwork{
-		link: link,
+		link:  link,
+		peers: make(map[string]peer),
 	}, nil
 }
 
@@ -74,21 +71,20 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	n.mu.Lock()
 	defer n.mu.Unlock()
 
-	// Reconstruct the list of peers, ensuring that the last endpoint change time is preserved for any existing peers.
+	// Create or update peer structs based on the config.
-	existingPeersByPublicKey := map[string]peer{}
+	newPeersSet := make(map[string]struct{}, len(config.Peers))
-	for _, p := range n.peers {
+	for _, pc := range config.Peers {
-		existingPeersByPublicKey[p.config.PublicKey.String()] = p
+		if p, ok := n.peers[pc.PublicKey.String()]; ok {
-	}
+			p.updateConfig(pc)
-	n.peers = make([]peer, len(config.Peers))
-	for i, peerConfig := range config.Peers {
-		n.peers[i] = peer{
-			config: peerConfig,
-		}
-		existingPeer, ok := existingPeersByPublicKey[peerConfig.PublicKey.String()]
-		if ok && existingPeer.config.Endpoint == peerConfig.Endpoint {
-			n.peers[i].lastEndpointChangeTime = existingPeer.lastEndpointChangeTime
 		} else {
-			n.peers[i].lastEndpointChangeTime = time.Now()
+			n.peers[pc.PublicKey.String()] = newPeer(pc)
+		}
+		newPeersSet[pc.PublicKey.String()] = struct{}{}
+	}
+	// Delete peers that are no longer in the config.
+	for k := range n.peers {
+		if _, ok := newPeersSet[k]; !ok {
+			delete(n.peers, k)
 		}
 	}
 
@@ -96,7 +92,6 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	if err != nil {
 		return fmt.Errorf("create WireGuard client: %w", err)
 	}
-	//goland:noinspection GoUnhandledErrorResult
 	defer wg.Close()
 
 	wgConfig, err := config.toDeviceConfig()
@@ -109,6 +104,11 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	}
 	slog.Info("Configured WireGuard interface.", "name", n.link.Attrs().Name)
 
+	if err = n.updatePeersFromWireGuard(); err != nil {
+		return err
+	}
+	slog.Debug("Updated peers status from WireGuard interface.", "name", n.link.Attrs().Name)
+
 	machinePrefix := netip.PrefixFrom(MachineIP(config.Subnet), config.Subnet.Bits())
 	managementPrefix, err := addrToSingleIPPrefix(config.ManagementIP)
 	if err != nil {
@@ -141,6 +141,31 @@ func (n *WireGuardNetwork) Configure(config Config) error {
 	return nil
 }
 
+// updatePeersFromWireGuard updates the peers status from the WireGuard device peers.
+// mu lock must be held before calling this method.
+func (n *WireGuardNetwork) updatePeersFromWireGuard() error {
+	wg, err := wgctrl.New()
+	if err != nil {
+		return fmt.Errorf("create WireGuard client: %w", err)
+	}
+	defer wg.Close()
+
+	dev, err := wg.Device(n.link.Attrs().Name)
+	if err != nil {
+		return fmt.Errorf("get WireGuard device %q: %w", n.link.Attrs().Name, err)
+	}
+
+	for _, wgPeer := range dev.Peers {
+		if p, ok := n.peers[secret.Secret(wgPeer.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)
+		}
+	}
+	return nil
+}
+
 // updateAddresses assigns addresses to the WireGuard interface and removes old ones.
 // It also removes any other addresses that have been added out of band.
 func (n *WireGuardNetwork) updateAddresses(addrs []netip.Prefix) error {
@@ -178,7 +203,7 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 	// Build a set of compacted IP ranges for all peers.
 	var ipsetBuilder netipx.IPSetBuilder
 	for _, p := range n.peers {
-		prefixes, err := p.prefixes()
+		prefixes, err := p.config.prefixes()
 		if err != nil {
 			return fmt.Errorf("get peer addresses: %w", err)
 		}
@@ -235,18 +260,8 @@ func (n *WireGuardNetwork) updatePeerRoutes() error {
 }
 
 func (n *WireGuardNetwork) Run(ctx context.Context) error {
+	// 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.
 	<-ctx.Done()
 	return nil
 }
-
-func (p peer) prefixes() ([]netip.Prefix, error) {
-	managePrefix, err := addrToSingleIPPrefix(p.config.ManagementIP)
-	if err != nil {
-		return nil, fmt.Errorf("parse management IP: %w", err)
-	}
-	prefixes := []netip.Prefix{managePrefix}
-	if p.config.Subnet != nil {
-		prefixes = append(prefixes, *p.config.Subnet)
-	}
-	return prefixes, nil
-}