all: imp code

Author: EugeneOne1

internal/cmd/config.go

@@ -181,8 +181,8 @@ type configuration struct {
 	EnableEDNSSubnet bool `yaml:"edns"`
 
 	// PendingRequestsEnabled controls whether the server should track duplicate
-	// queries and only send the first of them to the upstream server.  It used
-	// to mitigate the cache poisoning attacks.
+	// queries and only send the first of them to the upstream server.  It is
+	// used to mitigate the cache poisoning attacks.
 	PendingRequestsEnabled bool `yaml:"pending-requests-enabled"`
 
 	// DNS64 defines whether DNS64 functionality is enabled or not.

proxy/config.go

@@ -58,10 +58,12 @@ type Config struct {
 	// constructor will be used.
 	MessageConstructor MessageConstructor
 
-	// PendingRequests limits the number of identical requests sent to the
-	// upstream server.  If nil, the default implementation will be used.  Use:
-	//  - [DefaultPendingRequests] to enable the pending requests feature.
-	//  - [EmptyPendingRequests] to disable the pending requests feature.
+	// PendingRequests is used to mitigate the cache poisoning attacks by
+	// tracking identical requests and returning the same response for them,
+	// peforming a single lookup.  If nil, the default implementation will be
+	// used.  Use:
+	//  - [DefaultPendingRequests] to enable pending identical requests.
+	//  - [EmptyPendingRequests] to disable pending identical requests.
 	PendingRequests PendingRequests
 
 	// BeforeRequestHandler is an optional custom handler called before each DNS

proxy/pending.go

@@ -42,10 +42,10 @@ type pendingRequest struct {
 	// channel is closed.
 	resolveErr error
 
-	// cloneCtx is a clone of the DNSContext that was used to create the
+	// cloneDNSCtx is a clone of the DNSContext that was used to create the
 	// pendingRequest and store its result.  It must only be accessed for
 	// reading after the finish channel is closed.
-	cloneCtx *DNSContext
+	cloneDNSCtx *DNSContext
 }
 
 // NewDefaultPendingRequests creates a new instance of DefaultPendingRequests.
@@ -77,23 +77,25 @@ func (pr *DefaultPendingRequests) queue(
 	}
 
 	pendingVal, exists := pr.storage.LoadOrStore(string(key), req)
-	if exists {
-		pending := pendingVal.(*pendingRequest)
-		<-pending.finish
+	if !exists {
+		return false, nil
+	}
 
-		if pending.cloneCtx != nil {
-			// TODO(e.burkov):  Add cloner for DNS messages.
-			dctx.Res = pending.cloneCtx.Res.Copy().SetReply(dctx.Req)
-			dctx.Upstream = pending.cloneCtx.Upstream
+	pending := pendingVal.(*pendingRequest)
+	<-pending.finish
 
-			// TODO(a.garipov):  !! Decide how to treat query durations.
-			dctx.queryStatistics = pending.cloneCtx.queryStatistics
-		}
+	origDNSCtx := pending.cloneDNSCtx
 
-		return exists, pending.resolveErr
+	// TODO(a.garipov):  Perhaps, statistics should be calculated separately for
+	// each request.
+	dctx.queryStatistics = origDNSCtx.queryStatistics
+	dctx.Upstream = origDNSCtx.Upstream
+	if origDNSCtx.Res != nil {
+		// TODO(e.burkov):  Add cloner for DNS messages.
+		dctx.Res = origDNSCtx.Res.Copy().SetReply(dctx.Req)
 	}
 
-	return false, nil
+	return exists, pending.resolveErr
 }
 
 // done implements the [PendingRequests] interface for [DefaultPendingRequests].
@@ -114,14 +116,16 @@ func (pr *DefaultPendingRequests) done(ctx context.Context, dctx *DNSContext, er
 	pending := pendingVal.(*pendingRequest)
 	pending.resolveErr = err
 
-	cloneCtx := &DNSContext{}
+	cloneCtx := &DNSContext{
+		Upstream:        dctx.Upstream,
+		queryStatistics: dctx.queryStatistics,
+	}
+
 	if dctx.Res != nil {
 		cloneCtx.Res = dctx.Res.Copy()
-		cloneCtx.Upstream = dctx.Upstream
-		cloneCtx.queryStatistics = dctx.queryStatistics
 	}
 
-	pending.cloneCtx = cloneCtx
+	pending.cloneDNSCtx = cloneCtx
 
 	pr.storage.Delete(string(key))
 	close(pending.finish)

proxy/pending_internal_test.go proxy/pending_test.goproxy/pending_internal_test.go renamed to proxy/pending_test.go

@@ -1,45 +1,45 @@
-package proxy
+package proxy_test
 
 import (
-	"context"
 	"net"
+	"net/netip"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
+	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
+	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/miekg/dns"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
-// testPendingRequests is a mock implementation of [PendingRequests] for tests.
-//
-// TODO(e.burkov):  Think of a better way to test [PendingRequests].
-type testPendingRequests struct {
-	onQueue func(ctx context.Context, dctx *DNSContext) (exists bool, err error)
-	onDone  func(ctx context.Context, dctx *DNSContext, err error)
-}
+// TODO(e.burkov):  Merge those with the ones in internal tests and move to
+// dnsproxytest.
 
-// type check
-var _ PendingRequests = (*testPendingRequests)(nil)
+const (
+	// testTimeout is the common timeout for tests and contexts.
+	testTimeout = 1 * time.Second
 
-// queue implements the [proxy.PendingRequests] interface for
-// *testPendingRequests.
-func (p *testPendingRequests) queue(
-	ctx context.Context,
-	dctx *DNSContext,
-) (exists bool, err error) {
-	return p.onQueue(ctx, dctx)
-}
+	// testCacheSize is the default size of the cache in bytes.
+	testCacheSize = 64 * 1024
+)
 
-// done implements the [proxy.PendingRequests] interface for
-// *testPendingRequests.
-func (p *testPendingRequests) done(ctx context.Context, dctx *DNSContext, err error) {
-	p.onDone(ctx, dctx, err)
-}
+var (
+	// localhostAnyPort is a localhost address with an arbitrary port.
+	localhostAnyPort = netip.AddrPortFrom(netutil.IPv4Localhost(), 0)
+
+	// testTrustedProxies is a set of trusted proxies that includes all
+	// addresses used in tests.
+	testTrustedProxies = netutil.SliceSubnetSet{
+		netip.MustParsePrefix("0.0.0.0/0"),
+		netip.MustParsePrefix("::0/0"),
+	}
+)
 
 // assertEqualResponses is a helper function that checks if two DNS messages are
 // equal, excluding their ID.
@@ -66,6 +66,7 @@ func assertEqualResponses(tb testing.TB, expected, actual *dns.Msg) {
 	assert.Equal(tb, expected.Extra, actual.Extra)
 }
 
+// TODO(e.burkov):  Consider unexporting the [proxy.PendingRequests] interface.
 func TestPendingRequests(t *testing.T) {
 	t.Parallel()
 
@@ -77,57 +78,54 @@ func TestPendingRequests(t *testing.T) {
 	once := &sync.Once{}
 	u := &dnsproxytest.FakeUpstream{
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
-			loadWG.Wait()
 			once.Do(func() {
 				resp = (&dns.Msg{}).SetReply(req)
 			})
 
 			// Only allow a single request to be processed.
 			require.NotNil(testutil.PanicT{}, resp)
 
+			loadWG.Wait()
+
 			return resp, nil
 		},
 		OnAddress: func() (addr string) { return "" },
 		OnClose:   func() (err error) { return nil },
 	}
 
-	pending := NewDefaultPendingRequests()
-	testPending := &testPendingRequests{
-		onQueue: func(ctx context.Context, dctx *DNSContext) (exists bool, err error) {
-			loadWG.Done()
-
-			return pending.queue(ctx, dctx)
-		},
-		onDone: pending.done,
-	}
-
-	p := mustNew(t, &Config{
+	p, err := proxy.New(&proxy.Config{
 		Logger:                 slogutil.NewDiscardLogger(),
 		UDPListenAddr:          []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)},
 		TCPListenAddr:          []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)},
-		UpstreamConfig:         &UpstreamConfig{Upstreams: []upstream.Upstream{u}},
-		TrustedProxies:         defaultTrustedProxies,
+		UpstreamConfig:         &proxy.UpstreamConfig{Upstreams: []upstream.Upstream{u}},
+		TrustedProxies:         testTrustedProxies,
 		RatelimitSubnetLenIPv4: 24,
 		RatelimitSubnetLenIPv6: 64,
 		Ratelimit:              0,
 		CacheEnabled:           true,
-		CacheSizeBytes:         defaultCacheSize,
+		CacheSizeBytes:         testCacheSize,
 		EnableEDNSClientSubnet: true,
-		PendingRequests:        testPending,
+		PendingRequests:        proxy.NewDefaultPendingRequests(),
+		RequestHandler: func(prx *proxy.Proxy, dctx *proxy.DNSContext) (err error) {
+			loadWG.Done()
+
+			return prx.Resolve(dctx)
+		},
 	})
+	require.NoError(t, err)
 
 	ctx := testutil.ContextWithTimeout(t, testTimeout)
-	err := p.Start(ctx)
+	err = p.Start(ctx)
 	require.NoError(t, err)
 	testutil.CleanupAndRequireSuccess(t, func() (err error) {
 		ctx = testutil.ContextWithTimeout(t, testTimeout)
 
 		return p.Shutdown(ctx)
 	})
 
-	addr := p.Addr(ProtoTCP).String()
+	addr := p.Addr(proxy.ProtoTCP).String()
 	client := &dns.Client{
-		Net:     string(ProtoTCP),
+		Net:     string(proxy.ProtoTCP),
 		Timeout: testTimeout,
 	}
 
@@ -138,11 +136,11 @@ func TestPendingRequests(t *testing.T) {
 	for i := range reqsNum {
 		resolveWG.Add(1)
 
+		req := (&dns.Msg{}).SetQuestion("domain.example.", dns.TypeA)
+
 		go func() {
 			defer resolveWG.Done()
 
-			req := (&dns.Msg{}).SetQuestion("domain.example.", dns.TypeA)
-
 			reqCtx := testutil.ContextWithTimeout(t, testTimeout)
 			responses[i], _, errs[i] = client.ExchangeContext(reqCtx, req, addr)
 		}()

proxy/proxy.go

@@ -693,9 +693,6 @@ func (p *Proxy) Resolve(dctx *DNSContext) (err error) {
 
 	dctx.calcFlagsAndSize()
 
-	// Also don't lookup the cache for responses with DNSSEC checking disabled
-	// since only validated responses are cached and those may be not the
-	// desired result for user specifying CD flag.
 	cacheWorks := p.cacheWorks(dctx)
 	if cacheWorks {
 		// Only add pending requests if the cache is enabled, since this is a
@@ -771,6 +768,9 @@ func (p *Proxy) cacheWorks(dctx *DNSContext) (ok bool) {
 		// TODO(e.burkov):  It probably should be decided after resolve.
 		reason = "custom upstreams cache is not configured"
 	case dctx.Req.CheckingDisabled:
+		// Also don't lookup the cache for responses with DNSSEC checking
+		// disabled since only validated responses are cached and those may be
+		// not the desired result for user specifying CD flag.
 		reason = "dnssec check disabled"
 	default:
 		return true

proxy/stats.go

@@ -74,12 +74,14 @@ func cachedQueryStatistics(addr string) (s *QueryStatistics) {
 	}
 }
 
-// Main returns the DNS query statistics for the upstream DNS servers.
+// Main returns the DNS query statistics for the upstream DNS servers.  us and
+// its items must not be modified.
 func (s *QueryStatistics) Main() (us []*UpstreamStatistics) {
 	return s.main
 }
 
-// Fallback returns the DNS query statistics for the fallback DNS servers.
+// Fallback returns the DNS query statistics for the fallback DNS servers. us
+// and its items must not be modified.
 func (s *QueryStatistics) Fallback() (us []*UpstreamStatistics) {
 	return s.fallback
 }

proxy/stats_test.go

@@ -17,10 +17,6 @@ import (
 )
 
 func TestCollectQueryStats(t *testing.T) {
-	const (
-		listenIP = "127.0.0.1"
-	)
-
 	var (
 		testReq = &dns.Msg{
 			Question: []dns.Question{{
@@ -34,8 +30,6 @@ func TestCollectQueryStats(t *testing.T) {
 			netip.MustParsePrefix("0.0.0.0/0"),
 			netip.MustParsePrefix("::0/0"),
 		}
-
-		localhostAnyPort = netip.MustParseAddrPort(netutil.JoinHostPort(listenIP, 0))
 	)
 
 	ups := &dnsproxytest.FakeUpstream{