Repository: Azure/azure-sdk-for-go (go)

Issue:
Cosmos needs a constructor taking TokenCrendentail
Add in a client constructor for AAD support


Pull request:
Cosmos DB: Add AAD authentication
This PR adds support for `azcore.TokenCredential` as authentication method to the Cosmos DB SDK.

## How to create a client with AAD

Using the default credential:
```
cred, err := azidentity.NewDefaultAzureCredential(nil)
if err != nil {
	panic(err)
}

client, err := azcosmos.NewClient(endpoint, cred, nil)
if err != nil {
	panic(err)
}
```

Using a client secret:
```
cred, err := azidentity.NewClientSecretCredential("tenantId", "clientId", "clientSecret")
if err != nil {
	panic(err)
}

client, err := azcosmos.NewClient(endpoint, cred, nil)
if err != nil {
	panic(err)
}
```

## API changes

There is only a single API change, which adds the new client constructor:

```
func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (*Client, error)
```

APIView https://apiview.dev/Assemblies/Review/ca456a2f7ed64556b3fb57ecebc0b9d2/cc654e57f8044ed8a14c5dbf421ab49b?diffRevisionId=54a33d355ba146cb8e595448aa039540&doc=False&diffOnly=True

(There are other bogus differences due to reordering of existing properties)

Closes #16219

Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index cc7d47a5f4..655b5666b8 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -14,43 +14,43 @@
 
 # Catch all for non-code project files and unowned files | folders
 *                             @rickwinter @jhendrixMSFT
-/sdk/                         @rickwinter @chlowell @richardpark-msft @seankane-msft
+/sdk/                         @rickwinter @chlowell @richardpark-msft
 
 # Samples
-/samples/                     @rickwinter @chlowell @jhendrixMSFT @richardpark-msft @seankane-msft
+/samples/                     @rickwinter @chlowell @jhendrixMSFT @richardpark-msft
 
 # PRLabel: %Azure.Core
-/sdk/azcore/                  @rickwinter @chlowell @jhendrixMSFT @richardpark-msft @seankane-msft
+/sdk/azcore/                  @rickwinter @chlowell @jhendrixMSFT @richardpark-msft
 
 # PRLabel: %Azure.Identity
 /sdk/azidentity/              @chlowell @jhendrixMSFT @rickwinter
 
 # PRLabel: %Internal
-/sdk/internal/                @rickwinter @chlowell @jhendrixMSFT @richardpark-msft @seankane-msft
+/sdk/internal/                @rickwinter @chlowell @jhendrixMSFT @richardpark-msft
 
 # PRLabel: %App Configuration
-/sdk/appconfig/               @antkmsft @seankane-msft @jhendrixMSFT
+/sdk/appconfig/               @antkmsft @jhendrixMSFT
 
 # PRLabel: %Cosmos
 /sdk/data/azcosmos/           @ealsur @kirankumarkolli
 
 # PRLabel: %Tables
-/sdk/data/aztables/           @seankane-msft @jhendrixMSFT 
+/sdk/data/aztables/           @jhendrixMSFT
 
 # PRLabel: %KeyVault
-/sdk/keyvault/                @seankane-msft @jhendrixMSFT
+/sdk/keyvault/                @chlowell @jhendrixMSFT
 
 # PRLabel: %Service Bus
-/sdk/messaging/azservicebus   @richardpark-msft @jhendrixMSFT 
+/sdk/messaging/azservicebus   @richardpark-msft @jhendrixMSFT
 
 # PRLabel: %Service Bus
-/sdk/messaging/internal       @richardpark-msft @jhendrixMSFT 
+/sdk/messaging/internal       @richardpark-msft @jhendrixMSFT
 
 # PRLabel: %Mgmt
 /sdk/resourcemanager/         @ArcturusZhang @lirenhe @tadelesh
 
 # PRLabel: %Storage
-/sdk/storage/                 @jhendrixMSFT @seankane-msft
+/sdk/storage/                 @mohsha-msft @zezha-msft @jhendrixMSFT
 
 ###########
 # Eng Sys
diff --git a/sdk/azcore/CHANGELOG.md b/sdk/azcore/CHANGELOG.md
index 0e87951b90..7c04a18624 100644
--- a/sdk/azcore/CHANGELOG.md
+++ b/sdk/azcore/CHANGELOG.md
@@ -6,6 +6,7 @@
 
 ### Breaking Changes
 * Renamed `cloud.Configuration.LoginEndpoint` to `.ActiveDirectoryAuthorityHost`
