Merge pull request #4 from chaosnative/body-modify-fix

Body Encoding Issue Fix
chaosnative · Sep 1, 2022 · cc9259c · cc9259c
2 parents 5ac41e0 + 56ec1d4
commit cc9259c
Show file tree

Hide file tree

Showing 7 changed files with 77 additions and 18 deletions.
diff --git a/toxics/httputils/body.go b/toxics/httputils/body.go
@@ -63,7 +63,36 @@ var StatusBodyTemplate = map[int]string{
 		<body><center><h1>504 Gateway Timeout</h1></body></html>`,
 }
 
-func EditResponseBody(r *http.Response, body string) {
-	r.Body = io.NopCloser(strings.NewReader(body))
+func EditResponseBody(r *http.Response, body, encoding, contentType string) {
+	compressedBody := encodeBody(r, []byte(body), encoding)
+
 	r.ContentLength = int64(len(body))
+	r.Header.Set("Content-Type", getContentType(contentType))
+	r.Body = io.NopCloser(strings.NewReader(string(compressedBody)))
+}
+
+func getContentType(contentType string) string {
+	if contentType == "" {
+		contentType = "text/plain"
+	}
+	return contentType
+}
+
+func encodeBody(r *http.Response, body []byte, encoding string) []byte {
+	r.Header.Set("Content-Encoding", "")
+	switch encoding {
+	case "gzip":
+		compressedBody, err := Gzip(body)
+		if err == nil {
+			r.Header.Set("Content-Encoding", "gzip")
+			return compressedBody
+		}
+	case "deflate":
+		compressedBody, err := Deflate(body)
+		if err == nil {
+			r.Header.Set("Content-Encoding", "deflate")
+			return compressedBody
+		}
+	}
+	return body
 }
diff --git a/toxics/httputils/body_encoding.go b/toxics/httputils/body_encoding.go
@@ -0,0 +1,30 @@
+package httputils
+
+import (
+	"bytes"
+	"compress/gzip"
+	"compress/zlib"
+)
+
+func Gzip(body []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	if _, err := gz.Write(body); err != nil {
+		gz.Close()
+		return nil, err
+	}
+	gz.Close()
+	return buf.Bytes(), nil
+}
+
+// Deflate compresses the body using the DEFLATE algorithm.
+func Deflate(body []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	z := zlib.NewWriter(&buf)
+	if _, err := z.Write(body); err != nil {
+		z.Close()
+		return nil, err
+	}
+	z.Close()
+	return buf.Bytes(), nil
+}
diff --git a/toxics/httputils/body_test.go b/toxics/httputils/body_test.go
@@ -32,7 +32,7 @@ func TestEditResponseBody(t *testing.T) {
 
 	AssertBodyNotEqual(t, body, []byte(checkBody))
 
-	httputils.EditResponseBody(resp, checkBody)
+	httputils.EditResponseBody(resp, checkBody, "", "")
 
 	body, _ = io.ReadAll(resp.Body)
 

diff --git a/toxics/httputils/status_code.go b/toxics/httputils/status_code.go
@@ -35,16 +35,8 @@ func SetHttpStatusCode(r *http.Response, statusCode int) {
 	// if the status code is not recognized, do not change it
 }
 
-func SetResponseBody(r *http.Response, statusCode int, body string) {
-	if body != "" {
-		EditResponseBody(r, body)
-	} else {
-		setErrorResponseBody(r, statusCode)
-	}
-}
-
-func setErrorResponseBody(r *http.Response, statusCode int) {
+func SetErrorResponseBody(r *http.Response, statusCode int) {
 	if _, exists := StatusBodyTemplate[statusCode]; statusCode >= 200 && statusCode < 600 && exists {
-		EditResponseBody(r, StatusBodyTemplate[statusCode])
+		EditResponseBody(r, StatusBodyTemplate[statusCode], "", "text/html")
 	}
 }
diff --git a/toxics/httputils/status_code_test.go b/toxics/httputils/status_code_test.go
@@ -44,7 +44,7 @@ func TestSetResponseBodyWithBody(t *testing.T) {
 	body, _ := io.ReadAll(resp.Body)
 	AssertBodyNotEqual(t, body, []byte(checkBody))
 
-	httputils.SetResponseBody(resp, 200, checkBody)
+	httputils.EditResponseBody(resp, checkBody, "", "text/plain")
 
 	body, _ = io.ReadAll(resp.Body)
 	AssertBodyEqual(t, body, []byte(checkBody))
@@ -57,7 +57,7 @@ func TestSetResponseBodyWithStatusCodeBody(t *testing.T) {
 	body, _ := io.ReadAll(resp.Body)
 	AssertBodyNotEqual(t, body, []byte(status500))
 
-	httputils.SetResponseBody(resp, 500, "")
+	httputils.SetErrorResponseBody(resp, 500)
 
 	body, _ = io.ReadAll(resp.Body)
 	AssertBodyEqual(t, body, []byte(status500))

diff --git a/toxics/modify_body.go b/toxics/modify_body.go
@@ -11,14 +11,16 @@ import (
 )
 
 type ModifyBodyToxic struct {
-	Body string `json:"body"`
+	Body             string `json:"body"`
+	ContentEnconding string `json:"content_encoding"`
+	ContentType      string `json:"content_type"`
 }
 
 func (t *ModifyBodyToxic) ModifyResponseBody(resp *http.Response) {
 	if t.Body == "" {
 		return
 	}
-	httputils.EditResponseBody(resp, t.Body)
+	httputils.EditResponseBody(resp, t.Body, t.ContentEnconding, t.ContentType)
 }
 
 func (t *ModifyBodyToxic) Pipe(stub *ToxicStub) {

diff --git a/toxics/status_code.go b/toxics/status_code.go
@@ -14,13 +14,19 @@ type StatusCodeToxic struct {
 	StatusCode         int    `json:"status_code"`
 	ModifyResponseBody int    `json:"modify_response_body"`
 	ResponseBody       string `json:"response_body"`
+	ContentEnconding   string `json:"content_encoding"`
+	ContentType        string `json:"content_type"`
 }
 
 func (t *StatusCodeToxic) ModifyResponseCode(resp *http.Response) {
 	httputils.SetHttpStatusCode(resp, t.StatusCode)
 
 	if t.ModifyResponseBody == 1 {
-		httputils.SetResponseBody(resp, t.StatusCode, t.ResponseBody)
+		if t.ResponseBody != "" {
+			httputils.EditResponseBody(resp, t.ResponseBody, t.ContentEnconding, t.ContentType)
+		} else {
+			httputils.SetErrorResponseBody(resp, t.StatusCode)
+		}
 	}
 }