rewrite code a bit

Signed-off-by: Naman Sood <mail@nsood.in>

errorcatcher/errorcatcher.go

@@ -0,0 +1,84 @@
+package errorcatcher
+
+import (
+	"net/http"
+	"prose/watcher"
+	"strconv"
+
+	"github.com/aymerick/raymond"
+)
+
+// errorCatcher is a wrapper for http.ResponseWriter that
+// captures 4xx and 5xx status codes and handles them in
+// a custom manner
+type errorCatcher struct {
+	r         *http.Request
+	w         http.ResponseWriter
+	templates watcher.AutoMap[string, *raymond.Template]
+	handled   bool
+}
+
+func New(w http.ResponseWriter, r *http.Request, templates watcher.AutoMap[string, *raymond.Template]) *errorCatcher {
+	return &errorCatcher{
+		r:         r,
+		w:         w,
+		templates: templates,
+	}
+}
+
+func (ec *errorCatcher) Header() http.Header {
+	return ec.w.Header()
+}
+
+func (ec *errorCatcher) Write(buf []byte) (int, error) {
+	// if we have already sent a response, pretend that this was successful
+	if ec.handled {
+		return len(buf), nil
+	}
+	return ec.w.Write(buf)
+}
+
+func (ec *errorCatcher) WriteHeader(statusCode int) {
+	if ec.handled {
+		return
+	}
+
+	if statusCode == http.StatusNotFound {
+		tpl, _ := ec.templates.Get("notfound.html")
+		page, err := tpl.Exec(map[string]string{
+			"path": ec.r.URL.Path,
+		})
+		// if we don't have a page to write, return before
+		// we toggle the flag so we fall back to the original
+		// error page
+		if err != nil {
+			return
+		}
+		ec.w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		ec.w.WriteHeader(statusCode)
+		ec.w.Write([]byte(page))
+		ec.handled = true
+		return
+	}
+
+	if statusCode >= 400 && statusCode < 600 {
+		ctx := map[string]string{
+			"code": strconv.Itoa(statusCode),
+		}
+		tpl, _ := ec.templates.Get("error.html")
+		page, err := tpl.Exec(ctx)
+		// if we don't have a page to write, return before
+		// we toggle the flag so we fall back to the original
+		// error page
+		if err != nil {
+			return
+		}
+		ec.w.Header().Set("Content-Type", "text/html; charset=utf-8")
+		ec.w.WriteHeader(statusCode)
+		ec.w.Write([]byte(page))
+		ec.handled = true
+		return
+	}
+
+	ec.w.WriteHeader(statusCode)
+}