Repository: Azure/azure-sdk-for-go
Pull request #20446 (2023-03-21 21:37:59)

Merge azcore beta branch to main
There are two ARM features here: CAE and cross-tenant auth, along with a couple new `TokenRequestOptions` fields that will eventually support similar features for other clients as well. I reverted the tracing commit because that feature will remain in beta.


Patch:
diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md
index f42e5a0e81..1b5ac8057a 100644
--- a/sdk/azcore/CHANGELOG.md
+++ b/sdk/azcore/CHANGELOG.md
@@ -1,6 +1,6 @@
 # Release History
 
-## 1.4.1 (Unreleased)
+## 1.5.0 (Unreleased)
 
 ### Features Added
 
@@ -11,6 +11,11 @@
 
 ### Other Changes
 
+## 1.5.0-beta.1 (2023-03-02)
+
+### Features Added
+* This release includes the features added in v1.4.0-beta.1
+
 ## 1.4.0 (2023-03-02)
 > This release doesn't include features added in v1.4.0-beta.1. They will return in v1.5.0-beta.1.
 
diff --git a/sdk/azcore/arm/policy/policy.go b/sdk/azcore/arm/policy/policy.go
index 7a700d661e..83cf91e3ec 100644
--- a/sdk/azcore/arm/policy/policy.go
+++ b/sdk/azcore/arm/policy/policy.go
@@ -14,6 +14,12 @@ import (
 
 // BearerTokenOptions configures the bearer token policy's behavior.
 type BearerTokenOptions struct {
+	// AuxiliaryTenants are additional tenant IDs for authenticating cross-tenant requests.
+	// The policy will add a token from each of these tenants to every request. The
+	// authenticating user or service principal must be a guest in these tenants, and the
+	// policy's credential must support multitenant authentication.
+	AuxiliaryTenants []string
+
 	// Scopes contains the list of permission scopes required for the token.
 	Scopes []string
 }
@@ -44,6 +50,12 @@ type RegistrationOptions struct {
 type ClientOptions struct {
 	policy.ClientOptions
 
+	// AuxiliaryTenants are additional tenant IDs for authenticating cross-tenant requests.
+	// The client will add a token from each of these tenants to every request. The
+	// authenticating user or service principal must be a guest in these tenants, and the
+	// client's credential must support multitenant authentication.
+	AuxiliaryTenants []string
+
 	// DisableRPRegistration disables the auto-RP registration policy. Defaults to false.
 	DisableRPRegistration bool
 }
diff --git a/sdk/azcore/arm/runtime/pipeline.go b/sdk/azcore/arm/runtime/pipeline.go
index 8da2153307..266c74b17b 100644
--- a/sdk/azcore/arm/runtime/pipeline.go
+++ b/sdk/azcore/arm/runtime/pipeline.go
@@ -28,7 +28,10 @@ func NewPipeline(module, version string, cred azcore.TokenCredential, plOpts azr
 	if err != nil {
 		return azruntime.Pipeline{}, err
 	}
-	authPolicy := NewBearerTokenPolicy(cred, &armpolicy.BearerTokenOptions{Scopes: []string{conf.Audience + "/.default"}})
+	authPolicy := NewBearerTokenPolicy(cred, &armpolicy.BearerTokenOptions{
+		AuxiliaryTenants: options.AuxiliaryTenants,
+		Scopes:           []string{conf.Audience + "/.default"},
+	})
 	perRetry := make([]azpolicy.Policy, len(plOpts.PerRetry), len(plOpts.PerRetry)+1)
 	copy(perRetry, plOpts.PerRetry)
 	plOpts.PerRetry = append(perRetry, authPolicy)
diff --git a/sdk/azcore/arm/runtime/policy_bearer_token.go b/sdk/azcore/arm/runtime/policy_bearer_token.go
index 10a3606e9d..83f1bf86e6 100644
--- a/sdk/azcore/arm/runtime/policy_bearer_token.go
+++ b/sdk/azcore/arm/runtime/policy_bearer_token.go
@@ -5,6 +5,7 @@ package runtime
 
 import (
 	"context"
+	"encoding/base64"
 	"fmt"
 	"net/http"
 	"strings"
@@ -14,9 +15,13 @@ import (
 	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"
+	azruntime "github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/temporal"
 )
 
+const headerAuxiliaryAuthorization = "x-ms-authorization-auxiliary"
+
+// acquiringResourceState holds data for an auxiliary token request
 type acquiringResourceState struct {
 	ctx    context.Context
 	p      *BearerTokenPolicy
@@ -26,7 +31,10 @@ type acquiringResourceState struct {
 // acquire acquires or updates the resource; only one
 // thread/goroutine at a time ever calls this function
 func acquire(state acquiringResourceState) (newResource azcore.AccessToken, newExpiration time.Time, err error) {
-	tk, err := state.p.cred.GetToken(state.ctx, azpolicy.TokenRequestOptions{Scopes: state.p.options.Scopes})
+	tk, err := state.p.cred.GetToken(state.ctx, azpolicy.TokenRequestOptions{
+		Scopes:   state.p.scopes,
+		TenantID: state.tenant,
+	})
 	if err != nil {
 		return azcore.AccessToken{}, time.Time{}, err
 	}
@@ -35,13 +43,10 @@ func acquire(state acquiringResourceState) (newResource azcore.AccessToken, newE
 
 // 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 *temporal.Resource[azcore.AccessToken, acquiringResourceState]
-	// auxResources are additional resources that are required for cross-tenant applications
 	auxResources map[string]*temporal.Resource[azcore.AccessToken, acquiringResourceState]
-	// the following fields are read-only
-	cred    azcore.TokenCredential
-	options armpolicy.BearerTokenOptions
+	btp          *azruntime.BearerTokenPolicy
+	cred         azcore.TokenCredential
+	scopes       []string
 }
 
 // NewBearerTokenPolicy creates a policy object that authorizes requests with bearer tokens.
@@ -51,36 +56,90 @@ func NewBearerTokenPolicy(cred azcore.TokenCredential, opts *armpolicy.BearerTok
 	if opts == nil {
 		opts = &armpolicy.BearerTokenOptions{}
 	}
-	p := &BearerTokenPolicy{
-		cred:         cred,
-		options:      *opts,
-		mainResource: temporal.NewResource(acquire),
+	p := &BearerTokenPolicy{cred: cred}
+	p.auxResources = make(map[string]*temporal.Resource[azcore.AccessToken, acquiringResourceState], len(opts.AuxiliaryTenants))
+	for _, t := range opts.AuxiliaryTenants {
+		p.auxResources[t] = temporal.NewResource(acquire)
 	}
+	p.scopes = make([]string, len(opts.Scopes))
+	copy(p.scopes, opts.Scopes)
+	p.btp = azruntime.NewBearerTokenPolicy(cred, opts.Scopes, &azpolicy.BearerTokenOptions{
+		AuthorizationHandler: azpolicy.AuthorizationHandler{
+			OnChallenge: p.onChallenge,
+			OnRequest:   p.onRequest,
+		},
+	})
 	return p
 }
 
-// Do authorizes a request with a bearer token
-func (b *BearerTokenPolicy) Do(req *azpolicy.Request) (*http.Response, error) {
+func (b *BearerTokenPolicy) onChallenge(req *azpolicy.Request, res *http.Response, authNZ func(azpolicy.TokenRequestOptions) error) error {
+	challenge := res.Header.Get(shared.HeaderWWWAuthenticate)
+	claims, err := parseChallenge(challenge)
+	if err != nil {
+		// the challenge contains claims we can't parse
+		return err
+	} else if claims != "" {
+		// request a new token having the specified claims, send the request again
+		return authNZ(azpolicy.TokenRequestOptions{Claims: claims, Scopes: b.scopes})
+	}
+	// auth challenge didn't include claims, so this is a simple authorization failure
+	return azruntime.NewResponseError(res)
+}
+
+// onRequest authorizes requests with one or more bearer tokens
+func (b *BearerTokenPolicy) onRequest(req *azpolicy.Request, authNZ func(azpolicy.TokenRequestOptions) error) error {
+	// authorize the request with a token for the primary tenant
+	err := authNZ(azpolicy.TokenRequestOptions{Scopes: b.scopes})
+	if err != nil || len(b.auxResources) == 0 {
+		return err
+	}
+	// add tokens for auxiliary tenants
 	as := acquiringResourceState{
 		ctx: req.Raw().Context(),
 		p:   b,
 	}
-	tk, err := b.mainResource.Get(as)
-	if err != nil {
-		return nil, err
-	}
-	req.Raw().Header.Set(shared.HeaderAuthorization, shared.BearerTokenPrefix+tk.Token)
-	auxTokens := []string{}
+	auxTokens := make([]string, 0, len(b.auxResources))
 	for tenant, er := range b.auxResources {
 		as.tenant = tenant
 		auxTk, err := er.Get(as)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		auxTokens = append(auxTokens, fmt.Sprintf("%s%s", shared.BearerTokenPrefix, auxTk.Token))
 	}
-	if len(auxTokens) > 0 {
-		req.Raw().Header.Set(shared.HeaderAuxiliaryAuthorization, strings.Join(auxTokens, ", "))
+	req.Raw().Header.Set(headerAuxiliaryAuthorization, strings.Join(auxTokens, ", "))
+	return nil
+}
+
+// Do authorizes a request with a bearer token
+func (b *BearerTokenPolicy) Do(req *azpolicy.Request) (*http.Response, error) {
+	return b.btp.Do(req)
+}
+
+// parseChallenge parses claims from an authentication challenge issued by ARM so a client can request a token
+// that will satisfy conditional access policies. It returns a non-nil error when the given value contains
+// claims it can't parse. If the value contains no claims, it returns an empty string and a nil error.
+func parseChallenge(wwwAuthenticate string) (string, error) {
+	claims := ""
+	var err error
+	for _, param := range strings.Split(wwwAuthenticate, ",") {
+		if _, after, found := strings.Cut(param, "claims="); found {
+			if claims != "" {
+				// The header contains multiple challenges, at least two of which specify claims. The specs allow this
+				// but it's unclear what a client should do in this case and there's as yet no concrete example of it.
+				err = fmt.Errorf("found multiple claims challenges in %q", wwwAuthenticate)
+				break
+			}
+			// trim stuff that would get an error from RawURLEncoding; claims may or may not be padded
+			claims = strings.Trim(after, `\"=`)
+			// we don't return this error because it's something unhelpful like "illegal base64 data at input byte 42"
+			if b, decErr := base64.RawURLEncoding.DecodeString(claims); decErr == nil {
+				claims = string(b)
+			} else {
+				err = fmt.Errorf("failed to parse claims from %q", wwwAuthenticate)
+				break
+			}
+		}
 	}
-	return req.Next()
+	return claims, err
 }
diff --git a/sdk/azcore/internal/exported/exported.go b/sdk/azcore/internal/exported/exported.go
index 2ffbc0e474..2ef55e3677 100644
--- a/sdk/azcore/internal/exported/exported.go
+++ b/sdk/azcore/internal/exported/exported.go
@@ -71,8 +71,16 @@ type AccessToken struct {
 // TokenRequestOptions contain specific parameter that may be used by credentials types when attempting to get a token.
 // Exported as policy.TokenRequestOptions.
 type TokenRequestOptions struct {
+	// Claims are any additional claims required for the token to satisfy a conditional access policy, such as a
+	// service may return in a claims challenge following an authorization failure. If a service returned the
+	// claims value base64 encoded, it must be decoded before setting this field.
+	Claims string
 	// Scopes contains the list of permission scopes required for the token.
 	Scopes []string
+
+	// TenantID identifies the tenant from which to request the token. azidentity credentials authenticate in
+	// their configured default tenants when this field isn't set.
+	TenantID string
 }
 
 // TokenCredential represents a credential capable of providing an OAuth token.
diff --git a/sdk/azcore/internal/shared/constants.go b/sdk/azcore/internal/shared/constants.go
index 07879a3015..98c6d3fbf4 100644
--- a/sdk/azcore/internal/shared/constants.go
+++ b/sdk/azcore/internal/shared/constants.go
@@ -21,6 +21,7 @@ const (
 	HeaderOperationLocation      = "Operation-Location"
 	HeaderRetryAfter             = "Retry-After"
 	HeaderUserAgent              = "User-Agent"
+	HeaderWWWAuthenticate        = "WWW-Authenticate"
 	HeaderXMSClientRequestID     = "x-ms-client-request-id"
 )
 
@@ -31,5 +32,5 @@ const (
 	Module = "azcore"
 
 	// Version is the semantic version (see http://semver.org) of this module.
-	Version = "v1.4.1"
+	Version = "v1.5.0"
 )


Test patch:
diff --git a/sdk/azcore/arm/runtime/policy_bearer_token_test.go b/sdk/azcore/arm/runtime/policy_bearer_token_test.go
index 04de55b595..1ab06ae00d 100644
--- a/sdk/azcore/arm/runtime/policy_bearer_token_test.go
+++ b/sdk/azcore/arm/runtime/policy_bearer_token_test.go
@@ -5,9 +5,9 @@ package runtime
 
 import (
 	"context"
+	"errors"
 	"strings"
 
-	"errors"
 	"net/http"
 	"testing"
 	"time"
@@ -17,7 +17,9 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/shared"
 	azpolicy "github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/errorinfo"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
+	"github.com/stretchr/testify/require"
 )
 
 const (
@@ -90,10 +92,10 @@ func TestBearerPolicy_SuccessGetToken(t *testing.T) {
 func TestBearerPolicy_CredentialFailGetToken(t *testing.T) {
 	srv, close := mock.NewTLSServer()
 	defer close()
-	expectedErr := errors.New("oops")
+	expectedErr := "oops"
 	failCredential := mockCredential{}
 	failCredential.getTokenImpl = func(ctx context.Context, options azpolicy.TokenRequestOptions) (azcore.AccessToken, error) {
-		return azcore.AccessToken{}, expectedErr
+		return azcore.AccessToken{}, errors.New(expectedErr)
 	}
 	b := NewBearerTokenPolicy(failCredential, nil)
 	pipeline := newTestPipeline(&azpolicy.ClientOptions{
@@ -104,16 +106,11 @@ func TestBearerPolicy_CredentialFailGetToken(t *testing.T) {
 		PerRetryPolicies: []azpolicy.Policy{b},
 	})
 	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
-	if err != nil {
-		t.Fatal(err)
-	}
+	require.NoError(t, err)
 	resp, err := pipeline.Do(req)
-	if err != expectedErr {
-		t.Fatalf("unexpected error: %v", err)
-	}
-	if resp != nil {
-		t.Fatal("expected nil response")
-	}
+	require.EqualError(t, err, expectedErr)
+	require.Nil(t, resp)
+	require.Implements(t, (*errorinfo.NonRetriable)(nil), err)
 }
 
 func TestBearerTokenPolicy_TokenExpired(t *testing.T) {
@@ -165,41 +162,124 @@ func TestBearerPolicy_GetTokenFailsNoDeadlock(t *testing.T) {
 	}
 }
 
-func TestBearerTokenWithAuxiliaryTenants(t *testing.T) {
-	t.Skip("this feature isn't implemented yet")
+func TestAuxiliaryTenants(t *testing.T) {
 	srv, close := mock.NewTLSServer()
 	defer close()
-	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
-	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
-	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
-	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
-	srv.AppendResponse()
-	retryOpts := azpolicy.RetryOptions{
-		MaxRetryDelay: 500 * time.Millisecond,
-		RetryDelay:    50 * time.Millisecond,
-	}
+	srv.SetResponse(mock.WithStatusCode(http.StatusOK))
+	primary := "primary"
+	auxTenants := []string{"aux1", "aux2", "aux3"}
+	expectCache := false
 	b := NewBearerTokenPolicy(
-		mockCredential{},
-		&armpolicy.BearerTokenOptions{
-			Scopes: []string{scope},
-			//AuxiliaryTenants: []string{"tenant1", "tenant2", "tenant3"},
+		mockCredential{
+			// getTokenImpl returns a token whose value equals the requested tenant so the test can validate how the policy handles tenants
+			// i.e., primary tenant token goes in Authorization header and aux tenant tokens go in x-ms-authorization-auxiliary
+			getTokenImpl: func(ctx context.Context, options azpolicy.TokenRequestOptions) (azcore.AccessToken, error) {
+				require.False(t, expectCache, "client should have used a cached token instead of requesting another")
+				tenant := primary
+				if options.TenantID != "" {
+					tenant = options.TenantID
+				}
+				return azcore.AccessToken{Token: tenant, ExpiresOn: time.Now().Add(time.Hour).UTC()}, nil
+			},
 		},
+		&armpolicy.BearerTokenOptions{AuxiliaryTenants: auxTenants, Scopes: []string{scope}},
 	)
-	pipeline := newTestPipeline(&azpolicy.ClientOptions{Transport: srv, Retry: retryOpts, PerRetryPolicies: []azpolicy.Policy{b}})
-	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
-	if err != nil {
-		t.Fatalf("Unexpected error: %v", err)
-	}
-	resp, err := pipeline.Do(req)
-	if err != nil {
-		t.Fatalf("Unexpected error: %v", err)
-	}
-	if resp.StatusCode != http.StatusOK {
-		t.Fatalf("unexpected status code: %d", resp.StatusCode)
+	pipeline := newTestPipeline(&azpolicy.ClientOptions{Transport: srv, PerRetryPolicies: []azpolicy.Policy{b}})
+	expected := strings.Split(shared.BearerTokenPrefix+strings.Join(auxTenants, ","+shared.BearerTokenPrefix), ",")
+	for i := 0; i < 3; i++ {
+		if i == 1 {
+			// policy should have a cached token after the first iteration
+			expectCache = true
+		}
+		req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+		require.NoError(t, err)
+		resp, err := pipeline.Do(req)
+		require.NoError(t, err)
+		require.Equal(t, shared.BearerTokenPrefix+primary, resp.Request.Header.Get(shared.HeaderAuthorization), "Authorization header must contain primary tenant token")
+		actual := strings.Split(resp.Request.Header.Get(headerAuxiliaryAuthorization), ", ")
+		// auxiliary tokens may appear in arbitrary order
+		require.ElementsMatch(t, expected, actual)
 	}
-	expectedHeader := strings.Repeat(shared.BearerTokenPrefix+tokenValue+", ", 3)
-	expectedHeader = expectedHeader[:len(expectedHeader)-2]
-	if auxH := resp.Request.Header.Get(shared.HeaderAuxiliaryAuthorization); auxH != expectedHeader {
-		t.Fatalf("unexpected auxiliary authorization header %s", auxH)
+}
+
+func TestBearerTokenPolicyChallengeParsing(t *testing.T) {
+	for _, test := range []struct {
+		challenge, desc, expectedClaims string
+		err                             error
+	}{
+		{
+			desc: "no challenge",
+		},
+		{
+			desc:      "no claims",
+			challenge: `Bearer authorization_uri="https://login.windows.net/", error="invalid_token", error_description="The authentication failed because of missing 'Authorization' header."`,
+			err:       (*azcore.ResponseError)(nil),
+		},
+		{
+			desc:      "parsing error",
+			challenge: `Bearer claims="invalid"`,
+			// the specific error type isn't important but it must be nonretriable
+			err: (errorinfo.NonRetriable)(nil),
+		},
+		// CAE claims challenges. Position of the "claims" parameter within the challenge shouldn't affect parsing.
+		{
+			desc:           "insufficient claims",
+			challenge:      `Bearer realm="", authorization_uri="https://login.microsoftonline.com/common/oauth2/authorize", client_id="00000003-0000-0000-c000-000000000000", error="insufficient_claims", claims="eyJhY2Nlc3NfdG9rZW4iOiB7ImZvbyI6ICJiYXIifX0="`,
+			expectedClaims: `{"access_token": {"foo": "bar"}}`,
+		},
+		{
+			desc:           "insufficient claims",
+			challenge:      `Bearer claims="eyJhY2Nlc3NfdG9rZW4iOiB7ImZvbyI6ICJiYXIifX0=", realm="", authorization_uri="https://login.microsoftonline.com/common/oauth2/authorize", client_id="00000003-0000-0000-c000-000000000000", error="insufficient_claims"`,
+			expectedClaims: `{"access_token": {"foo": "bar"}}`,
+		},
+		{
+			desc:           "sessions revoked",
+			challenge:      `Bearer authorization_uri="https://login.windows.net/", error="invalid_token", error_description="User session has been revoked", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0="`,
+			expectedClaims: `{"access_token":{"nbf":{"essential":true, "value":"1603742800"}}}`,
+		},
+		{
+			desc:           "sessions revoked",
+			challenge:      `Bearer authorization_uri="https://login.windows.net/", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwgInZhbHVlIjoiMTYwMzc0MjgwMCJ9fX0=", error="invalid_token", error_description="User session has been revoked"`,
+			expectedClaims: `{"access_token":{"nbf":{"essential":true, "value":"1603742800"}}}`,
+		},
+		{
+			desc:           "IP policy",
+			challenge:      `Bearer authorization_uri="https://login.windows.net/", error="invalid_token", error_description="Tenant IP Policy validate failed.", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwidmFsdWUiOiIxNjEwNTYzMDA2In0sInhtc19ycF9pcGFkZHIiOnsidmFsdWUiOiIxLjIuMy40In19fQ"`,
+			expectedClaims: `{"access_token":{"nbf":{"essential":true,"value":"1610563006"},"xms_rp_ipaddr":{"value":"1.2.3.4"}}}`,
+		},
+		{
+			desc:           "IP policy",
+			challenge:      `Bearer authorization_uri="https://login.windows.net/", error="invalid_token", claims="eyJhY2Nlc3NfdG9rZW4iOnsibmJmIjp7ImVzc2VudGlhbCI6dHJ1ZSwidmFsdWUiOiIxNjEwNTYzMDA2In0sInhtc19ycF9pcGFkZHIiOnsidmFsdWUiOiIxLjIuMy40In19fQ", error_description="Tenant IP Policy validate failed."`,
+			expectedClaims: `{"access_token":{"nbf":{"essential":true,"value":"1610563006"},"xms_rp_ipaddr":{"value":"1.2.3.4"}}}`,
+		},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			srv, close := mock.NewServer()
+			defer close()
+			srv.SetResponse(mock.WithHeader(shared.HeaderWWWAuthenticate, test.challenge), mock.WithStatusCode(http.StatusUnauthorized))
+			calls := 0
+			cred := mockCredential{
+				getTokenImpl: func(ctx context.Context, actual azpolicy.TokenRequestOptions) (azcore.AccessToken, error) {
+					calls += 1
+					if calls == 2 && test.expectedClaims != "" {
+						require.Equal(t, test.expectedClaims, actual.Claims)
+					}
+					return azcore.AccessToken{Token: "...", ExpiresOn: time.Now().Add(time.Hour).UTC()}, nil
+				},
+			}
+			b := NewBearerTokenPolicy(cred, &armpolicy.BearerTokenOptions{Scopes: []string{scope}})
+			pipeline := newTestPipeline(&azpolicy.ClientOptions{Transport: srv, PerRetryPolicies: []azpolicy.Policy{b}})
+			req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+			require.NoError(t, err)
+			_, err = pipeline.Do(req)
+			if test.err != nil {
+				require.ErrorAs(t, err, &test.err)
+			} else {
+				require.NoError(t, err)
+			}
+			if test.expectedClaims != "" {
+				require.Equal(t, 2, calls, "policy should have requested a new token upon receiving the challenge")
+			}
+		})
 	}
 }
diff --git a/sdk/azcore/runtime/policy_bearer_token_test.go b/sdk/azcore/runtime/policy_bearer_token_test.go
index 01ec151dd1..b9ba0175ca 100644
--- a/sdk/azcore/runtime/policy_bearer_token_test.go
+++ b/sdk/azcore/runtime/policy_bearer_token_test.go
@@ -149,7 +149,7 @@ func TestBearerTokenPolicy_AuthZHandler(t *testing.T) {
 	challenge := "Scheme parameters..."
 	srv, close := mock.NewTLSServer(mock.WithTransformAllRequestsToTestServerUrl())
 	defer close()
-	srv.AppendResponse(mock.WithStatusCode(401), mock.WithHeader("WWW-Authenticate", challenge))
+	srv.AppendResponse(mock.WithStatusCode(401), mock.WithHeader(shared.HeaderWWWAuthenticate, challenge))
 	srv.AppendResponse(mock.WithStatusCode(200))
 
 	req, err := NewRequest(context.Background(), "GET", "https://localhost")
@@ -167,7 +167,7 @@ func TestBearerTokenPolicy_AuthZHandler(t *testing.T) {
 	handler.OnChallenge = func(r *policy.Request, res *http.Response, f func(policy.TokenRequestOptions) error) error {
 		require.Equal(t, req.Raw().URL, r.Raw().URL)
 		handler.onChallengeCalls++
-		require.Equal(t, challenge, res.Header.Get("WWW-Authenticate"))
+		require.Equal(t, challenge, res.Header.Get(shared.HeaderWWWAuthenticate))
 		return nil
 	}
 
@@ -185,7 +185,7 @@ func TestBearerTokenPolicy_AuthZHandler(t *testing.T) {
 func TestBearerTokenPolicy_AuthZHandlerErrors(t *testing.T) {
 	srv, close := mock.NewTLSServer(mock.WithTransformAllRequestsToTestServerUrl())
 	defer close()
-	srv.SetResponse(mock.WithStatusCode(401), mock.WithHeader("WWW-Authenticate", "..."))
+	srv.SetResponse(mock.WithStatusCode(401), mock.WithHeader(shared.HeaderWWWAuthenticate, "..."))
 
 	req, err := NewRequest(context.Background(), "GET", "https://localhost")
 	require.NoError(t, err)