From a5025e35ea6f189e87060988d71362a868963bc7 Mon Sep 17 00:00:00 2001
From: Cuong Manh Le <cuong@windscribe.com>
Date: Wed, 27 Mar 2024 17:45:17 +0700
Subject: [PATCH] cmd/cli: add internal domain test query during self-check

So it's clear that client could be reached ctrld's listener or not.
---
 cmd/cli/cli.go       | 26 +++++++++++++++++++++++---
 cmd/cli/dns_proxy.go | 22 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/cmd/cli/cli.go b/cmd/cli/cli.go
index d438510..aab4440 100644
--- a/cmd/cli/cli.go
+++ b/cmd/cli/cli.go
@@ -46,6 +46,9 @@ import (
 	"github.com/Control-D-Inc/ctrld/internal/router"
 )
 
+// selfCheckInternalTestDomain is used for testing ctrld self response to clients.
+const selfCheckInternalTestDomain = "ctrld" + loopTestDomain
+
 var (
 	version = "dev"
 	commit  = "none"
@@ -1552,6 +1555,13 @@ func defaultIfaceName() string {
 // selfCheckStatus performs the end-to-end DNS test by sending query to ctrld listener.
 // It returns a boolean to indicate whether the check is succeeded, the actual status
 // of ctrld service, and an additional error if any.
+//
+// We perform two tests:
+//
+// - Internal testing, ensuring query could be sent from client -> ctrld.
+// - External testing, ensuring query could be sent from ctrld -> upstream.
+//
+// Self-check is considered success only if both tests are ok.
 func selfCheckStatus(s service.Service) (bool, service.Status, error) {
 	status, err := s.Status()
 	if err != nil {
@@ -1634,8 +1644,9 @@ func selfCheckStatus(s service.Service) (bool, service.Status, error) {
 	})
 	v.WatchConfig()
 	var (
-		lastAnswer *dns.Msg
-		lastErr    error
+		lastAnswer     *dns.Msg
+		lastErr        error
+		internalTested bool
 	)
 	for i := 0; i < maxAttempts; i++ {
 		mu.Lock()
@@ -1648,6 +1659,9 @@ func selfCheckStatus(s service.Service) (bool, service.Status, error) {
 		mu.Unlock()
 		lc := cfg.FirstListener()
 		domain = cfg.FirstUpstream().VerifyDomain()
+		if !internalTested {
+			domain = selfCheckInternalTestDomain
+		}
 		if domain == "" {
 			continue
 		}
@@ -1657,7 +1671,13 @@ func selfCheckStatus(s service.Service) (bool, service.Status, error) {
 		m.RecursionDesired = true
 		r, _, exErr := exchangeContextWithTimeout(c, time.Second, m, net.JoinHostPort(lc.IP, strconv.Itoa(lc.Port)))
 		if r != nil && r.Rcode == dns.RcodeSuccess && len(r.Answer) > 0 {
-			mainLog.Load().Debug().Msgf("self-check against %q succeeded", domain)
+			internalTested = domain == selfCheckInternalTestDomain
+			if internalTested {
+				mainLog.Load().Debug().Msgf("internal self-check against %q succeeded", domain)
+				continue // internal domain test ok, continue with external test.
+			} else {
+				mainLog.Load().Debug().Msgf("external self-check against %q succeeded", domain)
+			}
 			return true, status, nil
 		}
 		// Return early if this is a connection refused.
diff --git a/cmd/cli/dns_proxy.go b/cmd/cli/dns_proxy.go
index 4cd4641..790a993 100644
--- a/cmd/cli/dns_proxy.go
+++ b/cmd/cli/dns_proxy.go
@@ -101,6 +101,11 @@ func (p *prog) serveDNS(listenerNum string) error {
 		go p.detectLoop(m)
 		q := m.Question[0]
 		domain := canonicalName(q.Name)
+		if domain == selfCheckInternalTestDomain {
+			answer := resolveInternalDomainTestQuery(ctx, domain, m)
+			_ = w.WriteMsg(answer)
+			return
+		}
 		if _, ok := p.cacheFlushDomainsMap[domain]; ok && p.cache != nil {
 			p.cache.Purge()
 			ctrld.Log(ctx, mainLog.Load().Debug(), "received query %q, local cache is purged", domain)
@@ -953,3 +958,20 @@ func isWanClient(na net.Addr) bool {
 		!ip.IsLinkLocalMulticast() &&
 		!tsaddr.CGNATRange().Contains(ip)
 }
+
+// resolveInternalDomainTestQuery resolves internal test domain query, returning the answer to the caller.
+func resolveInternalDomainTestQuery(ctx context.Context, domain string, m *dns.Msg) *dns.Msg {
+	ctrld.Log(ctx, mainLog.Load().Debug(), "internal domain test query")
+	q := m.Question[0]
+	answer := new(dns.Msg)
+	rrStr := fmt.Sprintf("%s A %s", domain, net.IPv4zero)
+	if q.Qtype == dns.TypeAAAA {
+		rrStr = fmt.Sprintf("%s AAAA %s", domain, net.IPv6zero)
+	}
+	rr, err := dns.NewRR(rrStr)
+	if err == nil {
+		answer.Answer = append(answer.Answer, rr)
+	}
+	answer.SetReply(m)
+	return answer
+}