Repository: Azure/azure-sdk-for-go
Pull request #19902 (2023-01-31 19:44:10)

CAE support for ARM clients
This is the azcore side of #14931. I'm targeting the beta branch because the end-to-end scenario requires an azidentity beta. ARM is the only SDK service that supports CAE (continuous access evaluation) today. Briefly, it works like this: ARM evaluates security policies for the access token attached to each request. When a token doesn't satisfy policy, ARM responds 401 with an authentication challenge. The client must retry the request with a new access token having the claims specified by the challenge. So, this PR adds challenge handling to ARM's auth policy, and a `TokenRequestOptions.Claims` field that enables the policy to request claims when calling `GetToken()`.


Patch:
diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md
index bffd9c1491..eacb278dfd 100644
--- a/sdk/azcore/CHANGELOG.md
+++ b/sdk/azcore/CHANGELOG.md
@@ -1,16 +1,16 @@
 # Release History
 
-## 1.3.1 (Unreleased)
+## 1.4.0-beta.1 (2023-02-02)
 
 ### Features Added
 * Added support for ARM cross-tenant authentication. Set the `AuxiliaryTenants` field of `arm.ClientOptions` to enable.
-* Added `TenantID` field to `policy.TokenRequestOptions`.
+* Added `Claims` and `TenantID` fields to `policy.TokenRequestOptions`.
+* ARM bearer token policy handles CAE challenges.
 
-### Breaking Changes
-
-### Bugs Fixed
+## 1.3.1 (2023-02-02)
 
 ### Other Changes
+* Update dependencies to latest versions.
 
 ## 1.3.0 (2023-01-06)
 
diff --git a/sdk/azcore/arm/runtime/policy_bearer_token.go b/sdk/azcore/arm/runtime/policy_bearer_token.go
index 4415085ee5..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"
@@ -64,12 +65,27 @@ func NewBearerTokenPolicy(cred azcore.TokenCredential, opts *armpolicy.BearerTok
 	copy(p.scopes, opts.Scopes)
 	p.btp = azruntime.NewBearerTokenPolicy(cred, opts.Scopes, &azpolicy.BearerTokenOptions{
 		AuthorizationHandler: azpolicy.AuthorizationHandler{
-			OnRequest: p.onRequest,
+			OnChallenge: p.onChallenge,
+			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, 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
@@ -99,3 +115,31 @@ 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
+}
diff --git a/sdk/azcore/go.mod b/sdk/azcore/go.mod
index 88738481fb..61645d496c 100644
--- a/sdk/azcore/go.mod
+++ b/sdk/azcore/go.mod
@@ -3,14 +3,14 @@ module github.com/Azure/azure-sdk-for-go/sdk/azcore
 go 1.18
 
 require (
-	github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.1
+	github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2
 	github.com/stretchr/testify v1.7.0
-	golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4
+	golang.org/x/net v0.5.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	golang.org/x/text v0.3.7 // indirect
+	golang.org/x/text v0.6.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
 )
diff --git a/sdk/azcore/go.sum b/sdk/azcore/go.sum
index 43fe8dcae6..e6e3bd4fa2 100644
--- a/sdk/azcore/go.sum
+++ b/sdk/azcore/go.sum
@@ -1,5 +1,5 @@
-github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.1 h1:Oj853U9kG+RLTCQXpjvOnrv0WaZHxgmZz1TlLywgOPY=
-github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.1/go.mod h1:eWRD7oawr1Mu1sLCawqVc0CUiF43ia3qQMxLscsKQ9w=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2 h1:+5VZ72z0Qan5Bog5C+ZkgSqUbeVUd9wgtHOrIKuc5b8=
+github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2/go.mod h1:eWRD7oawr1Mu1sLCawqVc0CUiF43ia3qQMxLscsKQ9w=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
@@ -7,10 +7,10 @@ github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZN
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4 h1:HVyaeDAYux4pnY+D/SiwmLOR36ewZ4iGQIIrtnuCjFA=
-golang.org/x/net v0.0.0-20220425223048-2871e0cb64e4/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
-golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
-golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/net v0.5.0 h1:GyT4nK/YDHSqa1c4753ouYCDajOYKTja9Xb/OHtgvSw=
+golang.org/x/net v0.5.0/go.mod h1:DivGGAXEgPSlEBzxGzZI+ZLohi+xUj054jfeKui00ws=
+golang.org/x/text v0.6.0 h1:3XmdazWV+ubf7QgHSTWeykHOci5oeekaGJBLkrkaw4k=
+golang.org/x/text v0.6.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
diff --git a/sdk/azcore/internal/exported/exported.go b/sdk/azcore/internal/exported/exported.go
index 843da05b85..2ef55e3677 100644
--- a/sdk/azcore/internal/exported/exported.go
+++ b/sdk/azcore/internal/exported/exported.go
@@ -71,6 +71,10 @@ 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
 
diff --git a/sdk/azcore/internal/shared/constants.go b/sdk/azcore/internal/shared/constants.go
index 0a67d85899..b98de7939b 100644
--- a/sdk/azcore/internal/shared/constants.go
+++ b/sdk/azcore/internal/shared/constants.go
@@ -12,14 +12,16 @@ const (
 )
 
 const (
-	HeaderAuthorization     = "Authorization"
-	HeaderAzureAsync        = "Azure-AsyncOperation"
-	HeaderContentLength     = "Content-Length"
-	HeaderContentType       = "Content-Type"
-	HeaderLocation          = "Location"
-	HeaderOperationLocation = "Operation-Location"
-	HeaderRetryAfter        = "Retry-After"
-	HeaderUserAgent         = "User-Agent"
+	HeaderAuthorization          = "Authorization"
+	HeaderAuxiliaryAuthorization = "x-ms-authorization-auxiliary"
+	HeaderAzureAsync             = "Azure-AsyncOperation"
+	HeaderContentLength          = "Content-Length"
+	HeaderContentType            = "Content-Type"
+	HeaderLocation               = "Location"
+	HeaderOperationLocation      = "Operation-Location"
+	HeaderRetryAfter             = "Retry-After"
+	HeaderUserAgent              = "User-Agent"
+	HeaderWWWAuthenticate        = "WWW-Authenticate"
 )
 
 const BearerTokenPrefix = "Bearer "
@@ -29,5 +31,5 @@ const (
 	Module = "azcore"
 
 	// Version is the semantic version (see http://semver.org) of this module.
-	Version = "v1.3.1"
+	Version = "v1.4.0-beta.1"
 )


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 738425a535..1ab06ae00d 100644
--- a/sdk/azcore/arm/runtime/policy_bearer_token_test.go
+++ b/sdk/azcore/arm/runtime/policy_bearer_token_test.go
@@ -201,3 +201,85 @@ func TestAuxiliaryTenants(t *testing.T) {
 		require.ElementsMatch(t, expected, actual)
 	}
 }
+
+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)