[probes.http] Allow setting request scheme through target labels (clo…

…udprober#608)
jumpojoy · Oct 28, 2023 · d21462c · d21462c
1 parent 0684857
commit d21462c
Show file tree

Hide file tree

Showing 6 changed files with 217 additions and 137 deletions.
diff --git a/probes/http/http.go b/probes/http/http.go
@@ -67,11 +67,10 @@ type Probe struct {
 	redirectFunc  func(req *http.Request, via []*http.Request) error
 
 	// book-keeping params
-	targets  []endpoint.Endpoint
-	protocol string
-	method   string
-	url      string
-	oauthTS  oauth2.TokenSource
+	targets []endpoint.Endpoint
+	method  string
+	url     string
+	oauthTS oauth2.TokenSource
 
 	// How often to resolve targets (in probe counts), it's the minimum of
 	targetsUpdateInterval time.Duration
@@ -184,7 +183,6 @@ func (p *Probe) Init(name string, opts *options.Options) error {
 			totalDuration, p.opts.Interval)
 	}
 
-	p.protocol = strings.ToLower(p.c.GetProtocol().String())
 	p.method = p.c.GetMethod().String()
 
 	p.url = p.c.GetRelativeUrl()
@@ -433,14 +431,19 @@ 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.resolveFirst(target) {
+
+			// If we're resolving target first, url.Host will be an IP address.
+			// In that case, we need to set ServerName in TLSClientConfig to
+			// the actual hostname.
+			if p.schemeForTarget(target) == "https" && p.resolveFirst(target) {
 				if t.TLSClientConfig == nil {
 					t.TLSClientConfig = &tls.Config{}
 				}
 				if t.TLSClientConfig.ServerName == "" {
 					t.TLSClientConfig.ServerName = hostForTarget(target)
 				}
 			}
+
 			clients[i] = &http.Client{Transport: t}
 		} else {
 			clients[i] = &http.Client{Transport: p.baseTransport}

diff --git a/probes/http/http_test.go b/probes/http/http_test.go
@@ -205,7 +205,6 @@ func TestProbeVariousMethods(t *testing.T) {
 		want  string
 	}{
 		{&configpb.ProbeConf{}, "total: 1, success: 1"},
-		{&configpb.ProbeConf{Protocol: configpb.ProbeConf_HTTPS.Enum()}, "total: 1, success: 1"},
 		{&configpb.ProbeConf{RequestsPerProbe: proto.Int32(1)}, "total: 1, success: 1"},
 		{&configpb.ProbeConf{RequestsPerProbe: proto.Int32(4)}, "total: 4, success: 4"},
 		{&configpb.ProbeConf{Method: mpb("GET")}, "total: 1, success: 1"},
@@ -776,6 +775,7 @@ func TestClientsForTarget(t *testing.T) {
 	tests := []struct {
 		name                string
 		conf                *configpb.ProbeConf
+		https               bool
 		baseTransport       *http.Transport
 		target              endpoint.Endpoint
 		wantNumClients      int
@@ -798,20 +798,20 @@ func TestClientsForTarget(t *testing.T) {
 		{
 			name: "2_clients_https",
 			conf: &configpb.ProbeConf{
-				Protocol:         configpb.ProbeConf_HTTPS.Enum(),
 				RequestsPerProbe: proto.Int32(2),
 			},
+			https:          true,
 			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),
+			https:         true,
 			target: endpoint.Endpoint{
 				Name: "cloudprober.org",
 				IP:   net.ParseIP("1.2.3.4"),
@@ -822,15 +822,15 @@ func TestClientsForTarget(t *testing.T) {
 		{
 			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",
+					"__cp_scheme__": "https", // Note: using target label here.
+					"fqdn":          "manugarg.com",
 				},
 			},
 			wantNumClients:      2,
@@ -840,6 +840,11 @@ func TestClientsForTarget(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			if tt.https {
+				tt.conf.SchemeType = &configpb.ProbeConf_Scheme_{
+					Scheme: configpb.ProbeConf_HTTPS,
+				}
+			}
 			p := &Probe{
 				baseTransport: tt.baseTransport,
 				c:             tt.conf,