From 6b88618a358cf1477e07b317eea959dc11ea2988 Mon Sep 17 00:00:00 2001
From: Baptiste Courtois <b.courtois@criteo.com>
Date: Sat, 6 Jun 2020 18:27:39 +0200
Subject: [PATCH] [dns] Allow to specify query Class

Signed-off-by: Baptiste Courtois <b.courtois@criteo.com>
---
 CONFIGURATION.md   |  1 +
 config/config.go   |  1 +
 prober/dns.go      | 17 +++++++++++++++--
 prober/dns_test.go | 23 +++++++++++++++++++++--
 4 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/CONFIGURATION.md b/CONFIGURATION.md
index ec406858..0ab0d21c 100644
--- 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
index 148de844..9730bae5 100644
--- a/config/config.go
+++ b/config/config.go
@@ -175,6 +175,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.
diff --git a/prober/dns.go b/prober/dns.go
index 9c26c9bb..edb5b282 100644
--- 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
index b4f0472c..75dc7c33 100644
--- 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,