From 9a2696d216c3c873690c72dc301d4686eed265a4 Mon Sep 17 00:00:00 2001
From: JessonChan <chinachendongpo@gmail.com>
Date: Wed, 9 Mar 2016 15:56:18 +0800
Subject: [PATCH] accept asta's idea see the talk

https://github.com/astaxie/beego/pull/1719
---
 context/context.go | 13 ++++---------
 controller.go      |  2 +-
 2 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/context/context.go b/context/context.go
index 83d5a866..f7c16e31 100644
--- a/context/context.go
+++ b/context/context.go
@@ -175,16 +175,14 @@ func (ctx *Context) CheckXSRFCookie() bool {
 //started set to true if response was written to then don't execute other handler
 type Response struct {
 	http.ResponseWriter
-	Started     bool
-	Status      int
-	wroteHeader bool
+	Started bool
+	Status  int
 }
 
 func (r *Response) reset(rw http.ResponseWriter) {
 	r.ResponseWriter = rw
 	r.Status = 0
 	r.Started = false
-	r.wroteHeader = false
 }
 
 // Write writes the data to the connection as part of an HTTP reply,
@@ -192,11 +190,6 @@ func (r *Response) reset(rw http.ResponseWriter) {
 // started means the response has sent out.
 func (w *Response) Write(p []byte) (int, error) {
 	w.Started = true
-	if !w.wroteHeader {
-		w.ResponseWriter.WriteHeader(w.Status)
-		//prevent multiple response.WriteHeader calls
-		w.wroteHeader = true
-	}
 	return w.ResponseWriter.Write(p)
 }
 
@@ -204,10 +197,12 @@ func (w *Response) Write(p []byte) (int, error) {
 // and sets `started` to true.
 func (w *Response) WriteHeader(code int) {
 	if w.Status > 0 {
+		//prevent multiple response.WriteHeader calls
 		return
 	}
 	w.Status = code
 	w.Started = true
+	w.ResponseWriter.WriteHeader(w.Status)
 }
 
 // Hijack hijacker for http
diff --git a/controller.go b/controller.go
index a2943d42..2ad9d7cd 100644
--- a/controller.go
+++ b/controller.go
@@ -286,7 +286,7 @@ func (c *Controller) Abort(code string) {
 
 // CustomAbort stops controller handler and show the error data, it's similar Aborts, but support status code and body.
 func (c *Controller) CustomAbort(status int, body string) {
-	c.Ctx.ResponseWriter.WriteHeader(status)
+	c.Ctx.Output.Status = status
 	// first panic from ErrorMaps, is is user defined error functions.
 	if _, ok := ErrorMaps[body]; ok {
 		panic(body)