agent: Allow custom HTTP headers and method for checks (#3106)

This patch adds support for custom headers and method for HTTP checks. Fixes #2474 Fixes #2657 Fixes #3106
hashicorp · Jun 6, 2017 · 6f2ab43 · 6f2ab43
1 parent 2b4f31d
commit 6f2ab43
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 14 deletions.
diff --git a/api/agent.go b/api/agent.go
@@ -67,17 +67,19 @@ type AgentCheckRegistration struct {
 
 // AgentServiceCheck is used to define a node or service level check
 type AgentServiceCheck struct {
-	Script            string `json:",omitempty"`
-	DockerContainerID string `json:",omitempty"`
-	Shell             string `json:",omitempty"` // Only supported for Docker.
-	Interval          string `json:",omitempty"`
-	Timeout           string `json:",omitempty"`
-	TTL               string `json:",omitempty"`
-	HTTP              string `json:",omitempty"`
-	TCP               string `json:",omitempty"`
-	Status            string `json:",omitempty"`
-	Notes             string `json:",omitempty"`
-	TLSSkipVerify     bool   `json:",omitempty"`
+	Script            string              `json:",omitempty"`
+	DockerContainerID string              `json:",omitempty"`
+	Shell             string              `json:",omitempty"` // Only supported for Docker.
+	Interval          string              `json:",omitempty"`
+	Timeout           string              `json:",omitempty"`
+	TTL               string              `json:",omitempty"`
+	HTTP              string              `json:",omitempty"`
+	Header            map[string][]string `json:",omitempty"`
+	Method            string              `json:",omitempty"`
+	TCP               string              `json:",omitempty"`
+	Status            string              `json:",omitempty"`
+	Notes             string              `json:",omitempty"`
+	TLSSkipVerify     bool                `json:",omitempty"`
 
 	// In Consul 0.7 and later, checks that are associated with a service
 	// may also contain this optional DeregisterCriticalServiceAfter field,

diff --git a/command/agent/agent.go b/command/agent/agent.go
@@ -1575,6 +1575,8 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *CheckType, persist
 				Notify:        &a.state,
 				CheckID:       check.CheckID,
 				HTTP:          chkType.HTTP,
+				Header:        chkType.Header,
+				Method:        chkType.Method,
 				Interval:      chkType.Interval,
 				Timeout:       chkType.Timeout,
 				Logger:        a.logger,

diff --git a/command/agent/check.go b/command/agent/check.go
@@ -57,6 +57,8 @@ type CheckType struct {
 
 	Script            string
 	HTTP              string
+	Header            map[string][]string
+	Method            string
 	TCP               string
 	Interval          time.Duration
 	DockerContainerID string
@@ -347,6 +349,8 @@ type CheckHTTP struct {
 	Notify        CheckNotifier
 	CheckID       types.CheckID
 	HTTP          string
+	Header        map[string][]string
+	Method        string
 	Interval      time.Duration
 	Timeout       time.Duration
 	Logger        *log.Logger
@@ -429,15 +433,25 @@ func (c *CheckHTTP) run() {
 
 // check is invoked periodically to perform the HTTP check
 func (c *CheckHTTP) check() {
-	req, err := http.NewRequest("GET", c.HTTP, nil)
+	method := c.Method
+	if method == "" {
+		method = "GET"
+	}
+
+	req, err := http.NewRequest(method, c.HTTP, nil)
 	if err != nil {
 		c.Logger.Printf("[WARN] agent: http request failed '%s': %s", c.HTTP, err)
 		c.Notify.UpdateCheck(c.CheckID, api.HealthCritical, err.Error())
 		return
 	}
 
-	req.Header.Set("User-Agent", UserAgent)
-	req.Header.Set("Accept", "text/plain, text/*, */*")
+	req.Header = http.Header(c.Header)
+	if req.Header.Get("User-Agent") == "" {
+		req.Header.Set("User-Agent", UserAgent)
+	}
+	if req.Header.Get("Accept") == "" {
+		req.Header.Set("Accept", "text/plain, text/*, */*")
+	}
 
 	resp, err := c.httpClient.Do(req)
 	if err != nil {

diff --git a/command/agent/config.go b/command/agent/config.go
@@ -1437,6 +1437,20 @@ func FixupCheckType(raw interface{}) error {
 		case "service_id":
 			rawMap["serviceid"] = v
 			delete(rawMap, k)
+		case "header":
+			vm, ok := v.(map[string]interface{})
+			if ok {
+				m := map[string][]string{}
+				for k, vv := range vm {
+					vs, ok := vv.([]interface{})
+					if ok {
+						for _, s := range vs {
+							m[k] = append(m[k], s.(string))
+						}
+					}
+				}
+				rawMap["header"] = m
+			}
 		case "docker_container_id":
 			rawMap["DockerContainerID"] = v
 			delete(rawMap, k)

diff --git a/command/agent/config_test.go b/command/agent/config_test.go
@@ -724,6 +724,8 @@ func TestDecodeConfig(t *testing.T) {
 							"Notes": "j",
 							"Script": "k",
 							"HTTP": "l",
+							"Header": {"a":["b"], "c":["d", "e"]},
+							"Method": "x",
 							"TCP": "m",
 							"DockerContainerID": "n",
 							"Shell": "o",
@@ -752,6 +754,8 @@ func TestDecodeConfig(t *testing.T) {
 							Notes:             "j",
 							Script:            "k",
 							HTTP:              "l",
+							Header:            map[string][]string{"a": []string{"b"}, "c": []string{"d", "e"}},
+							Method:            "x",
 							TCP:               "m",
 							DockerContainerID: "n",
 							Shell:             "o",
@@ -784,6 +788,8 @@ func TestDecodeConfig(t *testing.T) {
 								"Notes": "j",
 								"Script": "k",
 								"HTTP": "l",
+								"Header": {"a":["b"], "c":["d", "e"]},
+								"Method": "x",
 								"TCP": "m",
 								"DockerContainerID": "n",
 								"Shell": "o",
@@ -800,6 +806,8 @@ func TestDecodeConfig(t *testing.T) {
 								"Notes": "jj",
 								"Script": "kk",
 								"HTTP": "ll",
+								"Header": {"aa":["bb"], "cc":["dd", "ee"]},
+								"Method": "xx",
 								"TCP": "mm",
 								"DockerContainerID": "nn",
 								"Shell": "oo",
@@ -830,6 +838,8 @@ func TestDecodeConfig(t *testing.T) {
 								Notes:             "j",
 								Script:            "k",
 								HTTP:              "l",
+								Header:            map[string][]string{"a": []string{"b"}, "c": []string{"d", "e"}},
+								Method:            "x",
 								TCP:               "m",
 								DockerContainerID: "n",
 								Shell:             "o",
@@ -846,6 +856,8 @@ func TestDecodeConfig(t *testing.T) {
 								Notes:             "jj",
 								Script:            "kk",
 								HTTP:              "ll",
+								Header:            map[string][]string{"aa": []string{"bb"}, "cc": []string{"dd", "ee"}},
+								Method:            "xx",
 								TCP:               "mm",
 								DockerContainerID: "nn",
 								Shell:             "oo",
@@ -879,6 +891,8 @@ func TestDecodeConfig(t *testing.T) {
 								"Notes": "j",
 								"Script": "k",
 								"HTTP": "l",
+								"Header": {"a":["b"], "c":["d", "e"]},
+								"Method": "x",
 								"TCP": "m",
 								"DockerContainerID": "n",
 								"Shell": "o",
@@ -904,6 +918,8 @@ func TestDecodeConfig(t *testing.T) {
 								"Notes": "jj",
 								"Script": "kk",
 								"HTTP": "ll",
+								"Header": {"aa":["bb"], "cc":["dd", "ee"]},
+								"Method": "xx",
 								"TCP": "mm",
 								"DockerContainerID": "nn",
 								"Shell": "oo",
@@ -933,6 +949,8 @@ func TestDecodeConfig(t *testing.T) {
 							Notes:             "j",
 							Script:            "k",
 							HTTP:              "l",
+							Header:            map[string][]string{"a": []string{"b"}, "c": []string{"d", "e"}},
+							Method:            "x",
 							TCP:               "m",
 							DockerContainerID: "n",
 							Shell:             "o",
@@ -958,6 +976,8 @@ func TestDecodeConfig(t *testing.T) {
 							Notes:             "jj",
 							Script:            "kk",
 							HTTP:              "ll",
+							Header:            map[string][]string{"aa": []string{"bb"}, "cc": []string{"dd", "ee"}},
+							Method:            "xx",
 							TCP:               "mm",
 							DockerContainerID: "nn",
 							Shell:             "oo",
@@ -985,6 +1005,8 @@ func TestDecodeConfig(t *testing.T) {
 						"script": "d",
 						"shell": "e",
 						"http": "f",
+						"Header": {"a":["b"], "c":["d", "e"]},
+						"Method": "x",
 						"tcp": "g",
 						"docker_container_id": "h",
 						"tls_skip_verify": true,
@@ -1006,6 +1028,8 @@ func TestDecodeConfig(t *testing.T) {
 						Script:            "d",
 						Shell:             "e",
 						HTTP:              "f",
+						Header:            map[string][]string{"a": []string{"b"}, "c": []string{"d", "e"}},
+						Method:            "x",
 						TCP:               "g",
 						DockerContainerID: "h",
 						TLSSkipVerify:     true,
@@ -1031,6 +1055,8 @@ func TestDecodeConfig(t *testing.T) {
 							"script": "g",
 							"shell": "h",
 							"http": "i",
+							"Header": {"a":["b"], "c":["d", "e"]},
+							"Method": "x",
 							"tcp": "j",
 							"docker_container_id": "k",
 							"tls_skip_verify": true,
@@ -1049,6 +1075,8 @@ func TestDecodeConfig(t *testing.T) {
 							"script": "gg",
 							"shell": "hh",
 							"http": "ii",
+							"Header": {"aa":["bb"], "cc":["dd", "ee"]},
+							"Method": "xx",
 							"tcp": "jj",
 							"docker_container_id": "kk",
 							"tls_skip_verify": false,
@@ -1071,6 +1099,8 @@ func TestDecodeConfig(t *testing.T) {
 						Script:            "g",
 						Shell:             "h",
 						HTTP:              "i",
+						Header:            map[string][]string{"a": []string{"b"}, "c": []string{"d", "e"}},
+						Method:            "x",
 						TCP:               "j",
 						DockerContainerID: "k",
 						TLSSkipVerify:     true,
@@ -1089,6 +1119,8 @@ func TestDecodeConfig(t *testing.T) {
 						Script:            "gg",
 						Shell:             "hh",
 						HTTP:              "ii",
+						Header:            map[string][]string{"aa": []string{"bb"}, "cc": []string{"dd", "ee"}},
+						Method:            "xx",
 						TCP:               "jj",
 						DockerContainerID: "kk",
 						TLSSkipVerify:     false,

diff --git a/command/agent/structs.go b/command/agent/structs.go
@@ -62,6 +62,8 @@ type CheckDefinition struct {
 	//
 	Script                         string
 	HTTP                           string
+	Header                         map[string][]string
+	Method                         string
 	TCP                            string
 	Interval                       time.Duration
 	DockerContainerID              string