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/api/plugin_resource.go

@@ -161,7 +161,6 @@ func (hs *HTTPServer) flushStream(stream callResourceClientResponseStream, w htt
 
 		// Expected that headers and status are only part of first stream
 		if processedStreams == 0 {
-			var hasContentType bool
 			for k, values := range resp.Headers {
 				// Convert the keys to the canonical format of MIME headers.
 				// This ensures that we can safely add/overwrite headers
@@ -169,15 +168,6 @@ func (hs *HTTPServer) flushStream(stream callResourceClientResponseStream, w htt
 				// and be sure they won't be present multiple times in the response.
 				k = textproto.CanonicalMIMEHeaderKey(k)
 
-				switch k {
-				case "Set-Cookie":
-					// Due to security reasons we don't want to forward
-					// cookies from a backend plugin to clients/browsers.
-					continue
-				case "Content-Type":
-					hasContentType = true
-				}
-
 				for _, v := range values {
 					// TODO: Figure out if we should use Set here instead
 					// nolint:gocritic
@@ -185,13 +175,6 @@ func (hs *HTTPServer) flushStream(stream callResourceClientResponseStream, w htt
 				}
 			}
 
-			// Make sure a content type always is returned in response
-			if !hasContentType && resp.Status != http.StatusNoContent {
-				w.Header().Set("Content-Type", "application/json")
-			}
-
-			proxyutil.SetProxyResponseHeaders(w.Header())
-
 			w.WriteHeader(resp.Status)
 		}
 

pkg/api/plugins_test.go

@@ -388,30 +388,9 @@ func TestMakePluginResourceRequest(t *testing.T) {
 		}
 	}
 
-	require.Equal(t, resp.Header().Get("Content-Type"), "application/json")
-	require.Equal(t, "sandbox", resp.Header().Get("Content-Security-Policy"))
-}
-
-func TestMakePluginResourceRequestSetCookieNotPresent(t *testing.T) {
-	hs := HTTPServer{
-		Cfg: setting.NewCfg(),
-		log: log.New(),
-		pluginClient: &fakePluginClient{
-			headers: map[string][]string{"Set-Cookie": {"monster"}},
-		},
-	}
-	req := httptest.NewRequest(http.MethodGet, "/", nil)
-	resp := httptest.NewRecorder()
-	pCtx := backend.PluginContext{}
-	err := hs.makePluginResourceRequest(resp, req, pCtx)
-	require.NoError(t, err)
-
-	for {
-		if resp.Flushed {
-			break
-		}
-	}
-	require.Empty(t, resp.Header().Values("Set-Cookie"), "Set-Cookie header should not be present")
+	res := resp.Result()
+	require.NoError(t, res.Body.Close())
+	require.Equal(t, http.StatusOK, res.StatusCode)
 }
 
 func TestMakePluginResourceRequestContentTypeUnique(t *testing.T) {

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 {
+			// 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 {

pkg/services/pluginsintegration/clientmiddleware/caching_middleware.go

@@ -123,7 +123,7 @@ func (m *CachingMiddleware) CallResource(ctx context.Context, req *backend.CallR
 		return m.next.CallResource(ctx, req, sender)
 	}
 	// Otherwise, intercept the responses in a wrapped sender so we can cache them first
-	cacheSender := cachedSenderFunc(func(res *backend.CallResourceResponse) error {
+	cacheSender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
 		cr.UpdateCacheFn(ctx, res)
 		return sender.Send(res)
 	})
@@ -150,9 +150,3 @@ func (m *CachingMiddleware) PublishStream(ctx context.Context, req *backend.Publ
 func (m *CachingMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
 	return m.next.RunStream(ctx, req, sender)
 }
-
-type cachedSenderFunc func(res *backend.CallResourceResponse) error
-
-func (fn cachedSenderFunc) Send(res *backend.CallResourceResponse) error {
-	return fn(res)
-}

pkg/services/pluginsintegration/clientmiddleware/resource_response_middleware.go

@@ -0,0 +1,65 @@
+package clientmiddleware
+
+import (
+	"context"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/util/proxyutil"
+)
+
+// NewResourceResponseMiddleware creates a new plugins.ClientMiddleware
+// that will enforce HTTP header rules for backend.CallResourceResponse's.
+func NewResourceResponseMiddleware() plugins.ClientMiddleware {
+	return plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
+		return &ResourceResponseMiddleware{
+			next: next,
+		}
+	})
+}
+
+type ResourceResponseMiddleware struct {
+	next plugins.Client
+}
+
+func (m *ResourceResponseMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	return m.next.QueryData(ctx, req)
+}
+
+func (m *ResourceResponseMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	if req == nil || sender == nil {
+		return m.next.CallResource(ctx, req, sender)
+	}
+
+	processedStreams := 0
+	wrappedSender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
+		if processedStreams == 0 {
+			proxyutil.SetProxyResponseHeaders(res.Headers)
+		}
+
+		processedStreams++
+		return sender.Send(res)
+	})
+
+	return m.next.CallResource(ctx, req, wrappedSender)
+}
+
+func (m *ResourceResponseMiddleware) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	return m.next.CheckHealth(ctx, req)
+}
+
+func (m *ResourceResponseMiddleware) CollectMetrics(ctx context.Context, req *backend.CollectMetricsRequest) (*backend.CollectMetricsResult, error) {
+	return m.next.CollectMetrics(ctx, req)
+}
+
+func (m *ResourceResponseMiddleware) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
+	return m.next.SubscribeStream(ctx, req)
+}
+
+func (m *ResourceResponseMiddleware) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
+	return m.next.PublishStream(ctx, req)
+}
+
+func (m *ResourceResponseMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
+	return m.next.RunStream(ctx, req, sender)
+}

