mox!

dns/mock.go

@@ -0,0 +1,156 @@
+package dns
+
+import (
+	"context"
+	"fmt"
+	"net"
+)
+
+// MockResolver is a Resolver used for testing.
+// Set DNS records in the fields, which map FQDNs (with trailing dot) to values.
+type MockResolver struct {
+	PTR   map[string][]string
+	A     map[string][]string
+	AAAA  map[string][]string
+	TXT   map[string][]string
+	MX    map[string][]*net.MX
+	CNAME map[string]string
+	Fail  map[Mockreq]struct{}
+}
+
+type Mockreq struct {
+	Type string // E.g. "cname", "txt", "mx", "ptr", etc.
+	Name string
+}
+
+var _ Resolver = MockResolver{}
+
+func (r MockResolver) nxdomain(s string) *net.DNSError {
+	return &net.DNSError{
+		Err:        "no record",
+		Name:       s,
+		Server:     "localhost",
+		IsNotFound: true,
+	}
+}
+
+func (r MockResolver) servfail(s string) *net.DNSError {
+	return &net.DNSError{
+		Err:         "temp error",
+		Name:        s,
+		Server:      "localhost",
+		IsTemporary: true,
+	}
+}
+
+func (r MockResolver) LookupCNAME(ctx context.Context, name string) (string, error) {
+	if _, ok := r.Fail[Mockreq{"cname", name}]; ok {
+		return "", r.servfail(name)
+	}
+	if cname, ok := r.CNAME[name]; ok {
+		return cname, nil
+	}
+	return "", r.nxdomain(name)
+}
+
+func (r MockResolver) LookupAddr(ctx context.Context, ip string) ([]string, error) {
+	if _, ok := r.Fail[Mockreq{"ptr", ip}]; ok {
+		return nil, r.servfail(ip)
+	}
+	l, ok := r.PTR[ip]
+	if !ok {
+		return nil, r.nxdomain(ip)
+	}
+	return l, nil
+}
+
+func (r MockResolver) LookupNS(ctx context.Context, name string) ([]*net.NS, error) {
+	return nil, r.servfail("ns not implemented")
+}
+
+func (r MockResolver) LookupPort(ctx context.Context, network, service string) (port int, err error) {
+	return 0, r.servfail("port not implemented")
+}
+
+func (r MockResolver) LookupSRV(ctx context.Context, service, proto, name string) (string, []*net.SRV, error) {
+	return "", nil, r.servfail("srv not implemented")
+}
+
+func (r MockResolver) LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error) {
+	if _, ok := r.Fail[Mockreq{"ipaddr", host}]; ok {
+		return nil, r.servfail(host)
+	}
+	addrs, err := r.LookupHost(ctx, host)
+	if err != nil {
+		return nil, err
+	}
+	ips := make([]net.IPAddr, len(addrs))
+	for i, a := range addrs {
+		ip := net.ParseIP(a)
+		if ip == nil {
+			return nil, fmt.Errorf("malformed ip %q", a)
+		}
+		ips[i] = net.IPAddr{IP: ip}
+	}
+	return ips, nil
+}
+
+func (r MockResolver) LookupHost(ctx context.Context, host string) (addrs []string, err error) {
+	if _, ok := r.Fail[Mockreq{"host", host}]; ok {
+		return nil, r.servfail(host)
+	}
+	addrs = append(addrs, r.A[host]...)
+	addrs = append(addrs, r.AAAA[host]...)
+	if len(addrs) > 0 {
+		return addrs, nil
+	}
+	if cname, ok := r.CNAME[host]; ok {
+		return []string{cname}, nil
+	}
+	return nil, r.nxdomain(host)
+}
+
+func (r MockResolver) LookupIP(ctx context.Context, network, host string) ([]net.IP, error) {
+	if _, ok := r.Fail[Mockreq{"ip", host}]; ok {
+		return nil, r.servfail(host)
+	}
+	var ips []net.IP
+	switch network {
+	case "ip", "ip4":
+		for _, ip := range r.A[host] {
+			ips = append(ips, net.ParseIP(ip))
+		}
+	}
+	switch network {
+	case "ip", "ip6":
+		for _, ip := range r.AAAA[host] {
+			ips = append(ips, net.ParseIP(ip))
+		}
+	}
+	if len(ips) == 0 {
+		return nil, r.nxdomain(host)
+	}
+	return ips, nil
+}
+
+func (r MockResolver) LookupMX(ctx context.Context, name string) ([]*net.MX, error) {
+	if _, ok := r.Fail[Mockreq{"mx", name}]; ok {
+		return nil, r.servfail(name)
+	}
+	l, ok := r.MX[name]
+	if !ok {
+		return nil, r.nxdomain(name)
+	}
+	return l, nil
+}
+
+func (r MockResolver) LookupTXT(ctx context.Context, name string) ([]string, error) {
+	if _, ok := r.Fail[Mockreq{"txt", name}]; ok {
+		return nil, r.servfail(name)
+	}
+	l, ok := r.TXT[name]
+	if !ok {
+		return nil, r.nxdomain(name)
+	}
+	return l, nil
+}