Pull request 243: Fix an issue with processing DNS messages that are …

…split into multiple QUIC frames

Merge in DNS/dnsproxy from fix-315-quic-large-packets to master

Fix #315

Squashed commit of the following:

commit 3321494
Author: Andrey Meshkov <am@adguard.com>
Date:   Tue Feb 7 18:37:38 2023 +0300

    fix review comments

commit 878f754
Author: Andrey Meshkov <am@adguard.com>
Date:   Tue Feb 7 18:24:58 2023 +0300

    fix review comments

commit 17e1c51
Author: Andrey Meshkov <am@adguard.com>
Date:   Tue Feb 7 18:14:17 2023 +0300

    Fix an issue with processing DNS messages that are split into multiple QUIC frames

    Fix #315

proxy/server_quic.go

@@ -157,7 +157,7 @@ func (p *Proxy) handleQUICStream(stream quic.Stream, conn quic.Connection) {
 	// indicated as error here.  Instead, we should check the number of bytes
 	// received.
 	buf := *bufPtr
-	n, err := stream.Read(buf)
+	n, err := readAll(stream, buf)
 
 	// Note that io.EOF does not really mean that there's any error, this is
 	// just a signal that there will be no data to read anymore from this
@@ -318,6 +318,7 @@ func logShortQUICRead(err error) {
 
 // isQUICNonCrit returns true if err is a non-critical error, most probably
 // related to the current QUIC implementation.
+//
 // TODO(ameshkov): re-test when updating quic-go.
 func isQUICNonCrit(err error) (ok bool) {
 	if err == nil {
@@ -414,3 +415,29 @@ func (v *quicAddrValidator) requiresValidation(addr net.Addr) (ok bool) {
 	// will require address validation.
 	return true
 }
+
+// readAll reads from r until an error or io.EOF into the specified buffer buf.
+// A successful call returns err == nil, not err == io.EOF.  If the buffer is
+// too small, it returns error io.ErrShortBuffer.  This function has some
+// similarities to io.ReadAll, but it reads to the specified buffer and not
+// allocates (and grows) a new one.  Also, it is completely different from
+// io.ReadFull as that one reads the exact number of bytes (buffer length) and
+// readAll reads until io.EOF or until the buffer is filled.
+func readAll(r io.Reader, buf []byte) (n int, err error) {
+	for {
+		if n == len(buf) {
+			return n, io.ErrShortBuffer
+		}
+
+		var read int
+		read, err = r.Read(buf[n:])
+		n += read
+
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+			}
+			return n, err
+		}
+	}
+}

proxy/server_quic_test.go

@@ -5,9 +5,12 @@ import (
 	"crypto/tls"
 	"crypto/x509"
 	"io"
+	"net"
 	"testing"
+	"time"
 
 	"github.com/AdguardTeam/dnsproxy/proxyutil"
+	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/miekg/dns"
 	"github.com/quic-go/quic-go"
 	"github.com/stretchr/testify/require"
@@ -17,6 +20,7 @@ func TestQuicProxy(t *testing.T) {
 	// Prepare the proxy server.
 	serverConfig, caPem := createServerTLSConfig(t)
 	dnsProxy := createTestProxy(t, serverConfig)
+	testutil.CleanupAndRequireSuccess(t, dnsProxy.Stop)
 
 	// Start listening.
 	err := dnsProxy.Start()
@@ -33,13 +37,12 @@ func TestQuicProxy(t *testing.T) {
 	// Create a DNS-over-QUIC client connection.
 	addr := dnsProxy.Addr(ProtoQUIC)
 
-	// Open QUIC connection.
+	// Open a QUIC connection.
 	conn, err := quic.DialAddrEarly(addr.String(), tlsConfig, nil)
 	require.NoError(t, err)
-	defer func() {
-		// TODO(ameshkov): check the error here.
-		_ = conn.CloseWithError(DoQCodeNoError, "")
-	}()
+	testutil.CleanupAndRequireSuccess(t, func() (err error) {
+		return conn.CloseWithError(DoQCodeNoError, "")
+	})
 
 	// Send several test messages.
 	for i := 0; i < 10; i++ {
@@ -48,23 +51,83 @@ func TestQuicProxy(t *testing.T) {
 		// Send a message encoded for a draft version as well.
 		sendTestQUICMessage(t, conn, DoQv1Draft)
 	}
+}
 
-	// Stop the proxy.
-	err = dnsProxy.Stop()
-	if err != nil {
-		t.Fatalf("cannot stop the DNS proxy: %s", err)
+func TestQuicProxy_largePackets(t *testing.T) {
+	// Prepare the proxy server.
+	serverConfig, caPem := createServerTLSConfig(t)
+	dnsProxy := createTestProxy(t, serverConfig)
+
+	// Make sure the request does not go to any real upstream.
+	dnsProxy.RequestHandler = func(p *Proxy, d *DNSContext) (err error) {
+		resp := &dns.Msg{}
+		resp.SetReply(d.Req)
+		resp.Answer = []dns.RR{
+			&dns.A{
+				Hdr: dns.RR_Header{
+					Name:   d.Req.Question[0].Name,
+					Rrtype: dns.TypeA,
+					Class:  dns.ClassINET,
+				},
+				A: net.IP{8, 8, 8, 8},
+			},
+		}
+		d.Res = resp
+
+		return nil
 	}
+
+	testutil.CleanupAndRequireSuccess(t, dnsProxy.Stop)
+
+	// Start listening.
+	err := dnsProxy.Start()
+	require.NoError(t, err)
+
+	roots := x509.NewCertPool()
+	roots.AppendCertsFromPEM(caPem)
+	tlsConfig := &tls.Config{
+		ServerName: tlsServerName,
+		RootCAs:    roots,
+		NextProtos: append([]string{NextProtoDQ}, compatProtoDQ...),
+	}
+
+	// Create a DNS-over-QUIC client connection.
+	addr := dnsProxy.Addr(ProtoQUIC)
+
+	// Open a QUIC connection.
+	conn, err := quic.DialAddrEarly(addr.String(), tlsConfig, nil)
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, func() (err error) {
+		return conn.CloseWithError(DoQCodeNoError, "")
+	})
+
+	// Create a test message large enough to take multiple QUIC frames.
+	msg := createTestMessage()
+	msg.Extra = []dns.RR{
+		&dns.OPT{
+			Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT, Class: 4096},
+			Option: []dns.EDNS0{
+				&dns.EDNS0_PADDING{Padding: make([]byte, 4096)},
+			},
+		},
+	}
+
+	resp := sendQUICMessage(t, msg, conn, DoQv1)
+	requireResponse(t, msg, resp)
 }
 
-// sendTestQUICMessage send a test message to the specified QUIC connection.
-func sendTestQUICMessage(t *testing.T, conn quic.Connection, doqVersion DoQVersion) {
-	// Open a new stream.
+// sendQUICMessage sends msg to the specified QUIC connection.
+func sendQUICMessage(
+	t *testing.T,
+	msg *dns.Msg,
+	conn quic.Connection,
+	doqVersion DoQVersion,
+) (resp *dns.Msg) {
+	// Open a new QUIC stream to write there a test DNS query.
 	stream, err := conn.OpenStreamSync(context.Background())
 	require.NoError(t, err)
-	defer stream.Close()
+	testutil.CleanupAndRequireSuccess(t, stream.Close)
 
-	// Prepare a test message.
-	msg := createTestMessage()
 	packedMsg, err := msg.Pack()
 	require.NoError(t, err)
 
@@ -74,7 +137,7 @@ func sendTestQUICMessage(t *testing.T, conn quic.Connection, doqVersion DoQVersi
 	}
 
 	// Send the DNS query to the stream.
-	_, err = stream.Write(buf)
+	err = writeQUICStream(buf, stream)
 	require.NoError(t, err)
 
 	// Close closes the write-direction of the stream and sends
@@ -90,14 +153,47 @@ func sendTestQUICMessage(t *testing.T, conn quic.Connection, doqVersion DoQVersi
 	require.Greater(t, n, minDNSPacketSize)
 
 	// Unpack the DNS response.
-	reply := new(dns.Msg)
+	resp = new(dns.Msg)
 	if doqVersion == DoQv1 {
-		err = reply.Unpack(respBytes[2:])
+		err = resp.Unpack(respBytes[2:])
 	} else {
-		err = reply.Unpack(respBytes)
+		err = resp.Unpack(respBytes)
 	}
 	require.NoError(t, err)
 
-	// Check the response
-	requireResponse(t, msg, reply)
+	return resp
+}
+
+// writeQUICStream writes buf to the specified QUIC stream in chunks.  This way
+// it is possible to test how the server deals with chunked DNS messages.
+func writeQUICStream(buf []byte, stream quic.Stream) (err error) {
+	// Send the DNS query to the stream and split it into chunks of up
+	// to 400 bytes.  400 is an arbitrary chosen value.
+	chunkSize := 400
+	for i := 0; i < len(buf); i += chunkSize {
+		chunkStart := i
+		chunkEnd := i + chunkSize
+		if chunkEnd > len(buf) {
+			chunkEnd = len(buf)
+		}
+
+		_, err = stream.Write(buf[chunkStart:chunkEnd])
+		if err != nil {
+			return err
+		}
+
+		if len(buf) > chunkSize {
+			// Emulate network latency.
+			time.Sleep(time.Millisecond)
+		}
+	}
+
+	return nil
+}
+
+// sendTestQUICMessage send a test message to the specified QUIC connection.
+func sendTestQUICMessage(t *testing.T, conn quic.Connection, doqVersion DoQVersion) {
+	msg := createTestMessage()
+	resp := sendQUICMessage(t, msg, conn, doqVersion)
+	requireResponse(t, msg, resp)
 }