Add configuration flag to set if RecursionDesired should be set on health checkers in Forward-plugin (#3679)

* Make the RD-flag in health-checks in the Forward-plugin configurable

Introduces a new configuration flag; `health_check_non_recursive`. This
flag makes the health-checker do non-recursive requests when checking
the health of upstream servers.

Signed-off-by: Geir Haugom <ghagit@haugom.org>
Signed-off-by: Christian Tryti <ctryti@gmail.com>

* Changes after feedback from reviewer

* Better tests of health-checks with and without recursion
* Removed the health_check_non_recursive configuration in favor of
extending the existing health_check configuration. Now supports an
optional `no_rec` argument.

Signed-off-by: Christian Tryti <ctryti@gmail.com>

* Add new test that checks setup of health_check.

Signed-off-by: Christian Tryti <ctryti@gmail.com>

7 files changed, 135 insertions, 26 deletions

diff --git a/plugin/forward/README.md b/plugin/forward/README.md
index 77d8e181e..92b6bd168 100644
--- a/plugin/forward/README.md
+++ b/plugin/forward/README.md
@@ -51,7 +51,7 @@ forward FROM TO... {
     tls CERT KEY CA
     tls_servername NAME
     policy random|round_robin|sequential
-    health_check DURATION
+    health_check DURATION [no_rec]
     max_concurrent MAX
 }
 ~~~
@@ -85,7 +85,10 @@ forward FROM TO... {
   * `random` is a policy that implements random upstream selection.
   * `round_robin` is a policy that selects hosts based on round robin ordering.
   * `sequential` is a policy that selects hosts based on sequential ordering.
-* `health_check`, use a different **DURATION** for health checking, the default duration is 0.5s.
+* `health_check` configure the behaviour of health checking of the upstream servers
+  * `<duration>` - use a different duration for health checking, the default duration is 0.5s.
+  * `no_rec` - optional argument that sets the RecursionDesired-flag of the dns-query used in health checking to `false`.
+    The flag is default `true`.
 * `max_concurrent` **MAX** will limit the number of concurrent queries to **MAX**.  Any new query that would
   raise the number of concurrent queries above the **MAX** will result in a SERVFAIL response. This
   response does not count as a health failure. When choosing a value for **MAX**, pick a number
diff --git a/plugin/forward/forward.go b/plugin/forward/forward.go
index f6dd939e9..20a3fcf38 100644
--- a/plugin/forward/forward.go
+++ b/plugin/forward/forward.go
@@ -52,7 +52,7 @@ type Forward struct {
 
 // New returns a new Forward.
 func New() *Forward {
-	f := &Forward{maxfails: 2, tlsConfig: new(tls.Config), expire: defaultExpire, p: new(policy.Random), from: ".", hcInterval: hcInterval}
+	f := &Forward{maxfails: 2, tlsConfig: new(tls.Config), expire: defaultExpire, p: new(policy.Random), from: ".", hcInterval: hcInterval, opts: options{forceTCP: false, preferUDP: false, hcRecursionDesired: true}}
 	return f
 }
 
@@ -224,8 +224,9 @@ var (
 
 // options holds various options that can be set.
 type options struct {
-	forceTCP  bool
-	preferUDP bool
+	forceTCP           bool
+	preferUDP          bool
+	hcRecursionDesired bool
 }
 
 const defaultTimeout = 5 * time.Second
diff --git a/plugin/forward/health.go b/plugin/forward/health.go
index 4512d96c8..2b59039a6 100644
--- a/plugin/forward/health.go
+++ b/plugin/forward/health.go
@@ -14,13 +14,18 @@ import (
 type HealthChecker interface {
 	Check(*Proxy) error
 	SetTLSConfig(*tls.Config)
+	SetRecursionDesired(bool)
+	GetRecursionDesired() bool
 }
 
 // dnsHc is a health checker for a DNS endpoint (DNS, and DoT).
-type dnsHc struct{ c *dns.Client }
+type dnsHc struct {
+	c                *dns.Client
+	recursionDesired bool
+}
 
 // NewHealthChecker returns a new HealthChecker based on transport.
-func NewHealthChecker(trans string) HealthChecker {
+func NewHealthChecker(trans string, recursionDesired bool) HealthChecker {
 	switch trans {
 	case transport.DNS, transport.TLS:
 		c := new(dns.Client)
@@ -28,7 +33,7 @@ func NewHealthChecker(trans string) HealthChecker {
 		c.ReadTimeout = 1 * time.Second
 		c.WriteTimeout = 1 * time.Second
 
-		return &dnsHc{c: c}
+		return &dnsHc{c: c, recursionDesired: recursionDesired}
 	}
 
 	log.Warningf("No healthchecker for transport %q", trans)
@@ -40,7 +45,14 @@ func (h *dnsHc) SetTLSConfig(cfg *tls.Config) {
 	h.c.TLSConfig = cfg
 }
 
-// For HC we send to . IN NS +norec message to the upstream. Dial timeouts and empty
+func (h *dnsHc) SetRecursionDesired(recursionDesired bool) {
+	h.recursionDesired = recursionDesired
+}
+func (h *dnsHc) GetRecursionDesired() bool {
+	return h.recursionDesired
+}
+
+// For HC we send to . IN NS +[no]rec message to the upstream. Dial timeouts and empty
 // replies are considered fails, basically anything else constitutes a healthy upstream.
 
 // Check is used as the up.Func in the up.Probe.
@@ -59,6 +71,7 @@ func (h *dnsHc) Check(p *Proxy) error {
 func (h *dnsHc) send(addr string) error {
 	ping := new(dns.Msg)
 	ping.SetQuestion(".", dns.TypeNS)
+	ping.MsgHdr.RecursionDesired = h.recursionDesired
 
 	m, _, err := h.c.Exchange(ping, addr)
 	// If we got a header, we're alright, basically only care about I/O errors 'n stuff.
diff --git a/plugin/forward/health_test.go b/plugin/forward/health_test.go
index b1785eeae..7dd0c337a 100644
--- a/plugin/forward/health_test.go
+++ b/plugin/forward/health_test.go
@@ -14,10 +14,15 @@ import (
 )
 
 func TestHealth(t *testing.T) {
-	const expected = 0
+	const expected = 1
 	i := uint32(0)
+	q := uint32(0)
 	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
-		if r.Question[0].Name == "." {
+		if atomic.LoadUint32(&q) == 0 { //drop the first query to trigger health-checking
+			atomic.AddUint32(&q, 1)
+			return
+		}
+		if r.Question[0].Name == "." && r.RecursionDesired == true {
 			atomic.AddUint32(&i, 1)
 		}
 		ret := new(dns.Msg)
@@ -39,7 +44,43 @@ func TestHealth(t *testing.T) {
 	time.Sleep(1 * time.Second)
 	i1 := atomic.LoadUint32(&i)
 	if i1 != expected {
-		t.Errorf("Expected number of health checks to be %d, got %d", expected, i1)
+		t.Errorf("Expected number of health checks with RecursionDesired==true to be %d, got %d", expected, i1)
+	}
+}
+
+func TestHealthNoRecursion(t *testing.T) {
+	const expected = 1
+	i := uint32(0)
+	q := uint32(0)
+	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		if atomic.LoadUint32(&q) == 0 { //drop the first query to trigger health-checking
+			atomic.AddUint32(&q, 1)
+			return
+		}
+		if r.Question[0].Name == "." && r.RecursionDesired == false {
+			atomic.AddUint32(&i, 1)
+		}
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	p := NewProxy(s.Addr, transport.DNS)
+	p.health.SetRecursionDesired(false)
+	f := New()
+	f.SetProxy(p)
+	defer f.OnShutdown()
+
+	req := new(dns.Msg)
+	req.SetQuestion("example.org.", dns.TypeA)
+
+	f.ServeDNS(context.TODO(), &test.ResponseWriter{}, req)
+
+	time.Sleep(1 * time.Second)
+	i1 := atomic.LoadUint32(&i)
+	if i1 != expected {
+		t.Errorf("Expected number of health checks with RecursionDesired==false to be %d, got %d", expected, i1)
 	}
 }
 
diff --git a/plugin/forward/proxy.go b/plugin/forward/proxy.go
index 60dfa47a6..2fe0f3839 100644
--- a/plugin/forward/proxy.go
+++ b/plugin/forward/proxy.go
@@ -29,7 +29,7 @@ func NewProxy(addr, trans string) *Proxy {
 		probe:     up.New(),
 		transport: newTransport(addr),
 	}
-	p.health = NewHealthChecker(trans)
+	p.health = NewHealthChecker(trans, true)
 	runtime.SetFinalizer(p, (*Proxy).finalizer)
 	return p
 }
diff --git a/plugin/forward/setup.go b/plugin/forward/setup.go
index dadf535c7..cf65c1ba4 100644
--- a/plugin/forward/setup.go
+++ b/plugin/forward/setup.go
@@ -121,6 +121,7 @@ func parseStanza(c *caddy.Controller) (*Forward, error) {
 			f.proxies[i].SetTLSConfig(f.tlsConfig)
 		}
 		f.proxies[i].SetExpire(f.expire)
+		f.proxies[i].health.SetRecursionDesired(f.opts.hcRecursionDesired)
 	}
 
 	return f, nil
@@ -161,6 +162,16 @@ func parseBlock(c *caddy.Controller, f *Forward) error {
 			return fmt.Errorf("health_check can't be negative: %d", dur)
 		}
 		f.hcInterval = dur
+
+		for c.NextArg() {
+			switch hcOpts := c.Val(); hcOpts {
+			case "no_rec":
+				f.opts.hcRecursionDesired = false
+			default:
+				return fmt.Errorf("health_check: unknown option %s", hcOpts)
+			}
+		}
+
 	case "force_tcp":
 		if c.NextArg() {
 			return c.ArgErr()
diff --git a/plugin/forward/setup_test.go b/plugin/forward/setup_test.go
index c2c2f6759..0ac5c3e18 100644
--- a/plugin/forward/setup_test.go
+++ b/plugin/forward/setup_test.go
@@ -21,21 +21,22 @@ func TestSetup(t *testing.T) {
 		expectedErr     string
 	}{
 		// positive
-		{"forward . 127.0.0.1", false, ".", nil, 2, options{}, ""},
-		{"forward . 127.0.0.1 {\nexcept miek.nl\n}\n", false, ".", nil, 2, options{}, ""},
-		{"forward . 127.0.0.1 {\nmax_fails 3\n}\n", false, ".", nil, 3, options{}, ""},
-		{"forward . 127.0.0.1 {\nforce_tcp\n}\n", false, ".", nil, 2, options{forceTCP: true}, ""},
-		{"forward . 127.0.0.1 {\nprefer_udp\n}\n", false, ".", nil, 2, options{preferUDP: true}, ""},
-		{"forward . 127.0.0.1 {\nforce_tcp\nprefer_udp\n}\n", false, ".", nil, 2, options{preferUDP: true, forceTCP: true}, ""},
-		{"forward . 127.0.0.1:53", false, ".", nil, 2, options{}, ""},
-		{"forward . 127.0.0.1:8080", false, ".", nil, 2, options{}, ""},
-		{"forward . [::1]:53", false, ".", nil, 2, options{}, ""},
-		{"forward . [2003::1]:53", false, ".", nil, 2, options{}, ""},
+		{"forward . 127.0.0.1", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 {\nexcept miek.nl\n}\n", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 {\nmax_fails 3\n}\n", false, ".", nil, 3, options{hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 {\nforce_tcp\n}\n", false, ".", nil, 2, options{forceTCP: true, hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 {\nprefer_udp\n}\n", false, ".", nil, 2, options{preferUDP: true, hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 {\nforce_tcp\nprefer_udp\n}\n", false, ".", nil, 2, options{preferUDP: true, forceTCP: true, hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1:53", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1:8080", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . [::1]:53", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . [2003::1]:53", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
+		{"forward . 127.0.0.1 \n", false, ".", nil, 2, options{hcRecursionDesired: true}, ""},
 		// negative
-		{"forward . a27.0.0.1", true, "", nil, 0, options{}, "not an IP"},
-		{"forward . 127.0.0.1 {\nblaatl\n}\n", true, "", nil, 0, options{}, "unknown property"},
+		{"forward . a27.0.0.1", true, "", nil, 0, options{hcRecursionDesired: true}, "not an IP"},
+		{"forward . 127.0.0.1 {\nblaatl\n}\n", true, "", nil, 0, options{hcRecursionDesired: true}, "unknown property"},
 		{`forward . ::1
-		forward com ::2`, true, "", nil, 0, options{}, "plugin"},
+		forward com ::2`, true, "", nil, 0, options{hcRecursionDesired: true}, "plugin"},
 	}
 
 	for i, test := range tests {
@@ -210,3 +211,42 @@ func TestSetupMaxConcurrent(t *testing.T) {
 		}
 	}
 }
+
+func TestSetupHealthCheck(t *testing.T) {
+	tests := []struct {
+		input       string
+		shouldErr   bool
+		expectedVal bool
+		expectedErr string
+	}{
+		// positive
+		{"forward . 127.0.0.1\n", false, true, ""},
+		{"forward . 127.0.0.1 {\nhealth_check 0.5s\n}\n", false, true, ""},
+		{"forward . 127.0.0.1 {\nhealth_check 0.5s no_rec\n}\n", false, false, ""},
+		// negative
+		{"forward . 127.0.0.1 {\nhealth_check no_rec\n}\n", true, true, "time: invalid duration"},
+		{"forward . 127.0.0.1 {\nhealth_check 0.5s rec\n}\n", true, true, "health_check: unknown option rec"},
+	}
+
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		f, err := parseForward(c)
+
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %d: expected error but found %s for input %s", i, err, test.input)
+		}
+
+		if err != nil {
+			if !test.shouldErr {
+				t.Errorf("Test %d: expected no error but found one for input %s, got: %v", i, test.input, err)
+			}
+
+			if !strings.Contains(err.Error(), test.expectedErr) {
+				t.Errorf("Test %d: expected error to contain: %v, found error: %v, input: %s", i, test.expectedErr, err, test.input)
+			}
+		}
+		if !test.shouldErr && (f.opts.hcRecursionDesired != test.expectedVal || f.proxies[0].health.GetRecursionDesired() != test.expectedVal) {
+			t.Errorf("Test %d: expected: %t, got: %d", i, test.expectedVal, f.maxConcurrent)
+		}
+	}
+}