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")
@@ -808,15 +809,15 @@ func updateNATEndpoints(nodes map[string]*Node, peers map[string]*Peer, conf *wi
 			if isPublic(peer.Endpoint.IP) && !n.Endpoint.Equal(peer.Endpoint) {
 				natEndpoints[string(n.Key)] = peer.Endpoint
 			}
-			n.Endpoint = 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
+			//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.
@@ -173,9 +177,37 @@ func NewTopology(nodes map[string]*Node, peers map[string]*Peer, granularity Gra
 		}
 	}
 
+	// Copy the host node DiscoveredEndpoints first.
+	for key := range nodes[hostname].DiscoveredEndpoints {
+		t.discoveredEndpoints[key] = nodes[hostname].DiscoveredEndpoints[key]
+	}
+	// Now that the topology is ordered, update the discoveredEndpoints map
+	// add new ones by going through the ordered topology: segments, nodes
+	for _, segment := range t.segments {
+		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 +222,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 +231,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,