From ff4e5588d87cf92b2df01b0ab16c7ebd8ec02d43 Mon Sep 17 00:00:00 2001
From: timothy-spencer <timothy.spencer@gsa.gov>
Date: Mon, 25 Mar 2019 10:32:29 -0700
Subject: [PATCH] incorporate suggestions from @benfdking

---
 http.go      | 24 +++++++++++++++++++++++-
 http_test.go | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/http.go b/http.go
index 3f9ba57..b45b3c1 100644
--- a/http.go
+++ b/http.go
@@ -24,9 +24,15 @@ func (s *Server) ListenAndServe() {
 	}
 }
 
-// gcpHealthcheck handles healthcheck queries from GCP
+// Used with gcpHealthcheck()
+const userAgentHeader = "User-Agent"
+const googleHealthCheckUserAgent = "GoogleHC/1.0"
+const rootPath = "/"
+
+// gcpHealthcheck handles healthcheck queries from GCP.
 func gcpHealthcheck(h http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Check for liveness and readiness:  used for Google App Engine
 		if r.URL.EscapedPath() == "/liveness_check" {
 			w.WriteHeader(http.StatusOK)
 			w.Write([]byte("OK"))
@@ -37,6 +43,22 @@ func gcpHealthcheck(h http.Handler) http.Handler {
 			w.Write([]byte("OK"))
 			return
 		}
+
+		// Check for GKE ingress healthcheck:  The ingress requires the root
+		// path of the target to return a 200 (OK) to indicate the service's good health. This can be quite a challenging demand
+		// depending on the application's path structure. This middleware filters out the requests from the health check by
+		//
+		// 1. checking that the request path is indeed the root path
+		// 2. ensuring that the User-Agent is "GoogleHC/1.0", the health checker
+		// 3. ensuring the request method is "GET"
+		if r.URL.Path == rootPath &&
+			r.Header.Get(userAgentHeader) == googleHealthCheckUserAgent &&
+			r.Method == http.MethodGet {
+
+			w.WriteHeader(http.StatusOK)
+			return
+		}
+
 		h.ServeHTTP(w, r)
 	})
 }
diff --git a/http_test.go b/http_test.go
index 8495fa8..2bc3d5d 100644
--- a/http_test.go
+++ b/http_test.go
@@ -54,3 +54,36 @@ func TestGCPHealthcheckNotHealthcheck(t *testing.T) {
 
 	assert.Equal(t, "test", rw.Body.String())
 }
+
+func TestGCPHealthcheckIngress(t *testing.T) {
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("test"))
+	}
+
+	h := gcpHealthcheck(http.HandlerFunc(handler))
+	rw := httptest.NewRecorder()
+	r, _ := http.NewRequest("GET", "/", nil)
+	r.RemoteAddr = "127.0.0.1"
+	r.Host = "test-server"
+	r.Header.Set(userAgentHeader, googleHealthCheckUserAgent)
+	h.ServeHTTP(rw, r)
+
+	assert.Equal(t, 200, rw.Code)
+	assert.Equal(t, "", rw.Body.String())
+}
+
+func TestGCPHealthcheckNotIngress(t *testing.T) {
+	handler := func(w http.ResponseWriter, req *http.Request) {
+		w.Write([]byte("test"))
+	}
+
+	h := gcpHealthcheck(http.HandlerFunc(handler))
+	rw := httptest.NewRecorder()
+	r, _ := http.NewRequest("GET", "/foo", nil)
+	r.RemoteAddr = "127.0.0.1"
+	r.Host = "test-server"
+	r.Header.Set(userAgentHeader, googleHealthCheckUserAgent)
+	h.ServeHTTP(rw, r)
+
+	assert.Equal(t, "test", rw.Body.String())
+}