From 5d65416227637b99cd947f3f437d8efeaf7fff30 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Thu, 25 Jan 2024 17:48:34 +0700
Subject: [PATCH] internal/clientinfo: fill empty hostname based on MAC address

An interface may have multiple MAC addresses, that leads to the problem
when looking up hostname for its multiple <ip, mac> pairs, because the
"ip" map, which storing "mac => ip" mapping can only store 1 entry. It
ends up returns an empty hostname for a known MAC address.

Fixing this by filling empty hostname based on clients which is already
listed, ensuring all clients with the same MAC address will have the
same hostname information.
---
 internal/clientinfo/client_info.go      |  7 +++++++
 internal/clientinfo/client_info_test.go | 28 +++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/internal/clientinfo/client_info.go b/internal/clientinfo/client_info.go
index 7f41fca..1a775ea 100644
--- a/internal/clientinfo/client_info.go
+++ b/internal/clientinfo/client_info.go
@@ -384,6 +384,7 @@ func (t *Table) ListClients() []*Client {
 			}
 		}
 	}
+	clientsByMAC := make(map[string]*Client)
 	for ip := range ipMap {
 		c := ipMap[ip]
 		for _, e := range t.lookupMacAll(ip) {
@@ -397,6 +398,7 @@ func (t *Table) ListClients() []*Client {
 		for _, e := range t.lookupHostnameAll(ip, c.Mac) {
 			if c.Hostname == "" && e.name != "" {
 				c.Hostname = e.name
+				clientsByMAC[c.Mac] = c
 			}
 			if e.name != "" {
 				c.Source[e.src] = struct{}{}
@@ -405,6 +407,11 @@ func (t *Table) ListClients() []*Client {
 	}
 	clients := make([]*Client, 0, len(ipMap))
 	for _, c := range ipMap {
+		// If we found a client with empty hostname, use hostname from
+		// an existed client which has the same MAC address.
+		if cFromMac := clientsByMAC[c.Mac]; cFromMac != nil && c.Hostname == "" {
+			c.Hostname = cFromMac.Hostname
+		}
 		clients = append(clients, c)
 	}
 	return clients
diff --git a/internal/clientinfo/client_info_test.go b/internal/clientinfo/client_info_test.go
index e6575f2..b5bdfa5 100644
--- a/internal/clientinfo/client_info_test.go
+++ b/internal/clientinfo/client_info_test.go
@@ -44,3 +44,31 @@ func TestTable_LookupRFC1918IPv4(t *testing.T) {
 		t.Fatalf("unexpected result, want: %s, got: %s", rfc1918IPv4, got)
 	}
 }
+
+func TestTable_ListClients(t *testing.T) {
+	mac := "74:56:3c:44:eb:5e"
+	ipv6_1 := "2405:4803:a04b:4190:fbe9:cd14:d522:bbae"
+	ipv6_2 := "2405:4803:a04b:4190:fbe9:cd14:d522:bbab"
+	table := &Table{}
+
+	// NDP init.
+	table.ndp = &ndpDiscover{}
+	table.ndp.mac.Store(ipv6_1, mac)
+	table.ndp.mac.Store(ipv6_2, mac)
+	table.ndp.ip.Store(mac, ipv6_1)
+	table.ndp.ip.Store(mac, ipv6_2)
+	table.ipResolvers = append(table.ipResolvers, table.ndp)
+	table.macResolvers = append(table.macResolvers, table.ndp)
+
+	hostname := "foo"
+	// mdns init.
+	table.mdns = &mdns{}
+	table.mdns.name.Store(ipv6_2, hostname)
+	table.hostnameResolvers = append(table.hostnameResolvers, table.mdns)
+
+	for _, c := range table.ListClients() {
+		if c.Hostname != hostname {
+			t.Fatalf("missing hostname for client: %v", c)
+		}
+	}
+}