From ca416a2ebbacad164dcea3243f4e3d6e24ecf3f8 Mon Sep 17 00:00:00 2001 From: Joel Speed Date: Sun, 14 Jun 2020 16:42:05 +0100 Subject: [PATCH] Add HealthCheck middleware --- go.mod | 1 + go.sum | 3 + pkg/middleware/healthcheck.go | 48 ++++++++++ pkg/middleware/healthcheck_test.go | 112 ++++++++++++++++++++++++ pkg/middleware/middleware_suite_test.go | 16 ++++ 5 files changed, 180 insertions(+) create mode 100644 pkg/middleware/healthcheck.go create mode 100644 pkg/middleware/healthcheck_test.go create mode 100644 pkg/middleware/middleware_suite_test.go diff --git a/go.mod b/go.mod index 62f41227..00d57333 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 098d3828..dd62bce2 100644 --- a/go.sum +++ b/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= diff --git a/pkg/middleware/healthcheck.go b/pkg/middleware/healthcheck.go new file mode 100644 index 00000000..ea1b533a --- /dev/null +++ b/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 +} diff --git a/pkg/middleware/healthcheck_test.go b/pkg/middleware/healthcheck_test.go new file mode 100644 index 00000000..8db4d57b --- /dev/null +++ b/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", + }), + ) +}) diff --git a/pkg/middleware/middleware_suite_test.go b/pkg/middleware/middleware_suite_test.go new file mode 100644 index 00000000..9972ce8f --- /dev/null +++ b/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") +}