Move healthcheck middleware to a separate package

add Cache-Control: no-store to status response

2 files changed, 75 insertions, 0 deletions

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