AGDNS-4409 Imp upstream tests

go.mod

@@ -4,7 +4,7 @@ go 1.26.6
 
 require (
 	github.com/AdguardTeam/dnscrypt v0.0.2
-	github.com/AdguardTeam/dnsproxy v0.84.0
+	github.com/AdguardTeam/dnsproxy v0.84.1
 	github.com/AdguardTeam/golibs v0.35.15
 	github.com/AdguardTeam/urlfilter v0.23.4
 	github.com/NYTimes/gziphandler v1.1.1

go.sum

@@ -8,6 +8,8 @@ github.com/AdguardTeam/dnscrypt v0.0.2 h1:kIFCyjYofKHug7RwPGOjVga81P0/zOQdbUN+7p
 github.com/AdguardTeam/dnscrypt v0.0.2/go.mod h1:qCFs51rLfNzEDZqb6nz1tocVLnEearJ1zng6O3I6ecA=
 github.com/AdguardTeam/dnsproxy v0.84.0 h1:bcTORgQL0L5cwcTz9f+iB0JPZVbCHlCv7phso9y3lrU=
 github.com/AdguardTeam/dnsproxy v0.84.0/go.mod h1:nZlktYebWcLO613q1j2N0bZ5BEnHPuqXAyW6O9vIhDU=
+github.com/AdguardTeam/dnsproxy v0.84.1 h1:Ge1UdyVNswKn0BKrV5qp+o/V0tF0N6wYcxaGYujDeEg=
+github.com/AdguardTeam/dnsproxy v0.84.1/go.mod h1:nZlktYebWcLO613q1j2N0bZ5BEnHPuqXAyW6O9vIhDU=
 github.com/AdguardTeam/golibs v0.35.15 h1:KjksvvzlQqxdNvrz+O3FgMbZKZAQdcoCkgBz0mR7wFQ=
 github.com/AdguardTeam/golibs v0.35.15/go.mod h1:lScOIavWuRFL0vls26jKfQXaITP7OKO1qw75WYp5nYk=
 github.com/AdguardTeam/urlfilter v0.23.4 h1:3cwt5xj7lpK2t3sWtr9WI8mmsKP+RtpkdnEy7lrn8cg=

internal/aghtest/interface.go

@@ -16,9 +16,7 @@ import (
 	nextagh "github.com/AdguardTeam/AdGuardHome/internal/next/agh"
 	"github.com/AdguardTeam/AdGuardHome/internal/rdns"
 	"github.com/AdguardTeam/AdGuardHome/internal/whois"
-	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/testutil"
-	"github.com/miekg/dns"
 )
 
 // FSWatcher is a fake [aghos.FSWatcher] implementation for tests.
@@ -149,34 +147,6 @@ func (e *Exchanger) Exchange(
 	return e.OnExchange(ctx, ip)
 }
 
-// UpstreamMock is a fake [upstream.Upstream] implementation for tests.
-//
-// TODO(a.garipov): Replace with all uses of Upstream with UpstreamMock and
-// rename it to just Upstream.
-type UpstreamMock struct {
-	OnAddress  func() (addr string)
-	OnExchange func(req *dns.Msg) (resp *dns.Msg, err error)
-	OnClose    func() (err error)
-}
-
-// type check
-var _ upstream.Upstream = (*UpstreamMock)(nil)
-
-// Address implements the [upstream.Upstream] interface for *UpstreamMock.
-func (u *UpstreamMock) Address() (addr string) {
-	return u.OnAddress()
-}
-
-// Exchange implements the [upstream.Upstream] interface for *UpstreamMock.
-func (u *UpstreamMock) Exchange(req *dns.Msg) (resp *dns.Msg, err error) {
-	return u.OnExchange(req)
-}
-
-// Close implements the [upstream.Upstream] interface for *UpstreamMock.
-func (u *UpstreamMock) Close() (err error) {
-	return u.OnClose()
-}
-
 // ConfigModifier is a fake [agh.ConfigModifier] implementation for tests.
 type ConfigModifier struct {
 	OnApply func(ctx context.Context)

internal/aghtest/upstream.go

@@ -7,78 +7,86 @@ import (
 	"net"
 	"net/netip"
 	"strings"
+	"testing"
 
-	"github.com/AdguardTeam/dnsproxy/upstream"
+	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
 	"github.com/AdguardTeam/golibs/errors"
+	"github.com/AdguardTeam/golibs/testutil"
 	"github.com/miekg/dns"
 )
 
-// Additional Upstream Testing Utilities
+// NewExchangingUpstream returns a new test upstream that responds with the
+// provided CNAME, A, and AAAA records.
+func NewExchangingUpstream(
+	tb testing.TB,
+	cname map[string][]string,
+	ipv4 map[string][]net.IP,
+	ipv6 map[string][]net.IP,
+) (ups *dnsproxytest.Upstream) {
+	tb.Helper()
 
-// Upstream is a mock implementation of upstream.Upstream.
-//
-// TODO(a.garipov): Replace with UpstreamMock and rename it to just Upstream.
-type Upstream struct {
-	// CName is a map of hostname to canonical name.
-	CName map[string][]string
-	// IPv4 is a map of hostname to IPv4.
-	IPv4 map[string][]net.IP
-	// IPv6 is a map of hostname to IPv6.
-	IPv6 map[string][]net.IP
-}
+	ups = &dnsproxytest.Upstream{
+		OnAddress: func() (addr string) { return "upstream.example" },
+		OnClose:   func() (err error) { return nil },
+		OnExchange: func(m *dns.Msg) (resp *dns.Msg, err error) {
+			resp = new(dns.Msg).SetReply(m)
 
-var _ upstream.Upstream = (*Upstream)(nil)
+			if len(m.Question) == 0 {
+				return nil, fmt.Errorf("question should not be empty")
+			}
 
-// Exchange implements the [upstream.Upstream] interface for *Upstream.
-//
-// TODO(a.garipov): Split further into handlers.
-func (u *Upstream) Exchange(m *dns.Msg) (resp *dns.Msg, err error) {
-	resp = new(dns.Msg).SetReply(m)
+			q := m.Question[0]
 
-	if len(m.Question) == 0 {
-		return nil, fmt.Errorf("question should not be empty")
+			resp.Answer = append(resp.Answer, cnameAnswers(q.Name, cname)...)
+			resp.Answer = append(resp.Answer, ipAnswers(q.Name, q.Qtype, ipv4, ipv6)...)
+
+			if len(resp.Answer) == 0 {
+				resp.SetRcode(m, dns.RcodeNameError)
+			}
+
+			return resp, nil
+		},
 	}
 
-	q := m.Question[0]
-	name := q.Name
-	for _, cname := range u.CName[name] {
-		resp.Answer = append(resp.Answer, &dns.CNAME{
+	return ups
+}
+
+// cnameAnswers returns CNAME records for name from the cname map.
+func cnameAnswers(name string, cname map[string][]string) (answers []dns.RR) {
+	for _, t := range cname[name] {
+		answers = append(answers, &dns.CNAME{
 			Hdr:    dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME},
-			Target: cname,
+			Target: t,
 		})
 	}
 
-	qtype := q.Qtype
-	hdr := dns.RR_Header{
-		Name:   name,
-		Rrtype: qtype,
-	}
+	return answers
+}
 
+// ipAnswers returns A or AAAA records for name from the corresponding IP map,
+// depending on qtype.
+func ipAnswers(name string, qtype uint16, ipv4, ipv6 map[string][]net.IP) (ans []dns.RR) {
+	var ips []net.IP
 	switch qtype {
 	case dns.TypeA:
-		for _, ip := range u.IPv4[name] {
-			resp.Answer = append(resp.Answer, &dns.A{Hdr: hdr, A: ip})
-		}
+		ips = ipv4[name]
 	case dns.TypeAAAA:
-		for _, ip := range u.IPv6[name] {
-			resp.Answer = append(resp.Answer, &dns.AAAA{Hdr: hdr, AAAA: ip})
+		ips = ipv6[name]
+	default:
+		return nil
+	}
+
+	hdr := dns.RR_Header{Name: name, Rrtype: qtype}
+	for _, ip := range ips {
+		switch qtype {
+		case dns.TypeA:
+			ans = append(ans, &dns.A{Hdr: hdr, A: ip})
+		case dns.TypeAAAA:
+			ans = append(ans, &dns.AAAA{Hdr: hdr, AAAA: ip})
 		}
 	}
-	if len(resp.Answer) == 0 {
-		resp.SetRcode(m, dns.RcodeNameError)
-	}
 
-	return resp, nil
-}
-
-// Address implements [upstream.Upstream] interface for *Upstream.
-func (u *Upstream) Address() string {
-	return "todo.upstream.example"
-}
-
-// Close implements [upstream.Upstream] interface for *Upstream.
-func (u *Upstream) Close() (err error) {
-	return nil
+	return ans
 }
 
 // MatchedResponse is a test helper that returns a response with answer if req
@@ -162,22 +170,24 @@ func mustAnsAAAA(respHdr dns.RR_Header, s string) (ans []dns.RR) {
 	}}
 }
 
-// NewUpstreamMock returns an [*UpstreamMock], fields OnAddress and OnClose of
-// which are set to stubs that return "upstream.example" and nil respectively.
-// The field OnExchange is set to onExc.
-func NewUpstreamMock(onExc func(req *dns.Msg) (resp *dns.Msg, err error)) (u *UpstreamMock) {
-	return &UpstreamMock{
-		OnAddress:  func() (addr string) { return "upstream.example" },
-		OnExchange: onExc,
-		OnClose:    func() (err error) { return nil },
+// NewUpstream returns an *dnsproxytest.Upstream, fields OnAddress and
+// OnClose of which are set to stubs that return "upstream.example" and nil
+// respectively. The field OnExchange is a stub that panics on call.
+func NewUpstream() (u *dnsproxytest.Upstream) {
+	return &dnsproxytest.Upstream{
+		OnAddress: func() (addr string) { return "upstream.example" },
+		OnClose:   func() (err error) { return nil },
+		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
+			panic(testutil.UnexpectedCall(req))
+		},
 	}
 }
 
-// NewBlockUpstream returns an [*UpstreamMock] that works like an upstream that
-// supports hash-based safe-browsing/adult-blocking feature.  If shouldBlock is
-// true, hostname's actual hash is returned, blocking it.  Otherwise, it returns
-// a different hash.
-func NewBlockUpstream(hostname string, shouldBlock bool) (u *UpstreamMock) {
+// NewBlockUpstream returns an *dnsproxytest.Upstream that works like an
+// upstream that supports hash-based safe-browsing/adult-blocking feature.  If
+// shouldBlock is true, hostname's actual hash is returned, blocking it.
+// Otherwise, it returns a different hash.
+func NewBlockUpstream(hostname string, shouldBlock bool) (u *dnsproxytest.Upstream) {
 	hash := sha256.Sum256([]byte(hostname))
 	hashStr := hex.EncodeToString(hash[:])
 	if !shouldBlock {
@@ -197,7 +207,7 @@ func NewBlockUpstream(hostname string, shouldBlock bool) (u *UpstreamMock) {
 		Answer: []dns.RR{ans},
 	}
 
-	return &UpstreamMock{
+	return &dnsproxytest.Upstream{
 		OnAddress: func() (addr string) { return "sbpc.upstream.example" },
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			resp = respTmpl.Copy()
@@ -210,14 +220,14 @@ func NewBlockUpstream(hostname string, shouldBlock bool) (u *UpstreamMock) {
 	}
 }
 
-// ErrUpstream is the error returned from the [*UpstreamMock] created by
-// [NewErrorUpstream].
+// ErrUpstream is the error returned from the *dnsproxytest.Upstream created
+// by NewErrorUpstream.
 const ErrUpstream errors.Error = "test upstream error"
 
-// NewErrorUpstream returns an [*UpstreamMock] that returns [ErrUpstream] from
-// its Exchange method.
-func NewErrorUpstream() (u *UpstreamMock) {
-	return &UpstreamMock{
+// NewErrorUpstream returns an *dnsproxytest.Upstream that returns
+// ErrUpstream from its Exchange method.
+func NewErrorUpstream() (u *dnsproxytest.Upstream) {
+	return &dnsproxytest.Upstream{
 		OnAddress: func() (addr string) { return "error.upstream.example" },
 		OnExchange: func(_ *dns.Msg) (resp *dns.Msg, err error) {
 			return nil, ErrUpstream

internal/dnsforward/dnsforward_internal_test.go

@@ -32,6 +32,7 @@ import (
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering/hashprefix"
 	"github.com/AdguardTeam/AdGuardHome/internal/filtering/safesearch"
 	"github.com/AdguardTeam/AdGuardHome/internal/schedule"
+	"github.com/AdguardTeam/dnsproxy/dnsproxytest"
 	"github.com/AdguardTeam/dnsproxy/proxy"
 	"github.com/AdguardTeam/dnsproxy/upstream"
 	"github.com/AdguardTeam/golibs/logutil/slogutil"
@@ -294,7 +295,7 @@ func createGoogleATestMessage() *dns.Msg {
 }
 
 func newGoogleUpstream() (u upstream.Upstream) {
-	return &aghtest.UpstreamMock{
+	return &dnsproxytest.Upstream{
 		OnAddress: func() (addr string) { return "google.upstream.example" },
 		OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
 			return cmp.Or(
@@ -651,11 +652,13 @@ func TestSafeSearch(t *testing.T) {
 	s := createTestServer(t, filterConf, forwardConf, testTLSManager)
 
 	pt := testutil.NewPanicT(t)
-	ups := aghtest.NewUpstreamMock(func(req *dns.Msg) (resp *dns.Msg, err error) {
+	ups := aghtest.NewUpstream()
+	ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
 		assert.Equal(pt, googleSafeSearch, req.Question[0].Name)
 
 		return aghtest.MatchedResponse(req, dns.TypeA, googleSafeSearch, "1.2.3.4"), nil
-	})
+	}
+
 	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
 
 	startDeferStop(t, s)
@@ -830,14 +833,15 @@ func TestServerCustomClientUpstream(t *testing.T) {
 		testTLSManager,
 	)
 
-	ups := aghtest.NewUpstreamMock(func(req *dns.Msg) (resp *dns.Msg, err error) {
+	ups := aghtest.NewUpstream()
+	ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
 		upsCalledCounter.Add(1)
 
 		return cmp.Or(
 			aghtest.MatchedResponse(req, dns.TypeA, "host", "192.168.0.1"),
 			new(dns.Msg).SetRcode(req, dns.RcodeNameError),
 		), nil
-	})
+	}
 
 	customUpsConf := proxy.NewCustomUpstreamConfig(
 		&proxy.UpstreamConfig{
@@ -909,10 +913,8 @@ func TestBlockCNAMEProtectionEnabled(t *testing.T) {
 		},
 		testTLSManager,
 	)
-	testUpstm := &aghtest.Upstream{
-		CName: testCNAMEs,
-		IPv4:  testIPv4,
-	}
+
+	testUpstm := aghtest.NewExchangingUpstream(t, testCNAMEs, testIPv4, nil)
 
 	// TODO(m.kazantsev):  Get rid of this manual assignment of upstreams across
 	// the whole project.
@@ -951,12 +953,10 @@ func TestBlockCNAME(t *testing.T) {
 		forwardConf,
 		testTLSManager,
 	)
-	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{
-		&aghtest.Upstream{
-			CName: testCNAMEs,
-			IPv4:  testIPv4,
-		},
-	}
+
+	ups := aghtest.NewExchangingUpstream(t, testCNAMEs, testIPv4, nil)
+
+	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
 	startDeferStop(t, s)
 
 	addr := s.dnsProxy.Addr(proxy.ProtoUDP).String()
@@ -1028,12 +1028,10 @@ func TestClientRulesForCNAMEMatching(t *testing.T) {
 		forwardConf,
 		testTLSManager,
 	)
-	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{
-		&aghtest.Upstream{
-			CName: testCNAMEs,
-			IPv4:  testIPv4,
-		},
-	}
+
+	ups := aghtest.NewExchangingUpstream(t, testCNAMEs, testIPv4, nil)
+
+	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
 	startDeferStop(t, s)
 
 	addr := s.dnsProxy.Addr(proxy.ProtoUDP)
@@ -1373,12 +1371,14 @@ func TestRewrite(t *testing.T) {
 		ServePlainDNS: true,
 	}))
 
-	ups := aghtest.NewUpstreamMock(func(req *dns.Msg) (resp *dns.Msg, err error) {
+	ups := aghtest.NewUpstream()
+	ups.OnExchange = func(req *dns.Msg) (resp *dns.Msg, err error) {
 		return cmp.Or(
 			aghtest.MatchedResponse(req, dns.TypeA, "example.org", "4.3.2.1"),
 			new(dns.Msg).SetRcode(req, dns.RcodeNameError),
 		), nil
-	})
+	}
+
 	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
 	startDeferStop(t, s)
 

internal/dnsforward/requesthandler_internal_test.go

@@ -73,22 +73,21 @@ func TestServer_ServeDNS(t *testing.T) {
 	err = s.Prepare(testutil.ContextWithTimeout(t, testTimeout), &forwardConf)
 	require.NoError(t, err)
 
-	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{
-		&aghtest.Upstream{
-			CName: map[string][]string{
-				"cname.exception.": {"cname.specific."},
-				"should.block.":    {"blocked.domain."},
-				"allowed.first.":   {"allowed.domain.", "blocked.domain."},
-				"blocked.first.":   {"blocked.domain.", "allowed.domain."},
-			},
-			IPv4: map[string][]net.IP{
-				"a.exception.": {{0, 0, 0, 1}},
-			},
-			IPv6: map[string][]net.IP{
-				"aaaa.exception.": {net.ParseIP("::1")},
-			},
-		},
+	cNames := map[string][]string{
+		"cname.exception.": {"cname.specific."},
+		"should.block.":    {"blocked.domain."},
+		"allowed.first.":   {"allowed.domain.", "blocked.domain."},
+		"blocked.first.":   {"blocked.domain.", "allowed.domain."},
 	}
+
+	ups := aghtest.NewExchangingUpstream(
+		t,
+		cNames,
+		map[string][]net.IP{"a.exception.": {{0, 0, 0, 1}}},
+		map[string][]net.IP{"aaaa.exception.": {net.ParseIP("::1")}},
+	)
+
+	s.conf.UpstreamConfig.Upstreams = []upstream.Upstream{ups}
 	startDeferStop(t, s)
 
 	testCases := []struct {