From dd9f2465bee4f3411da606c0923902980533b8ed Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Thu, 18 Apr 2024 22:13:20 +0700
Subject: [PATCH] internal/clientinfo: map ::1 to the right host MAC address

So queries originating from host using ::1 as source will be recognized
properly, and treated the same as other queries from host itself.
---
 cmd/cli/dns_proxy.go               |  6 +-----
 internal/clientinfo/client_info.go | 13 +++++++++++++
 internal/clientinfo/dhcp.go        | 10 ++++++----
 internal/clientinfo/hostsfile.go   |  2 +-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/cmd/cli/dns_proxy.go b/cmd/cli/dns_proxy.go
index 0477320..d6d4142 100644
--- a/cmd/cli/dns_proxy.go
+++ b/cmd/cli/dns_proxy.go
@@ -820,11 +820,7 @@ func (p *prog) getClientInfo(remoteIP string, msg *dns.Msg) *ctrld.ClientInfo {
 	// If this is a query from self, but ci.IP is not loopback IP,
 	// try using hostname mapping for lookback IP if presents.
 	if ci.Self {
-		loopbackIP := "127.0.0.1"
-		if isV6 {
-			loopbackIP = "::1"
-		}
-		if name := p.ciTable.LookupHostname(loopbackIP, ""); name != "" {
+		if name := p.ciTable.LocalHostname(isV6); name != "" {
 			ci.Hostname = name
 		}
 	}
diff --git a/internal/clientinfo/client_info.go b/internal/clientinfo/client_info.go
index 225b9cb..646994c 100644
--- a/internal/clientinfo/client_info.go
+++ b/internal/clientinfo/client_info.go
@@ -14,6 +14,11 @@ import (
 	"github.com/Control-D-Inc/ctrld/internal/controld"
 )
 
+const (
+	ipV4Loopback = "127.0.0.1"
+	ipv6Loopback = "::1"
+)
+
 // IpResolver is the interface for retrieving IP from Mac.
 type IpResolver interface {
 	fmt.Stringer
@@ -322,6 +327,14 @@ func (t *Table) LookupRFC1918IPv4(mac string) string {
 	return ""
 }
 
+// LocalHostname returns the localhost hostname associated with loopback IP.
+func (t *Table) LocalHostname(v6 bool) string {
+	if v6 {
+		return t.LookupHostname(ipv6Loopback, "")
+	}
+	return t.LookupHostname(ipV4Loopback, "")
+}
+
 type macEntry struct {
 	mac string
 	src string
diff --git a/internal/clientinfo/dhcp.go b/internal/clientinfo/dhcp.go
index 9d1f339..147ad29 100644
--- a/internal/clientinfo/dhcp.go
+++ b/internal/clientinfo/dhcp.go
@@ -353,8 +353,8 @@ func (d *dhcp) addSelf() {
 		return
 	}
 	hostname = normalizeHostname(hostname)
-	d.ip2name.Store("127.0.0.1", hostname)
-	d.ip2name.Store("::1", hostname)
+	d.ip2name.Store(ipV4Loopback, hostname)
+	d.ip2name.Store(ipv6Loopback, hostname)
 	found := false
 	interfaces.ForeachInterface(func(i interfaces.Interface, prefixes []netip.Prefix) {
 		mac := i.HardwareAddr.String()
@@ -375,15 +375,17 @@ func (d *dhcp) addSelf() {
 			d.mac.Store(ip.String(), mac)
 			d.ip.Store(mac, ip.String())
 			if ip.To4() != nil {
-				d.mac.Store("127.0.0.1", mac)
+				d.mac.Store(ipV4Loopback, mac)
 			} else {
-				d.mac.Store("::1", mac)
+				d.mac.Store(ipv6Loopback, mac)
 			}
 			d.mac2name.Store(mac, hostname)
 			d.ip2name.Store(ip.String(), hostname)
 			// If we have self IP set, and this IP is it, use this IP only.
 			if ip.String() == d.selfIP {
 				found = true
+				d.mac.Store(ipV4Loopback, mac)
+				d.mac.Store(ipv6Loopback, mac)
 			}
 		}
 	})
diff --git a/internal/clientinfo/hostsfile.go b/internal/clientinfo/hostsfile.go
index c758f3b..d96229d 100644
--- a/internal/clientinfo/hostsfile.go
+++ b/internal/clientinfo/hostsfile.go
@@ -95,7 +95,7 @@ func (hf *hostsFile) LookupHostnameByIP(ip string) string {
 	hf.mu.Lock()
 	defer hf.mu.Unlock()
 	if names := hf.m[ip]; len(names) > 0 {
-		isLoopback := ip == "127.0.0.1" || ip == "::1"
+		isLoopback := ip == ipV4Loopback || ip == ipv6Loopback
 		for _, hostname := range names {
 			name := normalizeHostname(hostname)
 			// Ignoring ipv4/ipv6 loopback entry.