Merge pull request #42 from tim-beatham/41-bugfix-fluctuating-ips

41 bugfix fluctuating ips

pkg/mesh/config.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"slices"
+	"strings"
 	"time"
 
 	"github.com/tim-beatham/wgmesh/pkg/conf"
@@ -52,7 +53,7 @@ func (m *WgMeshConfigApplyer) convertMeshNode(node MeshNode, device *wgtypes.Dev
 	allowedips := make([]net.IPNet, 1)
 	allowedips[0] = *node.GetWgHost()
 
-	clients, ok := peerToClients[node.GetWgHost().String()]
+	clients, ok := peerToClients[pubKey.String()]
 
 	if ok {
 		allowedips = append(allowedips, clients...)
@@ -154,62 +155,100 @@ func (m *WgMeshConfigApplyer) getRoutes(meshProvider MeshProvider) map[string][]
 	return routes
 }
 
-func (m *WgMeshConfigApplyer) updateWgConf(mesh MeshProvider) error {
+// getCorrespondignPeer: gets the peer corresponding to the client
-	snap, err := mesh.GetMesh()
+func (m *WgMeshConfigApplyer) getCorrespondingPeer(peers []MeshNode, client MeshNode) MeshNode {
-
+	hashFunc := func(mn MeshNode) int {
-	if err != nil {
+		pubKey, _ := mn.GetPublicKey()
-		return err
+		return lib.HashString(pubKey.String())
 	}
 
-	nodes := lib.MapValues(snap.GetNodes())
+	peer := lib.ConsistentHash(peers, client, hashFunc, hashFunc)
-	peerConfigs := make([]wgtypes.PeerConfig, len(nodes))
+	return peer
-
+}
-	peers := lib.Filter(nodes, func(mn MeshNode) bool {
-		return mn.GetType() == conf.PEER_ROLE
-	})
-
-	var count int = 0
 
+func (m *WgMeshConfigApplyer) getClientConfig(mesh MeshProvider, peers []MeshNode, clients []MeshNode) (*wgtypes.Config, error) {
 	self, err := m.meshManager.GetSelf(mesh.GetMeshId())
 
 	if err != nil {
-		return err
+		return nil, err
 	}
 
+	peer := m.getCorrespondingPeer(peers, self)
+
+	pubKey, _ := peer.GetPublicKey()
+
+	keepAlive := time.Duration(m.config.KeepAliveWg) * time.Second
+	endpoint, err := net.ResolveUDPAddr("udp", peer.GetWgEndpoint())
+
+	if err != nil {
+		return nil, err
+	}
+
+	allowedips := make([]net.IPNet, 1)
+	_, ipnet, _ := net.ParseCIDR("::/0")
+	allowedips[0] = *ipnet
+
+	peerCfgs := make([]wgtypes.PeerConfig, 1)
+
+	peerCfgs[0] = wgtypes.PeerConfig{
+		PublicKey:                   pubKey,
+		Endpoint:                    endpoint,
+		PersistentKeepaliveInterval: &keepAlive,
+		AllowedIPs:                  allowedips,
+	}
+
+	cfg := wgtypes.Config{
+		Peers: peerCfgs,
+	}
+
+	return &cfg, err
+}
+
+func (m *WgMeshConfigApplyer) getPeerConfig(mesh MeshProvider, peers []MeshNode, clients []MeshNode, dev *wgtypes.Device) (*wgtypes.Config, error) {
 	peerToClients := make(map[string][]net.IPNet)
 	routes := m.getRoutes(mesh)
 	installedRoutes := make([]lib.Route, 0)
+	peerConfigs := make([]wgtypes.PeerConfig, 0)
+	self, err := m.meshManager.GetSelf(mesh.GetMeshId())
 
-	dev, _ := mesh.GetDevice()
+	if err != nil {
+		return nil, err
+	}
 
-	for _, n := range nodes {
+	for _, n := range clients {
-		if NodeEquals(n, self) {
+		if len(peers) > 0 {
-			continue
+			peer := m.getCorrespondingPeer(peers, n)
-		}
+			pubKey, _ := peer.GetPublicKey()
-
+			clients, ok := peerToClients[pubKey.String()]
-		if n.GetType() == conf.CLIENT_ROLE && len(peers) > 0 && self.GetType() == conf.CLIENT_ROLE {
-			hashFunc := func(mn MeshNode) int {
-				pubKey, _ := mn.GetPublicKey()
-				return lib.HashString(pubKey.String())
-			}
-
-			peer := lib.ConsistentHash(peers, n, hashFunc, hashFunc)
-
-			clients, ok := peerToClients[peer.GetWgHost().String()]
 
 			if !ok {
 				clients = make([]net.IPNet, 0)
-				peerToClients[peer.GetWgHost().String()] = clients
+				peerToClients[pubKey.String()] = clients
 			}
 
-			peerToClients[peer.GetWgHost().String()] = append(clients, *n.GetWgHost())
+			peerToClients[pubKey.String()] = append(clients, *n.GetWgHost())
+
+			if NodeEquals(self, peer) {
+				cfg, err := m.convertMeshNode(n, dev, peerToClients, routes)
+
+				if err != nil {
+					return nil, err
+				}
+
+				peerConfigs = append(peerConfigs, *cfg)
+			}
+		}
+	}
+
+	for _, n := range peers {
+		if NodeEquals(n, self) {
 			continue
 		}
 
 		peer, err := m.convertMeshNode(n, dev, peerToClients, routes)
 
 		if err != nil {
-			return err
+			return nil, err
 		}
 
 		for _, route := range peer.AllowedIPs {
@@ -224,21 +263,66 @@ func (m *WgMeshConfigApplyer) updateWgConf(mesh MeshProvider) error {
 			}
 		}
 
-		peerConfigs[count] = *peer
+		peerConfigs = append(peerConfigs, *peer)
-		count++
 	}
 
 	cfg := wgtypes.Config{
-		Peers: peerConfigs,
+		Peers:        peerConfigs,
+		ReplacePeers: true,
 	}
 
-	err = m.meshManager.GetClient().ConfigureDevice(dev.Name, cfg)
+	err = m.routeInstaller.InstallRoutes(dev.Name, installedRoutes...)
+	return &cfg, err
+}
+
+func (m *WgMeshConfigApplyer) updateWgConf(mesh MeshProvider) error {
+	snap, err := mesh.GetMesh()
 
 	if err != nil {
 		return err
 	}
 
-	return m.routeInstaller.InstallRoutes(dev.Name, installedRoutes...)
+	nodes := lib.MapValues(snap.GetNodes())
+	dev, _ := mesh.GetDevice()
+
+	slices.SortFunc(nodes, func(a, b MeshNode) int {
+		return strings.Compare(string(a.GetType()), string(b.GetType()))
+	})
+
+	peers := lib.Filter(nodes, func(mn MeshNode) bool {
+		return mn.GetType() == conf.PEER_ROLE
+	})
+
+	clients := lib.Filter(nodes, func(mn MeshNode) bool {
+		return mn.GetType() == conf.CLIENT_ROLE
+	})
+
+	self, err := m.meshManager.GetSelf(mesh.GetMeshId())
+
+	if err != nil {
+		return err
+	}
+
+	var cfg *wgtypes.Config = nil
+
+	switch self.GetType() {
+	case conf.PEER_ROLE:
+		cfg, err = m.getPeerConfig(mesh, peers, clients, dev)
+	case conf.CLIENT_ROLE:
+		cfg, err = m.getClientConfig(mesh, peers, clients)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	err = m.meshManager.GetClient().ConfigureDevice(dev.Name, *cfg)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
 }
 
 func (m *WgMeshConfigApplyer) ApplyConfig() error {
@@ -267,7 +351,8 @@ func (m *WgMeshConfigApplyer) RemovePeers(meshId string) error {
 	}
 
 	m.meshManager.GetClient().ConfigureDevice(dev.Name, wgtypes.Config{
-		Peers: make([]wgtypes.PeerConfig, 0),
+		Peers:        make([]wgtypes.PeerConfig, 0),
+		ReplacePeers: true,
 	})
 
 	return nil