From 4094f97a2d5b652b83115fd5e187be4081834010 Mon Sep 17 00:00:00 2001
From: Konstantin Demin <rockdrilla@gmail.com>
Date: Sun, 7 Jul 2024 16:11:02 +0300
Subject: [PATCH] rework address mapping

---
 addr-map.go   | 92 ++++++++++++++++++++++++---------------------------
 prometheus.go | 32 ++++++++++++++++--
 2 files changed, 72 insertions(+), 52 deletions(-)

diff --git a/addr-map.go b/addr-map.go
index 25e102d..b74ee71 100644
--- a/addr-map.go
+++ b/addr-map.go
@@ -79,68 +79,62 @@ func addrMapGet(srcIp net.IP, dstCidr *net.IPNet, ttl uint32) net.IP {
 		log.Fatalf("addrMapGet(): src/dst size mismatch: %v vs %v", addrlen, len(dstCidr.IP))
 	}
 
+	var hkey any
+	switch addrlen {
+	case net.IPv4len:
+		hkey = binary.NativeEndian.Uint32(srcIp)
+	case net.IPv6len:
+		hkey = srcIp.To16().String()
+	}
+
 	var curr AddrMap
 	curr.SrcAddr = make([]byte, addrlen)
 	curr.DstAddr = make([]byte, addrlen)
 	copy(curr.DstAddr, srcIp)
 	curr.Ttl = ttl
 
-	for {
-		_, err := rand.Read(curr.SrcAddr)
-		if err != nil {
-			log.Fatalf("rand.Read(): error %v", err)
-		}
+	_, err := rand.Read(curr.SrcAddr)
+	if err != nil {
+		log.Fatalf("rand.Read(): error %v", err)
+	}
 
-		// adjust random bytes to dstCidr
-		for i := range addrlen / 4 {
-			a := binary.NativeEndian.Uint32(dstCidr.IP[i*4:])
-			b := binary.NativeEndian.Uint32(curr.SrcAddr[i*4:])
-			m := binary.NativeEndian.Uint32(dstCidr.Mask[i*4:])
-			a += (b & ^m)
-			binary.NativeEndian.PutUint32(curr.SrcAddr[i*4:], a)
-		}
+	// adjust random bytes to dstCidr
+	for i := range addrlen / 4 {
+		a := binary.NativeEndian.Uint32(dstCidr.IP[i*4:])
+		b := binary.NativeEndian.Uint32(curr.SrcAddr[i*4:])
+		m := binary.NativeEndian.Uint32(dstCidr.Mask[i*4:])
+		a += (b & ^m)
+		binary.NativeEndian.PutUint32(curr.SrcAddr[i*4:], a)
+	}
 
-		curr.Created = time.Now()
+	curr.Created = time.Now()
 
-		var hkey, xprev any
-		var loaded bool
+	var xprev any
+	var loaded bool
+	switch addrlen {
+	case net.IPv4len:
+		xprev, loaded = addr4.LoadOrStore(hkey, curr)
+	case net.IPv6len:
+		xprev, loaded = addr6.LoadOrStore(hkey, curr)
+	}
+	if !loaded {
+		// early return
+		return curr.SrcAddr
+	}
+
+	prev, ok := xprev.(AddrMap)
+	if !ok {
+		log.Fatalf("addrMapGet(): wrong value type from sync.Map")
+	}
+
+	copy(curr.SrcAddr, prev.SrcAddr)
+	if prev.GetTtl() < int32(curr.Ttl) {
 		switch addrlen {
 		case net.IPv4len:
-			hkey = binary.NativeEndian.Uint32(curr.SrcAddr)
-			xprev, loaded = addr4.LoadOrStore(hkey, curr)
+			addr4.Store(hkey, curr)
 		case net.IPv6len:
-			hkey = binary.NativeEndian.Uint64(curr.SrcAddr[net.IPv6len/2:])
-			xprev, loaded = addr6.LoadOrStore(hkey, curr)
+			addr6.Store(hkey, curr)
 		}
-		if !loaded {
-			// early return
-			return curr.SrcAddr
-		}
-
-		prev, ok := xprev.(AddrMap)
-		if !ok {
-			log.Fatalf("addrMapGet(): wrong value type from sync.Map")
-		}
-
-		if !net.IP.Equal(curr.SrcAddr, prev.SrcAddr) {
-			// generate next random address
-			continue
-		}
-		if !net.IP.Equal(curr.DstAddr, prev.DstAddr) {
-			// generate next random address
-			continue
-		}
-
-		if prev.GetTtl() < int32(curr.Ttl) {
-			switch addrlen {
-			case net.IPv4len:
-				addr4.Store(hkey, curr)
-			case net.IPv6len:
-				addr6.Store(hkey, curr)
-			}
-		}
-
-		break
 	}
 
 	return curr.SrcAddr
diff --git a/prometheus.go b/prometheus.go
index 6c2975d..46e80b7 100644
--- a/prometheus.go
+++ b/prometheus.go
@@ -16,11 +16,35 @@ var (
 	_promRegistry    *prometheus.Registry
 	_promHttpHandler http.Handler
 
-	opsProcessed = prometheus.NewCounter(prometheus.CounterOpts{
+	promOpsProcessed = prometheus.NewCounter(prometheus.CounterOpts{
 		Name: "processed_ops_total",
 		Help: "The total number of processed requests",
 	})
 
+	promAddr4Count = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Name: "ipv4_mapped_addr_count",
+		Help: "The total number of IPv4-mapped addresses",
+	}, func() float64 {
+		var x uint32
+		addr4.Range(func(key, value any) bool {
+			x++
+			return true
+		})
+		return float64(x)
+	})
+
+	promAddr6Count = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Name: "ipv6_mapped_addr_count",
+		Help: "The total number of IPv6-mapped addresses",
+	}, func() float64 {
+		var x uint32
+		addr6.Range(func(key, value any) bool {
+			x++
+			return true
+		})
+		return float64(x)
+	})
+
 	labelStringReplacer *strings.Replacer = strings.NewReplacer(
 		"\"", "",
 		"'", "",
@@ -29,7 +53,9 @@ var (
 
 func setupPrometheus(r *gin.Engine) {
 	_promRegistry = prometheus.NewRegistry()
-	_promRegistry.MustRegister(opsProcessed)
+	_promRegistry.MustRegister(promOpsProcessed)
+	_promRegistry.MustRegister(promAddr4Count)
+	_promRegistry.MustRegister(promAddr6Count)
 
 	_promHttpHandler = promhttp.HandlerFor(_promRegistry, promhttp.HandlerOpts{
 		Registry:          _promRegistry,
@@ -55,7 +81,7 @@ func promSanitizeLabel(str string, fallback string) string {
 }
 
 func promHttpHandler(c *gin.Context) {
-	opsProcessed.Inc()
+	promOpsProcessed.Inc()
 
 	_promHttpHandler.ServeHTTP(c.Writer, c.Request)
 }