[probes.http] Fix how we determine whether to change TLS config serve…

…rname.
cloudprober · Jul 12, 2023 · d725262 · d725262
1 parent e02536a
commit d725262
Show file tree

Hide file tree

Showing 4 changed files with 139 additions and 8 deletions.
diff --git a/probes/http/http.go b/probes/http/http.go
@@ -446,7 +446,7 @@ func (p *Probe) clientsForTarget(target endpoint.Endpoint) []*http.Client {
 		// RoundTripper implementation.
 		if ht, ok := p.baseTransport.(*http.Transport); ok {
 			t := ht.Clone()
-			if p.c.GetProtocol() == configpb.ProbeConf_HTTPS && p.c.GetResolveFirst() {
+			if p.c.GetProtocol() == configpb.ProbeConf_HTTPS && p.resolveFirst(target) {
 				if t.TLSClientConfig == nil {
 					t.TLSClientConfig = &tls.Config{}
 				}

diff --git a/probes/http/http_test.go b/probes/http/http_test.go
@@ -771,3 +771,86 @@ func TestProbeInitRedirectsNotSet(t *testing.T) {
 		t.Errorf("expected redirectFunc to be nil, found redirectFunc was initialized")
 	}
 }
+
+func TestClientsForTarget(t *testing.T) {
+	tests := []struct {
+		name                string
+		conf                *configpb.ProbeConf
+		baseTransport       *http.Transport
+		target              endpoint.Endpoint
+		wantNumClients      int
+		tlsConfigServerName string
+	}{
+		{
+			name:           "default",
+			conf:           &configpb.ProbeConf{},
+			baseTransport:  http.DefaultTransport.(*http.Transport),
+			target:         endpoint.Endpoint{Name: "cloudprober.org"},
+			wantNumClients: 1,
+		},
+		{
+			name:           "2_clients",
+			conf:           &configpb.ProbeConf{RequestsPerProbe: proto.Int32(2)},
+			baseTransport:  http.DefaultTransport.(*http.Transport),
+			target:         endpoint.Endpoint{Name: "cloudprober.org"},
+			wantNumClients: 2,
+		},
+		{
+			name: "2_clients_https",
+			conf: &configpb.ProbeConf{
+				Protocol:         configpb.ProbeConf_HTTPS.Enum(),
+				RequestsPerProbe: proto.Int32(2),
+			},
+			baseTransport:  http.DefaultTransport.(*http.Transport),
+			target:         endpoint.Endpoint{Name: "cloudprober.org"},
+			wantNumClients: 2,
+		},
+		{
+			name: "2_clients_https_server_name",
+			conf: &configpb.ProbeConf{
+				Protocol:         configpb.ProbeConf_HTTPS.Enum(),
+				RequestsPerProbe: proto.Int32(2),
+			},
+			baseTransport: http.DefaultTransport.(*http.Transport),
+			target: endpoint.Endpoint{
+				Name: "cloudprober.org",
+				IP:   net.ParseIP("1.2.3.4"),
+			},
+			wantNumClients:      2,
+			tlsConfigServerName: "cloudprober.org",
+		},
+		{
+			name: "2_clients_https_server_name_from_fqdn",
+			conf: &configpb.ProbeConf{
+				Protocol:         configpb.ProbeConf_HTTPS.Enum(),
+				RequestsPerProbe: proto.Int32(2),
+			},
+			baseTransport: http.DefaultTransport.(*http.Transport),
+			target: endpoint.Endpoint{
+				Name: "cloudprober.org",
+				IP:   net.ParseIP("1.2.3.4"),
+				Labels: map[string]string{
+					"fqdn": "manugarg.com",
+				},
+			},
+			wantNumClients:      2,
+			tlsConfigServerName: "manugarg.com",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := &Probe{
+				baseTransport: tt.baseTransport,
+				c:             tt.conf,
+			}
+			gotClients := p.clientsForTarget(tt.target)
+			assert.Equal(t, tt.wantNumClients, len(gotClients), "number of clients is not as expected")
+
+			for _, c := range gotClients {
+				tlsConfig := c.Transport.(*http.Transport).TLSClientConfig
+				assert.Equal(t, tt.tlsConfigServerName, tlsConfig.ServerName, "TLS config server name is not as expected")
+			}
+		})
+	}
+}
diff --git a/probes/http/request.go b/probes/http/request.go
@@ -72,6 +72,13 @@ func relURLForTarget(target endpoint.Endpoint, probeURL string) string {
 	return ""
 }
 
+func (p *Probe) resolveFirst(target endpoint.Endpoint) bool {
+	if p.c.ResolveFirst != nil {
+		return p.c.GetResolveFirst()
+	}
+	return target.IP != nil
+}
+
 func (p *Probe) httpRequestForTarget(target endpoint.Endpoint) *http.Request {
 	// Prepare HTTP.Request for Client.Do
 	port := int(p.c.GetPort())
@@ -83,13 +90,7 @@ func (p *Probe) httpRequestForTarget(target endpoint.Endpoint) *http.Request {
 	urlHost := urlHostForTarget(target)
 	ipForLabel := ""
 
-	resolveFirst := false
-	if p.c.ResolveFirst != nil {
-		resolveFirst = p.c.GetResolveFirst()
-	} else {
-		resolveFirst = target.IP != nil
-	}
-	if resolveFirst {
+	if p.resolveFirst(target) {
 		ip, err := target.Resolve(p.opts.IPVersion, p.opts.Targets)
 		if err != nil {
 			p.l.Error("target: ", target.Name, ", resolve error: ", err.Error())

diff --git a/probes/http/request_test.go b/probes/http/request_test.go
@@ -434,5 +434,52 @@ func TestRequestHasConfiguredHeaders(t *testing.T) {
 	val, ok = req.Header[testHeadersName]
 	assert.True(t, ok, "Configured header (via 'headers' setting) is not present in target request")
 	assert.Contains(t, val, testHeadersValue)
+}
 
+func TestResolveFirst(t *testing.T) {
+	tests := []struct {
+		name   string
+		target endpoint.Endpoint
+		conf   *configpb.ProbeConf
+		want   bool
+	}{
+		{
+			name:   "resolve_first_false",
+			target: endpoint.Endpoint{Name: "cloudprober.org"},
+			conf:   &configpb.ProbeConf{},
+			want:   false,
+		},
+		{
+			name:   "resolve_first_true",
+			target: endpoint.Endpoint{Name: "cloudprober.org"},
+			conf:   &configpb.ProbeConf{ResolveFirst: proto.Bool(true)},
+			want:   true,
+		},
+		{
+			name: "resolve_first_true_with_ip",
+			target: endpoint.Endpoint{
+				Name: "cloudprober.org",
+				IP:   net.ParseIP("1.1.1.1"),
+			},
+			conf: &configpb.ProbeConf{},
+			want: true,
+		},
+		{
+			name: "resolve_first_explcitly_false",
+			target: endpoint.Endpoint{
+				Name: "cloudprober.org",
+				IP:   net.ParseIP("1.1.1.1"),
+			},
+			conf: &configpb.ProbeConf{ResolveFirst: proto.Bool(false)},
+			want: false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			p := &Probe{
+				c: tt.conf,
+			}
+			assert.Equal(t, tt.want, p.resolveFirst(tt.target), "resolveFirst is not as expected")
+		})
+	}
 }