LBAC for datasources: Move validation of rules from datasources to LBAC Rules (#94622)

* FIX: Remove the checks for lbac rules inside of datasources

* Remove json validation for lbac rules

* Preserve lbac rules in updates

* Refactored test to remove the table structure

* refactor: change to allow naming and concise override instead of complex branching

* refactor to make sure we set an empty field for updates

* bugfix

* check for datasources.JsonData

* fix merge

* add datasource to check for field presence only

* add function call for readability

pkg/api/datasources.go

@@ -6,14 +6,10 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
-	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 
-	"github.com/prometheus/prometheus/promql/parser"
-	"golang.org/x/exp/slices"
-
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 
 	"github.com/grafana/grafana/pkg/api/datasource"
@@ -22,7 +18,6 @@ import (
 	"github.com/grafana/grafana/pkg/apimachinery/identity"
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/infra/log"
-	ac "github.com/grafana/grafana/pkg/services/accesscontrol"
 	contextmodel "github.com/grafana/grafana/pkg/services/contexthandler/model"
 	"github.com/grafana/grafana/pkg/services/datasources"
 	"github.com/grafana/grafana/pkg/services/featuremgmt"
@@ -339,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(ctx context.Context, jsonData *simplejson.Json, cfg *setting.Cfg, features featuremgmt.FeatureToggles) error {
+func validateJSONData(jsonData *simplejson.Json, cfg *setting.Cfg) error {
 	if jsonData == nil {
 		return nil
 	}
@@ -356,69 +351,9 @@ func validateJSONData(ctx context.Context, jsonData *simplejson.Json, cfg *setti
 		}
 	}
 
-	// Prevent adding a data source team header with a name that matches the auth proxy header name
-	if features.IsEnabled(ctx, featuremgmt.FlagTeamHttpHeaders) {
-		err := validateTeamHTTPHeaderJSON(jsonData)
-		if err != nil {
-			return err
-		}
-	}
-
 	return nil
 }
 
-// we only allow for now the following headers to be added to a data source team header
-var validHeaders = []string{"X-Prom-Label-Policy"}
-
-func validateTeamHTTPHeaderJSON(jsonData *simplejson.Json) error {
-	teamHTTPHeadersJSON, err := datasources.GetTeamHTTPHeaders(jsonData)
-	if err != nil {
-		datasourcesLogger.Error("Unable to marshal TeamHTTPHeaders")
-		return errors.New("validation error, invalid format of TeamHTTPHeaders")
-	}
-	if teamHTTPHeadersJSON == nil {
-		return nil
-	}
-	// whitelisting ValidHeaders
-	// each teams headers
-	for _, teamheaders := range teamHTTPHeadersJSON.Headers {
-		for _, header := range teamheaders {
-			if !slices.ContainsFunc(validHeaders, func(v string) bool {
-				return http.CanonicalHeaderKey(v) == http.CanonicalHeaderKey(header.Header)
-			}) {
-				datasourcesLogger.Error("Cannot add a data source team header that is different than", "headerName", header.Header)
-				return errors.New("validation error, invalid header name specified")
-			}
-			if !validateLBACHeader(header.Value) {
-				datasourcesLogger.Error("Cannot add a data source team header value with invalid value", "headerValue", header.Value)
-				return errors.New("validation error, invalid header value syntax")
-			}
-		}
-	}
-	return nil
-}
-
-// validateLBACHeader returns true if the header value matches the syntax
-// 1234:{ name!="value",foo!~"bar" }
-func validateLBACHeader(headervalue string) bool {
-	exp := `^\d+:(.+)`
-	pattern, err := regexp.Compile(exp)
-	if err != nil {
-		return false
-	}
-	match := pattern.FindSubmatch([]byte(strings.TrimSpace(headervalue)))
-	if match == nil || len(match) < 2 {
-		return false
-	}
-	_, err = parser.ParseMetricSelector(string(match[1]))
-	return err == nil
-}
-
-func evaluateTeamHTTPHeaderPermissions(hs *HTTPServer, c *contextmodel.ReqContext, scope string) (bool, error) {
-	ev := ac.EvalPermission(datasources.ActionPermissionsWrite, ac.Scope(scope))
-	return hs.AccessControl.Evaluate(c.Req.Context(), c.SignedInUser, ev)
-}
-
 // swagger:route POST /datasources datasources addDataSource
 //
 // Create a data source.
@@ -453,10 +388,20 @@ func (hs *HTTPServer) AddDataSource(c *contextmodel.ReqContext) response.Respons
 		}
 	}
 
