git » spf » master » tree

[master] / internal / dnstest / dns.go 

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
// DNS resolver for testing purposes.
//
// In the future, when go fuzz can make use of _test.go files, we can rename
// this file dns_test.go and remove this extra package entirely.
// Until then, unfortunately this is the most reasonable way to share these
// helpers between go and fuzz tests.
package dnstest

import (
	"context"
	"net"
	"strings"
)

// Testing DNS resolver.
//
// Not exported since this is not part of the public API and only used
// internally on tests.
type TestResolver struct {
	Txt    map[string][]string
	Mx     map[string][]*net.MX
	Ip     map[string][]net.IP
	Addr   map[string][]string
	Cname  map[string]string
	Errors map[string]error
}

func NewResolver() *TestResolver {
	return &TestResolver{
		Txt:    map[string][]string{},
		Mx:     map[string][]*net.MX{},
		Ip:     map[string][]net.IP{},
		Addr:   map[string][]string{},
		Cname:  map[string]string{},
		Errors: map[string]error{},
	}
}

var nxDomainErr = &net.DNSError{
	Err:        "domain not found (for testing)",
	IsNotFound: true,
}

func (r *TestResolver) LookupTXT(ctx context.Context, domain string) (txts []string, err error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	domain = strings.ToLower(domain)
	domain = strings.TrimRight(domain, ".")
	if cname, ok := r.Cname[domain]; ok {
		return r.LookupTXT(ctx, cname)
	}
	if _, ok := r.Txt[domain]; !ok && r.Errors[domain] == nil {
		return nil, nxDomainErr
	}
	return r.Txt[domain], r.Errors[domain]
}

func (r *TestResolver) LookupMX(ctx context.Context, domain string) (mxs []*net.MX, err error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	domain = strings.ToLower(domain)
	domain = strings.TrimRight(domain, ".")
	if cname, ok := r.Cname[domain]; ok {
		return r.LookupMX(ctx, cname)
	}
	if _, ok := r.Mx[domain]; !ok && r.Errors[domain] == nil {
		return nil, nxDomainErr
	}
	return r.Mx[domain], r.Errors[domain]
}

func (r *TestResolver) LookupIPAddr(ctx context.Context, host string) (as []net.IPAddr, err error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	host = strings.ToLower(host)
	host = strings.TrimRight(host, ".")
	if cname, ok := r.Cname[host]; ok {
		return r.LookupIPAddr(ctx, cname)
	}
	if _, ok := r.Ip[host]; !ok && r.Errors[host] == nil {
		return nil, nxDomainErr
	}
	return ipsToAddrs(r.Ip[host]), r.Errors[host]
}

func ipsToAddrs(ips []net.IP) []net.IPAddr {
	as := []net.IPAddr{}
	for _, ip := range ips {
		as = append(as, net.IPAddr{IP: ip, Zone: ""})
	}
	return as
}

func (r *TestResolver) LookupAddr(ctx context.Context, host string) (addrs []string, err error) {
	if ctx.Err() != nil {
		return nil, ctx.Err()
	}
	host = strings.ToLower(host)
	host = strings.TrimRight(host, ".")
	if cname, ok := r.Cname[host]; ok {
		return r.LookupAddr(ctx, cname)
	}
	if _, ok := r.Addr[host]; !ok && r.Errors[host] == nil {
		return nil, nxDomainErr
	}
	return r.Addr[host], r.Errors[host]
}