diff --git a/response_writer.go b/response_writer.go index 7088fcf..9899395 100644 --- a/response_writer.go +++ b/response_writer.go @@ -12,6 +12,10 @@ import ( "net/http" ) +const ( + NoWritten = -1 +) + type ( ResponseWriter interface { http.ResponseWriter @@ -27,30 +31,29 @@ type ( responseWriter struct { http.ResponseWriter - status int - size int - written bool + status int + size int } ) func (w *responseWriter) reset(writer http.ResponseWriter) { w.ResponseWriter = writer w.status = 200 - w.written = false + w.size = NoWritten } func (w *responseWriter) WriteHeader(code int) { if code > 0 { w.status = code - if w.written { + if w.Written() { log.Println("[GIN] WARNING. Headers were already written!") } } } func (w *responseWriter) WriteHeaderNow() { - if !w.written { - w.written = true + if !w.Written() { + w.size = 0 w.ResponseWriter.WriteHeader(w.status) } } @@ -58,7 +61,7 @@ func (w *responseWriter) WriteHeaderNow() { func (w *responseWriter) Write(data []byte) (n int, err error) { w.WriteHeaderNow() n, err = w.ResponseWriter.Write(data) - w.size = n + w.size += n return } @@ -71,7 +74,7 @@ func (w *responseWriter) Size() int { } func (w *responseWriter) Written() bool { - return w.written + return w.size != NoWritten } // Implements the http.Hijacker interface