From a990e60193565d8433e21141e182093b3b602763 Mon Sep 17 00:00:00 2001
From: Paul Tyng <ptyng@underarmour.com>
Date: Mon, 19 Jun 2017 08:23:42 -0400
Subject: [PATCH] Set default User-Agent on http probe

If unspecified in probe definition, User-Agent will be set to
`kube-probe/<version major.minor>` on httpGet probe types
instead of the default Go User-Agent.
---
 pkg/probe/http/BUILD        |  1 +
 pkg/probe/http/http.go      |  9 +++++++
 pkg/probe/http/http_test.go | 50 +++++++++++++++++++++++++++++--------
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/pkg/probe/http/BUILD b/pkg/probe/http/BUILD
index c7e22d2c769..8109f37279f 100644
--- a/pkg/probe/http/BUILD
+++ b/pkg/probe/http/BUILD
@@ -14,6 +14,7 @@ go_library(
     tags = ["automanaged"],
     deps = [
         "//pkg/probe:go_default_library",
+        "//pkg/version:go_default_library",
         "//vendor/github.com/golang/glog:go_default_library",
         "//vendor/k8s.io/apimachinery/pkg/util/net:go_default_library",
     ],
diff --git a/pkg/probe/http/http.go b/pkg/probe/http/http.go
index d0744e792ea..b9821be05fc 100644
--- a/pkg/probe/http/http.go
+++ b/pkg/probe/http/http.go
@@ -26,6 +26,7 @@ import (
 
 	utilnet "k8s.io/apimachinery/pkg/util/net"
 	"k8s.io/kubernetes/pkg/probe"
+	"k8s.io/kubernetes/pkg/version"
 
 	"github.com/golang/glog"
 )
@@ -68,6 +69,14 @@ func DoHTTPProbe(url *url.URL, headers http.Header, client HTTPGetInterface) (pr
 		// Convert errors into failures to catch timeouts.
 		return probe.Failure, err.Error(), nil
 	}
+	if _, ok := headers["User-Agent"]; !ok {
+		if headers == nil {
+			headers = http.Header{}
+		}
+		// explicitly set User-Agent so it's not set to default Go value
+		v := version.Get()
+		headers.Set("User-Agent", fmt.Sprintf("kube-probe/%s.%s", v.Major, v.Minor))
+	}
 	req.Header = headers
 	if headers.Get("Host") != "" {
 		req.Host = headers.Get("Host")
diff --git a/pkg/probe/http/http_test.go b/pkg/probe/http/http_test.go
index c1cf9a2c2c7..6e9efb86d29 100644
--- a/pkg/probe/http/http_test.go
+++ b/pkg/probe/http/http_test.go
@@ -40,12 +40,25 @@ func TestHTTPProbeChecker(t *testing.T) {
 		}
 	}
 
+	// Echo handler that returns the contents of request headers in the body
+	headerEchoHandler := func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		output := ""
+		for k, arr := range r.Header {
+			for _, v := range arr {
+				output += fmt.Sprintf("%s: %s\n", k, v)
+			}
+		}
+		w.Write([]byte(output))
+	}
+
 	prober := New()
 	testCases := []struct {
 		handler    func(w http.ResponseWriter, r *http.Request)
 		reqHeaders http.Header
 		health     probe.Result
 		accBody    string
+		notBody    string
 	}{
 		// The probe will be filled in below.  This is primarily testing that an HTTP GET happens.
 		{
@@ -54,23 +67,35 @@ func TestHTTPProbeChecker(t *testing.T) {
 			accBody: "ok body",
 		},
 		{
-			// Echo handler that returns the contents of request headers in the body
-			handler: func(w http.ResponseWriter, r *http.Request) {
-				w.WriteHeader(200)
-				output := ""
-				for k, arr := range r.Header {
-					for _, v := range arr {
-						output += fmt.Sprintf("%s: %s\n", k, v)
-					}
-				}
-				w.Write([]byte(output))
-			},
+			handler: headerEchoHandler,
 			reqHeaders: http.Header{
 				"X-Muffins-Or-Cupcakes": {"muffins"},
 			},
 			health:  probe.Success,
 			accBody: "X-Muffins-Or-Cupcakes: muffins",
 		},
+		{
+			handler: headerEchoHandler,
+			reqHeaders: http.Header{
+				"User-Agent": {"foo/1.0"},
+			},
+			health:  probe.Success,
+			accBody: "User-Agent: foo/1.0",
+		},
+		{
+			handler: headerEchoHandler,
+			reqHeaders: http.Header{
+				"User-Agent": {""},
+			},
+			health:  probe.Success,
+			notBody: "User-Agent",
+		},
+		{
+			handler:    headerEchoHandler,
+			reqHeaders: http.Header{},
+			health:     probe.Success,
+			accBody:    "User-Agent: kube-probe/",
+		},
 		{
 			// Echo handler that returns the contents of Host in the body
 			handler: func(w http.ResponseWriter, r *http.Request) {
@@ -130,6 +155,9 @@ func TestHTTPProbeChecker(t *testing.T) {
 				if !strings.Contains(output, test.accBody) {
 					t.Errorf("Expected response body to contain %v, got %v", test.accBody, output)
 				}
+				if test.notBody != "" && strings.Contains(output, test.notBody) {
+					t.Errorf("Expected response not to contain %v, got %v", test.notBody, output)
+				}
 			}
 		}()
 	}