From fbb72869b3b3a96631dd8c98d755efa7bc63c362 Mon Sep 17 00:00:00 2001
From: "Noam Y. Tenne" <noam@10ne.org>
Date: Fri, 24 May 2019 07:13:57 +0200
Subject: [PATCH] Fix for #1334 (#1335)

* echo.context.cjson should encode the JSON before writing the status code #1334 :

`response.Write` automatically sets status to `200` if a response code wasn't committed yet. This is convenient, but it ignores the fact that `response.Status` is a public field that may be set separately/before `response.Write` has been called
A `response.Status` is by default `0`, or `200` if the response was reset, so `response.Write` should fallback to `200` only if a code wasn't set yet.

* echo.context.cjson should encode the JSON before writing the status code #1334 :

Writing the response code before encoding the payload is prone to error.
If JSON encoding fails, the response code is already committed, the server is able to only modify the response body to reflect the error and the user receives an awkward response where the status is successful but the body reports an error.
Instead - set the desired code on `c.response.Status`. If writing eventually takes place, the desired code is committed. If an error occurs, the server can still change the response.
---
 context.go       |  2 +-
 context_test.go  | 29 +++++++++++++++++++++++++++++
 response.go      |  5 ++++-
 response_test.go | 19 +++++++++++++++++++
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/context.go b/context.go
index 98cf50bc..065f5815 100644
--- a/context.go
+++ b/context.go
@@ -437,7 +437,7 @@ func (c *context) json(code int, i interface{}, indent string) error {
 		enc.SetIndent("", indent)
 	}
 	c.writeContentType(MIMEApplicationJSONCharsetUTF8)
-	c.response.WriteHeader(code)
+	c.response.Status = code
 	return enc.Encode(i)
 }
 
diff --git a/context_test.go b/context_test.go
index 6caa2704..fe542c81 100644
--- a/context_test.go
+++ b/context_test.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"math"
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
@@ -374,6 +375,34 @@ func TestContext(t *testing.T) {
 	assert.Equal(0, len(c.QueryParams()))
 }
 
+func TestContext_JSON_CommitsCustomResponseCode(t *testing.T) {
+	e := New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec).(*context)
+	err := c.JSON(http.StatusCreated, user{1, "Jon Snow"})
+
+	assert := testify.New(t)
+	if assert.NoError(err) {
+		assert.Equal(http.StatusCreated, rec.Code)
+		assert.Equal(MIMEApplicationJSONCharsetUTF8, rec.Header().Get(HeaderContentType))
+		assert.Equal(userJSON+"\n", rec.Body.String())
+	}
+}
+
+func TestContext_JSON_DoesntCommitResponseCodePrematurely(t *testing.T) {
+	e := New()
+	req := httptest.NewRequest(http.MethodGet, "/", nil)
+	rec := httptest.NewRecorder()
+	c := e.NewContext(req, rec).(*context)
+	err := c.JSON(http.StatusCreated, map[string]float64{"a": math.NaN()})
+
+	assert := testify.New(t)
+	if assert.Error(err) {
+		assert.False(c.response.Committed)
+	}
+}
+
 func TestContextCookie(t *testing.T) {
 	e := New()
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
diff --git a/response.go b/response.go
index eb2d988d..ca7405c5 100644
--- a/response.go
+++ b/response.go
@@ -67,7 +67,10 @@ func (r *Response) WriteHeader(code int) {
 // Write writes the data to the connection as part of an HTTP reply.
 func (r *Response) Write(b []byte) (n int, err error) {
 	if !r.Committed {
-		r.WriteHeader(http.StatusOK)
+		if r.Status == 0 {
+			r.Status = http.StatusOK
+		}
+		r.WriteHeader(r.Status)
 	}
 	n, err = r.Writer.Write(b)
 	r.Size += int64(n)
diff --git a/response_test.go b/response_test.go
index 8f08250f..bc570a50 100644
--- a/response_test.go
+++ b/response_test.go
@@ -22,3 +22,22 @@ func TestResponse(t *testing.T) {
 	res.Write([]byte("test"))
 	assert.Equal(t, "echo", rec.Header().Get(HeaderServer))
 }
+
+func TestResponse_Write_FallsBackToDefaultStatus(t *testing.T) {
+	e := New()
+	rec := httptest.NewRecorder()
+	res := &Response{echo: e, Writer: rec}
+
+	res.Write([]byte("test"))
+	assert.Equal(t, http.StatusOK, rec.Code)
+}
+
+func TestResponse_Write_UsesSetResponseCode(t *testing.T) {
+	e := New()
+	rec := httptest.NewRecorder()
+	res := &Response{echo: e, Writer: rec}
+
+	res.Status = http.StatusBadRequest
+	res.Write([]byte("test"))
+	assert.Equal(t, http.StatusBadRequest, rec.Code)
+}