plugin/header: Add support for query modification (#5548) (#5556)

6 files changed, 174 insertions, 49 deletions

diff --git a/plugin/header/README.md b/plugin/header/README.md
index 30ae90315..9a855c722 100644
--- a/plugin/header/README.md
+++ b/plugin/header/README.md
@@ -2,22 +2,24 @@
 
 ## Name
 
-*header* - modifies the header for responses.
+*header* - modifies the header for queries and responses.
 
 ## Description
 
-*header* ensures that the flags are in the desired state for responses. The modifications are made transparently for
-the client.
+*header* ensures that the flags are in the desired state for queries and responses.
+The modifications are made transparently for the client and subsequent plugins.
 
 ## Syntax
 
 ~~~
 header {
-    ACTION FLAGS...
-    ACTION FLAGS...
+    [SELECTOR] ACTION FLAGS...
+    [SELECTOR] ACTION FLAGS...
 }
 ~~~
 
+* **SELECTOR** defines if the action should be applied on `query` or `response`. In future CoreDNS version the selector will be mandatory. For backwards compatibility the action will be applied on `response` if the selector is undefined.
+
 * **ACTION** defines the state for DNS message header flags. Actions are evaluated in the order they are defined so last one has the
   most precedence. Allowed values are:
     * `set`
@@ -34,7 +36,7 @@ Make sure recursive available `ra` flag is set in all the responses:
 ~~~ corefile
 . {
     header {
-        set ra
+        response set ra
     }
 }
 ~~~
@@ -44,8 +46,18 @@ Make sure "recursion available" `ra` and "authoritative answer" `aa` flags are s
 ~~~ corefile
 . {
     header {
-        set ra aa
-        clear rd
+        response set ra aa
+        response clear rd
+    }
+}
+~~~
+
+Make sure "recursion desired" `rd` is set for all subsequent plugins::
+
+~~~ corefile
+. {
+    header {
+        query set rd
     }
 }
 ~~~
diff --git a/plugin/header/handler.go b/plugin/header/handler.go
index a002c4fb1..e11eb03c4 100644
--- a/plugin/header/handler.go
+++ b/plugin/header/handler.go
@@ -8,15 +8,18 @@ import (
 	"github.com/miekg/dns"
 )
 
-// Header modifies dns.MsgHdr in the responses
+// Header modifies flags of dns.MsgHdr in queries and / or responses
 type Header struct {
-	Rules []Rule
-	Next  plugin.Handler
+	QueryRules    []Rule
+	ResponseRules []Rule
+	Next          plugin.Handler
 }
 
 // ServeDNS implements the plugin.Handler interface.
 func (h Header) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
-	wr := ResponseHeaderWriter{ResponseWriter: w, Rules: h.Rules}
+	applyRules(r, h.QueryRules)
+
+	wr := ResponseHeaderWriter{ResponseWriter: w, Rules: h.ResponseRules}
 	return plugin.NextOrFailure(h.Name(), h.Next, ctx, &wr, r)
 }
 
