1 files changed, 75 insertions, 0 deletions
diff --git a/plugin/cache/do_test.go b/plugin/cache/do_test.go
new file mode 100644
index 000000000..3cf87cabe
--- /dev/null
+++ b/plugin/cache/do_test.go
@@ -0,0 +1,75 @@
+package cache
+
+import (
+	"context"
+	"testing"
+
+	"github.com/coredns/coredns/plugin"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+
+	"github.com/miekg/dns"
+)
+
+func TestDo(t *testing.T) {
+	// cache sets Do and requests that don't have them.
+	c := New()
+	c.Next = echoHandler()
+	req := new(dns.Msg)
+	req.SetQuestion("example.org.", dns.TypeA)
+	rec := dnstest.NewRecorder(&test.ResponseWriter{})
+
+	// No DO set.
+	c.ServeDNS(context.TODO(), rec, req)
+	reply := rec.Msg
+	opt := reply.Extra[len(reply.Extra)-1]
+	if x, ok := opt.(*dns.OPT); !ok {
+		t.Fatalf("Expected OPT RR, got %T", x)
+	}
+	if !opt.(*dns.OPT).Do() {
+		t.Errorf("Expected DO bit to be set, got false")
+	}
+	if x := opt.(*dns.OPT).UDPSize(); x != defaultUDPBufSize {
+		t.Errorf("Expected %d bufsize, got %d", defaultUDPBufSize, x)
+	}
+
+	// Do set - so left alone.
+	const mysize = defaultUDPBufSize * 2
+	setDo(req)
+	// set bufsize to something else than default to see cache doesn't touch it
+	req.Extra[len(req.Extra)-1].(*dns.OPT).SetUDPSize(mysize)
+	c.ServeDNS(context.TODO(), rec, req)
+	reply = rec.Msg
+	opt = reply.Extra[len(reply.Extra)-1]
+	if x, ok := opt.(*dns.OPT); !ok {
+		t.Fatalf("Expected OPT RR, got %T", x)
+	}
+	if !opt.(*dns.OPT).Do() {
+		t.Errorf("Expected DO bit to be set, got false")
+	}
+	if x := opt.(*dns.OPT).UDPSize(); x != mysize {
+		t.Errorf("Expected %d bufsize, got %d", mysize, x)
+	}
+
+	// edns0 set, but not DO, so _not_ left alone.
+	req.Extra[len(req.Extra)-1].(*dns.OPT).SetDo(false)
+	c.ServeDNS(context.TODO(), rec, req)
+	reply = rec.Msg
+	opt = reply.Extra[len(reply.Extra)-1]
+	if x, ok := opt.(*dns.OPT); !ok {
+		t.Fatalf("Expected OPT RR, got %T", x)
+	}
+	if !opt.(*dns.OPT).Do() {
+		t.Errorf("Expected DO bit to be set, got false")
+	}
+	if x := opt.(*dns.OPT).UDPSize(); x != defaultUDPBufSize {
+		t.Errorf("Expected %d bufsize, got %d", defaultUDPBufSize, x)
+	}
+}
+
+func echoHandler() plugin.Handler {
+	return plugin.HandlerFunc(func(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
+		w.WriteMsg(r)
+		return dns.RcodeSuccess, nil
+	})
+}