plugin/file/cache: Add metadata for wildcard record responses (#5308)

For responses synthesized by known wildcard records, publish metadata containing the wildcard record name

Signed-off-by: Chris O'Haver <cohaver@infoblox.com>

5 files changed, 93 insertions, 3 deletions

diff --git a/plugin/cache/cache.go b/plugin/cache/cache.go
index fc42866ac..d1989f35b 100644
--- a/plugin/cache/cache.go
+++ b/plugin/cache/cache.go
@@ -114,6 +114,9 @@ type ResponseWriter struct {
 	ad         bool // When true the original request had the AD bit set.
 	prefetch   bool // When true write nothing back to the client.
 	remoteAddr net.Addr
+
+	wildcardFunc func() string // function to retrieve wildcard name that synthesized the result.
+
 }
 
 // newPrefetchResponseWriter returns a Cache ResponseWriter to be used in
@@ -202,6 +205,9 @@ func (w *ResponseWriter) set(m *dns.Msg, key uint64, mt response.Type, duration
 	switch mt {
 	case response.NoError, response.Delegation:
 		i := newItem(m, w.now(), duration)
+		if w.wildcardFunc != nil {
+			i.wildcard = w.wildcardFunc()
+		}
 		if w.pcache.Add(key, i) {
 			evictions.WithLabelValues(w.server, Success, w.zonesMetricLabel).Inc()
 		}
@@ -212,6 +218,9 @@ func (w *ResponseWriter) set(m *dns.Msg, key uint64, mt response.Type, duration
 
 	case response.NameError, response.NoData, response.ServerError:
 		i := newItem(m, w.now(), duration)
+		if w.wildcardFunc != nil {
+			i.wildcard = w.wildcardFunc()
+		}
 		if w.ncache.Add(key, i) {
 			evictions.WithLabelValues(w.server, Denial, w.zonesMetricLabel).Inc()
 		}
diff --git a/plugin/cache/cache_test.go b/plugin/cache/cache_test.go
index 69bea61a4..8b279e265 100644
--- a/plugin/cache/cache_test.go
+++ b/plugin/cache/cache_test.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/plugin/pkg/dnstest"
 	"github.com/coredns/coredns/plugin/pkg/response"
 	"github.com/coredns/coredns/plugin/test"
@@ -578,3 +579,59 @@ func TestComputeTTL(t *testing.T) {
 		}
 	}
 }
+
+func TestCacheWildcardMetadata(t *testing.T) {
+	c := New()
+	qname := "foo.bar.example.org."
+	wildcard := "*.bar.example.org."
+	c.Next = wildcardMetadataBackend(qname, wildcard)
+
+	req := new(dns.Msg)
+	req.SetQuestion(qname, dns.TypeA)
+
+	// 1. Test writing wildcard metadata retrieved from backend to the cache
+
+	ctx := metadata.ContextWithMetadata(context.TODO())
+	w := dnstest.NewRecorder(&test.ResponseWriter{})
+	c.ServeDNS(ctx, w, req)
+	if c.pcache.Len() != 1 {
+		t.Errorf("Msg should have been cached")
+	}
+	_, k := key(qname, w.Msg, response.NoError)
+	i, _ := c.pcache.Get(k)
+	if i.(*item).wildcard != wildcard {
+		t.Errorf("expected wildcard reponse to enter cache with cache item's wildcard = %q, got %q", wildcard, i.(*item).wildcard)
+	}
+
+	// 2. Test retrieving the cached item from cache and writing its wildcard value to metadata
+
+	// reset context and response writer
+	ctx = metadata.ContextWithMetadata(context.TODO())
+	w = dnstest.NewRecorder(&test.ResponseWriter{})
+
+	c.ServeDNS(ctx, &test.ResponseWriter{}, req)
+	f := metadata.ValueFunc(ctx, "zone/wildcard")
+	if f == nil {
+		t.Fatal("expected metadata func for wildcard response retrieved from cache, got nil")
+	}
+	if f() != wildcard {
+		t.Errorf("after retrieving wildcard item from cache, expected \"zone/wildcard\" metadata value to be %q, got %q", wildcard, i.(*item).wildcard)
+	}
+}
+
+// wildcardMetadataBackend mocks a backend that reponds with a response for qname synthesized by wildcard
+// and sets the zone/wildcard metadata value
+func wildcardMetadataBackend(qname, wildcard string) plugin.Handler {
+	return plugin.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+		m := new(dns.Msg)
+		m.SetReply(r)
+		m.Response, m.RecursionAvailable = true, true
+		m.Answer = []dns.RR{test.A(qname + " 300 IN A 127.0.0.1")}
+		metadata.SetValueFunc(ctx, "zone/wildcard", func() string {
+			return wildcard
+		})
+		w.WriteMsg(m)
+
+		return dns.RcodeSuccess, nil
+	})
+}
diff --git a/plugin/cache/handler.go b/plugin/cache/handler.go
index e2b4155ee..f443f3d47 100644
--- a/plugin/cache/handler.go
+++ b/plugin/cache/handler.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/plugin/metrics"
 	"github.com/coredns/coredns/request"
 
