Repository: Azure/azure-sdk-for-go
Pull request #17940 (2022-05-13 14:59:09)

Cosmos DB - Refactor AAD policy
Related to https://github.com/Azure/azure-sdk-for-go/issues/17752

Refactor the internal AAD policy using the new types exposed in https://github.com/Azure/azure-sdk-for-go/pull/17798

Instead of rewriting the Auth header from the azcore AAD policy (which was a temporal workaround), we can now correctly create the header ourselves with the desired format.


Patch:
diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go
index 200097565b..e5b7e496b8 100644
--- a/sdk/data/azcosmos/cosmos_client.go
+++ b/sdk/data/azcosmos/cosmos_client.go
@@ -34,7 +34,7 @@ func (c *Client) Endpoint() string {
 // cred - The credential used to authenticate with the cosmos service.
 // options - Optional Cosmos client options.  Pass nil to accept default values.
 func NewClientWithKey(endpoint string, cred KeyCredential, o *ClientOptions) (*Client, error) {
-	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{newSharedKeyCredPolicy(cred)}, o)}, nil
+	return &Client{endpoint: endpoint, pipeline: newPipeline(newSharedKeyCredPolicy(cred), o)}, nil
 }
 
 // NewClient creates a new instance of Cosmos client with Azure AD access token authentication. It uses the default pipeline configuration.
@@ -46,10 +46,10 @@ func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (
 	if err != nil {
 		return nil, err
 	}
-	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{azruntime.NewBearerTokenPolicy(cred, scope, nil), &cosmosBearerTokenPolicy{}}, o)}, nil
+	return &Client{endpoint: endpoint, pipeline: newPipeline(newCosmosBearerTokenPolicy(cred, scope, nil), o)}, nil
 }
 
-func newPipeline(authPolicy []policy.Policy, options *ClientOptions) azruntime.Pipeline {
+func newPipeline(authPolicy policy.Policy, options *ClientOptions) azruntime.Pipeline {
 	if options == nil {
 		options = &ClientOptions{}
 	}
@@ -61,7 +61,9 @@ func newPipeline(authPolicy []policy.Policy, options *ClientOptions) azruntime.P
 					enableContentResponseOnWrite: options.EnableContentResponseOnWrite,
 				},
 			},
-			PerRetry: authPolicy,
+			PerRetry: []policy.Policy{
+				authPolicy,
+			},
 		},
 		&options.ClientOptions)
 }
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token.go b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
index dfba0e9b3c..da71fa52a2 100644
--- a/sdk/data/azcosmos/cosmos_policy_bearer_token.go
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
@@ -4,25 +4,62 @@
 package azcosmos
 
 import (
-	"errors"
 	"fmt"
 	"net/http"
+	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+	"github.com/Azure/azure-sdk-for-go/sdk/internal/temporal"
 )
 
-const lenBearerTokenPrefix = len("Bearer ")
-
+// cosmosBearerTokenPolicy authorizes requests with bearer tokens acquired from a TokenCredential.
+// Copy of sdk/azcore/runtime/policy_bearer_token.go using Cosmos header format
 type cosmosBearerTokenPolicy struct {
+	// mainResource is the resource to be retrieved using the tenant specified in the credential
+	mainResource *temporal.Resource[azcore.AccessToken, acquiringResourceState]
+	// the following fields are read-only
+	cred   azcore.TokenCredential
+	scopes []string
 }
 
