cleanup usage for less repeated code, and enable cwe fix (#3112)

Co-authored-by: Alex Meijer <ameijer@users.noreply.github.com>

core/pkg/protocol/http.go

@@ -1,9 +1,9 @@
 package protocol
 
 import (
-	"fmt"
 	"net/http"
 
+	"github.com/opencost/opencost/core/pkg/log"
 	"github.com/opencost/opencost/core/pkg/util/json"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
@@ -53,6 +53,25 @@ func (hp HTTPProtocol) InternalServerError(message string) HTTPError {
 	}
 }
 
+func (hp HTTPProtocol) NotImplemented(message string) HTTPError {
+	if message == "" {
+		message = "Not Implemented"
+	}
+	return HTTPError{
+		StatusCode: http.StatusNotImplemented,
+		Body:       message,
+	}
+}
+func (hp HTTPProtocol) Forbidden(message string) HTTPError {
+	if message == "" {
+		message = "Forbidden"
+	}
+	return HTTPError{
+		StatusCode: http.StatusForbidden,
+		Body:       message,
+	}
+}
+
 // NotFound creates a NotFound HTTPError
 func (hp HTTPProtocol) NotFound() HTTPError {
 	return HTTPError{
@@ -85,65 +104,84 @@ func (hp HTTPProtocol) ToResponse(data interface{}, err error) *HTTPResponse {
 		Data: data,
 	}
 }
+func (hp HTTPProtocol) WriteRawOK(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "application/json")
+	w.Header().Set("Content-Length", "0")
+	w.WriteHeader(http.StatusOK)
+}
+
+func (hp HTTPProtocol) WriteRawNoContent(w http.ResponseWriter) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusNoContent)
+}
+
+// WriteJSONData uses json content-type and json encoder with no data envelope allowing to remove
+// xss CWE as well as backwards compatibility to exisitng FE expectations
+func (hp HTTPProtocol) WriteJSONData(w http.ResponseWriter, data interface{}) {
+	w.Header().Set("Content-Type", "application/json")
+	status := http.StatusOK
+	w.WriteHeader(status)
+	if err := json.NewEncoder(w).Encode(data); err != nil {
+		log.Error("Failed to encode JSON response: " + err.Error())
+	}
+}
+
+// WriteRawError uses json content-type and outputs raw error message for backwards compatibility to existing
+// frontend expectations.
+func (hp HTTPProtocol) WriteRawError(w http.ResponseWriter, httpStatusCode int, err string) {
+	// I know this isn't json, but its what we've done and don't want to break frontned while we fix CWE
+	w.Header().Set("Content-Type", "application/json")
+	http.Error(w, err, httpStatusCode)
+}
+
+// WriteEncodedError writes an error response in the format of HTTPResponse
+func (hp HTTPProtocol) WriteEncodedError(w http.ResponseWriter, httpStatusCode int, errorResponse interface{}) {
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(httpStatusCode)
+	if err := json.NewEncoder(w).Encode(errorResponse); err != nil {
+		log.Error("Failed to encode error response: " + err.Error())
+	}
+}
 
 // WriteData wraps the data payload in an HTTPResponse and writes the resulting response using the
 // http.ResponseWriter
 func (hp HTTPProtocol) WriteData(w http.ResponseWriter, data interface{}) {
+	w.Header().Set("Content-Type", "application/json")
 	status := http.StatusOK
-	resp, err := json.Marshal(&HTTPResponse{
-		Code: status,
-		Data: data,
-	})
-	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Code:    status,
-			Message: fmt.Sprintf("Error: %s", err),
-		})
-	}
-
 	w.WriteHeader(status)
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(data); err != nil {
+		log.Error("Failed to encode response: " + err.Error())
+	}
 }
 
 // WriteDataWithWarning writes the data payload similiar to WriteData except it provides an additional warning message.
 func (hp HTTPProtocol) WriteDataWithWarning(w http.ResponseWriter, data interface{}, warning string) {
+	w.Header().Set("Content-Type", "application/json")
 	status := http.StatusOK
-	resp, err := json.Marshal(&HTTPResponse{
+	resp := &HTTPResponse{
 		Code:    status,
 		Data:    data,
 		Warning: warning,
-	})
-	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Code:    status,
-			Message: fmt.Sprintf("Error: %s", err),
-		})
 	}
