MSAL for On-Behalf-Of token requests

go.mod

@@ -3,6 +3,7 @@ module github.com/grafana/azure-data-explorer-datasource
 go 1.17
 
 require (
+	github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0
 	github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.13.2
 	github.com/google/go-cmp v0.5.7
 	github.com/grafana/grafana-azure-sdk-go v1.1.0
@@ -17,7 +18,6 @@ require (
 require (
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v0.22.0 // indirect
 	github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.1 // indirect
-	github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0 // indirect
 	github.com/apache/arrow/go/arrow v0.0.0-20211112161151-bc219186db40 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect

pkg/azuredx/azureauth/aad_client.go

@@ -0,0 +1,69 @@
+package azureauth
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"regexp"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+	"github.com/AzureAD/microsoft-authentication-library-for-go/apps/confidential"
+	"github.com/grafana/grafana-azure-sdk-go/azcredentials"
+	"github.com/grafana/grafana-azure-sdk-go/azsettings"
+)
+
+// Abstraction over confidential.Client from MSAL for Go
+type aadClient interface {
+	AcquireTokenOnBehalfOf(ctx context.Context, userAssertion string, scopes []string) (confidential.AuthResult, error)
+}
+
+func newAADClient(credentials *azcredentials.AzureClientSecretCredentials, httpClient *http.Client) (aadClient, error) {
+	authority, err := resolveAuthorityForCloud(credentials.AzureCloud)
+	if err != nil {
+		return nil, fmt.Errorf("invalid Azure credentials: %w", err)
+	}
+
+	if !validTenantId(credentials.TenantId) {
+		return nil, errors.New("invalid tenantId")
+	}
+
+	clientCredential, err := confidential.NewCredFromSecret(credentials.ClientSecret)
+	if err != nil {
+		return nil, err
+	}
+
+	authorityOpts := confidential.WithAuthority(runtime.JoinPaths(string(authority), credentials.TenantId))
+	httpClientOPts := confidential.WithHTTPClient(httpClient)
+
+	client, err := confidential.New(credentials.ClientId, clientCredential, authorityOpts, httpClientOPts)
+	if err != nil {
+		return nil, err
+	}
+
+	return client, nil
+}
+
+func resolveAuthorityForCloud(cloudName string) (azidentity.AuthorityHost, error) {
+	// Known Azure clouds
+	switch cloudName {
+	case azsettings.AzurePublic:
+		return azidentity.AzurePublicCloud, nil
+	case azsettings.AzureChina:
+		return azidentity.AzureChina, nil
+	case azsettings.AzureUSGovernment:
+		return azidentity.AzureGovernment, nil
+	default:
+		err := fmt.Errorf("the Azure cloud '%s' not supported", cloudName)
+		return "", err
+	}
+}
+
+func validTenantId(tenantId string) bool {
+	match, err := regexp.MatchString("^[0-9a-zA-Z-.]+$", tenantId)
+	if err != nil {
+		return false
+	}
+	return match
+}

pkg/azuredx/azureauth/auth.go

@@ -2,17 +2,11 @@ package azureauth
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"net/http"
-	"net/url"
-	"strconv"
-	"strings"
 	"time"
 
-	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/grafana/grafana-azure-sdk-go/azcredentials"
 	"github.com/grafana/grafana-azure-sdk-go/azsettings"
 	"github.com/grafana/grafana-azure-sdk-go/aztokenprovider"
@@ -41,27 +35,22 @@ type ServiceCredentials interface {
 }
 
 type ServiceCredentialsImpl struct {
-	models.DatasourceSettings
-	// HTTPDo is the http.Client Do method.
-	HTTPDo        func(req *http.Request) (*http.Response, error)
-	authority     azidentity.AuthorityHost
+	OnBehalfOf   bool
+	QueryTimeout time.Duration
+
 	tokenProvider aztokenprovider.AzureTokenProvider
 	tokenCache    *cache
+	aadClient     aadClient
 	scopes        []string
 }
 
 func NewServiceCredentials(settings *models.DatasourceSettings, azureSettings *azsettings.AzureSettings,
-	client *http.Client) (ServiceCredentials, error) {
+	httpClient *http.Client) (ServiceCredentials, error) {
 	azureCloud, err := normalizeAzureCloud(settings.AzureCloud)
 	if err != nil {
 		return nil, fmt.Errorf("invalid Azure credentials: %w", err)
 	}
 
-	authority, err := resolveAuthorityForCloud(azureCloud)
-	if err != nil {
-		return nil, fmt.Errorf("invalid Azure credentials: %w", err)
-	}
-
 	credentials := &azcredentials.AzureClientSecretCredentials{
 		AzureCloud:   azureCloud,
 		TenantId:     settings.TenantID,
@@ -74,18 +63,23 @@ func NewServiceCredentials(settings *models.DatasourceSettings, azureSettings *a
 		return nil, fmt.Errorf("invalid Azure configuration: %w", err)
 	}
 
+	aadClient, err := newAADClient(credentials, httpClient)
+	if err != nil {
+		return nil, fmt.Errorf("invalid Azure configuration: %w", err)
+	}
+
 	scopes, err := getAzureScopes(credentials, settings.ClusterURL)
 	if err != nil {
 		return nil, fmt.Errorf("invalid Azure configuration: %w", err)
 	}
 
 	return &ServiceCredentialsImpl{
-		DatasourceSettings: *settings,
-		HTTPDo:             client.Do,
-		authority:          authority,
-		tokenProvider:      tokenProvider,
-		tokenCache:         newCache(),
-		scopes:             scopes,
+		OnBehalfOf:    settings.OnBehalfOf,
+		QueryTimeout:  settings.QueryTimeout,
+		tokenProvider: tokenProvider,
+		tokenCache:    newCache(),
+		aadClient:     aadClient,
+		scopes:        scopes,
 	}, nil
 }
 
@@ -147,71 +141,10 @@ func (c *ServiceCredentialsImpl) queryDataOnBehalfOf(ctx context.Context, req *b
 	return "Bearer " + onBehalfOfToken, nil
 }
 
-// OnBehalfOf resolves a token which impersonates the subject of userToken.
-// UserToken has to be an ID token. See the following link for more detail.
-// https://docs.microsoft.com/en-us/azure/active-directory/develop/v2-oauth2-on-behalf-of-flow
-func (c *ServiceCredentialsImpl) onBehalfOf(ctx context.Context, userToken string) (onBehalfOfToken string, expire time.Time, err error) {
-	params := make(url.Values)
-	params.Set("client_id", c.ClientID)
-	params.Set("client_secret", c.Secret)
-	params.Set("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer")
-	params.Set("assertion", userToken)
-	params.Set("scope", strings.Join(c.scopes, " "))
-	params.Set("requested_token_use", "on_behalf_of")
-	reqBody := strings.NewReader(params.Encode())
-
-	// https://docs.microsoft.com/en-us/azure/active-directory/develop/active-directory-v2-protocols#endpoints
-	tokenURL := fmt.Sprintf("%s%s/oauth2/v2.0/token", c.authority, url.PathEscape(c.TenantID))
-	req, err := http.NewRequestWithContext(ctx, "POST", tokenURL, reqBody)
-	if err != nil {
-		return "", time.Time{}, fmt.Errorf("on-behalf-of grant request <%q> instantiation: %w", tokenURL, err)
-	}
-	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
-	req.Header.Set("Accept", "application/json")
-
-	// HTTP Exchange
-	reqStart := time.Now()
-	resp, err := c.HTTPDo(req)
+func (c *ServiceCredentialsImpl) onBehalfOf(ctx context.Context, idToken string) (onBehalfOfToken string, expire time.Time, err error) {
+	result, err := c.aadClient.AcquireTokenOnBehalfOf(ctx, idToken, c.scopes)
 	if err != nil {
-		return "", time.Time{}, fmt.Errorf("on-behalf-of grant POST <%q>: %w", tokenURL, err)
-	}
-	body, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return "", time.Time{}, fmt.Errorf("on-behalf-of grant POST <%q> response: %w", tokenURL, err)
-	}
-	onBehalfOfLatencySeconds.WithLabelValues(strconv.Itoa(resp.StatusCode)).Observe(float64(time.Since(reqStart)) / float64(time.Second))
-	if resp.StatusCode/100 != 2 {
-		var deny struct {
-			Desc string `json:"error_description"`
-		}
-		_ = json.Unmarshal(body, &deny)
-		return "", time.Time{}, fmt.Errorf("on-behalf-of grant POST <%q> status %q: %q", tokenURL, resp.Status, deny.Desc)
-	}
-
-	// Parse Essentials
-	var grant struct {
-		Token  string `json:"access_token"`
-		Expire int    `json:"expires_in"`
-	}
-	if err := json.Unmarshal(body, &grant); err != nil {
-		return "", time.Time{}, fmt.Errorf("malformed response from on-behalf-of grant POST <%q>: %w", tokenURL, err)
-	}
-
-	expire = time.Now().Add(time.Duration(grant.Expire) * time.Second)
-	return grant.Token, expire, nil
-}
-
-func resolveAuthorityForCloud(cloudName string) (azidentity.AuthorityHost, error) {
-	// Known Azure clouds
-	switch cloudName {
-	case azsettings.AzurePublic:
-		return azidentity.AzurePublicCloud, nil
-	case azsettings.AzureChina:
-		return azidentity.AzureChina, nil
-	case azsettings.AzureUSGovernment:
-		return azidentity.AzureGovernment, nil
-	default:
-		err := fmt.Errorf("the Azure cloud '%s' not supported", cloudName)
-		return "", err
+		return "", time.Time{}, err
 	}
+	return result.AccessToken, result.ExpiresOn.UTC(), nil
 }