From dcafad3ced1d87333a7abdaecafc984813df5392 Mon Sep 17 00:00:00 2001
From: Manu Mtz-Almeida <manu.valladolid@gmail.com>
Date: Mon, 18 Aug 2014 05:24:48 +0200
Subject: [PATCH] Deferring WriteHeader. Part 1

---
 gin.go             | 10 ++++++++--
 logger.go          |  1 -
 render/render.go   |  6 ++----
 response_writer.go | 33 ++++++++++++++++++++++++---------
 4 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/gin.go b/gin.go
index 106e1b9c..99c5cbce 100644
--- a/gin.go
+++ b/gin.go
@@ -45,10 +45,15 @@ type (
 
 func (engine *Engine) handle404(w http.ResponseWriter, req *http.Request) {
 	c := engine.createContext(w, req, nil, engine.finalNoRoute)
-	c.Writer.setStatus(404)
+	// set 404 by default, useful for logging
+	c.Writer.WriteHeader(404)
 	c.Next()
 	if !c.Writer.Written() {
-		c.Data(404, MIMEPlain, []byte("404 page not found"))
+		if c.Writer.Status() == 404 {
+			c.Data(-1, MIMEPlain, []byte("404 page not found"))
+		} else {
+			c.Writer.WriteHeaderNow()
+		}
 	}
 	engine.cache.Put(c)
 }
@@ -166,6 +171,7 @@ func (group *RouterGroup) Handle(method, p string, handlers []HandlerFunc) {
 	group.engine.router.Handle(method, p, func(w http.ResponseWriter, req *http.Request, params httprouter.Params) {
 		c := group.engine.createContext(w, req, params, handlers)
 		c.Next()
+		c.Writer.WriteHeaderNow()
 		group.engine.cache.Put(c)
 	})
 }
diff --git a/logger.go b/logger.go
index 2667d131..8cbc24be 100644
--- a/logger.go
+++ b/logger.go
@@ -47,7 +47,6 @@ func Logger() HandlerFunc {
 		if len(requester) == 0 {
 			requester = c.Request.Header.Get("X-Forwarded-For")
 		}
-
 		// if the requester is still empty, use the hard-coded address from the socket
 		if len(requester) == 0 {
 			requester = c.Request.RemoteAddr
diff --git a/render/render.go b/render/render.go
index 2915ddc8..9080acea 100644
--- a/render/render.go
+++ b/render/render.go
@@ -35,10 +35,8 @@ var (
 )
 
 func writeHeader(w http.ResponseWriter, code int, contentType string) {
-	if code >= 0 {
-		w.Header().Set("Content-Type", contentType)
-		w.WriteHeader(code)
-	}
+	w.Header().Set("Content-Type", contentType)
+	w.WriteHeader(code)
 }
 
 func (_ jsonRender) Render(w http.ResponseWriter, code int, data ...interface{}) error {
diff --git a/response_writer.go b/response_writer.go
index cf02e90c..3d9de82d 100644
--- a/response_writer.go
+++ b/response_writer.go
@@ -1,6 +1,7 @@
 package gin
 
 import (
+	"log"
 	"net/http"
 )
 
@@ -9,9 +10,7 @@ type (
 		http.ResponseWriter
 		Status() int
 		Written() bool
-
-		// private
-		setStatus(int)
+		WriteHeaderNow()
 	}
 
 	responseWriter struct {
@@ -27,14 +26,30 @@ func (w *responseWriter) reset(writer http.ResponseWriter) {
 	w.written = false
 }
 
-func (w *responseWriter) setStatus(code int) {
-	w.status = code
+func (w *responseWriter) WriteHeader(code int) {
+	if code != 0 {
+		w.status = code
+		if w.written {
+			log.Println("[GIN] WARNING. Headers were already written!")
+		}
+	}
 }
 
-func (w *responseWriter) WriteHeader(code int) {
-	w.status = code
-	w.written = true
-	w.ResponseWriter.WriteHeader(code)
+func (w *responseWriter) WriteHeaderNow() {
+	if !w.written {
+		w.written = true
+		w.ResponseWriter.WriteHeader(w.status)
+	}
+}
+
+func (w *responseWriter) Write(data []byte) (n int, err error) {
+	if !w.written {
+		if w.status != 0 {
+			w.ResponseWriter.WriteHeader(w.status)
+		}
+		w.written = true
+	}
+	return w.ResponseWriter.Write(data)
 }
 
 func (w *responseWriter) Status() int {