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/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)
+}