From fdf60eeed15e650dd1d13de43fbad53f9c3e6f0f Mon Sep 17 00:00:00 2001
From: Stanislav Chzhen <s.chzhen@adguard.com>
Date: Fri, 10 Nov 2023 18:22:09 +0300
Subject: [PATCH] Pull request 2067: 6399-fix-domain-specific-upstreams-check

Closes #6399.

Co-authored-by: Eugene Burkov <E.Burkov@AdGuard.COM>

Squashed commit of the following:

commit 7af223caf83565e8dab34636dcda810a8a7732c7
Merge: b5dd3903f 4b9947da8
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 10 18:14:54 2023 +0300

    Merge branch 'master' into 6399-fix-domain-specific-upstreams-check

commit b5dd3903fe7f86d9dba217acabdc66dd39834a85
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 10 18:06:25 2023 +0300

    dnsforward: imp docs

commit 02b00d2630b2d6d1585c8716d3fd7c785ae19169
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 10 17:35:00 2023 +0300

    dnsforward: add todo

commit 7c542905607ade22b32481ad58b646267aa1787d
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 10 16:42:50 2023 +0300

    dnsforward: fix typo

commit 86259571293aa789836b6b7b2ed8935505ae3734
Author: Stanislav Chzhen <s.chzhen@adguard.com>
Date:   Fri Nov 10 16:35:35 2023 +0300

    dnsforward: fix domain specific upstreams check
---
 internal/dnsforward/http.go      | 50 +++++++++++++++-----------------
 internal/dnsforward/http_test.go | 25 ++++++++++++++++
 2 files changed, 49 insertions(+), 26 deletions(-)

diff --git a/internal/dnsforward/http.go b/internal/dnsforward/http.go
index 99ced3ef..19f06f37 100644
--- a/internal/dnsforward/http.go
+++ b/internal/dnsforward/http.go
@@ -702,18 +702,28 @@ func (err domainSpecificTestError) Error() (msg string) {
 }
 
 // checkDNS parses line, creates DNS upstreams using opts, and checks if the
-// upstreams are exchanging correctly.  It returns a map where key is an
-// upstream address and value is "OK", if the upstream exchanges correctly, or
-// text of the error.
+// upstreams are exchanging correctly.  It saves the result into a sync.Map
+// where key is an upstream address and value is "OK", if the upstream
+// exchanges correctly, or text of the error.  It is intended to be used as a
+// goroutine.
+//
+// TODO(s.chzhen):  Separate to a different structure/file.
 func (s *Server) checkDNS(
 	line string,
 	opts *upstream.Options,
 	check healthCheckFunc,
-) (result map[string]string) {
-	result = map[string]string{}
+	wg *sync.WaitGroup,
+	m *sync.Map,
+) {
+	defer wg.Done()
+	defer log.OnPanic("dnsforward: checking upstreams")
+
 	upstreams, domains, err := separateUpstream(line)
 	if err != nil {
-		return nil
+		err = fmt.Errorf("wrong upstream format: %w", err)
+		m.Store(line, err.Error())
+
+		return
 	}
 
 	specific := len(domains) > 0
@@ -723,7 +733,7 @@ func (s *Server) checkDNS(
 		useDefault, err = validateUpstream(upstreamAddr, domains)
 		if err != nil {
 			err = fmt.Errorf("wrong upstream format: %w", err)
-			result[upstreamAddr] = err.Error()
+			m.Store(upstreamAddr, err.Error())
 
 			continue
 		}
@@ -736,13 +746,11 @@ func (s *Server) checkDNS(
 
 		err = s.checkUpstreamAddr(upstreamAddr, specific, opts, check)
 		if err != nil {
-			result[upstreamAddr] = err.Error()
+			m.Store(upstreamAddr, err.Error())
 		} else {
-			result[upstreamAddr] = "OK"
+			m.Store(upstreamAddr, "OK")
 		}
 	}
-
-	return result
 }
 
 // checkUpstreamAddr creates the DNS upstream using opts and information from
@@ -813,34 +821,24 @@ func (s *Server) handleTestUpstreamDNS(w http.ResponseWriter, r *http.Request) {
 	wg := &sync.WaitGroup{}
 	m := &sync.Map{}
 
-	// TODO(s.chzhen):  Separate to a different structure/file.
-	worker := func(upstreamLine string, check healthCheckFunc) {
-		defer log.OnPanic("dnsforward: checking upstreams")
-
-		res := s.checkDNS(upstreamLine, opts, check)
-		for ups, status := range res {
-			m.Store(ups, status)
-		}
-
-		wg.Done()
-	}
-
 	wg.Add(len(req.Upstreams) + len(req.FallbackDNS) + len(req.PrivateUpstreams))
 
 	for _, ups := range req.Upstreams {
-		go worker(ups, checkDNSUpstreamExc)
+		go s.checkDNS(ups, opts, checkDNSUpstreamExc, wg, m)
 	}
 	for _, ups := range req.FallbackDNS {
-		go worker(ups, checkDNSUpstreamExc)
+		go s.checkDNS(ups, opts, checkDNSUpstreamExc, wg, m)
 	}
 	for _, ups := range req.PrivateUpstreams {
-		go worker(ups, checkPrivateUpstreamExc)
+		go s.checkDNS(ups, opts, checkPrivateUpstreamExc, wg, m)
 	}
 
 	wg.Wait()
 
 	result := map[string]string{}
 	m.Range(func(k, v any) bool {
+		// TODO(e.burkov):  The upstreams used for both common and private
+		// resolving should be reported separately.
 		ups := k.(string)
 		status := v.(string)
 
diff --git a/internal/dnsforward/http_test.go b/internal/dnsforward/http_test.go
index dfd1862a..1418a4a1 100644
--- a/internal/dnsforward/http_test.go
+++ b/internal/dnsforward/http_test.go
@@ -479,6 +479,8 @@ func TestServer_HandleTestUpstreamDNS(t *testing.T) {
 		Host:   newLocalUpstreamListener(t, 0, badHandler).String(),
 	}).String()
 
+	goodAndBadUps := strings.Join([]string{goodUps, badUps}, " ")
+
 	const (
 		upsTimeout = 100 * time.Millisecond
 
@@ -605,6 +607,29 @@ func TestServer_HandleTestUpstreamDNS(t *testing.T) {
 				`dns: id mismatch`,
 		},
 		name: "multiple_domain_specific_upstreams",
+	}, {
+		body: map[string]any{
+			"upstream_dns": []string{"[/domain.example/]/]1.2.3.4"},
+		},
+		wantResp: map[string]any{
+			"[/domain.example/]/]1.2.3.4": `wrong upstream format: ` +
+				`bad upstream for domain "[/domain.example/]/]1.2.3.4": ` +
+				`duplicated separator`,
+		},
+		name: "bad_specification",
+	}, {
+		body: map[string]any{
+			"upstream_dns":     []string{"[/domain.example/]" + goodAndBadUps},
+			"fallback_dns":     []string{"[/domain.example/]" + goodAndBadUps},
+			"private_upstream": []string{"[/domain.example/]" + goodAndBadUps},
+		},
+		wantResp: map[string]any{
+			goodUps: "OK",
+			badUps: `WARNING: couldn't communicate ` +
+				`with upstream: exchanging with ` + badUps + ` over tcp: ` +
+				`dns: id mismatch`,
+		},
+		name: "all_different",
 	}}
 
 	for _, tc := range testCases {