From a95d50c0afc72a7737153aaee4e4b340c9343884 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Mon, 13 Jan 2025 20:03:56 +0700
Subject: [PATCH] cmd/cli: ensure set/reset DNS is done before checking OS
 resolver

Otherwise, new DNS settings could be reverted by dns watchers, causing
the checking will be always false.
---
 cmd/cli/dns_proxy.go        |  2 +-
 cmd/cli/prog.go             |  2 +-
 cmd/cli/resolvconf.go       |  2 +-
 cmd/cli/upstream_monitor.go | 23 ++++++++++++++++++-----
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/cmd/cli/dns_proxy.go b/cmd/cli/dns_proxy.go
index 623c7d2..631b0e3 100644
--- a/cmd/cli/dns_proxy.go
+++ b/cmd/cli/dns_proxy.go
@@ -552,7 +552,7 @@ func (p *prog) proxy(ctx context.Context, req *proxyRequest) *proxyResponse {
 			continue
 		}
 		if p.um.isDown(upstreams[n]) {
-			ctrld.Log(ctx, mainLog.Load().Warn(), "%s is down", upstreams[n])
+			ctrld.Log(ctx, mainLog.Load().Debug(), "%s is down", upstreams[n])
 			continue
 		}
 		answer := resolve(n, upstreamConfig, req.msg)
diff --git a/cmd/cli/prog.go b/cmd/cli/prog.go
index 9ff04bc..9de6b48 100644
--- a/cmd/cli/prog.go
+++ b/cmd/cli/prog.go
@@ -728,7 +728,7 @@ func (p *prog) dnsWatchdog(iface *net.Interface, nameservers []string, allIfaces
 			mainLog.Load().Debug().Msg("stop dns watchdog")
 			return
 		case <-ticker.C:
-			if p.leakingQuery.Load() {
+			if p.leakingQuery.Load() || p.um.isChecking(upstreamOS) {
 				return
 			}
 			if dnsChanged(iface, ns) {
diff --git a/cmd/cli/resolvconf.go b/cmd/cli/resolvconf.go
index 6df7be6..6bd8c2a 100644
--- a/cmd/cli/resolvconf.go
+++ b/cmd/cli/resolvconf.go
@@ -40,7 +40,7 @@ func (p *prog) watchResolvConf(iface *net.Interface, ns []netip.Addr, setDnsFn f
 			mainLog.Load().Debug().Msgf("stopping watcher for %s", resolvConfPath)
 			return
 		case event, ok := <-watcher.Events:
-			if p.leakingQuery.Load() {
+			if p.leakingQuery.Load() || p.um.isChecking(upstreamOS) {
 				return
 			}
 			if !ok {
diff --git a/cmd/cli/upstream_monitor.go b/cmd/cli/upstream_monitor.go
index 3400b60..86c191d 100644
--- a/cmd/cli/upstream_monitor.go
+++ b/cmd/cli/upstream_monitor.go
@@ -60,6 +60,14 @@ func (um *upstreamMonitor) isDown(upstream string) bool {
 	return um.down[upstream]
 }
 
+// isChecking reports whether the given upstream is being checked.
+func (um *upstreamMonitor) isChecking(upstream string) bool {
+	um.mu.Lock()
+	defer um.mu.Unlock()
+
+	return um.checking[upstream]
+}
+
 // reset marks an upstream as up and set failed queries counter to zero.
 func (um *upstreamMonitor) reset(upstream string) {
 	um.mu.Lock()
@@ -86,9 +94,10 @@ func (p *prog) checkUpstream(upstream string, uc *ctrld.UpstreamConfig) {
 		p.um.mu.Unlock()
 	}()
 
-	if uc.Type == ctrld.ResolverTypeOS {
-		ns := ctrld.InitializeOsResolver()
-		mainLog.Load().Debug().Msgf("re-initializing OS resolver with nameservers: %v", ns)
+	isOsResolver := uc.Type == ctrld.ResolverTypeOS
+	if isOsResolver {
+		p.resetDNS()
+		defer p.setDNS()
 	}
 	resolver, err := ctrld.NewResolver(uc)
 	if err != nil {
@@ -105,12 +114,16 @@ func (p *prog) checkUpstream(upstream string, uc *ctrld.UpstreamConfig) {
 		ctx, cancel := context.WithTimeout(context.Background(), timeout)
 		defer cancel()
 		uc.ReBootstrap()
+		if isOsResolver {
+			ctrld.InitializeOsResolver()
+		}
 		_, err := resolver.Resolve(ctx, msg)
 		return err
 	}
+	mainLog.Load().Warn().Msgf("upstream %q is offline", uc.Endpoint)
 	for {
 		if err := check(); err == nil {
-			mainLog.Load().Debug().Msgf("upstream %q is online", uc.Endpoint)
+			mainLog.Load().Warn().Msgf("upstream %q is online", uc.Endpoint)
 			p.um.reset(upstream)
 			if p.leakingQuery.CompareAndSwap(true, false) {
 				p.leakingQueryMu.Lock()
@@ -120,7 +133,7 @@ func (p *prog) checkUpstream(upstream string, uc *ctrld.UpstreamConfig) {
 			}
 			return
 		} else {
-			mainLog.Load().Debug().Msgf("upstream %q is offline: %v", uc.Endpoint, err)
+			mainLog.Load().Debug().Msgf("checked upstream %q failed: %v", uc.Endpoint, err)
 		}
 		time.Sleep(checkUpstreamBackoffSleep)
 	}