-	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg, hs.Features); err != nil {
+	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to add datasource", err)
 	}
 
+	// It's forbidden to update the rules from the datasource api.
+	// team HTTP headers update have to be done through `updateDatasourceLBACRules`
+	if hs.Features != nil && hs.Features.IsEnabled(c.Req.Context(), featuremgmt.FlagTeamHttpHeaders) {
+		if cmd.JsonData != nil {
+			if _, ok := cmd.JsonData.CheckGet("teamHttpHeaders"); ok {
+				return response.Error(http.StatusForbidden, "Cannot create datasource with team HTTP headers, need to use updateDatasourceLBACRules API", nil)
+			}
+		}
+	}
+
 	dataSource, err := hs.DataSourcesService.AddDataSource(c.Req.Context(), &cmd)
 	if err != nil {
 		if errors.Is(err, datasources.ErrDataSourceNameExists) || errors.Is(err, datasources.ErrDataSourceUidExists) {
@@ -518,7 +463,7 @@ func (hs *HTTPServer) UpdateDataSourceByID(c *contextmodel.ReqContext) response.
 	if resp := validateURL(cmd.Type, cmd.URL); resp != nil {
 		return resp
 	}
-	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg, hs.Features); err != nil {
+	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to update datasource", err)
 	}
 	ds, err := hs.getRawDataSourceById(c.Req.Context(), cmd.ID, cmd.OrgID)
@@ -529,12 +474,6 @@ func (hs *HTTPServer) UpdateDataSourceByID(c *contextmodel.ReqContext) response.
 		return response.Error(http.StatusInternalServerError, "Failed to update datasource", err)
 	}
 
-	// check if LBAC rules have been modified
-	hasAccess, errAccess := checkTeamHTTPHeaderPermissions(hs, c, ds, cmd)
-	if !hasAccess {
-		return response.Error(http.StatusForbidden, fmt.Sprintf("You'll need additional permissions to perform this action. Permissions needed: %s", datasources.ActionPermissionsWrite), errAccess)
-	}
-
 	return hs.updateDataSourceByID(c, ds, cmd)
 }
 
@@ -564,7 +503,7 @@ func (hs *HTTPServer) UpdateDataSourceByUID(c *contextmodel.ReqContext) response
 	if resp := validateURL(cmd.Type, cmd.URL); resp != nil {
 		return resp
 	}
-	if err := validateJSONData(c.Req.Context(), cmd.JsonData, hs.Cfg, hs.Features); err != nil {
+	if err := validateJSONData(cmd.JsonData, hs.Cfg); err != nil {
 		return response.Error(http.StatusBadRequest, "Failed to update datasource", err)
 	}
 
@@ -577,36 +516,9 @@ func (hs *HTTPServer) UpdateDataSourceByUID(c *contextmodel.ReqContext) response
 	}
 	cmd.ID = ds.ID
 
-	// check if LBAC rules have been modified
-	hasAccess, errAccess := checkTeamHTTPHeaderPermissions(hs, c, ds, cmd)
-	if !hasAccess {
-		return response.Error(http.StatusForbidden, fmt.Sprintf("You'll need additional permissions to perform this action. Permissions needed: %s", datasources.ActionPermissionsWrite), errAccess)
-	}
-
 	return hs.updateDataSourceByID(c, ds, cmd)
 }
 
