AuthTypeChange: Adds new API to update the AuthType dynamically

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

@@ -6,10 +6,110 @@ namespace Microsoft.Azure.Cosmos
 {
     using System;
     using System.Globalization;
+    using System.Runtime.CompilerServices;
+    using System.Security.AccessControl;
+    using System.Threading;
     using System.Threading.Tasks;
+    using global::Azure;
+    using global::Azure.Core;
+    using HdrHistogram.Encoding;
+    using Microsoft.Azure.Cosmos.Authorization;
     using Microsoft.Azure.Cosmos.Tracing;
     using Microsoft.Azure.Documents;
     using Microsoft.Azure.Documents.Collections;
+    using static Microsoft.Azure.Cosmos.Query.Core.Metrics.ServerSideMetricsTokenizer;
+
+#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member
+#pragma warning disable SA1600 // Elements should be documented
+#pragma warning disable SA1649 // File name should match first type name
+    public sealed class CosmosAuthorizationTokenProvider : IDisposable
+    {
+        private CosmosAuthorizationTokenProvider provider;
+
+        private CosmosAuthorizationTokenProvider()
+        {
+        }
+
+        public static CosmosAuthorizationTokenProvider FromAzureKeyCredential(AzureKeyCredential authKeyOrResourceTokenCredential)
+        {
+            return new CosmosAuthorizationTokenProvider()
+            {
+                authorizationTokenProvider = new AzureKeyCredentialAuthorizationTokenProvider(authKeyOrResourceTokenCredential)
+            };
+        }
+
+        public static CosmosAuthorizationTokenProvider FromTokenCredential(string accountEndpoint, 
+            TokenCredential tokenCredential, 
+            TimeSpan? tokenCredentialBackgroundRefreshInterval = null)
+        {
+            return new CosmosAuthorizationTokenProvider()
+            {
+                authorizationTokenProvider = new AuthorizationTokenProviderTokenCredential(tokenCredential, new Uri(accountEndpoint), tokenCredentialBackgroundRefreshInterval)
+            };
+        }
+
+        internal AuthorizationTokenProvider authorizationTokenProvider { get; set; }
+
+        public void Update(CosmosAuthorizationTokenProvider tokenProvider)
+        {
+            if (tokenProvider == null) throw new ArgumentNullException(nameof(tokenProvider));
+
+            Volatile.Write(ref this.provider, tokenProvider);
+        }
+
+        public void Dispose()
+        {
+            if (this.authorizationTokenProvider != null)
+            {
+                this.authorizationTokenProvider?.Dispose();
+                this.authorizationTokenProvider = null;
+            }
+        }
+    }
+
+    internal class ComposedAuthorizationTokenProvider : AuthorizationTokenProvider
+    {
+        private CosmosAuthorizationTokenProvider authorizationTokenProvider;
+
+        public ComposedAuthorizationTokenProvider(CosmosAuthorizationTokenProvider authorizationTokenProvider)
+        {
+            if (authorizationTokenProvider?.authorizationTokenProvider == null) throw new ArgumentNullException(nameof(authorizationTokenProvider));
+
+            this.authorizationTokenProvider = authorizationTokenProvider;
+        }
+
+        public override ValueTask AddAuthorizationHeaderAsync(INameValueCollection headersCollection, Uri requestAddress, string verb, AuthorizationTokenType tokenType)
+        {
+            return this.authorizationTokenProvider.authorizationTokenProvider.AddAuthorizationHeaderAsync(headersCollection, requestAddress, verb, tokenType);
+        }
+
+        public override void Dispose()
+        {
+            if (this.authorizationTokenProvider == null)
+            {
+                this.authorizationTokenProvider.Dispose();
+                this.authorizationTokenProvider = null;
+            }
+        }
+
+        public override ValueTask<(string token, string payload)> GetUserAuthorizationAsync(string resourceAddress, string resourceType, string requestVerb, INameValueCollection headers, AuthorizationTokenType tokenType)
+        {
+            return this.authorizationTokenProvider.authorizationTokenProvider.GetUserAuthorizationAsync(resourceAddress, resourceType, requestVerb, headers, tokenType);
+        }
+
+        public override ValueTask<string> GetUserAuthorizationTokenAsync(string resourceAddress, string resourceType, string requestVerb, INameValueCollection headers, AuthorizationTokenType tokenType, ITrace trace)
+        {
+            return this.authorizationTokenProvider.authorizationTokenProvider.GetUserAuthorizationTokenAsync(resourceAddress, resourceType, requestVerb, headers, tokenType, trace);
+        }
+
+        public override void TraceUnauthorized(DocumentClientException dce, string authorizationToken, string payload)
+        {
+            this.authorizationTokenProvider.authorizationTokenProvider.TraceUnauthorized(dce, authorizationToken, payload);
+        }
+    }
+#pragma warning restore SA1649 // File name should match first type name
+#pragma warning restore SA1600 // Elements should be documented
+#pragma warning restore CS1591 // Missing XML comment for publicly visible type or member
 
     internal abstract class AuthorizationTokenProvider : ICosmosAuthorizationTokenProvider, IAuthorizationTokenProvider, IDisposable
     {

Microsoft.Azure.Cosmos/src/CosmosClient.cs

@@ -328,6 +328,40 @@ public CosmosClient(
         {
         }
 
+        /// <summary>
+        /// Creates a new CosmosClient with the account endpoint and given CosmosAuthorizationTokenProvider
+        /// </summary>
+        /// <param name="accountEndpoint">The cosmos service endpoint to use.</param>
+        /// <param name="tokenProvider"><see cref="CosmosAuthorizationTokenProvider"/>authorization token provider.</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[
+        ///   CosmosAuthorizationTokenProvider authorizationTokenProvider = CosmosAuthorizationTokenProvider.FromTokenCredential(endpoint, miTokenCredential);
+        ///   using (CosmosClient client = new CosmosClient(
+        ///       endpoint,
+        ///       miTokenCredential))
+        ///   {
+        ///        // do some work 
+        ///
+        ///        // Change 
+        ///        AzureKeyCredential masterKeyCredential = new AzureKeyCredential(authKey);
+        ///        authorizationTokenProvider.Update(CosmosAuthorizationTokenProvider.FromAzureKeyCredential(masterKeyCredential));
+        ///
+        ///        // do some more work 
+        /// }
+        /// ]]>
+        /// </code>
+        /// </example>
+        public CosmosClient(
+            string accountEndpoint,
+            CosmosAuthorizationTokenProvider tokenProvider,
+            CosmosClientOptions clientOptions = null)
+            : this(accountEndpoint, new ComposedAuthorizationTokenProvider(tokenProvider), clientOptions)
+        {
+        }
+
         /// <summary>
         /// Used by Compute
         /// Creates a new CosmosClient with the AuthorizationTokenProvider

Microsoft.Azure.Cosmos/tests/Microsoft.Azure.Cosmos.EmulatorTests/ClientTests.cs

@@ -157,6 +157,30 @@ public async Task InitTaskThreadSafe()
             }
         }
 
+        [TestMethod]
+        public async Task ValidateCredentialTypeChangeTest()
+        {
+            string authKey = ConfigurationManager.AppSettings["MasterKey"];
+            string endpoint = ConfigurationManager.AppSettings["GatewayEndpoint"];
+
+            LocalEmulatorTokenCredential miTokenCredential = new LocalEmulatorTokenCredential(authKey);
+            CosmosAuthorizationTokenProvider authorizationTokenProvider = CosmosAuthorizationTokenProvider.FromTokenCredential(endpoint, miTokenCredential);
+
+            using (CosmosClient client = new CosmosClient(
+                    endpoint,
+                    miTokenCredential))
+            {
+                // do some work 
+                await Task.Yield();
+
+                // Change 
+                AzureKeyCredential masterKeyCredential = new AzureKeyCredential(authKey);
+                authorizationTokenProvider.Update(CosmosAuthorizationTokenProvider.FromAzureKeyCredential(masterKeyCredential));
+
+                // do some work 
+                await Task.Yield();
+            }
+        }
 
         [TestMethod]
         public async Task ValidateAzureKeyCredentialDirectModeUpdateAsync()