Shared Azure middleware between Azure Monitor and Prometheus datasources (#46002)

* Scopes in Azure middleware

* Enable Azure middleware without feature flag

* Use common Azure middleware in Azure Monitor

* Apply feature flag to JsonData configuration of Azure auth

* Enforce feature flag in Prometheus datasource

* Prometheus provider tests

* Datasource service tests

* Fix http client provider tests

* Pass sdkhttpclient.Options by reference

* Add middleware to httpclient.Options

* Remove dependency on Grafana settings

* Unit-tests updated

* Fix ds_proxy_test

* Fix service_test

pkg/tsdb/prometheus/promclient/provider.go

@@ -4,6 +4,8 @@ import (
 	"strings"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/prometheus/middleware"
 	"github.com/grafana/grafana/pkg/util/maputil"
 
@@ -19,6 +21,8 @@ type Provider struct {
 	jsonData       map[string]interface{}
 	httpMethod     string
 	clientProvider httpclient.Provider
+	cfg            *setting.Cfg
+	features       featuremgmt.FeatureToggles
 	log            log.Logger
 }
 
@@ -26,6 +30,8 @@ func NewProvider(
 	settings backend.DataSourceInstanceSettings,
 	jsonData map[string]interface{},
 	clientProvider httpclient.Provider,
+	cfg *setting.Cfg,
+	features featuremgmt.FeatureToggles,
 	log log.Logger,
 ) *Provider {
 	httpMethod, _ := maputil.GetStringOptional(jsonData, "httpMethod")
@@ -34,6 +40,8 @@ func NewProvider(
 		jsonData:       jsonData,
 		httpMethod:     httpMethod,
 		clientProvider: clientProvider,
+		cfg:            cfg,
+		features:       features,
 		log:            log,
 	}
 }
@@ -53,7 +61,7 @@ func (p *Provider) GetClient(headers map[string]string) (apiv1.API, error) {
 	}
 
 	// Azure authentication
-	err = p.configureAzureAuthentication(opts)
+	err = p.configureAzureAuthentication(&opts)
 	if err != nil {
 		return nil, err
 	}

pkg/tsdb/prometheus/promclient/provider_azure.go

@@ -2,13 +2,22 @@ package promclient
 
 import (
 	"fmt"
+	"net/url"
+	"path"
 
 	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/tsdb/azuremonitor/azcredentials"
+	"github.com/grafana/grafana/pkg/tsdb/azuremonitor/azhttpclient"
 	"github.com/grafana/grafana/pkg/util/maputil"
 )
 
-func (p *Provider) configureAzureAuthentication(opts sdkhttpclient.Options) error {
+func (p *Provider) configureAzureAuthentication(opts *sdkhttpclient.Options) error {
+	// Azure authentication is experimental (#35857)
+	if !p.features.IsEnabled(featuremgmt.FlagPrometheusAzureAuth) {
+		return nil
+	}
+
 	credentials, err := azcredentials.FromDatasourceData(p.jsonData, p.settings.DecryptedSecureJSONData)
 	if err != nil {
 		err = fmt.Errorf("invalid Azure credentials: %s", err)
@@ -16,16 +25,24 @@ func (p *Provider) configureAzureAuthentication(opts sdkhttpclient.Options) erro
 	}
 
 	if credentials != nil {
-		opts.CustomOptions["_azureCredentials"] = credentials
-
-		resourceId, err := maputil.GetStringOptional(p.jsonData, "azureEndpointResourceId")
+		resourceIdStr, err := maputil.GetStringOptional(p.jsonData, "azureEndpointResourceId")
 		if err != nil {
 			return err
+		} else if resourceIdStr == "" {
+			err := fmt.Errorf("endpoint resource ID (audience) not provided")
+			return err
 		}
 
-		if resourceId != "" {
-			opts.CustomOptions["azureEndpointResourceId"] = resourceId
+		resourceId, err := url.Parse(resourceIdStr)
+		if err != nil || resourceId.Scheme == "" || resourceId.Host == "" {
+			err := fmt.Errorf("endpoint resource ID (audience) '%s' invalid", resourceIdStr)
+			return err
 		}
+
+		resourceId.Path = path.Join(resourceId.Path, ".default")
+		scopes := []string{resourceId.String()}
+
+		azhttpclient.AddAzureAuthentication(opts, p.cfg.Azure, credentials, scopes)
 	}
 
 	return nil

pkg/tsdb/prometheus/promclient/provider_azure_test.go

@@ -0,0 +1,150 @@
+package promclient
+
+import (
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestConfigureAzureAuthentication(t *testing.T) {
+	cfg := &setting.Cfg{}
+	settings := backend.DataSourceInstanceSettings{}
+
+	t.Run("given feature flag enabled", func(t *testing.T) {
+		features := featuremgmt.WithFeatures(featuremgmt.FlagPrometheusAzureAuth)
+
+		t.Run("should set Azure middleware when JsonData contains valid credentials", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod": "POST",
+				"azureCredentials": map[string]interface{}{
+					"authType": "msi",
+				},
+				"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			require.NoError(t, err)
+
+			require.NotNil(t, opts.Middlewares)
+			assert.Len(t, opts.Middlewares, 1)
+		})
+
+		t.Run("should not set Azure middleware when JsonData doesn't contain valid credentials", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod": "POST",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			require.NoError(t, err)
+
+			assert.NotContains(t, opts.CustomOptions, "_azureCredentials")
+		})
+
+		t.Run("should return error when JsonData contains invalid credentials", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod":       "POST",
+				"azureCredentials": "invalid",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			assert.Error(t, err)
+		})
+
+		t.Run("should set Azure middleware when JsonData contains credentials and valid audience", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod": "POST",
+				"azureCredentials": map[string]interface{}{
+					"authType": "msi",
+				},
+				"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			require.NoError(t, err)
+
+			require.NotNil(t, opts.Middlewares)
+			assert.Len(t, opts.Middlewares, 1)
+		})
+
+		t.Run("should not set Azure middleware when JsonData doesn't contain credentials", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod":              "POST",
+				"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			require.NoError(t, err)
+
+			if opts.Middlewares != nil {
+				assert.Len(t, opts.Middlewares, 0)
+			}
+		})
+
+		t.Run("should return error when JsonData contains invalid audience", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod": "POST",
+				"azureCredentials": map[string]interface{}{
+					"authType": "msi",
+				},
+				"azureEndpointResourceId": "invalid",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			assert.Error(t, err)
+		})
+	})
+
+	t.Run("given feature flag not enabled", func(t *testing.T) {
+		features := featuremgmt.WithFeatures()
+
+		t.Run("should not set Azure Credentials even when JsonData contains credentials", func(t *testing.T) {
+			jsonData := map[string]interface{}{
+				"httpMethod": "POST",
+				"azureCredentials": map[string]interface{}{
+					"authType": "msi",
+				},
+				"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5",
+			}
+
+			var p = NewProvider(settings, jsonData, nil, cfg, features, nil)
+
+			var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
+
+			err := p.configureAzureAuthentication(opts)
+			require.NoError(t, err)
+
+			if opts.Middlewares != nil {
+				assert.Len(t, opts.Middlewares, 0)
+			}
+		})
+	})
+}

