Add HealthCheck middleware

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")
+}