Handle discovered Endpoints in topology to enable NAT 2 NAT

Author: JulienVdG

pkg/mesh/mesh.go

@@ -389,6 +389,7 @@ func (m *Mesh) handleLocal(n *Node) {
 		PersistentKeepalive: n.PersistentKeepalive,
 		Subnet:              n.Subnet,
 		WireGuardIP:         m.wireGuardIP,
+		DiscoveredEndpoints: n.DiscoveredEndpoints,
 	}
 	if !nodesAreEqual(n, local) {
 		level.Debug(m.logger).Log("msg", "local node differs from backend")
@@ -469,7 +470,7 @@ func (m *Mesh) applyTopology() {
 		return
 	}
 	oldConf := wireguard.Parse(oldConfRaw)
-	natEndpoints := updateNATEndpoints(nodes, peers, oldConf, m.logger)
+	natEndpoints := discoverNATEndpoints(nodes, peers, oldConf, m.logger)
 	nodes[m.hostname].DiscoveredEndpoints = natEndpoints
 	m.nodes[m.hostname].DiscoveredEndpoints = natEndpoints
 	t, err := NewTopology(nodes, peers, m.granularity, m.hostname, nodes[m.hostname].Endpoint.Port, m.priv, m.subnet, nodes[m.hostname].PersistentKeepalive)
@@ -792,9 +793,8 @@ func linkByIndex(index int) (netlink.Link, error) {
 	return link, nil
 }
 
-// updateNATEndpoints ensures that nodes and peers behind NAT update
-// their endpoints from the WireGuard configuration so they can roam.
-func updateNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wireguard.Conf, logger log.Logger) map[string]*wireguard.Endpoint {
+// discoverNATEndpoints uses the node's WireGuard configuration to returns a list of the most recently discovered endpoints for all nodes and peers behind NAT so that they can roam.
+func discoverNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wireguard.Conf, logger log.Logger) map[string]*wireguard.Endpoint {
 	natEndpoints := make(map[string]*wireguard.Endpoint)
 	keys := make(map[string]*wireguard.Peer)
 	for i := range conf.Peers {
@@ -808,15 +808,13 @@ func updateNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wi
 			if !n.Endpoint.Equal(peer.Endpoint) {
 				natEndpoints[string(n.Key)] = peer.Endpoint
 			}
-			n.Endpoint = peer.Endpoint
 		}
 	}
 	for _, p := range peers {
 		if peer, ok := keys[string(p.PublicKey)]; ok && p.PersistentKeepalive > 0 {
 			if !p.Endpoint.Equal(peer.Endpoint) {
 				natEndpoints[string(p.PublicKey)] = peer.Endpoint
 			}
-			p.Endpoint = peer.Endpoint
 		}
 	}
 	level.Debug(logger).Log("msg", "Discovered WireGuard NAT Endpoints", "DiscoveredEndpoints", natEndpoints)

pkg/mesh/topology.go

@@ -55,12 +55,15 @@ type Topology struct {
 	// the IP is the 0th address in the subnet, i.e. the CIDR
 	// is equal to the Kilo subnet.
 	wireGuardCIDR *net.IPNet
+	// discoveredEndpoints is the updated map of valid discovered Endpoints
+	discoveredEndpoints map[string]*wireguard.Endpoint
 }
 
 type segment struct {
-	allowedIPs []*net.IPNet
-	endpoint   *wireguard.Endpoint
-	key        []byte
+	allowedIPs          []*net.IPNet
+	endpoint            *wireguard.Endpoint
+	key                 []byte
+	persistentKeepalive int
 	// Location is the logical location of this segment.
 	location string
 
@@ -106,7 +109,7 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 		localLocation = nodeLocationPrefix + hostname
 	}
 
