From 52e4ab705c1f89e73efca2e91bb4f9a873440eb7 Mon Sep 17 00:00:00 2001
From: Seth Hoenig <seth.a.hoenig@gmail.com>
Date: Wed, 3 Aug 2022 10:06:35 -0500
Subject: [PATCH] nsd: add support for setting headers on nomad service http
 checks

This PR enables setting of the headers block on services registered
into Nomad's service provider. Works just like the existing support
in Consul checks.
---
 client/serviceregistration/checks/client.go   |  1 +
 .../serviceregistration/checks/client_test.go | 33 ++++++++++++-------
 client/serviceregistration/checks/result.go   | 12 ++++---
 nomad/structs/services.go                     |  5 ---
 nomad/structs/services_test.go                |  6 ++--
 5 files changed, 35 insertions(+), 22 deletions(-)

diff --git a/client/serviceregistration/checks/client.go b/client/serviceregistration/checks/client.go
index c26431436..b17b26603 100644
--- a/client/serviceregistration/checks/client.go
+++ b/client/serviceregistration/checks/client.go
@@ -162,6 +162,7 @@ func (c *checker) checkHTTP(ctx context.Context, qc *QueryContext, q *Query) *st
 		qr.Status = structs.CheckFailure
 		return qr
 	}
+	request.Header = q.Headers
 	request = request.WithContext(ctx)
 
 	result, err := c.httpClient.Do(request)
diff --git a/client/serviceregistration/checks/client_test.go b/client/serviceregistration/checks/client_test.go
index 82e589b11..71482ee28 100644
--- a/client/serviceregistration/checks/client_test.go
+++ b/client/serviceregistration/checks/client_test.go
@@ -261,6 +261,14 @@ func TestChecker_Do_HTTP_extras(t *testing.T) {
 			method:  "HEAD",
 			headers: makeHeaders(agent),
 		},
+		{
+			name:   "extra headers",
+			method: "GET",
+			headers: makeHeaders(encoding, agent,
+				[2]string{"X-My-Header", "hello"},
+				[2]string{"Authorization", "Basic ZWxhc3RpYzpjaGFuZ2VtZQ=="},
+			),
+		},
 	}
 
 	for _, tc := range cases {
@@ -286,19 +294,22 @@ func TestChecker_Do_HTTP_extras(t *testing.T) {
 			Protocol:    "http",
 			Path:        "/",
 			Method:      tc.method,
+			Headers:     tc.headers,
 		}
 
-		logger := testlog.HCLogger(t)
-		c := New(logger)
-		ctx := context.Background()
-		result := c.Do(ctx, qc, q)
-		must.Eq(t, http.StatusOK, result.StatusCode,
-			must.Sprintf("test.URL: %s", ts.URL),
-			must.Sprintf("headers: %v", tc.headers),
-		)
-		must.Eq(t, tc.method, method)
-		must.Eq(t, tc.body, string(body))
-		must.Eq(t, tc.headers, headers)
+		t.Run(tc.name, func(t *testing.T) {
+			logger := testlog.HCLogger(t)
+			c := New(logger)
+			ctx := context.Background()
+			result := c.Do(ctx, qc, q)
+			must.Eq(t, http.StatusOK, result.StatusCode,
+				must.Sprintf("test.URL: %s", ts.URL),
+				must.Sprintf("headers: %v", tc.headers),
+			)
+			must.Eq(t, tc.method, method)
+			must.Eq(t, tc.body, string(body))
+			must.Eq(t, tc.headers, headers)
+		})
 	}
 }
 
diff --git a/client/serviceregistration/checks/result.go b/client/serviceregistration/checks/result.go
index 8ef5859ce..07b78fb77 100644
--- a/client/serviceregistration/checks/result.go
+++ b/client/serviceregistration/checks/result.go
@@ -1,8 +1,10 @@
 package checks
 
 import (
+	"net/http"
 	"time"
 
+	"github.com/hashicorp/nomad/helper"
 	"github.com/hashicorp/nomad/nomad/structs"
 )
 
@@ -18,9 +20,10 @@ func GetCheckQuery(c *structs.ServiceCheck) *Query {
 		Timeout:     c.Timeout,
 		AddressMode: c.AddressMode,
 		PortLabel:   c.PortLabel,
+		Protocol:    protocol,
 		Path:        c.Path,
 		Method:      c.Method,
-		Protocol:    protocol,
+		Headers:     helper.CopyMap(c.Header),
 	}
 }
 
@@ -35,9 +38,10 @@ type Query struct {
 	AddressMode string // host, driver, or alloc
 	PortLabel   string // label or value
 
-	Protocol string // http checks only (http or https)
-	Path     string // http checks only
-	Method   string // http checks only
+	Protocol string      // http checks only (http or https)
+	Path     string      // http checks only
+	Method   string      // http checks only
+	Headers  http.Header // http checks only
 }
 
 // A QueryContext contains allocation and service parameters necessary for
diff --git a/nomad/structs/services.go b/nomad/structs/services.go
index 5d06ef9e1..f3ad50c65 100644
--- a/nomad/structs/services.go
+++ b/nomad/structs/services.go
@@ -351,11 +351,6 @@ func (sc *ServiceCheck) validateNomad() error {
 			return fmt.Errorf("method type %q not supported in Nomad http check", sc.Method)
 		}
 
-		// todo(shoenig) support headers
-		if len(sc.Header) > 0 {
-			return fmt.Errorf("http checks may not set headers in Nomad services")
-		}
-
 		// todo(shoenig) support body
 		if len(sc.Body) > 0 {
 			return fmt.Errorf("http checks may not set Body in Nomad services")
diff --git a/nomad/structs/services_test.go b/nomad/structs/services_test.go
index 3833f5c2b..8cdd01d30 100644
--- a/nomad/structs/services_test.go
+++ b/nomad/structs/services_test.go
@@ -323,9 +323,11 @@ func TestServiceCheck_validateNomad(t *testing.T) {
 				Timeout:  1 * time.Second,
 				Path:     "/health",
 				Method:   "GET",
-				Header:   map[string][]string{"foo": {"bar"}},
+				Header: map[string][]string{
+					"foo": {"bar"},
+					"baz": nil,
+				},
 			},
-			exp: `http checks may not set headers in Nomad services`,
 		},
 		{
 			name: "http with body",