From cc6ae290f84661e9b9e7e7ad6c06520f1499eb9b Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Thu, 7 Mar 2024 22:31:26 +0700
Subject: [PATCH] internal/clientinfo: use last seen IP for NDP discovery

---
 internal/clientinfo/ndp.go       | 24 ++++++++++++++++++------
 internal/clientinfo/ndp_linux.go |  3 +--
 internal/clientinfo/ndp_test.go  | 11 +++++++----
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/internal/clientinfo/ndp.go b/internal/clientinfo/ndp.go
index 600b54c..81edc66 100644
--- a/internal/clientinfo/ndp.go
+++ b/internal/clientinfo/ndp.go
@@ -69,6 +69,21 @@ func (nd *ndpDiscover) List() []string {
 	return ips
 }
 
+// saveInfo saves ip and mac info to mapping table.
+// Last seen ip address will override the old one,
+func (nd *ndpDiscover) saveInfo(ip, mac string) {
+	// Store ip => map mapping,
+	nd.mac.Store(ip, mac)
+	// If there is old ip => mac mapping, delete it.
+	old, ok := nd.ip.Load(mac)
+	if ok {
+		oldIP := old.(string)
+		nd.mac.Delete(oldIP)
+	}
+	// Store mac => ip mapping.
+	nd.ip.Store(mac, ip)
+}
+
 // listen listens on ipv6 link local for Neighbor Solicitation message
 // to update new neighbors information to ndp table.
 func (nd *ndpDiscover) listen(ctx context.Context) {
@@ -111,8 +126,7 @@ func (nd *ndpDiscover) listen(ctx context.Context) {
 		for _, opt := range am.Options {
 			if lla, ok := opt.(*ndp.LinkLayerAddress); ok {
 				mac := lla.Addr.String()
-				nd.mac.Store(fromIP, mac)
-				nd.ip.Store(mac, fromIP)
+				nd.saveInfo(fromIP, mac)
 			}
 		}
 	}
@@ -127,8 +141,7 @@ func (nd *ndpDiscover) scanWindows(r io.Reader) {
 			continue
 		}
 		if mac := parseMAC(fields[1]); mac != "" {
-			nd.mac.Store(fields[0], mac)
-			nd.ip.Store(mac, fields[0])
+			nd.saveInfo(fields[0], mac)
 		}
 	}
 }
@@ -147,8 +160,7 @@ func (nd *ndpDiscover) scanUnix(r io.Reader) {
 			if idx := strings.IndexByte(ip, '%'); idx != -1 {
 				ip = ip[:idx]
 			}
-			nd.mac.Store(ip, mac)
-			nd.ip.Store(mac, ip)
+			nd.saveInfo(ip, mac)
 		}
 	}
 }
diff --git a/internal/clientinfo/ndp_linux.go b/internal/clientinfo/ndp_linux.go
index 713a7e3..cc99675 100644
--- a/internal/clientinfo/ndp_linux.go
+++ b/internal/clientinfo/ndp_linux.go
@@ -17,8 +17,7 @@ func (nd *ndpDiscover) scan() {
 	for _, n := range neighs {
 		ip := n.IP.String()
 		mac := n.HardwareAddr.String()
-		nd.mac.Store(ip, mac)
-		nd.ip.Store(mac, ip)
+		nd.saveInfo(ip, mac)
 	}
 
 }
diff --git a/internal/clientinfo/ndp_test.go b/internal/clientinfo/ndp_test.go
index c8cd398..5a20555 100644
--- a/internal/clientinfo/ndp_test.go
+++ b/internal/clientinfo/ndp_test.go
@@ -45,12 +45,15 @@ ff02::c                                       33-33-00-00-00-0c  Permanent
 	nd.scanWindows(r)
 
 	count := 0
+	expectedCount := 5
 	nd.mac.Range(func(key, value any) bool {
 		count++
 		return true
 	})
-	if count != 6 {
-		t.Errorf("unexpected count, want 6, got: %d", count)
+	// There are 2 entries for 60-57-47-21-dd-00 in the table, but (*ndpDiscover).saveInfo
+	// only saves the last one, that's why the expected count number is 5.
+	if count != expectedCount {
+		t.Errorf("unexpected count, want %d, got: %d", expectedCount, count)
 	}
 
 	count = 0
@@ -58,7 +61,7 @@ ff02::c                                       33-33-00-00-00-0c  Permanent
 		count++
 		return true
 	})
-	if count != 5 {
-		t.Errorf("unexpected count, want 5, got: %d", count)
+	if count != expectedCount {
+		t.Errorf("unexpected count, want %d, got: %d", expectedCount, count)
 	}
 }