plugin/rewrite - extend edns0 local variable support with metadata (#1928)

* - add support of metadata values for edns0 local variables

* - comments from review.

* - simplify label check. Add UT

* - enhance check for Labels, add UT
- remove IsMetadataSet

* - edns0 variable - if variable is not found just ignore the rewrite.

9 files changed, 132 insertions, 22 deletions

diff --git a/plugin/metadata/metadata_test.go b/plugin/metadata/metadata_test.go
index 7ded05c03..08d9b6734 100644
--- a/plugin/metadata/metadata_test.go
+++ b/plugin/metadata/metadata_test.go
@@ -30,8 +30,8 @@ func (m *testHandler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns
 
 func TestMetadataServeDNS(t *testing.T) {
 	expectedMetadata := []testProvider{
-		testProvider{"test/key1": func() string { return "testvalue1" }},
-		testProvider{"test/key2": func() string { return "two" }, "test/key3": func() string { return "testvalue3" }},
+		{"test/key1": func() string { return "testvalue1" }},
+		{"test/key2": func() string { return "two" }, "test/key3": func() string { return "testvalue3" }},
 	}
 	// Create fake Providers based on expectedMetadata
 	providers := []Provider{}
@@ -52,6 +52,9 @@ func TestMetadataServeDNS(t *testing.T) {
 
 	for _, expected := range expectedMetadata {
 		for label, expVal := range expected {
+			if !IsLabel(label) {
+				t.Errorf("Expected label %s is not considered a valid label", label)
+			}
 			val := ValueFunc(nctx, label)
 			if val() != expVal() {
 				t.Errorf("Expected value %s for %s, but got %s", expVal(), label, val())
@@ -59,3 +62,26 @@ func TestMetadataServeDNS(t *testing.T) {
 		}
 	}
 }
+
+func TestLabelFormat(t *testing.T) {
+	labels := []struct {
+		label   string
+		isValid bool
+	}{
+		{"plugin/LABEL", true},
+		{"p/LABEL", true},
+		{"plugin/L", true},
+		{"LABEL", false},
+		{"plugin.LABEL", false},
+		{"/NO-PLUGIN-NOT-ACCEPTED", false},
+		{"ONLY-PLUGIN-NOT-ACCEPTED/", false},
+		{"PLUGIN/LABEL/SUB-LABEL", false},
+		{"/", false},
+	}
+
+	for _, test := range labels {
+		if IsLabel(test.label) != test.isValid {
+			t.Errorf("Label %v is expected to have this validaty : %v - and has the opposite", test.label, test.isValid)
+		}
+	}
+}
diff --git a/plugin/metadata/provider.go b/plugin/metadata/provider.go
index eb7bb9755..276a44127 100644
--- a/plugin/metadata/provider.go
+++ b/plugin/metadata/provider.go
@@ -32,6 +32,7 @@ package metadata
 
 import (
 	"context"
+	"strings"
 
 	"github.com/coredns/coredns/request"
 )
@@ -48,6 +49,21 @@ type Provider interface {
 // Func is the type of function in the metadata, when called they return the value of the label.
 type Func func() string
 
+// IsLabel check that the provided name looks like a valid label name
+func IsLabel(label string) bool {
+	p := strings.Index(label, "/")
+	if p <= 0 || p >= len(label)-1 {
+		// cannot accept namespace empty nor label empty
+		return false
+	}
+	if strings.LastIndex(label, "/") != p {
+		// several slash in the Label
+		return false
+	}
+	return true
+
+}
+
 // Labels returns all metadata keys stored in the context. These label names should be named
 // as: plugin/NAME, where NAME is something descriptive.
 func Labels(ctx context.Context) []string {
diff --git a/plugin/rewrite/README.md b/plugin/rewrite/README.md
index 680e69722..feec7dc81 100644
--- a/plugin/rewrite/README.md
+++ b/plugin/rewrite/README.md
@@ -209,12 +209,24 @@ rewrites the first local option with code 0xffee, setting the data to "abcd". Eq
 * A variable data is specified with a pair of curly brackets `{}`. Following are the supported variables:
   {qname}, {qtype}, {client_ip}, {client_port}, {protocol}, {server_ip}, {server_port}.
 
-Example:
+* If the metadata plugin is enabled, then labels are supported as variables if they are presented within curly brackets. 
+the variable data will be filled with the value associated with that label. If that label is not provided, 
+the variable will be silently substitute by an empty string.  
+
+Examples:
 
 ~~~
 rewrite edns0 local set 0xffee {client_ip}
 ~~~
 
+The following example uses metadata and an imaginary "some-plugin" that would provide "some-label" as metadata information.
+
+~~~
+metadata
+some-plugin
+rewrite edns0 local set 0xffee {some-plugin/some-label}
+~~~
+
 ### EDNS0_NSID
 
 This has no fields; it will add an NSID option with an empty string for the NSID. If the option already exists
diff --git a/plugin/rewrite/class.go b/plugin/rewrite/class.go
index d72557884..dceb4429b 100644
--- a/plugin/rewrite/class.go
+++ b/plugin/rewrite/class.go
@@ -1,6 +1,7 @@
 package rewrite
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -29,7 +30,7 @@ func newClassRule(nextAction string, args ...string) (Rule, error) {
 }
 
 // Rewrite rewrites the the current request.
-func (rule *classRule) Rewrite(state request.Request) Result {
+func (rule *classRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if rule.fromClass > 0 && rule.toClass > 0 {
 		if state.Req.Question[0].Qclass == rule.fromClass {
 			state.Req.Question[0].Qclass = rule.toClass
diff --git a/plugin/rewrite/edns0.go b/plugin/rewrite/edns0.go
index 68084a773..074dd3ebf 100644
--- a/plugin/rewrite/edns0.go
+++ b/plugin/rewrite/edns0.go
@@ -2,12 +2,14 @@
 package rewrite
 
 import (
+	"context"
 	"encoding/hex"
 	"fmt"
 	"net"
 	"strconv"
 	"strings"
 
+	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
@@ -46,7 +48,7 @@ func setupEdns0Opt(r *dns.Msg) *dns.OPT {
 }
 
 // Rewrite will alter the request EDNS0 NSID option
-func (rule *edns0NsidRule) Rewrite(state request.Request) Result {
+func (rule *edns0NsidRule) Rewrite(ctx context.Context, state request.Request) Result {
 	o := setupEdns0Opt(state.Req)
 
 	for _, s := range o.Option {
@@ -74,7 +76,7 @@ func (rule *edns0NsidRule) Mode() string { return rule.mode }
 func (rule *edns0NsidRule) GetResponseRule() ResponseRule { return ResponseRule{} }
 
 // Rewrite will alter the request EDNS0 local options.
-func (rule *edns0LocalRule) Rewrite(state request.Request) Result {
+func (rule *edns0LocalRule) Rewrite(ctx context.Context, state request.Request) Result {
 	o := setupEdns0Opt(state.Req)
 
 	for _, s := range o.Option {
@@ -174,7 +176,7 @@ func newEdns0VariableRule(mode, action, code, variable string) (*edns0VariableRu
 }
 
 // ruleData returns the data specified by the variable.
-func (rule *edns0VariableRule) ruleData(state request.Request) ([]byte, error) {
+func (rule *edns0VariableRule) ruleData(ctx context.Context, state request.Request) ([]byte, error) {
 
 	switch rule.variable {
 	case queryName:
@@ -199,12 +201,17 @@ func (rule *edns0VariableRule) ruleData(state request.Request) ([]byte, error) {
 		return []byte(state.Proto()), nil
 	}
 
+	fetcher := metadata.ValueFunc(ctx, rule.variable[1:len(rule.variable)-1])
+	if fetcher != nil {
+		return []byte(fetcher()), nil
+	}
+
 	return nil, fmt.Errorf("unable to extract data for variable %s", rule.variable)
 }
 
 // Rewrite will alter the request EDNS0 local options with specified variables.
-func (rule *edns0VariableRule) Rewrite(state request.Request) Result {
-	data, err := rule.ruleData(state)
+func (rule *edns0VariableRule) Rewrite(ctx context.Context, state request.Request) Result {
+	data, err := rule.ruleData(ctx, state)
 	if err != nil || data == nil {
 		return RewriteIgnored
 	}
@@ -249,6 +256,10 @@ func isValidVariable(variable string) bool {
 		serverPort:
 		return true
 	}
+	// we cannot validate the labels of metadata - but we can verify it has the syntax of a label
+	if strings.HasPrefix(variable, "{") && strings.HasSuffix(variable, "}") && metadata.IsLabel(variable[1:len(variable)-1]) {
+		return true
+	}
 	return false
 }
 
@@ -310,7 +321,7 @@ func (rule *edns0SubnetRule) fillEcsData(state request.Request, ecs *dns.EDNS0_S
 }
 
 // Rewrite will alter the request EDNS0 subnet option.
-func (rule *edns0SubnetRule) Rewrite(state request.Request) Result {
+func (rule *edns0SubnetRule) Rewrite(ctx context.Context, state request.Request) Result {
 	o := setupEdns0Opt(state.Req)
 
 	for _, s := range o.Option {
diff --git a/plugin/rewrite/name.go b/plugin/rewrite/name.go
index 45e67267d..816b6cdf7 100644
--- a/plugin/rewrite/name.go
+++ b/plugin/rewrite/name.go
@@ -1,6 +1,7 @@
 package rewrite
 
 import (
+	"context"
 	"fmt"
 	"regexp"
 	"strconv"
@@ -56,7 +57,7 @@ const (
 
 // Rewrite rewrites the current request based upon exact match of the name
 // in the question section of the request.
-func (rule *nameRule) Rewrite(state request.Request) Result {
+func (rule *nameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if rule.From == state.Name() {
 		state.Req.Question[0].Name = rule.To
 		return RewriteDone
@@ -65,7 +66,7 @@ func (rule *nameRule) Rewrite(state request.Request) Result {
 }
 
 // Rewrite rewrites the current request when the name begins with the matching string.
-func (rule *prefixNameRule) Rewrite(state request.Request) Result {
+func (rule *prefixNameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if strings.HasPrefix(state.Name(), rule.Prefix) {
 		state.Req.Question[0].Name = rule.Replacement + strings.TrimLeft(state.Name(), rule.Prefix)
 		return RewriteDone
@@ -74,7 +75,7 @@ func (rule *prefixNameRule) Rewrite(state request.Request) Result {
 }
 
 // Rewrite rewrites the current request when the name ends with the matching string.
-func (rule *suffixNameRule) Rewrite(state request.Request) Result {
+func (rule *suffixNameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if strings.HasSuffix(state.Name(), rule.Suffix) {
 		state.Req.Question[0].Name = strings.TrimRight(state.Name(), rule.Suffix) + rule.Replacement
 		return RewriteDone
@@ -84,7 +85,7 @@ func (rule *suffixNameRule) Rewrite(state request.Request) Result {
 
 // Rewrite rewrites the current request based upon partial match of the
 // name in the question section of the request.
-func (rule *substringNameRule) Rewrite(state request.Request) Result {
+func (rule *substringNameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if strings.Contains(state.Name(), rule.Substring) {
 		state.Req.Question[0].Name = strings.Replace(state.Name(), rule.Substring, rule.Replacement, -1)
 		return RewriteDone
@@ -94,7 +95,7 @@ func (rule *substringNameRule) Rewrite(state request.Request) Result {
 
 // Rewrite rewrites the current request when the name in the question
 // section of the request matches a regular expression.
-func (rule *regexNameRule) Rewrite(state request.Request) Result {
+func (rule *regexNameRule) Rewrite(ctx context.Context, state request.Request) Result {
 	regexGroups := rule.Pattern.FindStringSubmatch(state.Name())
 	if len(regexGroups) == 0 {
 		return RewriteIgnored
diff --git a/plugin/rewrite/rewrite.go b/plugin/rewrite/rewrite.go
index 7389c68eb..b4bdd5e20 100644
--- a/plugin/rewrite/rewrite.go
+++ b/plugin/rewrite/rewrite.go
@@ -42,7 +42,7 @@ func (rw Rewrite) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 	state := request.Request{W: w, Req: r}
 
 	for _, rule := range rw.Rules {
-		switch result := rule.Rewrite(state); result {
+		switch result := rule.Rewrite(ctx, state); result {
 		case RewriteDone:
 			respRule := rule.GetResponseRule()
 			if respRule.Active == true {
@@ -71,7 +71,7 @@ func (rw Rewrite) Name() string { return "rewrite" }
 // Rule describes a rewrite rule.
 type Rule interface {
 	// Rewrite rewrites the current request.
-	Rewrite(state request.Request) Result
+	Rewrite(ctx context.Context, state request.Request) Result
 	// Mode returns the processing mode stop or continue.
 	Mode() string
 	// GetResponseRule returns the rule to rewrite response with, if any.
diff --git a/plugin/rewrite/rewrite_test.go b/plugin/rewrite/rewrite_test.go
index d67f8d0bf..f70933c3c 100644
--- a/plugin/rewrite/rewrite_test.go
+++ b/plugin/rewrite/rewrite_test.go
@@ -7,8 +7,10 @@ import (
 	"testing"
 
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/plugin/pkg/dnstest"
 	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
 )
@@ -434,12 +436,32 @@ func optsEqual(a, b []dns.EDNS0) bool {
 	return true
 }
 
+type testProvider map[string]metadata.Func
+
+func (tp testProvider) Metadata(ctx context.Context, state request.Request) context.Context {
+	for k, v := range tp {
+		metadata.SetValueFunc(ctx, k, v)
+	}
+	return ctx
+}
+
 func TestRewriteEDNS0LocalVariable(t *testing.T) {
 	rw := Rewrite{
 		Next:     plugin.HandlerFunc(msgPrinter),
 		noRevert: true,
 	}
 
+	expectedMetadata := []metadata.Provider{
+		testProvider{"test/label": func() string { return "my-value" }},
+		testProvider{"test/empty": func() string { return "" }},
+	}
+
+	meta := metadata.Metadata{
+		Zones:     []string{"."},
+		Providers: expectedMetadata,
+		Next:      &rw,
+	}
+
 	// test.ResponseWriter has the following values:
 	// 		The remote will always be 10.240.0.1 and port 40212.
 	// 		The local address is always 127.0.0.1 and port 53.
@@ -492,9 +514,26 @@ func TestRewriteEDNS0LocalVariable(t *testing.T) {
 			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte{0x7F, 0x00, 0x00, 0x01}}},
 			true,
 		},
+		{
+			[]dns.EDNS0{},
+			[]string{"local", "set", "0xffee", "{test/label}"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("my-value")}},
+			true,
+		},
+		{
+			[]dns.EDNS0{},
+			[]string{"local", "set", "0xffee", "{test/empty}"},
+			[]dns.EDNS0{&dns.EDNS0_LOCAL{Code: 0xffee, Data: []byte("")}},
+			true,
+		},
+		{
+			[]dns.EDNS0{},
+			[]string{"local", "set", "0xffee", "{test/does-not-exist}"},
+			nil,
+			false,
+		},
 	}
 
-	ctx := context.TODO()
 	for i, tc := range tests {
 		m := new(dns.Msg)
 		m.SetQuestion("example.com.", dns.TypeA)
@@ -506,16 +545,19 @@ func TestRewriteEDNS0LocalVariable(t *testing.T) {
 		}
 		rw.Rules = []Rule{r}
 
+		ctx := context.TODO()
 		rec := dnstest.NewRecorder(&test.ResponseWriter{})
-		rw.ServeDNS(ctx, rec, m)
+		meta.ServeDNS(ctx, rec, m)
 
 		resp := rec.Msg
 		o := resp.IsEdns0()
-		o.SetDo(tc.doBool)
 		if o == nil {
-			t.Errorf("Test %d: EDNS0 options not set", i)
+			if tc.toOpts != nil {
+				t.Errorf("Test %d: EDNS0 options not set", i)
+			}
 			continue
 		}
+		o.SetDo(tc.doBool)
 		if o.Do() != tc.doBool {
 			t.Errorf("Test %d: Expected %v but got %v", i, tc.doBool, o.Do())
 		}
diff --git a/plugin/rewrite/type.go b/plugin/rewrite/type.go
index 5aa0aeae7..eebb1c02e 100644
--- a/plugin/rewrite/type.go
+++ b/plugin/rewrite/type.go
@@ -2,6 +2,7 @@
 package rewrite
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -30,7 +31,7 @@ func newTypeRule(nextAction string, args ...string) (Rule, error) {
 }
 
 // Rewrite rewrites the the current request.
-func (rule *typeRule) Rewrite(state request.Request) Result {
+func (rule *typeRule) Rewrite(ctx context.Context, state request.Request) Result {
 	if rule.fromType > 0 && rule.toType > 0 {
 		if state.QType() == rule.fromType {
 			state.Req.Question[0].Qtype = rule.toType