@@ -37,7 +38,7 @@ func (c *Cache) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 	ttl := 0
 	i := c.getIgnoreTTL(now, state, server)
 	if i == nil {
-		crr := &ResponseWriter{ResponseWriter: w, Cache: c, state: state, server: server, do: do, ad: ad}
+		crr := &ResponseWriter{ResponseWriter: w, Cache: c, state: state, server: server, do: do,  ad: ad, wildcardFunc: wildcardFunc(ctx)}
 		return c.doRefresh(ctx, state, crr)
 	}
 	ttl = i.ttl(now)
@@ -63,12 +64,29 @@ func (c *Cache) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 		cw := newPrefetchResponseWriter(server, state, c)
 		go c.doPrefetch(ctx, state, cw, i, now)
 	}
+
+	if i.wildcard != "" {
+		// Set wildcard source record name to metadata
+		metadata.SetValueFunc(ctx, "zone/wildcard", func() string {
+			return i.wildcard
+		})
+	}
+
 	resp := i.toMsg(r, now, do, ad)
 	w.WriteMsg(resp)
-
 	return dns.RcodeSuccess, nil
 }
 
+func wildcardFunc(ctx context.Context) func() string {
+	return func() string {
+		// Get wildcard source record name from metadata
+		if f := metadata.ValueFunc(ctx, "zone/wildcard"); f != nil {
+			return f()
+		}
+		return ""
+	}
+}
+
 func (c *Cache) doPrefetch(ctx context.Context, state request.Request, cw *ResponseWriter, i *item, now time.Time) {
 	cachePrefetches.WithLabelValues(cw.server, c.zonesMetricLabel).Inc()
 	c.doRefresh(ctx, state, cw)
diff --git a/plugin/cache/item.go b/plugin/cache/item.go
index 27bd4ccbb..6b51a5bad 100644
--- a/plugin/cache/item.go
+++ b/plugin/cache/item.go
@@ -19,6 +19,7 @@ type item struct {
 	Answer             []dns.RR
 	Ns                 []dns.RR
 	Extra              []dns.RR
+	wildcard           string
 
 	origTTL uint32
 	stored  time.Time
diff --git a/plugin/file/lookup.go b/plugin/file/lookup.go
index 08eed7d61..a999a2a79 100644
--- a/plugin/file/lookup.go
+++ b/plugin/file/lookup.go
@@ -6,6 +6,7 @@ import (
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin/file/rrutil"
 	"github.com/coredns/coredns/plugin/file/tree"
+	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
@@ -214,7 +215,10 @@ func (z *Zone) Lookup(ctx context.Context, state request.Request, qname string)
 
 	// Found wildcard.
 	if wildElem != nil {
-		auth := ap.ns(do)
+		// set metadata value for the wildcard record that synthesized the result
+		metadata.SetValueFunc(ctx, "zone/wildcard", func() string {
+			return wildElem.Name()
+		})
 
 		if rrs := wildElem.TypeForWildcard(dns.TypeCNAME, qname); len(rrs) > 0 && qtype != dns.TypeCNAME {
 			ctx = context.WithValue(ctx, dnsserver.LoopKey{}, loop+1)
@@ -233,6 +237,7 @@ func (z *Zone) Lookup(ctx context.Context, state request.Request, qname string)
 			return nil, ret, nil, NoData
 		}
 
+		auth := ap.ns(do)
 		if do {
 			// An NSEC is needed to say no longer name exists under this wildcard.
 			if deny, found := tr.Prev(qname); found {