From 586f15ddcc8174a5da66189ee4a329a180bb29e5 Mon Sep 17 00:00:00 2001
From: Ryan Deivert <ryandeivert@gmail.com>
Date: Sun, 25 Feb 2018 19:39:53 -0800
Subject: [PATCH] [api] using io.WriteString for writing responses instead of
 http.ResponseWriter.Write(...)

---
 handlers/distributed/distributed.go |  3 ++-
 handlers/response/response.go       | 35 ++++++++++++++++++++---------
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/handlers/distributed/distributed.go b/handlers/distributed/distributed.go
index 8a93038..4482333 100644
--- a/handlers/distributed/distributed.go
+++ b/handlers/distributed/distributed.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"time"
@@ -62,7 +63,7 @@ func DistributedQueryRead(respWriter http.ResponseWriter, request *http.Request)
 			return errors.New("no queries in list: %s")
 		}
 
-		respWriter.Write([]byte(distributedQuery.ToJSON()))
+		io.WriteString(respWriter, distributedQuery.ToJSON())
 		err = dyndb.DeleteDistributedQuery(distributedQuery, dynSvc)
 		if err != nil {
 			return fmt.Errorf("could not delete query: %s", err)
diff --git a/handlers/response/response.go b/handlers/response/response.go
index 26a6948..15c8efc 100644
--- a/handlers/response/response.go
+++ b/handlers/response/response.go
@@ -3,7 +3,10 @@ package response
 import (
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
+
+	"github.com/oktasecuritylabs/sgt/logger"
 )
 
 const (
@@ -22,32 +25,42 @@ type sgtBaseResponse struct {
 // storing custom endpoint response values
 type SGTCustomResponse map[string]interface{}
 
-func getResponseJSON(response interface{}) []byte {
+func writeResponseJSON(respWriter http.ResponseWriter, response interface{}) {
+
+	// Set the header type
+	respWriter.Header().Set("Content-Type", "application/json")
 
 	respJSON, err := json.Marshal(response)
 	if err != nil {
-		// If there is an error marshaling the interface, return a basic error
+		// If there is an error marshaling the interface, write a basic error
 		errString := fmt.Sprintf("response failed to marshal to json: %s", err)
-		newResp := sgtBaseResponse{Message: errString, Status: errString}
-		respJSON, _ = json.Marshal(newResp)
+		logger.Error(errString)
+		http.Error(respWriter, errString, http.StatusInternalServerError)
+		return
+	}
+
+	// Write the response to the http.ResponseWriter using io.WriteString
+	_, err = io.WriteString(respWriter, string(respJSON))
+
+	if err != nil {
+		// If there is an error writing the repsonse, write an error
+		errString := fmt.Sprintf("failed to write response: %s", err)
+		logger.Error(errString)
+		http.Error(respWriter, errString, http.StatusInternalServerError)
 	}
-	return respJSON
 }
 
 // WriteError will write the passed error to the http response writer
 func WriteError(respWriter http.ResponseWriter, errorString string) {
-	respWriter.Header().Set("Content-Type", "application/json")
-	respWriter.Write(getResponseJSON(sgtBaseResponse{Message: errorString, Status: statusError}))
+	writeResponseJSON(respWriter, sgtBaseResponse{Message: errorString, Status: statusError})
 }
 
 // WriteSuccess will write the a success status and optional message to the http response writer
 func WriteSuccess(respWriter http.ResponseWriter, optionalMessage string) {
-	respWriter.Header().Set("Content-Type", "application/json")
-	respWriter.Write(getResponseJSON(sgtBaseResponse{Message: optionalMessage, Status: statusSuccess}))
+	writeResponseJSON(respWriter, sgtBaseResponse{Message: optionalMessage, Status: statusSuccess})
 }
 
 // WriteCustomJSON will write the custom response to the http response writer as json
 func WriteCustomJSON(respWriter http.ResponseWriter, resp interface{}) {
-	respWriter.Header().Set("Content-Type", "application/json")
-	respWriter.Write(getResponseJSON(resp))
+	writeResponseJSON(respWriter, resp)
 }