Add ARM-specific bearer token policy

sdk/azcore/CHANGELOG.md

@@ -9,11 +9,15 @@
 * `runtime.NewPipeline` has a new signature that simplifies implementing custom authentication
 * `arm/runtime.RegistrationOptions` embeds `policy.ClientOptions`
 * Contents in the `log` package have been slightly renamed.
+* Removed `AuthenticationOptions` in favor of `policy.BearerTokenOptions`
+* Changed parameters for `NewBearerTokenPolicy()`
+* Moved policy config options out of `arm/runtime` and into `arm/policy`
 
 ### Features Added
 * Updating Documentation
 * Added string typdef `arm.Endpoint` to provide a hint toward expected ARM client endpoints
 * `azcore.ClientOptions` contains common pipeline configuration settings
+* Added support for multi-tenant authorization in `arm/runtime`
 
 ### Bug Fixes
 * Fixed a potential panic when creating the default Transporter.

sdk/azcore/arm/policy/policy.go

@@ -0,0 +1,44 @@
+//go:build go1.16
+// +build go1.16
+
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package policy
+
+import (
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+// BearerTokenOptions configures the bearer token policy's behavior.
+type BearerTokenOptions struct {
+	// Scopes contains the list of permission scopes required for the token.
+	Scopes []string
+	// AuxiliaryTenants contains a list of additional tenant IDs to be used to authenticate
+	// in cross-tenant applications.
+	AuxiliaryTenants []string
+}
+
+// RegistrationOptions configures the registration policy's behavior.
+// All zero-value fields will be initialized with their default values.
+type RegistrationOptions struct {
+	policy.ClientOptions
+
+	// MaxAttempts is the total number of times to attempt automatic registration
+	// in the event that an attempt fails.
+	// The default value is 3.
+	// Set to a value less than zero to disable the policy.
+	MaxAttempts int
+
+	// PollingDelay is the amount of time to sleep between polling intervals.
+	// The default value is 15 seconds.
+	// A value less than zero means no delay between polling intervals (not recommended).
+	PollingDelay time.Duration
+
+	// PollingDuration is the amount of time to wait before abandoning polling.
+	// The default valule is 5 minutes.
+	// NOTE: Setting this to a small value might cause the policy to prematurely fail.
+	PollingDuration time.Duration
+}

sdk/azcore/arm/runtime/pipeline.go

@@ -9,9 +9,10 @@ package runtime
 import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
+	armpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/pipeline"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared"
-	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	azpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 )
 
@@ -25,19 +26,16 @@ func NewPipeline(module, version string, cred azcore.TokenCredential, options *a
 	if len(ep) == 0 {
 		ep = arm.AzurePublicCloud
 	}
-	perCallPolicies := []policy.Policy{}
+	perCallPolicies := []azpolicy.Policy{}
 	if !options.DisableRPRegistration {
-		regRPOpts := RegistrationOptions{ClientOptions: options.ClientOptions}
+		regRPOpts := armpolicy.RegistrationOptions{ClientOptions: options.ClientOptions}
 		perCallPolicies = append(perCallPolicies, NewRPRegistrationPolicy(string(ep), cred, &regRPOpts))
 	}
-	perRetryPolicies := []policy.Policy{
-		azruntime.NewBearerTokenPolicy(cred, azruntime.AuthenticationOptions{
-			TokenRequest: policy.TokenRequestOptions{
-				Scopes: []string{shared.EndpointToScope(string(ep))},
-			},
+	perRetryPolicies := []azpolicy.Policy{
+		NewBearerTokenPolicy(cred, &armpolicy.BearerTokenOptions{
+			Scopes:           []string{shared.EndpointToScope(string(ep))},
 			AuxiliaryTenants: options.AuxiliaryTenants,
-		},
-		),
+		}),
 	}
 	return azruntime.NewPipeline(module, version, perCallPolicies, perRetryPolicies, &options.ClientOptions)
 }

sdk/azcore/arm/runtime/policy_bearer_token.go

@@ -0,0 +1,98 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package runtime
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	armpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/arm/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared"
+	azpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+type acquiringResourceState struct {
+	ctx    context.Context
+	p      *BearerTokenPolicy
+	tenant string
+}
+
+// acquire acquires or updates the resource; only one
+// thread/goroutine at a time ever calls this function
+func acquire(state interface{}) (newResource interface{}, newExpiration time.Time, err error) {
+	s := state.(acquiringResourceState)
+	tk, err := s.p.cred.GetToken(s.ctx, azpolicy.TokenRequestOptions{
+		Scopes:   s.p.options.Scopes,
+		TenantID: s.tenant,
+	})
+	if err != nil {
+		return nil, time.Time{}, err
+	}
+	return tk, tk.ExpiresOn, nil
+}
+
+// BearerTokenPolicy authorizes requests with bearer tokens acquired from a TokenCredential.
+type BearerTokenPolicy struct {
+	// mainResource is the resource to be retreived using the tenant specified in the credential
+	mainResource *shared.ExpiringResource
+	// auxResources are additional resources that are required for cross-tenant applications
+	auxResources map[string]*shared.ExpiringResource
+	// the following fields are read-only
+	cred    azcore.TokenCredential
+	options armpolicy.BearerTokenOptions
+}
+
+// NewBearerTokenPolicy creates a policy object that authorizes requests with bearer tokens.
+// cred: an azcore.TokenCredential implementation such as a credential object from azidentity
+// opts: optional settings. Pass nil to accept default values; this is the same as passing a zero-value options.
+func NewBearerTokenPolicy(cred azcore.TokenCredential, opts *armpolicy.BearerTokenOptions) *BearerTokenPolicy {
+	if opts == nil {
+		opts = &armpolicy.BearerTokenOptions{}
+	}
+	p := &BearerTokenPolicy{
+		cred:         cred,
+		options:      *opts,
+		mainResource: shared.NewExpiringResource(acquire),
+	}
+	if len(opts.AuxiliaryTenants) > 0 {
+		p.auxResources = map[string]*shared.ExpiringResource{}
+	}
+	for _, t := range opts.AuxiliaryTenants {
+		p.auxResources[t] = shared.NewExpiringResource(acquire)
+
+	}
+	return p
+}
+
+// Do authorizes a request with a bearer token
+func (b *BearerTokenPolicy) Do(req *azpolicy.Request) (*http.Response, error) {
+	as := acquiringResourceState{
+		ctx: req.Raw().Context(),
+		p:   b,
+	}
+	tk, err := b.mainResource.GetResource(as)
+	if err != nil {
+		return nil, err
+	}
+	if token, ok := tk.(*azcore.AccessToken); ok {
+		req.Raw().Header.Set(shared.HeaderAuthorization, shared.BearerTokenPrefix+token.Token)
+	}
+	auxTokens := []string{}
+	for tenant, er := range b.auxResources {
+		as.tenant = tenant
+		auxTk, err := er.GetResource(as)
+		if err != nil {
+			return nil, err
+		}
+		auxTokens = append(auxTokens, fmt.Sprintf("%s%s", shared.BearerTokenPrefix, auxTk.(*azcore.AccessToken).Token))
+	}
+	if len(auxTokens) > 0 {
+		req.Raw().Header.Set(shared.HeaderAuxiliaryAuthorization, strings.Join(auxTokens, ", "))
+	}
+	return req.Next()
+}