Enable token cache for client assertion login flow (#20297)

* Add client assert code for upcoming change

* cache token for client assertion

* work around AssemblyInformationalVersion issue

* changelog

* Update README.md

Co-authored-by: dingmeng-xue <[email protected]>

Author: isra-fel

src/Accounts/Accounts/ChangeLog.md

@@ -19,7 +19,9 @@
 -->
 
 ## Upcoming Release
-* Enabled caching tokens when logging in with a service principal. This could reduce network traffic and improve performance.
+* Enabled caching tokens when logging in with a service principal or client assertion. [#20013]
+    - This could reduce extra network traffic and improve performance.
+    - It also fixed the incorrectly short lifespan of tokens.
 * Upgraded target framework of Microsoft.Identity.Client to net461 [#20189]
 * Stored `ServicePrincipalSecret` and `CertificatePassword` into `AzKeyStore`.
 * Updated the reference of Azure PowerShell Common to 1.3.65-preview.

src/Accounts/Authentication/Identity/AsyncLockWithValue.cs

@@ -0,0 +1,194 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+//
+
+using System;
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Azure.PowerShell.Authenticators.Identity
+{
+    /// <summary>
+    /// Primitive that combines async lock and value cache
+    /// </summary>
+    /// <typeparam name="T"></typeparam>
+    internal sealed class AsyncLockWithValue<T>
+    {
+        private readonly object _syncObj = new object();
+        private Queue<TaskCompletionSource<Lock>> _waiters;
+        private bool _isLocked;
+        private bool _hasValue;
+        private T _value;
+
+        /// <summary>
+        /// Method that either returns cached value or acquire a lock.
+        /// If one caller has acquired a lock, other callers will be waiting for the lock to be released.
+        /// If value is set, lock is released and all waiters get that value.
+        /// If value isn't set, the next waiter in the queue will get the lock.
+        /// </summary>
+        /// <param name="async"></param>
+        /// <param name="cancellationToken"></param>
+        /// <returns></returns>
+        public async ValueTask<Lock> GetLockOrValueAsync(bool async, CancellationToken cancellationToken = default)
+        {
+            TaskCompletionSource<Lock> valueTcs;
+            lock (_syncObj)
+            {
+                // If there is a value, just return it
+                if (_hasValue)
+                {
+                    return new Lock(_value);
+                }
+
+                // If lock isn't acquire yet, acquire it and return to the caller
+                if (!_isLocked)
+                {
+                    _isLocked = true;
+                    return new Lock(this);
+                }
+
+                // Check cancellationToken before instantiating waiter
+                cancellationToken.ThrowIfCancellationRequested();
+
+                // If lock is already taken, create a waiter and wait either until value is set or lock can be acquired by this waiter
+                if(_waiters is null)
+                {
+                    _waiters = new Queue<TaskCompletionSource<Lock>>();
+                }
+                // if async == false, valueTcs will be waited only in this thread and only synchronously, so RunContinuationsAsynchronously isn't needed.
+                valueTcs = new TaskCompletionSource<Lock>(async ? TaskCreationOptions.RunContinuationsAsynchronously : TaskCreationOptions.None);
+                _waiters.Enqueue(valueTcs);
+            }
+
+            try
+            {
+                if (async)
+                {
+                    return await valueTcs.Task.AwaitWithCancellation(cancellationToken);
+                }
+
+#pragma warning disable AZC0104 // Use EnsureCompleted() directly on asynchronous method return value.
+#pragma warning disable AZC0111 // DO NOT use EnsureCompleted in possibly asynchronous scope.
+                valueTcs.Task.Wait(cancellationToken);
+                return valueTcs.Task.EnsureCompleted();
+#pragma warning restore AZC0111 // DO NOT use EnsureCompleted in possibly asynchronous scope.
+#pragma warning restore AZC0104 // Use EnsureCompleted() directly on asynchronous method return value.
+            }
+            catch (OperationCanceledException)
+            {
+                // Throw OperationCanceledException only if another thread hasn't set a value to this waiter
+                // by calling either Reset or SetValue
+                if (valueTcs.TrySetCanceled(cancellationToken))
+                {
+                    throw;
+                }
+
+                return valueTcs.Task.Result;
+            }
+        }
+
+        /// <summary>
+        /// Set value to the cache and to all the waiters
+        /// </summary>
+        /// <param name="value"></param>
+        private void SetValue(T value)
+        {
+            Queue<TaskCompletionSource<Lock>> waiters;
+            lock (_syncObj)
+            {
+                _value = value;
+                _hasValue = true;
+                _isLocked = false;
+                if (_waiters == default)
+                {
+                    return;
+                }
+
+                waiters = _waiters;
+                _waiters = default;
+            }
+
+            while (waiters.Count > 0)
+            {
+                waiters.Dequeue().TrySetResult(new Lock(value));
+            }
+        }
+
+        /// <summary>
+        /// Release the lock and allow next waiter acquire it
+        /// </summary>
+        private void Reset()
+        {
+            TaskCompletionSource<Lock> nextWaiter = UnlockOrGetNextWaiter();
+            while (nextWaiter != default && !nextWaiter.TrySetResult(new Lock(this)))
+            {
+                nextWaiter = UnlockOrGetNextWaiter();
+            }
+        }
+
+        private TaskCompletionSource<Lock> UnlockOrGetNextWaiter()
+        {
+            lock (_syncObj)
+            {
+                if (!_isLocked)
+                {
+                    return default;
+                }
+
+                if (_waiters == default)
+                {
+                    _isLocked = false;
+                    return default;
+                }
+
+                while (_waiters.Count > 0)
+                {
+                    var nextWaiter = _waiters.Dequeue();
+                    if (!nextWaiter.Task.IsCompleted)
+                    {
+                        // Return the waiter only if it wasn't canceled already
+                        return nextWaiter;
+                    }
+                }
+
+                _isLocked = false;
+                return default;
+            }
+        }
+
+        public readonly struct Lock : IDisposable
+        {
+            private readonly AsyncLockWithValue<T> _owner;
+            public bool HasValue => _owner == default;
+            public T Value { get; }
+
+            public Lock(T value)
+            {
+                _owner = default;
+                Value = value;
+            }
+
+            public Lock(AsyncLockWithValue<T> owner)
+            {
+                _owner = owner;
+                Value = default;
+            }
+
+            public void SetValue(T value) => _owner.SetValue(value);
+
+            public void Dispose() => _owner?.Reset();
+        }
+    }
+}

src/Accounts/Authentication/Identity/Constants.cs

@@ -0,0 +1,57 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+//
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+
+namespace Microsoft.Azure.PowerShell.Authenticators.Identity
+{
+    internal class Constants
+    {
+        public const string OrganizationsTenantId = "organizations";
+
+        public const string AdfsTenantId = "adfs";
+
+        // TODO: Currently this is piggybacking off the Azure CLI client ID, but needs to be switched once the Developer Sign On application is available
+        public const string DeveloperSignOnClientId = "04b07795-8ddb-461a-bbee-02f9e1bf7b46";
+
+        public static string SharedTokenCacheFilePath { get { return Path.Combine(DefaultMsalTokenCacheDirectory, DefaultMsalTokenCacheName); } }
+
+        public const int SharedTokenCacheAccessRetryCount = 100;
+
+        public static readonly TimeSpan SharedTokenCacheAccessRetryDelay = TimeSpan.FromMilliseconds(600);
+
+        public const string DefaultRedirectUrl = "http://localhost";
+
+        public static readonly string DefaultMsalTokenCacheDirectory = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData), ".IdentityService");
+
+        public const string DefaultMsalTokenCacheKeychainService = "Microsoft.Developer.IdentityService";
+
+        public const string DefaultMsalTokenCacheKeychainAccount = "MSALCache";
+
+        public const string DefaultMsalTokenCacheKeyringLabel = "MSALCache";
+
+        public const string DefaultMsalTokenCacheKeyringSchema = "msal.cache";
+
+        public const string DefaultMsalTokenCacheKeyringCollection = "default";
+
+        public static readonly KeyValuePair<string, string> DefaultMsaltokenCacheKeyringAttribute1 = new KeyValuePair<string, string>("MsalClientID", "Microsoft.Developer.IdentityService");
+
+        public static readonly KeyValuePair<string, string> DefaultMsaltokenCacheKeyringAttribute2 = new KeyValuePair<string, string>("Microsoft.Developer.IdentityService", "1.0.0.0");
+
+        public const string DefaultMsalTokenCacheName = "msal.cache";
+    }
+}

src/Accounts/Authentication/Identity/CredentialPipeline.cs

@@ -0,0 +1,58 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+//
+
+using Azure.Core;
+using Azure.Core.Pipeline;
+using Azure.Identity;
+using System;
+
+namespace Microsoft.Azure.PowerShell.Authenticators.Identity
+{
+    internal class CredentialPipeline
+    {
+        private static readonly Lazy<CredentialPipeline> s_singleton = new Lazy<CredentialPipeline>(() => new CredentialPipeline(new TokenCredentialOptions()));
+
+        private CredentialPipeline(TokenCredentialOptions options)
+        {
+            AuthorityHost = options.AuthorityHost;
+
+            HttpPipeline = HttpPipelineBuilder.Build(options, Array.Empty<HttpPipelinePolicy>(), Array.Empty<HttpPipelinePolicy>(), new CredentialResponseClassifier());
+        }
+
+        public CredentialPipeline(Uri authorityHost, HttpPipeline httpPipeline)
+        {
+            AuthorityHost = authorityHost;
+
+            HttpPipeline = httpPipeline;
+        }
+
+        public static CredentialPipeline GetInstance(TokenCredentialOptions options)
+        {
+            return options is null ? s_singleton.Value : new CredentialPipeline(options);
+        }
+
+        public Uri AuthorityHost { get; }
+
+        public HttpPipeline HttpPipeline { get; }
+
+        private class CredentialResponseClassifier : ResponseClassifier
+        {
+            public override bool IsRetriableResponse(HttpMessage message)
+            {
+                return base.IsRetriableResponse(message) || message.Response.Status == 404;
+            }
+        }
+    }
+}