BearerTokenPolicy handles CAE claims challenges by default

sdk/azcore/CHANGELOG.md

@@ -1,10 +1,10 @@
 # Release History
 
-## 1.14.1 (Unreleased)
+## 1.15.0 (2024-10-07)
 
 ### Features Added
 
-### Breaking Changes
+* `BearerTokenPolicy` handles CAE claims challenges
 
 ### Bugs Fixed
 

sdk/azcore/arm/runtime/policy_bearer_token.go

@@ -5,7 +5,6 @@ package runtime
 
 import (
 	"context"
-	"encoding/base64"
 	"fmt"
 	"net/http"
 	"strings"
@@ -66,31 +65,16 @@ func NewBearerTokenPolicy(cred azcore.TokenCredential, opts *armpolicy.BearerTok
 	p.btp = azruntime.NewBearerTokenPolicy(cred, opts.Scopes, &azpolicy.BearerTokenOptions{
 		InsecureAllowCredentialWithHTTP: opts.InsecureAllowCredentialWithHTTP,
 		AuthorizationHandler: azpolicy.AuthorizationHandler{
-			OnChallenge: p.onChallenge,
-			OnRequest:   p.onRequest,
+			OnRequest: p.onRequest,
 		},
 	})
 	return p
 }
 
-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, EnableCAE: true, 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{EnableCAE: true, Scopes: b.scopes})
+	err := authNZ(azpolicy.TokenRequestOptions{Scopes: b.scopes})
 	if err != nil || len(b.auxResources) == 0 {
 		return err
 	}
@@ -116,31 +100,3 @@ func (b *BearerTokenPolicy) onRequest(req *azpolicy.Request, authNZ func(azpolic
 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 claims, err
-}

sdk/azcore/arm/runtime/policy_bearer_token_test.go

@@ -205,88 +205,6 @@ func TestAuxiliaryTenants(t *testing.T) {
 	}
 }
 
