From bdd57c01da7cfa642836155314cf99e918463662 Mon Sep 17 00:00:00 2001
From: Serhiy Mytrovtsiy <mitrovtsiy@ukr.net>
Date: Fri, 20 Nov 2020 23:25:37 +0100
Subject: [PATCH] - fix superfluous response.WriteHeader (setting a few times
 header when JsonResponse)

---
 response.go      | 16 ++++++----------
 response_test.go | 31 +++++++++++++++++++++++++++++++
 server.go        |  2 +-
 3 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/response.go b/response.go
index 135b120..23dc618 100644
--- a/response.go
+++ b/response.go
@@ -20,7 +20,7 @@ func (e HttpError) Error() string {
 }
 
 // RenderJSON sends data as json.
-func RenderJSON(w http.ResponseWriter, data interface{}) {
+func RenderJSON(w http.ResponseWriter, code int, data interface{}) {
 	buf := &bytes.Buffer{}
 	enc := json.NewEncoder(buf)
 	enc.SetEscapeHTML(true)
@@ -32,21 +32,19 @@ func RenderJSON(w http.ResponseWriter, data interface{}) {
 		}
 	}
 
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.WriteHeader(code)
 	_, _ = w.Write(buf.Bytes())
 }
 
 // JsonResponse - write a response with application/json Content-Type header.
 func JsonResponse(w http.ResponseWriter, data interface{}) {
-	w.Header().Set("Content-Type", "application/json; charset=utf-8")
-	w.WriteHeader(http.StatusOK)
-	RenderJSON(w, data)
+	RenderJSON(w, http.StatusOK, data)
 }
 
 // OKResponse - write a OK response with application/json Content-Type header.
 func OkResponse(w http.ResponseWriter) {
-	w.Header().Set("Content-Type", "application/json; charset=utf-8")
-	w.WriteHeader(http.StatusOK)
-	RenderJSON(w, struct {
+	RenderJSON(w, http.StatusOK, struct {
 		OK bool `json:"ok"`
 	}{
 		OK: true,
@@ -71,9 +69,7 @@ func ErrorResponse(w http.ResponseWriter, r *http.Request, code int, error error
 
 	log.Printf("[DEBUG] %s - %s - %d (%s) - %s - %s", r.Method, uri, code, http.StatusText(code), err, msg)
 
-	w.Header().Set("Content-Type", "application/json; charset=utf-8")
-	w.WriteHeader(code)
-	RenderJSON(w, err)
+	RenderJSON(w, code, err)
 }
 
 // NotFound - return a error page for not found
diff --git a/response_test.go b/response_test.go
index 093effc..534f681 100644
--- a/response_test.go
+++ b/response_test.go
@@ -10,6 +10,14 @@ import (
 	"testing"
 )
 
+type custom struct {
+	test string
+}
+
+func (c *custom) MarshalJSON() ([]byte, error) {
+	return nil, errors.New("test")
+}
+
 func TestJsonError(t *testing.T) {
 	handler := func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path == "/empty_error" {
@@ -71,6 +79,11 @@ func TestJsonResponse(t *testing.T) {
 		} else if r.URL.Path == "/string" {
 			JsonResponse(w, []byte("OK"))
 			return
+		} else if r.URL.Path == "/multiply-header-set" {
+			JsonResponse(w, &custom{
+				test: "test",
+			})
+			return
 		}
 		OkResponse(w)
 	}
@@ -80,8 +93,10 @@ func TestJsonResponse(t *testing.T) {
 		w := httptest.NewRecorder()
 		handler(w, req)
 		resp := w.Result()
+
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		require.Equal(t, "application/json; charset=utf-8", resp.Header.Get("Content-Type"))
+
 		body, err := ioutil.ReadAll(resp.Body)
 		require.NoError(t, err)
 		require.Empty(t, body)
@@ -92,8 +107,10 @@ func TestJsonResponse(t *testing.T) {
 		w := httptest.NewRecorder()
 		handler(w, req)
 		resp := w.Result()
+
 		require.Equal(t, http.StatusOK, resp.StatusCode)
 		require.Equal(t, "application/json; charset=utf-8", resp.Header.Get("Content-Type"))
+
 		body, err := ioutil.ReadAll(resp.Body)
 		require.NoError(t, err)
 
@@ -102,6 +119,20 @@ func TestJsonResponse(t *testing.T) {
 
 		require.Equal(t, append(b, '\n'), body)
 	})
+
+	t.Run("multiply header set", func(t *testing.T) {
+		req := httptest.NewRequest("GET", "http://test/multiply-header-set", nil)
+		w := httptest.NewRecorder()
+		handler(w, req)
+		resp := w.Result()
+
+		require.Equal(t, http.StatusInternalServerError, resp.StatusCode)
+		require.Equal(t, "text/plain; charset=utf-8", resp.Header.Get("Content-Type"))
+
+		body, err := ioutil.ReadAll(resp.Body)
+		require.NoError(t, err)
+		require.Equal(t, "json: error calling MarshalJSON for type *rest.custom: test\n", string(body))
+	})
 }
 
 func TestNotFound(t *testing.T) {
diff --git a/server.go b/server.go
index 66aeaec..86f3fe7 100644
--- a/server.go
+++ b/server.go
@@ -22,7 +22,7 @@ type Server struct {
 func handler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "text/plain")
 	w.WriteHeader(http.StatusOK)
-	w.Write([]byte("."))
+	_, _ = w.Write([]byte("."))
 }
 
 // Run - will initialize server and run it on provided port.