From 4c4444b16033c7b000492f57119901cba9870b5c Mon Sep 17 00:00:00 2001
From: Roy Lou <roylou@gmail.com>
Date: Fri, 15 Apr 2016 00:02:29 +0800
Subject: [PATCH] Write header immediately in AbortWithStatus()

Otherwise, caller needs to invoke WriteHeaderNow himself after
AbortWithStatus(), which is error-prone.

Also modified ErrorLoggerT() such that it always writes log to response
body. Otherwise calling AbortWithStatus() will fail to write body because
c.Writer.Written() is set true by WriteHeaderNow().
---
 context.go       | 1 +
 context_test.go  | 2 --
 logger.go        | 9 +++------
 logger_test.go   | 2 +-
 recovery_test.go | 2 +-
 5 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/context.go b/context.go
index b043d1b4..3c1c5e2a 100644
--- a/context.go
+++ b/context.go
@@ -115,6 +115,7 @@ func (c *Context) Abort() {
 // For example, a failed attempt to authentificate a request could use: context.AbortWithStatus(401).
 func (c *Context) AbortWithStatus(code int) {
 	c.Status(code)
+	c.Writer.WriteHeaderNow()
 	c.Abort()
 }
 
diff --git a/context_test.go b/context_test.go
index 322c4829..56e11567 100644
--- a/context_test.go
+++ b/context_test.go
@@ -545,7 +545,6 @@ func TestContextAbortWithStatus(t *testing.T) {
 	c, w, _ := CreateTestContext()
 	c.index = 4
 	c.AbortWithStatus(401)
-	c.Writer.WriteHeaderNow()
 
 	assert.Equal(t, c.index, abortIndex)
 	assert.Equal(t, c.Writer.Status(), 401)
@@ -596,7 +595,6 @@ func TestContextTypedError(t *testing.T) {
 func TestContextAbortWithError(t *testing.T) {
 	c, w, _ := CreateTestContext()
 	c.AbortWithError(401, errors.New("bad input")).SetMeta("some input")
-	c.Writer.WriteHeaderNow()
 
 	assert.Equal(t, w.Code, 401)
 	assert.Equal(t, c.index, abortIndex)
diff --git a/logger.go b/logger.go
index c5d4c3e2..a1776a99 100644
--- a/logger.go
+++ b/logger.go
@@ -28,12 +28,9 @@ func ErrorLogger() HandlerFunc {
 func ErrorLoggerT(typ ErrorType) HandlerFunc {
 	return func(c *Context) {
 		c.Next()
-		// avoid writting if we already wrote into the response body
-		if !c.Writer.Written() {
-			errors := c.Errors.ByType(typ)
-			if len(errors) > 0 {
-				c.JSON(-1, errors)
-			}
+		errors := c.Errors.ByType(typ)
+		if len(errors) > 0 {
+			c.JSON(-1, errors)
 		}
 	}
 }
diff --git a/logger_test.go b/logger_test.go
index c1471fe9..2ad1f474 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -116,7 +116,7 @@ func TestErrorLogger(t *testing.T) {
 
 	w = performRequest(router, "GET", "/print")
 	assert.Equal(t, w.Code, 500)
-	assert.Equal(t, w.Body.String(), "hola!")
+	assert.Equal(t, w.Body.String(), "hola!{\"error\":\"this is an error\"}\n")
 }
 
 func TestSkippingPaths(t *testing.T) {
diff --git a/recovery_test.go b/recovery_test.go
index 94c71a11..4545ba3c 100644
--- a/recovery_test.go
+++ b/recovery_test.go
@@ -39,5 +39,5 @@ func TestPanicWithAbort(t *testing.T) {
 	// RUN
 	w := performRequest(router, "GET", "/recovery")
 	// TEST
-	assert.Equal(t, w.Code, 500) // NOT SURE
+	assert.Equal(t, w.Code, 400)
 }