From d6768c4c399a64549aeb222dc26578a246858343 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Tue, 18 Jul 2023 20:52:51 +0700
Subject: [PATCH] internal/clientinfo: use default route IP as self client info

---
 cmd/ctrld/dns_proxy.go             |  3 +++
 cmd/ctrld/prog.go                  | 30 +++++++++++++++++++++++++++++-
 internal/clientinfo/client_info.go | 11 ++++++++---
 internal/clientinfo/dhcp.go        |  9 +++++++++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/cmd/ctrld/dns_proxy.go b/cmd/ctrld/dns_proxy.go
index abacd85..0586151 100644
--- a/cmd/ctrld/dns_proxy.go
+++ b/cmd/ctrld/dns_proxy.go
@@ -549,6 +549,9 @@ func (p *prog) getClientInfo(ip, mac string) *ctrld.ClientInfo {
 	} else {
 		ci.IP = ip
 		ci.Mac = p.ciTable.LookupMac(ip)
+		if ip == "127.0.0.1" || ip == "::1" {
+			ci.IP = p.ciTable.LookupIP(ci.Mac)
+		}
 	}
 	ci.Hostname = p.ciTable.LookupHostname(ci.IP, ci.Mac)
 	return ci
diff --git a/cmd/ctrld/prog.go b/cmd/ctrld/prog.go
index 4e011b2..dd0b5f0 100644
--- a/cmd/ctrld/prog.go
+++ b/cmd/ctrld/prog.go
@@ -11,6 +11,7 @@ import (
 	"syscall"
 
 	"github.com/kardianos/service"
+	"tailscale.com/net/interfaces"
 
 	"github.com/Control-D-Inc/ctrld"
 	"github.com/Control-D-Inc/ctrld/internal/clientinfo"
@@ -109,7 +110,7 @@ func (p *prog) run() {
 		go uc.Ping()
 	}
 
-	p.ciTable = clientinfo.NewTable(&cfg)
+	p.ciTable = clientinfo.NewTable(&cfg, defaultRouteIP())
 	if leaseFile := p.cfg.Service.DHCPLeaseFile; leaseFile != "" {
 		mainLog.Debug().Msgf("watching custom lease file: %s", leaseFile)
 		format := ctrld.LeaseFileFormat(p.cfg.Service.DHCPLeaseFileFormat)
@@ -307,3 +308,30 @@ func errAddrInUse(err error) bool {
 	}
 	return errors.Is(opErr.Err, syscall.EADDRINUSE)
 }
+
+// defaultRouteIP returns IP string of the default route if present, prefer IPv4 over IPv6.
+func defaultRouteIP() string {
+	if dr, err := interfaces.DefaultRoute(); err == nil {
+		if netIface, err := netInterface(dr.InterfaceName); err == nil {
+			addrs, _ := netIface.Addrs()
+			do := func(v4 bool) net.IP {
+				for _, addr := range addrs {
+					if netIP, ok := addr.(*net.IPNet); ok && netIP.IP.IsPrivate() {
+						if v4 {
+							return netIP.IP.To4()
+						}
+						return netIP.IP
+					}
+				}
+				return nil
+			}
+			if ip := do(true); ip != nil {
+				return ip.String()
+			}
+			if ip := do(false); ip != nil {
+				return ip.String()
+			}
+		}
+	}
+	return ""
+}
diff --git a/internal/clientinfo/client_info.go b/internal/clientinfo/client_info.go
index ad17121..6d220e1 100644
--- a/internal/clientinfo/client_info.go
+++ b/internal/clientinfo/client_info.go
@@ -49,10 +49,15 @@ type Table struct {
 	mdns   *mdns
 	cfg    *ctrld.Config
 	quitCh chan struct{}
+	selfIP string
 }
 
-func NewTable(cfg *ctrld.Config) *Table {
-	return &Table{cfg: cfg, quitCh: make(chan struct{})}
+func NewTable(cfg *ctrld.Config, selfIP string) *Table {
+	return &Table{
+		cfg:    cfg,
+		quitCh: make(chan struct{}),
+		selfIP: selfIP,
+	}
 }
 
 func (t *Table) AddLeaseFile(name string, format ctrld.LeaseFileFormat) {
@@ -88,7 +93,7 @@ func (t *Table) Init() {
 		}
 	}
 	if t.discoverDHCP() {
-		t.dhcp = &dhcp{}
+		t.dhcp = &dhcp{selfIP: t.selfIP}
 		ctrld.ProxyLog.Debug().Msg("start dhcp discovery")
 		if err := t.dhcp.refresh(); err != nil {
 			ctrld.ProxyLog.Error().Err(err).Msg("could not init DHCP discover")
diff --git a/internal/clientinfo/dhcp.go b/internal/clientinfo/dhcp.go
index 5ca28ec..1aca295 100644
--- a/internal/clientinfo/dhcp.go
+++ b/internal/clientinfo/dhcp.go
@@ -25,6 +25,7 @@ type dhcp struct {
 	mac      sync.Map // ip  => mac
 
 	watcher *fsnotify.Watcher
+	selfIP  string
 }
 
 func (d *dhcp) refresh() error {
@@ -229,6 +230,7 @@ func (d *dhcp) addSelf() {
 	hostname = normalizeHostname(hostname)
 	d.ip2name.Store("127.0.0.1", hostname)
 	d.ip2name.Store("::1", hostname)
+	found := false
 	interfaces.ForeachInterface(func(i interfaces.Interface, prefixes []netip.Prefix) {
 		mac := i.HardwareAddr.String()
 		// Skip loopback interfaces, info was stored above.
@@ -237,6 +239,9 @@ func (d *dhcp) addSelf() {
 		}
 		addrs, _ := i.Addrs()
 		for _, addr := range addrs {
+			if found {
+				return
+			}
 			ipNet, ok := addr.(*net.IPNet)
 			if !ok {
 				continue
@@ -251,6 +256,10 @@ func (d *dhcp) addSelf() {
 			}
 			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
+			}
 		}
 	})
 }