plugin/forward: Fix UDP forwarding so a malformed upstream datagram wont block valid ones later (#8287)

* plugin/plugin: Fix UDP forwarding so a malformed upstream datagram wont block valid ones later

This PR fixes UDP forwarding so a malformed upstream datagram
does not prevent CoreDNS from accepting a subsequent valid response
before the existing read deadline.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Address review and add additonal tests

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Cover branch

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

---------

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

plugin/pkg/proxy/connect.go

@@ -186,6 +186,12 @@ func (p *Proxy) lookupDNS(_ctx context.Context, state request.Request, opts Opti
 	for {
 		ret, err = pc.c.ReadMsg()
 		if err != nil {
+			if p.transport.transportTypeFromConn(pc) == typeUDP &&
+				((ret == nil && errors.Is(err, dns.ErrShortRead)) ||
+					(ret != nil && ret.Id != state.Req.Id)) {
+				continue
+			}
+
 			if ret != nil && (state.Req.Id == ret.Id) && p.transport.transportTypeFromConn(pc) == typeUDP && shouldTruncateResponse(err) {
 				// For UDP, if the error is an overflow, we probably have an upstream misbehaving in some way.
 				// (e.g. sending >512 byte responses without an eDNS0 OPT RR).

plugin/pkg/proxy/proxy_test.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"net"
 	"runtime"
+	"strings"
 	"testing"
 	"time"
 
@@ -269,3 +270,124 @@ func TestShouldTruncateResponse(t *testing.T) {
 		})
 	}
 }
+
+func TestProxyMalformedUDPThenValid(t *testing.T) {
+	tests := []struct {
+		name      string
+		malformed func(*dns.Msg) []byte
+	}{
+		{
+			name: "short packet without complete header",
+			malformed: func(r *dns.Msg) []byte {
+				badID := r.Id + 1
+
+				return []byte{
+					byte(badID >> 8),
+					byte(badID),
+					0x00,
+				}
+			},
+		},
+		{
+			name: "partial malformed response with mismatched ID",
+			malformed: func(r *dns.Msg) []byte {
+				badID := r.Id + 1
+
+				// Complete DNS header claiming one question, followed by an
+				// incomplete one-byte label. ReadMsg returns a partial message
+				// containing badID together with an unpacking error.
+				return []byte{
+					byte(badID >> 8),
+					byte(badID),
+					0x81,
+					0x80,
+					0x00,
+					0x01,
+					0x00,
+					0x00,
+					0x00,
+					0x00,
+					0x00,
+					0x00,
+					0x01,
+				}
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			serverResult := make(chan error, 1)
+
+			s := dnstest.NewServer(func(w dns.ResponseWriter, r *dns.Msg) {
+				if _, err := w.Write(tc.malformed(r)); err != nil {
+					serverResult <- err
+					return
+				}
+
+				// Immediately send a completely valid response for the same query.
+				reply := new(dns.Msg)
+				reply.SetReply(r)
+				reply.Answer = append(
+					reply.Answer,
+					test.A("variant.example. IN A 192.0.2.55"),
+				)
+
+				if err := w.WriteMsg(reply); err != nil {
+					serverResult <- err
+					return
+				}
+
+				serverResult <- nil
+			})
+			defer s.Close()
+
+			p := NewProxy(
+				"TestProxyMalformedUDPThenValid",
+				s.Addr,
+				transport.DNS,
+			)
+			p.readTimeout = time.Second
+			p.Start(5 * time.Second)
+			defer p.Stop()
+
+			q := new(dns.Msg)
+			q.SetQuestion("variant.example.", dns.TypeA)
+
+			req := request.Request{
+				Req: q,
+				W:   &test.ResponseWriter{},
+			}
+
+			resp, _, _, err := p.Connect(
+				context.Background(),
+				req,
+				Options{PreferUDP: true},
+			)
+
+			if serverErr := <-serverResult; serverErr != nil {
+				t.Fatalf("upstream failed to send responses: %v", serverErr)
+			}
+
+			// Expected: ignore the malformed UDP datagram and read the valid response.
+			if err != nil {
+				t.Fatalf(
+					"valid response after malformed UDP datagram was not accepted: %v",
+					err,
+				)
+			}
+
+			if resp == nil {
+				t.Fatal("expected valid response, got nil")
+			}
+
+			if len(resp.Answer) != 1 {
+				t.Fatalf("expected one answer, got %d", len(resp.Answer))
+			}
+
+			if got := resp.Answer[0].String(); !strings.Contains(got, "192.0.2.55") {
+				t.Fatalf("expected 192.0.2.55, got %q", got)
+			}
+		})
+	}
+}