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
2024-10-25 10:07:53 +01:00
parent 226dcdde0f
commit b1e1297bb3
5 changed files with 284 additions and 244 deletions
@@ -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.
@@ -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
+}
@@ -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) {