-	t := Topology{key: key, port: port, hostname: hostname, location: localLocation, persistentKeepalive: persistentKeepalive, privateIP: nodes[hostname].InternalIP, subnet: nodes[hostname].Subnet, wireGuardCIDR: subnet}
+	t := Topology{key: key, port: port, hostname: hostname, location: localLocation, persistentKeepalive: persistentKeepalive, privateIP: nodes[hostname].InternalIP, subnet: nodes[hostname].Subnet, wireGuardCIDR: subnet, discoveredEndpoints: make(map[string]*wireguard.Endpoint)}
 	for location := range topoMap {
 		// Sort the location so the result is stable.
 		sort.Slice(topoMap[location], func(i, j int) bool {
@@ -134,14 +137,15 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 			hostnames = append(hostnames, node.Name)
 		}
 		t.segments = append(t.segments, &segment{
-			allowedIPs: allowedIPs,
-			endpoint:   topoMap[location][leader].Endpoint,
-			key:        topoMap[location][leader].Key,
-			location:   location,
-			cidrs:      cidrs,
-			hostnames:  hostnames,
-			leader:     leader,
-			privateIPs: privateIPs,
+			allowedIPs:          allowedIPs,
+			endpoint:            topoMap[location][leader].Endpoint,
+			key:                 topoMap[location][leader].Key,
+			persistentKeepalive: topoMap[location][leader].PersistentKeepalive,
+			location:            location,
+			cidrs:               cidrs,
+			hostnames:           hostnames,
+			leader:              leader,
+			privateIPs:          privateIPs,
 		})
 	}
 	// Sort the Topology segments so the result is stable.
@@ -159,6 +163,10 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 	// We need to defensively deduplicate peer allowed IPs. If two peers claim the same IP,
 	// the WireGuard configuration could flap, causing the interface to churn.
 	t.peers = deduplicatePeerIPs(t.peers)
+	// Copy the host node DiscoveredEndpoints in the topology as a starting point.
+	for key := range nodes[hostname].DiscoveredEndpoints {
+		t.discoveredEndpoints[key] = nodes[hostname].DiscoveredEndpoints[key]
+	}
 	// Allocate IPs to the segment leaders in a stable, coordination-free manner.
 	a := newAllocator(*subnet)
 	for _, segment := range t.segments {
@@ -171,11 +179,33 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 		if t.leader && segment.location == t.location {
 			t.wireGuardCIDR = &net.IPNet{IP: ipNet.IP, Mask: subnet.Mask}
 		}
+
+		// Now that the topology is ordered, update the discoveredEndpoints map
+		// add new ones by going through the ordered topology: segments, nodes
+		for _, node := range topoMap[segment.location] {
+			for key := range node.DiscoveredEndpoints {
+				if _, ok := t.discoveredEndpoints[key]; !ok {
+					t.discoveredEndpoints[key] = node.DiscoveredEndpoints[key]
+				}
+			}
+		}
 	}
 
 	return &t, nil
 }
 
+func (t *Topology) updateEndpoint(endpoint *wireguard.Endpoint, key []byte, persistentKeepalive int) *wireguard.Endpoint {
+	// Do not update non-nat peers
+	if persistentKeepalive == 0 {
+		return endpoint
+	}
+	e, ok := t.discoveredEndpoints[string(key)]
+	if ok {
+		return e
+	}
+	return endpoint
+}
+
 // Conf generates a WireGuard configuration file for a given Topology.
 func (t *Topology) Conf() *wireguard.Conf {
 	c := &wireguard.Conf{
@@ -190,7 +220,7 @@ func (t *Topology) Conf() *wireguard.Conf {
 		}
 		peer := &wireguard.Peer{
 			AllowedIPs:          s.allowedIPs,
-			Endpoint:            s.endpoint,
+			Endpoint:            t.updateEndpoint(s.endpoint, s.key, s.persistentKeepalive),
 			PersistentKeepalive: t.persistentKeepalive,
 			PublicKey:           s.key,
 		}
@@ -199,7 +229,7 @@ func (t *Topology) Conf() *wireguard.Conf {
 	for _, p := range t.peers {
 		peer := &wireguard.Peer{
 			AllowedIPs:          p.AllowedIPs,
-			Endpoint:            p.Endpoint,
+			Endpoint:            t.updateEndpoint(p.Endpoint, p.PublicKey, p.PersistentKeepalive),
 			PersistentKeepalive: t.persistentKeepalive,
 			PresharedKey:        p.PresharedKey,
 			PublicKey:           p.PublicKey,