From 3bb4d6f0133dbe526e1917bebe2dc7a98ced6384 Mon Sep 17 00:00:00 2001
From: Christoph Portmann <c.portmann@salamworld.com>
Date: Tue, 8 Jul 2014 23:20:55 +0300
Subject: [PATCH] beego/context: Fix ignored Header in case SetStatus has been
 called before

---
 context/output.go | 9 ++++++++-
 router.go         | 5 +++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/context/output.go b/context/output.go
index b46c6a2d..8f59c94d 100644
--- a/context/output.go
+++ b/context/output.go
@@ -71,6 +71,14 @@ func (output *BeegoOutput) Body(content []byte) {
 	} else {
 		output.Header("Content-Length", strconv.Itoa(len(content)))
 	}
+
+	// Write status code if it has been set manually
+	// Set it to 0 afterwards to prevent "multiple response.WriteHeader calls"
+	if output.Status != 0 {
+		output.Context.ResponseWriter.WriteHeader(output.Status)
+		output.Status = 0
+	}
+
 	output_writer.Write(content)
 	switch output_writer.(type) {
 	case *gzip.Writer:
@@ -270,7 +278,6 @@ func (output *BeegoOutput) ContentType(ext string) {
 // SetStatus sets response status code.
 // It writes response header directly.
 func (output *BeegoOutput) SetStatus(status int) {
-	output.Context.ResponseWriter.WriteHeader(status)
 	output.Status = status
 }
 
diff --git a/router.go b/router.go
index e43e55fe..eba65c51 100644
--- a/router.go
+++ b/router.go
@@ -762,6 +762,11 @@ Admin:
 			Info("beego:" + r.URL.Path + " 404" + " +" + timeend.String())
 		}
 	}
+
+	// Call WriteHeader if status code has been set changed
+	if context.Output.Status != 0 {
+		w.writer.WriteHeader(context.Output.Status)
+	}
 }
 
 func (p *ControllerRegistor) recoverPanic(rw http.ResponseWriter, r *http.Request) {