1 files changed, 90 insertions, 30 deletions

diff --git a/echo/error_handler.go b/echo/error_handler.go
index bfbcfb6..bb4102b 100644
--- a/echo/error_handler.go
+++ b/echo/error_handler.go
@@ -7,19 +7,94 @@ import (
         "io/fs"
         "net/http"
 
+        "github.com/elnormous/contenttype"
         "github.com/labstack/echo/v4"
 )
 
-// Copied from echo and tweaked to make our errors nicer
+// TODO: This should allow plugging in other content types
+// TODO: This should also be refactored into something prettier
 func ErrorHandler(templates fs.FS, funcs template.FuncMap) func(error, echo.Context) {
+        handleHtml := func(c echo.Context, he *echo.HTTPError) error {
+                t, err := template.New("").Funcs(funcs).ParseFS(
+                        templates,
+                        "404.tpl",
+                        "40x.tpl",
+                        "50x.tpl",
+                        "header.tpl",
+                        "footer.tpl",
+                )
+                if err != nil {
+                        return err
+                }
+
+                path := "50x.tpl"
+                if he.Code == 404 {
+                        path = "404.tpl"
+                } else if he.Code >= 400 && he.Code <= 499 {
+                        path = "40x.tpl"
+                }
+
+                buf := bytes.Buffer{}
+                if err = t.ExecuteTemplate(&buf, path, nil); err != nil {
+                        err = c.String(he.Code, fmt.Sprintf("%s", he.Message))
+                }
+
+                return c.HTMLBlob(he.Code, buf.Bytes())
+        }
+
+        handlePlain := func(c echo.Context, he *echo.HTTPError) error {
+                return c.String(he.Code, fmt.Sprintf("%s", he.Message))
+        }
+
+        handleJson := func(c echo.Context, he *echo.HTTPError) error {
+                code := he.Code
+                message := he.Message
+                if m, ok := he.Message.(string); ok {
+                        if c.Echo().Debug {
+                                message = echo.Map{"message": m, "error": he.Error()}
+                        } else {
+                                message = echo.Map{"message": m}
+                        }
+                }
+
+                return c.JSON(code, message)
+        }
+
+        errorWhileErroring := func(c echo.Context, err interface{}) {
+                c.Echo().Logger.Error(err)
+                if c.Echo().Debug {
+                        c.JSON(http.StatusInternalServerError, &echo.HTTPError{
+                                Code:    http.StatusInternalServerError,
+                                Message: fmt.Sprintf("Error while processing error page. %w", err),
+                        })
+                } else {
+                        c.JSON(http.StatusInternalServerError, &echo.HTTPError{
+                                Code:    http.StatusInternalServerError,
+                                Message: http.StatusText(http.StatusInternalServerError),
+                        })
+                }
+        }
+
+        handlers := map[string]func(echo.Context, *echo.HTTPError) error{
+                "text/plain":       handlePlain,
+                "text/html":        handleHtml,
+                "text/json":        handleJson,
+                "application/json": handleJson,
+        }
+
+        // This is hand maintained here because order is important for negotiation,
+        // especially in the case of */*
+        hIndex := []contenttype.MediaType{
+                contenttype.NewMediaType("text/json"),
+                contenttype.NewMediaType("application/json"),
+                contenttype.NewMediaType("text/plain"),
+                contenttype.NewMediaType("text/html"),
+        }
+
         return func(err error, c echo.Context) {
                 defer func() {
                         if r := recover(); r != nil {
-                                if c.Echo().Debug {
-                                        c.String(http.StatusInternalServerError, fmt.Sprintf("Error while processing error page. %s", r))
-                                } else {
-                                        c.String(http.StatusInternalServerError, "Error while processing error page.")
-                                }
+                                errorWhileErroring(c, r)
                         }
                 }()
 
@@ -37,41 +112,26 @@ func ErrorHandler(templates fs.FS, funcs template.FuncMap) func(error, echo.Cont
                         }
                 }
 
-                t, err := template.New("").Funcs(funcs).ParseFS(
-                        templates,
-                        "404.tpl",
-                        "40x.tpl",
-                        "50x.tpl",
-                        "header.tpl",
-                        "footer.tpl",
-                )
+                ct, _, err := contenttype.GetAcceptableMediaType(c.Request(), hIndex)
                 if err != nil {
-                        he = &echo.HTTPError{
-                                Code:    http.StatusInternalServerError,
-                                Message: http.StatusText(http.StatusInternalServerError),
-                        }
+                        c.Echo().Logger.Error("Error negotiating content type in error handler, using json")
+                        ct = contenttype.NewMediaType("text/json")
                 }
 
-                path := "50x.tpl"
-                if he.Code == 404 {
-                        path = "404.tpl"
-                } else if he.Code >= 400 && he.Code <= 499 {
-                        path = "40x.tpl"
+                handle, ok := handlers[ct.String()]
+                if !ok {
+                        c.Echo().Logger.Errorf("Error handler content type %s is unknown", ct.String())
+                        handle = handleJson
                 }
 
-                // Send response
                 if !c.Response().Committed {
                         if c.Request().Method == http.MethodHead { // Issue #608
                                 err = c.NoContent(he.Code)
                         } else {
-                                buf := bytes.Buffer{}
-                                if err = t.ExecuteTemplate(&buf, path, nil); err != nil {
-                                        err = c.String(he.Code, fmt.Sprintf("%s", he.Message))
-                                }
-                                c.HTMLBlob(he.Code, buf.Bytes())
+                                err = handle(c, he)
                         }
                         if err != nil {
-                                c.Echo().Logger.Error(err)
+                                errorWhileErroring(c, err)
                         }
                 }
         }