feat: support type query in dns target

README.md

@@ -104,8 +104,12 @@ examples:
 
 Resolve hostname via DNS and check the host exists or not.
 
+You can specify record type as a `type` query.
+Supported type is `A`, `AAAA`, `CNAME`, `MX`, `NS`, and `TXT`.
+
 examples:
 - `dns:example.com`
+- `dns:example.com?type=AAAA`
 
 #### exec
 

probe/dns.go

@@ -2,6 +2,7 @@ package probe
 
 import (
 	"context"
+	"errors"
 	"net"
 	"net/url"
 	"strings"
@@ -10,44 +11,122 @@ import (
 	"github.com/macrat/ayd/store"
 )
 
+var (
+	ErrUnsupportedDNSType = errors.New("unsupported DNS type")
+)
+
+func dnsResolveAuto(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	addrs, err := r.LookupHost(ctx, target)
+	return strings.Join(addrs, "\n"), err
+}
+
+func dnsResolveA(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	ips, err := r.LookupIP(ctx, "ip4", target)
+	addrs := make([]string, len(ips))
+	for i, x := range ips {
+		addrs[i] = x.String()
+	}
+	return strings.Join(addrs, "\n"), err
+}
+
+func dnsResolveAAAA(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	ips, err := r.LookupIP(ctx, "ip6", target)
+	addrs := make([]string, len(ips))
+	for i, x := range ips {
+		addrs[i] = x.String()
+	}
+	return strings.Join(addrs, "\n"), err
+}
+
+func dnsResolveCNAME(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	return r.LookupCNAME(ctx, target)
+}
+
+func dnsResolveMX(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	mxs, err := r.LookupMX(ctx, target)
+	addrs := make([]string, len(mxs))
+	for i, x := range mxs {
+		addrs[i] = x.Host
+	}
+	return strings.Join(addrs, "\n"), err
+}
+
+func dnsResolveNS(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	nss, err := r.LookupNS(ctx, target)
+	addrs := make([]string, len(nss))
+	for i, x := range nss {
+		addrs[i] = x.Host
+	}
+	return strings.Join(addrs, "\n"), err
+}
+
+func dnsResolveTXT(r *net.Resolver, ctx context.Context, target string) (string, error) {
+	texts, err := r.LookupTXT(ctx, target)
+	return strings.Join(texts, "\n"), err
+}
+
 type DNSProbe struct {
-	target *url.URL
+	target  *url.URL
+	resolve func(r *net.Resolver, ctx context.Context, target string) (string, error)
 }
 
 func NewDNSProbe(u *url.URL) (DNSProbe, error) {
-	if u.Opaque != "" {
-		return DNSProbe{&url.URL{Scheme: "dns", Opaque: u.Opaque}}, nil
-	} else {
-		return DNSProbe{&url.URL{Scheme: "dns", Opaque: u.Hostname()}}, nil
+	p := DNSProbe{target: &url.URL{Scheme: "dns", Opaque: u.Opaque}}
+	if u.Opaque == "" {
+		p.target.Opaque = u.Hostname()
+	}
+	switch strings.ToUpper(u.Query().Get("type")) {
+	case "":
+		p.resolve = dnsResolveAuto
+	case "A":
+		p.target.RawQuery = url.Values{"type": {"A"}}.Encode()
+		p.resolve = dnsResolveA
+	case "AAAA":
+		p.target.RawQuery = url.Values{"type": {"AAAA"}}.Encode()
+		p.resolve = dnsResolveAAAA
+	case "CNAME":
+		p.target.RawQuery = url.Values{"type": {"CNAME"}}.Encode()
+		p.resolve = dnsResolveCNAME
+	case "MX":
+		p.target.RawQuery = url.Values{"type": {"MX"}}.Encode()
+		p.resolve = dnsResolveMX
+	case "NS":
+		p.target.RawQuery = url.Values{"type": {"NS"}}.Encode()
+		p.resolve = dnsResolveNS
+	case "TXT":
+		p.target.RawQuery = url.Values{"type": {"TXT"}}.Encode()
+		p.resolve = dnsResolveTXT
+	default:
+		return DNSProbe{}, ErrUnsupportedDNSType
 	}
+	return p, nil
 }
 
 func (p DNSProbe) Target() *url.URL {
 	return p.target
 }
 
 func (p DNSProbe) Check(ctx context.Context, r Reporter) {
-	resolver := &net.Resolver{}
+	var resolver net.Resolver
 
 	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
 	defer cancel()
 
 	st := time.Now()
-	addrs, err := resolver.LookupHost(ctx, p.target.Opaque)
+	msg, err := p.resolve(&resolver, ctx, p.target.Opaque)
 	d := time.Now().Sub(st)
 
 	rec := store.Record{
 		CheckedAt: st,
 		Target:    p.target,
+		Status:    store.STATUS_HEALTHY,
+		Message:   msg,
 		Latency:   d,
 	}
 
 	if err != nil {
 		rec.Status = store.STATUS_FAILURE
 		rec.Message = err.Error()
-	} else {
-		rec.Status = store.STATUS_HEALTHY
-		rec.Message = strings.Join(addrs, "\n")
 	}
 
 	r.Report(timeoutOr(ctx, rec))

probe/dns_test.go

@@ -11,6 +11,18 @@ func TestDNSProbe(t *testing.T) {
 
 	AssertProbe(t, []ProbeTest{
 		{"dns:localhost", store.STATUS_HEALTHY, `(127\.0\.0\.1|::1)` + "\n" + `(127\.0\.0\.1|::1)`},
+
+		{"dns:localhost?type=AAAA", store.STATUS_HEALTHY, `::1`},
+		{"dns:localhost?type=A", store.STATUS_HEALTHY, `127.0.0.1`},
+
+		{"dns:example.com?type=CNAME", store.STATUS_HEALTHY, `example.com.`},
+
+		{"dns:example.com?type=MX", store.STATUS_HEALTHY, `.`},
+
+		{"dns:example.com?type=NS", store.STATUS_HEALTHY, `[a-z]\.iana-servers\.net\.(` + "\n" + `[a-z]\.iana-servers\.net\.)*`},
+
+		{"dns:example.com?type=TXT", store.STATUS_HEALTHY, "v=spf1 -all\n[0-9a-z]{32}"},
+
 		{"dns:of-course-definitely-no-such-host", store.STATUS_FAILURE, `lookup of-course-definitely-no-such-host(:| ).+`},
 	})