Pull request: AGDNS-3523-add-context

Merge in GO/dnsproxy from AGDNS-3523-add-context to master

Squashed commit of the following:

commit 59b61ec
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu Mar 12 09:20:57 2026 +0700

    proxy: start ctx

commit fe4cbfb
Merge: 10c620a b7d4b05
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Wed Mar 11 07:59:19 2026 +0700

    Merge remote-tracking branch 'origin/master' into AGDNS-3523-add-context

commit 10c620a
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue Mar 10 12:21:06 2026 +0700

    proxy: fix udp

commit 6dbe120
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Fri Mar 6 12:58:47 2026 +0700

    proxy: start ctx

commit 88cf3e7
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Thu Mar 5 11:24:26 2026 +0700

    proxy: req ctx constructor

commit f8d8e26
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue Mar 3 08:53:28 2026 +0700

    proxy: todo

commit 6e32f1e
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Mon Mar 2 09:27:03 2026 +0700

    proxy: add context

commit 9e42af8
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Sun Mar 1 09:42:30 2026 +0700

    proxy: imp code

commit 29e498d
Author: Dimitry Kolyshev <dkolyshev@adguard.com>
Date:   Tue Feb 24 13:32:27 2026 +0700

    all: add proxy context

Author: Mizzick

internal/middleware/middleware.go

@@ -60,9 +60,7 @@ var _ proxy.Middleware = (*Default)(nil)
 // and resolves the DNS request within proxyCtx.  It only calls h if the request
 // isn't handled by any of the internal handlers.
 func (mw *Default) Wrap(h proxy.Handler) (wrapped proxy.Handler) {
-	f := func(p *proxy.Proxy, proxyCtx *proxy.DNSContext) (err error) {
-		ctx := context.TODO()
-
+	f := func(ctx context.Context, p *proxy.Proxy, proxyCtx *proxy.DNSContext) (err error) {
 		mw.logger.DebugContext(ctx, "handling request", "req", &proxyCtx.Req.Question[0])
 
 		if proxyCtx.Res = mw.haltAAAA(ctx, proxyCtx.Req); proxyCtx.Res != nil {
@@ -73,7 +71,7 @@ func (mw *Default) Wrap(h proxy.Handler) (wrapped proxy.Handler) {
 			return nil
 		}
 
-		return h.ServeDNS(p, proxyCtx)
+		return h.ServeDNS(ctx, p, proxyCtx)
 	}
 
 	return proxy.HandlerFunc(f)

proxy/bogusnxdomain_internal_test.go

@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/AdguardTeam/dnsproxy/upstream"
+	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/AdguardTeam/golibs/testutil/servicetest"
 	"github.com/miekg/dns"
 	"github.com/stretchr/testify/assert"
@@ -89,7 +90,7 @@ func TestProxy_IsBogusNXDomain(t *testing.T) {
 		u.ans = tc.ans
 
 		t.Run(tc.name, func(t *testing.T) {
-			err := prx.Resolve(d)
+			err := prx.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d)
 			require.NoError(t, err)
 			require.NotNil(t, d.Res)
 

proxy/cache_internal_test.go

@@ -395,7 +395,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 			A: net.IP{4, 3, 2, 1},
 		}}
 
-		err := dnsProxy.Resolve(d)
+		err := dnsProxy.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d)
 		require.NoError(t, err)
 
 		ci, expired, key := dnsProxy.cache.get(d.Req)
@@ -419,8 +419,8 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) {
 			A: net.IP{4, 3, 2, 1},
 		}}
 
-		err := dnsProxy.Resolve(d)
-		assert.Nil(t, err)
+		err := dnsProxy.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d)
+		require.NoError(t, err)
 
 		ci, expired, key := dnsProxy.cache.get(d.Req)
 		assert.False(t, expired)

proxy/config.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/AdguardTeam/dnsproxy/upstream"
+	"github.com/AdguardTeam/golibs/contextutil"
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/netutil"
 	"github.com/ameshkov/dnscrypt/v2"
@@ -55,6 +56,10 @@ type Config struct {
 	// performing a single lookup.  If nil, it will be enabled by default.
 	PendingRequests *PendingRequestsConfig
 
+	// RequestContext is a context constructor that returns contexts for
+	// requests.  If not set, the proxy uses [contextutil.EmptyConstructor].
+	RequestContext contextutil.Constructor
+
 	// RequestHandler is an optional custom handler for DNS requests.  It's used
 	// instead of DefaultHandler if set.  In case of [ErrDrop] error returned
 	// from this handler, the proxy will not send any response to the client.

proxy/dns64_internal_test.go

@@ -375,7 +375,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
 				Addr: localCliAddr,
 			}
 
-			err = p.handleDNSRequest(dctx)
+			err = p.handleDNSRequest(testutil.ContextWithTimeout(t, defaultTimeout), dctx)
 			require.NoError(t, err)
 
 			res := dctx.Res

