From a007394f6003f6fa8bc2ac5d2687bbeed39838fa Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Wed, 7 Aug 2024 22:55:31 +0700
Subject: [PATCH] cmd/cli: ensure goroutines that check DNS terminated

So changes to DNS after ctrld stopped won't be reverted by the goroutine
itself. The problem happens rarely on darwin, because networksetup
command won't propagate config to /etc/resolv.conf if there is no
changes between multiple running.
---
 cmd/cli/cli.go        |  4 ++++
 cmd/cli/prog.go       | 53 +++++++++++++++++++++++++++----------------
 cmd/cli/resolvconf.go |  5 +++-
 3 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/cmd/cli/cli.go b/cmd/cli/cli.go
index dcefd11..2317840 100644
--- a/cmd/cli/cli.go
+++ b/cmd/cli/cli.go
@@ -1333,6 +1333,10 @@ func run(appCallback *AppCallback, stopCh chan struct{}) {
 
 	close(waitCh)
 	<-stopCh
+
+	// Wait goroutines which watches/manipulates DNS settings terminated,
+	// ensuring that changes to DNS since here won't be reverted.
+	p.dnsWg.Wait()
 	for _, f := range p.onStopped {
 		f()
 	}
diff --git a/cmd/cli/prog.go b/cmd/cli/prog.go
index f9b9f06..8ee0df1 100644
--- a/cmd/cli/prog.go
+++ b/cmd/cli/prog.go
@@ -78,6 +78,7 @@ type prog struct {
 	csSetDnsDone    chan struct{}
 	csSetDnsOk      bool
 	dnsWatchDogOnce sync.Once
+	dnsWg           sync.WaitGroup
 
 	cfg                  *ctrld.Config
 	localUpstreams       []string
@@ -596,7 +597,11 @@ func (p *prog) setDNS() {
 		for i := range nameservers {
 			servers[i] = netip.MustParseAddr(nameservers[i])
 		}
-		go watchResolvConf(netIface, servers, setResolvConf)
+		p.dnsWg.Add(1)
+		go func() {
+			defer p.dnsWg.Done()
+			p.watchResolvConf(netIface, servers, setResolvConf)
+		}()
 	}
 	if allIfaces {
 		withEachPhysicalInterfaces(netIface.Name, "set DNS", func(i *net.Interface) error {
@@ -604,7 +609,11 @@ func (p *prog) setDNS() {
 		})
 	}
 	if p.dnsWatchdogEnabled() {
-		go p.dnsWatchdog(netIface, nameservers, allIfaces)
+		p.dnsWg.Add(1)
+		go func() {
+			defer p.dnsWg.Done()
+			p.dnsWatchdog(netIface, nameservers, allIfaces)
+		}()
 	}
 }
 
@@ -639,24 +648,30 @@ func (p *prog) dnsWatchdog(iface *net.Interface, nameservers []string, allIfaces
 		slices.Sort(ns)
 		ticker := time.NewTicker(p.dnsWatchdogDuration())
 		logger := mainLog.Load().With().Str("iface", iface.Name).Logger()
-		for range ticker.C {
-			if dnsChanged(iface, ns) {
-				logger.Debug().Msg("DNS settings were changed, re-applying settings")
-				if err := setDNS(iface, ns); err != nil {
-					mainLog.Load().Error().Err(err).Str("iface", iface.Name).Msgf("could not re-apply DNS settings")
-				}
-			}
-			if allIfaces {
-				withEachPhysicalInterfaces(iface.Name, "re-applying DNS", func(i *net.Interface) error {
-					if dnsChanged(i, ns) {
-						if err := setDnsIgnoreUnusableInterface(i, nameservers); err != nil {
-							mainLog.Load().Error().Err(err).Str("iface", i.Name).Msgf("could not re-apply DNS settings")
-						} else {
-							mainLog.Load().Debug().Msgf("re-applying DNS for interface %q successfully", i.Name)
-						}
+		for {
+			select {
+			case <-p.stopCh:
+				mainLog.Load().Debug().Msg("stop dns watchdog")
+				return
+			case <-ticker.C:
+				if dnsChanged(iface, ns) {
+					logger.Debug().Msg("DNS settings were changed, re-applying settings")
+					if err := setDNS(iface, ns); err != nil {
+						mainLog.Load().Error().Err(err).Str("iface", iface.Name).Msgf("could not re-apply DNS settings")
 					}
-					return nil
-				})
+				}
+				if allIfaces {
+					withEachPhysicalInterfaces(iface.Name, "re-applying DNS", func(i *net.Interface) error {
+						if dnsChanged(i, ns) {
+							if err := setDnsIgnoreUnusableInterface(i, nameservers); err != nil {
+								mainLog.Load().Error().Err(err).Str("iface", i.Name).Msgf("could not re-apply DNS settings")
+							} else {
+								mainLog.Load().Debug().Msgf("re-applying DNS for interface %q successfully", i.Name)
+							}
+						}
+						return nil
+					})
+				}
 			}
 		}
 	})
diff --git a/cmd/cli/resolvconf.go b/cmd/cli/resolvconf.go
index 0b53af0..6196487 100644
--- a/cmd/cli/resolvconf.go
+++ b/cmd/cli/resolvconf.go
@@ -10,7 +10,7 @@ import (
 
 // watchResolvConf watches any changes to /etc/resolv.conf file,
 // and reverting to the original config set by ctrld.
-func watchResolvConf(iface *net.Interface, ns []netip.Addr, setDnsFn func(iface *net.Interface, ns []netip.Addr) error) {
+func (p *prog) watchResolvConf(iface *net.Interface, ns []netip.Addr, setDnsFn func(iface *net.Interface, ns []netip.Addr) error) {
 	resolvConfPath := "/etc/resolv.conf"
 	// Evaluating symbolics link to watch the target file that /etc/resolv.conf point to.
 	if rp, _ := filepath.EvalSymlinks(resolvConfPath); rp != "" {
@@ -34,6 +34,9 @@ func watchResolvConf(iface *net.Interface, ns []netip.Addr, setDnsFn func(iface
 
 	for {
 		select {
+		case <-p.stopCh:
+			mainLog.Load().Debug().Msgf("stopping watcher for %s", resolvConfPath)
+			return
 		case event, ok := <-watcher.Events:
 			if !ok {
 				return