From 7b939e0468b07959f04c45436165f5a4d0465f9a Mon Sep 17 00:00:00 2001
From: Tim Beatham <tcb1@st-andrews.ac.uk>
Date: Tue, 21 Nov 2023 20:42:43 +0000
Subject: [PATCH] 24-keepalive-holepunch

Added the ability to hole punch NAT
---
 pkg/lib/hashing.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 pkg/mesh/config.go | 17 +++++++++--------
 2 files changed, 55 insertions(+), 8 deletions(-)
 create mode 100644 pkg/lib/hashing.go

diff --git a/pkg/lib/hashing.go b/pkg/lib/hashing.go
new file mode 100644
index 0000000..8bb40ab
--- /dev/null
+++ b/pkg/lib/hashing.go
@@ -0,0 +1,46 @@
+package lib
+
+import (
+	"hash/fnv"
+	"sort"
+)
+
+type consistentHashRecord[V any] struct {
+	record V
+	value  int
+}
+
+func HashString(value string) int {
+	f := fnv.New32a()
+	f.Write([]byte(value))
+	return int(f.Sum32())
+}
+
+// ConsistentHash implementation. Traverse the values until we find a key
+// less than ours.
+func ConsistentHash[V any](values []V, client V, keyFunc func(V) int) V {
+	if len(values) == 0 {
+		panic("values is empty")
+	}
+
+	vs := Map(values, func(v V) consistentHashRecord[V] {
+		return consistentHashRecord[V]{
+			v,
+			keyFunc(v),
+		}
+	})
+
+	sort.SliceStable(vs, func(i, j int) bool {
+		return vs[i].value < vs[j].value
+	})
+
+	ourKey := keyFunc(client)
+
+	for _, record := range vs {
+		if ourKey < record.value {
+			return record.record
+		}
+	}
+
+	return vs[0].record
+}
diff --git a/pkg/mesh/config.go b/pkg/mesh/config.go
index 00f47a7..55c0035 100644
--- a/pkg/mesh/config.go
+++ b/pkg/mesh/config.go
@@ -2,7 +2,6 @@ package mesh
 
 import (
 	"fmt"
-	"hash/fnv"
 	"net"
 	"time"
 
@@ -88,14 +87,16 @@ func (m *WgMeshConfigApplyer) updateWgConf(mesh MeshProvider) error {
 	}
 
 	for _, n := range nodes {
+		if NodeEquals(n, self) {
+			continue
+		}
+
 		if n.GetType() == conf.CLIENT_ROLE && len(peers) > 0 {
-			a := fnv.New32a()
-			a.Write([]byte(n.GetHostEndpoint()))
-			sum := a.Sum32()
+			peer := lib.ConsistentHash(peers, n, func(mn MeshNode) int {
+				return lib.HashString(mn.GetWgHost().String())
+			})
 
-			responsiblePeer := peers[int(sum)%len(peers)]
-
-			if responsiblePeer.GetHostEndpoint() != self.GetHostEndpoint() {
+			if !NodeEquals(peer, self) {
 				dev, err := mesh.GetDevice()
 
 				if err != nil {
@@ -103,7 +104,7 @@ func (m *WgMeshConfigApplyer) updateWgConf(mesh MeshProvider) error {
 				}
 
 				rtnl.AddRoute(dev.Name, lib.Route{
-					Gateway:     responsiblePeer.GetWgHost().IP,
+					Gateway:     peer.GetWgHost().IP,
 					Destination: *n.GetWgHost(),
 				})