all: imp test code

Author: EugeneOne1

proxy/serverquic_internal_test.go

@@ -6,7 +6,6 @@ import (
 	"crypto/x509"
 	"io"
 	"net"
-	"sync/atomic"
 	"testing"
 	"time"
 
@@ -17,6 +16,7 @@ import (
 	"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"
 )
 
@@ -143,10 +143,9 @@ func TestQuicProxy_largePackets(t *testing.T) {
 	requireResponse(t, msg, resp)
 }
 
-func TestQuicProxy_truncatedRequestUsesOnlyReadBytes(t *testing.T) {
+func TestQuicProxy_truncatedRequest(t *testing.T) {
 	serverConfig, caPem := newTLSConfig(t)
 
-	var handled atomic.Bool
 	conf := &Config{
 		Logger:                 slogutil.NewDiscardLogger(),
 		QUICListenAddr:         []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
@@ -156,27 +155,15 @@ func TestQuicProxy_truncatedRequestUsesOnlyReadBytes(t *testing.T) {
 		RatelimitSubnetLenIPv4: 24,
 		RatelimitSubnetLenIPv6: 64,
 		RequestHandler: &TestRequestHandler{
-			OnHandle: func(p *Proxy, dctx *DNSContext) (err error) {
-				handled.Store(true)
-				resp := (&dns.Msg{}).SetReply(dctx.Req)
-				resp.Answer = []dns.RR{&dns.A{
-					Hdr: dns.RR_Header{
-						Name:   dctx.Req.Question[0].Name,
-						Rrtype: dns.TypeA,
-						Class:  dns.ClassINET,
-					},
-					A: net.IP{8, 8, 8, 8},
-				}}
-				dctx.Res = resp
-
-				return nil
+			OnHandle: func(p *Proxy, d *DNSContext) (_ error) {
+				panic(testutil.UnexpectedCall(p, d))
 			},
 		},
 	}
 
 	dnsProxy := mustNew(t, conf)
 
-	req := &dns.Msg{
+	req := (&dns.Msg{
 		MsgHdr: dns.MsgHdr{
 			Id:               0,
 			RecursionDesired: true,
@@ -186,8 +173,7 @@ func TestQuicProxy_truncatedRequestUsesOnlyReadBytes(t *testing.T) {
 			Qtype:  dns.TypeA,
 			Qclass: dns.ClassINET,
 		}},
-	}
-	req.SetEdns0(4096, false)
+	}).SetEdns0(4096, false)
 
 	packed, err := req.Pack()
 	require.NoError(t, err)
@@ -203,36 +189,39 @@ func TestQuicProxy_truncatedRequestUsesOnlyReadBytes(t *testing.T) {
 
 	servicetest.RequireRun(t, dnsProxy, testTimeout)
 
+	addr := dnsProxy.Addr(ProtoQUIC)
+
 	roots := x509.NewCertPool()
-	roots.AppendCertsFromPEM(caPem)
+	require.True(t, roots.AppendCertsFromPEM(caPem))
+
 	tlsConfig := &tls.Config{
 		ServerName: tlsServerName,
 		RootCAs:    roots,
 		NextProtos: append([]string{NextProtoDQ}, compatProtoDQ...),
 	}
 
-	addr := dnsProxy.Addr(ProtoQUIC)
-	conn, err := quic.DialAddrEarly(context.Background(), addr.String(), tlsConfig, nil)
+	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) - 5
+	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(context.Background())
-	require.NoError(t, err)
-	_, err = stream.Write(reqBuf)
+	stream, err := conn.OpenStreamSync(ctx)
 	require.NoError(t, err)
-	_ = stream.Close()
+	testutil.CleanupAndRequireSuccess(t, stream.Close)
 
-	require.Never(t, func() bool {
-		return handled.Load()
-	}, testTimeout, 10*time.Millisecond)
+	n, err := stream.Write(reqBuf)
+	require.NoError(t, err)
+	assert.Greater(t, n, minDNSPacketSize)
 }
 
 // newTestResponse creates a test response for the specified request.

upstream/doq_internal_test.go

@@ -232,7 +232,7 @@ func TestUpstreamDoQ_0RTT(t *testing.T) {
 	require.True(t, conns[1].is0RTT())
 }
 
-func TestUpstreamDoQ_readMsg_partialRead(t *testing.T) {
+func TestUpstreamDoQ_ReadMsg_partialRead(t *testing.T) {
 	oldReq := createHostTestMessage("old.example")
 	oldResp := (&dns.Msg{}).SetReply(oldReq)
 	oldResp.Answer = []dns.RR{&dns.A{
@@ -242,23 +242,34 @@ func TestUpstreamDoQ_readMsg_partialRead(t *testing.T) {
 			Class:  dns.ClassINET,
 			Ttl:    60,
 		},
-		A: net.IPv4(1, 1, 1, 1),
+		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() any {
+		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{
@@ -268,11 +279,12 @@ func TestUpstreamDoQ_readMsg_partialRead(t *testing.T) {
 			Class:  dns.ClassINET,
 			Ttl:    60,
 		},
-		A: net.IPv4(2, 2, 2, 2),
+		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)
@@ -282,20 +294,14 @@ func TestUpstreamDoQ_readMsg_partialRead(t *testing.T) {
 		chunk: 1,
 	}
 
-	p := &dnsOverQUIC{
-		addr:        &url.URL{Host: "example:853"},
-		bytesPool:   pool,
-		bytesPoolMu: &sync.Mutex{},
-	}
-
-	got, err := p.readMsg(stream)
+	got, err := doq.readMsg(stream)
 	require.NoError(t, err)
 	require.Len(t, got.Answer, 1)
-	require.Equal(t, dns.Fqdn("new.example"), got.Question[0].Name)
 
-	a, ok := got.Answer[0].(*dns.A)
-	require.True(t, ok, "wrong answer type: %T", got.Answer[0])
-	require.Equal(t, net.IPv4(2, 2, 2, 2).To4(), a.A)
+	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)
 }
 
 type testQUICStream struct {