diff --git a/response_writer.go b/response_writer.go index cf02e90c..897d3307 100644 --- a/response_writer.go +++ b/response_writer.go @@ -9,6 +9,8 @@ type ( http.ResponseWriter Status() int Written() bool + // Before allows for a function to be called before the ResponseWriter has been written to. + Before(BeforeFunc) // private setStatus(int) @@ -16,15 +18,19 @@ type ( responseWriter struct { http.ResponseWriter - status int - written bool + status int + written bool + beforeFuncs []BeforeFunc } + + BeforeFunc func(ResponseWriter) ) func (w *responseWriter) reset(writer http.ResponseWriter) { w.ResponseWriter = writer w.status = 0 w.written = false + w.beforeFuncs = w.beforeFuncs[:0] } func (w *responseWriter) setStatus(code int) { @@ -32,6 +38,7 @@ func (w *responseWriter) setStatus(code int) { } func (w *responseWriter) WriteHeader(code int) { + w.callBefore() w.status = code w.written = true w.ResponseWriter.WriteHeader(code) @@ -44,3 +51,13 @@ func (w *responseWriter) Status() int { func (w *responseWriter) Written() bool { return w.written } + +func (w *responseWriter) Before(before BeforeFunc) { + w.beforeFuncs = append(w.beforeFuncs, before) +} + +func (w *responseWriter) callBefore() { + for i := len(w.beforeFuncs) - 1; i >= 0; i-- { + w.beforeFuncs[i](w) + } +}