Handle discovered Endpoints in topology to enable NAT 2 NAT

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
+// 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.
-// their endpoints from the WireGuard configuration so they can roam.
+func discoverNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wireguard.Conf, logger log.Logger) map[string]*wireguard.Endpoint {
-func updateNATEndpoints(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,7 +808,6 @@ 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 {
@@ -816,7 +815,6 @@ func updateNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wi
 			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
+	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 {
@@ -137,6 +140,7 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 			allowedIPs:          allowedIPs,
 			endpoint:            topoMap[location][leader].Endpoint,
 			key:                 topoMap[location][leader].Key,
+			persistentKeepalive: topoMap[location][leader].PersistentKeepalive,
 			location:            location,
 			cidrs:               cidrs,
 			hostnames:           hostnames,
@@ -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,

pkg/mesh/topology_test.go

@@ -129,6 +129,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["a"].Location,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -139,6 +140,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["b"].Location,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet, nodes["c"].Subnet},
 						hostnames:           []string{"b", "c"},
@@ -149,6 +151,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -175,6 +178,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["a"].Location,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -185,6 +189,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["b"].Location,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet, nodes["c"].Subnet},
 						hostnames:           []string{"b", "c"},
@@ -195,6 +200,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -221,6 +227,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["a"].Location,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -231,6 +238,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            logicalLocationPrefix + nodes["b"].Location,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet, nodes["c"].Subnet},
 						hostnames:           []string{"b", "c"},
@@ -241,6 +249,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -267,6 +276,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["a"].Name,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -277,6 +287,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["b"].Name,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet},
 						hostnames:           []string{"b"},
@@ -287,6 +298,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["c"].Endpoint,
 						key:                 nodes["c"].Key,
+						persistentKeepalive: nodes["c"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["c"].Name,
 						cidrs:               []*net.IPNet{nodes["c"].Subnet},
 						hostnames:           []string{"c"},
@@ -297,6 +309,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w4, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -323,6 +336,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["a"].Name,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -333,6 +347,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["b"].Name,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet},
 						hostnames:           []string{"b"},
@@ -343,6 +358,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["c"].Endpoint,
 						key:                 nodes["c"].Key,
+						persistentKeepalive: nodes["c"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["c"].Name,
 						cidrs:               []*net.IPNet{nodes["c"].Subnet},
 						hostnames:           []string{"c"},
@@ -353,6 +369,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w4, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -379,6 +396,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["a"].Name,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -389,6 +407,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["b"].Name,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet},
 						hostnames:           []string{"b"},
@@ -399,6 +418,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["c"].Endpoint,
 						key:                 nodes["c"].Key,
+						persistentKeepalive: nodes["c"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["c"].Name,
 						cidrs:               []*net.IPNet{nodes["c"].Subnet},
 						hostnames:           []string{"c"},
@@ -409,6 +429,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w4, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},
@@ -435,6 +456,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["a"].Subnet, nodes["a"].InternalIP, {IP: w1, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["a"].Endpoint,
 						key:                 nodes["a"].Key,
+						persistentKeepalive: nodes["a"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["a"].Name,
 						cidrs:               []*net.IPNet{nodes["a"].Subnet},
 						hostnames:           []string{"a"},
@@ -445,6 +467,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["b"].Subnet, nodes["b"].InternalIP, {IP: w2, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["b"].Endpoint,
 						key:                 nodes["b"].Key,
+						persistentKeepalive: nodes["b"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["b"].Name,
 						cidrs:               []*net.IPNet{nodes["b"].Subnet},
 						hostnames:           []string{"b"},
@@ -455,6 +478,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["c"].Subnet, nodes["c"].InternalIP, {IP: w3, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["c"].Endpoint,
 						key:                 nodes["c"].Key,
+						persistentKeepalive: nodes["c"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["c"].Name,
 						cidrs:               []*net.IPNet{nodes["c"].Subnet},
 						hostnames:           []string{"c"},
@@ -465,6 +489,7 @@ func TestNewTopology(t *testing.T) {
 						allowedIPs:          []*net.IPNet{nodes["d"].Subnet, {IP: w4, Mask: net.CIDRMask(32, 32)}},
 						endpoint:            nodes["d"].Endpoint,
 						key:                 nodes["d"].Key,
+						persistentKeepalive: nodes["d"].PersistentKeepalive,
 						location:            nodeLocationPrefix + nodes["d"].Name,
 						cidrs:               []*net.IPNet{nodes["d"].Subnet},
 						hostnames:           []string{"d"},