Plugins: Refactor cleaning of call resource response headers (#67145)

First part of #66889 moving cleaning of call resource response headers 
within plugin management client.

pkg/plugins/manager/client/client.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"net/textproto"
 	"strings"
 
@@ -16,6 +17,12 @@ import (
 	"github.com/grafana/grafana/pkg/plugins/manager/registry"
 )
 
+const (
+	setCookieHeaderName   = "Set-Cookie"
+	contentTypeHeaderName = "Content-Type"
+	defaultContentType    = "application/json"
+)
+
 var _ plugins.Client = (*Service)(nil)
 
 type Service struct {
@@ -99,13 +106,22 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq
 	}, totalBytes, func() error {
 		removeConnectionHeaders(req.Headers)
 		removeHopByHopHeaders(req.Headers)
+		removeNonAllowedHeaders(req.Headers)
 
+		processedStreams := 0
 		wrappedSender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
-			if res != nil && len(res.Headers) > 0 {
-				removeConnectionHeaders(res.Headers)
-				removeHopByHopHeaders(res.Headers)
+			// Expected that headers and status are only part of first stream
+			if processedStreams == 0 && res != nil {
+				if len(res.Headers) > 0 {
+					removeConnectionHeaders(res.Headers)
+					removeHopByHopHeaders(res.Headers)
+					removeNonAllowedHeaders(res.Headers)
+				}
+
+				ensureContentTypeHeader(res)
 			}
 
+			processedStreams++
 			return sender.Send(res)
 		})
 
@@ -293,6 +309,33 @@ func removeHopByHopHeaders(h map[string][]string) {
 	}
 }
 
+func removeNonAllowedHeaders(h map[string][]string) {
+	for k := range h {
+		if textproto.CanonicalMIMEHeaderKey(k) == setCookieHeaderName {
+			delete(h, k)
+		}
+	}
+}
+
+// ensureContentTypeHeader makes sure a content type always is returned in response.
+func ensureContentTypeHeader(res *backend.CallResourceResponse) {
+	if res == nil {
+		return
+	}
+
+	var hasContentType bool
+	for k := range res.Headers {
+		if textproto.CanonicalMIMEHeaderKey(k) == contentTypeHeaderName {
+			hasContentType = true
+			break
+		}
+	}
+
+	if !hasContentType && res.Status != http.StatusNoContent {
+		res.Headers[contentTypeHeaderName] = []string{defaultContentType}
+	}
+}
+
 type callResourceResponseSenderFunc func(res *backend.CallResourceResponse) error
 
 func (fn callResourceResponseSenderFunc) Send(res *backend.CallResourceResponse) error {

pkg/plugins/manager/client/client_test.go

@@ -136,7 +136,6 @@ func TestCallResource(t *testing.T) {
 		res := responses[0]
 		require.Equal(t, http.StatusOK, res.Status)
 		require.Equal(t, []byte(backendResponse), res.Body)
-		require.Len(t, res.Headers, 1)
 		require.Equal(t, "should not be deleted", actualReq.Headers["X-Custom"][0])
 	})
 
@@ -200,9 +199,123 @@ func TestCallResource(t *testing.T) {
 		res := responses[0]
 		require.Equal(t, http.StatusOK, res.Status)
 		require.Equal(t, []byte(backendResponse), res.Body)
-		require.Len(t, res.Headers, 1)
 		require.Equal(t, "should not be deleted", actualReq.Headers["X-Custom"][0])
 	})
+
+	t.Run("Should remove non-allowed response headers", func(t *testing.T) {
+		resHeaders := map[string][]string{
+			setCookieHeaderName: {"monster"},
+			"X-Custom":          {"should not be deleted"},
+		}
+
+		req := &backend.CallResourceRequest{
+			PluginContext: backend.PluginContext{
+				PluginID: "pid",
+			},
+		}
+
+		responses := []*backend.CallResourceResponse{}
+		sender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
+			responses = append(responses, res)
+			return nil
+		})
+
+		p.RegisterClient(&fakePluginBackend{
+			crr: func(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+				return sender.Send(&backend.CallResourceResponse{
+					Headers: resHeaders,
+					Status:  http.StatusOK,
+					Body:    []byte(backendResponse),
+				})
+			},
+		})
+		err := registry.Add(context.Background(), p)
+		require.NoError(t, err)
+
+		client := ProvideService(registry, &config.Cfg{})
+
+		err = client.CallResource(context.Background(), req, sender)
+		require.NoError(t, err)
+
+		require.Len(t, responses, 1)
+		res := responses[0]
+		require.Equal(t, http.StatusOK, res.Status)
+		require.Equal(t, []byte(backendResponse), res.Body)
+		require.Empty(t, res.Headers[setCookieHeaderName])
+		require.Equal(t, "should not be deleted", res.Headers["X-Custom"][0])
+	})
+
+	t.Run("Should ensure content type header", func(t *testing.T) {
+		tcs := []struct {
+			contentType    string
+			responseStatus int
+			expContentType string
+		}{
+			{
+				contentType:    "",
+				responseStatus: http.StatusOK,
+				expContentType: defaultContentType,
+			},
+			{
+				contentType:    "text/plain",
+				responseStatus: http.StatusOK,
+				expContentType: "text/plain",
+			},
+			{
+				contentType:    "",
+				responseStatus: http.StatusNoContent,
+				expContentType: "",
+			},
+		}
+
+		for _, tc := range tcs {
+			t.Run(fmt.Sprintf("content type=%s, status=%d, exp=%s", tc.contentType, tc.responseStatus, tc.expContentType), func(t *testing.T) {
+				resHeaders := map[string][]string{}
+
+				if tc.contentType != "" {
+					resHeaders[contentTypeHeaderName] = []string{tc.contentType}
+				}
+
+				req := &backend.CallResourceRequest{
+					PluginContext: backend.PluginContext{
+						PluginID: "pid",
+					},
+				}
+
+				responses := []*backend.CallResourceResponse{}
+				sender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
+					responses = append(responses, res)
+					return nil
+				})
+
+				p.RegisterClient(&fakePluginBackend{
+					crr: func(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+						return sender.Send(&backend.CallResourceResponse{
+							Headers: resHeaders,
+							Status:  tc.responseStatus,
+							Body:    []byte(backendResponse),
+						})
+					},
+				})
+				err := registry.Add(context.Background(), p)
+				require.NoError(t, err)
+
+				client := ProvideService(registry, &config.Cfg{})
+
+				err = client.CallResource(context.Background(), req, sender)
+				require.NoError(t, err)
+
+				require.Len(t, responses, 1)
+				res := responses[0]
+
+				if tc.expContentType != "" {
+					require.Equal(t, tc.expContentType, res.Headers[contentTypeHeaderName][0])
+				} else {
+					require.Empty(t, res.Headers[contentTypeHeaderName])
+				}
+			})
+		}
+	})
 }
 
 type fakePluginBackend struct {