From 48203226c45145d8bbcb2dffbfc4d4d156faacac Mon Sep 17 00:00:00 2001
From: Christian Haudum <christian.haudum@gmail.com>
Date: Thu, 2 Dec 2021 08:29:20 +0100
Subject: [PATCH] Return 404 Not Found responses as JSON

Signed-off-by: Christian Haudum <christian.haudum@gmail.com>
---
 docs/sources/upgrading/_index.md | 2 --
 pkg/loki/loki.go                 | 1 +
 pkg/util/server/error.go         | 6 +++++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/docs/sources/upgrading/_index.md b/docs/sources/upgrading/_index.md
index deb166d57b12..563e8848aa83 100644
--- a/docs/sources/upgrading/_index.md
+++ b/docs/sources/upgrading/_index.md
@@ -49,8 +49,6 @@ The response body has the following schema:
 }
 ```
 
-"404 Not Found" errors still return responses with content type `text/plain`.
-
 ### Promtail
 
 #### `gcplog` labels have changed
diff --git a/pkg/loki/loki.go b/pkg/loki/loki.go
index f4f57ea2ae51..2b5f8ceb16aa 100644
--- a/pkg/loki/loki.go
+++ b/pkg/loki/loki.go
@@ -304,6 +304,7 @@ func (t *Loki) Run(opts RunOpts) error {
 
 	t.serviceMap = serviceMap
 	t.Server.HTTP.Path("/services").Methods("GET").Handler(http.HandlerFunc(t.servicesHandler))
+	t.Server.HTTP.NotFoundHandler = http.HandlerFunc(serverutil.NotFoundHandler)
 
 	// get all services, create service manager and tell it to start
 	var servs []services.Service
diff --git a/pkg/util/server/error.go b/pkg/util/server/error.go
index 4b55f6d87ede..6b423d1440a4 100644
--- a/pkg/util/server/error.go
+++ b/pkg/util/server/error.go
@@ -34,9 +34,13 @@ type ErrorResponseBody struct {
 	Message string `json:"message"`
 }
 
+func NotFoundHandler(w http.ResponseWriter, r *http.Request) {
+	JSONError(w, 404, "page not found")
+}
+
 func JSONError(w http.ResponseWriter, code int, message string, args ...interface{}) {
-	w.WriteHeader(code)
 	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.WriteHeader(code)
 	json.NewEncoder(w).Encode(ErrorResponseBody{
 		Code:    code,
 		Status:  "error",