AdGuardHome/internal/aghtest/upstream.go

228 lines
6.1 KiB
Go
Raw Normal View History

package aghtest
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"net"
"net/netip"
"strings"
"github.com/AdguardTeam/dnsproxy/upstream"
"github.com/AdguardTeam/golibs/errors"
"github.com/miekg/dns"
)
// Additional Upstream Testing Utilities
// 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
}
var _ upstream.Upstream = (*Upstream)(nil)
// 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)
if len(m.Question) == 0 {
return nil, fmt.Errorf("question should not be empty")
}
q := m.Question[0]
name := q.Name
for _, cname := range u.CName[name] {
resp.Answer = append(resp.Answer, &dns.CNAME{
Hdr: dns.RR_Header{Name: name, Rrtype: dns.TypeCNAME},
Target: cname,
})
}
qtype := q.Qtype
hdr := dns.RR_Header{
Name: name,
Rrtype: qtype,
}
switch qtype {
case dns.TypeA:
for _, ip := range u.IPv4[name] {
resp.Answer = append(resp.Answer, &dns.A{Hdr: hdr, A: ip})
}
case dns.TypeAAAA:
for _, ip := range u.IPv6[name] {
resp.Answer = append(resp.Answer, &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
}
// MatchedResponse is a test helper that returns a response with answer if req
// has question type qt, and target targ. Otherwise, it returns nil.
//
// req must not be nil and req.Question must have a length of 1. Answer is
// interpreted in the following ways:
//
// - For A and AAAA queries, answer must be an IP address of the corresponding
// protocol version.
//
// - For PTR queries, answer should be a domain name in the response.
//
// If the answer does not correspond to the question type, MatchedResponse panics.
// Panics are used instead of [testing.TB], because the helper is intended to
// use in [UpstreamMock.OnExchange] callbacks, which are usually called in a
// separate goroutine.
//
// TODO(a.garipov): Consider adding version with DNS class as well.
func MatchedResponse(req *dns.Msg, qt uint16, targ, answer string) (resp *dns.Msg) {
if req == nil || len(req.Question) != 1 {
panic(fmt.Errorf("bad req: %+v", req))
}
q := req.Question[0]
targ = dns.Fqdn(targ)
if q.Qclass != dns.ClassINET || q.Qtype != qt || q.Name != targ {
return nil
}
respHdr := dns.RR_Header{
Name: targ,
Rrtype: qt,
Class: dns.ClassINET,
Ttl: 60,
}
resp = new(dns.Msg).SetReply(req)
switch qt {
case dns.TypeA:
resp.Answer = mustAnsA(respHdr, answer)
case dns.TypeAAAA:
resp.Answer = mustAnsAAAA(respHdr, answer)
case dns.TypePTR:
resp.Answer = []dns.RR{&dns.PTR{
Hdr: respHdr,
Ptr: answer,
}}
default:
panic(fmt.Errorf("aghtest: bad question type: %s", dns.Type(qt)))
}
return resp
}
// mustAnsA returns valid answer records if s is a valid IPv4 address.
// Otherwise, mustAnsA panics.
func mustAnsA(respHdr dns.RR_Header, s string) (ans []dns.RR) {
ip, err := netip.ParseAddr(s)
if err != nil || !ip.Is4() {
panic(fmt.Errorf("aghtest: bad A answer: %+v", s))
}
return []dns.RR{&dns.A{
Hdr: respHdr,
A: ip.AsSlice(),
}}
}
// mustAnsAAAA returns valid answer records if s is a valid IPv6 address.
// Otherwise, mustAnsAAAA panics.
func mustAnsAAAA(respHdr dns.RR_Header, s string) (ans []dns.RR) {
ip, err := netip.ParseAddr(s)
if err != nil || !ip.Is6() {
panic(fmt.Errorf("aghtest: bad AAAA answer: %+v", s))
}
return []dns.RR{&dns.AAAA{
Hdr: respHdr,
AAAA: ip.AsSlice(),
}}
}
// 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 },
}
}
// 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) {
hash := sha256.Sum256([]byte(hostname))
hashStr := hex.EncodeToString(hash[:])
if !shouldBlock {
hashStr = hex.EncodeToString(hash[:])[:2] + strings.Repeat("ab", 28)
}
ans := &dns.TXT{
Hdr: dns.RR_Header{
Name: "",
Rrtype: dns.TypeTXT,
Class: dns.ClassINET,
Ttl: 60,
},
Txt: []string{hashStr},
}
respTmpl := &dns.Msg{
Answer: []dns.RR{ans},
}
return &UpstreamMock{
OnAddress: func() (addr string) { return "sbpc.upstream.example" },
OnExchange: func(req *dns.Msg) (resp *dns.Msg, err error) {
resp = respTmpl.Copy()
resp.SetReply(req)
resp.Answer[0].(*dns.TXT).Hdr.Name = req.Question[0].Name
return resp, nil
},
OnClose: func() (err error) { return nil },
}
Pull request: 2704 local addresses vol.2 Merge in DNS/adguard-home from 2704-local-addresses-vol.2 to master Updates #2704. Updates #2829. Squashed commit of the following: commit 507d038c2709de59246fc0b65c3c4ab8e38d1990 Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 31 14:33:05 2021 +0300 aghtest: fix file name commit 8e19f99337bee1d88ad6595adb96f9bb23fa3c41 Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 31 14:06:43 2021 +0300 aghnet: rm redundant mutexes commit 361fa418b33ed160ca20862be1c455ab9378c03f Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 31 13:45:30 2021 +0300 all: fix names, docs commit 14034f4f0230d7aaa3645054946ae5c278089a99 Merge: 35e265cc a72ce1cf Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 31 13:38:15 2021 +0300 Merge branch 'master' into 2704-local-addresses-vol.2 commit 35e265cc8cd308ef1fda414b58c0217cb5f258e4 Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 31 13:33:35 2021 +0300 aghnet: imp naming commit 7a7edac7208a40697d7bc50682b923a144e28e2b Author: Eugene Burkov <e.burkov@adguard.com> Date: Tue Mar 30 20:59:54 2021 +0300 changelog: oops, nope yet commit d26a5d2513daf662ac92053b5e235189a64cc022 Author: Eugene Burkov <e.burkov@adguard.com> Date: Tue Mar 30 20:55:53 2021 +0300 all: some renaming for the glory of semantics commit 9937fa619452b0742616217b975e3ff048d58acb Author: Eugene Burkov <e.burkov@adguard.com> Date: Mon Mar 29 15:34:42 2021 +0300 all: log changes commit d8d9e6dfeea8474466ee25f27021efdd3ddb1592 Author: Eugene Burkov <e.burkov@adguard.com> Date: Fri Mar 26 18:32:23 2021 +0300 all: imp localresolver, imp cutting off own addresses commit 344140df449b85925f19b460fd7dc7c08e29c35a Author: Eugene Burkov <e.burkov@adguard.com> Date: Fri Mar 26 14:53:33 2021 +0300 all: imp code quality commit 1c5c0babec73b125044e23dd3aa75d8eefc19b28 Author: Eugene Burkov <e.burkov@adguard.com> Date: Thu Mar 25 20:44:08 2021 +0300 all: fix go.mod commit 0b9fb3c2369a752e893af8ddc45a86bb9fb27ce5 Author: Eugene Burkov <e.burkov@adguard.com> Date: Thu Mar 25 20:38:51 2021 +0300 all: add error handling commit a7a2e51f57fc6f8f74b95a264ad345cd2a9e026e Merge: c13be634 27f4f052 Author: Eugene Burkov <e.burkov@adguard.com> Date: Thu Mar 25 19:48:36 2021 +0300 Merge branch 'master' into 2704-local-addresses-vol.2 commit c13be634f47bcaed9320a732a51c0e4752d0dad0 Author: Eugene Burkov <e.burkov@adguard.com> Date: Thu Mar 25 18:52:28 2021 +0300 all: cover rdns with tests, imp aghnet functionality commit 48bed9025944530c613ee53e7961d6d5fbabf8be Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 24 20:18:07 2021 +0300 home: make rdns great again commit 1dbacfc8d5b6895807797998317fe3cc814617c1 Author: Eugene Burkov <e.burkov@adguard.com> Date: Wed Mar 24 16:07:52 2021 +0300 all: imp external client restriction commit 1208a319a7f4ffe7b7fa8956f245d7a19437c0a4 Author: Eugene Burkov <e.burkov@adguard.com> Date: Mon Mar 22 15:26:45 2021 +0300 all: finish local ptr processor commit c8827fc3db289e1a5d7a11d057743bab39957b02 Author: Eugene Burkov <e.burkov@adguard.com> Date: Tue Mar 2 13:41:22 2021 +0300 all: imp ipdetector, add local ptr processor
2021-03-31 13:00:47 +01:00
}
// ErrUpstream is the error returned from the [*UpstreamMock] 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{
OnAddress: func() (addr string) { return "error.upstream.example" },
OnExchange: func(_ *dns.Msg) (resp *dns.Msg, err error) {
return nil, errors.Error("test upstream error")
},
OnClose: func() (err error) { return nil },
}
}