Datasources: provide generic function to extract custom headers (#66738)

pkg/api/datasources.go

@@ -340,7 +340,7 @@ func validateJSONData(jsonData *simplejson.Json, cfg *setting.Cfg) error {
 	}
 
 	for key, value := range jsonData.MustMap() {
-		if strings.HasPrefix(key, "httpHeaderName") {
+		if strings.HasPrefix(key, datasources.CustomHeaderName) {
 			header := fmt.Sprint(value)
 			if http.CanonicalHeaderKey(header) == http.CanonicalHeaderKey(cfg.AuthProxyHeaderName) {
 				datasourcesLogger.Error("Forbidden to add a data source header with a name equal to auth proxy header name", "headerName", key)

pkg/services/datasources/datasources.go

@@ -54,6 +54,11 @@ type DataSourceService interface {
 	// DecryptedPassword decrypts the encrypted datasource password and returns the
 	// decrypted value.
 	DecryptedPassword(ctx context.Context, ds *DataSource) (string, error)
+
+	// CustomHeaders returns a map of custom headers the user might have
+	// configured for this Datasource. Not every datasource can has the option
+	// to configure those.
+	CustomHeaders(ctx context.Context, ds *DataSource) (map[string]string, error)
 }
 
 // CacheService interface for retrieving a cached datasource.

pkg/services/datasources/fakes/fake_datasource_service.go

@@ -133,3 +133,7 @@ func (s *FakeDataSourceService) DecryptedBasicAuthPassword(ctx context.Context,
 func (s *FakeDataSourceService) DecryptedPassword(ctx context.Context, ds *datasources.DataSource) (string, error) {
 	return "", nil
 }
+
+func (s *FakeDataSourceService) CustomHeaders(ctx context.Context, ds *datasources.DataSource) (map[string]string, error) {
+	return nil, nil
+}

pkg/services/datasources/models.go

@@ -28,6 +28,10 @@ const (
 	DS_ES_OPEN_DISTRO = "grafana-es-open-distro-datasource"
 	DS_ES_OPENSEARCH  = "grafana-opensearch-datasource"
 	DS_AZURE_MONITOR  = "grafana-azure-monitor-datasource"
+	// CustomHeaderName is the prefix that is used to store the name of a custom header.
+	CustomHeaderName = "httpHeaderName"
+	// CustomHeaderValue is the prefix that is used to store the value of a custom header.
+	CustomHeaderValue = "httpHeaderValue"
 )
 
 type DsAccess string

pkg/services/datasources/service/datasource.go

@@ -561,8 +561,8 @@ func (s *Service) getCustomHeaders(jsonData *simplejson.Json, decryptedValues ma
 	index := 0
 	for {
 		index++
-		headerNameSuffix := fmt.Sprintf("httpHeaderName%d", index)
-		headerValueSuffix := fmt.Sprintf("httpHeaderValue%d", index)
+		headerNameSuffix := fmt.Sprintf("%s%d", datasources.CustomHeaderName, index)
+		headerValueSuffix := fmt.Sprintf("%s%d", datasources.CustomHeaderValue, index)
 
 		key := jsonData.Get(headerNameSuffix).MustString()
 		if key == "" {
@@ -651,3 +651,12 @@ func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
 	limits.Set(orgQuotaTag, cfg.Quota.Org.DataSource)
 	return limits, nil
 }
+
+// CustomerHeaders returns the custom headers specified in the datasource. The context is used for the decryption operation that might use the store, so consider setting an acceptable timeout for your use case.
+func (s *Service) CustomHeaders(ctx context.Context, ds *datasources.DataSource) (map[string]string, error) {
+	values, err := s.SecretsService.DecryptJsonData(ctx, ds.SecureJsonData)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get custom headers: %w", err)
+	}
+	return s.getCustomHeaders(ds.JsonData, values), nil
+}

pkg/services/datasources/service/datasource_test.go

@@ -707,6 +707,78 @@ func TestService_GetDecryptedValues(t *testing.T) {
 	})
 }
 
+func TestDataSource_CustomHeaders(t *testing.T) {
+	sqlStore := db.InitTestDB(t)
+	secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
+	secretsStore := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
+	quotaService := quotatest.New(false, nil)
+	dsService, err := ProvideService(sqlStore, secretsService, secretsStore, nil, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService(), quotaService)
+	require.NoError(t, err)
+
+	dsService.cfg = setting.NewCfg()
+
+	testValue := "HeaderValue1"
+
+	encryptedValue, err := secretsService.Encrypt(context.Background(), []byte(testValue), secrets.WithoutScope())
+	require.NoError(t, err)
+
+	testCases := []struct {
+		name             string
+		jsonData         *simplejson.Json
+		secureJsonData   map[string][]byte
+		expectedHeaders  map[string]string
+		expectedErrorMsg string
+	}{
+		{
+			name: "valid custom headers",
+			jsonData: simplejson.NewFromAny(map[string]interface{}{
+				"httpHeaderName1": "X-Test-Header1",
+			}),
+			secureJsonData: map[string][]byte{
+				"httpHeaderValue1": encryptedValue,
+			},
+			expectedHeaders: map[string]string{
+				"X-Test-Header1": testValue,
+			},
+		},
+		{
+			name: "missing header value",
+			jsonData: simplejson.NewFromAny(map[string]interface{}{
+				"httpHeaderName1": "X-Test-Header1",
+			}),
+			secureJsonData:  map[string][]byte{},
+			expectedHeaders: map[string]string{},
+		},
+		{
+			name: "non customer header value",
+			jsonData: simplejson.NewFromAny(map[string]interface{}{
+				"someotherheader": "X-Test-Header1",
+			}),
+			secureJsonData:  map[string][]byte{},
+			expectedHeaders: map[string]string{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ds := &datasources.DataSource{
+				JsonData:       tc.jsonData,
+				SecureJsonData: tc.secureJsonData,
+			}
+
+			headers, err := dsService.CustomHeaders(context.Background(), ds)
+
+			if tc.expectedErrorMsg != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tc.expectedErrorMsg)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tc.expectedHeaders, headers)
+			}
+		})
+	}
+}
+
 const caCert string = `-----BEGIN CERTIFICATE-----
 MIIDATCCAemgAwIBAgIJAMQ5hC3CPDTeMA0GCSqGSIb3DQEBCwUAMBcxFTATBgNV
 BAMMDGNhLWs4cy1zdGhsbTAeFw0xNjEwMjcwODQyMjdaFw00NDAzMTQwODQyMjda