plugin/template: Persist truncated state to client if CNAME lookup response is truncated (#4713)

* persist truncated state to client if cname lookup response is truncated

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

plugin/template/cname_test.go

@@ -0,0 +1,96 @@
+package template
+
+import (
+	"context"
+	"regexp"
+	"testing"
+	gotmpl "text/template"
+
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
+
+	"github.com/miekg/dns"
+)
+
+func TestTruncatedCNAME(t *testing.T) {
+	up := &Upstub{
+		Qclass:    dns.ClassINET,
+		Truncated: true,
+		Case: test.Case{
+			Qname: "cname.test.",
+			Qtype: dns.TypeA,
+			Rcode: dns.RcodeSuccess,
+			Answer: []dns.RR{
+				test.CNAME("cname.test. 600 IN CNAME test.up"),
+				test.A("test.up. 600 IN A 1.2.3.4"),
+			},
+		},
+	}
+
+	handler := Handler{
+		Zones: []string{"."},
+		Templates: []template{{
+			regex:    []*regexp.Regexp{regexp.MustCompile("^cname\\.test\\.$")},
+			answer:   []*gotmpl.Template{gotmpl.Must(gotmpl.New("answer").Parse(up.Answer[0].String()))},
+			qclass:   dns.ClassINET,
+			qtype:    dns.TypeA,
+			zones:    []string{"test."},
+			upstream: up,
+		}},
+	}
+
+	r := &dns.Msg{Question: []dns.Question{{Name: up.Qname, Qclass: up.Qclass, Qtype: up.Qtype}}}
+	w := dnstest.NewRecorder(&test.ResponseWriter{})
+
+	_, err := handler.ServeDNS(context.TODO(), w, r)
+
+	if err != nil {
+		t.Fatalf("Unexpecetd error %q", err)
+	}
+	if w.Msg == nil {
+		t.Fatalf("Unexpecetd empty response.")
+	}
+	if !w.Msg.Truncated {
+		t.Error("Expected reply to be marked truncated.")
+	}
+	err = test.SortAndCheck(w.Msg, up.Case)
+	if err != nil {
+		t.Error(err)
+	}
+}
+
+// Upstub implements an Upstreamer that returns a set response for test purposes
+type Upstub struct {
+	test.Case
+	Truncated bool
+	Qclass    uint16
+}
+
+// Lookup returns a set response
+func (t *Upstub) Lookup(ctx context.Context, state request.Request, name string, typ uint16) (*dns.Msg, error) {
+	var answer []dns.RR
+	// if query type is not CNAME, remove any CNAME with same name as qname from the answer
+	if t.Qtype != dns.TypeCNAME {
+		for _, a := range t.Answer {
+			if c, ok := a.(*dns.CNAME); ok && c.Header().Name == t.Qname {
+				continue
+			}
+			answer = append(answer, a)
+		}
+	} else {
+		answer = t.Answer
+	}
+
+	return &dns.Msg{
+		MsgHdr: dns.MsgHdr{
+			Response:  true,
+			Truncated: t.Truncated,
+			Rcode:     t.Rcode,
+		},
+		Question: []dns.Question{{Name: t.Qname, Qtype: t.Qtype, Qclass: t.Qclass}},
+		Answer:   answer,
+		Extra:    t.Extra,
+		Ns:       t.Ns,
+	}, nil
+}

plugin/template/template.go

@@ -11,7 +11,6 @@ import (
 	"github.com/coredns/coredns/plugin/metadata"
 	"github.com/coredns/coredns/plugin/metrics"
 	"github.com/coredns/coredns/plugin/pkg/fall"
-	"github.com/coredns/coredns/plugin/pkg/upstream"
 	"github.com/coredns/coredns/request"
 
 	"github.com/miekg/dns"
@@ -35,7 +34,12 @@ type template struct {
 	qclass     uint16
 	qtype      uint16
 	fall       fall.F
-	upstream   *upstream.Upstream
+	upstream   Upstreamer
+}
+
+// Upstreamer looks up targets of CNAME templates
+type Upstreamer interface {
+	Lookup(ctx context.Context, state request.Request, name string, typ uint16) (*dns.Msg, error)
 }
 
 type templateData struct {
@@ -100,10 +104,12 @@ func (h Handler) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg)
 			}
 			msg.Answer = append(msg.Answer, rr)
 			if template.upstream != nil && (state.QType() == dns.TypeA || state.QType() == dns.TypeAAAA) && rr.Header().Rrtype == dns.TypeCNAME {
-				up, _ := template.upstream.Lookup(ctx, state, rr.(*dns.CNAME).Target, state.QType())
+				if up, err := template.upstream.Lookup(ctx, state, rr.(*dns.CNAME).Target, state.QType()); err == nil && up != nil {
+					msg.Truncated = up.Truncated
 					msg.Answer = append(msg.Answer, up.Answer...)
 				}
 			}
+		}
 		for _, additional := range template.additional {
 			rr, err := executeRRTemplate(metrics.WithServer(ctx), "additional", additional, data)
 			if err != nil {