pkg/tsdb/prometheus/promclient/provider_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/grafana/grafana/pkg/tsdb/prometheus/promclient"
 
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/setting"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
@@ -138,9 +139,11 @@ func setup(jsonData ...string) *testContext {
 	var jd map[string]interface{}
 	_ = json.Unmarshal(rawData, &jd)
 
+	cfg := &setting.Cfg{}
 	settings := backend.DataSourceInstanceSettings{URL: "test-url", JSONData: rawData}
+	features := featuremgmt.WithFeatures()
 	hp := &fakeHttpClientProvider{}
-	p := promclient.NewProvider(settings, jd, hp, nil)
+	p := promclient.NewProvider(settings, jd, hp, cfg, features, nil)
 
 	return &testContext{
 		httpProvider:       hp,

pkg/tsdb/prometheus/prometheus.go

@@ -7,15 +7,16 @@ import (
 	"fmt"
 	"regexp"
 
-	"github.com/grafana/grafana/pkg/tsdb/prometheus/promclient"
-
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 	"github.com/grafana/grafana/pkg/infra/httpclient"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/intervalv2"
+	"github.com/grafana/grafana/pkg/tsdb/prometheus/promclient"
 	"github.com/grafana/grafana/pkg/util/maputil"
 	apiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 )
@@ -32,16 +33,16 @@ type Service struct {
 	tracer             tracing.Tracer
 }
 
-func ProvideService(httpClientProvider httpclient.Provider, tracer tracing.Tracer) *Service {
+func ProvideService(httpClientProvider httpclient.Provider, cfg *setting.Cfg, features featuremgmt.FeatureToggles, tracer tracing.Tracer) *Service {
 	plog.Debug("initializing")
 	return &Service{
 		intervalCalculator: intervalv2.NewCalculator(),
-		im:                 datasource.NewInstanceManager(newInstanceSettings(httpClientProvider)),
+		im:                 datasource.NewInstanceManager(newInstanceSettings(httpClientProvider, cfg, features)),
 		tracer:             tracer,
 	}
 }
 
-func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.InstanceFactoryFunc {
+func newInstanceSettings(httpClientProvider httpclient.Provider, cfg *setting.Cfg, features featuremgmt.FeatureToggles) datasource.InstanceFactoryFunc {
 	return func(settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
 		var jsonData map[string]interface{}
 		err := json.Unmarshal(settings.JSONData, &jsonData)
@@ -49,7 +50,7 @@ func newInstanceSettings(httpClientProvider httpclient.Provider) datasource.Inst
 			return nil, fmt.Errorf("error reading settings: %w", err)
 		}
 
-		p := promclient.NewProvider(settings, jsonData, httpClientProvider, plog)
+		p := promclient.NewProvider(settings, jsonData, httpClientProvider, cfg, features, plog)
 		pc, err := promclient.NewProviderCache(p)
 		if err != nil {
 			return nil, err