-func getEncodedString(jsonData *simplejson.Json, key string) string {
-	if jsonData == nil {
-		return ""
-	}
-	jsonValues, exists := jsonData.CheckGet(key)
-	if !exists {
-		return ""
-	}
-	val, _ := jsonValues.Encode()
-	return string(val)
-}
-
-func checkTeamHTTPHeaderPermissions(hs *HTTPServer, c *contextmodel.ReqContext, ds *datasources.DataSource, cmd datasources.UpdateDataSourceCommand) (bool, error) {
-	currentTeamHTTPHeaders := getEncodedString(ds.JsonData, "teamHttpHeaders")
-	newTeamHTTPHeaders := getEncodedString(cmd.JsonData, "teamHttpHeaders")
-	if (currentTeamHTTPHeaders != "" || newTeamHTTPHeaders != "") && currentTeamHTTPHeaders != newTeamHTTPHeaders {
-		return evaluateTeamHTTPHeaderPermissions(hs, c, datasources.ScopePrefix+ds.UID)
-	}
-	return true, nil
-}
-
 func (hs *HTTPServer) updateDataSourceByID(c *contextmodel.ReqContext, ds *datasources.DataSource, cmd datasources.UpdateDataSourceCommand) response.Response {
 	if ds.ReadOnly {
 		return response.Error(http.StatusForbidden, "Cannot update read-only data source", nil)

pkg/api/datasources_test.go

@@ -223,104 +223,58 @@ 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) {
+func TestAddDataSourceTeamHTTPHeaders(t *testing.T) {
 	tenantID := "1234"
-	testcases := []struct {
-		desc string
-		data datasources.TeamHTTPHeaders
-		want int
-	}{
-		{
-			desc: "We should only allow for headers being X-Prom-Label-Policy",
-			data: datasources.TeamHTTPHeaders{
-				Headers: datasources.TeamHeaders{
-					tenantID: []datasources.TeamHTTPHeader{
-						{
-							Header: "Authorization",
-							Value:  "foo!=bar",
-						},
-					},
-				}},
-			want: 400,
+	hs := &HTTPServer{
+		DataSourcesService: &dataSourcesServiceMock{
+			expectedDatasource: &datasources.DataSource{},
 		},
-		{
-			desc: "Allowed header but no team id",
-			data: datasources.TeamHTTPHeaders{
-				Headers: datasources.TeamHeaders{"": []datasources.TeamHTTPHeader{
-					{
-						Header: "X-Prom-Label-Policy",
-						Value:  "foo=bar",
-					},
-				},
-				}},
-			want: 400,
-		},
-		{
-			desc: "Allowed team id and header name with invalid header values ",
-			data: datasources.TeamHTTPHeaders{
-				Headers: datasources.TeamHeaders{tenantID: []datasources.TeamHTTPHeader{
-					{
-						Header: "X-Prom-Label-Policy",
-						Value:  "Bad value",
-					},
-				},
-				}},
-			want: 400,
-		},
-		// Complete valid case, with team id, header name and header value
-		{
-			desc: "Allowed header and header values ",
-			data: datasources.TeamHTTPHeaders{
-				Headers: datasources.TeamHeaders{tenantID: []datasources.TeamHTTPHeader{
-					{
-						Header: "X-Prom-Label-Policy",
-						Value:  `1234:{ name!="value",foo!~"bar" }`,
-					},
-				},
-				}},
-			want: 200,
+		Cfg:                  setting.NewCfg(),
+		Features:             featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders),
+		accesscontrolService: actest.FakeService{},
+		AccessControl: actest.FakeAccessControl{
+			ExpectedEvaluate: true,
+			ExpectedErr:      nil,
 		},
 	}
-	for _, tc := range testcases {
-		t.Run(tc.desc, func(t *testing.T) {
-			hs := &HTTPServer{
-				DataSourcesService: &dataSourcesServiceMock{
-					expectedDatasource: &datasources.DataSource{},
+	sc := setupScenarioContext(t, fmt.Sprintf("/api/datasources/%s", tenantID))
+	hs.Cfg.AuthProxy.Enabled = true
+
+	jsonData := simplejson.New()
+	jsonData.Set("teamHttpHeaders", datasources.TeamHTTPHeaders{
+		Headers: datasources.TeamHeaders{
+			tenantID: []datasources.TeamHTTPHeader{
+				{
+					Header: "Authorization",
+					Value:  "foo!=bar",
 				},
-				Cfg:                  setting.NewCfg(),
-				Features:             featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders),
-				accesscontrolService: actest.FakeService{},
-				AccessControl: actest.FakeAccessControl{
-					ExpectedEvaluate: true,
-					ExpectedErr:      nil,
-				},
-			}
-			sc := setupScenarioContext(t, fmt.Sprintf("/api/datasources/%s", tenantID))
-			hs.Cfg.AuthProxy.Enabled = true
-
-			jsonData := simplejson.New()
-			jsonData.Set("teamHttpHeaders", tc.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{
-					{Action: datasources.ActionPermissionsWrite, Scope: datasources.ScopeAll},
-				})
-				return hs.AddDataSource(c)
-			}))
-
-			sc.fakeReqWithParams("PUT", sc.url, map[string]string{}).exec()
-
-			assert.Equal(t, tc.want, sc.resp.Code)
+			},
+		},
+	})
+	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{
+			{Action: datasources.ActionPermissionsWrite, Scope: datasources.ScopeAll},
+		})
+		return hs.AddDataSource(c)
+	}))
+
+	sc.fakeReqWithParams("PUT", sc.url, map[string]string{}).exec()
+	assert.Equal(t, http.StatusForbidden, sc.resp.Code)
+
+	// Parse the JSON response
+	var response map[string]string
+	err := json.Unmarshal(sc.resp.Body.Bytes(), &response)
+	assert.NoError(t, err, "Failed to parse JSON response")
+
+	// Check the error message in the JSON response
+	assert.Equal(t, "Cannot create datasource with team HTTP headers, need to use updateDatasourceLBACRules API", response["message"])
 }
 
 // Updating data sources with URLs not specifying protocol should work.
