From c2db5bf60385eb811b2f1c4ff6aa3ad32ffcdf8f Mon Sep 17 00:00:00 2001
From: Baptiste Courtois <b.courtois@criteo.com>
Date: Sun, 14 Jun 2020 00:00:20 +0200
Subject: [PATCH] [dns] validate query_class and query_type config values

This allows to catch simple config errors earlier.

Signed-off-by: Baptiste Courtois <b.courtois@criteo.com>
---
 config/config.go                      | 12 ++++++++++++
 config/config_test.go                 |  8 ++++++++
 config/testdata/invalid-dns-class.yml |  8 ++++++++
 config/testdata/invalid-dns-type.yml  |  8 ++++++++
 4 files changed, 36 insertions(+)
 create mode 100644 config/testdata/invalid-dns-class.yml
 create mode 100644 config/testdata/invalid-dns-type.yml

diff --git a/config/config.go b/config/config.go
index 9730bae5..3fddb3c5 100644
--- 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"
 )
@@ -233,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
index 4fc45298..21c81360 100644
--- 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
new file mode 100644
index 00000000..2893b317
--- /dev/null
+++ 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
new file mode 100644
index 00000000..6c01c445
--- /dev/null
+++ 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