-
 	w.WriteHeader(status)
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		log.Error("Failed to encode response with warning: " + err.Error())
+	}
 }
 
 // WriteDataWithMessage writes the data payload similiar to WriteData except it provides an additional string message.
 func (hp HTTPProtocol) WriteDataWithMessage(w http.ResponseWriter, data interface{}, message string) {
+	w.Header().Set("Content-Type", "application/json")
 	status := http.StatusOK
-	resp, err := json.Marshal(&HTTPResponse{
+	resp := &HTTPResponse{
 		Code:    status,
 		Data:    data,
 		Message: message,
-	})
-	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Code:    status,
-			Message: fmt.Sprintf("Error: %s", err),
-		})
 	}
-
 	w.WriteHeader(status)
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		log.Error("Failed to encode response with message: " + err.Error())
+	}
 }
 
 // WriteProtoWithMessage uses the protojson package to convert proto3 response to json response and
@@ -151,45 +189,41 @@ func (hp HTTPProtocol) WriteDataWithMessage(w http.ResponseWriter, data interfac
 // EmitUnpopulated to true it returns default values in the Json response payload. If error is
 // encountered it sent InternalServerError and the error why the json conversion failed.
 func (hp HTTPProtocol) WriteProtoWithMessage(w http.ResponseWriter, data proto.Message) {
+	w.Header().Set("Content-Type", "application/json")
 	m := protojson.MarshalOptions{
 		EmitUnpopulated: true,
 	}
 	status := http.StatusOK
-	resp, err := m.Marshal(data)
+	w.WriteHeader(status)
+	b, err := m.Marshal(data)
 	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Message: fmt.Sprintf("Error: %s", err),
-		})
+		hp.WriteError(w, hp.InternalServerError(err.Error()))
+		log.Error("Failed to marshal proto to json: " + err.Error())
+		return
 	}
 
-	w.WriteHeader(status)
-	w.Write(resp)
+	w.Write(b)
 }
 
 // WriteDataWithMessageAndWarning writes the data payload similiar to WriteData except it provides a warning and additional message string.
 func (hp HTTPProtocol) WriteDataWithMessageAndWarning(w http.ResponseWriter, data interface{}, message string, warning string) {
+	w.Header().Set("Content-Type", "application/json")
 	status := http.StatusOK
-	resp, err := json.Marshal(&HTTPResponse{
+	resp := &HTTPResponse{
 		Code:    status,
 		Data:    data,
 		Message: message,
 		Warning: warning,
-	})
-	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Code:    status,
-			Message: fmt.Sprintf("Error: %s", err),
-		})
 	}
-
 	w.WriteHeader(status)
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		log.Error("Failed to encode response with message and warning: " + err.Error())
+	}
 }
 
 // WriteError wraps the HTTPError in a HTTPResponse and writes it via http.ResponseWriter
 func (hp HTTPProtocol) WriteError(w http.ResponseWriter, err HTTPError) {
+	w.Header().Set("Content-Type", "application/json")
 	status := err.StatusCode
 	if status == 0 {
 		status = http.StatusInternalServerError
@@ -200,21 +234,17 @@ func (hp HTTPProtocol) WriteError(w http.ResponseWriter, err HTTPError) {
 		Code:    status,
 		Message: err.Body,
 	})
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(resp); err != nil {
+		log.Error("Failed to encode error response: " + err.Error())
+	}
 }
 
 // WriteResponse writes the provided HTTPResponse instance via http.ResponseWriter
 func (hp HTTPProtocol) WriteResponse(w http.ResponseWriter, r *HTTPResponse) {
+	w.Header().Set("Content-Type", "application/json")
 	status := r.Code
-	resp, err := json.Marshal(r)
-	if err != nil {
-		status = http.StatusInternalServerError
-		resp, _ = json.Marshal(&HTTPResponse{
-			Code:    status,
-			Message: fmt.Sprintf("Error: %s", err),
-		})
-	}
-
 	w.WriteHeader(status)
-	w.Write(resp)
+	if err := json.NewEncoder(w).Encode(r); err != nil {
+		log.Error("Failed to encode response: " + err.Error())
+	}
 }