Pull request #420: 484-fix-quic-buffers

Merge in GO/dnsproxy from 484-fix-quic-buffers to master

* commit '70dfcdc18d5992a507781932567ad48a7e91dad7':
  upstream: imp test
  upstream: imp code
  all: imp code, docs
  all: imp test code
  Fix DoQ server parsing to avoid stale pooled bytes
  Fix DoQ response parsing and add partial read regression test

Author: EugeneOne1

proxy/proxy_internal_test.go

@@ -63,6 +63,8 @@ var defaultTrustedProxies netutil.SubnetSet = netutil.SliceSubnetSet{
 // mustNew wraps [New] function failing the test on error.
 //
 // TODO(e.burkov):  Move into the proxytest package.
+//
+// TODO(e.burkov):  Put default values into configuration.
 func mustNew(tb testing.TB, conf *Config) (p *Proxy) {
 	tb.Helper()
 

proxy/serverquic.go

@@ -248,9 +248,9 @@ func (p *Proxy) handleQUICStream(ctx context.Context, stream *quic.Stream, conn
 	// draft DNS messages were not prefixed with the message length.
 	packetLen := binary.BigEndian.Uint16(buf[:2])
 	if packetLen == uint16(n-2) {
-		err = req.Unpack(buf[2:])
+		err = req.Unpack(buf[2:n])
 	} else {
-		err = req.Unpack(buf)
+		err = req.Unpack(buf[:n])
 		doqVersion = DoQv1Draft
 	}
 

proxy/serverquic_internal_test.go

@@ -11,14 +11,16 @@ import (
 
 	"github.com/AdguardTeam/dnsproxy/proxyutil"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
+	"github.com/AdguardTeam/golibs/syncutil"
 	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/AdguardTeam/golibs/testutil/servicetest"
 	"github.com/miekg/dns"
 	"github.com/quic-go/quic-go"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-func TestQuicProxy(t *testing.T) {
+func TestProxy_quic(t *testing.T) {
 	serverConfig, caPem := newTLSConfig(t)
 
 	roots := x509.NewCertPool()
@@ -83,7 +85,7 @@ func TestQuicProxy(t *testing.T) {
 	})
 }
 
-func TestQuicProxy_largePackets(t *testing.T) {
+func TestProxy_quicLargePackets(t *testing.T) {
 	reqHandler := &TestHandler{
 		OnHandle: func(_ *Proxy, d *DNSContext) (err error) {
 			d.Res = newTestResponse(d)
@@ -141,6 +143,88 @@ func TestQuicProxy_largePackets(t *testing.T) {
 	requireResponse(t, msg, resp)
 }
 
+func TestProxy_quicTruncatedRequest(t *testing.T) {
+	serverConfig, caPem := newTLSConfig(t)
+
+	conf := &Config{
+		Logger:                 slogutil.NewDiscardLogger(),
+		QUICListenAddr:         []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
+		TLSConfig:              serverConfig,
+		UpstreamConfig:         newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr),
+		TrustedProxies:         defaultTrustedProxies,
+		RatelimitSubnetLenIPv4: 24,
+		RatelimitSubnetLenIPv6: 64,
+		RequestHandler: &TestHandler{
+			OnHandle: func(p *Proxy, d *DNSContext) (_ error) {
+				panic(testutil.UnexpectedCall(p, d))
+			},
+		},
+	}
+
+	dnsProxy := mustNew(t, conf)
+
+	req := (&dns.Msg{
+		MsgHdr: dns.MsgHdr{
+			Id:               0,
+			RecursionDesired: true,
+		},
+		Question: []dns.Question{{
+			Name:   "example.org.",
+			Qtype:  dns.TypeA,
+			Qclass: dns.ClassINET,
+		}},
+	}).SetEdns0(4096, false)
+
+	packed, err := req.Pack()
+	require.NoError(t, err)
+
+	fullBuf := proxyutil.AddPrefix(packed)
+
+	dnsProxy.bytesPool = syncutil.NewPool(func() (v *[]byte) {
+		b := make([]byte, 2+dns.MaxMsgSize)
+		copy(b, fullBuf)
+
+		return &b
+	})
+
+	servicetest.RequireRun(t, dnsProxy, testTimeout)
+
+	addr := dnsProxy.Addr(ProtoQUIC)
+
+	roots := x509.NewCertPool()
+	require.True(t, roots.AppendCertsFromPEM(caPem))
+
+	tlsConfig := &tls.Config{
+		ServerName: tlsServerName,
+		RootCAs:    roots,
+		NextProtos: append([]string{NextProtoDQ}, compatProtoDQ...),
+	}
+
+	ctx := testutil.ContextWithTimeout(t, testTimeout)
+
+	conn, err := quic.DialAddrEarly(ctx, addr.String(), tlsConfig, nil)
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, func() (err error) {
+		return conn.CloseWithError(DoQCodeNoError, "")
+	})
+
+	truncLen := len(packed) - len(packed)/2
+	require.Greater(t, truncLen, 0)
+
+	truncated := packed[:truncLen]
+	reqBuf := proxyutil.AddPrefix(truncated)
+	require.Greater(t, len(reqBuf), minDNSPacketSize)
+
+	stream, err := conn.OpenStreamSync(ctx)
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, stream.Close)
+
+	n, err := stream.Write(reqBuf)
+	require.NoError(t, err)
+
+	assert.Greater(t, n, minDNSPacketSize)
+}
+
 // newTestResponse creates a test response for the specified request.
 func newTestResponse(d *DNSContext) (resp *dns.Msg) {
 	resp = &dns.Msg{}

upstream/doq.go

@@ -3,7 +3,9 @@ package upstream
 import (
 	"context"
 	"crypto/tls"
+	"encoding/binary"
 	"fmt"
+	"io"
 	"log/slog"
 	"net"
 	"net/url"
@@ -15,6 +17,7 @@ import (
 	"github.com/AdguardTeam/dnsproxy/proxyutil"
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
+	"github.com/AdguardTeam/golibs/validate"
 	"github.com/miekg/dns"
 	"github.com/quic-go/quic-go"
 	"github.com/quic-go/quic-go/http3"
@@ -95,6 +98,23 @@ type dnsOverQUIC struct {
 	timeout time.Duration
 }
 
+// quicStream is the interface of QUIC stream used by readMsg to simplify
+// testing.
+//
+// Note, that this interface is implemented by [*quic.Stream].
+type quicStream interface {
+	// Read reads data from the stream.  If the stream was canceled, the error
+	// is a [quic.StreamError].
+	Read(p []byte) (n int, err error)
+
+	// CancelRead aborts receiving on this stream.  See
+	// [quic.ReceiveStream.CancelRead] for more details.
+	CancelRead(code quic.StreamErrorCode)
+}
+
+// type check
+var _ quicStream = (*quic.Stream)(nil)
+
 // newDoQ returns the DNS-over-QUIC Upstream.
 func newDoQ(addr *url.URL, opts *Options) (u Upstream, err error) {
 	addPort(addr, defaultPortDoQ)
@@ -397,29 +417,42 @@ func (p *dnsOverQUIC) closeConnWithError(conn *quic.Conn, err error) {
 }
 
 // readMsg reads the incoming DNS message from the QUIC stream.
-func (p *dnsOverQUIC) readMsg(stream *quic.Stream) (m *dns.Msg, err error) {
+func (p *dnsOverQUIC) readMsg(stream quicStream) (m *dns.Msg, err error) {
+	defer func() { err = errors.Annotate(err, "from %s: %w", p.addr) }()
+
+	var lenBuf [2]byte
+	_, err = io.ReadFull(stream, lenBuf[:])
+	if err != nil {
+		return nil, fmt.Errorf("reading response length: %w", err)
+	}
+
+	msgLen := binary.BigEndian.Uint16(lenBuf[:])
+	err = validate.Positive("response length", msgLen)
+	if err != nil {
+		// Don't wrap the error, since it's informative enough as is.
+		return nil, err
+	}
+
 	pool := p.getBytesPool()
 	bufPtr := pool.Get().(*[]byte)
-
 	defer pool.Put(bufPtr)
 
-	respBuf := *bufPtr
-	n, err := stream.Read(respBuf)
-	if err != nil && n == 0 {
-		return nil, fmt.Errorf("reading response from %s: %w", p.addr, err)
+	respBuf := (*bufPtr)[:msgLen]
+
+	_, err = io.ReadFull(stream, respBuf)
+	if err != nil {
+		return nil, err
 	}
 
 	stream.CancelRead(0)
 
 	// All DNS messages (queries and responses) sent over DoQ connections MUST
 	// be encoded as a 2-octet length field followed by the message content as
 	// specified in [RFC1035].
-	// IMPORTANT: Note, that we ignore this prefix here as this implementation
-	// does not support receiving multiple messages over a single connection.
-	m = new(dns.Msg)
-	err = m.Unpack(respBuf[2:])
+	m = &dns.Msg{}
+	err = m.Unpack(respBuf)
 	if err != nil {
-		return nil, fmt.Errorf("unpacking response from %s: %w", p.addr, err)
+		return nil, fmt.Errorf("from %s: unpacking response: %w", p.addr, err)
 	}
 
 	return m, nil

upstream/doq_internal_test.go

@@ -24,7 +24,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestUpstreamDoQ(t *testing.T) {
+func TestDNSOverQUIC(t *testing.T) {
 	tlsConf, rootCAs := createServerTLSConfig(t, "127.0.0.1")
 
 	srv := startDoQServer(t, tlsConf, 0)
@@ -77,7 +77,7 @@ func TestUpstreamDoQ(t *testing.T) {
 	checkRaceCondition(u)
 }
 
-func TestUpstream_Exchange_quicServerCloseConn(t *testing.T) {
+func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
 	// Use the same tlsConf for all servers to preserve the data necessary for
 	// 0-RTT connections.
 	tlsConf, rootCAs := createServerTLSConfig(t, "127.0.0.1")
@@ -129,7 +129,7 @@ func TestUpstream_Exchange_quicServerCloseConn(t *testing.T) {
 	wg.Wait()
 }
 
-func TestUpstreamDoQ_serverRestart(t *testing.T) {
+func TestDNSOverQUIC_serverRestart(t *testing.T) {
 	t.Parallel()
 
 	// Use the same tlsConf for all servers to preserve the data necessary for
@@ -182,7 +182,7 @@ func TestUpstreamDoQ_serverRestart(t *testing.T) {
 	})
 }
 
-func TestUpstreamDoQ_0RTT(t *testing.T) {
+func TestDNSOverQUIC_0RTT(t *testing.T) {
 	tlsConf, rootCAs := createServerTLSConfig(t, "127.0.0.1")
 
 	srv := startDoQServer(t, tlsConf, 0)
@@ -232,6 +232,112 @@ func TestUpstreamDoQ_0RTT(t *testing.T) {
 	require.True(t, conns[1].is0RTT())
 }
 
+func TestDNSOverQUIC_ReadMsg_partialRead(t *testing.T) {
+	oldReq := createHostTestMessage("old.example")
+	oldResp := (&dns.Msg{}).SetReply(oldReq)
+	oldResp.Answer = []dns.RR{&dns.A{
+		Hdr: dns.RR_Header{
+			Name:   dns.Fqdn("old.example"),
+			Rrtype: dns.TypeA,
+			Class:  dns.ClassINET,
+			Ttl:    60,
+		},
+		A: net.IP{192, 0, 2, 1},
+	}}
+
+	oldPacked, err := oldResp.Pack()
+	require.NoError(t, err)
+
+	oldBuf := make([]byte, 2+len(oldPacked))
+	binary.BigEndian.PutUint16(oldBuf, uint16(len(oldPacked)))
+	copy(oldBuf[2:], oldPacked)
+
+	pool := &sync.Pool{
+		New: func() (buf any) {
+			b := make([]byte, dns.MaxMsgSize)
+			copy(b, oldBuf)
+
+			return &b
+		},
+	}
+
+	ups, err := newDoQ(&url.URL{Host: "example:853"}, &Options{})
+	require.NoError(t, err)
+	testutil.CleanupAndRequireSuccess(t, ups.Close)
+
+	doq := testutil.RequireTypeAssert[*dnsOverQUIC](t, ups)
+	doq.bytesPool = pool
+
+	newAnsIP := net.IP{192, 0, 2, 2}
+
+	newReq := createHostTestMessage("new.example")
+	newResp := (&dns.Msg{}).SetReply(newReq)
+	newResp.Answer = []dns.RR{&dns.A{
+		Hdr: dns.RR_Header{
+			Name:   dns.Fqdn("new.example"),
+			Rrtype: dns.TypeA,
+			Class:  dns.ClassINET,
+			Ttl:    60,
+		},
+		A: newAnsIP,
+	}}
+
+	newPacked, err := newResp.Pack()
+	require.NoError(t, err)
+
+	newBuf := make([]byte, 2+len(newPacked))
+	binary.BigEndian.PutUint16(newBuf, uint16(len(newPacked)))
+	copy(newBuf[2:], newPacked)
+
+	stream := &testQUICStream{
+		onRead: func(p []byte) (n int, err error) {
+			if len(newBuf) == 0 {
+				return 0, io.EOF
+			}
+
+			n = min(1, len(newBuf), len(p))
+
+			copy(p, newBuf[:n])
+			newBuf = newBuf[n:]
+
+			return n, nil
+		},
+		onCancelRead: func(code quic.StreamErrorCode) {
+			assert.Equal(t, quic.StreamErrorCode(0), code)
+		},
+	}
+
+	got, err := doq.readMsg(stream)
+	require.NoError(t, err)
+	require.Len(t, got.Answer, 1)
+
+	assert.Equal(t, dns.Fqdn("new.example"), got.Question[0].Name)
+
+	a := testutil.RequireTypeAssert[*dns.A](t, got.Answer[0])
+
+	assert.Equal(t, newAnsIP, a.A)
+}
+
+// testQUICStream is a mock implementation of the [quicStream] interface for
+// tests.
+type testQUICStream struct {
+	onRead       func(p []byte) (n int, err error)
+	onCancelRead func(code quic.StreamErrorCode)
+}
+
+// type check
+var _ quicStream = (*testQUICStream)(nil)
+
+// Read implements the [quicStream] interface for *testQUICStream.
+func (s *testQUICStream) Read(p []byte) (n int, err error) {
+	return s.onRead(p)
+}
+
+// CancelRead implements the [quicStream] interface for *testQUICStream.
+func (s *testQUICStream) CancelRead(code quic.StreamErrorCode) {
+	s.onCancelRead(code)
+}
+
 // testDoHServer is an instance of a test DNS-over-QUIC server.
 type testDoQServer struct {
 	// listener is the QUIC connections listener.