From cebfd12d5c5f84ec8295bad6d82a9346448df89f Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Wed, 6 Dec 2023 15:09:13 +0700
Subject: [PATCH] internal/clientinfo: ensure RFC1918 address is chosen over
 others

---
 internal/clientinfo/dhcp.go      | 32 +++++++++++++++++++++++---------
 internal/clientinfo/dhcp_test.go | 12 ++++++++++++
 2 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/internal/clientinfo/dhcp.go b/internal/clientinfo/dhcp.go
index a103263..ebbeb77 100644
--- a/internal/clientinfo/dhcp.go
+++ b/internal/clientinfo/dhcp.go
@@ -8,6 +8,7 @@ import (
 	"net"
 	"net/netip"
 	"os"
+	"sort"
 	"strings"
 	"sync"
 
@@ -138,20 +139,33 @@ func (d *dhcp) lookupIPByHostname(name string, v6 bool) string {
 	if d == nil {
 		return ""
 	}
-	var ip string
+	var (
+		rfc1918Addrs []netip.Addr
+		others       []netip.Addr
+	)
 	d.ip2name.Range(func(key, value any) bool {
-		if value == name {
-			if addr, err := netip.ParseAddr(key.(string)); err == nil && addr.Is6() == v6 {
-				ip = addr.String()
-				if addr.IsLoopback() { // Continue searching if this is loopback address.
-					return true
-				}
-				return false
+		if value != name {
+			return true
+		}
+		if addr, err := netip.ParseAddr(key.(string)); err == nil && addr.Is6() == v6 {
+			if addr.IsPrivate() {
+				rfc1918Addrs = append(rfc1918Addrs, addr)
+			} else {
+				others = append(others, addr)
 			}
 		}
 		return true
 	})
-	return ip
+	result := [][]netip.Addr{rfc1918Addrs, others}
+	for _, addrs := range result {
+		if len(addrs) > 0 {
+			sort.Slice(addrs, func(i, j int) bool {
+				return addrs[i].Less(addrs[j])
+			})
+			return addrs[0].String()
+		}
+	}
+	return ""
 }
 
 // AddLeaseFile adds given lease file for reading/watching clients info.
diff --git a/internal/clientinfo/dhcp_test.go b/internal/clientinfo/dhcp_test.go
index af3a168..359f441 100644
--- a/internal/clientinfo/dhcp_test.go
+++ b/internal/clientinfo/dhcp_test.go
@@ -86,3 +86,15 @@ lease 192.168.1.2 {
 		})
 	}
 }
+
+func Test_dhcp_lookupIPByHostname(t *testing.T) {
+	d := &dhcp{}
+	want := "192.168.1.123"
+	d.ip2name.Store(want, "foo")
+	d.ip2name.Store("127.0.0.1", "foo")
+	d.ip2name.Store("169.254.123.123", "foo")
+
+	if got := d.lookupIPByHostname("foo", false); got != want {
+		t.Fatalf("unexpected result, want: %s, got: %s", want, got)
+	}
+}