AdguardTeam/dnsproxy
1
0
Fork 0
mirror of https://github.com/AdguardTeam/dnsproxy.git synced 2026-10-12 02:01:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

AGH-15 exchange-ctx (18)

Browse Source
This commit is contained in:
Fedor Setrakov
2026-09-30 15:06:35 +03:00
committed by GitHub
parent caf3e7edcc
commit 8daf162e78
36 changed files with 551 additions and 309 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dnsproxytest/interface.go
+3 -3
View File
@@ -14,7 +14,7 @@ import (
// Upstream is a mock [upstream.Upstream] implementation for tests.
type Upstream struct {
OnAddress func() (addr string)
OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
OnExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
OnClose func() (err error)
}
@@ -27,8 +27,8 @@ func (u *Upstream) Address() (addr string) {
}
// Exchange implements the [upstream.Upstream] interface for *Upstream.
func (u *Upstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
return u.OnExchange(req)
func (u *Upstream) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return u.OnExchange(ctx, req)
}
// Close implements the [upstream.Upstream] interface for *Upstream.
fastip/fastest.go
+5 -3
View File
@@ -4,6 +4,7 @@
package fastip
import (
"context"
"log/slog"
"net"
"net/netip"
@@ -92,10 +93,11 @@ func New(c *Config) (f *FastestAddr) {
// the fastest IP address. The fastest IP address is considered to be the first
// one successfully dialed and other addresses are removed from the answer.
func (f *FastestAddr) ExchangeFastest(
ctx context.Context,
req *dns.Msg,
ups []upstream.Upstream,
) (resp *dns.Msg, u upstream.Upstream, err error) {
replies, err := upstream.ExchangeAll(ups, req)
replies, err := upstream.ExchangeAll(ctx, ups, req)
if err != nil {
return nil, nil, err
}
@@ -112,11 +114,11 @@ func (f *FastestAddr) ExchangeFastest(
ips := ipSet.Values()
host := strings.ToLower(req.Question[0].Name)
if pingRes := f.pingAll(host, ips); pingRes != nil {
if pingRes := f.pingAll(ctx, host, ips); pingRes != nil {
return f.prepareReply(pingRes, replies)
}
f.logger.Debug("no fastest ip found, using the first response", "host", host)
f.logger.DebugContext(ctx, "no fastest ip found, using the first response", "host", host)
return replies[0].Resp, replies[0].Upstream, nil
}
fastip/fastest_test.go
+14 -7
View File
@@ -1,6 +1,7 @@
package fastip_test
import (
"context"
"net"
"net/netip"
"strings"
@@ -8,6 +9,7 @@ import (
"github.com/AdguardTeam/dnsproxy/dnsproxytest"
"github.com/AdguardTeam/dnsproxy/fastip"
proxytest "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
"github.com/AdguardTeam/dnsproxy/internal/nettest"
"github.com/AdguardTeam/dnsproxy/upstream"
"github.com/AdguardTeam/golibs/logutil/slogutil"
@@ -30,16 +32,19 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
t.Parallel()
u := &dnsproxytest.Upstream{
OnAddress: func() (addr string) { return "bad_upstream" },
OnExchange: func(_ *dns.Msg) (resp *dns.Msg, err error) { return nil, assert.AnError },
OnClose: func() (err error) { return nil },
OnAddress: func() (addr string) { return "bad_upstream" },
OnExchange: func(_ context.Context, _ *dns.Msg) (resp *dns.Msg, err error) {
return nil, assert.AnError
},
OnClose: func() (err error) { return nil },
}
f := fastip.New(&fastip.Config{
Logger: l,
PingWaitTimeout: fastip.DefaultPingWaitTimeout,
})
resp, up, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{u})
ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
resp, up, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{u})
require.Error(t, err)
assert.ErrorIs(t, err, assert.AnError)
@@ -67,7 +72,8 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
alive := newTestAUpstream(t, []*dns.A{newTestRec(t, aliveAddr)})
dead := newTestAUpstream(t, []*dns.A{newTestRec(t, netip.MustParseAddr("192.0.2.1"))})
rep, ups, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{dead, alive})
ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
rep, ups, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{dead, alive})
require.NoError(t, err)
assert.Equal(t, ups, alive)
@@ -96,7 +102,8 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
newTestRec(t, netip.MustParseAddr("127.0.0.3")),
})
resp, _, err := f.ExchangeFastest(newTestReq(t), []upstream.Upstream{ups})
ctx := testutil.ContextWithTimeout(t, proxytest.Timeout)
resp, _, err := f.ExchangeFastest(ctx, newTestReq(t), []upstream.Upstream{ups})
require.NoError(t, err)
require.NotNil(t, resp)
@@ -112,7 +119,7 @@ func TestFastestAddr_ExchangeFastest(t *testing.T) {
func newTestAUpstream(tb testing.TB, recs []*dns.A) (ups *dnsproxytest.Upstream) {
tb.Helper()
onExchange := func(m *dns.Msg) (resp *dns.Msg, err error) {
onExchange := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = &dns.Msg{}
resp.SetReply(m)
fastip/ping.go
+34 -14
View File
@@ -1,6 +1,7 @@
package fastip
import (
"context"
"net/netip"
"time"
@@ -29,6 +30,7 @@ type pingResult struct {
// Returns scheduled flag which indicates that some goroutines have been
// scheduled.
func (f *FastestAddr) schedulePings(
ctx context.Context,
resCh chan *pingResult,
ips []netip.Addr,
host string,
@@ -38,7 +40,7 @@ func (f *FastestAddr) schedulePings(
if cached == nil {
scheduled = true
for _, port := range f.pingPorts {
go f.pingDoTCP(host, netip.AddrPortFrom(ip, uint16(port)), resCh)
go f.pingDoTCP(ctx, host, netip.AddrPortFrom(ip, uint16(port)), resCh)
}
continue
@@ -58,7 +60,7 @@ func (f *FastestAddr) schedulePings(
// pingAll pings all ips concurrently and returns as soon as the fastest one is
// found or the timeout is exceeded.
func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
func (f *FastestAddr) pingAll(ctx context.Context, host string, ips []netip.Addr) (pr *pingResult) {
ipN := len(ips)
switch ipN {
case 0:
@@ -71,22 +73,23 @@ func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
}
resCh := make(chan *pingResult, ipN*len(f.pingPorts))
pr, scheduled := f.schedulePings(resCh, ips, host)
pr, scheduled := f.schedulePings(ctx, resCh, ips, host)
if !scheduled {
if pr != nil {
f.logger.Debug(
f.logger.DebugContext(
ctx,
"pinging all returns cached response",
"host", host,
"addr", pr.addrPort,
)
} else {
f.logger.Debug("pinging all returns nothing", "host", host)
f.logger.DebugContext(ctx, "pinging all returns nothing", "host", host)
}
return pr
}
res := f.firstSuccessRes(resCh, host)
res := f.firstSuccessRes(ctx, resCh, host)
if res == nil {
// In case of timeout return cached or nil.
return pr
@@ -103,12 +106,17 @@ func (f *FastestAddr) pingAll(host string, ips []netip.Addr) (pr *pingResult) {
// firstSuccessRes waits and returns the first successful ping result or nil in
// case of timeout.
func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res *pingResult) {
func (f *FastestAddr) firstSuccessRes(
ctx context.Context,
resCh chan *pingResult,
host string,
) (res *pingResult) {
after := time.After(f.pingWaitTimeout)
for {
select {
case res = <-resCh:
f.logger.Debug(
f.logger.DebugContext(
ctx,
"pinging all got result",
"host", host,
"addr", res.addrPort,
@@ -121,7 +129,7 @@ func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res
return res
case <-after:
f.logger.Debug("pinging all timed out", "host", host)
f.logger.DebugContext(ctx, "pinging all timed out", "host", host)
return nil
}
@@ -129,18 +137,25 @@ func (f *FastestAddr) firstSuccessRes(resCh chan *pingResult, host string) (res
}
// pingDoTCP sends the result of dialing the specified address into resCh.
func (f *FastestAddr) pingDoTCP(host string, addrPort netip.AddrPort, resCh chan *pingResult) {
func (f *FastestAddr) pingDoTCP(
ctx context.Context,
host string,
addrPort netip.AddrPort,
resCh chan *pingResult,
) {
l := f.logger.With("host", host, "addr", addrPort)
l.Debug("open tcp connection")
l.DebugContext(ctx, "open tcp connection")
start := time.Now()
// TODO(f.setrakov): Respect context timeout.
conn, err := f.pinger.Dial(bootstrap.NetworkTCP, addrPort.String())
elapsed := time.Since(start)
success := err == nil
if success {
if cErr := conn.Close(); cErr != nil {
l.Debug("closing tcp connection", slogutil.KeyError, cErr)
l.DebugContext(ctx, "closing tcp connection", slogutil.KeyError, cErr)
}
}
@@ -154,10 +169,15 @@ func (f *FastestAddr) pingDoTCP(host string, addrPort netip.AddrPort, resCh chan
addr := addrPort.Addr().Unmap()
if success {
l.Debug("tcp ping success", "elapsed", elapsed)
l.DebugContext(ctx, "tcp ping success", "elapsed", elapsed)
f.cacheAddSuccessful(addr, latency)
} else {
l.Debug("tcp ping failed to connect", "elapsed", elapsed, slogutil.KeyError, err)
l.DebugContext(
ctx,
"tcp ping failed to connect",
"elapsed", elapsed,
slogutil.KeyError, err,
)
f.cacheAddFailure(addr)
}
}
fastip/ping_internal_test.go
+12 -5
View File
@@ -7,6 +7,7 @@ import (
"syscall"
"testing"
"github.com/AdguardTeam/dnsproxy/internal/dnsproxytest"
"github.com/AdguardTeam/golibs/logutil/slogutil"
"github.com/AdguardTeam/golibs/netutil"
"github.com/AdguardTeam/golibs/testutil"
@@ -30,7 +31,9 @@ func TestFastestAddr_PingAll_timeout(t *testing.T) {
}
ip := netutil.IPv4Localhost()
res := f.pingAll("", []netip.Addr{ip, ip})
ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
require.Nil(t, res)
waitCh <- unit{}
@@ -52,7 +55,8 @@ func TestFastestAddr_PingAll_timeout(t *testing.T) {
return nil
}
res := f.pingAll("", []netip.Addr{ip1, ip2})
ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
res := f.pingAll(ctx, "", []netip.Addr{ip1, ip2})
require.NotNil(t, res)
assert.True(t, res.success)
@@ -83,7 +87,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
f := New(&Config{Logger: slogutil.NewDiscardLogger()})
f.cacheAddFailure(ip)
res := f.pingAll("", []netip.Addr{ip, ip})
ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
require.Nil(t, res)
})
@@ -93,7 +98,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
f := New(&Config{Logger: slogutil.NewDiscardLogger()})
f.cacheAddSuccessful(ip, lat)
res := f.pingAll("", []netip.Addr{ip, ip})
ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
res := f.pingAll(ctx, "", []netip.Addr{ip, ip})
require.NotNil(t, res)
assert.True(t, res.success)
assert.Equal(t, lat, res.latency)
@@ -125,7 +131,8 @@ func TestFastestAddr_PingAll_cache(t *testing.T) {
return nil
}
res := f.pingAll("", ips)
ctx := testutil.ContextWithTimeout(t, dnsproxytest.Timeout)
res := f.pingAll(ctx, "", ips)
require.NotNil(t, res)
assert.True(t, res.success)
proxy/cache_internal_test.go
+6 -3
View File
@@ -2,6 +2,7 @@ package proxy
import (
"cmp"
"context"
"net"
"net/netip"
"strings"
@@ -26,9 +27,11 @@ const testUpsAddr = "https://upstream.address"
// upstreamWithAddr is a [dnsproxytest.Upstream] that is only expected to be
// used to get its address.
var upstreamWithAddr = &testUpstream{
OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnAddress: func() (addr string) { return testUpsAddr },
OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
panic(testutil.UnexpectedCall(ctx, m))
},
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnAddress: func() (addr string) { return testUpsAddr },
}
// newTestCache is a helper that returns new cache and fills its config with
proxy/dns64.go
+6 -4
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"fmt"
"net"
"net/netip"
@@ -278,6 +279,7 @@ func (p *Proxy) synthRR(rr dns.RR, soaTTL uint32) (result dns.RR) {
// nil, if the request was not performed. If origResp is not nil, origReq must
// also not be nil.
func (p *Proxy) performDNS64(
ctx context.Context,
origReq *dns.Msg,
origResp *dns.Msg,
upstreams []upstream.Upstream,
@@ -292,17 +294,17 @@ func (p *Proxy) performDNS64(
}
host := origReq.Question[0].Name
p.logger.Debug("received an empty aaaa response, checking dns64", "host", host)
p.logger.DebugContext(ctx, "received an empty aaaa response, checking dns64", "host", host)
dns64Resp, u, err := p.exchangeUpstreams(dns64Req, upstreams)
dns64Resp, u, err := p.exchangeUpstreams(ctx, dns64Req, upstreams)
if err != nil {
p.logger.Error("dns64 request failed", slogutil.KeyError, err)
p.logger.ErrorContext(ctx, "dns64 request failed", slogutil.KeyError, err)
return nil
}
if dns64Resp != nil && p.synthDNS64(origReq, origResp, dns64Resp) {
p.logger.Debug("synthesized aaaa response", "host", host)
p.logger.DebugContext(ctx, "synthesized aaaa response", "host", host)
return u
}
proxy/dns64_internal_test.go
+9 -6
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"net"
"net/netip"
"sync"
@@ -21,7 +22,7 @@ const ipv4OnlyFqdn = "ipv4.only."
func TestProxy_HandleDNSRequest_dns64Race(t *testing.T) {
ans := newRR(t, ipv4OnlyFqdn, dns.TypeA, 3600, dnsproxytest.IPv4)
ups := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(req)
if req.Question[0].Qtype == dns.TypeA {
resp.Answer = []dns.RR{dns.Copy(ans)}
@@ -33,9 +34,11 @@ func TestProxy_HandleDNSRequest_dns64Race(t *testing.T) {
OnClose: func() (err error) { return nil },
}
localUps := &testUpstream{
OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
OnAddress: func() (addr string) { return "fake.address" },
OnClose: func() (err error) { return nil },
OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
panic(testutil.UnexpectedCall(ctx, m))
},
OnAddress: func() (addr string) { return "fake.address" },
OnClose: func() (err error) { return nil },
}
dnsProxy := mustNew(t, &Config{
@@ -193,7 +196,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
pt := testutil.NewPanicT(t)
newUps := func(answers answerMap) (u upstream.Upstream) {
return &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
q := req.Question[0]
require.Contains(pt, answers, q.Qtype)
@@ -213,7 +216,7 @@ func TestProxy_Resolve_dns64(t *testing.T) {
localRR := newRR(t, ptr64Domain, dns.TypePTR, 3600, domainPointed)
localUps := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
require.Equal(pt, req.Question[0].Name, ptr64Domain)
resp = (&dns.Msg{}).SetReply(req)
resp.Answer = []dns.RR{localRR}
proxy/exchange.go
+12 -7
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"fmt"
"time"
@@ -15,16 +16,17 @@ import (
// response, the upstream that successfully resolved the request, and the error
// if any.
func (p *Proxy) exchangeUpstreams(
ctx context.Context,
req *dns.Msg,
ups []upstream.Upstream,
) (resp *dns.Msg, u upstream.Upstream, err error) {
switch p.upstreamMode {
case UpstreamModeParallel:
return upstream.ExchangeParallel(ups, req)
return upstream.ExchangeParallel(ctx, ups, req)
case UpstreamModeFastestAddr:
switch req.Question[0].Qtype {
case dns.TypeA, dns.TypeAAAA:
return p.fastestAddr.ExchangeFastest(req, ups)
return p.fastestAddr.ExchangeFastest(ctx, req, ups)
default:
// Go on to the load-balancing mode.
}
@@ -34,7 +36,7 @@ func (p *Proxy) exchangeUpstreams(
if len(ups) == 1 {
u = ups[0]
resp, _, err = p.exchange(u, req)
resp, _, err = p.exchange(ctx, u, req)
if err != nil {
return nil, nil, err
}
@@ -50,7 +52,7 @@ func (p *Proxy) exchangeUpstreams(
u = ups[i]
var elapsed time.Duration
resp, elapsed, err = p.exchange(u, req)
resp, elapsed, err = p.exchange(ctx, u, req)
if err == nil {
p.updateRTT(u.Address(), elapsed)
@@ -72,11 +74,12 @@ func (p *Proxy) exchangeUpstreams(
// exchange returns the result of the DNS request exchange with the given
// upstream and the elapsed time. u and req must not be nil.
func (p *Proxy) exchange(
ctx context.Context,
u upstream.Upstream,
req *dns.Msg,
) (resp *dns.Msg, dur time.Duration, err error) {
startTime := p.time.Now()
resp, err = u.Exchange(req)
resp, err = u.Exchange(ctx, req)
// Don't use [time.Since] because it uses [time.Now].
dur = p.time.Now().Sub(startTime)
@@ -84,7 +87,8 @@ func (p *Proxy) exchange(
addr := u.Address()
q := &req.Question[0]
if err != nil {
p.logger.Error(
p.logger.ErrorContext(
ctx,
"exchange failed",
"upstream", addr,
"question", q,
@@ -92,7 +96,8 @@ func (p *Proxy) exchange(
slogutil.KeyError, err,
)
} else {
p.logger.Debug(
p.logger.DebugContext(
ctx,
"exchange successfully finished",
"upstream", addr,
"question", q,
proxy/exchange_internal_test.go
+18 -13
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"math/rand/v2"
"net"
"net/netip"
@@ -25,7 +26,7 @@ func newUpstreamWithErrorRate(rate uint, name string) (u upstream.Upstream) {
var n uint
return &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
n++
if n%rate == 0 {
return nil, assert.AnError
@@ -52,10 +53,10 @@ type measuredUpstream struct {
var _ upstream.Upstream = measuredUpstream{}
// Exchange implements the [upstream.Upstream] interface for measuredUpstream.
func (u measuredUpstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (u measuredUpstream) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
u.stats[u.Address()]++
return u.Upstream.Exchange(req)
return u.Upstream.Exchange(ctx, req)
}
func TestProxy_Exchange_loadBalance(t *testing.T) {
@@ -88,7 +89,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
}
fastUps := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
currentNow = zeroTime.Add(testRTT / 100)
return (&dns.Msg{}).SetReply(req), nil
@@ -97,7 +98,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
}
slowerUps := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
currentNow = zeroTime.Add(testRTT / 10)
return (&dns.Msg{}).SetReply(req), nil
@@ -106,7 +107,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
}
slowestUps := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
currentNow = zeroTime.Add(testRTT / 2)
return (&dns.Msg{}).SetReply(req), nil
@@ -116,20 +117,24 @@ func TestProxy_Exchange_loadBalance(t *testing.T) {
}
err1Ups := &testUpstream{
OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
OnAddress: func() (addr string) { return "error1" },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnExchange: func(_ context.Context, _ *dns.Msg) (r *dns.Msg, err error) {
return nil, assert.AnError
},
OnAddress: func() (addr string) { return "error1" },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
}
err2Ups := &testUpstream{
OnExchange: func(_ *dns.Msg) (r *dns.Msg, err error) { return nil, assert.AnError },
OnAddress: func() (addr string) { return "error2" },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnExchange: func(_ context.Context, _ *dns.Msg) (r *dns.Msg, err error) {
return nil, assert.AnError
},
OnAddress: func() (addr string) { return "error2" },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
}
singleError := &sync.Once{}
// fastestUps responds with an error on the first request.
fastestUps := &testUpstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
singleError.Do(func() { err = assert.AnError })
currentNow = zeroTime.Add(testRTT / 200)
proxy/optimisticresolver.go
+12 -7
View File
@@ -15,7 +15,7 @@ type cachingResolver interface {
// resolved and the response may be cached.
//
// TODO(e.burkov): Find out when ok can be false with nil err.
replyFromUpstream(dctx *DNSContext) (ok bool, err error)
replyFromUpstream(ctx context.Context, dctx *DNSContext) (ok bool, err error)
// cacheResp caches the response from dctx.
cacheResp(dctx *DNSContext)
@@ -46,10 +46,13 @@ func newOptimisticResolver(cr cachingResolver) (s *optimisticResolver) {
// with the same key at the same period of time. It runs in a separate
// goroutine. Do not pass the *DNSContext which is used elsewhere since it
// isn't intended to be used concurrently.
//
// TODO(e.burkov): Pass the context.
func (s *optimisticResolver) resolveOnce(dctx *DNSContext, key []byte, l *slog.Logger) {
defer slogutil.RecoverAndLog(context.TODO(), l)
func (s *optimisticResolver) resolveOnce(
ctx context.Context,
dctx *DNSContext,
key []byte,
l *slog.Logger,
) {
defer slogutil.RecoverAndLog(ctx, l)
keyHexed := hex.EncodeToString(key)
if _, ok := s.reqs.LoadOrStore(keyHexed, unit{}); ok {
@@ -57,9 +60,11 @@ func (s *optimisticResolver) resolveOnce(dctx *DNSContext, key []byte, l *slog.L
}
defer s.reqs.Delete(keyHexed)
ok, err := s.cr.replyFromUpstream(dctx)
// NOTE: Discard the context deadline to prevent cancellation during
// refresh.
ok, err := s.cr.replyFromUpstream(context.WithoutCancel(ctx), dctx)
if err != nil {
l.Debug("resolving request for optimistic cache", slogutil.KeyError, err)
l.DebugContext(ctx, "resolving request for optimistic cache", slogutil.KeyError, err)
}
if ok {
proxy/optimisticresolver_internal_test.go
+27 -12
View File
@@ -2,25 +2,30 @@ package proxy
import (
"bytes"
"context"
"log/slog"
"sync"
"testing"
"github.com/AdguardTeam/golibs/errors"
"github.com/AdguardTeam/golibs/testutil"
"github.com/stretchr/testify/assert"
)
// testCachingResolver is a stub implementation of the cachingResolver interface
// to simplify testing.
type testCachingResolver struct {
onReplyFromUpstream func(dctx *DNSContext) (ok bool, err error)
onReplyFromUpstream func(ctx context.Context, dctx *DNSContext) (ok bool, err error)
onCacheResp func(dctx *DNSContext)
}
// replyFromUpstream implements the cachingResolver interface for
// *testCachingResolver.
func (tcr *testCachingResolver) replyFromUpstream(dctx *DNSContext) (ok bool, err error) {
return tcr.onReplyFromUpstream(dctx)
func (tcr *testCachingResolver) replyFromUpstream(
ctx context.Context,
dctx *DNSContext,
) (ok bool, err error) {
return tcr.onReplyFromUpstream(ctx, dctx)
}
// cacheResp implements the cachingResolver interface for *testCachingResolver.
@@ -33,7 +38,7 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
var timesResolved, timesSet int
tcr := &testCachingResolver{
onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) {
onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
timesResolved++
return true, nil
@@ -52,7 +57,8 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
sameKey := []byte{1, 2, 3}
// Start the primary goroutine.
go s.resolveOnce(nil, sameKey, testLogger)
ctx := testutil.ContextWithTimeout(t, testTimeout)
go s.resolveOnce(ctx, nil, sameKey, testLogger)
// Block until the primary goroutine reaches the resolve function.
<-out
@@ -61,7 +67,8 @@ func TestOptimisticResolver_ResolveOnce(t *testing.T) {
const secondaryNum = 10
for range secondaryNum {
wg.Go(func() {
s.resolveOnce(nil, sameKey, testLogger)
resolveCtx := testutil.ContextWithTimeout(t, testTimeout)
s.resolveOnce(resolveCtx, nil, sameKey, testLogger)
})
}
@@ -90,10 +97,14 @@ func TestOptimisticResolver_ResolveOnce_unsuccessful(t *testing.T) {
cached := false
s := newOptimisticResolver(&testCachingResolver{
onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) { return true, rErr },
onCacheResp: func(_ *DNSContext) { cached = true },
onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
return true, rErr
},
onCacheResp: func(_ *DNSContext) { cached = true },
})
s.resolveOnce(nil, key, l)
ctx := testutil.ContextWithTimeout(t, testTimeout)
s.resolveOnce(ctx, nil, key, l)
assert.True(t, cached)
assert.Contains(t, logOutput.String(), rErr.Error())
@@ -102,10 +113,14 @@ func TestOptimisticResolver_ResolveOnce_unsuccessful(t *testing.T) {
t.Run("not_ok", func(t *testing.T) {
cached := false
s := newOptimisticResolver(&testCachingResolver{
onReplyFromUpstream: func(_ *DNSContext) (ok bool, err error) { return false, nil },
onCacheResp: func(_ *DNSContext) { cached = true },
onReplyFromUpstream: func(_ context.Context, _ *DNSContext) (ok bool, err error) {
return false, nil
},
onCacheResp: func(_ *DNSContext) { cached = true },
})
s.resolveOnce(nil, key, testLogger)
ctx := testutil.ContextWithTimeout(t, testTimeout)
s.resolveOnce(ctx, nil, key, testLogger)
assert.False(t, cached)
})
proxy/pending_test.go
+1 -1
View File
@@ -64,7 +64,7 @@ func TestPendingRequests(t *testing.T) {
once := &sync.Once{}
u := &dnsproxytest.Upstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
once.Do(func() {
resp = (&dns.Msg{}).SetReply(req)
})
proxy/proxy.go
+14 -12
View File
@@ -791,7 +791,7 @@ func (p *Proxy) selectUpstreams(d *DNSContext) (upstreams []upstream.Upstream, i
// replyFromUpstream tries to resolve the request via configured upstream
// servers. It returns true if the response actually came from an upstream. d
// must not be nil.
func (p *Proxy) replyFromUpstream(d *DNSContext) (ok bool, err error) {
func (p *Proxy) replyFromUpstream(ctx context.Context, d *DNSContext) (ok bool, err error) {
req := d.Req
upstreams, isPrivate := p.selectUpstreams(d)
@@ -809,40 +809,42 @@ func (p *Proxy) replyFromUpstream(d *DNSContext) (ok bool, err error) {
wrapped := upstreamsWithStats(upstreams)
// Perform the DNS request.
resp, u, err := p.exchangeUpstreams(req, wrapped)
if dns64Ups := p.performDNS64(req, resp, wrapped); dns64Ups != nil {
resp, u, err := p.exchangeUpstreams(ctx, req, wrapped)
if dns64Ups := p.performDNS64(ctx, req, resp, wrapped); dns64Ups != nil {
u = dns64Ups
} else if p.isBogusNXDomain(resp) {
p.logger.Debug("response contains bogus-nxdomain ip")
p.logger.DebugContext(ctx, "response contains bogus-nxdomain ip")
resp = p.messages.NewMsgNXDOMAIN(req)
}
var wrappedFallbacks []upstream.Upstream
if err != nil && !isPrivate && p.fallbacks != nil {
p.logger.Debug("using fallback", slogutil.KeyError, err)
p.logger.DebugContext(ctx, "using fallback", slogutil.KeyError, err)
src = "fallback"
// upstreams mustn't appear empty since they have been validated when
// Upstreams mustn't appear empty since they have been validated when
// creating proxy.
upstreams = p.fallbacks.getUpstreamsForDomain(req.Question[0].Name)
wrappedFallbacks = upstreamsWithStats(upstreams)
resp, u, err = upstream.ExchangeParallel(wrappedFallbacks, req)
// NOTE: Use [context.WithoutCancel] to have a separate timeout on
// retry.
resp, u, err = upstream.ExchangeParallel(context.WithoutCancel(ctx), wrappedFallbacks, req)
}
if err != nil {
p.logger.Debug("resolving err", "src", src, slogutil.KeyError, err)
p.logger.DebugContext(ctx, "resolving err", "src", src, slogutil.KeyError, err)
}
if resp != nil {
p.logger.Debug("resolved", "upstream", u.Address(), "src", src)
p.logger.DebugContext(ctx, "resolved", "upstream", u.Address(), "src", src)
}
unwrapped, stats := collectQueryStats(p.upstreamMode, u, wrapped, wrappedFallbacks)
d.queryStatistics = stats
ctx := context.TODO()
p.handleExchangeResult(ctx, d, req, resp, unwrapped)
return resp != nil, err
@@ -927,7 +929,7 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
}
defer func() { p.pendingRequests.done(ctx, dctx, err) }()
if p.replyFromCache(dctx) {
if p.replyFromCache(ctx, dctx) {
// Complete the response from cache.
filterMsg(dctx.Res, dctx.Res, dctx.adBit, dctx.doBit, 0)
dctx.scrub()
@@ -937,7 +939,7 @@ func (p *Proxy) Resolve(ctx context.Context, dctx *DNSContext) (err error) {
}
var ok bool
ok, err = p.replyFromUpstream(dctx)
ok, err = p.replyFromUpstream(ctx, dctx)
// Don't cache the responses having CD flag, just like Dnsmasq does. It
// prevents the cache from being poisoned with unvalidated answers which may
proxy/proxy_internal_test.go
+35 -22
View File
@@ -82,7 +82,7 @@ func newTestUpstreamConfig(tb testing.TB, ups ...upstream.Upstream) (u *Upstream
func newTestUpstream(tb testing.TB) (uc upstream.Upstream) {
tb.Helper()
onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return dnsproxytest.NewTestResponse(req), nil
}
@@ -94,7 +94,7 @@ func newTestUpstream(tb testing.TB) (uc upstream.Upstream) {
// be replaced with stub implementations.
func newTestUpstreamWithExchange(
tb testing.TB,
onExchange func(req *dns.Msg) (resp *dns.Msg, err error),
onExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error),
) (uc upstream.Upstream) {
tb.Helper()
@@ -220,7 +220,7 @@ func TestProxy_Resolve_dnssecCache(t *testing.T) {
}
u := &testUpstream{
OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(m)
q := m.Question[0]
@@ -348,9 +348,12 @@ func TestProxy_Resolve_dnssecCache(t *testing.T) {
func TestProxy_HandleDNSRequest_exchangeWithReservedDomains(t *testing.T) {
t.Parallel()
emptyUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
return (&dns.Msg{}).SetReply(req), nil
})
emptyUpstream := newTestUpstreamWithExchange(
t,
func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return (&dns.Msg{}).SetReply(req), nil
},
)
const (
host1 = "example-1.test"
@@ -422,9 +425,12 @@ func TestProxy_HandleDNSRequest_exchangeWithReservedDomains(t *testing.T) {
func TestProxy_HandleDNSRequest_oneByOneUpstreamsExchange(t *testing.T) {
t.Parallel()
errUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
return nil, assert.AnError
})
errUpstream := newTestUpstreamWithExchange(
t,
func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return nil, assert.AnError
},
)
dnsProxy := mustNew(t, &Config{
Logger: testLogger,
@@ -477,18 +483,21 @@ func TestProxy_HandleDNSRequest_fallback(t *testing.T) {
pt := testutil.NewPanicT(t)
successExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
successExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
testutil.RequireSend(pt, responseCh, req.Id, dnsproxytest.Timeout)
return dnsproxytest.NewTestResponse(req), nil
}
successUpstream := newTestUpstreamWithExchange(t, successExchange)
failUpstream := newTestUpstreamWithExchange(t, func(req *dns.Msg) (resp *dns.Msg, err error) {
testutil.RequireSend(pt, failCh, req.Id, dnsproxytest.Timeout)
failUpstream := newTestUpstreamWithExchange(
t,
func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
testutil.RequireSend(pt, failCh, req.Id, dnsproxytest.Timeout)
return nil, assert.AnError
})
return nil, assert.AnError
},
)
upsConf := &UpstreamConfig{
DomainReservedUpstreams: map[string][]upstream.Upstream{
@@ -636,7 +645,7 @@ func TestProxy_Resolve_customUpstreamConfigCache(t *testing.T) {
var count int
exchangeFunc := func(m *dns.Msg) (resp *dns.Msg, err error) {
exchangeFunc := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = &dns.Msg{}
resp.SetReply(m)
resp.Answer = append(resp.Answer, &dns.A{
@@ -992,7 +1001,7 @@ func TestProxy_Resolve_withOptimisticResolver(t *testing.T) {
p.initCache()
out, in := make(chan unit), make(chan unit)
p.shortFlighter.cr = &testCachingResolver{
onReplyFromUpstream: func(dctx *DNSContext) (ok bool, err error) {
onReplyFromUpstream: func(_ context.Context, dctx *DNSContext) (ok bool, err error) {
dctx.Res = buildResp(dctx.Req, nonOptimisticTTL)
return true, nil
@@ -1077,7 +1086,7 @@ func TestProxy_ValidateRequest(t *testing.T) {
}
ups := &testUpstream{
OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = &dns.Msg{}
resp.SetReply(m)
@@ -1206,7 +1215,7 @@ func (h *testHandler) ServeDNS(ctx context.Context, p *Proxy, dctx *DNSContext)
// TODO(m.kazantsev): Use [dnsproxytest.Upstream].
type testUpstream struct {
OnAddress func() (addr string)
OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
OnExchange func(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
OnClose func() (err error)
}
@@ -1216,8 +1225,12 @@ type testUpstream struct {
// The handler dereferences ans, ecsIP, and ecsReqIP at call time, so tests
// can reassign those variables between sub-tests and the handler will see
// the current values.
func newECSReplyHandler(ans *[]dns.RR, ecsIP, ecsReqIP *net.IP) func(*dns.Msg) (*dns.Msg, error) {
return func(m *dns.Msg) (resp *dns.Msg, err error) {
func newECSReplyHandler(
ans *[]dns.RR,
ecsIP *net.IP,
ecsReqIP *net.IP,
) func(context.Context, *dns.Msg) (*dns.Msg, error) {
return func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(m)
if ans != nil && *ans != nil {
@@ -1245,8 +1258,8 @@ func newECSReplyHandler(ans *[]dns.RR, ecsIP, ecsReqIP *net.IP) func(*dns.Msg) (
var _ upstream.Upstream = (*testUpstream)(nil)
// Exchange implements the upstream.Upstream interface for *testUpstream.
func (u *testUpstream) Exchange(m *dns.Msg) (resp *dns.Msg, err error) {
return u.OnExchange(m)
func (u *testUpstream) Exchange(ctx context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
return u.OnExchange(ctx, m)
}
// Address implements the upstream.Upstream interface for *testUpstream.
proxy/proxy_test.go
+10 -7
View File
@@ -1,6 +1,7 @@
package proxy_test
import (
"context"
"net"
"net/netip"
"os"
@@ -44,7 +45,7 @@ var (
//
// TODO(f.setrakov): DRY with internal version.
func newTestUpstreamConfig(tb testing.TB) (uc *proxy.UpstreamConfig) {
onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
pt := testutil.NewPanicT(tb)
require.Len(pt, req.Question, 1)
@@ -103,7 +104,7 @@ func sendTestMessages(tb testing.TB, conn *dns.Conn) {
func TestProxy_Resolve_badResponse(t *testing.T) {
dnsProxy := mustStartDefaultProxy(t)
onExchange := func(m *dns.Msg) (resp *dns.Msg, err error) {
onExchange := func(_ context.Context, m *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(m)
resp.Answer = append(resp.Answer, &dns.A{
Hdr: dns.RR_Header{
@@ -257,7 +258,7 @@ func TestProxy_Resolve_cache(t *testing.T) {
OnAddress: func() (addr string) { return "stub" },
OnClose: func() (err error) { return nil },
}
ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
ups.OnExchange = func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(req)
resp.Answer = append(resp.Answer, &dns.A{
Hdr: dns.RR_Header{
@@ -349,9 +350,11 @@ func TestProxy_Start_closeOnFail(t *testing.T) {
tcpAddr := testutil.RequireTypeAssert[*net.TCPAddr](t, l.Addr())
ups := &dnsproxytest.Upstream{
OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
panic(testutil.UnexpectedCall(ctx, m))
},
OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
}
p, err := proxy.New(&proxy.Config{
@@ -389,7 +392,7 @@ func TestProxy_ServeDNS_formatError(t *testing.T) {
OnAddress: func() (addr string) { return testIPv4.String() },
OnClose: func() (err error) { return nil },
}
ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
ups.OnExchange = func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
panic(testutil.UnexpectedCall(req))
}
proxy/proxycache.go
+5 -3
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"net"
"slices"
)
@@ -17,7 +18,7 @@ func (p *Proxy) cacheForContext(d *DNSContext) (c *cache) {
// replyFromCache tries to get the response from general or subnet cache. In
// case the cache is present in d, it's used first. Returns true on success. d
// must not be nil.
func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
func (p *Proxy) replyFromCache(ctx context.Context, d *DNSContext) (hit bool) {
dctxCache := p.cacheForContext(d)
var ci *cacheItem
@@ -41,7 +42,8 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
d.Res = ci.m
d.queryStatistics = cachedQueryStatistics(ci.u)
p.logger.Debug(
p.logger.DebugContext(
ctx,
"replying from cache",
"source", cacheSource,
"ecs_enabled", p.enableEDNSClientSubnet,
@@ -59,7 +61,7 @@ func (p *Proxy) replyFromCache(d *DNSContext) (hit bool) {
minCtxClone.Req = d.Req.Copy()
}
go p.shortFlighter.resolveOnce(minCtxClone, key, p.logger)
go p.shortFlighter.resolveOnce(ctx, minCtxClone, key, p.logger)
}
return hit
proxy/servertcp_test.go
+6 -3
View File
@@ -1,6 +1,7 @@
package proxy_test
import (
"context"
"crypto/tls"
"crypto/x509"
"encoding/binary"
@@ -100,9 +101,11 @@ func TestProxy_handleDNSRequest_emptyTCPMessage(t *testing.T) {
t.Parallel()
u := &dnsproxytest.Upstream{
OnExchange: func(m *dns.Msg) (_ *dns.Msg, _ error) { panic(testutil.UnexpectedCall(m)) },
OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
OnClose: func() (err error) { return nil },
OnExchange: func(ctx context.Context, m *dns.Msg) (_ *dns.Msg, _ error) {
panic(testutil.UnexpectedCall(ctx, m))
},
OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
OnClose: func() (err error) { return nil },
}
upsConf := &proxy.UpstreamConfig{
Upstreams: []upstream.Upstream{u},
proxy/stats.go
+3 -2
View File
@@ -1,6 +1,7 @@
package proxy
import (
"context"
"fmt"
"time"
@@ -25,9 +26,9 @@ type upstreamWithStats struct {
var _ upstream.Upstream = (*upstreamWithStats)(nil)
// Exchange implements the [upstream.Upstream] for *upstreamWithStats.
func (u *upstreamWithStats) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (u *upstreamWithStats) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
start := time.Now()
resp, err = u.upstream.Exchange(req)
resp, err = u.upstream.Exchange(ctx, req)
u.err = err
u.queryDuration = time.Since(start)
proxy/stats_test.go
+3 -2
View File
@@ -1,6 +1,7 @@
package proxy_test
import (
"context"
"net"
"testing"
@@ -27,7 +28,7 @@ func TestDNSContext_QueryStatistics(t *testing.T) {
}
ups := &dnsproxytest.Upstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return (&dns.Msg{}).SetReply(req), nil
},
OnAddress: func() (addr string) { return "upstream" },
@@ -35,7 +36,7 @@ func TestDNSContext_QueryStatistics(t *testing.T) {
}
failUps := &dnsproxytest.Upstream{
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
return nil, errors.Error("exchange error")
},
OnAddress: func() (addr string) { return "fail.upstream" },
upstream/dnscrypt.go
+12 -3
View File
@@ -54,8 +54,7 @@ var _ Upstream = (*dnsCrypt)(nil)
func (p *dnsCrypt) Address() string { return p.addr.String() }
// Exchange implements the [Upstream] interface for *dnsCrypt.
func (p *dnsCrypt) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
ctx := context.Background()
func (p *dnsCrypt) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
if p.timeout > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, p.timeout)
@@ -103,7 +102,8 @@ func (p *dnsCrypt) exchangeDNSCrypt(ctx context.Context, req *dns.Msg) (resp *dn
resp, err = client.ExchangeContext(ctx, req, resolverInfo)
if resp != nil && resp.Truncated {
q := &req.Question[0]
p.logger.Debug(
p.logger.DebugContext(
ctx,
"dnscrypt received truncated, falling back to tcp",
"addr", p.addr,
"question", q,
@@ -114,6 +114,15 @@ func (p *dnsCrypt) exchangeDNSCrypt(ctx context.Context, req *dns.Msg) (resp *dn
Proto: dnscrypt.ProtoTCP,
})
// NOTE: Give the TCP retry its own timeout.
if p.timeout > 0 {
ctx = context.WithoutCancel(ctx)
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, p.timeout)
defer cancel()
}
resp, err = tcpClient.ExchangeContext(ctx, req, resolverInfo)
}
if err == nil && resp != nil && resp.Id != req.Id {
upstream/dnscrypt_internal_test.go
+11 -5
View File
@@ -111,7 +111,7 @@ func TestUpstreamDNSCrypt(t *testing.T) {
// Test that it responds properly
for range 10 {
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
}
}
@@ -161,8 +161,9 @@ func TestDNSCrypt_Exchange_truncated(t *testing.T) {
req := (&dns.Msg{}).SetQuestion("unit-test2.dns.adguard.com.", dns.TypeTXT)
ctx := testutil.ContextWithTimeout(t, testTimeout)
// Check that response is not truncated (even though it's huge).
res, err := u.Exchange(req)
res, err := u.Exchange(ctx, req)
require.NoError(t, err)
assert.False(t, res.Truncated)
@@ -191,7 +192,8 @@ func TestDNSCrypt_Exchange_deadline(t *testing.T) {
req := (&dns.Msg{}).SetQuestion("unit-test2.dns.adguard.com.", dns.TypeTXT)
res, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
res, err := u.Exchange(ctx, req)
require.ErrorIs(t, err, context.DeadlineExceeded)
assert.Nil(t, res)
@@ -219,8 +221,10 @@ func TestDNSCrypt_Exchange_dialFail(t *testing.T) {
require.True(t, t.Run("dial_fail", func(t *testing.T) {
testutil.CleanupAndRequireSuccess(t, u.Close)
ctx := testutil.ContextWithTimeout(t, testTimeout)
var res *dns.Msg
res, err = u.Exchange(req)
res, err = u.Exchange(ctx, req)
require.Error(t, err)
assert.Nil(t, res)
@@ -241,8 +245,10 @@ func TestDNSCrypt_Exchange_dialFail(t *testing.T) {
})
require.NoError(t, err)
ctx := testutil.ContextWithTimeout(t, testTimeout)
var res *dns.Msg
res, err = u.Exchange(req)
res, err = u.Exchange(ctx, req)
require.ErrorIs(t, err, validationErr)
assert.Nil(t, res)
upstream/doh.go
+61 -36
View File
@@ -153,16 +153,16 @@ var _ Upstream = (*dnsOverHTTPS)(nil)
func (p *dnsOverHTTPS) Address() string { return p.addrRedacted }
// Exchange implements the [Upstream] interface for *dnsOverHTTPS.
func (p *dnsOverHTTPS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (p *dnsOverHTTPS) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
// Check if there was already an active client before sending the request.
// We'll only attempt to re-connect if there was one.
client, isCached, err := p.getClient()
client, isCached, err := p.getClient(ctx)
if err != nil {
return nil, fmt.Errorf("failed to init http client: %w", err)
}
// Make the first attempt to send the DNS query.
resp, err = p.exchangeHTTPS(client, req)
resp, err = p.exchangeHTTPS(ctx, client, req)
// Make up to 2 attempts to re-create the HTTP client and send the request
// again. There are several cases (mostly, with QUIC) where this workaround
@@ -170,17 +170,17 @@ func (p *dnsOverHTTPS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
// the case when the connection was closed (due to inactivity for example)
// AND the server refuses to open a 0-RTT connection.
for i := 0; isCached && p.shouldRetry(err) && i < 2; i++ {
client, err = p.resetClient(err)
client, err = p.resetClient(ctx, err)
if err != nil {
return nil, fmt.Errorf("failed to reset http client: %w", err)
}
resp, err = p.exchangeHTTPS(client, req)
resp, err = p.exchangeHTTPS(ctx, client, req)
}
if err != nil {
// If the request failed anyway, make sure we don't use this client.
_, resErr := p.resetClient(err)
_, resErr := p.resetClient(ctx, err)
return nil, errors.WithDeferred(err, resErr)
}
@@ -217,14 +217,18 @@ func (p *dnsOverHTTPS) closeClient(client *http.Client) (err error) {
// exchangeHTTPS logs the request and its result and calls exchangeHTTPSClient.
// client and req must not be nil.
func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *dns.Msg, err error) {
func (p *dnsOverHTTPS) exchangeHTTPS(
ctx context.Context,
client *http.Client,
req *dns.Msg,
) (resp *dns.Msg, err error) {
n := networkTCP
if isHTTP3(client) {
n = networkUDP
}
logBegin(p.logger, p.addrRedacted, n, req)
defer func() { logFinish(p.logger, p.addrRedacted, n, err) }()
logBegin(ctx, p.logger, p.addrRedacted, n, req)
defer func() { logFinish(ctx, p.logger, p.addrRedacted, n, err) }()
buf, err := req.Pack()
if err != nil {
@@ -238,7 +242,7 @@ func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *d
// See https://www.rfc-editor.org/rfc/rfc8484.html.
binary.BigEndian.PutUint16(buf, 0)
resp, err = p.exchangeHTTPSClient(client, buf)
resp, err = p.exchangeHTTPSClient(ctx, client, buf)
if err != nil {
return nil, fmt.Errorf("exchanging: %w", err)
}
@@ -264,6 +268,7 @@ func (p *dnsOverHTTPS) exchangeHTTPS(client *http.Client, req *dns.Msg) (resp *d
// http.Client instance. buf is the packed DNS message that will be sent to the
// resolver. client must not be nil.
func (p *dnsOverHTTPS) exchangeHTTPSClient(
ctx context.Context,
client *http.Client,
buf []byte,
) (resp *dns.Msg, err error) {
@@ -287,7 +292,7 @@ func (p *dnsOverHTTPS) exchangeHTTPSClient(
RawQuery: q.Encode(),
}
httpReq, err := http.NewRequest(method, u.String(), nil)
httpReq, err := http.NewRequestWithContext(ctx, method, u.String(), nil)
if err != nil {
return nil, fmt.Errorf("creating http request to %s: %w", p.addrRedacted, err)
}
@@ -301,7 +306,7 @@ func (p *dnsOverHTTPS) exchangeHTTPSClient(
if err != nil {
return nil, fmt.Errorf("requesting %s: %w", p.addrRedacted, err)
}
defer slogutil.CloseAndLog(httpReq.Context(), p.logger, httpResp.Body, slog.LevelDebug)
defer slogutil.CloseAndLog(ctx, p.logger, httpResp.Body, slog.LevelDebug)
limitBody := ioutil.LimitReader(httpResp.Body, dns.MaxMsgSize)
body, err := io.ReadAll(limitBody)
@@ -359,7 +364,10 @@ func (p *dnsOverHTTPS) shouldRetry(err error) (ok bool) {
// resetClient triggers re-creation of the *http.Client that is used by this
// upstream. This method accepts the error that caused resetting client as
// depending on the error we may also reset the QUIC config.
func (p *dnsOverHTTPS) resetClient(resetErr error) (client *http.Client, err error) {
func (p *dnsOverHTTPS) resetClient(
ctx context.Context,
resetErr error,
) (client *http.Client, err error) {
p.clientMu.Lock()
defer p.clientMu.Unlock()
@@ -372,12 +380,16 @@ func (p *dnsOverHTTPS) resetClient(resetErr error) (client *http.Client, err err
if oldClient != nil {
closeErr := p.closeClient(oldClient)
if closeErr != nil {
p.logger.Warn("failed to close the old http client", slogutil.KeyError, closeErr)
p.logger.WarnContext(
ctx,
"failed to close the old http client",
slogutil.KeyError, closeErr,
)
}
}
p.logger.Debug("recreating the http client", slogutil.KeyError, resetErr)
p.client, err = p.createClient()
p.logger.DebugContext(ctx, "recreating the http client", slogutil.KeyError, resetErr)
p.client, err = p.createClient(ctx)
return p.client, err
}
@@ -403,7 +415,7 @@ func (p *dnsOverHTTPS) resetQUICConfig() {
// getClient gets or lazily initializes an HTTP client (and transport) that will
// be used for this DoH resolver.
func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
func (p *dnsOverHTTPS) getClient(ctx context.Context) (c *http.Client, isCached bool, err error) {
startTime := time.Now()
p.clientMu.Lock()
@@ -420,8 +432,8 @@ func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
return nil, false, fmt.Errorf("timeout exceeded: %s", elapsed)
}
p.logger.Debug("creating a new http client")
p.client, err = p.createClient()
p.logger.DebugContext(ctx, "creating a new http client")
p.client, err = p.createClient(ctx)
return p.client, false, err
}
@@ -430,8 +442,8 @@ func (p *dnsOverHTTPS) getClient() (c *http.Client, isCached bool, err error) {
// will depend on whether HTTP3 is allowed and provided by this upstream. Note,
// that we'll attempt to establish a QUIC connection when creating the client in
// order to check whether HTTP3 is supported.
func (p *dnsOverHTTPS) createClient() (*http.Client, error) {
transport, err := p.createTransport()
func (p *dnsOverHTTPS) createClient(ctx context.Context) (*http.Client, error) {
transport, err := p.createTransport(ctx)
if err != nil {
return nil, fmt.Errorf("initializing http transport: %w", err)
}
@@ -455,7 +467,7 @@ func (p *dnsOverHTTPS) createClient() (*http.Client, error) {
// that this function will first attempt to establish a QUIC connection (if
// HTTP3 is enabled in the upstream options). If this attempt is successful,
// it returns an HTTP3 transport, otherwise it returns the H1/H2 transport.
func (p *dnsOverHTTPS) createTransport() (t http.RoundTripper, err error) {
func (p *dnsOverHTTPS) createTransport(ctx context.Context) (t http.RoundTripper, err error) {
dialContext, err := p.getDialer()
if err != nil {
return nil, fmt.Errorf("bootstrapping %s: %w", p.addrRedacted, err)
@@ -465,14 +477,14 @@ func (p *dnsOverHTTPS) createTransport() (t http.RoundTripper, err error) {
// connection is established successfully, we'll be using HTTP3 for this
// upstream.
tlsConf := p.tlsConf.Clone()
transportH3, err := p.createTransportH3(tlsConf, dialContext)
transportH3, err := p.createTransportH3(ctx, tlsConf, dialContext)
if err == nil {
p.logger.Debug("using http/3 for this upstream, quic was faster")
p.logger.DebugContext(ctx, "using http/3 for this upstream, quic was faster")
return transportH3, nil
}
p.logger.Debug("got error, switching to http/2 for this upstream", slogutil.KeyError, err)
p.logger.DebugContext(ctx, "got error, switching to http/2 for this upstream", slogutil.KeyError, err)
if !p.supportsHTTP() {
return nil, errors.Error("HTTP1/1 and HTTP2 are not supported by this upstream")
@@ -558,6 +570,7 @@ func (h *http3Transport) Close() (err error) {
// parallel (one for TLS, the other one for QUIC) and if QUIC is faster it will
// create the [*http3.Transport] instance.
func (p *dnsOverHTTPS) createTransportH3(
ctx context.Context,
tlsConfig *tls.Config,
dialContext bootstrap.DialHandler,
) (roundTripper http.RoundTripper, err error) {
@@ -565,7 +578,7 @@ func (p *dnsOverHTTPS) createTransportH3(
return nil, errors.Error("HTTP3 support is not enabled")
}
addr, err := p.probeH3(tlsConfig, dialContext)
addr, err := p.probeH3(ctx, tlsConfig, dialContext)
if err != nil {
return nil, err
}
@@ -594,13 +607,14 @@ func (p *dnsOverHTTPS) createTransportH3(
// upstream. If the test is successful it will return the address that we
// should use to establish the QUIC connections.
func (p *dnsOverHTTPS) probeH3(
ctx context.Context,
tlsConfig *tls.Config,
dialContext bootstrap.DialHandler,
) (addr string, err error) {
// We're using bootstrapped address instead of what's passed to the function
// it does not create an actual connection, but it helps us determine
// what IP is actually reachable (when there are v4/v6 addresses).
rawConn, err := dialContext(context.Background(), "udp", "")
rawConn, err := dialContext(ctx, "udp", "")
if err != nil {
return "", fmt.Errorf("failed to dial: %w", err)
}
@@ -635,8 +649,8 @@ func (p *dnsOverHTTPS) probeH3(
// Run probeQUIC and probeTLS in parallel and see which one is faster.
chQUIC := make(chan error, 1)
chTLS := make(chan error, 1)
go p.probeQUIC(addr, probeTLSCfg, chQUIC)
go p.probeTLS(dialContext, probeTLSCfg, chTLS)
go p.probeQUIC(ctx, addr, probeTLSCfg, chQUIC)
go p.probeTLS(ctx, dialContext, probeTLSCfg, chTLS)
select {
case quicErr := <-chQUIC:
@@ -650,7 +664,7 @@ func (p *dnsOverHTTPS) probeH3(
case tlsErr := <-chTLS:
if tlsErr != nil {
// Return immediately, TLS failed.
p.logger.Debug("probing tls", slogutil.KeyError, tlsErr)
p.logger.DebugContext(ctx, "probing tls", slogutil.KeyError, tlsErr)
return addr, nil
}
@@ -661,14 +675,20 @@ func (p *dnsOverHTTPS) probeH3(
// probeQUIC attempts to establish a QUIC connection to the specified address.
// We run probeQUIC and probeTLS in parallel and see which one is faster.
func (p *dnsOverHTTPS) probeQUIC(addr string, tlsConfig *tls.Config, ch chan error) {
func (p *dnsOverHTTPS) probeQUIC(
ctx context.Context,
addr string,
tlsConfig *tls.Config,
ch chan error,
) {
startTime := time.Now()
t := p.timeout
if t == 0 {
t = dialTimeout
}
ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(t))
ctx, cancel := context.WithTimeout(ctx, t)
defer cancel()
conn, err := quic.DialAddrEarly(ctx, addr, tlsConfig, p.getQUICConfig())
@@ -683,15 +703,20 @@ func (p *dnsOverHTTPS) probeQUIC(addr string, tlsConfig *tls.Config, ch chan err
ch <- nil
elapsed := time.Since(startTime)
p.logger.Debug("quic connection established", "elapsed", elapsed)
p.logger.DebugContext(ctx, "quic connection established", "elapsed", elapsed)
}
// probeTLS attempts to establish a TLS connection to the specified address. We
// run probeQUIC and probeTLS in parallel and see which one is faster.
func (p *dnsOverHTTPS) probeTLS(dialContext bootstrap.DialHandler, tlsConfig *tls.Config, ch chan error) {
func (p *dnsOverHTTPS) probeTLS(
ctx context.Context,
dialContext bootstrap.DialHandler,
tlsConfig *tls.Config,
ch chan error,
) {
startTime := time.Now()
conn, err := tlsDial(dialContext, tlsConfig)
conn, err := tlsDial(ctx, dialContext, tlsConfig)
if err != nil {
ch <- fmt.Errorf("opening TLS connection: %w", err)
return
@@ -703,7 +728,7 @@ func (p *dnsOverHTTPS) probeTLS(dialContext bootstrap.DialHandler, tlsConfig *tl
ch <- nil
elapsed := time.Since(startTime)
p.logger.Debug("tls connection established", "elapsed", elapsed)
p.logger.DebugContext(ctx, "tls connection established", "elapsed", elapsed)
}
// supportsH3 returns true if HTTP/3 is supported by this upstream.
upstream/doh_internal_test.go
+12 -9
View File
@@ -97,7 +97,7 @@ func TestUpstreamDoH(t *testing.T) {
// Test that it responds properly.
for range 10 {
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
}
doh := testutil.RequireTypeAssert[*dnsOverHTTPS](t, u)
@@ -106,7 +106,7 @@ func TestUpstreamDoH(t *testing.T) {
doh.client = nil
// Force it to establish the connection again.
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
// Check that TLS session was resumed properly.
require.True(t, lastState.DidResume)
@@ -194,7 +194,7 @@ func TestUpstreamDoH_raceReconnect(t *testing.T) {
require.NoError(t, err)
testutil.CleanupAndRequireSuccess(t, u.Close)
checkRaceCondition(u)
checkRaceCondition(t, u)
})
}
}
@@ -238,7 +238,7 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
})
require.NoError(t, err)
checkUpstream(t, u, upsAddr)
checkUpstream(t, u, upsAddr, testTimeout)
})
require.False(t, t.Failed())
testutil.CleanupAndRequireSuccess(t, u.Close)
@@ -249,12 +249,13 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
port: int(addr.Port()),
})
checkUpstream(t, u, upsAddr)
checkUpstream(t, u, upsAddr, testTimeout)
})
require.False(t, t.Failed())
t.Run("retry", func(t *testing.T) {
_, err := u.Exchange(createTestMessage())
ctx := testutil.ContextWithTimeout(t, testTimeout)
_, err := u.Exchange(ctx, createTestMessage())
require.Error(t, err)
_ = startDoHServer(t, testDoHServerOptions{
@@ -262,7 +263,7 @@ func TestUpstreamDoH_serverRestart(t *testing.T) {
port: int(addr.Port()),
})
checkUpstream(t, u, upsAddr)
checkUpstream(t, u, upsAddr, testTimeout)
})
})
}
@@ -291,7 +292,8 @@ func TestUpstreamDoH_0RTT(t *testing.T) {
req := createTestMessage()
// Trigger connection to a DoH3 server.
resp, err := uh.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
resp, err := uh.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, resp)
@@ -307,7 +309,8 @@ func TestUpstreamDoH_0RTT(t *testing.T) {
}()
// Trigger second connection.
resp, err = uh.Exchange(req)
ctx = testutil.ContextWithTimeout(t, testTimeout)
resp, err = uh.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, resp)
upstream/doq.go
+32 -24
View File
@@ -167,7 +167,7 @@ var _ Upstream = (*dnsOverQUIC)(nil)
func (p *dnsOverQUIC) Address() string { return p.addr.String() }
// Exchange implements the [Upstream] interface for *dnsOverQUIC.
func (p *dnsOverQUIC) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (p *dnsOverQUIC) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
// When sending queries over a QUIC connection, the DNS Message ID MUST be
// set to 0. The stream mapping for DoQ allows for unambiguous correlation
// of queries and responses, so the Message ID field is not required.
@@ -184,39 +184,43 @@ func (p *dnsOverQUIC) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
}()
// Gets or opens a QUIC connection to use for this query.
conn, cached, err := p.getConnection()
conn, cached, err := p.getConnection(ctx)
if err != nil {
return nil, fmt.Errorf("getting conn: %w", err)
}
// Make the first attempt to send the DNS query.
resp, err = p.exchangeQUIC(req, conn)
resp, err = p.exchangeQUIC(ctx, req, conn)
// Failure to use a cached connection should be handled gracefully as this
// connection could have been closed by the server or simply be broken due
// to how UDP NAT works. In this case the connection should be re-created.
if cached && err != nil {
p.logger.Debug("recreating the quic connection and retrying", slogutil.KeyError, err)
p.logger.DebugContext(
ctx,
"recreating the quic connection and retrying",
slogutil.KeyError, err,
)
// Close the active connection to make sure the cached connection is
// cleaned up.
p.closeConnWithError(conn, err)
p.closeConnWithError(ctx, conn, err)
// Get or re-create the QUIC connection in order to make the second
// attempt.
conn, _, err = p.getConnection()
conn, _, err = p.getConnection(ctx)
if err != nil {
return nil, fmt.Errorf("getting new conn: %w", err)
}
// Retry sending the request through the new connection.
resp, err = p.exchangeQUIC(req, conn)
resp, err = p.exchangeQUIC(ctx, req, conn)
}
if err != nil {
// If we're unable to exchange messages, make sure the connection is
// closed and signal about an internal error.
p.closeConnWithError(conn, err)
p.closeConnWithError(ctx, conn, err)
}
return resp, err
@@ -238,18 +242,22 @@ func (p *dnsOverQUIC) Close() (err error) {
// exchangeQUIC attempts to open a new QUIC stream, send the DNS message
// through it and return the response it got from the server.
func (p *dnsOverQUIC) exchangeQUIC(req *dns.Msg, conn *quic.Conn) (resp *dns.Msg, err error) {
func (p *dnsOverQUIC) exchangeQUIC(
ctx context.Context,
req *dns.Msg,
conn *quic.Conn,
) (resp *dns.Msg, err error) {
addr := p.Address()
logBegin(p.logger, addr, networkUDP, req)
defer func() { logFinish(p.logger, addr, networkUDP, err) }()
logBegin(ctx, p.logger, addr, networkUDP, req)
defer func() { logFinish(ctx, p.logger, addr, networkUDP, err) }()
buf, err := req.Pack()
if err != nil {
return nil, fmt.Errorf("failed to pack DNS message for DoQ: %w", err)
}
stream, err := p.openStream(conn)
stream, err := p.openStream(ctx, conn)
if err != nil {
return nil, fmt.Errorf("opening stream: %w", err)
}
@@ -272,7 +280,7 @@ func (p *dnsOverQUIC) exchangeQUIC(req *dns.Msg, conn *quic.Conn) (resp *dns.Msg
// of the stream, but does not prevent reading from it.
err = stream.Close()
if err != nil {
p.logger.Debug("closing quic stream", slogutil.KeyError, err)
p.logger.DebugContext(ctx, "closing quic stream", slogutil.KeyError, err)
}
return p.readMsg(stream)
@@ -298,7 +306,7 @@ func (p *dnsOverQUIC) getBytesPool() (pool *sync.Pool) {
// getConnection opens or returns an existing *quic.Conn and indicates whether
// it opened a new connection or used an existing cached one.
func (p *dnsOverQUIC) getConnection() (conn *quic.Conn, cached bool, err error) {
func (p *dnsOverQUIC) getConnection(ctx context.Context) (conn *quic.Conn, cached bool, err error) {
p.connMu.Lock()
defer p.connMu.Unlock()
@@ -307,7 +315,7 @@ func (p *dnsOverQUIC) getConnection() (conn *quic.Conn, cached bool, err error)
return conn, true, nil
}
conn, err = p.openConnection()
conn, err = p.openConnection(ctx)
if err != nil {
return nil, false, err
}
@@ -337,8 +345,8 @@ func (p *dnsOverQUIC) resetQUICConfig() {
}
// openStream opens a new QUIC stream for the specified connection.
func (p *dnsOverQUIC) openStream(conn *quic.Conn) (*quic.Stream, error) {
ctx, cancel := p.withDeadline(context.Background())
func (p *dnsOverQUIC) openStream(ctx context.Context, conn *quic.Conn) (*quic.Stream, error) {
ctx, cancel := p.withDeadline(ctx)
defer cancel()
stream, err := conn.OpenStreamSync(ctx)
@@ -350,16 +358,16 @@ func (p *dnsOverQUIC) openStream(conn *quic.Conn) (*quic.Stream, error) {
}
// openConnection dials a new QUIC connection.
func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
func (p *dnsOverQUIC) openConnection(ctx context.Context) (conn *quic.Conn, err error) {
dialContext, err := p.getDialer()
if err != nil {
return nil, fmt.Errorf("bootstrapping %s: %w", p.addr, err)
}
// we're using bootstrapped address instead of what's passed to the function
// We're using bootstrapped address instead of what's passed to the function
// it does not create an actual connection, but it helps us determine
// what IP is actually reachable (when there're v4/v6 addresses).
rawConn, err := dialContext(context.Background(), "udp", "")
rawConn, err := dialContext(ctx, "udp", "")
if err != nil {
return nil, fmt.Errorf("dialing raw connection to %s: %w", p.addr, err)
}
@@ -367,7 +375,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
// It's never actually used.
err = rawConn.Close()
if err != nil {
p.logger.Debug("closing raw connection", "addr", p.addr, slogutil.KeyError, err)
p.logger.DebugContext(ctx, "closing raw connection", "addr", p.addr, slogutil.KeyError, err)
}
udpConn, ok := rawConn.(*net.UDPConn)
@@ -377,7 +385,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
addr := udpConn.RemoteAddr().String()
ctx, cancel := p.withDeadline(context.Background())
ctx, cancel := p.withDeadline(ctx)
defer cancel()
conn, err = quic.DialAddrEarly(ctx, addr, p.tlsConf.Clone(), p.getQUICConfig())
@@ -391,7 +399,7 @@ func (p *dnsOverQUIC) openConnection() (conn *quic.Conn, err error) {
// closeConnWithError closes the active connection with error to make sure that
// new queries were processed in another connection. We can do that in the case
// of a fatal error.
func (p *dnsOverQUIC) closeConnWithError(conn *quic.Conn, err error) {
func (p *dnsOverQUIC) closeConnWithError(ctx context.Context, conn *quic.Conn, err error) {
p.connMu.Lock()
defer p.connMu.Unlock()
@@ -407,7 +415,7 @@ func (p *dnsOverQUIC) closeConnWithError(conn *quic.Conn, err error) {
err = conn.CloseWithError(code, "")
if err != nil {
p.logger.Error("failed to close the conn", slogutil.KeyError, err)
p.logger.ErrorContext(ctx, "failed to close the conn", slogutil.KeyError, err)
}
// If the connection that's being closed is cached, reset the cache.
upstream/doq_internal_test.go
+17 -14
View File
@@ -49,7 +49,7 @@ func TestDNSOverQUIC(t *testing.T) {
// Test that it responds properly
for range 10 {
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
if conn == nil {
conn = uq.conn
@@ -63,7 +63,7 @@ func TestDNSOverQUIC(t *testing.T) {
_ = conn.CloseWithError(quic.ApplicationErrorCode(0), "")
// Try to establish it again.
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
// Make sure that the session has been resumed.
require.True(t, lastState.DidResume)
@@ -74,7 +74,7 @@ func TestDNSOverQUIC(t *testing.T) {
require.NoError(t, err)
testutil.CleanupAndRequireSuccess(t, u.Close)
checkRaceCondition(u)
checkRaceCondition(t, u)
}
func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
@@ -96,7 +96,7 @@ func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
testutil.CleanupAndRequireSuccess(t, u.Close)
// Test that the upstream works properly.
checkUpstream(t, u, address)
checkUpstream(t, u, address, testTimeout)
// Close all active connections.
err = srv.closeConns()
@@ -113,17 +113,17 @@ func TestDNSOverQUIC_Exchange_quicCloseConn(t *testing.T) {
wg := sync.WaitGroup{}
wg.Add(parallelQueries)
ctx := testutil.ContextWithTimeout(t, testTimeout)
for range 10 {
pt := testutil.PanicT{}
go func(t assert.TestingT) {
go func(ctx context.Context, t assert.TestingT) {
defer wg.Done()
req := createTestMessage()
_, errExch := u.Exchange(req)
_, errExch := u.Exchange(ctx, req)
assert.NoError(t, errExch)
}(pt)
}(ctx, pt)
}
wg.Wait()
@@ -160,7 +160,7 @@ func TestDNSOverQUIC_serverRestart(t *testing.T) {
)
require.NoError(t, err)
checkUpstream(t, u, upsStr)
checkUpstream(t, u, upsStr, testTimeout)
})
require.False(t, t.Failed())
testutil.CleanupAndRequireSuccess(t, u.Close)
@@ -168,17 +168,18 @@ func TestDNSOverQUIC_serverRestart(t *testing.T) {
t.Run("second_try", func(t *testing.T) {
_ = startDoQServer(t, tlsConf, int(addr.Port()))
checkUpstream(t, u, upsStr)
checkUpstream(t, u, upsStr, testTimeout)
})
require.False(t, t.Failed())
t.Run("retry", func(t *testing.T) {
_, err := u.Exchange(createTestMessage())
ctx := testutil.ContextWithTimeout(t, testTimeout)
_, err := u.Exchange(ctx, createTestMessage())
require.Error(t, err)
_ = startDoQServer(t, tlsConf, int(addr.Port()))
checkUpstream(t, u, upsStr)
checkUpstream(t, u, upsStr, testTimeout)
})
}
@@ -199,9 +200,10 @@ func TestDNSOverQUIC_0RTT(t *testing.T) {
uq := testutil.RequireTypeAssert[*dnsOverQUIC](t, u)
req := createTestMessage()
ctx := testutil.ContextWithTimeout(t, testTimeout)
// Trigger connection to a QUIC server.
resp, err := uq.Exchange(req)
resp, err := uq.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, resp)
@@ -216,8 +218,9 @@ func TestDNSOverQUIC_0RTT(t *testing.T) {
uq.conn = nil
}()
ctx = testutil.ContextWithTimeout(t, testTimeout)
// Trigger second connection.
resp, err = uq.Exchange(req)
resp, err = uq.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, resp)
upstream/dot.go
+46 -15
View File
@@ -49,6 +49,9 @@ type dnsOverTLS struct {
// This leads to weak performance for all exchanges coming across such
// connections.
conns []net.Conn
// timeout is the timeout for DNS request reads and writes.
timeout time.Duration
}
// newDoT returns the DNS-over-TLS Upstream.
@@ -81,6 +84,7 @@ func newDoT(addr *url.URL, opts *Options) (ups Upstream, err error) {
},
connsMu: &sync.Mutex{},
logger: opts.Logger,
timeout: opts.Timeout,
}
runtime.SetFinalizer(tlsUps, (*dnsOverTLS).Close)
@@ -95,28 +99,33 @@ var _ Upstream = (*dnsOverTLS)(nil)
func (p *dnsOverTLS) Address() string { return p.addr.String() }
// Exchange implements the [Upstream] interface for *dnsOverTLS.
func (p *dnsOverTLS) Exchange(req *dns.Msg) (reply *dns.Msg, err error) {
func (p *dnsOverTLS) Exchange(ctx context.Context, req *dns.Msg) (reply *dns.Msg, err error) {
h, err := p.getDialer()
if err != nil {
return nil, fmt.Errorf("getting conn to %s: %w", p.addr, err)
}
conn, err := p.conn(h)
conn, err := p.conn(ctx, h)
if err != nil {
return nil, fmt.Errorf("getting conn to %s: %w", p.addr, err)
}
reply, err = p.exchangeWithConn(conn, req)
reply, err = p.exchangeWithConn(ctx, conn, req)
if err != nil {
// The pooled connection might have been closed already, see
// https://github.com/AdguardTeam/dnsproxy/issues/3. The following
// connection from pool may also be malformed, so dial a new one.
err = errors.WithDeferred(err, conn.Close())
p.logger.Debug("dot got bad conn from pool", "addr", p.addr, slogutil.KeyError, err)
p.logger.DebugContext(
ctx,
"dot got bad conn from pool",
"addr", p.addr,
slogutil.KeyError, err,
)
// Retry.
conn, err = tlsDial(h, p.tlsConf.Clone())
conn, err = tlsDial(ctx, h, p.tlsConf.Clone())
if err != nil {
return nil, fmt.Errorf(
"dialing %s: connecting to %s: %w",
@@ -126,7 +135,7 @@ func (p *dnsOverTLS) Exchange(req *dns.Msg) (reply *dns.Msg, err error) {
)
}
reply, err = p.exchangeWithConn(conn, req)
reply, err = p.exchangeWithConn(ctx, conn, req)
if err != nil {
return reply, errors.WithDeferred(err, conn.Close())
}
@@ -157,11 +166,11 @@ func (p *dnsOverTLS) Close() (err error) {
// conn returns the first available connection from the pool if there is any, or
// dials a new one otherwise.
func (p *dnsOverTLS) conn(h bootstrap.DialHandler) (conn net.Conn, err error) {
func (p *dnsOverTLS) conn(ctx context.Context, h bootstrap.DialHandler) (conn net.Conn, err error) {
// Dial a new connection outside the lock, if needed.
defer func() {
if conn == nil {
conn, err = tlsDial(h, p.tlsConf.Clone())
conn, err = tlsDial(ctx, h, p.tlsConf.Clone())
err = errors.Annotate(err, "connecting to %s: %w", p.tlsConf.ServerName)
}
}()
@@ -178,14 +187,18 @@ func (p *dnsOverTLS) conn(h bootstrap.DialHandler) (conn net.Conn, err error) {
err = conn.SetDeadline(time.Now().Add(dialTimeout))
if err != nil {
p.logger.Debug("dot upstream setting deadline to conn from pool", slogutil.KeyError, err)
p.logger.DebugContext(
ctx,
"dot upstream setting deadline to conn from pool",
slogutil.KeyError, err,
)
// If deadLine can't be updated it means that connection was already
// closed.
return nil, nil
}
p.logger.Debug("dot upstream using existing conn", "raddr", conn.RemoteAddr())
p.logger.DebugContext(ctx, "dot upstream using existing conn", "raddr", conn.RemoteAddr())
return conn, nil
}
@@ -198,14 +211,26 @@ func (p *dnsOverTLS) putBack(conn net.Conn) {
}
// exchangeWithConn tries to exchange the query using conn.
func (p *dnsOverTLS) exchangeWithConn(conn net.Conn, req *dns.Msg) (reply *dns.Msg, err error) {
func (p *dnsOverTLS) exchangeWithConn(
ctx context.Context,
conn net.Conn,
req *dns.Msg,
) (reply *dns.Msg, err error) {
addr := p.Address()
logBegin(p.logger, addr, networkTCP, req)
defer func() { logFinish(p.logger, addr, networkTCP, err) }()
logBegin(ctx, p.logger, addr, networkTCP, req)
defer func() { logFinish(ctx, p.logger, addr, networkTCP, err) }()
dnsConn := dns.Conn{Conn: conn}
if p.timeout > 0 {
deadline := time.Now().Add(p.timeout)
err = dnsConn.SetDeadline(deadline)
if err != nil {
return nil, fmt.Errorf("setting connection deadline: %w", err)
}
}
err = dnsConn.WriteMsg(req)
if err != nil {
return nil, fmt.Errorf("sending request to %s: %w", addr, err)
@@ -223,10 +248,14 @@ func (p *dnsOverTLS) exchangeWithConn(conn net.Conn, req *dns.Msg) (reply *dns.M
// tlsDial is basically the same as tls.DialWithDialer, but we will call our own
// dialContext function to get connection.
func tlsDial(dialContext bootstrap.DialHandler, conf *tls.Config) (c *tls.Conn, err error) {
func tlsDial(
ctx context.Context,
dialContext bootstrap.DialHandler,
conf *tls.Config,
) (c *tls.Conn, err error) {
// We're using bootstrapped address instead of what's passed to the
// function.
rawConn, err := dialContext(context.Background(), networkTCP, "")
rawConn, err := dialContext(ctx, networkTCP, "")
if err != nil {
return nil, err
}
@@ -234,6 +263,8 @@ func tlsDial(dialContext bootstrap.DialHandler, conf *tls.Config) (c *tls.Conn,
// We want the timeout to cover the whole process: TCP connection and TLS
// handshake dialTimeout will be used as connection deadLine.
conn := tls.Client(rawConn, conf)
// TODO(f.setrakov): Make it configurable.
err = conn.SetDeadline(time.Now().Add(dialTimeout))
if err != nil {
// Must not happen in normal circumstances.
upstream/dot_internal_test.go
+18 -11
View File
@@ -1,6 +1,7 @@
package upstream
import (
"context"
"crypto/tls"
"crypto/x509"
"fmt"
@@ -38,7 +39,7 @@ func TestUpstream_dnsOverTLS(t *testing.T) {
// Test that it responds properly.
for range 10 {
checkUpstream(t, u, addr)
checkUpstream(t, u, addr, testTimeout)
}
}
@@ -69,8 +70,9 @@ func TestUpstream_dnsOverTLS_race(t *testing.T) {
wg.Go(func() {
pt := testutil.PanicT{}
ctx := testutil.ContextWithTimeout(t, testTimeout)
req := createTestMessage()
resp, uErr := u.Exchange(req)
resp, uErr := u.Exchange(ctx, req)
require.NoError(pt, uErr)
requireResponse(pt, req, resp)
})
@@ -112,7 +114,8 @@ func TestUpstream_dnsOverTLS_poolReconnect(t *testing.T) {
// Send the first test message.
req := createTestMessage()
reply, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
reply, err := u.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, reply)
@@ -123,7 +126,8 @@ func TestUpstream_dnsOverTLS_poolReconnect(t *testing.T) {
// Send the second test message.
req = createTestMessage()
reply, err = u.Exchange(req)
ctx = testutil.ContextWithTimeout(t, testTimeout)
reply, err = u.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, reply)
@@ -154,7 +158,8 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
// Send the first test message.
req := createTestMessage()
response, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
response, err := u.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, response)
@@ -167,11 +172,12 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
dialHandler, err := p.getDialer()
require.NoError(t, err)
usedConn, err := p.conn(dialHandler)
ctx = testutil.ContextWithTimeout(t, testTimeout)
usedConn, err := p.conn(ctx, dialHandler)
require.NoError(t, err)
require.Same(t, usedConn, conn)
response, err = p.exchangeWithConn(conn, req)
response, err = p.exchangeWithConn(ctx, conn, req)
require.NoError(t, err)
requireResponse(t, req, response)
@@ -185,11 +191,12 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
require.Len(t, p.conns, 1)
conn = p.conns[0]
usedConn, err = p.conn(dialHandler)
ctx = testutil.ContextWithTimeout(t, testTimeout)
usedConn, err = p.conn(ctx, dialHandler)
require.NoError(t, err)
require.Same(t, usedConn, conn)
response, err = p.exchangeWithConn(usedConn, req)
response, err = p.exchangeWithConn(ctx, usedConn, req)
require.NoError(t, err)
requireResponse(t, req, response)
@@ -198,7 +205,7 @@ func TestUpstream_dnsOverTLS_poolDeadline(t *testing.T) {
require.NoError(t, err)
// Connection with expired deadLine can't be used.
response, err = p.exchangeWithConn(usedConn, req)
response, err = p.exchangeWithConn(ctx, usedConn, req)
require.Error(t, err)
require.Nil(t, response)
}
@@ -297,7 +304,7 @@ func BenchmarkDoTUpstream(b *testing.B) {
b.RunParallel(func(p *testing.PB) {
for p.Next() {
_, _ = u.Exchange(<-reqChan)
_, _ = u.Exchange(context.Background(), <-reqChan)
}
})
})
upstream/parallel.go
+19 -9
View File
@@ -1,6 +1,7 @@
package upstream
import (
"context"
"fmt"
"slices"
@@ -21,13 +22,17 @@ const (
// ExchangeParallel returns the first successful response from one of u. It
// returns an error if all upstreams failed to exchange the request.
func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Upstream, err error) {
func ExchangeParallel(
ctx context.Context,
ups []Upstream,
req *dns.Msg,
) (reply *dns.Msg, resolved Upstream, err error) {
upsNum := len(ups)
switch upsNum {
case 0:
return nil, nil, ErrNoUpstreams
case 1:
return exchangeSingle(ups[0], req)
return exchangeSingle(ctx, ups[0], req)
default:
// Go on.
}
@@ -39,7 +44,7 @@ func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Up
//
// TODO(s.chzhen): Consider using buffer pool.
copyReq := req.Copy()
go exchangeAsync(f, copyReq, resCh)
go exchangeAsync(ctx, f, copyReq, resCh)
}
errs := []error{}
@@ -68,10 +73,11 @@ func ExchangeParallel(ups []Upstream, req *dns.Msg) (reply *dns.Msg, resolved Up
// exchangeSingle returns a successful response and resolver if a DNS lookup was
// successful.
func exchangeSingle(
ctx context.Context,
ups Upstream,
req *dns.Msg,
) (resp *dns.Msg, resolved Upstream, err error) {
resp, err = ups.Exchange(req)
resp, err = ups.Exchange(ctx, req)
if err != nil {
return nil, nil, err
}
@@ -91,14 +97,18 @@ type ExchangeAllResult struct {
// ExchangeAll returns the responses from all of u. It returns an error only if
// all upstreams failed to exchange the request.
func ExchangeAll(ups []Upstream, req *dns.Msg) (res []ExchangeAllResult, err error) {
func ExchangeAll(
ctx context.Context,
ups []Upstream,
req *dns.Msg,
) (res []ExchangeAllResult, err error) {
upsNum := len(ups)
switch upsNum {
case 0:
return nil, ErrNoUpstreams
case 1:
var reply *dns.Msg
reply, err = ups[0].Exchange(req)
reply, err = ups[0].Exchange(ctx, req)
if err != nil {
return nil, err
} else if reply == nil {
@@ -122,7 +132,7 @@ func ExchangeAll(ups []Upstream, req *dns.Msg) (res []ExchangeAllResult, err err
//
// TODO(s.chzhen): Consider using buffer pool.
copyReq := req.Copy()
go exchangeAsync(u, copyReq, resCh)
go exchangeAsync(ctx, u, copyReq, resCh)
}
// Wait for all exchanges to finish.
@@ -162,8 +172,8 @@ func receiveAsyncResult(resCh chan any) (res *ExchangeAllResult, err error) {
// exchangeAsync tries to resolve DNS request with one upstream and sends the
// result to respCh.
func exchangeAsync(u Upstream, req *dns.Msg, resCh chan any) {
reply, err := u.Exchange(req)
func exchangeAsync(ctx context.Context, u Upstream, req *dns.Msg, resCh chan any) {
reply, err := u.Exchange(ctx, req)
if err != nil {
resCh <- err
} else {
upstream/parallel_internal_test.go
+8 -4
View File
@@ -1,6 +1,7 @@
package upstream
import (
"context"
"fmt"
"net/netip"
"testing"
@@ -30,7 +31,8 @@ func TestExchangeParallel(t *testing.T) {
req := createTestMessage()
start := time.Now()
resp, u, err := ExchangeParallel(upstreams, req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
resp, u, err := ExchangeParallel(ctx, upstreams, req)
if err != nil {
t.Fatalf("no response from test upstreams: %s", err)
}
@@ -53,7 +55,8 @@ func TestExchangeParallelEmpty(t *testing.T) {
}
req := createTestMessage()
resp, up, err := ExchangeParallel(ups, req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
resp, up, err := ExchangeParallel(ctx, ups, req)
require.Error(t, err)
assert.Nil(t, resp)
@@ -79,7 +82,7 @@ type testUpstream struct {
var _ Upstream = (*testUpstream)(nil)
// Exchange implements the [Upstream] interface for *testUpstream.
func (u *testUpstream) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (u *testUpstream) Exchange(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
if u.sleep != 0 {
time.Sleep(u.sleep)
}
@@ -130,7 +133,8 @@ func TestExchangeAll(t *testing.T) {
}}
req := createHostTestMessage("test.org")
res, err := ExchangeAll(ups, req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
res, err := ExchangeAll(ctx, ups, req)
require.NoError(t, err)
require.Len(t, res, 2)
upstream/plain.go
+34 -13
View File
@@ -86,11 +86,21 @@ func (p *plainDNS) Address() string {
// dialExchange performs a DNS exchange with the specified dial handler.
// network must be either [networkUDP] or [networkTCP].
func (p *plainDNS) dialExchange(
ctx context.Context,
network network,
dial bootstrap.DialHandler,
req *dns.Msg,
) (resp *dns.Msg, err error) {
var cancel context.CancelFunc
if p.timeout > 0 {
ctx, cancel = context.WithTimeout(ctx, p.timeout)
defer cancel()
}
addr := p.Address()
// NOTE: It is important to set a client timeout to prevent the use of the
// default timeout settings.
client := &dns.Client{Timeout: p.timeout}
conn := &dns.Conn{}
@@ -101,25 +111,34 @@ func (p *plainDNS) dialExchange(
}
}()
logBegin(p.logger, addr, network, upstreamReq)
defer func() { logFinish(p.logger, addr, network, err) }()
logBegin(ctx, p.logger, addr, network, upstreamReq)
defer func() { logFinish(ctx, p.logger, addr, network, err) }()
ctx := context.Background()
conn.Conn, err = dial(ctx, network, "")
// NOTE: Discard context timeout to prevent deadline sharing between dial
// calls.
//
// TODO(f.setrakov): Consider supporting context timeout here.
conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
if err != nil {
return nil, fmt.Errorf("dialing %s over %s: %w", p.addr.Host, network, err)
}
defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)
resp, _, err = client.ExchangeWithConn(upstreamReq, conn)
resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
if isExpectedConnErr(err) {
conn.Conn, err = dial(ctx, network, "")
if p.timeout > 0 {
ctx = context.WithoutCancel(ctx)
ctx, cancel = context.WithTimeout(ctx, p.timeout)
defer cancel()
}
conn.Conn, err = dial(context.WithoutCancel(ctx), network, "")
if err != nil {
return nil, fmt.Errorf("dialing %s over %s again: %w", p.addr.Host, network, err)
}
defer func(c net.Conn) { err = errors.WithDeferred(err, c.Close()) }(conn.Conn)
resp, _, err = client.ExchangeWithConn(upstreamReq, conn)
resp, _, err = client.ExchangeWithConnContext(ctx, upstreamReq, conn)
}
if err != nil {
@@ -159,7 +178,7 @@ func isExpectedConnErr(err error) (is bool) {
}
// Exchange implements the [Upstream] interface for *plainDNS.
func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
func (p *plainDNS) Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
dial, err := p.getDialer()
if err != nil {
// Don't wrap the error since it's informative enough as is.
@@ -168,7 +187,7 @@ func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
addr := p.Address()
resp, err = p.dialExchange(p.net, dial, req)
resp, err = p.dialExchange(ctx, p.net, dial, req)
if p.net != networkUDP {
// The network is already TCP.
return resp, err
@@ -181,22 +200,24 @@ func (p *plainDNS) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
if errors.Is(err, errQuestion) {
// The upstream responds with malformed messages, so try TCP.
p.logger.Debug(
p.logger.DebugContext(
ctx,
"plain response is malformed, using tcp",
"addr", addr,
slogutil.KeyError, err,
)
return p.dialExchange(networkTCP, dial, req)
return p.dialExchange(ctx, networkTCP, dial, req)
} else if resp.Truncated {
// Fallback to TCP on truncated responses.
p.logger.Debug(
p.logger.DebugContext(
ctx,
"plain response is truncated, using tcp",
"question", &req.Question[0],
"addr", addr,
)
return p.dialExchange(networkTCP, dial, req)
return p.dialExchange(ctx, networkTCP, dial, req)
}
// There is either no error or the error isn't related to the received
upstream/plain_internal_test.go
+5 -3
View File
@@ -34,7 +34,7 @@ func TestUpstream_plainDNS(t *testing.T) {
testutil.CleanupAndRequireSuccess(t, u.Close)
for range 10 {
checkUpstream(t, u, addr)
checkUpstream(t, u, addr, testTimeout)
}
}
@@ -57,7 +57,8 @@ func TestUpstream_plainDNS_badID(t *testing.T) {
require.NoError(t, err)
testutil.CleanupAndRequireSuccess(t, u.Close)
resp, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
resp, err := u.Exchange(ctx, req)
var netErr net.Error
require.ErrorAs(t, err, &netErr)
@@ -132,7 +133,8 @@ func TestUpstream_plainDNS_fallbackToTCP(t *testing.T) {
require.NoError(t, err)
testutil.CleanupAndRequireSuccess(t, u.Close)
resp, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
resp, err := u.Exchange(ctx, req)
require.NoError(t, err)
requireResponse(t, req, resp)
upstream/resolver.go
+16 -10
View File
@@ -154,16 +154,14 @@ type ipResult struct {
// lookupNetIP performs a DNS lookup of host and returns the result. network
// must be either [bootstrap.NetworkIP4], [bootstrap.NetworkIP6], or
// [bootstrap.NetworkIP]. host must be in a lower-case FQDN form.
//
// TODO(e.burkov): Use context.
func (r *UpstreamResolver) lookupNetIP(
_ context.Context,
ctx context.Context,
network bootstrap.Network,
host string,
) (result *ipResult, err error) {
switch network {
case bootstrap.NetworkIP4, bootstrap.NetworkIP6:
return r.request(host, network)
return r.request(ctx, host, network)
case bootstrap.NetworkIP:
// Go on.
default:
@@ -171,8 +169,8 @@ func (r *UpstreamResolver) lookupNetIP(
}
resCh := make(chan any, 2)
go r.resolveAsync(resCh, host, bootstrap.NetworkIP4)
go r.resolveAsync(resCh, host, bootstrap.NetworkIP6)
go r.resolveAsync(ctx, resCh, host, bootstrap.NetworkIP4)
go r.resolveAsync(ctx, resCh, host, bootstrap.NetworkIP6)
var errs []error
result = &ipResult{}
@@ -199,7 +197,11 @@ func (r *UpstreamResolver) lookupNetIP(
//
// TODO(e.burkov): Consider NS and Extra sections when setting TTL. Check out
// what RFCs say about it.
func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipResult, err error) {
func (r *UpstreamResolver) request(
ctx context.Context,
host string,
n bootstrap.Network,
) (res *ipResult, err error) {
var qtype uint16
switch n {
case bootstrap.NetworkIP4:
@@ -224,7 +226,7 @@ func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipRes
// As per [Upstream.Exchange] documentation, the response is always returned
// if no error occurred.
resp, err := r.Exchange(req)
resp, err := r.Exchange(ctx, req)
if err != nil {
return res, err
}
@@ -251,8 +253,12 @@ func (r *UpstreamResolver) request(host string, n bootstrap.Network) (res *ipRes
// resolveAsync performs a single DNS lookup and sends the result to ch. It's
// intended to be used as a goroutine.
func (r *UpstreamResolver) resolveAsync(resCh chan<- any, host, network string) {
res, err := r.request(host, network)
func (r *UpstreamResolver) resolveAsync(
ctx context.Context,
resCh chan<- any,
host, network string,
) {
res, err := r.request(ctx, host, network)
if err != nil {
resCh <- err
} else {
upstream/resolver_test.go
+2 -2
View File
@@ -20,7 +20,7 @@ func TestNewUpstreamResolver(t *testing.T) {
ups := &dnsproxytest.Upstream{
OnAddress: func() (_ string) { panic(testutil.UnexpectedCall()) },
OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) },
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
OnExchange: func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(req)
resp.Answer = []dns.RR{&dns.A{
Hdr: dns.RR_Header{
@@ -133,7 +133,7 @@ func TestCachingResolver_staleness(t *testing.T) {
fqdn = "test.fully.qualified.name."
)
onExchange := func(req *dns.Msg) (resp *dns.Msg, err error) {
onExchange := func(_ context.Context, req *dns.Msg) (resp *dns.Msg, err error) {
resp = (&dns.Msg{}).SetReply(req)
hdr := dns.RR_Header{
upstream/upstream.go
+5 -5
View File
@@ -36,7 +36,7 @@ type Upstream interface {
// must not modify req as well as the caller must not modify it until the
// method returns. It shouldn't be called after closing. req must not be
// nil.
Exchange(req *dns.Msg) (resp *dns.Msg, err error)
Exchange(ctx context.Context, req *dns.Msg) (resp *dns.Msg, err error)
// Address returns the human-readable address of the upstream DNS resolver.
// It may differ from what was passed to [AddressToUpstream].
@@ -330,7 +330,7 @@ func addPort(u *url.URL, port uint16) {
// logBegin logs the start of DNS request resolution. It should be called right
// before dialing the connection to the upstream. n is the [network] that will
// be used to send the request.
func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
func logBegin(ctx context.Context, l *slog.Logger, addr string, n network, req *dns.Msg) {
var qtype dns.Type
var qname string
if len(req.Question) != 0 {
@@ -339,7 +339,7 @@ func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
}
l.DebugContext(
context.TODO(),
ctx,
"sending request",
"addr", addr,
"proto", n,
@@ -352,7 +352,7 @@ func logBegin(l *slog.Logger, addr string, n network, req *dns.Msg) {
// logFinish logs the end of DNS request resolution. It should be called right
// after receiving the response from the upstream or the failing action. n is
// the [network] that was used to send the request.
func logFinish(l *slog.Logger, addr string, n network, err error) {
func logFinish(ctx context.Context, l *slog.Logger, addr string, n network, err error) {
lvl := slog.LevelDebug
status := "ok"
@@ -364,7 +364,7 @@ func logFinish(l *slog.Logger, addr string, n network, err error) {
}
}
l.Log(context.TODO(), lvl, "response received", "addr", addr, "proto", n, "status", status)
l.Log(ctx, lvl, "response received", "addr", addr, "proto", n, "status", status)
}
// isTimeout returns true if err is a timeout error.
upstream/upstream_internal_test.go
+18 -10
View File
@@ -90,7 +90,9 @@ func TestUpstream_bootstrapTimeout(t *testing.T) {
req := createTestMessage()
start := time.Now()
_, rErr := u.Exchange(req)
ctx := testutil.ContextWithTimeout(t, testTimeout)
_, rErr := u.Exchange(ctx, req)
elapsed := time.Since(start)
// Require an error, since the bootstrap server cannot work.
@@ -229,7 +231,7 @@ func TestUpstreams(t *testing.T) {
require.NoErrorf(t, upsErr, "failed to generate upstream from address %s", test.address)
testutil.CleanupAndRequireSuccess(t, u.Close)
checkUpstream(t, u, test.address)
checkUpstream(t, u, test.address, testTimeout)
})
}
}
@@ -399,7 +401,7 @@ func TestUpstreamDoTBootstrap(t *testing.T) {
require.NoErrorf(t, err, "failed to generate upstream from address %s", tc.address)
testutil.CleanupAndRequireSuccess(t, u.Close)
checkUpstream(t, u, tc.address)
checkUpstream(t, u, tc.address, testTimeout)
})
}
}
@@ -431,6 +433,7 @@ func TestUpstreamsInvalidBootstrap(t *testing.T) {
}}
l := testLogger
const upstreamTimeout = 5 * time.Second
for _, tc := range upstreams {
t.Run(tc.address, func(t *testing.T) {
@@ -450,12 +453,12 @@ func TestUpstreamsInvalidBootstrap(t *testing.T) {
u, err := AddressToUpstream(tc.address, &Options{
Logger: l,
Bootstrap: rslv,
Timeout: testTimeout,
Timeout: upstreamTimeout,
})
require.NoErrorf(t, err, "failed to generate upstream from address %s", tc.address)
testutil.CleanupAndRequireSuccess(t, u.Close)
checkUpstream(t, u, tc.address)
checkUpstream(t, u, tc.address, upstreamTimeout)
})
}
@@ -527,7 +530,7 @@ func TestAddressToUpstream_StaticResolver(t *testing.T) {
testutil.CleanupAndRequireSuccess(t, u.Close)
assert.NotPanics(t, func() {
checkUpstream(t, u, tc.address)
checkUpstream(t, u, tc.address, testTimeout)
})
})
}
@@ -594,11 +597,13 @@ func TestAddPort(t *testing.T) {
}
// checkUpstream sends a test message to the upstream and checks the result.
func checkUpstream(tb testing.TB, u Upstream, addr string) {
func checkUpstream(tb testing.TB, u Upstream, addr string, timeout time.Duration) {
tb.Helper()
req := createTestMessage()
reply, err := u.Exchange(req)
ctx := testutil.ContextWithTimeout(tb, timeout)
reply, err := u.Exchange(ctx, req)
require.NoErrorf(tb, err, "couldn't talk to upstream %s", addr)
requireResponse(tb, req, reply)
@@ -606,7 +611,9 @@ func checkUpstream(tb testing.TB, u Upstream, addr string) {
// checkRaceCondition runs several goroutines in parallel and each of them calls
// checkUpstream several times.
func checkRaceCondition(u Upstream) {
func checkRaceCondition(tb testing.TB, u Upstream) {
tb.Helper()
wg := sync.WaitGroup{}
// The number of requests to run in every goroutine.
@@ -619,7 +626,8 @@ func checkRaceCondition(u Upstream) {
for range reqCount {
req := createTestMessage()
// Ignore exchange errors here, the point is to check for races.
_, _ = u.Exchange(req)
ctx := testutil.ContextWithTimeout(tb, testTimeout)
_, _ = u.Exchange(ctx, req)
}
}