@@ -493,35 +447,6 @@ func TestAPI_datasources_AccessControl(t *testing.T) {
 	}
 }
 
-func TestValidateLBACHeader(t *testing.T) {
-	testcases := []struct {
-		desc            string
-		teamHeaderValue string
-		want            bool
-	}{
-		{
-			desc:            "Should allow valid header",
-			teamHeaderValue: `1234:{ name!="value",foo!~"bar" }`,
-			want:            true,
-		},
-		{
-			desc:            "Should allow valid selector",
-			teamHeaderValue: `1234:{ name!="value",foo!~"bar/baz.foo" }`,
-			want:            true,
-		},
-		{
-			desc:            "Should return false for incorrect header value",
-			teamHeaderValue: `1234:!="value",foo!~"bar" }`,
-			want:            false,
-		},
-	}
-	for _, tc := range testcases {
-		t.Run(tc.desc, func(t *testing.T) {
-			assert.Equal(t, tc.want, validateLBACHeader(tc.teamHeaderValue))
-		})
-	}
-}
-
 type dataSourcesServiceMock struct {
 	datasources.DataSourceService
 

pkg/services/datasources/models.go

@@ -206,7 +206,7 @@ type UpdateDataSourceCommand struct {
 	UpdateSecretFn          UpdateSecretFn    `json:"-"`
 	IgnoreOldSecureJsonData bool              `json:"-"`
 
-	OnlyUpdateLBACRulesFromAPI bool `json:"-"`
+	AllowLBACRuleUpdates bool `json:"-"`
 }
 
 // DeleteDataSourceCommand will delete a DataSource based on OrgID as well as the UID (preferred), ID, or Name.

pkg/services/datasources/service/datasource.go

@@ -534,29 +534,15 @@ func (s *Service) UpdateDataSource(ctx context.Context, cmd *datasources.UpdateD
 			}
 		}
 
-		// TODO: we will eventually remove this check for moving the resource to it's separate API
-		if s.features != nil && s.features.IsEnabled(ctx, featuremgmt.FlagTeamHttpHeaders) && !cmd.OnlyUpdateLBACRulesFromAPI {
-			s.logger.Debug("Overriding LBAC rules with stored ones",
-				"reason", "update_lbac_rules_from_datasource_api",
-				"action", "use_updateLBACRules_API",
+		// preserve existing lbac rules when updating datasource if we're not updating lbac rules
+		// TODO: Refactor to store lbac rules separate from a datasource
+		if s.features != nil && s.features.IsEnabled(ctx, featuremgmt.FlagTeamHttpHeaders) && !cmd.AllowLBACRuleUpdates {
+			s.logger.Debug("Overriding LBAC rules with stored ones using updateLBACRules API",
+				"reason", "overriding_lbac_rules_from_datasource_api",
 				"datasource_id", dataSource.ID,
 				"datasource_uid", dataSource.UID)
 
-			if dataSource.JsonData != nil {
-				previousRules := dataSource.JsonData.Get("teamHttpHeaders").Interface()
-				if previousRules == nil {
-					if cmd.JsonData != nil {
-						cmd.JsonData.Del("teamHttpHeaders")
-					}
-				} else {
-					if cmd.JsonData == nil {
-						// It's fine to instantiate a new JsonData here
-						// Because it's done in the SQLStore.UpdateDataSource anyway
-						cmd.JsonData = simplejson.New()
-					}
-					cmd.JsonData.Set("teamHttpHeaders", previousRules)
-				}
-			}
+			cmd.JsonData = RetainExistingLBACRules(dataSource.JsonData, cmd.JsonData)
 		}
 
 		if cmd.Name != "" && cmd.Name != dataSource.Name {
@@ -1001,3 +987,30 @@ func (s *Service) CustomHeaders(ctx context.Context, ds *datasources.DataSource)
 	}
 	return s.getCustomHeaders(ds.JsonData, values), nil
 }
+
+func RetainExistingLBACRules(storedJsonData, cmdJsonData *simplejson.Json) *simplejson.Json {
+	// If there are no stored data, we should remove the key from the command json data
+	if storedJsonData == nil {
+		if cmdJsonData != nil {
+			cmdJsonData.Del("teamHttpHeaders")
+		}
+		return cmdJsonData
+	}
+
+	previousRules := storedJsonData.Get("teamHttpHeaders").Interface()
+	// If there are no previous rules, we should remove the key from the command json data
+	if previousRules == nil {
+		if cmdJsonData != nil {
+			cmdJsonData.Del("teamHttpHeaders")
+		}
+		return cmdJsonData
+	}
+
+	if cmdJsonData == nil {
+		// It's fine to instantiate a new JsonData here
+		// Because it's done in the SQLStore.UpdateDataSource anyway
+		cmdJsonData = simplejson.New()
+	}
+	cmdJsonData.Set("teamHttpHeaders", previousRules)
+	return cmdJsonData
+}

pkg/services/datasources/service/datasource_test.go

@@ -569,6 +569,196 @@ func TestService_UpdateDataSource(t *testing.T) {
 		require.True(t, mutateExecuted)
 		require.Equal(t, "test-datasource-updated", dsUpdated.Name)
 	})