proxy/exchange_internal_test.go

@@ -223,7 +223,9 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
 
 		t.Run(tc.name, func(t *testing.T) {
 			for range requestsNum {
-				_ = p.Resolve(&DNSContext{Req: req, Addr: cli})
+				ctx := testutil.ContextWithTimeout(t, defaultTimeout)
+
+				_ = p.Resolve(ctx, &DNSContext{Req: req, Addr: cli})
 			}
 
 			assert.Equal(t, wantStat, stats)

proxy/handler_internal_test.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"net"
 	"sync"
 	"testing"
@@ -17,13 +18,13 @@ func TestFilteringHandler(t *testing.T) {
 	blockResponse := false
 
 	reqHandler := &TestHandler{
-		OnHandle: func(p *Proxy, d *DNSContext) (err error) {
+		OnHandle: func(ctx context.Context, p *Proxy, d *DNSContext) (err error) {
 			m.Lock()
 			defer m.Unlock()
 
 			if !blockResponse {
 				// Use the default Resolve method if response is not blocked.
-				return p.Resolve(d)
+				return p.Resolve(ctx, d)
 			}
 
 			resp := dns.Msg{}

proxy/lookup.go

@@ -33,7 +33,7 @@ func (p *Proxy) lookupIPAddr(
 
 	// TODO(d.kolyshev): Investigate why the client address is not defined.
 	d := p.newDNSContext(ProtoUDP, req, netip.AddrPort{})
-	err := p.Resolve(d)
+	err := p.Resolve(ctx, d)
 	ch <- &lookupResult{
 		resp: d.Res,
 		err:  err,

proxy/pending_test.go

@@ -1,6 +1,7 @@
 package proxy_test
 
 import (
+	"context"
 	"net"
 	"net/netip"
 	"sync"
@@ -79,10 +80,10 @@ func TestPendingRequests(t *testing.T) {
 	workloadWG.Add(reqsNum)
 
 	reqHandler := &proxy.TestHandler{
-		OnHandle: func(p *proxy.Proxy, d *proxy.DNSContext) (err error) {
+		OnHandle: func(ctx context.Context, p *proxy.Proxy, d *proxy.DNSContext) (err error) {
 			workloadWG.Done()
 
-			return p.Resolve(d)
+			return p.Resolve(ctx, d)
 		},
 	}
 

proxy/proxy.go

@@ -20,6 +20,7 @@ import (
 	"github.com/AdguardTeam/dnsproxy/internal/dnsmsg"
 	proxynetutil "github.com/AdguardTeam/dnsproxy/internal/netutil"
 	"github.com/AdguardTeam/dnsproxy/upstream"
+	"github.com/AdguardTeam/golibs/contextutil"
 	"github.com/AdguardTeam/golibs/errors"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
 	"github.com/AdguardTeam/golibs/netutil"
@@ -112,6 +113,9 @@ type Proxy struct {
 	// requests for private addresses.
 	recDetector *recursionDetector
 
+	// reqCtx is a constructor for the request contexts.  It is never nil.
+	reqCtx contextutil.Constructor
+
 	// pendingRequests is a storage for duplicated requests.  It is used to
 	// prevent sending the same request to upstreams multiple times.
 	pendingRequests pendingRequests
@@ -215,6 +219,10 @@ func New(c *Config) (p *Proxy, err error) {
 			c.PrivateSubnets,
 			netutil.SubnetSetFunc(netutil.IsLocallyServed),
 		),
+		reqCtx: cmp.Or[contextutil.Constructor](
+			c.RequestContext,
+			contextutil.EmptyConstructor{},
+		),
 		requestHandler:   cmp.Or[Handler](c.RequestHandler, DefaultHandler{}),
 		upstreamRTTStats: map[string]upstreamRTTStats{},
 		rttLock:          sync.Mutex{},
@@ -338,7 +346,9 @@ func (p *Proxy) Start(ctx context.Context) (err error) {
 		return fmt.Errorf("configuring listeners: %w", errors.WithDeferred(err, closeErr))
 	}
 
-	p.serveListeners()
+	// Use context without cancel to prevent listeners' context from being
+	// canceled.
+	p.serveListeners(context.WithoutCancel(ctx))
 
 	p.started = true
 
@@ -675,12 +685,8 @@ func addDO(msg *dns.Msg) {
 const defaultUDPBufSize = 2048
 
 // Resolve is the default resolving method used by the DNS proxy to query
-// upstream servers.  It expects dctx is filled with the request, the client's
-//
-// TODO(e.burkov):  Add [context.Context].
-func (p *Proxy) Resolve(dctx *DNSContext) (err error) {
-	ctx := context.Background()
-
+// upstream servers.  It expects dctx is filled with the client's request.
+func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
 	if p.EnableEDNSClientSubnet {
 		dctx.processECS(p.EDNSAddr, p.logger)
 	}