Handle UDP responses that overflow with TC bit with test case (#6277)

Signed-off-by: SriHarshaBS001 <SriHarshaBS009@gmail.com>

1 files changed, 95 insertions, 0 deletions

diff --git a/plugin/pkg/proxy/proxy_test.go b/plugin/pkg/proxy/proxy_test.go
index 33a7170c0..03d10ce5f 100644
--- a/plugin/pkg/proxy/proxy_test.go
+++ b/plugin/pkg/proxy/proxy_test.go
@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"crypto/tls"
+	"errors"
 	"math"
 	"testing"
 	"time"
@@ -128,3 +129,97 @@ func TestProxyIncrementFails(t *testing.T) {
 		})
 	}
 }
+
+func TestCoreDNSOverflow(t *testing.T) {
+	s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+		ret := new(dns.Msg)
+		ret.SetReply(r)
+
+		answers := []dns.RR{
+			test.A("example.org. IN A 127.0.0.1"),
+			test.A("example.org. IN A 127.0.0.2"),
+			test.A("example.org. IN A 127.0.0.3"),
+			test.A("example.org. IN A 127.0.0.4"),
+			test.A("example.org. IN A 127.0.0.5"),
+			test.A("example.org. IN A 127.0.0.6"),
+			test.A("example.org. IN A 127.0.0.7"),
+			test.A("example.org. IN A 127.0.0.8"),
+			test.A("example.org. IN A 127.0.0.9"),
+			test.A("example.org. IN A 127.0.0.10"),
+			test.A("example.org. IN A 127.0.0.11"),
+			test.A("example.org. IN A 127.0.0.12"),
+			test.A("example.org. IN A 127.0.0.13"),
+			test.A("example.org. IN A 127.0.0.14"),
+			test.A("example.org. IN A 127.0.0.15"),
+			test.A("example.org. IN A 127.0.0.16"),
+			test.A("example.org. IN A 127.0.0.17"),
+			test.A("example.org. IN A 127.0.0.18"),
+			test.A("example.org. IN A 127.0.0.19"),
+			test.A("example.org. IN A 127.0.0.20"),
+		}
+		ret.Answer = answers
+		w.WriteMsg(ret)
+	})
+	defer s.Close()
+
+	p := NewProxy("TestCoreDNSOverflow", s.Addr, transport.DNS)
+	p.readTimeout = 10 * time.Millisecond
+	p.Start(5 * time.Second)
+	defer p.Stop()
+
+	// Test different connection modes
+	testConnection := func(proto string, options Options, expectTruncated bool) {
+		t.Helper()
+
+		queryMsg := new(dns.Msg)
+		queryMsg.SetQuestion("example.org.", dns.TypeA)
+
+		recorder := dnstest.NewRecorder(&test.ResponseWriter{})
+		request := request.Request{Req: queryMsg, W: recorder}
+
+		response, err := p.Connect(context.Background(), request, options)
+		if err != nil {
+			t.Errorf("Failed to connect to testdnsserver: %s", err)
+		}
+
+		if response.Truncated != expectTruncated {
+			t.Errorf("Expected truncated response for %s, but got TC flag %v", proto, response.Truncated)
+		}
+	}
+
+	// Test PreferUDP, expect truncated response
+	testConnection("PreferUDP", Options{PreferUDP: true}, true)
+
+	// Test ForceTCP, expect no truncated response
+	testConnection("ForceTCP", Options{ForceTCP: true}, false)
+
+	// Test No options specified, expect truncated response
+	testConnection("NoOptionsSpecified", Options{}, true)
+
+	// Test both TCP and UDP provided, expect no truncated response
+	testConnection("BothTCPAndUDP", Options{PreferUDP: true, ForceTCP: true}, false)
+}
+
+func TestShouldTruncateResponse(t *testing.T) {
+	testCases := []struct {
+		testname string
+		err      error
+		expected bool
+	}{
+		{"BadAlgorithm", dns.ErrAlg, false},
+		{"BufferSizeTooSmall", dns.ErrBuf, true},
+		{"OverflowUnpackingA", errors.New("overflow unpacking a"), true},
+		{"OverflowingHeaderSize", errors.New("overflowing header size"), true},
+		{"OverflowpackingA", errors.New("overflow packing a"), true},
+		{"ErrSig", dns.ErrSig, false},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.testname, func(t *testing.T) {
+			result := shouldTruncateResponse(tc.err)
+			if result != tc.expected {
+				t.Errorf("For testname '%v', expected %v but got %v", tc.testname, tc.expected, result)
+			}
+		})
+	}
+}