From c0ce61bf99978af89c09a0465dbc8abda95b54eb Mon Sep 17 00:00:00 2001
From: conorbroderick <cjayjayb@gmail.com>
Date: Mon, 14 Aug 2017 12:51:22 +0100
Subject: [PATCH] Added check for empty query name and handle case where
 targetAddr may empty

---
 config.go                       |  4 ++++
 config_test.go                  | 27 +++++++++++++++++++--------
 dns.go                          |  9 ++++-----
 testdata/blackbox-good.yml      |  3 ++-
 testdata/invalid-dns-module.yml |  8 ++++++++
 5 files changed, 37 insertions(+), 14 deletions(-)
 create mode 100644 testdata/invalid-dns-module.yml

diff --git a/config.go b/config.go
index 91880da..3d1d324 100644
--- a/config.go
+++ b/config.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"strings"
 	"sync"
@@ -155,6 +156,9 @@ func (s *DNSProbe) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	if err := checkOverflow(s.XXX, "dns probe"); err != nil {
 		return err
 	}
+	if s.QueryName == "" {
+		return errors.New("Query name must be set for DNS module")
+	}
 	return nil
 }
 
diff --git a/config_test.go b/config_test.go
index cea82c4..1379098 100644
--- a/config_test.go
+++ b/config_test.go
@@ -18,21 +18,32 @@ func TestLoadConfig(t *testing.T) {
 	}
 }
 
-func TestLoadBadConfig(t *testing.T) {
+func TestLoadBadConfigs(t *testing.T) {
 	sc := &SafeConfig{
 		C: &Config{},
 	}
-
-	expected := "unknown fields in dns probe: invalid_extra_field"
-
-	err := sc.reloadConfig("testdata/blackbox-bad.yml")
-	if err.Error() != expected {
-		t.Errorf("\nexpected:\n%v\ngot:\n%v", expected, err.Error())
+	tests := []struct {
+		ConfigFile    string
+		ExpectedError string
+	}{
+		{
+			ConfigFile:    "testdata/blackbox-bad.yml",
+			ExpectedError: "unknown fields in dns probe: invalid_extra_field",
+		},
+		{
+			ConfigFile:    "testdata/invalid-dns-module.yml",
+			ExpectedError: "Query name must be set for DNS module",
+		},
+	}
+	for i, test := range tests {
+		err := sc.reloadConfig(test.ConfigFile)
+		if err.Error() != test.ExpectedError {
+			t.Errorf("In case %v:\nExpected:\n%v\nGot:\n%v", i, test.ExpectedError, err.Error())
+		}
 	}
 }
 
 func TestHideConfigSecrets(t *testing.T) {
-
 	sc := &SafeConfig{
 		C: &Config{},
 	}
diff --git a/dns.go b/dns.go
index 6718faf..97c5ff2 100644
--- a/dns.go
+++ b/dns.go
@@ -110,15 +110,15 @@ func probeDNS(ctx context.Context, target string, module Module, registry *prome
 	}()
 
 	var ip *net.IPAddr
-	var err error
-
 	if module.DNS.TransportProtocol == "" {
 		module.DNS.TransportProtocol = "udp"
 	}
 	if module.DNS.TransportProtocol == "udp" || module.DNS.TransportProtocol == "tcp" {
-		targetAddr, port, _ := net.SplitHostPort(target)
-		if port == "" {
+		targetAddr, port, err := net.SplitHostPort(target)
+		if err != nil {
+			// Target only contains host so fallback to default port and set targetAddr as target.
 			port = "53"
+			targetAddr = target
 		}
 		ip, err = chooseProtocol(module.DNS.PreferredIPProtocol, targetAddr, registry)
 		if err != nil {
@@ -153,7 +153,6 @@ func probeDNS(ctx context.Context, target string, module Module, registry *prome
 
 	timeoutDeadline, _ := ctx.Deadline()
 	client.Timeout = timeoutDeadline.Sub(time.Now())
-
 	response, _, err := client.Exchange(msg, target)
 	if err != nil {
 		log.Warnf("Error while sending a DNS query: %s", err)
diff --git a/testdata/blackbox-good.yml b/testdata/blackbox-good.yml
index 9025394..1f145d2 100644
--- a/testdata/blackbox-good.yml
+++ b/testdata/blackbox-good.yml
@@ -47,6 +47,7 @@ modules:
     prober: dns
     timeout: 5s
     dns:
-      preferred_ip_protocol: ip6
+      query_name: example.com
+      preferred_ip_protocol: ip4
       validate_answer_rrs:
         fail_if_matches_regexp: [test]
diff --git a/testdata/invalid-dns-module.yml b/testdata/invalid-dns-module.yml
new file mode 100644
index 0000000..ae91e66
--- /dev/null
+++ b/testdata/invalid-dns-module.yml
@@ -0,0 +1,8 @@
+modules:
+  dns_test:
+    prober: dns
+    timeout: 5s
+    dns:
+      preferred_ip_protocol: ip6
+      validate_answer_rrs:
+        fail_if_matches_regexp: [test]
-- 
2.25.1