From e6e0ae094404dd158e65c20500b64a7bb57d2bcf Mon Sep 17 00:00:00 2001 From: Maksim Kazantsev Date: Wed, 19 Aug 2026 15:52:03 +0300 Subject: [PATCH] AGDNS-4409 Imp upstream tests --- dnsproxytest/dnsproxytest.go | 3 + .../interface.go | 64 ++++- internal/dnsproxytest/dnsproxytest.go | 6 +- .../middleware/middleware_internal_test.go | 2 +- proxy/bogusnxdomain_internal_test.go | 17 +- proxy/cache_internal_test.go | 14 +- proxy/dns64_internal_test.go | 9 +- proxy/exchange_internal_test.go | 15 +- proxy/handler_internal_test.go | 2 +- proxy/pending_test.go | 4 +- proxy/proxy_internal_test.go | 236 +++++++++++------- proxy/proxy_test.go | 2 +- proxy/requesthandler_internal_test.go | 19 -- proxy/serverhttps_internal_test.go | 2 +- proxy/serverquic_internal_test.go | 4 +- proxy/stats_test.go | 2 +- ratelimit/ratelimit_test.go | 24 +- upstream/resolver_test.go | 2 +- 18 files changed, 255 insertions(+), 172 deletions(-) create mode 100644 dnsproxytest/dnsproxytest.go rename {internal/dnsproxytest => dnsproxytest}/interface.go (62%) delete mode 100644 proxy/requesthandler_internal_test.go diff --git a/dnsproxytest/dnsproxytest.go b/dnsproxytest/dnsproxytest.go new file mode 100644 index 00000000..d273623a --- /dev/null +++ b/dnsproxytest/dnsproxytest.go @@ -0,0 +1,3 @@ +// Package dnsproxytest provides test utilities and mock implementations +// for the dnsproxy module interfaces. +package dnsproxytest diff --git a/internal/dnsproxytest/interface.go b/dnsproxytest/interface.go similarity index 62% rename from internal/dnsproxytest/interface.go rename to dnsproxytest/interface.go index 359c180e..5f700abb 100644 --- a/internal/dnsproxytest/interface.go +++ b/dnsproxytest/interface.go @@ -1,15 +1,17 @@ package dnsproxytest import ( - "github.com/AdguardTeam/dnsproxy/internal/dnsmsg" + "context" + + "github.com/AdguardTeam/dnsproxy/proxy" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/testutil" "github.com/miekg/dns" + "github.com/quic-go/quic-go" + "github.com/quic-go/quic-go/qlogwriter" ) // Upstream is a mock [upstream.Upstream] implementation for tests. -// -// TODO(e.burkov): Move to golibs. type Upstream struct { OnAddress func() (addr string) OnExchange func(req *dns.Msg) (resp *dns.Msg, err error) @@ -34,7 +36,55 @@ func (u *Upstream) Close() (err error) { return u.OnClose() } -// MessageConstructor is a mock [dnsmsg.MessageConstructor] implementation for +// Handler is a mock [proxy.Handler] implementation for tests. +type Handler struct { + OnHandle func(ctx context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) +} + +// type check +var _ proxy.Handler = (*Handler)(nil) + +// ServeDNS implements the [proxy.Handler] interface for *Handler. +func (h *Handler) ServeDNS(ctx context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) { + return h.OnHandle(ctx, p, dctx) +} + +// Middleware is a mock [proxy.Middleware] implementation for tests. +type Middleware struct { + OnWrap func(h proxy.Handler) (wrapped proxy.Handler) +} + +// type check +var _ proxy.Middleware = (*Middleware)(nil) + +// Wrap implements the [proxy.Middleware] interface for *Middleware. +func (m *Middleware) Wrap(h proxy.Handler) (wrapped proxy.Handler) { + return m.OnWrap(h) +} + +// QUICTracer is a mock [upstream.QUICTracer] implementation for tests. +type QUICTracer struct { + OnTraceForConnection func( + ctx context.Context, + isClient bool, + connID quic.ConnectionID, + ) (trace qlogwriter.Trace) +} + +// type check +var _ upstream.QUICTracer = (*QUICTracer)(nil) + +// TraceForConnection implements the [upstream.QUICTracer] interface for +// *QUICTracer. +func (t *QUICTracer) TraceForConnection( + ctx context.Context, + isClient bool, + connID quic.ConnectionID, +) (trace qlogwriter.Trace) { + return t.OnTraceForConnection(ctx, isClient, connID) +} + +// MessageConstructor is a mock [proxy.MessageConstructor] implementation for // tests. type MessageConstructor struct { OnNewMsgNXDOMAIN func(req *dns.Msg) (resp *dns.Msg) @@ -44,8 +94,8 @@ type MessageConstructor struct { OnNewMsgFORMERR func(req *dns.Msg) (resp *dns.Msg) } -// NewMessageConstructor creates a new *TestMessageConstructor with all it's -// methods set to panic. +// NewMessageConstructor creates a new *MessageConstructor with all its methods +// set to panic. func NewMessageConstructor() (c *MessageConstructor) { return &MessageConstructor{ OnNewMsgNXDOMAIN: func(req *dns.Msg) (_ *dns.Msg) { @@ -67,7 +117,7 @@ func NewMessageConstructor() (c *MessageConstructor) { } // type check -var _ dnsmsg.MessageConstructor = (*MessageConstructor)(nil) +var _ proxy.MessageConstructor = (*MessageConstructor)(nil) // NewMsgNXDOMAIN implements the [proxy.MessageConstructor] interface for // *MessageConstructor. diff --git a/internal/dnsproxytest/dnsproxytest.go b/internal/dnsproxytest/dnsproxytest.go index 70e440d4..643c107d 100644 --- a/internal/dnsproxytest/dnsproxytest.go +++ b/internal/dnsproxytest/dnsproxytest.go @@ -1,5 +1,5 @@ -// Package dnsproxytest provides a set of test utilities for the dnsproxy -// module. +// Package dnsproxytest provides test utilities and mock implementations +// for the dnsproxy module interfaces. package dnsproxytest import ( @@ -15,6 +15,8 @@ import ( // that can be used for testing. Note that there is theoretically a TOCTTOU // race here: the port may be reoccupied between the time it is released and the // time the caller binds to it. +// +// TODO(m.kazantsev): Move to the top-level dnsproxytest package. func NewFreePort(tb testing.TB) (p uint) { tb.Helper() diff --git a/internal/middleware/middleware_internal_test.go b/internal/middleware/middleware_internal_test.go index 944465ed..a50055ef 100644 --- a/internal/middleware/middleware_internal_test.go +++ b/internal/middleware/middleware_internal_test.go @@ -8,8 +8,8 @@ import ( "testing" "time" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/internal/dnsmsg" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/golibs/errors" "github.com/AdguardTeam/golibs/logutil/slogutil" "github.com/AdguardTeam/golibs/netutil" diff --git a/proxy/bogusnxdomain_internal_test.go b/proxy/bogusnxdomain_internal_test.go index f1d52283..fa0b650f 100644 --- a/proxy/bogusnxdomain_internal_test.go +++ b/proxy/bogusnxdomain_internal_test.go @@ -14,11 +14,19 @@ import ( ) func TestProxy_IsBogusNXDomain(t *testing.T) { + var ans []dns.RR + + onExchange := newECSReplyHandler(&ans, nil, nil) + u := newTestECSUpstream(onExchange) + + upsConf := newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr) + upsConf.Upstreams = []upstream.Upstream{u} + prx := mustNew(t, &Config{ Logger: testLogger, UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)}, TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)}, - UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr), + UpstreamConfig: upsConf, TrustedProxies: defaultTrustedProxies, CacheEnabled: true, BogusNXDomain: []netip.Prefix{ @@ -77,9 +85,6 @@ func TestProxy_IsBogusNXDomain(t *testing.T) { wantRcode: dns.RcodeSuccess, }} - u := testUpstream{} - prx.upstreamConf.Upstreams = []upstream.Upstream{&u} - servicetest.RequireRun(t, prx, testTimeout) d := &DNSContext{ @@ -87,9 +92,9 @@ func TestProxy_IsBogusNXDomain(t *testing.T) { } for _, tc := range testCases { - u.ans = tc.ans - t.Run(tc.name, func(t *testing.T) { + ans = tc.ans + err := prx.Resolve(testutil.ContextWithTimeout(t, defaultTimeout), d) require.NoError(t, err) require.NotNil(t, d.Res) diff --git a/proxy/cache_internal_test.go b/proxy/cache_internal_test.go index bb2cda1b..8a8a272d 100644 --- a/proxy/cache_internal_test.go +++ b/proxy/cache_internal_test.go @@ -9,7 +9,6 @@ import ( "testing" "time" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/netutil" "github.com/AdguardTeam/golibs/testutil" @@ -28,7 +27,7 @@ const testUpsAddr = "https://upstream.address" // upstreamWithAddr is a [dnsproxytest.Upstream] that is only expected to be // used to get its address. -var upstreamWithAddr = &dnsproxytest.Upstream{ +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 }, @@ -361,14 +360,17 @@ func TestCacheExpiration(t *testing.T) { } func TestCacheExpirationWithTTLOverride(t *testing.T) { - u := testUpstream{} + var ans []dns.RR + + onExchange := newECSReplyHandler(&ans, nil, nil) + u := newTestECSUpstream(onExchange) dnsProxy := mustNew(t, &Config{ Logger: testLogger, UDPListenAddr: []*net.UDPAddr{net.UDPAddrFromAddrPort(localhostAnyPort)}, TCPListenAddr: []*net.TCPAddr{net.TCPAddrFromAddrPort(localhostAnyPort)}, UpstreamConfig: &UpstreamConfig{ - Upstreams: []upstream.Upstream{&u}, + Upstreams: []upstream.Upstream{u}, }, TrustedProxies: defaultTrustedProxies, CacheEnabled: true, @@ -385,7 +387,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) { d.Req = newHostTestMessage("host") d.Addr = netip.AddrPort{} - u.ans = []dns.RR{&dns.A{ + ans = []dns.RR{&dns.A{ Hdr: dns.RR_Header{ Rrtype: dns.TypeA, Name: "host.", @@ -409,7 +411,7 @@ func TestCacheExpirationWithTTLOverride(t *testing.T) { d.Req = newHostTestMessage("host2") d.Addr = netip.AddrPort{} - u.ans = []dns.RR{&dns.A{ + ans = []dns.RR{&dns.A{ Hdr: dns.RR_Header{ Rrtype: dns.TypeA, Name: "host2.", diff --git a/proxy/dns64_internal_test.go b/proxy/dns64_internal_test.go index 88823a83..7d05159b 100644 --- a/proxy/dns64_internal_test.go +++ b/proxy/dns64_internal_test.go @@ -6,7 +6,6 @@ import ( "sync" "testing" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/netutil" "github.com/AdguardTeam/golibs/testutil" @@ -20,7 +19,7 @@ const ipv4OnlyFqdn = "ipv4.only." func TestDNS64Race(t *testing.T) { ans := newRR(t, ipv4OnlyFqdn, dns.TypeA, 3600, net.ParseIP("1.2.3.4")) - ups := &dnsproxytest.Upstream{ + ups := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { resp = (&dns.Msg{}).SetReply(req) if req.Question[0].Qtype == dns.TypeA { @@ -32,7 +31,7 @@ func TestDNS64Race(t *testing.T) { OnAddress: func() (addr string) { return "fake.address" }, OnClose: func() (err error) { return nil }, } - localUps := &dnsproxytest.Upstream{ + 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 }, @@ -180,7 +179,7 @@ func TestProxy_Resolve_dns64(t *testing.T) { pt := testutil.PanicT{} newUps := func(answers answerMap) (u upstream.Upstream) { - return &dnsproxytest.Upstream{ + return &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { q := req.Question[0] require.Contains(pt, answers, q.Qtype) @@ -200,7 +199,7 @@ func TestProxy_Resolve_dns64(t *testing.T) { } localRR := newRR(t, ptr64Domain, dns.TypePTR, 3600, domainPointed) - localUps := &dnsproxytest.Upstream{ + localUps := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { require.Equal(pt, req.Question[0].Name, ptr64Domain) resp = (&dns.Msg{}).SetReply(req) diff --git a/proxy/exchange_internal_test.go b/proxy/exchange_internal_test.go index 35035225..d2df6ef3 100644 --- a/proxy/exchange_internal_test.go +++ b/proxy/exchange_internal_test.go @@ -8,7 +8,6 @@ import ( "testing" "time" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/netutil" "github.com/AdguardTeam/golibs/testutil" @@ -24,7 +23,7 @@ import ( func newUpstreamWithErrorRate(rate uint, name string) (u upstream.Upstream) { var n uint - return &dnsproxytest.Upstream{ + return &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { n++ if n%rate == 0 { @@ -87,7 +86,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) { }, } - fastUps := &dnsproxytest.Upstream{ + fastUps := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { currentNow = zeroTime.Add(testRTT / 100) @@ -96,7 +95,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) { OnAddress: func() (addr string) { return "fast" }, OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) }, } - slowerUps := &dnsproxytest.Upstream{ + slowerUps := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { currentNow = zeroTime.Add(testRTT / 10) @@ -105,7 +104,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) { OnAddress: func() (addr string) { return "slower" }, OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) }, } - slowestUps := &dnsproxytest.Upstream{ + slowestUps := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { currentNow = zeroTime.Add(testRTT / 2) @@ -115,12 +114,12 @@ func TestProxy_Exchange_loadBalance(t *testing.T) { OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) }, } - err1Ups := &dnsproxytest.Upstream{ + 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()) }, } - err2Ups := &dnsproxytest.Upstream{ + 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()) }, @@ -128,7 +127,7 @@ func TestProxy_Exchange_loadBalance(t *testing.T) { singleError := &sync.Once{} // fastestUps responds with an error on the first request. - fastestUps := &dnsproxytest.Upstream{ + fastestUps := &testUpstream{ OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) { singleError.Do(func() { err = assert.AnError }) currentNow = zeroTime.Add(testRTT / 200) diff --git a/proxy/handler_internal_test.go b/proxy/handler_internal_test.go index 82d56a18..7bab265b 100644 --- a/proxy/handler_internal_test.go +++ b/proxy/handler_internal_test.go @@ -17,7 +17,7 @@ func TestFilteringHandler(t *testing.T) { m := &sync.RWMutex{} blockResponse := false - reqHandler := &TestHandler{ + reqHandler := &testHandler{ OnHandle: func(ctx context.Context, p *Proxy, d *DNSContext) (err error) { m.Lock() defer m.Unlock() diff --git a/proxy/pending_test.go b/proxy/pending_test.go index 52deb272..10ae8c5d 100644 --- a/proxy/pending_test.go +++ b/proxy/pending_test.go @@ -8,7 +8,7 @@ import ( "testing" "time" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/proxy" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/netutil" @@ -79,7 +79,7 @@ func TestPendingRequests(t *testing.T) { workloadWG := &sync.WaitGroup{} workloadWG.Add(reqsNum) - reqHandler := &proxy.TestHandler{ + reqHandler := &dnsproxytest.Handler{ OnHandle: func(ctx context.Context, p *proxy.Proxy, d *proxy.DNSContext) (err error) { workloadWG.Done() diff --git a/proxy/proxy_internal_test.go b/proxy/proxy_internal_test.go index bbd5eba3..e82558d1 100644 --- a/proxy/proxy_internal_test.go +++ b/proxy/proxy_internal_test.go @@ -1,6 +1,7 @@ package proxy import ( + "context" "crypto/rand" "crypto/rsa" "crypto/tls" @@ -15,7 +16,6 @@ import ( "testing" "time" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/upstream" glcache "github.com/AdguardTeam/golibs/cache" "github.com/AdguardTeam/golibs/contextutil" @@ -189,48 +189,6 @@ func firstIP(resp *dns.Msg) (ip net.IP) { return nil } -type testUpstream struct { - ans []dns.RR - - ecsIP net.IP - ecsReqIP net.IP - ecsReqMask int -} - -// type check -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) { - resp = &dns.Msg{} - resp.SetReply(m) - - if u.ans != nil { - resp.Answer = append(resp.Answer, u.ans...) - } - - ecs, _ := ecsFromMsg(m) - if ecs != nil { - u.ecsReqIP = ecs.IP - u.ecsReqMask, _ = ecs.Mask.Size() - } - if u.ecsIP != nil { - setECS(resp, u.ecsIP, 24) - } - - return resp, nil -} - -// Address implements the upstream.Upstream interface for *testUpstream. -func (u *testUpstream) Address() (addr string) { - return "" -} - -// Close implements the upstream.Upstream interface for *testUpstream. -func (u *testUpstream) Close() (err error) { - return nil -} - // newTestUpstreamConfigWithBoot creates a new UpstreamConfig with upstream // addresses and a bootstrapped resolver. func newTestUpstreamConfigWithBoot( @@ -462,7 +420,7 @@ func TestProxy_Resolve_dnssecCache(t *testing.T) { Signature: "c29tZSBycnNpZyByZWxhdGVkIHN0dWZm", } - u := &dnsproxytest.Upstream{ + u := &testUpstream{ OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) { resp = (&dns.Msg{}).SetReply(m) @@ -897,7 +855,7 @@ func TestResponseInRequest(t *testing.T) { func TestProxy_ReplyFromUpstream_badResponse(t *testing.T) { dnsProxy := mustStartDefaultProxy(t) - u := &dnsproxytest.Upstream{ + u := &testUpstream{ OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) { resp = (&dns.Msg{}).SetReply(m) resp.Answer = append(resp.Answer, &dns.A{ @@ -941,16 +899,17 @@ func TestExchangeCustomUpstreamConfig(t *testing.T) { prx := mustStartDefaultProxy(t) ansIP := net.IP{4, 3, 2, 1} - u := &testUpstream{ - ans: []dns.RR{&dns.A{ - Hdr: dns.RR_Header{ - Rrtype: dns.TypeA, - Name: "host.", - Ttl: 60, - }, - A: ansIP, - }}, - } + ans := []dns.RR{&dns.A{ + Hdr: dns.RR_Header{ + Rrtype: dns.TypeA, + Name: "host.", + Ttl: 60, + }, + A: ansIP, + }} + + onExchange := newECSReplyHandler(&ans, nil, nil) + u := newTestECSUpstream(onExchange) d := &DNSContext{ CustomUpstreamConfig: NewCustomUpstreamConfig( @@ -1002,7 +961,7 @@ func TestExchangeCustomUpstreamConfigCache(t *testing.T) { return resp, nil } - u := &dnsproxytest.Upstream{ + u := &testUpstream{ OnExchange: exchangeFunc, OnAddress: func() (addr string) { return "stub" }, OnClose: func() (_ error) { panic(testutil.UnexpectedCall()) }, @@ -1091,13 +1050,20 @@ func TestECSProxy(t *testing.T) { ip4323 = net.IP{4, 3, 2, 3} ) - u := &testUpstream{ - ans: []dns.RR{&dns.A{ - Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60}, - A: ip4321, - }}, - ecsIP: ip1230, - } + var ( + ans []dns.RR + ecsIP net.IP + ecsReqIP net.IP + ) + + onExchange := newECSReplyHandler(&ans, &ecsIP, &ecsReqIP) + u := newTestECSUpstream(onExchange) + + ans = []dns.RR{&dns.A{ + Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60}, + A: ip4321, + }} + ecsIP = ip1230 prx := mustNew(t, &Config{ Logger: testLogger, @@ -1125,7 +1091,7 @@ func TestECSProxy(t *testing.T) { require.NoError(t, err) assert.Equal(t, net.IP{4, 3, 2, 1}, firstIP(d.Res)) - assert.Equal(t, ip1230, u.ecsReqIP) + assert.Equal(t, ip1230, ecsReqIP) }) t.Run("serve_subnet_cache", func(t *testing.T) { @@ -1133,14 +1099,15 @@ func TestECSProxy(t *testing.T) { Req: newHostTestMessage("host"), Addr: netip.MustParseAddrPort("1.2.3.1:1234"), } - u.ans, u.ecsIP, u.ecsReqIP = nil, nil, nil + ans, ecsIP = nil, nil + ecsReqIP = nil ctx := testutil.ContextWithTimeout(t, defaultTimeout) err := prx.Resolve(ctx, d) require.NoError(t, err) assert.Equal(t, ip4321, firstIP(d.Res)) - assert.Nil(t, u.ecsReqIP) + assert.Nil(t, ecsReqIP) }) t.Run("another_subnet", func(t *testing.T) { @@ -1148,18 +1115,18 @@ func TestECSProxy(t *testing.T) { Req: newHostTestMessage("host"), Addr: netip.MustParseAddrPort("2.2.3.0:1234"), } - u.ans = []dns.RR{&dns.A{ + ans = []dns.RR{&dns.A{ Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60}, A: ip4322, }} - u.ecsIP = ip2230 + ecsIP = ip2230 ctx := testutil.ContextWithTimeout(t, defaultTimeout) err := prx.Resolve(ctx, d) require.NoError(t, err) assert.Equal(t, ip4322, firstIP(d.Res)) - assert.Equal(t, ip2230, u.ecsReqIP) + assert.Equal(t, ip2230, ecsReqIP) }) t.Run("cache_general", func(t *testing.T) { @@ -1167,18 +1134,19 @@ func TestECSProxy(t *testing.T) { Req: newHostTestMessage("host"), Addr: netip.MustParseAddrPort("127.0.0.1:1234"), } - u.ans = []dns.RR{&dns.A{ + ans = []dns.RR{&dns.A{ Hdr: dns.RR_Header{Rrtype: dns.TypeA, Name: "host.", Ttl: 60}, A: ip4323, }} - u.ecsIP, u.ecsReqIP = nil, nil + ecsIP = nil + ecsReqIP = nil ctx := testutil.ContextWithTimeout(t, defaultTimeout) err := prx.Resolve(ctx, d) require.NoError(t, err) assert.Equal(t, ip4323, firstIP(d.Res)) - assert.Nil(t, u.ecsReqIP) + assert.Nil(t, ecsReqIP) }) t.Run("serve_general_cache", func(t *testing.T) { @@ -1186,30 +1154,38 @@ func TestECSProxy(t *testing.T) { Req: newHostTestMessage("host"), Addr: netip.MustParseAddrPort("127.0.0.2:1234"), } - u.ans, u.ecsIP, u.ecsReqIP = nil, nil, nil + ans, ecsIP = nil, nil + ecsReqIP = nil ctx := testutil.ContextWithTimeout(t, defaultTimeout) err := prx.Resolve(ctx, d) require.NoError(t, err) assert.Equal(t, ip4323, firstIP(d.Res)) - assert.Nil(t, u.ecsReqIP) + assert.Nil(t, ecsReqIP) }) } func TestECSProxyCacheMinMaxTTL(t *testing.T) { clientIP := net.IP{1, 2, 3, 0} - u := &testUpstream{ - ans: []dns.RR{&dns.A{ - Hdr: dns.RR_Header{ - Rrtype: dns.TypeA, - Name: "host.", - Ttl: 10, - }, - A: net.IP{4, 3, 2, 1}, - }}, - ecsIP: clientIP, - } + + var ( + ans []dns.RR + ecsIP net.IP + ) + + onExchange := newECSReplyHandler(&ans, &ecsIP, nil) + u := newTestECSUpstream(onExchange) + + ans = []dns.RR{&dns.A{ + Hdr: dns.RR_Header{ + Rrtype: dns.TypeA, + Name: "host.", + Ttl: 10, + }, + A: net.IP{4, 3, 2, 1}, + }} + ecsIP = clientIP prx := mustNew(t, &Config{ Logger: testLogger, @@ -1249,7 +1225,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) { clientIP = net.IP{1, 2, 4, 0} d.Req = newHostTestMessage("host") d.Addr = netip.MustParseAddrPort("1.2.4.0:1234") - u.ans = []dns.RR{&dns.A{ + ans = []dns.RR{&dns.A{ Hdr: dns.RR_Header{ Rrtype: dns.TypeA, Name: "host.", @@ -1257,7 +1233,7 @@ func TestECSProxyCacheMinMaxTTL(t *testing.T) { }, A: net.IP{4, 3, 2, 1}, }} - u.ecsIP = clientIP + ecsIP = clientIP err = prx.Resolve(ctx, d) require.NoError(t, err) @@ -1406,7 +1382,7 @@ func TestProxy_validateRequest(t *testing.T) { netip.MustParsePrefix("203.0.113.0/8"), } - ups := &dnsproxytest.Upstream{ + ups := &testUpstream{ OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) { resp = &dns.Msg{} resp.SetReply(m) @@ -1514,3 +1490,87 @@ func TestProxy_validateRequest(t *testing.T) { }) } } + +// testHandler is a mock request handler implementation to simplify +// testing. +// +// TODO(m.kazantsev): Use [dnsproxytest.Handler]. +type testHandler struct { + OnHandle func(ctx context.Context, p *Proxy, dctx *DNSContext) (err error) +} + +// type check +var _ Handler = (*testHandler)(nil) + +// ServeDNS implements the [Handler] interface for *testHandler. +func (h *testHandler) ServeDNS(ctx context.Context, p *Proxy, dctx *DNSContext) (err error) { + return h.OnHandle(ctx, p, dctx) +} + +// testUpstream is a mock upstream implementation to simplify testing. +// +// TODO(m.kazantsev): Use [dnsproxytest.Upstream]. +type testUpstream struct { + OnAddress func() (addr string) + OnExchange func(req *dns.Msg) (resp *dns.Msg, err error) + OnClose func() (err error) +} + +// newTestECSUpstream creates a new test upstream with the given OnExchange +// handler. +func newTestECSUpstream(onExchange func(req *dns.Msg) (resp *dns.Msg, err error)) (u *testUpstream) { + return &testUpstream{ + OnAddress: func() string { return "" }, + OnClose: func() error { return nil }, + OnExchange: onExchange, + } +} + +// newECSReplyHandler creates an OnExchange handler that builds a DNS reply +// with optional answer records and ECS support. +// +// 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) { + resp = (&dns.Msg{}).SetReply(m) + + if ans != nil && *ans != nil { + resp.Answer = append(resp.Answer, *ans...) + } + + if ecsIP != nil && *ecsIP != nil { + setECS(resp, *ecsIP, 24) + } + + if ecsReqIP == nil { + return resp, nil + } + + ecs, _ := ecsFromMsg(m) + if ecs != nil { + *ecsReqIP = ecs.IP + } + + return resp, nil + } +} + +// type check +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) +} + +// Address implements the upstream.Upstream interface for *testUpstream. +func (u *testUpstream) Address() (addr string) { + return u.OnAddress() +} + +// Close implements the upstream.Upstream interface for *testUpstream. +func (u *testUpstream) Close() (err error) { + return u.OnClose() +} diff --git a/proxy/proxy_test.go b/proxy/proxy_test.go index e5942b52..72cda99a 100644 --- a/proxy/proxy_test.go +++ b/proxy/proxy_test.go @@ -7,8 +7,8 @@ import ( "testing" "time" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/internal/bootstrap" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/proxy" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/logutil/slogutil" diff --git a/proxy/requesthandler_internal_test.go b/proxy/requesthandler_internal_test.go deleted file mode 100644 index 3dbe1020..00000000 --- a/proxy/requesthandler_internal_test.go +++ /dev/null @@ -1,19 +0,0 @@ -package proxy - -import "context" - -// TestHandler is a mock request handler implementation to simplify -// testing. -// -// TODO(d.kolyshev): Move to internal/dnsproxytest. -type TestHandler struct { - OnHandle func(ctx context.Context, p *Proxy, dctx *DNSContext) (err error) -} - -// type check -var _ Handler = (*TestHandler)(nil) - -// ServeDNS implements the [Handler] interface for *TestHandler. -func (h *TestHandler) ServeDNS(ctx context.Context, p *Proxy, dctx *DNSContext) (err error) { - return h.OnHandle(ctx, p, dctx) -} diff --git a/proxy/serverhttps_internal_test.go b/proxy/serverhttps_internal_test.go index 2048ebc2..90c73656 100644 --- a/proxy/serverhttps_internal_test.go +++ b/proxy/serverhttps_internal_test.go @@ -77,7 +77,7 @@ func TestProxy_trustedProxies(t *testing.T) { doRequest := func(t *testing.T, addr, expectedClientIP netip.Addr) { var gotAddr netip.Addr - reqHandler := &TestHandler{ + reqHandler := &testHandler{ OnHandle: func(ctx context.Context, p *Proxy, d *DNSContext) (err error) { gotAddr = d.Addr.Addr() diff --git a/proxy/serverquic_internal_test.go b/proxy/serverquic_internal_test.go index ed1155bb..0ebc607b 100644 --- a/proxy/serverquic_internal_test.go +++ b/proxy/serverquic_internal_test.go @@ -83,7 +83,7 @@ func TestProxy_quic(t *testing.T) { } func TestProxy_quicLargePackets(t *testing.T) { - reqHandler := &TestHandler{ + reqHandler := &testHandler{ OnHandle: func(_ context.Context, _ *Proxy, d *DNSContext) (err error) { d.Res = newTestResponse(d) @@ -146,7 +146,7 @@ func TestProxy_quicTruncatedRequest(t *testing.T) { TLSConfig: serverConfig, UpstreamConfig: newTestUpstreamConfig(t, defaultTimeout, testDefaultUpstreamAddr), TrustedProxies: defaultTrustedProxies, - RequestHandler: &TestHandler{ + RequestHandler: &testHandler{ OnHandle: func(ctx context.Context, p *Proxy, d *DNSContext) (_ error) { panic(testutil.UnexpectedCall(ctx, p, d)) }, diff --git a/proxy/stats_test.go b/proxy/stats_test.go index c1478963..8c1385d2 100644 --- a/proxy/stats_test.go +++ b/proxy/stats_test.go @@ -5,7 +5,7 @@ import ( "net/netip" "testing" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/proxy" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/errors" diff --git a/ratelimit/ratelimit_test.go b/ratelimit/ratelimit_test.go index 34353a71..0214101c 100644 --- a/ratelimit/ratelimit_test.go +++ b/ratelimit/ratelimit_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/proxy" "github.com/AdguardTeam/dnsproxy/ratelimit" "github.com/AdguardTeam/golibs/logutil/slogutil" @@ -70,7 +71,7 @@ func TestMiddleware_Wrap(t *testing.T) { for _, tc := range testCases { called := 0 - mock := &TestHandler{ + mock := &dnsproxytest.Handler{ OnHandle: func(_ context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) { called++ @@ -117,7 +118,7 @@ func TestMiddleware_Wrap_allowlist(t *testing.T) { } called := 0 - mock := &TestHandler{ + mock := &dnsproxytest.Handler{ OnHandle: func(_ context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) { called++ @@ -162,22 +163,3 @@ func TestMiddleware_Wrap_allowlist(t *testing.T) { assert.Equal(t, 3, called) }) } - -// TestHandler is a mock request middleware implementation to simplify testing. -// -// TODO(d.kolyshev): Move to internal/dnsproxytest. -type TestHandler struct { - OnHandle func(ctx context.Context, p *proxy.Proxy, dctx *proxy.DNSContext) (err error) -} - -// type check -var _ proxy.Handler = (*TestHandler)(nil) - -// ServeDNS implements the [Handler] interface for *TestHandler. -func (h *TestHandler) ServeDNS( - ctx context.Context, - p *proxy.Proxy, - dctx *proxy.DNSContext, -) (err error) { - return h.OnHandle(ctx, p, dctx) -} diff --git a/upstream/resolver_test.go b/upstream/resolver_test.go index d7da868e..4d13139a 100644 --- a/upstream/resolver_test.go +++ b/upstream/resolver_test.go @@ -6,8 +6,8 @@ import ( "testing" "time" + "github.com/AdguardTeam/dnsproxy/dnsproxytest" "github.com/AdguardTeam/dnsproxy/internal/bootstrap" - "github.com/AdguardTeam/dnsproxy/internal/dnsproxytest" "github.com/AdguardTeam/dnsproxy/upstream" "github.com/AdguardTeam/golibs/errors" "github.com/AdguardTeam/golibs/testutil"