+
+	t.Run("Should update LBAC rules when updating from API", func(t *testing.T) {
+		dsService := initDSService(t)
+		dsService.features = featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders)
+
+		// Create a datasource with existing LBAC rules
+		existingRules := []interface{}{
+			map[string]interface{}{
+				"name":  "X-Grafana-Team",
+				"value": "team1",
+			},
+		}
+		jsonData := simplejson.NewFromAny(map[string]interface{}{
+			"teamHttpHeaders": existingRules,
+		})
+
+		ds, err := dsService.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+			OrgID:    1,
+			Name:     "test-datasource",
+			Type:     "prometheus",
+			JsonData: jsonData,
+		})
+		require.NoError(t, err)
+		// Verify that the datasource was created with the correct JsonData
+		createdDS, err := dsService.GetDataSource(context.Background(), &datasources.GetDataSourceQuery{
+			OrgID: ds.OrgID,
+			ID:    ds.ID,
+		})
+		require.NoError(t, err)
+		require.NotNil(t, createdDS.JsonData)
+		createdRules := createdDS.JsonData.Get("teamHttpHeaders").MustArray()
+		require.Equal(t, existingRules, createdRules)
+
+		// Update the datasource with new LBAC rules from API
+		newRules := []interface{}{
+			map[string]interface{}{
+				"name":  "X-Grafana-Team",
+				"value": "team2",
+			},
+		}
+		updateCmd := &datasources.UpdateDataSourceCommand{
+			ID:    ds.ID,
+			OrgID: ds.OrgID,
+			Name:  "updated-datasource",
+			Type:  "prometheus",
+			JsonData: simplejson.NewFromAny(map[string]interface{}{
+				"teamHttpHeaders": newRules,
+			}),
+			AllowLBACRuleUpdates: true,
+		}
+
+		updatedDS, err := dsService.UpdateDataSource(context.Background(), updateCmd)
+		require.NoError(t, err)
+
+		// Check if the LBAC rules are updated
+		updatedRules := updatedDS.JsonData.Get("teamHttpHeaders").MustArray()
+		require.Equal(t, newRules, updatedRules)
+	})
+	t.Run("Should preserve LBAC rules when not updating from API", func(t *testing.T) {
+		dsService := initDSService(t)
+		dsService.features = featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders)
+		// Create a datasource with existing LBAC rules
+		existingRules := []interface{}{
+			map[string]interface{}{
+				"name":  "X-Grafana-Team",
+				"value": "team1",
+			},
+		}
+		jsonData := simplejson.NewFromAny(map[string]interface{}{
+			"teamHttpHeaders": existingRules,
+		})
+
+		ds, err := dsService.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+			OrgID:    1,
+			Name:     "test-datasource",
+			Type:     "prometheus",
+			JsonData: jsonData,
+		})
+		require.NoError(t, err)
+		// Verify that the datasource was created with the correct JsonData
+		createdDS, err := dsService.GetDataSource(context.Background(), &datasources.GetDataSourceQuery{
+			OrgID: ds.OrgID,
+			ID:    ds.ID,
+		})
+		require.NoError(t, err)
+		require.NotNil(t, createdDS.JsonData)
+		createdRules := createdDS.JsonData.Get("teamHttpHeaders").MustArray()
+		require.Equal(t, existingRules, createdRules)
+
+		// Update the datasource without LBAC rules in the command
+		updateCmd := &datasources.UpdateDataSourceCommand{
+			ID:    ds.ID,
+			OrgID: ds.OrgID,
+			Name:  "updated-datasource",
+			Type:  "prometheus",
+			JsonData: simplejson.NewFromAny(map[string]interface{}{
+				"someOtherSetting": "value",
+			}),
+			AllowLBACRuleUpdates: false,
+		}
+
+		updatedDS, err := dsService.UpdateDataSource(context.Background(), updateCmd)
+		require.NoError(t, err)
+
+		// Check if the LBAC rules are preserved
+		updatedRules := updatedDS.JsonData.Get("teamHttpHeaders").MustArray()
+		require.Equal(t, existingRules, updatedRules)
+	})
+
+	t.Run("Should not remove stored rules without AllowLBACRuleUpdates", func(t *testing.T) {
+		dsService := initDSService(t)
+		dsService.features = featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders)
+
+		// Create a datasource with existing LBAC rules
+		existingRules := []interface{}{
+			map[string]interface{}{
+				"name":  "X-Grafana-Team",
+				"value": "team1",
+			},
+		}
+		jsonData := simplejson.NewFromAny(map[string]interface{}{
+			"teamHttpHeaders": existingRules,
+		})
+
+		ds, err := dsService.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+			OrgID:    1,
+			Name:     "test-datasource",
+			Type:     "prometheus",
+			JsonData: jsonData,
+		})
+		require.NoError(t, err)
+
+		// Update the datasource without any LBAC rules in the command
+		updateCmd := &datasources.UpdateDataSourceCommand{
+			ID:                   ds.ID,
+			OrgID:                ds.OrgID,
+			Name:                 "updated-datasource",
+			Type:                 "prometheus",
+			AllowLBACRuleUpdates: false,
+		}
+
+		updatedDS, err := dsService.UpdateDataSource(context.Background(), updateCmd)
+		require.NoError(t, err)
+
+		// Check if the LBAC rules are preserved
+		updatedRules := updatedDS.JsonData.Get("teamHttpHeaders").MustArray()
+		require.Equal(t, existingRules, updatedRules)
+	})
+
+	t.Run("Should not populate empty stored rules without AllowLBACRuleUpdates", func(t *testing.T) {
+		dsService := initDSService(t)
+		dsService.features = featuremgmt.WithFeatures(featuremgmt.FlagTeamHttpHeaders)
+
+		// Create a datasource with empty LBAC rules
+		jsonData := simplejson.New()
+
+		ds, err := dsService.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+			OrgID:    1,
+			Name:     "test-datasource",
+			Type:     "prometheus",
+			JsonData: jsonData,
+		})
+		require.NoError(t, err)
+
+		// Update the datasource with new LBAC rules but without AllowLBACRuleUpdates
+		newRules := []interface{}{
+			map[string]interface{}{
+				"name":  "X-Grafana-Team",
+				"value": "team2",
+			},
+		}
+		updateCmd := &datasources.UpdateDataSourceCommand{
+			ID:    ds.ID,
+			OrgID: ds.OrgID,
+			Name:  "updated-datasource",
+			Type:  "prometheus",
+			JsonData: simplejson.NewFromAny(map[string]interface{}{
+				"teamHttpHeaders": newRules,
+			}),
+			AllowLBACRuleUpdates: false,
+		}
+
+		updatedDS, err := dsService.UpdateDataSource(context.Background(), updateCmd)
+		require.NoError(t, err)
+
+		// Check if the LBAC rules are still empty
+		updatedRules, ok := updatedDS.JsonData.CheckGet("teamHttpHeaders")
+		require.False(t, ok)
+		require.Nil(t, updatedRules)
+	})
 }
 
 func TestService_DeleteDataSource(t *testing.T) {