[dns] Allow to specify query Class (#635)

* [dns] validate query_class and query_type config values This allows to catch simple config errors earlier. Signed-off-by: Baptiste Courtois <[email protected]>
prometheus · Jun 14, 2020 · a289888 · a289888
1 parent a8a525c
commit a289888
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 4 deletions.
diff --git a/CONFIGURATION.md b/CONFIGURATION.md
@@ -156,6 +156,7 @@ tls_config:
 query_name: <string>
 
 [ query_type: <string> | default = "ANY" ]
+[ query_class: <string> | default = "IN" ]
 
 # List of valid response codes.
 valid_rcodes:

diff --git a/config/config.go b/config/config.go
@@ -23,6 +23,7 @@ import (
 
 	yaml "gopkg.in/yaml.v3"
 
+	"github.com/miekg/dns"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/config"
 )
@@ -175,6 +176,7 @@ type DNSProbe struct {
 	IPProtocolFallback bool           `yaml:"ip_protocol_fallback,omitempty"`
 	SourceIPAddress    string         `yaml:"source_ip_address,omitempty"`
 	TransportProtocol  string         `yaml:"transport_protocol,omitempty"`
+	QueryClass         string         `yaml:"query_class,omitempty"` // Defaults to IN.
 	QueryName          string         `yaml:"query_name,omitempty"`
 	QueryType          string         `yaml:"query_type,omitempty"`   // Defaults to ANY.
 	ValidRcodes        []string       `yaml:"valid_rcodes,omitempty"` // Defaults to NOERROR.
@@ -232,6 +234,17 @@ func (s *DNSProbe) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if s.QueryName == "" {
 		return errors.New("query name must be set for DNS module")
 	}
+	if s.QueryClass != "" {
+		if _, ok := dns.StringToClass[s.QueryClass]; !ok {
+			return fmt.Errorf("query class '%s' is not valid", s.QueryClass)
+		}
+	}
+	if s.QueryType != "" {
+		if _, ok := dns.StringToType[s.QueryType]; !ok {
+			return fmt.Errorf("query type '%s' is not valid", s.QueryType)
+		}
+	}
+
 	return nil
 }
 

diff --git a/config/config_test.go b/config/config_test.go
@@ -51,6 +51,14 @@ func TestLoadBadConfigs(t *testing.T) {
 			ConfigFile:    "testdata/invalid-dns-module.yml",
 			ExpectedError: "error parsing config file: query name must be set for DNS module",
 		},
+		{
+			ConfigFile:    "testdata/invalid-dns-class.yml",
+			ExpectedError: "error parsing config file: query class 'X' is not valid",
+		},
+		{
+			ConfigFile:    "testdata/invalid-dns-type.yml",
+			ExpectedError: "error parsing config file: query type 'X' is not valid",
+		},
 		{
 			ConfigFile:    "testdata/invalid-http-header-match.yml",
 			ExpectedError: "error parsing config file: regexp must be set for HTTP header matchers",

diff --git a/config/testdata/invalid-dns-class.yml b/config/testdata/invalid-dns-class.yml
@@ -0,0 +1,8 @@
+modules:
+  dns_test:
+    prober: dns
+    timeout: 5s
+    dns:
+      query_name: example.com
+      query_class: X
+      query_type: A
diff --git a/config/testdata/invalid-dns-type.yml b/config/testdata/invalid-dns-type.yml
@@ -0,0 +1,8 @@
+modules:
+  dns_test:
+    prober: dns
+    timeout: 5s
+    dns:
+      query_name: example.com
+      query_class: CH
+      query_type: X
diff --git a/prober/dns.go b/prober/dns.go
@@ -141,6 +141,16 @@ func ProbeDNS(ctx context.Context, target string, module config.Module, registry
 	registry.MustRegister(probeDNSAuthorityRRSGauge)
 	registry.MustRegister(probeDNSAdditionalRRSGauge)
 
+	qc := uint16(dns.ClassINET)
+	if module.DNS.QueryClass != "" {
+		var ok bool
+		qc, ok = dns.StringToClass[module.DNS.QueryClass]
+		if !ok {
+			level.Error(logger).Log("msg", "Invalid query class", "Class seen", module.DNS.QueryClass, "Existing classes", dns.ClassToString)
+			return false
+		}
+	}
+
 	qt := dns.TypeANY
 	if module.DNS.QueryType != "" {
 		var ok bool
@@ -200,9 +210,12 @@ func ProbeDNS(ctx context.Context, target string, module config.Module, registry
 	}
 
 	msg := new(dns.Msg)
-	msg.SetQuestion(dns.Fqdn(module.DNS.QueryName), qt)
+	msg.Id = dns.Id()
+	msg.RecursionDesired = true
+	msg.Question = make([]dns.Question, 1)
+	msg.Question[0] = dns.Question{dns.Fqdn(module.DNS.QueryName), qt, qc}
 
-	level.Info(logger).Log("msg", "Making DNS query", "target", target, "dial_protocol", dialProtocol, "query", module.DNS.QueryName, "type", qt)
+	level.Info(logger).Log("msg", "Making DNS query", "target", target, "dial_protocol", dialProtocol, "query", module.DNS.QueryName, "type", qt, "class", qc)
 	timeoutDeadline, _ := ctx.Deadline()
 	client.Timeout = time.Until(timeoutDeadline)
 	response, _, err := client.Exchange(msg, target)

diff --git a/prober/dns_test.go b/prober/dns_test.go
@@ -181,7 +181,12 @@ func authoritativeDNSHandler(w dns.ResponseWriter, r *dns.Msg) {
 			panic(err)
 		}
 		m.Answer = append(m.Answer, a)
-
+	} else if r.Question[0].Qclass == dns.ClassCHAOS && r.Question[0].Qtype == dns.TypeTXT {
+		txt, err := dns.NewRR("example.com. 3600 CH TXT \"goCHAOS\"")
+		if err != nil {
+			panic(err)
+		}
+		m.Answer = append(m.Answer, txt)
 	} else {
 		a, err := dns.NewRR("example.com. 3600 IN A 127.0.0.1")
 		if err != nil {
@@ -243,7 +248,21 @@ func TestAuthoritativeDNSResponse(t *testing.T) {
 				QueryName:          "example.com",
 				QueryType:          "SOA",
 			}, true,
-		}, {
+		},
+		{
+			config.DNSProbe{
+				IPProtocol:         "ip4",
+				IPProtocolFallback: true,
+				QueryClass:         "CH",
+				QueryName:          "example.com",
+				QueryType:          "TXT",
+				ValidateAnswer: config.DNSRRValidator{
+					FailIfMatchesRegexp:    []string{".*IN.*"},
+					FailIfNotMatchesRegexp: []string{".*CH.*"},
+				},
+			}, true,
+		},
+		{
 			config.DNSProbe{
 				IPProtocol:         "ip4",
 				IPProtocolFallback: true,