fix(field-suggestions): Fix block-field-suggestions rule throwing Bod… (#19)

Fix bug on Block Field suggestions rule.

We found that the block field suggestion rule broke when the response
body did not contain the valid JSON (i.e. when proxying GraphiQL).

This MR makes sure to read the response body as plain bytes, try to
determine whether it can and should be modified and if that fails, set
the body back to the original bytes so that it can be read again later.

Co-authored-by: ldebruijn <[email protected]>

Author: ldebruijn

cmd/test_test.go

@@ -32,25 +32,36 @@ func NewProxy(cfg Config, blockFieldSuggestions *block_field_suggestions.BlockFi
 			KeepAlive: cfg.KeepAlive,
 		}).DialContext,
 	}
-	proxy.ModifyResponse = func(res *http.Response) error {
+	proxy.ModifyResponse = modifyResponse(blockFieldSuggestions)
+
+	return proxy, nil
+}
+
+func modifyResponse(blockFieldSuggestions *block_field_suggestions.BlockFieldSuggestionsHandler) func(res *http.Response) error {
+	return func(res *http.Response) error {
 		if !blockFieldSuggestions.Enabled() {
 			return nil
 		}
 
-		decoder := json.NewDecoder(res.Body)
+		// read raw response bytes
+		bodyBytes, _ := io.ReadAll(res.Body)
 		defer res.Body.Close()
 
 		var response map[string]interface{}
-		err := decoder.Decode(&response)
+		err := json.Unmarshal(bodyBytes, &response)
 		if err != nil {
 			// if we cannot decode just return
+			// make sure to set body back to original bytes
+			res.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
 			return nil
 		}
 
 		modified := blockFieldSuggestions.ProcessBody(response)
 		bts, err := json.Marshal(modified)
 		if err != nil {
 			// if we cannot marshall just return
+			// make sure to set body back to original bytes
+			res.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
 			return nil
 		}
 
@@ -61,6 +72,4 @@ func NewProxy(cfg Config, blockFieldSuggestions *block_field_suggestions.BlockFi
 
 		return nil
 	}
-
-	return proxy, nil
 }

internal/business/proxy/proxy.go

@@ -0,0 +1,144 @@
+package proxy
+
+import (
+	"github.com/ldebruijn/go-graphql-armor/internal/business/block_field_suggestions"
+	"github.com/stretchr/testify/assert"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func Test_modifyResponse(t *testing.T) {
+	type args struct {
+		blockFieldSuggestions *block_field_suggestions.BlockFieldSuggestionsHandler
+		response              *http.Response
+	}
+	tests := []struct {
+		name string
+		args args
+		want func(res *http.Response)
+	}{
+		{
+			name: "nothing if disabled",
+			args: args{
+				blockFieldSuggestions: func() *block_field_suggestions.BlockFieldSuggestionsHandler {
+					return block_field_suggestions.NewBlockFieldSuggestionsHandler(block_field_suggestions.Config{
+						Enabled: false,
+					})
+				}(),
+				response: func() *http.Response {
+					return &http.Response{
+						Status:        "200",
+						StatusCode:    200,
+						Body:          io.NopCloser(strings.NewReader("this is not valid json")),
+						Proto:         "HTTP/1.1",
+						ProtoMajor:    1,
+						ProtoMinor:    1,
+						ContentLength: 0,
+						Header:        map[string][]string{},
+					}
+				}(),
+			},
+			want: func(res *http.Response) {
+				body, _ := io.ReadAll(res.Body)
+				assert.Equal(t, 200, res.StatusCode)
+				assert.Equal(t, "200", res.Status)
+				assert.Equal(t, "this is not valid json", string(body))
+			},
+		},
+		{
+			name: "handles non-json gracefully",
+			args: args{
+				blockFieldSuggestions: func() *block_field_suggestions.BlockFieldSuggestionsHandler {
+					return block_field_suggestions.NewBlockFieldSuggestionsHandler(block_field_suggestions.Config{
+						Enabled: true,
+					})
+				}(),
+				response: func() *http.Response {
+					return &http.Response{
+						Status:        "200",
+						StatusCode:    200,
+						Body:          io.NopCloser(strings.NewReader("this is not valid json")),
+						Proto:         "HTTP/1.1",
+						ProtoMajor:    1,
+						ProtoMinor:    1,
+						ContentLength: 0,
+						Header:        map[string][]string{},
+					}
+				}(),
+			},
+			want: func(res *http.Response) {
+				body, _ := io.ReadAll(res.Body)
+				assert.Equal(t, 200, res.StatusCode)
+				assert.Equal(t, "200", res.Status)
+				assert.Equal(t, "this is not valid json", string(body))
+			},
+		},
+		{
+			name: "handles invalid-json gracefully",
+			args: args{
+				blockFieldSuggestions: func() *block_field_suggestions.BlockFieldSuggestionsHandler {
+					return block_field_suggestions.NewBlockFieldSuggestionsHandler(block_field_suggestions.Config{
+						Enabled: true,
+					})
+				}(),
+				response: func() *http.Response {
+					return &http.Response{
+						Status:        "200",
+						StatusCode:    200,
+						Body:          io.NopCloser(strings.NewReader("{ \"this\": \" is not valid json }")),
+						Proto:         "HTTP/1.1",
+						ProtoMajor:    1,
+						ProtoMinor:    1,
+						ContentLength: 0,
+						Header:        map[string][]string{},
+					}
+				}(),
+			},
+			want: func(res *http.Response) {
+				body, _ := io.ReadAll(res.Body)
+				assert.Equal(t, 200, res.StatusCode)
+				assert.Equal(t, "200", res.Status)
+				assert.Equal(t, "{ \"this\": \" is not valid json }", string(body))
+			},
+		},
+		{
+			name: "handles json gracefully",
+			args: args{
+				blockFieldSuggestions: func() *block_field_suggestions.BlockFieldSuggestionsHandler {
+					return block_field_suggestions.NewBlockFieldSuggestionsHandler(block_field_suggestions.Config{
+						Enabled: true,
+						Mask:    "[masked]",
+					})
+				}(),
+				response: func() *http.Response {
+					return &http.Response{
+						Status:        "200",
+						StatusCode:    200,
+						Body:          io.NopCloser(strings.NewReader("{ \"errors\": [{\"message\": \"Did you mean \"}] }")),
+						Proto:         "HTTP/1.1",
+						ProtoMajor:    1,
+						ProtoMinor:    1,
+						ContentLength: 0,
+						Header:        map[string][]string{},
+					}
+				}(),
+			},
+			want: func(res *http.Response) {
+				body, _ := io.ReadAll(res.Body)
+				assert.Equal(t, 200, res.StatusCode)
+				assert.Equal(t, "200", res.Status)
+				assert.Equal(t, "{\"errors\":[{\"message\":\"[masked]\"}]}", string(body))
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := modifyResponse(tt.args.blockFieldSuggestions)
+
+			_ = result(tt.args.response)
+			tt.want(tt.args.response)
+		})
+	}
+}