From 96dcf14d5d305c3b47ff5669d54fc72b8fe83d44 Mon Sep 17 00:00:00 2001
From: lemmi <lemmi@nerd2nerd.org>
Date: Wed, 15 Dec 2021 05:41:05 +0100
Subject: [PATCH] use netaddr.IPPrefix as map key

---
 main.go | 29 +++++++++++++----------------
 1 file changed, 13 insertions(+), 16 deletions(-)

diff --git a/main.go b/main.go
index e93ff1f..ae2249d 100644
--- a/main.go
+++ b/main.go
@@ -66,16 +66,16 @@ func (r route) String() string {
 
 type routeStats struct {
 	sync.Mutex
-	stats map[string]*route
+	stats map[netaddr.IPPrefix]*route
 }
 
 func newRouteStats() *routeStats {
 	return &routeStats{
-		stats: make(map[string]*route),
+		stats: make(map[netaddr.IPPrefix]*route),
 	}
 }
 
-func dstIPNet(rm rtnetlink.RouteMessage) netaddr.IPPrefix {
+func prefixFromRM(rm rtnetlink.RouteMessage) netaddr.IPPrefix {
 	switch rm.Family {
 	case unix.AF_INET, unix.AF_INET6:
 	default:
@@ -85,15 +85,13 @@ func dstIPNet(rm rtnetlink.RouteMessage) netaddr.IPPrefix {
 	return netaddr.IPPrefixFrom(ip, rm.DstLength)
 }
 
-func (rs *routeStats) add(rm rtnetlink.RouteMessage) {
+func (rs *routeStats) add(prefix netaddr.IPPrefix) {
 	rs.Lock()
 	defer rs.Unlock()
 
-	dst := dstIPNet(rm)
-	key := dst.String()
-	r := rs.stats[key]
+	r := rs.stats[prefix]
 	if r == nil {
-		r = &route{Dst: dst}
+		r = &route{Dst: prefix}
 		go func(r *route, rs *routeStats) {
 			names, err := net.LookupAddr(r.Dst.IP().String())
 			if err != nil || len(names) == 0 {
@@ -111,16 +109,14 @@ func (rs *routeStats) add(rm rtnetlink.RouteMessage) {
 		r.ReachableSince = time.Time{}
 	}
 
-	rs.stats[key] = r
+	rs.stats[prefix] = r
 }
 
-func (rs *routeStats) del(rm rtnetlink.RouteMessage) {
+func (rs *routeStats) del(prefix netaddr.IPPrefix) {
 	rs.Lock()
 	defer rs.Unlock()
 
-	dst := dstIPNet(rm)
-	key := dst.String()
-	r := rs.stats[key]
+	r := rs.stats[prefix]
 	if r == nil {
 		return
 	}
@@ -129,7 +125,7 @@ func (rs *routeStats) del(rm rtnetlink.RouteMessage) {
 	r.UnreachableSince = time.Time{}
 	r.ReachableSince = time.Now()
 
-	rs.stats[key] = r
+	rs.stats[prefix] = r
 }
 
 func (rs *routeStats) getAll() []route {
@@ -195,11 +191,12 @@ func monitor(done <-chan struct{}, rs *routeStats) {
 				rm.Protocol != 42 {
 				continue
 			}
+			prefix := prefixFromRM(rm)
 			switch m.Header.Type {
 			case unix.RTM_NEWROUTE:
-				rs.add(rm)
+				rs.add(prefix)
 			case unix.RTM_DELROUTE:
-				rs.del(rm)
+				rs.del(prefix)
 			}
 		}
 	}