middleware: Add an ErrorHandler middleware used to serve an alternate handler on a certain error code

common/middleware/errorhandler.go

@@ -0,0 +1,76 @@
+package middleware
+
+import (
+	"net/http"
+)
+
+func copyHeaders(src, dest http.Header) {
+	for k, v := range src {
+		dest[k] = v
+	}
+}
+
+// ErrorHandler lets you call an alternate http handler upon a certain response code.
+// Note it will assume a 200 if the wrapped handler does not write anything
+type ErrorHandler struct {
+	Code    int
+	Handler http.Handler
+}
+
+func (e ErrorHandler) Wrap(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		i := newErrorInterceptor(w, e.Code)
+		next.ServeHTTP(i, r)
+		if !i.gotCode {
+			i.WriteHeader(http.StatusOK)
+		}
+		if i.intercepted {
+			e.Handler.ServeHTTP(w, r)
+		}
+	})
+}
+
+// errorInterceptor wraps an underlying ResponseWriter and buffers all header changes, until it knows the return code.
+// It then passes everything through, unless the code matches the target code, in which case it will discard everything.
+type errorInterceptor struct {
+	originalWriter http.ResponseWriter
+	targetCode     int
+	headers        http.Header
+	gotCode        bool
+	intercepted    bool
+}
+
+func newErrorInterceptor(w http.ResponseWriter, code int) *errorInterceptor {
+	i := errorInterceptor{originalWriter: w, targetCode: code}
+	i.headers = make(http.Header)
+	copyHeaders(w.Header(), i.headers)
+	return &i
+}
+
+func (i errorInterceptor) Header() http.Header {
+	return i.headers
+}
+
+func (i errorInterceptor) WriteHeader(code int) {
+	if i.gotCode {
+		panic("errorInterceptor.WriteHeader() called twice")
+	}
+
+	i.gotCode = true
+	if code == i.targetCode {
+		i.intercepted = true
+	} else {
+		copyHeaders(i.headers, i.originalWriter.Header())
+		i.originalWriter.WriteHeader(code)
+	}
+}
+
+func (i errorInterceptor) Write(data []byte) (int, error) {
+	if !i.gotCode {
+		i.WriteHeader(http.StatusOK)
+	}
+	if !i.intercepted {
+		return i.originalWriter.Write(data)
+	}
+	return len(data), nil
+}