From a163be35844b22f2b842a3d8736c83e86a65dfdb Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Wed, 7 Feb 2024 14:29:20 +0700
Subject: [PATCH] cmd/cli: preserve static DNS on Windows/Mac

---
 cmd/cli/os_darwin.go  |  5 +++++
 cmd/cli/os_windows.go |  5 +++++
 cmd/cli/prog.go       | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/cmd/cli/os_darwin.go b/cmd/cli/os_darwin.go
index 5931819..1c61efd 100644
--- a/cmd/cli/os_darwin.go
+++ b/cmd/cli/os_darwin.go
@@ -44,6 +44,11 @@ func setDNS(iface *net.Interface, nameservers []string) error {
 
 // TODO(cuonglm): use system API
 func resetDNS(iface *net.Interface) error {
+	if ns := savedNameservers(iface); len(ns) > 0 {
+		if err := setDNS(iface, ns); err == nil {
+			return nil
+		}
+	}
 	cmd := "networksetup"
 	args := []string{"-setdnsservers", iface.Name, "empty"}
 
diff --git a/cmd/cli/os_windows.go b/cmd/cli/os_windows.go
index 694643f..f075ded 100644
--- a/cmd/cli/os_windows.go
+++ b/cmd/cli/os_windows.go
@@ -69,6 +69,11 @@ func resetDNS(iface *net.Interface) error {
 		}
 	})
 
+	if ns := savedNameservers(iface); len(ns) > 0 {
+		if err := setDNS(iface, ns); err == nil {
+			return nil
+		}
+	}
 	if ctrldnet.SupportsIPv6ListenLocal() {
 		if output, err := netsh("interface", "ipv6", "set", "dnsserver", strconv.Itoa(iface.Index), "dhcp"); err != nil {
 			mainLog.Load().Warn().Err(err).Msgf("failed to reset ipv6 DNS: %s", string(output))
diff --git a/cmd/cli/prog.go b/cmd/cli/prog.go
index 92eadc8..e25dcfb 100644
--- a/cmd/cli/prog.go
+++ b/cmd/cli/prog.go
@@ -462,6 +462,7 @@ func (p *prog) setDNS() {
 	if needRFC1918Listeners(lc) {
 		nameservers = append(nameservers, ctrld.Rfc1918Addresses()...)
 	}
+	saveCurrentDNS(netIface)
 	if err := setDNS(netIface, nameservers); err != nil {
 		logger.Error().Err(err).Msgf("could not set DNS for interface")
 		return
@@ -469,6 +470,7 @@ func (p *prog) setDNS() {
 	logger.Debug().Msg("setting DNS successfully")
 	if allIfaces {
 		withEachPhysicalInterfaces(netIface.Name, "set DNS", func(i *net.Interface) error {
+			saveCurrentDNS(i)
 			return setDNS(i, nameservers)
 		})
 	}
@@ -724,3 +726,35 @@ func requiredMultiNICsConfig() bool {
 		return false
 	}
 }
+
+// saveCurrentDNS saves the current DNS settings for restoring later.
+// Only works on Windows and Mac.
+func saveCurrentDNS(iface *net.Interface) {
+	switch runtime.GOOS {
+	case "windows", "darwin":
+	default:
+		return
+	}
+	ns := currentDNS(iface)
+	if len(ns) == 0 {
+		return
+	}
+	file := savedDnsSettingsFilePath(iface)
+	if err := os.WriteFile(file, []byte(strings.Join(ns, ",")), 0600); err != nil {
+		mainLog.Load().Err(err).Msgf("could not save DNS settings for iface: %s", iface.Name)
+	}
+}
+
+// savedDnsSettingsFilePath returns the path to saved DNS settings of the given interface.
+func savedDnsSettingsFilePath(iface *net.Interface) string {
+	return absHomeDir(".dns_" + iface.Name)
+}
+
+// savedNameservers returns the static DNS nameservers of the given interface.
+func savedNameservers(iface *net.Interface) []string {
+	file := savedDnsSettingsFilePath(iface)
+	if data, err := os.ReadFile(file); err != nil && len(data) > 0 {
+		return strings.Split(string(data), ",")
+	}
+	return nil
+}