-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.NewTLSServer()
-			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")
-			}
-		})
-	}
-}
-
 func TestBearerTokenPolicyRequiresHTTPS(t *testing.T) {
 	srv, close := mock.NewServer()
 	defer close()

sdk/azcore/internal/shared/constants.go

@@ -40,5 +40,5 @@ const (
 	Module = "azcore"
 
 	// Version is the semantic version (see http://semver.org) of this module.
-	Version = "v1.14.1"
+	Version = "v1.15.0"
 )

sdk/azcore/policy/policy.go

@@ -168,12 +168,12 @@ type AuthorizationHandler struct {
 	// token from its credential according to its configuration.
 	OnRequest func(*Request, func(TokenRequestOptions) error) error
 
-	// OnChallenge is called when the policy receives a 401 response, allowing the AuthorizationHandler to re-authorize the
-	// request according to an authentication challenge (the Response's WWW-Authenticate header). OnChallenge is responsible
-	// for parsing parameters from the challenge. Its func parameter will authorize the request with a token from the policy's
-	// given credential. Implementations that need to perform I/O should use the Request's context, available from
-	// Request.Raw().Context(). When OnChallenge returns nil, the policy will send the request again. When OnChallenge is nil,
-	// the policy will return any 401 response to the client.
+	// OnChallenge allows clients to implement custom HTTP authentication challenge handling. BearerTokenPolicy calls it upon
+	// receiving a 401 response containing multiple Bearer challenges or a challenge BearerTokenPolicy itself can't handle.
+	// OnChallenge is responsible for parsing challenge(s) (the Response's WWW-Authenticate header) and reauthorizing the
+	// Request accordingly. Its func argument authorizes the Request with a token from the policy's credential using the given
+	// TokenRequestOptions. OnChallenge should honor the Request's context, available from Request.Raw().Context(). When
+	// OnChallenge returns nil, the policy will send the Request again.
 	OnChallenge func(*Request, *http.Response, func(TokenRequestOptions) error) error
 }
 

sdk/azcore/runtime/policy_bearer_token.go

@@ -4,9 +4,12 @@
 package runtime
 
 import (
+	"encoding/base64"
 	"errors"
 	"net/http"
+	"regexp"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/internal/exported"
@@ -17,6 +20,11 @@ import (
 )
 
 // BearerTokenPolicy authorizes requests with bearer tokens acquired from a TokenCredential.
+// It handles [Continuous Access Evaluation] (CAE) challenges. Clients needing to handle
+// additional authentication challenges, or needing more control over authorization, should
+// provide a [policy.AuthorizationHandler] in [policy.BearerTokenOptions].
+//
+// [Continuous Access Evaluation]: https://learn.microsoft.com/entra/identity/conditional-access/concept-continuous-access-evaluation
 type BearerTokenPolicy struct {
 	// mainResource is the resource to be retreived using the tenant specified in the credential
 	mainResource *temporal.Resource[exported.AccessToken, acquiringResourceState]
@@ -63,6 +71,7 @@ func NewBearerTokenPolicy(cred exported.TokenCredential, scopes []string, opts *
 // authenticateAndAuthorize returns a function which authorizes req with a token from the policy's credential
 func (b *BearerTokenPolicy) authenticateAndAuthorize(req *policy.Request) func(policy.TokenRequestOptions) error {
 	return func(tro policy.TokenRequestOptions) error {
+		tro.EnableCAE = true
 		as := acquiringResourceState{p: b, req: req, tro: tro}
 		tk, err := b.mainResource.Get(as)
 		if err != nil {
@@ -101,17 +110,46 @@ func (b *BearerTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
 		return nil, err
 	}
 
+	res, err = b.handleChallenge(req, res, false)
+	return res, err
+}
+
+// handleChallenge handles authentication challenges either directly (for CAE challenges) or by calling
+// the AuthorizationHandler. It's a no-op when the response doesn't include an authentication challenge.
+// It will recurse at most once, to handle a CAE challenge following a non-CAE challenge handled by the
+// AuthorizationHandler.
+func (b *BearerTokenPolicy) handleChallenge(req *policy.Request, res *http.Response, recursed bool) (*http.Response, error) {
+	var err error
 	if res.StatusCode == http.StatusUnauthorized {
 		b.mainResource.Expire()
-		if res.Header.Get("WWW-Authenticate") != "" && b.authzHandler.OnChallenge != nil {
-			if err = b.authzHandler.OnChallenge(req, res, b.authenticateAndAuthorize(req)); err == nil {
-				res, err = req.Next()
+		if res.Header.Get(shared.HeaderWWWAuthenticate) != "" {
+			caeChallenge, parseErr := parseCAEChallenge(res)
+			if parseErr != nil {
+				return res, parseErr
+			}
+			switch {
+			case caeChallenge != nil:
+				tro := policy.TokenRequestOptions{
+					Claims: caeChallenge.params["claims"],
+					Scopes: b.scopes,
+				}
+				if err = b.authenticateAndAuthorize(req)(tro); err == nil {
+					res, err = req.Next()
+				}
+			case b.authzHandler.OnChallenge != nil && !recursed:
+				if err = b.authzHandler.OnChallenge(req, res, b.authenticateAndAuthorize(req)); err == nil {
+					if res, err = req.Next(); err == nil {
+						res, err = b.handleChallenge(req, res, true)
+					}
+				} else {
+					// don't retry challenge handling errors
+					err = errorinfo.NonRetriableError(err)
+				}
+			default:
+				// return the response to the pipeline
 			}
 		}
 	}
-	if err != nil {
-		err = errorinfo.NonRetriableError(err)
-	}
 	return res, err
 }
 
@@ -121,3 +159,65 @@ func checkHTTPSForAuth(req *policy.Request, allowHTTP bool) error {
 	}
 	return nil
 }
+
+// parseCAEChallenge returns a *authChallenge representing Response's CAE challenge (nil when Response has none).
+// If Response includes a CAE challenge having invalid claims, it returns a NonRetriableError.
+func parseCAEChallenge(res *http.Response) (*authChallenge, error) {
+	var (
+		caeChallenge *authChallenge
+		err          error
+	)
+	for _, c := range parseChallenges(res) {
+		if c.scheme == "Bearer" {
+			if claims := c.params["claims"]; claims != "" && c.params["error"] == "insufficient_claims" {
+				if b, de := base64.StdEncoding.DecodeString(claims); de == nil {
+					c.params["claims"] = string(b)
+					caeChallenge = &c
+				} else {
+					// don't include the decoding error because it's something
+					// unhelpful like "illegal base64 data at input byte 42"
+					err = errorinfo.NonRetriableError(errors.New("authentication challenge contains invalid claims: " + claims))
+				}
+				break
+			}
+		}
+	}
+	return caeChallenge, err
+}
+
+var (
+	challenge, challengeParams *regexp.Regexp
+	once                       = &sync.Once{}
+)
+
+type authChallenge struct {
+	scheme string
+	params map[string]string
+}
+
+// parseChallenges assumes authentication challenges have quoted parameter values
+func parseChallenges(res *http.Response) []authChallenge {
+	once.Do(func() {
+		// matches challenges having quoted parameters, capturing scheme and parameters
+		challenge = regexp.MustCompile(`(?:(\w+) ((?:\w+="[^"]*",?\s*)+))`)
+		// captures parameter names and values in a match of the above expression
+		challengeParams = regexp.MustCompile(`(\w+)="([^"]*)"`)
+	})
+	parsed := []authChallenge{}
+	// WWW-Authenticate can have multiple values, each containing multiple challenges
+	for _, h := range res.Header.Values(shared.HeaderWWWAuthenticate) {
+		for _, sm := range challenge.FindAllStringSubmatch(h, -1) {
+			// sm is [challenge, scheme, params] (see regexp documentation on submatches)
+			c := authChallenge{
+				params: make(map[string]string),
+				scheme: sm[1],
+			}
+			for _, sm := range challengeParams.FindAllStringSubmatch(sm[2], -1) {
+				// sm is [key="value", key, value] (see regexp documentation on submatches)
+				c.params[sm[1]] = sm[2]
+			}
+			parsed = append(parsed, c)
+		}
+	}
+	return parsed
+}