From ace8eb9a0793bdbbc83f5c16eb809a0ec994e4d3 Mon Sep 17 00:00:00 2001
From: Julien Viard de Galbert <jviarddegalbert@scaleway.com>
Date: Fri, 16 Apr 2021 15:34:29 +0200
Subject: [PATCH] Record discovered endpoints in node

---
 pkg/mesh/backend.go |  1 +
 pkg/mesh/mesh.go    | 18 ++++++++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/pkg/mesh/backend.go b/pkg/mesh/backend.go
index 6053da1..ad8c0ca 100644
--- a/pkg/mesh/backend.go
+++ b/pkg/mesh/backend.go
@@ -66,6 +66,7 @@ type Node struct {
 	PersistentKeepalive int
 	Subnet              *net.IPNet
 	WireGuardIP         *net.IPNet
+	DiscoveredEndpoints map[string]*wireguard.Endpoint
 }
 
 // Ready indicates whether or not the node is ready.
diff --git a/pkg/mesh/mesh.go b/pkg/mesh/mesh.go
index c3b49c9..9c2ef11 100644
--- a/pkg/mesh/mesh.go
+++ b/pkg/mesh/mesh.go
@@ -469,7 +469,9 @@ func (m *Mesh) applyTopology() {
 		return
 	}
 	oldConf := wireguard.Parse(oldConfRaw)
-	updateNATEndpoints(nodes, peers, oldConf)
+	natEndpoints := updateNATEndpoints(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)
 	if err != nil {
 		level.Error(m.logger).Log("error", err)
@@ -774,19 +776,31 @@ func linkByIndex(index int) (netlink.Link, error) {
 
 // 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) {
+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 {
 		keys[string(conf.Peers[i].PublicKey)] = conf.Peers[i]
 	}
 	for _, n := range nodes {
 		if peer, ok := keys[string(n.Key)]; ok && n.PersistentKeepalive > 0 {
+			level.Debug(logger).Log("msg", "WireGuard Update NAT Endpoint", "node", n.Name, "endpoint", peer.Endpoint, "former-endpoint", n.Endpoint, "same", n.Endpoint.Equal(peer.Endpoint))
+			// Should check location leader but only available in topology ... or have topology handle that list
+			// Better check wg latest-handshake
+			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)
+	return natEndpoints
 }