-func (b *cosmosBearerTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
-	currentAuthorization := req.Raw().Header.Get(headerAuthorization)
-	if currentAuthorization == "" {
-		return nil, errors.New("authorization header is missing")
+type acquiringResourceState struct {
+	req *policy.Request
+	p   *cosmosBearerTokenPolicy
+}
+
+// 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.req.Raw().Context(), policy.TokenRequestOptions{Scopes: state.p.scopes})
+	if err != nil {
+		return azcore.AccessToken{}, time.Time{}, err
 	}
+	return tk, tk.ExpiresOn, nil
+}
 
-	token := currentAuthorization[lenBearerTokenPrefix:]
-	req.Raw().Header.Set(headerAuthorization, fmt.Sprintf("type=aad&ver=1.0&sig=%v", token))
+// NewBearerTokenPolicy creates a policy object that authorizes requests with bearer tokens.
+// cred: an azcore.TokenCredential implementation such as a credential object from azidentity
+// scopes: the list of permission scopes required for the token.
+// opts: optional settings. Pass nil to accept default values; this is the same as passing a zero-value options.
+func newCosmosBearerTokenPolicy(cred azcore.TokenCredential, scopes []string, opts *policy.BearerTokenOptions) *cosmosBearerTokenPolicy {
+	return &cosmosBearerTokenPolicy{
+		cred:         cred,
+		scopes:       scopes,
+		mainResource: temporal.NewResource(acquire),
+	}
+}
+
+// Do authorizes a request with a bearer token
+func (b *cosmosBearerTokenPolicy) Do(req *policy.Request) (*http.Response, error) {
+	as := acquiringResourceState{
+		p:   b,
+		req: req,
+	}
+	tk, err := b.mainResource.Get(as)
+	if err != nil {
+		return nil, err
+	}
+	req.Raw().Header.Set(headerAuthorization, fmt.Sprintf("type=aad&ver=1.0&sig=%v", tk.Token))
 	return req.Next()
 }


Test patch:
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
index e5d3eecade..fb986ccc1d 100644
--- a/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
@@ -5,54 +5,148 @@ package azcosmos
 
 import (
 	"context"
+	"fmt"
+
+	"errors"
 	"net/http"
 	"testing"
+	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"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/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/mock"
 )
 
-func TestConvertBearerToken(t *testing.T) {
-	srv, close := mock.NewTLSServer()
-	defer close()
-	srv.SetResponse(mock.WithStatusCode(http.StatusOK))
+const (
+	tokenValue                = "***"
+	accessTokenRespSuccess    = `{"access_token": "` + tokenValue + `", "expires_in": 3600}`
+	accessTokenRespShortLived = `{"access_token": "` + tokenValue + `", "expires_in": 0}`
+	scope                     = "scope"
+)
 
-	verifier := bearerTokenVerify{}
-	pl := azruntime.NewPipeline("azcosmostest", "v1.0.0", azruntime.PipelineOptions{PerCall: []policy.Policy{&mockAuthPolicy{}, &cosmosBearerTokenPolicy{}, &verifier}}, &policy.ClientOptions{Transport: srv})
-	req, err := azruntime.NewRequest(context.Background(), http.MethodGet, srv.URL())
-	req.SetOperationValue(pipelineRequestOptions{
-		isWriteOperation: true,
-	})
+type mockCredential struct {
+	getTokenImpl func(ctx context.Context, options policy.TokenRequestOptions) (azcore.AccessToken, error)
+}
 
-	if err != nil {
-		t.Fatalf("unexpected error: %v", err)
+func (mc mockCredential) GetToken(ctx context.Context, options policy.TokenRequestOptions) (azcore.AccessToken, error) {
+	if mc.getTokenImpl != nil {
+		return mc.getTokenImpl(ctx, options)
 	}
+	return azcore.AccessToken{Token: "***", ExpiresOn: time.Now().Add(time.Hour)}, nil
+}
 
-	_, err = pl.Do(req)
-	if err != nil {
-		t.Fatalf("unexpected error: %v", err)
-	}
+func (mc mockCredential) Do(req *policy.Request) (*http.Response, error) {
+	return nil, nil
+}
 
-	if verifier.authHeaderContent != "type=aad&ver=1.0&sig=this is a test token" {
-		t.Fatalf("Expected auth header content to be 'type=aad&ver=1.0&sig=this is a test token', got %s", verifier.authHeaderContent)
+func defaultTestPipeline(srv policy.Transporter, scope string) runtime.Pipeline {
+	retryOpts := policy.RetryOptions{
+		MaxRetryDelay: 500 * time.Millisecond,
+		RetryDelay:    time.Millisecond,
 	}
+	b := newCosmosBearerTokenPolicy(mockCredential{}, []string{scope}, nil)
+	return runtime.NewPipeline(
+		"azcosmostest",
+		"v1.0.0",
+		runtime.PipelineOptions{PerRetry: []policy.Policy{b}},
+		&policy.ClientOptions{Retry: retryOpts, Transport: srv},
+	)
 }
 
-type bearerTokenVerify struct {
-	authHeaderContent string
+func TestBearerPolicy_SuccessGetToken(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK))
+	pipeline := defaultTestPipeline(srv, scope)
+	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp, err := pipeline.Do(req)
+	if err != nil {
+		t.Fatalf("Expected nil error but received one")
+	}
+	expectedToken := fmt.Sprintf("type=aad&ver=1.0&sig=%v", tokenValue)
+	if token := resp.Request.Header.Get(headerAuthorization); token != expectedToken {
+		t.Fatalf("expected token '%s', got '%s'", expectedToken, token)
+	}
 }
 
-func (p *bearerTokenVerify) Do(req *policy.Request) (*http.Response, error) {
-	p.authHeaderContent = req.Raw().Header.Get(headerAuthorization)
-
-	return req.Next()
+func TestBearerPolicy_CredentialFailGetToken(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	expectedErr := errors.New("oops")
+	failCredential := mockCredential{}
+	failCredential.getTokenImpl = func(ctx context.Context, options policy.TokenRequestOptions) (azcore.AccessToken, error) {
+		return azcore.AccessToken{}, expectedErr
+	}
+	b := newCosmosBearerTokenPolicy(failCredential, nil, nil)
+	pipeline := runtime.NewPipeline("azcosmostest", "v1.0.0", runtime.PipelineOptions{}, &policy.ClientOptions{
+		Transport: srv,
+		Retry: policy.RetryOptions{
+			RetryDelay: 10 * time.Millisecond,
+		},
+		PerRetryPolicies: []policy.Policy{b}})
+	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp, err := pipeline.Do(req)
+	if err != expectedErr {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if resp != nil {
+		t.Fatal("expected nil response")
+	}
 }
 
-type mockAuthPolicy struct{}
-
-func (p *mockAuthPolicy) Do(req *policy.Request) (*http.Response, error) {
-	req.Raw().Header.Set(headerAuthorization, "Bearer this is a test token")
+func TestBearerTokenPolicy_TokenExpired(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespShortLived)))
+	srv.AppendResponse(mock.WithStatusCode(http.StatusOK))
+	pipeline := defaultTestPipeline(srv, scope)
+	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, err = pipeline.Do(req)
+	if err != nil {
+		t.Fatalf("unexpected error %v", err)
+	}
+	_, err = pipeline.Do(req)
+	if err != nil {
+		t.Fatalf("unexpected error %v", err)
+	}
+}
 
-	return req.Next()
+func TestBearerPolicy_GetTokenFailsNoDeadlock(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.AppendResponse(mock.WithBody([]byte(accessTokenRespSuccess)))
+	retryOpts := policy.RetryOptions{
+		// use a negative try timeout to trigger a deadline exceeded error causing GetToken() to fail
+		TryTimeout:    -1 * time.Nanosecond,
+		MaxRetryDelay: 500 * time.Millisecond,
+		RetryDelay:    50 * time.Millisecond,
+		MaxRetries:    3,
+	}
+	b := newCosmosBearerTokenPolicy(mockCredential{}, nil, nil)
+	pipeline := runtime.NewPipeline("azcosmostest", "v1.0.0", runtime.PipelineOptions{}, &policy.ClientOptions{
+		Transport:        srv,
+		Retry:            retryOpts,
+		PerRetryPolicies: []policy.Policy{b}})
+	req, err := runtime.NewRequest(context.Background(), http.MethodGet, srv.URL())
+	if err != nil {
+		t.Fatal(err)
+	}
+	resp, err := pipeline.Do(req)
+	if err == nil {
+		t.Fatal("unexpected nil error")
+	}
+	if resp != nil {
+		t.Fatal("expected nil response")
+	}
 }