+* Renamed `cloud.AzurePublicCloud` to `cloud.AzurePublic`
 * Removed `AuxiliaryTenants` field from `arm/ClientOptions` and `arm/policy/BearerTokenOptions`
 * Removed `TokenRequestOptions.TenantID`
 * `Poller[T].PollUntilDone()` now takes an `options *PollUntilDoneOptions` param instead of `freq time.Duration`
diff --git a/sdk/azcore/arm/arm.go b/sdk/azcore/arm/arm.go
index 055353650c..3a0cb64429 100644
--- a/sdk/azcore/arm/arm.go
+++ b/sdk/azcore/arm/arm.go
@@ -17,7 +17,7 @@ func init() {
 		Audience: "https://management.core.usgovcloudapi.net",
 		Endpoint: "https://management.usgovcloudapi.net",
 	}
-	cloud.AzurePublicCloud.Services[cloud.ResourceManager] = cloud.ServiceConfiguration{
+	cloud.AzurePublic.Services[cloud.ResourceManager] = cloud.ServiceConfiguration{
 		Audience: "https://management.core.windows.net/",
 		Endpoint: "https://management.azure.com",
 	}
diff --git a/sdk/azcore/arm/runtime/pipeline.go b/sdk/azcore/arm/runtime/pipeline.go
index a68054705c..6b9a242cea 100644
--- a/sdk/azcore/arm/runtime/pipeline.go
+++ b/sdk/azcore/arm/runtime/pipeline.go
@@ -46,7 +46,7 @@ func NewPipeline(module, version string, cred azcore.TokenCredential, plOpts azr
 }
 
 func getConfiguration(o *azpolicy.ClientOptions) (cloud.ServiceConfiguration, error) {
-	c := cloud.AzurePublicCloud
+	c := cloud.AzurePublic
 	if !reflect.ValueOf(o.Cloud).IsZero() {
 		c = o.Cloud
 	}
diff --git a/sdk/azcore/cloud/cloud.go b/sdk/azcore/cloud/cloud.go
index 232792ca4c..9d077a3e12 100644
--- a/sdk/azcore/cloud/cloud.go
+++ b/sdk/azcore/cloud/cloud.go
@@ -15,8 +15,8 @@ var (
 	AzureGovernment = Configuration{
 		ActiveDirectoryAuthorityHost: "https://login.microsoftonline.us/", Services: map[ServiceName]ServiceConfiguration{},
 	}
-	// AzurePublicCloud contains configuration for Azure Public Cloud.
-	AzurePublicCloud = Configuration{
+	// AzurePublic contains configuration for Azure Public Cloud.
+	AzurePublic = Configuration{
 		ActiveDirectoryAuthorityHost: "https://login.microsoftonline.com/", Services: map[ServiceName]ServiceConfiguration{},
 	}
 )
diff --git a/sdk/azidentity/azidentity.go b/sdk/azidentity/azidentity.go
index 341601bdd9..7a3e4454eb 100644
--- a/sdk/azidentity/azidentity.go
+++ b/sdk/azidentity/azidentity.go
@@ -42,7 +42,7 @@ func setAuthorityHost(cc cloud.Configuration) (string, error) {
 		if len(cc.Services) > 0 {
 			return "", errors.New("missing ActiveDirectoryAuthorityHost for specified cloud")
 		}
-		host = cloud.AzurePublicCloud.ActiveDirectoryAuthorityHost
+		host = cloud.AzurePublic.ActiveDirectoryAuthorityHost
 		if envAuthorityHost := os.Getenv(azureAuthorityHost); envAuthorityHost != "" {
 			host = envAuthorityHost
 		}
diff --git a/sdk/data/azcosmos/CHANGELOG.md b/sdk/data/azcosmos/CHANGELOG.md
index 6bd050a4d5..e062a92998 100644
--- a/sdk/data/azcosmos/CHANGELOG.md
+++ b/sdk/data/azcosmos/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 ### Features Added
 * Added single partition query support.
+* Added Azure AD authentication support through `azcosmos.NewClient`
 
 ### Breaking Changes
 * This module now requires Go 1.18
diff --git a/sdk/data/azcosmos/README.md b/sdk/data/azcosmos/README.md
index ac0bd012e2..50792db7b4 100644
--- a/sdk/data/azcosmos/README.md
+++ b/sdk/data/azcosmos/README.md
@@ -44,13 +44,29 @@ The following section provides several code snippets covering some of the most c
 
 ### Create Cosmos Client
 
+The clients support different forms of authentication. The azcosmos library supports authorization via Azure Active Directory or an account key.
+
+**Using Azure Active Directory**
+
+```go
+import "github.com/Azure/azure-sdk-for-go/sdk/azidentity"
+
+cred, err := azidentity.NewDefaultAzureCredential(nil)
+handle(err)
+client, err := azcosmos.NewClient("myAccountEndpointURL", cred, nil)
+handle(err)
+```
+
+**Using account keys**
+
 ```go
 const (
     cosmosDbEndpoint = "someEndpoint"
     cosmosDbKey = "someKey"
 )
 
-cred, _ := azcosmos.NewKeyCredential(cosmosDbKey)
+cred, err := azcosmos.NewKeyCredential(cosmosDbKey)
+handle(err)
 client, err := azcosmos.NewClientWithKey(cosmosDbEndpoint, cred, nil)
 handle(err)
 ```
diff --git a/sdk/data/azcosmos/cosmos_client.go b/sdk/data/azcosmos/cosmos_client.go
index c9df942bf2..b1d62a73cb 100644
--- a/sdk/data/azcosmos/cosmos_client.go
+++ b/sdk/data/azcosmos/cosmos_client.go
@@ -7,9 +7,12 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"net/http"
+	"net/url"
 	"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/streaming"
@@ -26,15 +29,27 @@ func (c *Client) Endpoint() string {
 	return c.endpoint
 }
 
-// NewClientWithKey creates a new instance of Cosmos client with the specified values. It uses the default pipeline configuration.
+// NewClientWithKey creates a new instance of Cosmos client with shared key authentication. It uses the default pipeline configuration.
 // endpoint - The cosmos service endpoint to use.
 // 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(cred, o)}, nil
+	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{newSharedKeyCredPolicy(cred)}, o)}, nil
 }
 
-func newPipeline(cred KeyCredential, options *ClientOptions) azruntime.Pipeline {
+// NewClient creates a new instance of Cosmos client with Azure AD access token authentication. It uses the default pipeline configuration.
+// endpoint - The cosmos service endpoint to use.
+// cred - The credential used to authenticate with the cosmos service.
+// options - Optional Cosmos client options.  Pass nil to accept default values.
+func NewClient(endpoint string, cred azcore.TokenCredential, o *ClientOptions) (*Client, error) {
+	scope, err := createScopeFromEndpoint(endpoint)
+	if err != nil {
+		return nil, err
+	}
+	return &Client{endpoint: endpoint, pipeline: newPipeline([]policy.Policy{azruntime.NewBearerTokenPolicy(cred, scope, nil), &cosmosBearerTokenPolicy{}}, o)}, nil
+}
+
+func newPipeline(authPolicy []policy.Policy, options *ClientOptions) azruntime.Pipeline {
 	if options == nil {
 		options = &ClientOptions{}
 	}
@@ -42,15 +57,24 @@ func newPipeline(cred KeyCredential, options *ClientOptions) azruntime.Pipeline
 	return azruntime.NewPipeline("azcosmos", serviceLibVersion,
 		azruntime.PipelineOptions{
 			PerCall: []policy.Policy{
-				newSharedKeyCredPolicy(cred),
 				&headerPolicies{
 					enableContentResponseOnWrite: options.EnableContentResponseOnWrite,
 				},
 			},
+			PerRetry: authPolicy,
 		},
 		&options.ClientOptions)
 }
 
+func createScopeFromEndpoint(endpoint string) ([]string, error) {
+	u, err := url.Parse(endpoint)
+	if err != nil {
+		return nil, err
+	}
+
+	return []string{fmt.Sprintf("%s://%s/.default", u.Scheme, u.Hostname())}, nil
+}
+
 // NewDatabase returns a struct that represents a database and allows database level operations.
 // id - The id of the database.
 func (c *Client) NewDatabase(id string) (*DatabaseClient, error) {
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token.go b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
new file mode 100644
index 0000000000..dfba0e9b3c
--- /dev/null
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token.go
@@ -0,0 +1,28 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+const lenBearerTokenPrefix = len("Bearer ")
+
+type cosmosBearerTokenPolicy struct {
+}
+
+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")
+	}
+
+	token := currentAuthorization[lenBearerTokenPrefix:]
+	req.Raw().Header.Set(headerAuthorization, fmt.Sprintf("type=aad&ver=1.0&sig=%v", token))
+	return req.Next()
+}
diff --git a/sdk/data/azcosmos/doc.go b/sdk/data/azcosmos/doc.go
index 30ccb34271..b8689fa183 100644
--- a/sdk/data/azcosmos/doc.go
+++ b/sdk/data/azcosmos/doc.go
@@ -12,6 +12,19 @@ The azcosmos package is capable of:
 
 Creating the Client
 
+Types of Credentials
+The clients support different forms of authentication. The azcosmos library supports
+authorization via Azure Active Directory or an account key.
+
+Using Azure Active Directory
+To create a client, you can use any of the TokenCredential implementations provided by `azidentity`.
+	cred, err := azidentity.NewClientSecretCredential("tenantId", "clientId", "clientSecret")
+	handle(err)
+	client, err := azcosmos.NewClient("myAccountEndpointURL", cred, nil)
+	handle(err)
+
+
+Using account keys
 To create a client, you will need the account's endpoint URL and a key credential.
 
 	cred, err := azcosmos.NewKeyCredential("myAccountKey")
diff --git a/sdk/data/azcosmos/go.mod b/sdk/data/azcosmos/go.mod
index 53a61a5260..f1f166c940 100644
--- a/sdk/data/azcosmos/go.mod
+++ b/sdk/data/azcosmos/go.mod
@@ -5,14 +5,22 @@ go 1.18
 require (
 	github.com/Azure/azure-sdk-for-go v63.2.0+incompatible
 	github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0
+	github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0
 	github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2
 	github.com/stretchr/testify v1.7.0
 )
 
 require (
+	github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/golang-jwt/jwt v3.2.1+incompatible // indirect
+	github.com/google/uuid v1.1.1 // indirect
+	github.com/kylelemons/godebug v1.1.0 // indirect
+	github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 // indirect
 	golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f // indirect
+	golang.org/x/sys v0.0.0-20211019181941-9d821ace8654 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )
diff --git a/sdk/data/azcosmos/go.sum b/sdk/data/azcosmos/go.sum
index 11b87f3d51..edd062b165 100644
--- a/sdk/data/azcosmos/go.sum
+++ b/sdk/data/azcosmos/go.sum
@@ -2,13 +2,28 @@ github.com/Azure/azure-sdk-for-go v63.2.0+incompatible h1:OIqkK/zTGqVUuzpEvY0B1Y
 github.com/Azure/azure-sdk-for-go v63.2.0+incompatible/go.mod h1:9XXNKU+eRnpl9moKnB4QOLf1HestfXbmab5FXxiDBjc=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0 h1:D7l5jspkc4kwBYRWoZE4DQnu6LVpLwDsMZjBKS4wZLQ=
 github.com/Azure/azure-sdk-for-go/sdk/azcore v0.23.0/go.mod h1:w5pDIZuawUmY3Bj4tVx3Xb8KS96ToB0j315w9rqpAg0=
+github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0 h1:NVS/4LOQfkBpk+B1VopIzv1ptmYeEskA8w/3K/w7vjo=
+github.com/Azure/azure-sdk-for-go/sdk/azidentity v0.14.0/go.mod h1:RG0cZndeZM17StwohYclmcXSr4oOJ8b1I5hB8llIc6Y=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2 h1:Px2KVERcYEg2Lv25AqC2hVr0xUWaq94wuEObLIkYzmA=
 github.com/Azure/azure-sdk-for-go/sdk/internal v0.9.2/go.mod h1:CdSJQNNzZhCkwDaV27XV1w48ZBPtxe7mlrZAsPNxD5g=
+github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0 h1:WVsrXCnHlDDX8ls+tootqRE87/hL9S/g4ewig9RsD/c=
+github.com/AzureAD/microsoft-authentication-library-for-go v0.4.0/go.mod h1:Vt9sXTKwMyGcOxSmLDMnGPgqsUg7m8pe215qMLrDXw4=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dnaeon/go-vcr v1.1.0 h1:ReYa/UBrRyQdant9B4fNHGoCNKw6qh6P0fsdGmZpR7c=
 github.com/dnaeon/go-vcr v1.1.0/go.mod h1:M7tiix8f0r6mKKJ3Yq/kqU1OYf3MnfmBWVbPx/yU9ko=
+github.com/golang-jwt/jwt v3.2.1+incompatible h1:73Z+4BJcrTC+KczS6WvTPvRGOp1WmfEP4Q1lOd9Z/+c=
+github.com/golang-jwt/jwt v3.2.1+incompatible/go.mod h1:8pz2t5EyA70fFQQSrl6XZXzqecmYZeUEB8OUGHkxJ+I=
+github.com/golang-jwt/jwt/v4 v4.2.0 h1:besgBTC8w8HjP6NzQdxwKH9Z5oQMZ24ThTrHp3cZ8eU=
+github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY=
+github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
+github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/modocache/gover v0.0.0-20171022184752-b58185e213c5/go.mod h1:caMODM3PzxT8aQXRPkAt8xlV/e7d7w8GM5g0fa5F0D8=
+github.com/montanaflynn/stats v0.6.6/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
+github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4 h1:Qj1ukM4GlMWXNdMBuXcXfz/Kw9s1qm0CLY32QxuSImI=
+github.com/pkg/browser v0.0.0-20210115035449-ce105d075bb4/go.mod h1:N6UoU20jOqggOuDwUaBQpluzLNDqif3kq9z2wpdYEfQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -16,6 +31,8 @@ github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5Cc
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
+golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897 h1:pLI5jrR7OSLijeIDcmRxNmw2api+jEfxLoykJVice/E=
+golang.org/x/crypto v0.0.0-20201016220609-9e8e0b390897/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20201010224723-4f7140c49acb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
 golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f h1:OfiFi4JbukWwe3lzw+xunroH1mnC1e2Gy5cxNJApiSY=
@@ -23,6 +40,9 @@ golang.org/x/net v0.0.0-20211015210444-4f30a5c0130f/go.mod h1:9nx3DQGgdP8bBQD5qx
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20211019181941-9d821ace8654 h1:id054HUawV2/6IGm2IV8KZQjqtwAOo2CYlOToYqa0d0=
+golang.org/x/sys v0.0.0-20211019181941-9d821ace8654/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
@@ -31,6 +51,7 @@ golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGm
 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.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b h1:h8qDotaEPuJATrMmW04NCwg7v22aHH28wwpauUhK9Oo=


Test patch:
diff --git a/eng/common/testproxy/docker-start-proxy.ps1 b/eng/common/testproxy/docker-start-proxy.ps1
index c15fc37a7e..c34c9d38d8 100644
--- a/eng/common/testproxy/docker-start-proxy.ps1
+++ b/eng/common/testproxy/docker-start-proxy.ps1
@@ -38,8 +38,7 @@ catch {
 
 $SELECTED_IMAGE_TAG = $(Get-Content "$PSScriptRoot/target_version.txt" -Raw).Trim()
 $CONTAINER_NAME = "ambitious_azsdk_test_proxy"
-$LINUX_IMAGE_SOURCE = "azsdkengsys.azurecr.io/engsys/testproxy-lin:${SELECTED_IMAGE_TAG}"
-$WINDOWS_IMAGE_SOURCE = "azsdkengsys.azurecr.io/engsys/testproxy-win:${SELECTED_IMAGE_TAG}"
+$IMAGE_SOURCE = "azsdkengsys.azurecr.io/engsys/testproxy:${SELECTED_IMAGE_TAG}"
 
 if($VersionOverride) { 
     Write-Host "Overriding default target proxy version of '$SELECTED_IMAGE_TAG' with override $VersionOverride."
@@ -58,8 +57,6 @@ function Get-Proxy-Container(){
                 | Select-Object -First 1)
 }
 
-
-$SelectedImage = $LINUX_IMAGE_SOURCE
 $Initial = ""
 $AdditionalContainerArgs = "--add-host=host.docker.internal:host-gateway"
 
@@ -67,7 +64,6 @@ $AdditionalContainerArgs = "--add-host=host.docker.internal:host-gateway"
 # however, in CI, windows images default to _windows_ containers. We cannot swap them. We can tell if we're in a CI build by
 # checking for the environment variable TF_BUILD.
 if ($IsWindows -and $env:TF_BUILD){
-    $SelectedImage = $WINDOWS_IMAGE_SOURCE
     $Initial = "C:"
     $AdditionalContainerArgs = ""
 }
@@ -93,9 +89,9 @@ if ($Mode -eq "start"){
     else {
         $attempts = 0
         Write-Host "Attempting creation of Docker host $CONTAINER_NAME"
-        Write-Host "docker container create -v `"${root}:${Initial}/srv/testproxy`" $AdditionalContainerArgs -p 5001:5001 -p 5000:5000 --name $CONTAINER_NAME $SelectedImage"
+        Write-Host "docker container create -v `"${root}:${Initial}/srv/testproxy`" $AdditionalContainerArgs -p 5001:5001 -p 5000:5000 --name $CONTAINER_NAME $IMAGE_SOURCE"
         while($attempts -lt 3){
-            docker container create -v "${root}:${Initial}/srv/testproxy" $AdditionalContainerArgs -p 5001:5001 -p 5000:5000 --name $CONTAINER_NAME $SelectedImage
+            docker container create -v "${root}:${Initial}/srv/testproxy" $AdditionalContainerArgs -p 5001:5001 -p 5000:5000 --name $CONTAINER_NAME $IMAGE_SOURCE
 
             if($LASTEXITCODE -ne 0){
                 $attempts += 1
diff --git a/eng/common/testproxy/target_version.txt b/eng/common/testproxy/target_version.txt
index ac34389c60..a72b151d94 100644
--- a/eng/common/testproxy/target_version.txt
+++ b/eng/common/testproxy/target_version.txt
@@ -1,1 +1,1 @@
-1.0.0-dev.20220224.2
+1.0.0-dev.20220426.1
diff --git a/eng/pipelines/templates/steps/build-test.yml b/eng/pipelines/templates/steps/build-test.yml
index 17dc761ac5..b06ca963f8 100644
--- a/eng/pipelines/templates/steps/build-test.yml
+++ b/eng/pipelines/templates/steps/build-test.yml
@@ -48,6 +48,9 @@ steps:
 
   - ${{ if eq(parameters.TestProxy, true) }}:
     - template: /eng/common/testproxy/test-proxy-tool.yml
+      parameters:
+        targetVersion: '1.0.0-dev.20220224.2'
+
 
   - task: PowerShell@2
     displayName: 'Run Tests'
diff --git a/sdk/azcore/arm/runtime/pipeline_test.go b/sdk/azcore/arm/runtime/pipeline_test.go
index 239bce745a..056a1c3a2d 100644
--- a/sdk/azcore/arm/runtime/pipeline_test.go
+++ b/sdk/azcore/arm/runtime/pipeline_test.go
@@ -167,7 +167,7 @@ func TestPipelineWithCustomPolicies(t *testing.T) {
 }
 
 func TestPipelineAudience(t *testing.T) {
-	for _, c := range []cloud.Configuration{cloud.AzureChina, cloud.AzureGovernment, cloud.AzurePublicCloud} {
+	for _, c := range []cloud.Configuration{cloud.AzureChina, cloud.AzureGovernment, cloud.AzurePublic} {
 		srv, close := mock.NewServer()
 		defer close()
 		srv.AppendResponse(mock.WithStatusCode(200))
diff --git a/sdk/azidentity/azidentity_test.go b/sdk/azidentity/azidentity_test.go
index e936ba15ad..10c33ac5ab 100644
--- a/sdk/azidentity/azidentity_test.go
+++ b/sdk/azidentity/azidentity_test.go
@@ -156,7 +156,7 @@ func setEnvironmentVariables(t *testing.T, vars map[string]string) {
 }
 
 func Test_WellKnownHosts(t *testing.T) {
-	for _, cloud := range []cloud.Configuration{cloud.AzureChina, cloud.AzureGovernment, cloud.AzurePublicCloud} {
+	for _, cloud := range []cloud.Configuration{cloud.AzureChina, cloud.AzureGovernment, cloud.AzurePublic} {
 		host, err := setAuthorityHost(cloud)
 		if err != nil {
 			t.Fatal(err)
@@ -196,7 +196,7 @@ func Test_DefaultAuthorityHost(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	if authorityHost != cloud.AzurePublicCloud.ActiveDirectoryAuthorityHost {
+	if authorityHost != cloud.AzurePublic.ActiveDirectoryAuthorityHost {
 		t.Fatal("unexpected default host: " + authorityHost)
 	}
 }
diff --git a/sdk/data/azcosmos/cosmos_client_test.go b/sdk/data/azcosmos/cosmos_client_test.go
index 1eab3dcba3..6bda337e82 100644
--- a/sdk/data/azcosmos/cosmos_client_test.go
+++ b/sdk/data/azcosmos/cosmos_client_test.go
@@ -354,6 +354,22 @@ func TestSendQuery(t *testing.T) {
 	}
 }
 
+func TestCreateScopeFromEndpoint(t *testing.T) {
+	url := "https://foo.documents.azure.com:443/"
+	scope, err := createScopeFromEndpoint(url)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if scope[0] != "https://foo.documents.azure.com/.default" {
+		t.Errorf("Expected %v, but got %v", "https://foo.documents.azure.com/.default", scope[0])
+	}
+
+	if len(scope) != 1 {
+		t.Errorf("Expected %v, but got %v", 1, len(scope))
+	}
+}
+
 type pipelineVerifier struct {
 	requests []pipelineVerifierRequest
 }
diff --git a/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
new file mode 100644
index 0000000000..e5d3eecade
--- /dev/null
+++ b/sdk/data/azcosmos/cosmos_policy_bearer_token_test.go
@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"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/mock"
+)
+
+func TestConvertBearerToken(t *testing.T) {
+	srv, close := mock.NewTLSServer()
+	defer close()
+	srv.SetResponse(mock.WithStatusCode(http.StatusOK))
+
+	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,
+	})
+
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	_, err = pl.Do(req)
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	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)
+	}
+}
+
+type bearerTokenVerify struct {
+	authHeaderContent string
+}
+
+func (p *bearerTokenVerify) Do(req *policy.Request) (*http.Response, error) {
+	p.authHeaderContent = req.Raw().Header.Get(headerAuthorization)
+
+	return req.Next()
+}
+
+type mockAuthPolicy struct{}
+
+func (p *mockAuthPolicy) Do(req *policy.Request) (*http.Response, error) {
+	req.Raw().Header.Set(headerAuthorization, "Bearer this is a test token")
+
+	return req.Next()
+}
diff --git a/sdk/data/azcosmos/emulator_cosmos_aad_test.go b/sdk/data/azcosmos/emulator_cosmos_aad_test.go
new file mode 100644
index 0000000000..8a1502507f
--- /dev/null
+++ b/sdk/data/azcosmos/emulator_cosmos_aad_test.go
@@ -0,0 +1,140 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package azcosmos
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+)
+
+func TestAAD(t *testing.T) {
+	emulatorTests := newEmulatorTests(t)
+	client := emulatorTests.getClient(t)
+
+	database := emulatorTests.createDatabase(t, context.TODO(), client, "aadTest")
+	defer emulatorTests.deleteDatabase(t, context.TODO(), database)
+	properties := ContainerProperties{
+		ID: "aContainer",
+		PartitionKeyDefinition: PartitionKeyDefinition{
+			Paths: []string{"/id"},
+		},
+	}
+
+	_, err := database.CreateContainer(context.TODO(), properties, nil)
+	if err != nil {
+		t.Fatalf("Failed to create container: %v", err)
+	}
+
+	aadClient := emulatorTests.getAadClient(t)
+
+	item := map[string]string{
+		"id":    "1",
+		"value": "2",
+	}
+
+	container, _ := aadClient.NewContainer("aadTest", "aContainer")
+	pk := NewPartitionKeyString("1")
+
+	marshalled, err := json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	itemResponse, err := container.CreateItem(context.TODO(), pk, marshalled, nil)
+	if err != nil {
+		t.Fatalf("Failed to create item: %v", err)
+	}
+
+	if itemResponse.SessionToken == "" {
+		t.Fatalf("Session token is empty")
+	}
+
+	// No content on write by default
+	if len(itemResponse.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", itemResponse.Value)
+	}
+
+	itemResponse, err = container.ReadItem(context.TODO(), pk, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to read item: %v", err)
+	}
+
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	var itemResponseBody map[string]interface{}
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "2" {
+		t.Fatalf("Expected value to be 2, got %v", itemResponseBody["value"])
+	}
+
+	item["value"] = "3"
+	marshalled, err = json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+	itemResponse, err = container.ReplaceItem(context.TODO(), pk, "1", marshalled, &ItemOptions{EnableContentResponseOnWrite: true})
+	if err != nil {
+		t.Fatalf("Failed to replace item: %v", err)
+	}
+
+	// Explicitly requesting body on write
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "3" {
+		t.Fatalf("Expected value to be 3, got %v", itemResponseBody["value"])
+	}
+
+	item["value"] = "4"
+	marshalled, err = json.Marshal(item)
+	if err != nil {
+		t.Fatal(err)
+	}
+	itemResponse, err = container.UpsertItem(context.TODO(), pk, marshalled, &ItemOptions{EnableContentResponseOnWrite: true})
+	if err != nil {
+		t.Fatalf("Failed to upsert item: %v", err)
+	}
+
+	// Explicitly requesting body on write
+	if len(itemResponse.Value) == 0 {
+		t.Fatalf("Expected non-empty response, got %v", itemResponse.Value)
+	}
+
+	err = json.Unmarshal(itemResponse.Value, &itemResponseBody)
+	if err != nil {
+		t.Fatalf("Failed to unmarshal item response: %v", err)
+	}
+	if itemResponseBody["id"] != "1" {
+		t.Fatalf("Expected id to be 1, got %v", itemResponseBody["id"])
+	}
+	if itemResponseBody["value"] != "4" {
+		t.Fatalf("Expected value to be 4, got %v", itemResponseBody["value"])
+	}
+
+	itemResponse, err = container.DeleteItem(context.TODO(), pk, "1", nil)
+	if err != nil {
+		t.Fatalf("Failed to replace item: %v", err)
+	}
+
+	if len(itemResponse.Value) != 0 {
+		t.Fatalf("Expected empty response, got %v", itemResponse.Value)
+	}
+}
diff --git a/sdk/data/azcosmos/emulator_tests.go b/sdk/data/azcosmos/emulator_tests.go
index d0a1793995..c8994ab939 100644
--- a/sdk/data/azcosmos/emulator_tests.go
+++ b/sdk/data/azcosmos/emulator_tests.go
@@ -5,8 +5,14 @@ package azcosmos
 
 import (
 	"context"
+	"encoding/base64"
 	"os"
+	"strconv"
 	"testing"
+	"time"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
 )
 
 type emulatorTests struct {
@@ -36,6 +42,16 @@ func (e *emulatorTests) getClient(t *testing.T) *Client {
 	return client
 }
 
+func (e *emulatorTests) getAadClient(t *testing.T) *Client {
+	cred := &emulatorTokenCredential{}
+	client, err := NewClient(e.host, cred, nil)
+	if err != nil {
+		t.Fatalf("Failed to create client: %v", err)
+	}
+
+	return client
+}
+
 func (e *emulatorTests) createDatabase(
 	t *testing.T,
 	ctx context.Context,
@@ -64,3 +80,47 @@ func (e *emulatorTests) deleteDatabase(
 		t.Fatalf("Failed to delete database: %v", err)
 	}
 }
+
+type emulatorTokenCredential struct {
+}
+
+func (c *emulatorTokenCredential) GetToken(ctx context.Context, options policy.TokenRequestOptions) (*azcore.AccessToken, error) {
+	header := `{"typ":"JWT","alg":"RS256","x5t":"CosmosEmulatorPrimaryMaster","kid":"CosmosEmulatorPrimaryMaster"}`
+	unixNow := time.Now().Unix()
+	expiration := unixNow + 7200
+	payload := `{ 
+		"appid":"localhost", 
+		"aio":"",
+		"appidacr":"1",
+		"idp": "https://localhost:8081/",
+		"oid": "96313034-4739-43cb-93cd-74193adbe5b6",
+		"rh": "",
+		"sub": "localhost",
+		"tid": "EmulatorFederation",
+		"uti": "",
+		"ver": "1.0",
+		"scp": "user_impersonation",
+		"groups":[ 
+			"7ce1d003-4cb3-4879-b7c5-74062a35c66e",
+			"e99ff30c-c229-4c67-ab29-30a6aebc3e58",
+			"5549bb62-c77b-4305-bda9-9ec66b85d9e4",
+			"c44fd685-5c58-452c-aaf7-13ce75184f65",
+			"be895215-eab5-43b7-9536-9ef8fe130330"], 
+		"nbf":` + strconv.FormatInt(unixNow, 10) + `, 
+		"exp":` + strconv.FormatInt(expiration, 10) + `, 
+		"iat":` + strconv.FormatInt(unixNow, 10) + `,
+		"iss":"https://sts.fake-issuer.net/7b1999a1-dfd7-440e-8204-00170979b984",
+		"aud":"https://localhost.localhost" 
+	}`
+
+	headerBase64 := base64.RawURLEncoding.EncodeToString([]byte(header))
+	payloadBase64 := base64.RawURLEncoding.EncodeToString([]byte(payload))
+	masterKeyBase64 := base64.RawURLEncoding.EncodeToString([]byte("C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw=="))
+
+	token := headerBase64 + "." + payloadBase64 + "." + masterKeyBase64
+
+	return &azcore.AccessToken{
+		Token:     token,
+		ExpiresOn: time.Unix(expiration, 0),
+	}, nil
+}
diff --git a/sdk/data/azcosmos/example_test.go b/sdk/data/azcosmos/example_test.go
index 9d927a8f9f..32d5161468 100644
--- a/sdk/data/azcosmos/example_test.go
+++ b/sdk/data/azcosmos/example_test.go
@@ -12,9 +12,32 @@ import (
 	"time"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
+	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
 )
 
+func ExampleNewClient() {
+	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
+	if !ok {
+		panic("AZURE_COSMOS_ENDPOINT could not be found")
+	}
+
+	// Obtain a TokenCredential for the current environment
+	// Alternatively, you could use any of the other credential types
+	// For example, azidentity.NewClientSecretCredential("tenantId", "clientId", "clientSecret")
+	cred, err := azidentity.NewDefaultAzureCredential(nil)
+	if err != nil {
+		panic(err)
+	}
+
+	client, err := azcosmos.NewClient(endpoint, cred, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	fmt.Println(client)
+}
+
 func ExampleNewClientWithKey() {
 	endpoint, ok := os.LookupEnv("AZURE_COSMOS_ENDPOINT")
 	if !ok {