Merge pull request #271 from sjbarag/let-before-funcs-modify-response…

…-status-code

fix: let ResponseWriter.Before() callbacks change status

response_writer.go

@@ -36,17 +36,28 @@ func NewResponseWriter(rw http.ResponseWriter) ResponseWriter {
 
 type responseWriter struct {
 	http.ResponseWriter
-	status      int
-	size        int
-	beforeFuncs []beforeFunc
+	pendingStatus  int
+	status         int
+	size           int
+	beforeFuncs    []beforeFunc
+	callingBefores bool
 }
 
 func (rw *responseWriter) WriteHeader(s int) {
 	if rw.Written() {
 		return
 	}
-	rw.status = s
+
+	rw.pendingStatus = s
 	rw.callBefore()
+
+	// Any of the rw.beforeFuncs may have written a header,
+	// so check again to see if any work is necessary.
+	if rw.Written() {
+		return
+	}
+
+	rw.status = s
 	rw.ResponseWriter.WriteHeader(s)
 }
 
@@ -74,7 +85,11 @@ func (rw *responseWriter) ReadFrom(r io.Reader) (n int64, err error) {
 }
 
 func (rw *responseWriter) Status() int {
-	return rw.status
+	if rw.Written() {
+		return rw.status
+	}
+
+	return rw.pendingStatus
 }
 
 func (rw *responseWriter) Size() int {
@@ -90,6 +105,15 @@ func (rw *responseWriter) Before(before func(ResponseWriter)) {
 }
 
 func (rw *responseWriter) callBefore() {
+	// Don't recursively call before() functions, to avoid infinite looping if
+	// one of them calls rw.WriteHeader again.
+	if rw.callingBefores {
+		return
+	}
+
+	rw.callingBefores = true
+	defer func() { rw.callingBefores = false }()
+
 	for i := len(rw.beforeFuncs) - 1; i >= 0; i-- {
 		rw.beforeFuncs[i](rw)
 	}

response_writer_test.go

@@ -69,6 +69,34 @@ func TestResponseWriterBeforeFuncHasAccessToStatus(t *testing.T) {
 	expect(t, status, http.StatusCreated)
 }
 
+func TestResponseWriterBeforeFuncCanChangeStatus(t *testing.T) {
+	rec := httptest.NewRecorder()
+	rw := NewResponseWriter(rec)
+
+	// Always respond with 200.
+	rw.Before(func(w ResponseWriter) {
+		w.WriteHeader(http.StatusOK)
+	})
+
+	rw.WriteHeader(http.StatusBadRequest)
+	expect(t, rec.Code, http.StatusOK)
+}
+
+func TestResponseWriterBeforeFuncChangesStatusMultipleTimes(t *testing.T) {
+	rec := httptest.NewRecorder()
+	rw := NewResponseWriter(rec)
+
+	rw.Before(func(w ResponseWriter) {
+		w.WriteHeader(http.StatusInternalServerError)
+	})
+	rw.Before(func(w ResponseWriter) {
+		w.WriteHeader(http.StatusNotFound)
+	})
+
+	rw.WriteHeader(http.StatusOK)
+	expect(t, rec.Code, http.StatusNotFound)
+}
+
 func TestResponseWriterWritingString(t *testing.T) {
 	rec := httptest.NewRecorder()
 	rw := NewResponseWriter(rec)