Add HealthCheck middleware

go.mod

@@ -12,6 +12,7 @@ require (
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/fsnotify/fsnotify v1.4.9
 	github.com/go-redis/redis/v7 v7.2.0
+	github.com/justinas/alice v1.2.0
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/mbland/hmacauth v0.0.0-20170912233209-44256dfd4bfa
 	github.com/mitchellh/mapstructure v1.1.2

go.sum

@@ -102,6 +102,8 @@ github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1
 github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo=
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/julienschmidt/httprouter v1.2.0/go.mod h1:SYymIcj16QtmaHHD7aYtjjsJG7VTCxuUUipMqKk8s4w=
+github.com/justinas/alice v1.2.0 h1:+MHSA/vccVCF4Uq37S42jwlkvI2Xzl7zTPCN5BnZNVo=
+github.com/justinas/alice v1.2.0/go.mod h1:fN5HRH/reO/zrUflLfTN43t3vXvKzvZIENsNEe7i7qA=
 github.com/kisielk/errcheck v1.1.0/go.mod h1:EZBBE59ingxPouuu3KfxchcWSUPOHkagtvWXihfKN4Q=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
@@ -129,6 +131,7 @@ github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+W
 github.com/onsi/ginkgo v1.10.1/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.12.0 h1:Iw5WCbBcaAAd0fpRb1c9r5YCylv4XDoCSigm1zLevwU=
 github.com/onsi/ginkgo v1.12.0/go.mod h1:oUhWkIvk5aDxtKvDDuw8gItl8pKl42LzjC9KZE0HfGg=
+github.com/onsi/ginkgo v1.12.3 h1:+RYp9QczoWz9zfUyLP/5SLXQVhfr6gZOoKGfQqHuLZQ=
 github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.9.0 h1:R1uwffexN6Pr340GtYRIdZmAiN4J+iw6WG4wog1DUXg=

pkg/middleware/healthcheck.go

@@ -0,0 +1,48 @@
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/justinas/alice"
+)
+
+func NewHealthCheck(paths, userAgents []string) alice.Constructor {
+	return func(next http.Handler) http.Handler {
+		return healthCheck(paths, userAgents, next)
+	}
+}
+
+func healthCheck(paths, userAgents []string, next http.Handler) http.Handler {
+	// Use a map as a set to check health check paths
+	pathSet := make(map[string]struct{})
+	for _, path := range paths {
+		pathSet[path] = struct{}{}
+	}
+
+	// Use a map as a set to check health check paths
+	userAgentSet := make(map[string]struct{})
+	for _, userAgent := range userAgents {
+		userAgentSet[userAgent] = struct{}{}
+	}
+
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		if isHealthCheckRequest(pathSet, userAgentSet, req) {
+			rw.WriteHeader(http.StatusOK)
+			fmt.Fprintf(rw, "OK")
+			return
+		}
+
+		next.ServeHTTP(rw, req)
+	})
+}
+
+func isHealthCheckRequest(paths, userAgents map[string]struct{}, req *http.Request) bool {
+	if _, ok := paths[req.URL.EscapedPath()]; ok {
+		return true
+	}
+	if _, ok := userAgents[req.Header.Get("User-Agent")]; ok {
+		return true
+	}
+	return false
+}

pkg/middleware/healthcheck_test.go

@@ -0,0 +1,112 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/ginkgo/extensions/table"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("HealthCheck suite", func() {
+	type requestTableInput struct {
+		healthCheckPaths      []string
+		healthCheckUserAgents []string
+		requestString         string
+		headers               map[string]string
+		expectedStatus        int
+		expectedBody          string
+	}
+
+	DescribeTable("when serving a request",
+		func(in *requestTableInput) {
+			req := httptest.NewRequest("", in.requestString, nil)
+			for k, v := range in.headers {
+				req.Header.Add(k, v)
+			}
+
+			rw := httptest.NewRecorder()
+
+			handler := NewHealthCheck(in.healthCheckPaths, in.healthCheckUserAgents)(http.NotFoundHandler())
+			handler.ServeHTTP(rw, req)
+
+			Expect(rw.Code).To(Equal(in.expectedStatus))
+			Expect(rw.Body.String()).To(Equal(in.expectedBody))
+		},
+		Entry("when requesting the healthcheck path", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/ping",
+			headers:               map[string]string{},
+			expectedStatus:        200,
+			expectedBody:          "OK",
+		}),
+		Entry("when requesting a different path", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/different",
+			headers:               map[string]string{},
+			expectedStatus:        404,
+			expectedBody:          "404 page not found\n",
+		}),
+		Entry("with a request from the health check user agent", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/abc",
+			headers: map[string]string{
+				"User-Agent": "hc/1.0",
+			},
+			expectedStatus: 200,
+			expectedBody:   "OK",
+		}),
+		Entry("with a request from a different user agent", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/abc",
+			headers: map[string]string{
+				"User-Agent": "different",
+			},
+			expectedStatus: 404,
+			expectedBody:   "404 page not found\n",
+		}),
+		Entry("with multiple paths, request one of the healthcheck paths", &requestTableInput{
+			healthCheckPaths:      []string{"/ping", "/liveness_check", "/readiness_check"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/readiness_check",
+			headers:               map[string]string{},
+			expectedStatus:        200,
+			expectedBody:          "OK",
+		}),
+		Entry("with multiple paths, request none of the healthcheck paths", &requestTableInput{
+			healthCheckPaths:      []string{"/ping", "/liveness_check", "/readiness_check"},
+			healthCheckUserAgents: []string{"hc/1.0"},
+			requestString:         "http://example.com/readiness",
+			headers: map[string]string{
+				"User-Agent": "user",
+			},
+			expectedStatus: 404,
+			expectedBody:   "404 page not found\n",
+		}),
+		Entry("with multiple user agents, request from a health check user agent", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0", "GoogleHC/1.0"},
+			requestString:         "http://example.com/abc",
+			headers: map[string]string{
+				"User-Agent": "GoogleHC/1.0",
+			},
+			expectedStatus: 200,
+			expectedBody:   "OK",
+		}),
+		Entry("with multiple user agents, request from none of the health check user agents", &requestTableInput{
+			healthCheckPaths:      []string{"/ping"},
+			healthCheckUserAgents: []string{"hc/1.0", "GoogleHC/1.0"},
+			requestString:         "http://example.com/abc",
+			headers: map[string]string{
+				"User-Agent": "user",
+			},
+			expectedStatus: 404,
+			expectedBody:   "404 page not found\n",
+		}),
+	)
+})

pkg/middleware/middleware_suite_test.go

@@ -0,0 +1,16 @@
+package middleware
+
+import (
+	"testing"
+
+	"github.com/oauth2-proxy/oauth2-proxy/pkg/logger"
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+func TestMiddlewareSuite(t *testing.T) {
+	logger.SetOutput(GinkgoWriter)
+
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Middleware")
+}