Add tests to PersistentTokenCache (Azure#18887)

* Add tests to PersistentTokenCache

Author: christothes

sdk/identity/Azure.Identity/api/Azure.Identity.netstandard2.0.cs

@@ -196,8 +196,7 @@ public ManagedIdentityCredential(string clientId = null, Azure.Identity.TokenCre
     }
     public partial class PersistentTokenCache : Azure.Identity.TokenCache
     {
-        public PersistentTokenCache(Azure.Identity.PersistentTokenCacheOptions options) { }
-        public PersistentTokenCache(bool allowUnencryptedStorage = true) { }
+        public PersistentTokenCache(Azure.Identity.PersistentTokenCacheOptions options = null) { }
     }
     public partial class PersistentTokenCacheOptions
     {

sdk/identity/Azure.Identity/src/MsalCacheHelperWrapper.cs

@@ -0,0 +1,96 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Diagnostics;
+using System.Threading.Tasks;
+using Microsoft.Identity.Client;
+using Microsoft.Identity.Client.Extensions.Msal;
+
+namespace Azure.Identity
+{
+    internal class MsalCacheHelperWrapper
+    {
+        private MsalCacheHelper _helper;
+
+        /// <summary>
+        /// Default Constructor.
+        /// </summary>
+        public MsalCacheHelperWrapper()
+        {
+        }
+
+        /// <summary>
+        /// Creates a new instance of Microsoft.Identity.Client.Extensions.Msal.MsalCacheHelper.
+        /// To configure MSAL to use this cache persistence, call Microsoft.Identity.Client.Extensions.Msal.MsalCacheHelper.RegisterCache(Microsoft.Identity.Client.ITokenCache)
+        /// </summary>
+        /// <param name="storageCreationProperties"></param>
+        /// <param name="logger">Passing null uses a default logger</param>
+        /// <returns>A new instance of Microsoft.Identity.Client.Extensions.Msal.MsalCacheHelper.</returns>
+        public virtual async Task InitializeAsync(StorageCreationProperties storageCreationProperties, TraceSource logger = null)
+        {
+            _helper = await MsalCacheHelper.CreateAsync(storageCreationProperties, logger).ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// Performs a write -> read -> clear using the underlying persistence mechanism
+        /// and throws an Microsoft.Identity.Client.Extensions.Msal.MsalCachePersistenceException
+        /// if something goes wrong.
+        /// </summary>
+        /// <remarks>
+        /// Does not overwrite the token cache. Should never fail on Windows and Mac where
+        /// the cache accessors are guaranteed to exist by the OS.
+        /// </remarks>
+        public virtual void VerifyPersistence()
+        {
+            _helper.VerifyPersistence();
+        }
+
+        /// <summary>
+        /// Registers a token cache to synchronize with on disk storage.
+        /// </summary>
+        /// <param name="tokenCache"></param>
+        public virtual void RegisterCache(ITokenCache tokenCache)
+        {
+            _helper.RegisterCache(tokenCache);
+        }
+
+        /// <summary>
+        /// Unregisters a token cache so it no longer synchronizes with on disk storage.
+        /// </summary>
+        /// <param name="tokenCache"></param>
+        public virtual void UnregisterCache(ITokenCache tokenCache)
+        {
+            _helper.UnregisterCache(tokenCache);
+        }
+
+        /// <summary>
+        /// Clears the token store.
+        /// </summary>
+        public virtual void Clear()
+        {
+            _helper.Clear();
+        }
+
+        /// <summary>
+        /// Extracts the token cache data from the persistent store
+        /// </summary>
+        /// <remarks>
+        /// This method should be used with care. The data returned is unencrypted.
+        /// </remarks>
+        /// <returns>UTF-8 byte array of the unencrypted token cache</returns>
+        public virtual byte[] LoadUnencryptedTokenCache()
+        {
+            return _helper.LoadUnencryptedTokenCache();
+        }
+
+        /// <summary>
+        /// Saves an unencrypted, UTF-8 encoded byte array representing an MSAL token cache.
+        /// The save operation will persist the data in a secure location, as configured
+        /// in Microsoft.Identity.Client.Extensions.Msal.StorageCreationProperties
+        /// </summary>
+        public virtual void SaveUnencryptedTokenCache(byte[] tokenCache)
+        {
+            _helper.SaveUnencryptedTokenCache(tokenCache);
+        }
+    }
+}

sdk/identity/Azure.Identity/src/PersistentTokenCache.cs

@@ -16,44 +16,49 @@ public class PersistentTokenCache : TokenCache
     {
         // we are creating the MsalCacheHelper with a random guid based clientId to work around issue https://github.com/AzureAD/microsoft-authentication-extensions-for-dotnet/issues/98
         // This does not impact the functionality of the cacheHelper as the ClientId is only used to iterate accounts in the cache not for authentication purposes.
-        private static readonly string s_msalCacheClientId = Guid.NewGuid().ToString();
-
-        private static AsyncLockWithValue<MsalCacheHelper> cacheHelperLock = new AsyncLockWithValue<MsalCacheHelper>();
-        private static AsyncLockWithValue<MsalCacheHelper> s_ProtectedCacheHelperLock = new AsyncLockWithValue<MsalCacheHelper>();
-        private static AsyncLockWithValue<MsalCacheHelper> s_FallbackCacheHelperLock = new AsyncLockWithValue<MsalCacheHelper>();
+        internal static readonly string s_msalCacheClientId = Guid.NewGuid().ToString();
         private readonly bool _allowUnencryptedStorage;
         private readonly string _name;
-
-        /// <summary>
-        /// Creates a new instance of <see cref="PersistentTokenCache"/>.
-        /// </summary>
-        /// <param name="allowUnencryptedStorage"></param>
-        public PersistentTokenCache(bool allowUnencryptedStorage = true)
-        {
-            _allowUnencryptedStorage = allowUnencryptedStorage;
-        }
+        private static AsyncLockWithValue<MsalCacheHelperWrapper> cacheHelperLock = new AsyncLockWithValue<MsalCacheHelperWrapper>();
+        private readonly MsalCacheHelperWrapper _cacheHelperWrapper;
 
         /// <summary>
         /// Creates a new instance of <see cref="PersistentTokenCache"/> with the specified options.
         /// </summary>
         /// <param name="options">Options controlling the storage of the <see cref="PersistentTokenCache"/>.</param>
-        public PersistentTokenCache(PersistentTokenCacheOptions options)
+        public PersistentTokenCache(PersistentTokenCacheOptions options = null)
         {
             _allowUnencryptedStorage = options?.AllowUnencryptedStorage ?? false;
+            _name = options?.Name ?? Constants.DefaultMsalTokenCacheName;
+            _cacheHelperWrapper = new MsalCacheHelperWrapper();
+        }
 
-            _name = options?.Name;
+        internal PersistentTokenCache(PersistentTokenCacheOptions options, MsalCacheHelperWrapper cacheHelperWrapper)
+        {
+            _allowUnencryptedStorage = options?.AllowUnencryptedStorage ?? false;
+            _name = options?.Name ?? Constants.DefaultMsalTokenCacheName;
+            _cacheHelperWrapper = cacheHelperWrapper;
         }
 
         internal override async Task RegisterCache(bool async, ITokenCache tokenCache, CancellationToken cancellationToken)
         {
-            MsalCacheHelper cacheHelper = await GetCacheHelperAsync(async, cancellationToken).ConfigureAwait(false);
+            MsalCacheHelperWrapper cacheHelper = await GetCacheHelperAsync(async, cancellationToken).ConfigureAwait(false);
 
             cacheHelper.RegisterCache(tokenCache);
 
             await base.RegisterCache(async, tokenCache, cancellationToken).ConfigureAwait(false);
         }
 
-        private async Task<MsalCacheHelper> GetCacheHelperAsync(bool async, CancellationToken cancellationToken)
+        /// <summary>
+        /// Resets the <see cref="cacheHelperLock"/> so that tests can validate multiple calls to <see cref="RegisterCache"/>
+        /// This should only be used for testing.
+        /// </summary>
+        internal static void ResetWrapperCache()
+        {
+            cacheHelperLock = new AsyncLockWithValue<MsalCacheHelperWrapper>();
+        }
+
+        private async Task<MsalCacheHelperWrapper> GetCacheHelperAsync(bool async, CancellationToken cancellationToken)
         {
             using var asyncLock = await cacheHelperLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
 
@@ -62,19 +67,19 @@ private async Task<MsalCacheHelper> GetCacheHelperAsync(bool async, Cancellation
                 return asyncLock.Value;
             }
 
-            MsalCacheHelper cacheHelper;
+            MsalCacheHelperWrapper cacheHelper;
 
             try
             {
-                cacheHelper = string.IsNullOrEmpty(_name) ? await GetProtectedCacheHelperAsync(async, cancellationToken).ConfigureAwait(false) : await GetProtectedCacheHelperAsync(async, _name).ConfigureAwait(false);
+                cacheHelper = await GetProtectedCacheHelperAsync(async, _name).ConfigureAwait(false);
 
                 cacheHelper.VerifyPersistence();
             }
             catch (MsalCachePersistenceException)
             {
                 if (_allowUnencryptedStorage)
                 {
-                    cacheHelper = string.IsNullOrEmpty(_name) ? await GetFallbackCacheHelperAsync(async, cancellationToken).ConfigureAwait(false) : await GetFallbackCacheHelperAsync(async, _name).ConfigureAwait(false);
+                    cacheHelper = await GetFallbackCacheHelperAsync(async, _name).ConfigureAwait(false);
 
                     cacheHelper.VerifyPersistence();
                 }
@@ -89,69 +94,43 @@ private async Task<MsalCacheHelper> GetCacheHelperAsync(bool async, Cancellation
             return cacheHelper;
         }
 
-        private static async Task<MsalCacheHelper> GetProtectedCacheHelperAsync(bool async, CancellationToken cancellationToken)
-        {
-            using var asyncLock = await s_ProtectedCacheHelperLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
-
-            if (asyncLock.HasValue)
-            {
-                return asyncLock.Value;
-            }
-
-            MsalCacheHelper cacheHelper = await GetProtectedCacheHelperAsync(async, Constants.DefaultMsalTokenCacheName).ConfigureAwait(false);
-
-            asyncLock.SetValue(cacheHelper);
-
-            return cacheHelper;
-        }
-
-        private static async Task<MsalCacheHelper> GetProtectedCacheHelperAsync(bool async, string name)
+        private async Task<MsalCacheHelperWrapper> GetProtectedCacheHelperAsync(bool async, string name)
         {
             StorageCreationProperties storageProperties = new StorageCreationPropertiesBuilder(name, Constants.DefaultMsalTokenCacheDirectory, s_msalCacheClientId)
                 .WithMacKeyChain(Constants.DefaultMsalTokenCacheKeychainService, name)
                 .WithLinuxKeyring(Constants.DefaultMsalTokenCacheKeyringSchema, Constants.DefaultMsalTokenCacheKeyringCollection, name, Constants.DefaultMsaltokenCacheKeyringAttribute1, Constants.DefaultMsaltokenCacheKeyringAttribute2)
                 .Build();
 
-            MsalCacheHelper cacheHelper = await CreateCacheHelper(async, storageProperties).ConfigureAwait(false);
-
-            return cacheHelper;
-        }
-
-        private static async Task<MsalCacheHelper> GetFallbackCacheHelperAsync(bool async, CancellationToken cancellationToken)
-        {
-            using var asyncLock = await s_FallbackCacheHelperLock.GetLockOrValueAsync(async, cancellationToken).ConfigureAwait(false);
-
-            if (asyncLock.HasValue)
-            {
-                return asyncLock.Value;
-            }
-
-            MsalCacheHelper cacheHelper = await GetFallbackCacheHelperAsync(async, Constants.DefaultMsalTokenCacheName).ConfigureAwait(false);
-
-            asyncLock.SetValue(cacheHelper);
+            MsalCacheHelperWrapper cacheHelper = await InitializeCacheHelper(async, storageProperties).ConfigureAwait(false);
 
             return cacheHelper;
         }
 
-        private static async Task<MsalCacheHelper> GetFallbackCacheHelperAsync(bool async, string name)
+        private async Task<MsalCacheHelperWrapper> GetFallbackCacheHelperAsync(bool async, string name = Constants.DefaultMsalTokenCacheName)
         {
             StorageCreationProperties storageProperties = new StorageCreationPropertiesBuilder(name, Constants.DefaultMsalTokenCacheDirectory, s_msalCacheClientId)
                 .WithMacKeyChain(Constants.DefaultMsalTokenCacheKeychainService, name)
                 .WithLinuxUnprotectedFile()
                 .Build();
 
-            MsalCacheHelper cacheHelper = await CreateCacheHelper(async, storageProperties).ConfigureAwait(false);
+            MsalCacheHelperWrapper cacheHelper = await InitializeCacheHelper(async, storageProperties).ConfigureAwait(false);
 
             return cacheHelper;
         }
 
-        private static async Task<MsalCacheHelper> CreateCacheHelper(bool async, StorageCreationProperties storageProperties)
+        private async Task<MsalCacheHelperWrapper> InitializeCacheHelper(bool async, StorageCreationProperties storageProperties)
         {
-            return async
-                ? await MsalCacheHelper.CreateAsync(storageProperties).ConfigureAwait(false)
+            if (async)
+            {
+                await _cacheHelperWrapper.InitializeAsync(storageProperties).ConfigureAwait(false);
+            }
+            else
+            {
 #pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult(). Use the TaskExtensions.EnsureCompleted() extension method instead.
-                : MsalCacheHelper.CreateAsync(storageProperties).GetAwaiter().GetResult();
+                _cacheHelperWrapper.InitializeAsync(storageProperties).GetAwaiter().GetResult();
 #pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult(). Use the TaskExtensions.EnsureCompleted() extension method instead.
+            }
+            return _cacheHelperWrapper;
         }
     }
 }