pkg/services/pluginsintegration/clientmiddleware/resource_response_middleware_test.go

@@ -0,0 +1,41 @@
+package clientmiddleware
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/plugins/manager/client/clienttest"
+	"github.com/stretchr/testify/require"
+)
+
+func TestResourceResponseMiddleware(t *testing.T) {
+	t.Run("Should set proxy response headers when calling CallResource", func(t *testing.T) {
+		crResp := &backend.CallResourceResponse{
+			Status: http.StatusOK,
+			Headers: map[string][]string{
+				"X-Custom": {"Should not be deleted"},
+			},
+		}
+		cdt := clienttest.NewClientDecoratorTest(t,
+			clienttest.WithMiddlewares(NewResourceResponseMiddleware()),
+			clienttest.WithResourceResponses([]*backend.CallResourceResponse{crResp}),
+		)
+
+		var sentResponse *backend.CallResourceResponse
+		sender := callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
+			sentResponse = res
+			return nil
+		})
+
+		err := cdt.Decorator.CallResource(context.Background(), &backend.CallResourceRequest{
+			PluginContext: backend.PluginContext{},
+		}, sender)
+		require.NoError(t, err)
+
+		require.NotNil(t, sentResponse)
+		require.Equal(t, "sandbox", sentResponse.Headers["Content-Security-Policy"][0])
+		require.Equal(t, "Should not be deleted", sentResponse.Headers["X-Custom"][0])
+	})
+}

pkg/services/pluginsintegration/clientmiddleware/testing.go

@@ -2,12 +2,6 @@ package clientmiddleware
 
 import "github.com/grafana/grafana-plugin-sdk-go/backend"
 
-type callResourceResponseSenderFunc func(res *backend.CallResourceResponse) error
-
-func (fn callResourceResponseSenderFunc) Send(res *backend.CallResourceResponse) error {
-	return fn(res)
-}
-
 var nopCallResourceSender = callResourceResponseSenderFunc(func(res *backend.CallResourceResponse) error {
 	return nil
 })

pkg/services/pluginsintegration/clientmiddleware/utils.go

@@ -0,0 +1,11 @@
+package clientmiddleware
+
+import (
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+type callResourceResponseSenderFunc func(res *backend.CallResourceResponse) error
+
+func (fn callResourceResponseSenderFunc) Send(res *backend.CallResourceResponse) error {
+	return fn(res)
+}

pkg/services/pluginsintegration/pluginsintegration.go

@@ -114,6 +114,7 @@ func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthToken
 		clientmiddleware.NewClearAuthHeadersMiddleware(),
 		clientmiddleware.NewOAuthTokenMiddleware(oAuthTokenService),
 		clientmiddleware.NewCookiesMiddleware(skipCookiesNames),
+		clientmiddleware.NewResourceResponseMiddleware(),
 	}
 
 	// Placing the new service implementation behind a feature flag until it is known to be stable