prose/errorcatcher/errorcatcher.go

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