Add resolve_first functionality to DNS probes

Fixes cloudprober#369 PiperOrigin-RevId: 297769471
jumpojoy · Feb 28, 2020 · b68b665 · b68b665
1 parent 6291d83
commit b68b665
Show file tree

Hide file tree

Showing 6 changed files with 117 additions and 47 deletions.
diff --git a/go.mod b/go.mod
@@ -10,7 +10,7 @@ require (
 	github.com/aws/aws-sdk-go v1.25.37
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/golang/groupcache v0.0.0-20191027212112-611e8accdfc9 // indirect
-	github.com/golang/protobuf v1.3.3
+	github.com/golang/protobuf v1.3.4
 	github.com/google/go-cmp v0.3.1 // indirect
 	github.com/hashicorp/golang-lru v0.5.3 // indirect
 	github.com/hoisie/redis v0.0.0-20160730154456-b5c6e81454e0

diff --git a/go.sum b/go.sum
@@ -44,6 +44,8 @@ github.com/golang/protobuf v1.3.2 h1:6nsPYzhq5kReh6QImI3k5qWzO4PEbvbIW2cwSfR/6xs
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.3 h1:gyjaxf+svBWX08ZjK86iN9geUJF0H6gp2IRKX6Nf6/I=
 github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw=
+github.com/golang/protobuf v1.3.4 h1:87PNWwrRvUSnqS4dlcBU/ftvOIBep4sYuBLlh6rX2wk=
+github.com/golang/protobuf v1.3.4/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=

diff --git a/probes/dns/dns.go b/probes/dns/dns.go
@@ -197,7 +197,12 @@ func (p *Probe) validateResponse(resp *dns.Msg, target string, result *probeRunR
 	return true
 }
 
-func (p *Probe) runProbe(resultsChan chan<- statskeeper.ProbeResult) {
+// resolveFunc resolves the given host for the IP version.
+// This type is mainly used for testing. For all other cases, a nil function
+// should be passed to the runProbe function.
+type resolveFunc func(host string, ipVer int) (net.IP, error)
+
+func (p *Probe) runProbe(resultsChan chan<- statskeeper.ProbeResult, resolveF resolveFunc) {
 	// Refresh the list of targets to probe.
 	p.updateTargets()
 
@@ -221,8 +226,22 @@ func (p *Probe) runProbe(resultsChan chan<- statskeeper.ProbeResult) {
 				result.latency = metrics.NewFloat(0)
 			}
 
-			fullTarget := net.JoinHostPort(target.Name, "53")
 			result.total.Inc()
+
+			fullTarget := net.JoinHostPort(target.Name, "53")
+			if p.c.GetResolveFirst() {
+				if resolveF == nil {
+					resolveF = p.opts.Targets.Resolve
+				}
+				ip, err := resolveF(target.Name, p.opts.IPVersion)
+				if err != nil {
+					p.l.Warningf("Target(%s): Resolve error: %v", target.Name, err)
+					resultsChan <- result
+					return
+				}
+				fullTarget = net.JoinHostPort(ip.String(), "53")
+			}
+
 			resp, latency, err := p.client.Exchange(p.msg, fullTarget)
 
 			if err != nil {
@@ -266,6 +285,6 @@ func (p *Probe) Start(ctx context.Context, dataChan chan *metrics.EventMetrics)
 			return
 		default:
 		}
-		p.runProbe(resultsChan)
+		p.runProbe(resultsChan, nil)
 	}
 }
diff --git a/probes/dns/dns_test.go b/probes/dns/dns_test.go
@@ -15,6 +15,7 @@
 package dns
 
 import (
+	"fmt"
 	"net"
 	"testing"
 	"time"
@@ -49,7 +50,10 @@ type mockClient struct{}
 // Exchange implementation that returns an error status if the query is for
 // questionBad[Domain|Type]. This allows us to check if query parameters are
 // populated correctly.
-func (*mockClient) Exchange(in *dns.Msg, _ string) (*dns.Msg, time.Duration, error) {
+func (*mockClient) Exchange(in *dns.Msg, fullTarget string) (*dns.Msg, time.Duration, error) {
+	if fullTarget != "8.8.8.8:53" {
+		return nil, 0, fmt.Errorf("unexpected target: %v", fullTarget)
+	}
 	out := &dns.Msg{}
 	question := in.Question[0]
 	if question.Name == questionBadDomain+"." || int(question.Qtype) == int(questionBadType) {
@@ -67,12 +71,12 @@ func (*mockClient) Exchange(in *dns.Msg, _ string) (*dns.Msg, time.Duration, err
 func (*mockClient) setReadTimeout(time.Duration) {}
 func (*mockClient) setSourceIP(net.IP)           {}
 
-func runProbe(t *testing.T, testName string, p *Probe, total, success int64) {
+func runProbe(t *testing.T, testName string, p *Probe, resolveF resolveFunc, total, success int64) {
 	p.client = new(mockClient)
 	p.targets = p.opts.Targets.ListEndpoints()
 
 	resultsChan := make(chan statskeeper.ProbeResult, len(p.targets))
-	p.runProbe(resultsChan)
+	p.runProbe(resultsChan, resolveF)
 
 	// The resultsChan output iterates through p.targets in the same order.
 	for _, target := range p.targets {
@@ -100,7 +104,34 @@ func TestRun(t *testing.T) {
 	if err := p.Init("dns_test", opts); err != nil {
 		t.Fatalf("Error creating probe: %v", err)
 	}
-	runProbe(t, "basic", p, 1, 1)
+	runProbe(t, "basic", p, nil, 1, 1)
+}
+
+func TestResolveFirst(t *testing.T) {
+	p := &Probe{}
+	opts := options.DefaultOptions()
+	opts.Targets = targets.StaticTargets("foo")
+	opts.ProbeConf = &configpb.ProbeConf{ResolveFirst: proto.Bool(true)}
+	if err := p.Init("dns_test_resolve_first", opts); err != nil {
+		t.Fatalf("Error creating probe: %v", err)
+	}
+
+	t.Run("success", func(t *testing.T) {
+		resolveF := func(target string, ipVer int) (net.IP, error) {
+			if target == "foo" {
+				return net.ParseIP("8.8.8.8"), nil
+			}
+			return nil, fmt.Errorf("resolve error")
+		}
+		runProbe(t, "resolve_first_success", p, resolveF, 1, 1)
+	})
+
+	t.Run("error", func(t *testing.T) {
+		resolveF := func(target string, ipVer int) (net.IP, error) {
+			return nil, fmt.Errorf("resolve error")
+		}
+		runProbe(t, "resolve_first_error", p, resolveF, 1, 0)
+	})
 }
 
 func TestProbeType(t *testing.T) {
@@ -117,7 +148,7 @@ func TestProbeType(t *testing.T) {
 	if err := p.Init("dns_probe_type_test", opts); err != nil {
 		t.Fatalf("Error creating probe: %v", err)
 	}
-	runProbe(t, "probetype", p, 1, 0)
+	runProbe(t, "probetype", p, nil, 1, 0)
 }
 
 func TestBadName(t *testing.T) {
@@ -133,7 +164,7 @@ func TestBadName(t *testing.T) {
 	if err := p.Init("dns_bad_domain_test", opts); err != nil {
 		t.Fatalf("Error creating probe: %v", err)
 	}
-	runProbe(t, "baddomain", p, 1, 0)
+	runProbe(t, "baddomain", p, nil, 1, 0)
 }
 
 func TestAnswerCheck(t *testing.T) {
@@ -150,7 +181,7 @@ func TestAnswerCheck(t *testing.T) {
 		t.Fatalf("Error creating probe: %v", err)
 	}
 	// expect success minAnswers == num answers returned == 1.
-	runProbe(t, "matchminanswers", p, 1, 1)
+	runProbe(t, "matchminanswers", p, nil, 1, 1)
 
 	opts.ProbeConf = &configpb.ProbeConf{
 		MinAnswers: proto.Uint32(2),
@@ -159,7 +190,7 @@ func TestAnswerCheck(t *testing.T) {
 		t.Fatalf("Error creating probe: %v", err)
 	}
 	// expect failure because only one answer returned and two wanted.
-	runProbe(t, "toofewanswers", p, 1, 0)
+	runProbe(t, "toofewanswers", p, nil, 1, 0)
 }
 
 func TestValidator(t *testing.T) {
@@ -187,6 +218,6 @@ func TestValidator(t *testing.T) {
 		if err := p.Init("dns_probe_answer_"+tst.name, opts); err != nil {
 			t.Fatalf("Error creating probe: %v", err)
 		}
-		runProbe(t, tst.name, p, 1, tst.successCt)
+		runProbe(t, tst.name, p, nil, 1, tst.successCt)
 	}
 }
diff --git a/probes/dns/proto/config.pb.go b/probes/dns/proto/config.pb.go
diff --git a/probes/dns/proto/config.proto b/probes/dns/proto/config.proto
@@ -56,4 +56,9 @@ message ProbeConf {
   // Minimum number of answers expected. Default behavior is to return success
   // if DNS response status is NOERROR.
   optional uint32 min_answers = 4 [default = 0];
+
+  // Whether to resolve the target before making the request. If set to false,
+  // we hand over the target directly to the DNS client. Otherwise, we resolve
+  // the target first to an IP address.
+  optional bool resolve_first = 5 [default = false];
 }