Move healthcheck middleware to a separate package

add Cache-Control: no-store to status response

4 files changed, 77 insertions, 95 deletions

diff --git a/app.go b/app.go
index 4f187f9d..5aba60d3 100644
--- a/app.go
+++ b/app.go
@@ -30,6 +30,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/domain"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/errortracking"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/handlers"
+	health "gitlab.com/gitlab-org/gitlab-pages/internal/healthcheck"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/logging"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/netutil"
@@ -59,10 +60,6 @@ type theApp struct {
 	CustomHeaders  http.Header
 }
 
-func (a *theApp) isReady() bool {
-	return true
-}
-
 func (a *theApp) GetCertificate(ch *cryptotls.ClientHelloInfo) (*cryptotls.Certificate, error) {
 	if ch.ServerName == "" {
 		return nil, nil
@@ -127,11 +124,6 @@ func (a *theApp) tryAuxiliaryHandlers(w http.ResponseWriter, r *http.Request, ht
 		return true
 	}
 
-	if !a.isReady() {
-		httperrors.Serve503(w)
-		return true
-	}
-
 	if _, err := domain.GetLookupPath(r); err != nil {
 		if errors.Is(err, gitlab.ErrDiskDisabled) {
 			errortracking.CaptureErrWithReqAndStackTrace(err, r)
@@ -152,26 +144,6 @@ func (a *theApp) tryAuxiliaryHandlers(w http.ResponseWriter, r *http.Request, ht
 	return false
 }
 
-// healthCheckMiddleware is serving the application status check
-func (a *theApp) healthCheckMiddleware(handler http.Handler) http.Handler {
-	healthCheck := http.HandlerFunc(func(w http.ResponseWriter, _r *http.Request) {
-		if a.isReady() {
-			w.Write([]byte("success\n"))
-		} else {
-			http.Error(w, "not yet ready", http.StatusServiceUnavailable)
-		}
-	})
-
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if r.RequestURI == a.config.General.StatusPath {
-			healthCheck.ServeHTTP(w, r)
-			return
-		}
-
-		handler.ServeHTTP(w, r)
-	})
-}
-
 // auxiliaryMiddleware will handle status updates, not-ready requests and other
 // not static-content responses
 func (a *theApp) auxiliaryMiddleware(handler http.Handler) http.Handler {
@@ -250,7 +222,7 @@ func (a *theApp) buildHandlerPipeline() (http.Handler, error) {
 	handler = handlers.Ratelimiter(handler, &a.config.RateLimit)
 
 	// Health Check
-	handler = a.healthCheckMiddleware(handler)
+	handler = health.NewMiddleware(handler, a.config.General.StatusPath)
 
 	// Custom response headers
 	handler = customheaders.NewMiddleware(handler, a.CustomHeaders)
diff --git a/app_test.go b/app_test.go
index 3b1d6697..75747a2d 100644
--- a/app_test.go
+++ b/app_test.go
@@ -3,18 +3,14 @@ package main
 import (
 	"crypto/tls"
 	"fmt"
-	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
-	"time"
 
 	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/require"
 
-	"gitlab.com/gitlab-org/gitlab-pages/internal/config"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
-	"gitlab.com/gitlab-org/gitlab-pages/internal/source/gitlab"
 	"gitlab.com/gitlab-org/gitlab-pages/metrics"
 )
 
@@ -66,67 +62,6 @@ func newGetRequestWithScheme(t *testing.T, scheme string, withTLS bool) *http.Re
 	return req
 }
 
-func TestHealthCheckMiddleware(t *testing.T) {
-	tests := []struct {
-		name   string
-		path   string
-		status int
-		body   string
-	}{
-		{
-			name:   "Not a healthcheck request",
-			path:   "/foo/bar",
-			status: http.StatusOK,
-			body:   "Hello from inner handler",
-		},
-		{
-			name:   "Healthcheck request",
-			path:   "/-/healthcheck",
-			status: http.StatusOK,
-			body:   "success\n",
-		},
-	}
-
-	validCfg := config.GitLab{
-		InternalServer:     "server",
-		APISecretKey:       []byte("secret"),
-		ClientHTTPTimeout:  time.Second,
-		JWTTokenExpiration: time.Second,
-	}
-
-	source, err := gitlab.New(&validCfg)
-	require.NoError(t, err)
-
-	cfg := config.Config{
-		General: config.General{
-			StatusPath: "/-/healthcheck",
-		},
-	}
-
-	app := theApp{
-		config: &cfg,
-		source: source,
-	}
-
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-		io.WriteString(w, "Hello from inner handler")
-	})
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			r := httptest.NewRequest("GET", tc.path, nil)
-			rr := httptest.NewRecorder()
-
-			middleware := app.healthCheckMiddleware(handler)
-			middleware.ServeHTTP(rr, r)
-
-			require.Equal(t, tc.status, rr.Code)
-			require.Equal(t, tc.body, rr.Body.String())
-		})
-	}
-}
-
 func TestHandlePanicMiddleware(t *testing.T) {
 	next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		panic("on purpose")
diff --git a/internal/healthcheck/middleware.go b/internal/healthcheck/middleware.go
new file mode 100644
index 00000000..cd607f5a
--- /dev/null
+++ b/internal/healthcheck/middleware.go
@@ -0,0 +1,19 @@
+package healthcheck
+
+import (
+	"net/http"
+)
+
+// NewMiddleware is serving the application status check
+func NewMiddleware(handler http.Handler, statusPath string) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.RequestURI == statusPath {
+			w.Header().Set("Cache-Control", "no-store")
+			w.Write([]byte("success\n"))
+
+			return
+		}
+
+		handler.ServeHTTP(w, r)
+	})
+}
diff --git a/internal/healthcheck/middleware_test.go b/internal/healthcheck/middleware_test.go
new file mode 100644
index 00000000..dd189b92
--- /dev/null
+++ b/internal/healthcheck/middleware_test.go
@@ -0,0 +1,56 @@
+package healthcheck_test
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/config"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/healthcheck"
+)
+
+func TestHealthCheckMiddleware(t *testing.T) {
+	tests := []struct {
+		name string
+		path string
+		body string
+	}{
+		{
+			name: "Not a healthcheck request",
+			path: "/foo/bar",
+			body: "Hello from inner handler",
+		},
+		{
+			name: "Healthcheck request",
+			path: "/-/healthcheck",
+			body: "success\n",
+		},
+	}
+
+	cfg := config.Config{
+		General: config.General{
+			StatusPath: "/-/healthcheck",
+		},
+	}
+
+	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		io.WriteString(w, "Hello from inner handler")
+	})
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			r := httptest.NewRequest(http.MethodGet, tc.path, nil)
+			rr := httptest.NewRecorder()
+
+			middleware := healthcheck.NewMiddleware(handler, cfg.General.StatusPath)
+			middleware.ServeHTTP(rr, r)
+
+			require.Equal(t, http.StatusOK, rr.Code)
+			require.Equal(t, tc.body, rr.Body.String())
+		})
+	}
+}