diff --git a/plugin/header/header.go b/plugin/header/header.go
index 660b07146..830587d02 100644
--- a/plugin/header/header.go
+++ b/plugin/header/header.go
@@ -26,18 +26,7 @@ type ResponseHeaderWriter struct {
 
 // WriteMsg implements the dns.ResponseWriter interface.
 func (r *ResponseHeaderWriter) WriteMsg(res *dns.Msg) error {
-	// handle all supported flags
-	for _, rule := range r.Rules {
-		switch rule.Flag {
-		case authoritative:
-			res.Authoritative = rule.State
-		case recursionAvailable:
-			res.RecursionAvailable = rule.State
-		case recursionDesired:
-			res.RecursionDesired = rule.State
-		}
-	}
-
+	applyRules(res, r.Rules)
 	return r.ResponseWriter.WriteMsg(res)
 }
 
@@ -90,3 +79,17 @@ func newRules(key string, args []string) ([]Rule, error) {
 
 	return rules, nil
 }
+
+func applyRules(res *dns.Msg, rules []Rule) {
+	// handle all supported flags
+	for _, rule := range rules {
+		switch rule.Flag {
+		case authoritative:
+			res.Authoritative = rule.State
+		case recursionAvailable:
+			res.RecursionAvailable = rule.State
+		case recursionDesired:
+			res.RecursionDesired = rule.State
+		}
+	}
+}
diff --git a/plugin/header/header_test.go b/plugin/header/header_test.go
index bd80e8635..118265419 100644
--- a/plugin/header/header_test.go
+++ b/plugin/header/header_test.go
@@ -11,7 +11,7 @@ import (
 	"github.com/miekg/dns"
 )
 
-func TestHeader(t *testing.T) {
+func TestHeaderResponseRules(t *testing.T) {
 	wr := dnstest.NewRecorder(&test.ResponseWriter{})
 	next := plugin.HandlerFunc(func(ctx context.Context, writer dns.ResponseWriter, msg *dns.Msg) (int, error) {
 		writer.WriteMsg(msg)
@@ -25,8 +25,8 @@ func TestHeader(t *testing.T) {
 	}{
 		{
 			handler: Header{
-				Rules: []Rule{{Flag: recursionAvailable, State: true}},
-				Next:  next,
+				ResponseRules: []Rule{{Flag: recursionAvailable, State: true}},
+				Next:          next,
 			},
 			got: func(msg *dns.Msg) bool {
 				return msg.RecursionAvailable
@@ -35,18 +35,18 @@ func TestHeader(t *testing.T) {
 		},
 		{
 			handler: Header{
-				Rules: []Rule{{Flag: recursionAvailable, State: true}},
-				Next:  next,
+				ResponseRules: []Rule{{Flag: recursionAvailable, State: false}},
+				Next:          next,
 			},
 			got: func(msg *dns.Msg) bool {
 				return msg.RecursionAvailable
 			},
-			expected: true,
+			expected: false,
 		},
 		{
 			handler: Header{
-				Rules: []Rule{{Flag: recursionDesired, State: true}},
-				Next:  next,
+				ResponseRules: []Rule{{Flag: recursionDesired, State: true}},
+				Next:          next,
 			},
 			got: func(msg *dns.Msg) bool {
 				return msg.RecursionDesired
@@ -55,8 +55,8 @@ func TestHeader(t *testing.T) {
 		},
 		{
 			handler: Header{
-				Rules: []Rule{{Flag: authoritative, State: true}},
-				Next:  next,
+				ResponseRules: []Rule{{Flag: authoritative, State: true}},
+				Next:          next,
 			},
 			got: func(msg *dns.Msg) bool {
 				return msg.Authoritative
@@ -80,3 +80,73 @@ func TestHeader(t *testing.T) {
 		}
 	}
 }
+
+func TestHeaderQueryRules(t *testing.T) {
+	wr := dnstest.NewRecorder(&test.ResponseWriter{})
+	next := plugin.HandlerFunc(func(ctx context.Context, writer dns.ResponseWriter, msg *dns.Msg) (int, error) {
+		writer.WriteMsg(msg)
+		return dns.RcodeSuccess, nil
+	})
+
+	tests := []struct {
+		handler  plugin.Handler
+		got      func(msg *dns.Msg) bool
+		expected bool
+	}{
+		{
+			handler: Header{
+				QueryRules: []Rule{{Flag: recursionAvailable, State: true}},
+				Next:       next,
+			},
+			got: func(msg *dns.Msg) bool {
+				return msg.RecursionAvailable
+			},
+			expected: true,
+		},
+		{
+			handler: Header{
+				QueryRules: []Rule{{Flag: recursionDesired, State: true}},
+				Next:       next,
+			},
+			got: func(msg *dns.Msg) bool {
+				return msg.RecursionDesired
+			},
+			expected: true,
+		},
+		{
+			handler: Header{
+				QueryRules: []Rule{{Flag: recursionDesired, State: false}},
+				Next:       next,
+			},
+			got: func(msg *dns.Msg) bool {
+				return msg.RecursionDesired
+			},
+			expected: false,
+		},
+		{
+			handler: Header{
+				QueryRules: []Rule{{Flag: authoritative, State: true}},
+				Next:       next,
+			},
+			got: func(msg *dns.Msg) bool {
+				return msg.Authoritative
+			},
+			expected: true,
+		},
+	}
+
+	for i, tc := range tests {
+		m := new(dns.Msg)
+
+		_, err := tc.handler.ServeDNS(context.TODO(), wr, m)
+		if err != nil {
+			t.Errorf("Test %d: Expected no error, but got %s", i, err)
+			continue
+		}
+
+		if tc.got(m) != tc.expected {
+			t.Errorf("Test %d: Expected flag state=%t, but got %t", i, tc.expected, tc.got(m))
+			continue
+		}
+	}
+}
diff --git a/plugin/header/setup.go b/plugin/header/setup.go
index f84909563..3d6facf08 100644
--- a/plugin/header/setup.go
+++ b/plugin/header/setup.go
@@ -2,6 +2,7 @@ package header
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/coredns/caddy"
 	"github.com/coredns/coredns/core/dnsserver"
@@ -11,39 +12,63 @@ import (
 func init() { plugin.Register("header", setup) }
 
 func setup(c *caddy.Controller) error {
-	rules, err := parse(c)
+	queryRules, responseRules, err := parse(c)
 	if err != nil {
 		return plugin.Error("header", err)
 	}
 
 	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
 		return Header{
-			Rules: rules,
-			Next:  next,
+			QueryRules:    queryRules,
+			ResponseRules: responseRules,
+			Next:          next,
 		}
 	})
 
 	return nil
 }
 
-func parse(c *caddy.Controller) ([]Rule, error) {
+func parse(c *caddy.Controller) ([]Rule, []Rule, error) {
 	for c.Next() {
-		var all []Rule
+		var queryRules []Rule
+		var responseRules []Rule
+
 		for c.NextBlock() {
-			v := c.Val()
+			selector := strings.ToLower(c.Val())
+
+			var action string
+			if selector == "set" || selector == "clear" {
+				log.Warningf("The selector for header rule in line %d isn't explicit defined. "+
+					"Assume rule applies for selector 'response'. This syntax is deprecated. "+
+					"In future versions of CoreDNS the selector must be explicit defined.",
+					c.Line())
+
+				action = selector
+				selector = "response"
+			} else if selector == "query" || selector == "response" {
+				if c.NextArg() {
+					action = c.Val()
+				}
+			} else {
+				return nil, nil, fmt.Errorf("setting up rule: invalid selector=%s should be query or response", selector)
+			}
+
 			args := c.RemainingArgs()
-			// set up rules
-			rules, err := newRules(v, args)
+			rules, err := newRules(action, args)
 			if err != nil {
-				return nil, fmt.Errorf("seting up rule: %w", err)
+				return nil, nil, fmt.Errorf("setting up rule: %w", err)
+			}
+
+			if selector == "response" {
+				responseRules = append(responseRules, rules...)
+			} else {
+				queryRules = append(queryRules, rules...)
 			}
-			all = append(all, rules...)
 		}
 
-		// return combined rules
-		if len(all) > 0 {
-			return all, nil
+		if len(queryRules) > 0 || len(responseRules) > 0 {
+			return queryRules, responseRules, nil
 		}
 	}
-	return nil, c.ArgErr()
+	return nil, nil, c.ArgErr()
 }
diff --git a/plugin/header/setup_test.go b/plugin/header/setup_test.go
index 6a3a1af8e..36b7995e4 100644
--- a/plugin/header/setup_test.go
+++ b/plugin/header/setup_test.go
@@ -19,14 +19,26 @@ func TestSetupHeader(t *testing.T) {
 }`, true, "invalid length for flags, at least one should be provided"},
 		{`header {
 					foo
+}`, true, "invalid selector=foo should be query or response"},
+		{`header {
+					query foo
 }`, true, "invalid length for flags, at least one should be provided"},
 		{`header {
-					foo bar
+					query foo rd
 }`, true, "unknown flag action=foo, should be set or clear"},
 		{`header {
 					set ra
 }`, false, ""},
 		{`header {
+					clear ra
+		}`, false, ""},
+		{`header {
+					query set rd
+		}`, false, ""},
+		{`header {
+					response set aa
+		}`, false, ""},
+		{`header {
 			set ra aa
 			clear rd
 }`, false, ""},