Team LBAC: Add teamHeaders for datasource proxy requests (#76339)

* Add teamHeaders for datasource proxy requests

* adds validation for the teamHeaders

* added tests for applying teamHeaders

* remove previous implementation

* validation for header values being set to authproxy

* removed unnecessary checks

* newline

* Add middleware for injecting headers on the data source backend

* renamed feature toggle

* Get user teams from context

* Fix feature toggle name

* added test for validation of the auth headers and fixed evaluation to cover headers

* renaming of teamHeaders to teamHTTPHeaders

* use of header set for non-existing header and add for existing headers

* moves types into datasources

* fixed unchecked errors

* Refactor

* Add tests for data model

* Update pkg/api/datasources.go

Co-authored-by: Victor Cinaglia <victor@grafana.com>

* Update pkg/api/datasources.go

Co-authored-by: Victor Cinaglia <victor@grafana.com>

---------

Co-authored-by: Alexander Zobnin <alexanderzobnin@gmail.com>
Co-authored-by: Victor Cinaglia <victor@grafana.com>

docs/sources/setup-grafana/configure-grafana/feature-toggles/index.md

@@ -148,6 +148,7 @@ Experimental features might be changed or removed without prior notice.
 | `kubernetesPlaylists`                       | Use the kubernetes API in the frontend for playlists                                                                                                                                                                                                                              |
 | `navAdminSubsections`                       | Splits the administration section of the nav tree into subsections                                                                                                                                                                                                                |
 | `recoveryThreshold`                         | Enables feature recovery threshold (aka hysteresis) for threshold server-side expression                                                                                                                                                                                          |
+| `teamHttpHeaders`                           | Enables datasources to apply team headers to the client requests                                                                                                                                                                                                                  |
 | `awsDatasourcesNewFormStyling`              | Applies new form styling for configuration and query editors in AWS plugins                                                                                                                                                                                                       |
 | `cachingOptimizeSerializationMemoryUsage`   | If enabled, the caching backend gradually serializes query responses for the cache, comparing against the configured `[caching]max_value_mb` value as it goes. This can can help prevent Grafana from running out of memory while attempting to cache very large query responses. |
 | `pluginsInstrumentationStatusSource`        | Include a status source label for plugin request metrics and logs                                                                                                                                                                                                                 |

packages/grafana-data/src/types/featureToggles.gen.ts

@@ -141,6 +141,7 @@ export interface FeatureToggles {
   kubernetesPlaylists?: boolean;
   navAdminSubsections?: boolean;
   recoveryThreshold?: boolean;
+  teamHttpHeaders?: boolean;
   awsDatasourcesNewFormStyling?: boolean;
   cachingOptimizeSerializationMemoryUsage?: boolean;
   panelTitleSearchInV1?: boolean;

pkg/api/datasources.go

@@ -18,6 +18,7 @@ import (
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/services/auth/identity"
+	"github.com/grafana/grafana/pkg/services/contexthandler"
 	contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/setting"
@@ -333,7 +334,7 @@ func validateURL(cmdType string, url string) response.Response {
 // validateJSONData prevents the user from adding a custom header with name that matches the auth proxy header name.
 // This is done to prevent data source proxy from being used to circumvent auth proxy.
 // For more context take a look at CVE-2022-35957
-func validateJSONData(jsonData *simplejson.Json, cfg *setting.Cfg) error {
+func validateJSONData(ctx context.Context, jsonData *simplejson.Json, cfg *setting.Cfg) error {
 	if jsonData == nil || !cfg.AuthProxyEnabled {
 		return nil
 	}
@@ -347,6 +348,34 @@ func validateJSONData(jsonData *simplejson.Json, cfg *setting.Cfg) error {
 			}
 		}
 	}
+
+	// Prevent adding a data source team header with a name that matches the auth proxy header name
+	list := contexthandler.AuthHTTPHeaderListFromContext(ctx)
+	if list == nil {
+		return nil
+	}
+	teamHTTPHeadersJSON := datasources.TeamHTTPHeadersJSONData{}
+	if jsonData != nil {
+		jsonData, err := jsonData.MarshalJSON()
+		if err != nil {
+			return err
+		}
+		err = json.Unmarshal(jsonData, &teamHTTPHeadersJSON)
+		if err != nil {
+			return err
+		}
+		for _, headers := range teamHTTPHeadersJSON.TeamHTTPHeaders {
+			for _, header := range headers {
+				for _, name := range list.Items {
+					if http.CanonicalHeaderKey(header.Header) == http.CanonicalHeaderKey(name) {
+						datasourcesLogger.Error("Cannot add a data source team header with a used by our proxy header", "headerName", header.Header)
+						return errors.New("validation error, invalid header name specified")
+					}
+				}
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -387,7 +416,7 @@ func (hs *HTTPServer) AddDataSource(c *contextmodel.ReqContext) response.Respons
 			return resp
 		}
 	}
-	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
+	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to add datasource", err)
 	}
 
@@ -452,7 +481,7 @@ func (hs *HTTPServer) UpdateDataSourceByID(c *contextmodel.ReqContext) response.
 	if resp := validateURL(cmd.Type, cmd.URL); resp != nil {
 		return resp
 	}
-	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
+	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to update datasource", err)
 	}
 
@@ -492,7 +521,7 @@ func (hs *HTTPServer) UpdateDataSourceByUID(c *contextmodel.ReqContext) response
 	if resp := validateURL(cmd.Type, cmd.URL); resp != nil {
 		return resp
 	}
-	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
+	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to update datasource", err)
 	}
 

pkg/api/datasources_test.go

@@ -221,6 +221,47 @@ func TestUpdateDataSource_InvalidJSONData(t *testing.T) {
 	assert.Equal(t, 400, sc.resp.Code)
 }
 
+// Using a team HTTP header whose name matches the name specified for auth proxy header should fail
+func TestUpdateDataSourceTeamHTTPHeaders_InvalidJSONData(t *testing.T) {
+	hs := &HTTPServer{
+		DataSourcesService: &dataSourcesServiceMock{},
+		Cfg:                setting.NewCfg(),
+	}
+	sc := setupScenarioContext(t, "/api/datasources/1234")
+
+	data := datasources.TeamHTTPHeaders{
+		"1234": []datasources.TeamHTTPHeader{
+			// Authorization is used by the auth proxy
+			// As part of
+			// contexthandler.AuthHTTPHeaderListFromContext(ctx)
+			{
+				Header: "Authorization",
+				Value:  "Could be anything",
+			},
+		},
+	}
+
+	hs.Cfg.AuthProxyEnabled = true
+	jsonData := simplejson.New()
+	jsonData.Set("teamHTTPHeaders", data)
+
+	sc.m.Put(sc.url, routing.Wrap(func(c *contextmodel.ReqContext) response.Response {
+		c.Req.Body = mockRequestBody(datasources.AddDataSourceCommand{
+			Name:     "Test",
+			URL:      "localhost:5432",
+			Access:   "direct",
+			Type:     "test",
+			JsonData: jsonData,
+		})
+		c.SignedInUser = authedUserWithPermissions(1, 1, []ac.Permission{})
+		return hs.AddDataSource(c)
+	}))
+
+	sc.fakeReqWithParams("PUT", sc.url, map[string]string{}).exec()
+
+	assert.Equal(t, 400, sc.resp.Code)
+}
+
 // Updating data sources with URLs not specifying protocol should work.
 func TestUpdateDataSource_URLWithoutProtocol(t *testing.T) {
 	const name = "Test"

pkg/api/pluginproxy/ds_proxy.go

@@ -272,6 +272,15 @@ func (proxy *DataSourceProxy) director(req *http.Request) {
 	if proxy.features.IsEnabled(featuremgmt.FlagIdForwarding) {
 		proxyutil.ApplyForwardIDHeader(req, proxy.ctx.SignedInUser)
 	}
+
+	if proxy.features.IsEnabled(featuremgmt.FlagTeamHttpHeaders) {
+		err := proxyutil.ApplyTeamHTTPHeaders(req, proxy.ds, proxy.ctx.Teams)
+		if err != nil {
+			// NOTE: could downgrade the errors to warnings
+			ctxLogger.Error("Error applying teamHTTPHeaders", "error", err)
+			return
+		}
+	}
 }
 
 func (proxy *DataSourceProxy) validateRequest() error {

pkg/services/datasources/models.go

@@ -1,6 +1,7 @@
 package datasources
 
 import (
+	"encoding/json"
 	"time"
 
 	"github.com/grafana/grafana/pkg/components/simplejson"
@@ -65,6 +66,33 @@ type DataSource struct {
 	Updated time.Time `json:"updated,omitempty"`
 }
 
+type TeamHTTPHeadersJSONData struct {
+	TeamHTTPHeaders TeamHTTPHeaders `json:"teamHttpHeaders"`
+}
+
+type TeamHTTPHeaders map[string][]TeamHTTPHeader
+
+type TeamHTTPHeader struct {
+	Header string `json:"header"`
+	Value  string `json:"value"`
+}
+
+func (ds DataSource) TeamHTTPHeaders() (TeamHTTPHeaders, error) {
+	teamHTTPHeadersJSON := TeamHTTPHeadersJSONData{}
+	if ds.JsonData != nil {
+		jsonData, err := ds.JsonData.MarshalJSON()
+		if err != nil {
+			return nil, err
+		}
+		err = json.Unmarshal(jsonData, &teamHTTPHeadersJSON)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return teamHTTPHeadersJSON.TeamHTTPHeaders, nil
+}
+
 // AllowedCookies parses the jsondata.keepCookies and returns a list of
 // allowed cookies, otherwise an empty list.
 func (ds DataSource) AllowedCookies() []string {

pkg/services/datasources/models_test.go

@@ -58,3 +58,45 @@ func TestAllowedCookies(t *testing.T) {
 		})
 	}
 }
+
+func TestTeamHTTPHeaders(t *testing.T) {
+	testCases := []struct {
+		desc  string
+		given string
+		want  TeamHTTPHeaders
+	}{
+		{
+			desc:  "Usual json data with teamHttpHeaders",
+			given: `{"teamHttpHeaders": {"101": [{"header": "X-CUSTOM-HEADER", "value": "foo"}]}}`,
+			want: TeamHTTPHeaders{
+				"101": {
+					{Header: "X-CUSTOM-HEADER", Value: "foo"},
+				},
+			},
+		},
+		{
+			desc:  "Json data without teamHttpHeaders",
+			given: `{"foo": "bar"}`,
+			want:  nil,
+		},
+	}
+
+	for _, test := range testCases {
+		t.Run(test.desc, func(t *testing.T) {
+			jsonDataBytes := []byte(test.given)
+			jsonData, err := simplejson.NewJson(jsonDataBytes)
+			require.NoError(t, err)
+
+			ds := DataSource{
+				ID:       1235,
+				JsonData: jsonData,
+				UID:      "test",
+			}
+
+			actual, err := ds.TeamHTTPHeaders()
+			assert.NoError(t, err)
+			assert.Equal(t, test.want, actual)
+			assert.EqualValues(t, test.want, actual)
+		})
+	}
+}

pkg/services/featuremgmt/registry.go

@@ -862,6 +862,13 @@ var (
 			Owner:           grafanaAlertingSquad,
 			RequiresRestart: true,
 		},
+		{
+			Name:         "teamHttpHeaders",
+			Description:  "Enables datasources to apply team headers to the client requests",
+			Stage:        FeatureStageExperimental,
+			FrontendOnly: false,
+			Owner:        grafanaAuthnzSquad,
+		},
 		{
 			Name:         "awsDatasourcesNewFormStyling",
 			Description:  "Applies new form styling for configuration and query editors in AWS plugins",

pkg/services/featuremgmt/toggles_gen.csv

@@ -122,6 +122,7 @@ transformationsVariableSupport,experimental,@grafana/grafana-bi-squad,false,fals
 kubernetesPlaylists,experimental,@grafana/grafana-app-platform-squad,false,false,false,true
 navAdminSubsections,experimental,@grafana/grafana-frontend-platform,false,false,false,false
 recoveryThreshold,experimental,@grafana/alerting-squad,false,false,true,false
+teamHttpHeaders,experimental,@grafana/grafana-authnz-team,false,false,false,false
 awsDatasourcesNewFormStyling,experimental,@grafana/aws-datasources,false,false,false,true
 cachingOptimizeSerializationMemoryUsage,experimental,@grafana/grafana-operator-experience-squad,false,false,false,false
 panelTitleSearchInV1,experimental,@grafana/backend-platform,true,false,false,false

pkg/services/featuremgmt/toggles_gen.go

@@ -499,6 +499,10 @@ const (
 	// Enables feature recovery threshold (aka hysteresis) for threshold server-side expression
 	FlagRecoveryThreshold = "recoveryThreshold"
 
+	// FlagTeamHttpHeaders
+	// Enables datasources to apply team headers to the client requests
+	FlagTeamHttpHeaders = "teamHttpHeaders"
+
 	// FlagAwsDatasourcesNewFormStyling
 	// Applies new form styling for configuration and query editors in AWS plugins
 	FlagAwsDatasourcesNewFormStyling = "awsDatasourcesNewFormStyling"

pkg/services/pluginsintegration/clientmiddleware/team_headers_middleware.go

@@ -0,0 +1,131 @@
+package clientmiddleware
+
+import (
+	"context"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/infra/appcontext"
+	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/services/contexthandler"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/util/proxyutil"
+)
+
+// NewTeamHTTPHeaderMiddleware creates a new plugins.ClientMiddleware that will
+// set headers based on teams user is member of.
+func NewTeamHTTPHeadersMiddleware() plugins.ClientMiddleware {
+	return plugins.ClientMiddlewareFunc(func(next plugins.Client) plugins.Client {
+		return &TeamHTTPHeadersMiddleware{
+			next: next,
+		}
+	})
+}
+
+type TeamHTTPHeadersMiddleware struct {
+	next plugins.Client
+}
+
+func (m *TeamHTTPHeadersMiddleware) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
+	if req == nil {
+		return m.next.QueryData(ctx, req)
+	}
+
+	err := m.setHeaders(ctx, req.PluginContext, req)
+	if err != nil {
+		return nil, err
+	}
+
+	return m.next.QueryData(ctx, req)
+}
+
+func (m *TeamHTTPHeadersMiddleware) CallResource(ctx context.Context, req *backend.CallResourceRequest, sender backend.CallResourceResponseSender) error {
+	if req == nil {
+		return m.next.CallResource(ctx, req, sender)
+	}
+
+	err := m.setHeaders(ctx, req.PluginContext, req)
+	if err != nil {
+		return err
+	}
+
+	return m.next.CallResource(ctx, req, sender)
+}
+
+func (m *TeamHTTPHeadersMiddleware) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	if req == nil {
+		return m.next.CheckHealth(ctx, req)
+	}
+
+	// NOTE: might not be needed to set headers. we want to for now set these headers
+	err := m.setHeaders(ctx, req.PluginContext, req)
+	if err != nil {
+		return nil, err
+	}
+
+	return m.next.CheckHealth(ctx, req)
+}
+
+func (m *TeamHTTPHeadersMiddleware) CollectMetrics(ctx context.Context, req *backend.CollectMetricsRequest) (*backend.CollectMetricsResult, error) {
+	return m.next.CollectMetrics(ctx, req)
+}
+
+func (m *TeamHTTPHeadersMiddleware) SubscribeStream(ctx context.Context, req *backend.SubscribeStreamRequest) (*backend.SubscribeStreamResponse, error) {
+	return m.next.SubscribeStream(ctx, req)
+}
+
+func (m *TeamHTTPHeadersMiddleware) PublishStream(ctx context.Context, req *backend.PublishStreamRequest) (*backend.PublishStreamResponse, error) {
+	return m.next.PublishStream(ctx, req)
+}
+
+func (m *TeamHTTPHeadersMiddleware) RunStream(ctx context.Context, req *backend.RunStreamRequest, sender *backend.StreamSender) error {
+	return m.next.RunStream(ctx, req, sender)
+}
+
+func (m *TeamHTTPHeadersMiddleware) setHeaders(ctx context.Context, pCtx backend.PluginContext, req interface{}) error {
+	reqCtx := contexthandler.FromContext(ctx)
+	// if request not for a datasource or no HTTP request context skip middleware
+	if req == nil || pCtx.DataSourceInstanceSettings == nil || reqCtx == nil || reqCtx.Req == nil {
+		return nil
+	}
+
+	settings := pCtx.DataSourceInstanceSettings
+	jsonDataBytes, err := simplejson.NewJson(settings.JSONData)
+	if err != nil {
+		return err
+	}
+
+	ds := &datasources.DataSource{
+		ID:       settings.ID,
+		OrgID:    pCtx.OrgID,
+		JsonData: jsonDataBytes,
+		Updated:  settings.Updated,
+	}
+
+	signedInUser, err := appcontext.User(ctx)
+	if err != nil {
+		return nil // no user
+	}
+
+	teamHTTPHeaders, err := proxyutil.GetTeamHTTPHeaders(ds, signedInUser.GetTeams())
+	if err != nil {
+		return err
+	}
+
+	switch t := req.(type) {
+	case *backend.QueryDataRequest:
+		for key, value := range teamHTTPHeaders {
+			t.SetHTTPHeader(key, value)
+		}
+	case *backend.CheckHealthRequest:
+		for key, value := range teamHTTPHeaders {
+			t.SetHTTPHeader(key, value)
+		}
+	case *backend.CallResourceRequest:
+		for key, value := range teamHTTPHeaders {
+			t.SetHTTPHeader(key, value)
+		}
+	}
+
+	return nil
+}

pkg/services/pluginsintegration/pluginsintegration.go

@@ -179,6 +179,10 @@ func CreateMiddlewares(cfg *setting.Cfg, oAuthTokenService oauthtoken.OAuthToken
 		middlewares = append(middlewares, clientmiddleware.NewUserHeaderMiddleware())
 	}
 
+	if features.IsEnabled(featuremgmt.FlagTeamHttpHeaders) {
+		middlewares = append(middlewares, clientmiddleware.NewTeamHTTPHeadersMiddleware())
+	}
+
 	middlewares = append(middlewares, clientmiddleware.NewHTTPClientMiddleware())
 
 	return middlewares

pkg/util/proxyutil/proxyutil.go

@@ -5,9 +5,11 @@ import (
 	"net"
 	"net/http"
 	"sort"
+	"strconv"
 	"strings"
 
 	"github.com/grafana/grafana/pkg/services/auth/identity"
+	"github.com/grafana/grafana/pkg/services/datasources"
 )
 
 const (
@@ -130,3 +132,58 @@ func ApplyForwardIDHeader(req *http.Request, user identity.Requester) {
 		req.Header.Set(IDHeaderName, token)
 	}
 }
+
+func ApplyTeamHTTPHeaders(req *http.Request, ds *datasources.DataSource, teams []int64) error {
+	headers, err := GetTeamHTTPHeaders(ds, teams)
+	if err != nil {
+		return err
+	}
+
+	for header, value := range headers {
+		// check if headerv is already set in req.Header
+		if req.Header.Get(header) != "" {
+			req.Header.Add(header, value)
+		} else {
+			req.Header.Set(header, value)
+		}
+	}
+
+	return nil
+}
+
+func GetTeamHTTPHeaders(ds *datasources.DataSource, teams []int64) (map[string]string, error) {
+	teamHTTPHeadersMap := make(map[string]string)
+	teamHTTPHeaders, err := ds.TeamHTTPHeaders()
+	if err != nil {
+		return nil, err
+	}
+
+	for teamID, headers := range teamHTTPHeaders {
+		id, err := strconv.ParseInt(teamID, 10, 64)
+		if err != nil {
+			// FIXME: logging here
+			continue
+		}
+
+		if !contains(teams, id) {
+			continue
+		}
+
+		for _, header := range headers {
+			// TODO: handle multiple header values
+			// add tests for these cases
+			teamHTTPHeadersMap[header.Header] = header.Value
+		}
+	}
+
+	return teamHTTPHeadersMap, nil
+}
+
+func contains(slice []int64, value int64) bool {
+	for _, v := range slice {
+		if v == value {
+			return true
+		}
+	}
+	return false
+}

pkg/util/proxyutil/proxyutil_test.go

@@ -6,6 +6,8 @@ import (
 
 	"github.com/stretchr/testify/require"
 
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/user"
 )
 
@@ -203,3 +205,67 @@ func TestApplyUserHeader(t *testing.T) {
 		require.Equal(t, "admin", req.Header.Get("X-Grafana-User"))
 	})
 }
+
+func TestApplyteamHTTPHeaders(t *testing.T) {
+	t.Run("Should not apply team headers for users that are not part of the teams", func(t *testing.T) {
+		req, err := http.NewRequest(http.MethodGet, "/", nil)
+		require.NoError(t, err)
+		ds := &datasources.DataSource{
+			JsonData: simplejson.New(),
+		}
+		// add team headers
+		ds.JsonData.Set("teamHTTPHeaders", map[string]interface{}{
+			"1": []map[string]interface{}{
+				{
+					"header": "X-Team-Header",
+					"value":  "1",
+				},
+			},
+			"2": []map[string]interface{}{
+				{
+					"header": "X-Prom-Label-Policy",
+					"value":  "2",
+				},
+			},
+			// user is not part of this team
+			"3": []map[string]interface{}{
+				{
+					"header": "X-Custom-Label-Policy",
+					"value":  "3",
+				},
+			},
+		})
+
+		err = ApplyTeamHTTPHeaders(req, ds, []int64{1, 2})
+		require.NoError(t, err)
+		require.Contains(t, req.Header, "X-Team-Header")
+		require.Contains(t, req.Header, "X-Prom-Label-Policy")
+		require.NotContains(t, req.Header, "X-Custom-Label-Policy")
+	})
+	t.Run("Should apply team headers", func(t *testing.T) {
+		req, err := http.NewRequest(http.MethodGet, "/", nil)
+		require.NoError(t, err)
+		ds := &datasources.DataSource{
+			JsonData: simplejson.New(),
+		}
+		ds.JsonData.Set("teamHTTPHeaders", map[string]interface{}{
+			"1": []map[string]interface{}{
+				{
+					"header": "X-Team-Header",
+					"value":  "1",
+				},
+			},
+			"2": []map[string]interface{}{
+				{
+					"header": "X-Prom-Label-Policy",
+					"value":  "2",
+				},
+			},
+		})
+
+		err = ApplyTeamHTTPHeaders(req, ds, []int64{1, 2})
+		require.NoError(t, err)
+		require.Contains(t, req.Header, "X-Team-Header")
+		require.Contains(t, req.Header, "X-Prom-Label-Policy")
+	})
+}