KeyRefresh: Adds AzureKeyCredential support to enable key refresh sce…

…narios (#3276)

* Draft version of azureKeyCredential support (for key refresh scenario's)

* Defensive coding

* Adding tests

* Concurrent key update scenario fix

* Fixing test to cover for both master and resoruce tokens

* Incorporating feedback comment

* Addressing Fabian feedback comments

* Draft version of azureKeyCredential support (for key refresh scenario's)

* Defensive coding

* Adding tests

* Concurrent key update scenario fix

* Fixing test to cover for both master and resoruce tokens

* Addressing Fabian feedback comments

* Fixing merge rebase issues

* Adding a perf benchmark

* Pushing contract changes

* Fix: handling of concurrent access wile updating.

* Fix: one more concurrency related fix.

* Documentation suggestions

* Added direct mode diagnotics of AuthProvider LifeTime

* Update Microsoft.Azure.Cosmos/src/CosmosClient.cs

Co-authored-by: Matias Quaranta <[email protected]>

* Update Microsoft.Azure.Cosmos/src/CosmosClient.cs

Co-authored-by: Matias Quaranta <[email protected]>

* Update Microsoft.Azure.Cosmos/src/Fluent/CosmosClientBuilder.cs

Co-authored-by: Matias Quaranta <[email protected]>

* Addressing review comments

* Removing the int casting

* Addressing comments

* Adding the tracing contract test

* Base trace updated

Co-authored-by: Matias Quaranta <[email protected]>

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProvider.cs

@@ -13,6 +13,8 @@ namespace Microsoft.Azure.Cosmos
 
     internal abstract class AuthorizationTokenProvider : ICosmosAuthorizationTokenProvider, IAuthorizationTokenProvider, IDisposable
     {
+        private readonly DateTime creationTime = DateTime.UtcNow;
+
         public async Task AddSystemAuthorizationHeaderAsync(
             DocumentServiceRequest request, 
             string federationId, 
@@ -55,6 +57,11 @@ public abstract void TraceUnauthorized(
             string authorizationToken,
             string payload);
 
+        public virtual TimeSpan GetAge()
+        {
+            return DateTime.UtcNow.Subtract(this.creationTime);
+        }
+
         public static AuthorizationTokenProvider CreateWithResourceTokenOrAuthKey(string authKeyOrResourceToken)
         {
             if (string.IsNullOrEmpty(authKeyOrResourceToken))

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProviderMasterKey.cs

@@ -22,8 +22,7 @@ internal sealed class AuthorizationTokenProviderMasterKey : AuthorizationTokenPr
         private const string MacSignatureString = "to sign";
         private const string EnableAuthFailureTracesConfig = "enableAuthFailureTraces";
         private readonly Lazy<bool> enableAuthFailureTraces;
-        private readonly IComputeHash authKeyHashFunction;
-        private bool isDisposed = false;
+        private IComputeHash authKeyHashFunction;
 
         public AuthorizationTokenProviderMasterKey(IComputeHash computeHash)
         {
@@ -169,11 +168,14 @@ public override void TraceUnauthorized(
 
         public override void Dispose()
         {
-            if (!this.isDisposed)
-            {
-                this.authKeyHashFunction.Dispose();
-                this.isDisposed = true;
-            }
+            this.Dispose(disposing: true);
+
+            // This object will be cleaned up by the Dispose method.
+            // Therefore, you should call GC.SuppressFinalize to
+            // take this object off the finalization queue
+            // and prevent finalization code for this object
+            // from executing a second time.
+            GC.SuppressFinalize(this);
         }
 
         private static string NormalizeAuthorizationPayload(string input)
@@ -198,5 +200,28 @@ private static string NormalizeAuthorizationPayload(string input)
 
             return builder.ToString();
         }
+
+        // Dispose(bool disposing) executes in two distinct scenarios.
+        // If disposing equals true, the method has been called directly
+        // or indirectly by a user's code. Managed and unmanaged resources
+        // can be disposed.
+        // If disposing equals false, the method has been called by the
+        // runtime from inside the finalizer and you should not reference
+        // other objects. Only unmanaged resources can be disposed.
+        private void Dispose(bool disposing)
+        {
+            this.authKeyHashFunction?.Dispose();
+            this.authKeyHashFunction = null;
+        }
+
+        // Use C# finalizer syntax for finalization code.
+        // This finalizer will run only if the Dispose method does not get called.
+        // It gives your base class the opportunity to finalize.
+        ~AuthorizationTokenProviderMasterKey()
+        {
+            // Calling Dispose(disposing: false) is optimal in terms of
+            // readability and maintainability.
+            this.Dispose(disposing: false);
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Authorization/AuthorizationTokenProviderResourceToken.cs

@@ -70,6 +70,36 @@ public override void TraceUnauthorized(
 
         public override void Dispose()
         {
+            this.Dispose(disposing: true);
+
+            // This object will be cleaned up by the Dispose method.
+            // Therefore, you should call GC.SuppressFinalize to
+            // take this object off the finalization queue
+            // and prevent finalization code for this object
+            // from executing a second time.
+            GC.SuppressFinalize(this);
+        }
+
+        // Dispose(bool disposing) executes in two distinct scenarios.
+        // If disposing equals true, the method has been called directly
+        // or indirectly by a user's code. Managed and unmanaged resources
+        // can be disposed.
+        // If disposing equals false, the method has been called by the
+        // runtime from inside the finalizer and you should not reference
+        // other objects. Only unmanaged resources can be disposed.
+        private void Dispose(bool disposing)
+        {
+            // Do nothing
+        }
+
+        // Use C# finalizer syntax for finalization code.
+        // This finalizer will run only if the Dispose method does not get called.
+        // It gives your base class the opportunity to finalize.
+        ~AuthorizationTokenProviderResourceToken()
+        {
+            // Calling Dispose(disposing: false) is optimal in terms of
+            // readability and maintainability.
+            this.Dispose(disposing: false);
         }
     }
 }

Microsoft.Azure.Cosmos/src/Authorization/AzureKeyCredentialAuthorizationTokenProvider.cs

@@ -0,0 +1,129 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace Microsoft.Azure.Cosmos.Authorization
+{
+    using System;
+    using System.Threading;
+    using System.Threading.Tasks;
+    using global::Azure;
+    using Microsoft.Azure.Cosmos.Tracing;
+    using Microsoft.Azure.Documents;
+    using Microsoft.Azure.Documents.Collections;
+
+    internal class AzureKeyCredentialAuthorizationTokenProvider : AuthorizationTokenProvider
+    {
+        private readonly object refreshLock = new object();
+        private readonly AzureKeyCredential azureKeyCredential;
+
+        // keyObject is used to check for refresh 
+        private string currentKeyObject = null;
+
+        // Internal for unit testing
+        internal AuthorizationTokenProvider authorizationTokenProvider;
+
+        public AzureKeyCredentialAuthorizationTokenProvider(
+                AzureKeyCredential azureKeyCredential)
+        {
+            this.azureKeyCredential = azureKeyCredential ?? throw new ArgumentNullException(nameof(azureKeyCredential));
+            this.CheckAndRefreshTokenProvider();
+        }
+
+        public override ValueTask AddAuthorizationHeaderAsync(
+                INameValueCollection headersCollection, 
+                Uri requestAddress, 
+                string verb, 
+                AuthorizationTokenType tokenType)
+        {
+            this.CheckAndRefreshTokenProvider();
+            return this.authorizationTokenProvider.AddAuthorizationHeaderAsync(
+                headersCollection, 
+                requestAddress, 
+                verb, 
+                tokenType);
+        }
+
+        public override void Dispose()
+        {
+            this.authorizationTokenProvider?.Dispose();
+            this.authorizationTokenProvider = null;
+        }
+
+        public override ValueTask<(string token, string payload)> GetUserAuthorizationAsync(
+                string resourceAddress, 
+                string resourceType, 
+                string requestVerb, 
+                INameValueCollection headers, 
+                AuthorizationTokenType tokenType)
+        {
+            this.CheckAndRefreshTokenProvider();
+            return this.authorizationTokenProvider.GetUserAuthorizationAsync(
+                    resourceAddress, 
+                    resourceType, 
+                    requestVerb, 
+                    headers, 
+                    tokenType);
+        }
+
+        public override ValueTask<string> GetUserAuthorizationTokenAsync(
+                string resourceAddress, 
+                string resourceType, 
+                string requestVerb, 
+                INameValueCollection headers, 
+                AuthorizationTokenType tokenType, 
+                ITrace trace)
+        {
+            this.CheckAndRefreshTokenProvider();
+            return this.authorizationTokenProvider.GetUserAuthorizationTokenAsync(
+                    resourceAddress, 
+                    resourceType, 
+                    requestVerb, 
+                    headers, 
+                    tokenType, 
+                    trace);
+        }
+
+        public override void TraceUnauthorized(
+                DocumentClientException dce, 
+                string authorizationToken, 
+                string payload)
+        {
+            this.authorizationTokenProvider.TraceUnauthorized(
+                    dce, 
+                    authorizationToken, 
+                    payload);
+        }
+
+        public override TimeSpan GetAge()
+        {
+            return this.authorizationTokenProvider.GetAge();
+        }
+
+        private void CheckAndRefreshTokenProvider()
+        {
+            if (!Object.ReferenceEquals(this.currentKeyObject, this.azureKeyCredential.Key))
+            {
+                // Change is immediate for all new reqeust flows (pure compute work and should be very very quick)
+                // With-out lock possibility of concurrent updates (== #inflight reqeust's) but eventually only one will win
+                lock (this.refreshLock)
+                {
+                    // Process only if the authProvider is not yet exchanged
+                    if (!Object.ReferenceEquals(this.currentKeyObject, this.azureKeyCredential.Key))
+                    {
+                        AuthorizationTokenProvider newAuthProvider = AuthorizationTokenProvider.CreateWithResourceTokenOrAuthKey(this.azureKeyCredential.Key);
+                        AuthorizationTokenProvider currentAuthProvider = Interlocked.Exchange(ref this.authorizationTokenProvider, newAuthProvider);
+
+                        if (!Object.ReferenceEquals(newAuthProvider, currentAuthProvider))
+                        {
+                            // NewAuthProvider =>
+                            // 1. Credentials changed
+                            // 2. Dispose current token provider: NOT Disposed actively: There might be inflight usage => leaving it to GC finalizers
+                            Interlocked.Exchange(ref this.currentKeyObject, this.azureKeyCredential.Key);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -12,7 +12,9 @@ namespace Microsoft.Azure.Cosmos
     using System.Text;
     using System.Threading;
     using System.Threading.Tasks;
+    using global::Azure;
     using global::Azure.Core;
+    using Microsoft.Azure.Cosmos.Authorization;
     using Microsoft.Azure.Cosmos.Handlers;
     using Microsoft.Azure.Cosmos.Query.Core.Monads;
     using Microsoft.Azure.Cosmos.Query.Core.QueryPlan;
@@ -101,11 +103,11 @@ namespace Microsoft.Azure.Cosmos
     public class CosmosClient : IDisposable
     {
         internal readonly string Id = Guid.NewGuid().ToString();
+        private readonly object disposedLock = new object();
 
         private readonly string DatabaseRootUri = Paths.Databases_Root;
         private ConsistencyLevel? accountConsistencyLevel;
         private bool isDisposed = false;
-        private object disposedLock = new object();
 
         internal static int numberOfClientsCreated;
         internal static int NumberOfActiveClients;
@@ -237,6 +239,56 @@ public CosmosClient(
         {
         }
 
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint URI string and AzureKeyCredential.
+        /// AzureKeyCredential enables changing/updating master-key/ResourceToken while CosmosClient is still in use. 
+        /// 
+        /// CosmosClient is thread-safe. Its recommended to maintain a single instance of CosmosClient per lifetime 
+        /// of the application which enables efficient connection management and performance. Please refer to the
+        /// <see href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">performance guide</see>.
+        /// </summary>
+        /// <param name="accountEndpoint">The cosmos service endpoint to use</param>
+        /// <param name="authKeyOrResourceTokenCredential">AzureKeyCredential with master-key or resource token..</param>
+        /// <param name="clientOptions">(Optional) client options</param>
+        /// <example>
+        /// The CosmosClient is created with the AccountEndpoint, AccountKey or ResourceToken and configured to use "East US 2" region.
+        /// <code language="c#">
+        /// <![CDATA[
+        /// using Microsoft.Azure.Cosmos;
+        /// 
+        /// AzureKeyCredential keyCredential = new AzureKeyCredential("account-master-key/ResourceToken");
+        /// CosmosClient cosmosClient = new CosmosClient(
+        ///             "account-endpoint-from-portal", 
+        ///             keyCredential, 
+        ///             new CosmosClientOptions()
+        ///             {
+        ///                 ApplicationRegion = Regions.EastUS2,
+        ///             });
+        ///             
+        /// ....
+        /// 
+        /// // To udpate key/credentials 
+        /// keyCredential.Update("updated master-key/ResourceToken");
+        /// 
+        /// // Dispose cosmosClient at application exit
+        /// ]]>
+        /// </code>
+        /// </example>
+        /// <seealso cref="CosmosClientOptions"/>
+        /// <seealso cref="Fluent.CosmosClientBuilder"/>
+        /// <seealso href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">Performance Tips</seealso>
+        /// <seealso href="https://docs.microsoft.com/azure/cosmos-db/troubleshoot-dot-net-sdk">Diagnose and troubleshoot issues</seealso>
+        /// <remarks>AzureKeyCredential enables changing/updating master-key/ResourceToken whle CosmosClient is still in use.</remarks>
+        public CosmosClient(
+            string accountEndpoint,
+            AzureKeyCredential authKeyOrResourceTokenCredential,
+            CosmosClientOptions clientOptions = null)
+             : this(accountEndpoint,
+                     new AzureKeyCredentialAuthorizationTokenProvider(authKeyOrResourceTokenCredential),
+                     clientOptions)
+        {
+        }
+
         /// <summary>
         /// Creates a new CosmosClient with the account endpoint URI string and TokenCredential.
         /// 
@@ -339,6 +391,66 @@ public static async Task<CosmosClient> CreateAndInitializeAsync(string accountEn
             return cosmosClient;
         }
 
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint URI string and AzureKeyCredential.
+        /// AzureKeyCredential enables changing/updating master-key/ResourceToken while CosmosClient is still in use. 
+        /// 
+        /// In addition to that it initializes the client with containers provided i.e The SDK warms up the caches and 
+        /// connections before the first call to the service is made. Use this to obtain lower latency while startup of your application.
+        /// CosmosClient is thread-safe. Its recommended to maintain a single instance of CosmosClient per lifetime 
+        /// of the application which enables efficient connection management and performance. Please refer to the
+        /// <see href="https://docs.microsoft.com/azure/cosmos-db/performance-tips">performance guide</see>.
+        /// </summary>
+        /// <param name="accountEndpoint">The cosmos service endpoint to use</param>
+        /// <param name="authKeyOrResourceTokenCredential">AzureKeyCredential with master-key or resource token.</param>
+        /// <param name="containers">Containers to be initialized identified by it's database name and container name.</param>
+        /// <param name="cosmosClientOptions">(Optional) client options</param>
+        /// <param name="cancellationToken">(Optional) Cancellation Token</param>
+        /// <returns>
+        /// A CosmosClient object.
+        /// </returns>
+        /// <example>
+        /// The CosmosClient is created with the AccountEndpoint, AccountKey or ResourceToken and 2 containers in the account are initialized
+        /// <code language="c#">
+        /// <![CDATA[
+        /// using Microsoft.Azure.Cosmos;
+        /// List<(string, string)> containersToInitialize = new List<(string, string)>
+        /// { ("DatabaseName1", "ContainerName1"), ("DatabaseName2", "ContainerName2") };
+        /// 
+        /// AzureKeyCredential keyCredential = new AzureKeyCredential("account-master-key/ResourceToken");
+        /// CosmosClient cosmosClient = await CosmosClient.CreateAndInitializeAsync("account-endpoint-from-portal", 
+        ///                                                                         keyCredential,
+        ///                                                                         containersToInitialize)
+        ///             
+        /// ....
+        /// 
+        /// // To udpate key/credentials 
+        /// keyCredential.Update("updated master-key/ResourceToken");
+        /// 
+        /// // Dispose cosmosClient at application exit
+        /// ]]>
+        /// </code>
+        /// </example>
+        /// <remarks>AzureKeyCredential enables changing/updating master-key/ResourceToken whle CosmosClient is still in use.</remarks> 
+        public static async Task<CosmosClient> CreateAndInitializeAsync(string accountEndpoint,
+                                                                        AzureKeyCredential authKeyOrResourceTokenCredential,
+                                                                        IReadOnlyList<(string databaseId, string containerId)> containers,
+                                                                        CosmosClientOptions cosmosClientOptions = null,
+                                                                        CancellationToken cancellationToken = default)
+        {
+            if (containers == null)
+            {
+                throw new ArgumentNullException(nameof(containers));
+            }
+
+            CosmosClient cosmosClient = new CosmosClient(accountEndpoint,
+                                                         authKeyOrResourceTokenCredential,
+                                                         cosmosClientOptions);
+
+            await cosmosClient.InitializeContainersAsync(containers, cancellationToken);
+            return cosmosClient;
+        }
+
         /// <summary>
         /// Creates a new CosmosClient with the account endpoint URI string and TokenCredential.
         /// In addition to that it initializes the client with containers provided i.e The SDK warms up the caches and