plugin/forward: check TC correctly on reply. (#1670)

* plugin/forward: check TC correctly on reply.

Add test for this.

* Add proxy test as well

1 files changed, 114 insertions, 0 deletions

diff --git a/plugin/forward/truncated_test.go b/plugin/forward/truncated_test.go
new file mode 100644
index 000000000..4d8a0a25e
--- /dev/null
+++ b/plugin/forward/truncated_test.go
@@ -0,0 +1,114 @@
+package forward
+
+import (
+	"sync/atomic"
+	"testing"
+
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
+	"github.com/coredns/coredns/plugin/test"
+	"github.com/coredns/coredns/request"
+
+	"github.com/miekg/dns"
+)
+
+func TestLookupTruncated(t *testing.T) {
+	i := int32(0)
+	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		j := atomic.LoadInt32(&i)
+		atomic.AddInt32(&i, 1)
+
+		if j == 0 {
+			ret := new(dns.Msg)
+			ret.SetReply(r)
+			ret.Truncated = true
+			ret.Answer = append(ret.Answer, test.A("example.org. IN A 127.0.0.1"))
+			w.WriteMsg(ret)
+			return
+
+		}
+
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+		ret.Answer = append(ret.Answer, test.A("example.org. IN A 127.0.0.1"))
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	p := NewProxy(s.Addr, nil /* no TLS */)
+	f := New()
+	f.SetProxy(p)
+	defer f.Close()
+
+	state := request.Request{W: &test.ResponseWriter{}, Req: new(dns.Msg)}
+
+	resp, err := f.Lookup(state, "example.org.", dns.TypeA)
+	if err != nil {
+		t.Fatal("Expected to receive reply, but didn't")
+	}
+	// expect answer with TC
+	if !resp.Truncated {
+		t.Error("Expected to receive reply with TC bit set, but didn't")
+	}
+
+	resp, err = f.Lookup(state, "example.org.", dns.TypeA)
+	if err != nil {
+		t.Fatal("Expected to receive reply, but didn't")
+	}
+	// expect answer without TC
+	if resp.Truncated {
+		t.Error("Expected to receive reply without TC bit set, but didn't")
+	}
+}
+
+func TestForwardTruncated(t *testing.T) {
+	i := int32(0)
+	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		j := atomic.LoadInt32(&i)
+		atomic.AddInt32(&i, 1)
+
+		if j == 0 {
+			ret := new(dns.Msg)
+			ret.SetReply(r)
+			ret.Truncated = true
+			ret.Answer = append(ret.Answer, test.A("example.org. IN A 127.0.0.1"))
+			w.WriteMsg(ret)
+			return
+
+		}
+
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+		ret.Answer = append(ret.Answer, test.A("example.org. IN A 127.0.0.1"))
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	f := New()
+
+	p1 := NewProxy(s.Addr, nil /* no TLS */)
+	f.SetProxy(p1)
+	p2 := NewProxy(s.Addr, nil /* no TLS */)
+	f.SetProxy(p2)
+	defer f.Close()
+
+	state := request.Request{W: &test.ResponseWriter{}, Req: new(dns.Msg)}
+	state.Req.SetQuestion("example.org.", dns.TypeA)
+	resp, err := f.Forward(state)
+	if err != nil {
+		t.Fatal("Expected to receive reply, but didn't")
+	}
+
+	// expect answer with TC
+	if !resp.Truncated {
+		t.Error("Expected to receive reply with TC bit set, but didn't")
+	}
+
+	resp, err = f.Forward(state)
+	if err != nil {
+		t.Fatal("Expected to receive reply, but didn't")
+	}
+	// expect answer without TC
+	if resp.Truncated {
+		t.Error("Expected to receive reply without TC bit set, but didn't")
+	}
+}