From 60b69fbb3b58183ca0bfe77075b5957da95384e3 Mon Sep 17 00:00:00 2001
From: Jordon Smith <jordon@jordonsmith.uk>
Date: Tue, 25 Oct 2016 15:56:41 +0100
Subject: [PATCH] Fix XML double WriteHeader (#677)

* WriteHeader if not committed in test.Response

Call WriteHeader if response isn't committed yet in calls to Write to
behave the same as fasthttp and standard Responses.

* Fix XML double WriteHeader

XMLBlob was writing xml.Header to the reply before writing the header
triggering a WriteHeader(http.StatusOK).
---
 context.go       | 5 ++++-
 context_test.go  | 4 ++--
 test/response.go | 3 +++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/context.go b/context.go
index b63987eb..92581917 100644
--- a/context.go
+++ b/context.go
@@ -390,10 +390,13 @@ func (c *context) XML(code int, i interface{}) (err error) {
 }
 
 func (c *context) XMLBlob(code int, b []byte) (err error) {
+	c.response.Header().Set(HeaderContentType, MIMEApplicationXMLCharsetUTF8)
+	c.response.WriteHeader(code)
 	if _, err = c.response.Write([]byte(xml.Header)); err != nil {
 		return
 	}
-	return c.Blob(code, MIMEApplicationXMLCharsetUTF8, b)
+	_, err = c.response.Write(b)
+	return
 }
 
 func (c *context) Blob(code int, contentType string, b []byte) (err error) {
diff --git a/context_test.go b/context_test.go
index d7b584c7..a375d779 100644
--- a/context_test.go
+++ b/context_test.go
@@ -99,9 +99,9 @@ func TestContext(t *testing.T) {
 	// XML
 	rec = test.NewResponseRecorder()
 	c = e.NewContext(req, rec).(*context)
-	err = c.XML(http.StatusOK, user{1, "Jon Snow"})
+	err = c.XML(http.StatusCreated, user{1, "Jon Snow"})
 	if assert.NoError(t, err) {
-		assert.Equal(t, http.StatusOK, rec.Status())
+		assert.Equal(t, http.StatusCreated, rec.Status())
 		assert.Equal(t, MIMEApplicationXMLCharsetUTF8, rec.Header().Get(HeaderContentType))
 		assert.Equal(t, xml.Header+userXML, rec.Body.String())
 	}
diff --git a/test/response.go b/test/response.go
index dc614693..02809524 100644
--- a/test/response.go
+++ b/test/response.go
@@ -55,6 +55,9 @@ func (r *Response) WriteHeader(code int) {
 }
 
 func (r *Response) Write(b []byte) (n int, err error) {
+	if !r.committed {
+		r.WriteHeader(http.StatusOK)
+	}
 	n, err = r.writer.Write(b)
 	r.size += int64(n)
 	return