From d42aa6d868c84c83737779f2c0446513efdca317 Mon Sep 17 00:00:00 2001
From: Manu Mtz-Almeida <manu.valladolid@gmail.com>
Date: Fri, 4 Jul 2014 00:01:28 +0200
Subject: [PATCH] Fixes wrap around http.ResponseWriter

---
 README.md          |  4 ++++
 gin.go             |  8 ++++----
 response_writer.go | 47 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 response_writer.go

diff --git a/README.md b/README.md
index d80f07a9..97e80af2 100644
--- a/README.md
+++ b/README.md
@@ -286,6 +286,10 @@ func Logger() gin.HandlerFunc {
         // after request
         latency := time.Since(t)
         log.Print(latency)
+
+        // access the status we are sending
+        status := c.Writer.Status()
+        log.Println(status)
     }
 }
 
diff --git a/gin.go b/gin.go
index d1d42625..314894f2 100644
--- a/gin.go
+++ b/gin.go
@@ -40,7 +40,7 @@ type (
 	// manage the flow, validate the JSON of a request and render a JSON response for example.
 	Context struct {
 		Req      *http.Request
-		Writer   http.ResponseWriter
+		Writer   ResponseWriter
 		Keys     map[string]interface{}
 		Errors   ErrorMsgs
 		Params   httprouter.Params
@@ -92,7 +92,7 @@ func NewWithConfig(config Config) *Engine {
 
 	// Fill it with empty contexts
 	for i := 0; i < config.Preallocated; i++ {
-		engine.cache <- &Context{Engine: engine}
+		engine.cache <- &Context{Engine: engine, Writer: &responseWriter{}}
 	}
 	return engine
 }
@@ -171,7 +171,7 @@ func (engine *Engine) Run(addr string) {
 func (engine *Engine) createContext(w http.ResponseWriter, req *http.Request, params httprouter.Params, handlers []HandlerFunc) *Context {
 	select {
 	case c := <-engine.cache:
-		c.Writer = w
+		c.Writer.reset(w)
 		c.Req = req
 		c.Params = params
 		c.handlers = handlers
@@ -180,7 +180,7 @@ func (engine *Engine) createContext(w http.ResponseWriter, req *http.Request, pa
 		return c
 	default:
 		return &Context{
-			Writer:   w,
+			Writer:   &responseWriter{w, -1, false},
 			Req:      req,
 			Params:   params,
 			handlers: handlers,
diff --git a/response_writer.go b/response_writer.go
new file mode 100644
index 00000000..88c1b20d
--- /dev/null
+++ b/response_writer.go
@@ -0,0 +1,47 @@
+package gin
+
+import (
+	"net/http"
+)
+
+type (
+	ResponseWriter interface {
+		http.ResponseWriter
+		Status() int
+		Written() bool
+
+		// private
+		reset(http.ResponseWriter)
+		setStatus(int)
+	}
+
+	responseWriter struct {
+		http.ResponseWriter
+		status  int
+		written bool
+	}
+)
+
+func (w *responseWriter) reset(writer http.ResponseWriter) {
+	w.ResponseWriter = writer
+	w.status = 0
+	w.written = false
+}
+
+func (w *responseWriter) setStatus(code int) {
+	w.status = code
+}
+
+func (w *responseWriter) WriteHeader(code int) {
+	w.status = code
+	w.written = true
+	w.ResponseWriter.WriteHeader(code)
+}
+
+func (w *responseWriter) Status() int {
+	return w.status
+}
+
+func (w *responseWriter) Written() bool {
+	return w.written
+}