Update Azure.Identity to 1.4 and MSAL lib to 4.30.1

src/Accounts/Accounts.Test/AutosaveTests.cs

@@ -53,7 +53,6 @@ void ResetState()
             Environment.SetEnvironmentVariable("Azure_PS_Data_Collection", "false");
             PowerShellTokenCacheProvider tokenProvider = new InMemoryTokenCacheProvider();
             AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => tokenProvider);
-            AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => tokenProvider.GetTokenCache());
         }
 
         [Fact]

src/Accounts/Accounts.Test/AzureSessionTestInitializer.cs

@@ -29,7 +29,6 @@ public static void Initialize()
 
             PowerShellTokenCacheProvider tokenCacheProvider = new InMemoryTokenCacheProvider();
             AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => tokenCacheProvider);
-            AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => tokenCacheProvider.GetTokenCache());
             IAuthenticatorBuilder builder = new DefaultAuthenticatorBuilder();
             AzureSession.Instance.RegisterComponent(AuthenticatorBuilder.AuthenticatorBuilderKey, () => builder);
             AzureSession.Instance.RegisterComponent(nameof(AzureCredentialFactory), () => new AzureCredentialFactory());

src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs

@@ -598,12 +598,10 @@ public void OnImport()
                 {
                     provider = new InMemoryTokenCacheProvider();
                 }
-                var tokenCache = provider.GetTokenCache();
                 IAzureEventListenerFactory azureEventListenerFactory = new AzureEventListenerFactory();
                 AzureSession.Instance.RegisterComponent(nameof(CommonUtilities), () => new CommonUtilities());
                 AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => provider);
                 AzureSession.Instance.RegisterComponent(nameof(IAzureEventListenerFactory), () => azureEventListenerFactory);
-                AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => tokenCache);
                 AzureSession.Instance.RegisterComponent(nameof(AzureCredentialFactory), () => new AzureCredentialFactory());
                 AzureSession.Instance.RegisterComponent(nameof(MsalAccessTokenAcquirerFactory), () => new MsalAccessTokenAcquirerFactory());
 #if DEBUG

src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs

@@ -74,52 +74,19 @@ void DisableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextA
             FileUtilities.DataStore = session.DataStore;
             session.ARMContextSaveMode = ContextSaveMode.Process;
 
-            PowerShellTokenCacheProvider cacheProvider;
-            MemoryStream memoryStream = null;
             if (AzureSession.Instance.TryGetComponent(
                     PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey,
                     out PowerShellTokenCacheProvider originalTokenCacheProvider))
             {
                 if(originalTokenCacheProvider is SharedTokenCacheProvider)
                 {
-                    cacheProvider = new InMemoryTokenCacheProvider();
                     var token = originalTokenCacheProvider.ReadTokenData();
-                    if (token != null && token.Length > 0)
-                    {
-                        memoryStream = new MemoryStream(token);
-                    }
-                    cacheProvider.UpdateTokenDataWithoutFlush(token);
-                    cacheProvider.FlushTokenData();
+                    //must explicitely use type PowerShellTokenCacheProvider
+                    PowerShellTokenCacheProvider cacheProvider = new InMemoryTokenCacheProvider(token);
                     AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => cacheProvider, true);
                 }
-                else
-                {
-                    cacheProvider = originalTokenCacheProvider;
-                }
-            }
-            else
-            {
-                cacheProvider = new InMemoryTokenCacheProvider();
             }
 
-            PowerShellTokenCache newTokenCache = null;
-            if(AzureSession.Instance.TryGetComponent(nameof(PowerShellTokenCache), out PowerShellTokenCache tokenCache))
-            {
-                if(!tokenCache.IsPersistentCache)
-                {
-                    newTokenCache = tokenCache;
-                }
-                else
-                {
-                    newTokenCache = memoryStream == null ? null : PowerShellTokenCache.Deserialize(memoryStream);
-                }
-            }
-
-            if(newTokenCache == null)
-            {
-                newTokenCache = cacheProvider.GetTokenCache();
-            }
-            AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => newTokenCache, true);
             if(AzureSession.Instance.TryGetComponent(AuthenticatorBuilder.AuthenticatorBuilderKey, out IAuthenticatorBuilder builder))
             {
                 builder.Reset();

src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs

@@ -88,20 +88,19 @@ void EnableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextAu
             FileUtilities.DataStore = session.DataStore;
             session.ARMContextSaveMode = ContextSaveMode.CurrentUser;
 
-            AzureSession.Instance.TryGetComponent(nameof(PowerShellTokenCache), out PowerShellTokenCache originalTokenCache);
-            var stream = new MemoryStream();
-            originalTokenCache.Serialize(stream);
-            var tokenData = stream.ToArray();
-            //must use explicit interface type PowerShellTokenCacheProvider below instead of var, otherwise could not retrieve registered component
-            PowerShellTokenCacheProvider cacheProvider = new SharedTokenCacheProvider();
-            if (tokenData != null && tokenData.Length > 0)
+            AzureSession.Instance.TryGetComponent(nameof(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey), out PowerShellTokenCacheProvider originalCacheProvider);
+            if(originalCacheProvider is InMemoryTokenCacheProvider inMemoryTokenCacheProvider)
             {
-                cacheProvider.UpdateTokenDataWithoutFlush(tokenData);
-                cacheProvider.FlushTokenData();
+                var tokenData = inMemoryTokenCacheProvider.ReadTokenData();
+                //must use explicit interface type PowerShellTokenCacheProvider below instead of var, otherwise could not retrieve registered component
+                PowerShellTokenCacheProvider newCacheProvider = new SharedTokenCacheProvider();
+                if (tokenData != null && tokenData.Length > 0)
+                {
+                    newCacheProvider.UpdateTokenDataWithoutFlush(tokenData);
+                    newCacheProvider.FlushTokenData();
+                }
+                AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => newCacheProvider, true);
             }
-            var tokenCache = cacheProvider.GetTokenCache();
-            AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => cacheProvider, true);
-            AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => tokenCache, true);
 
 
             if (writeAutoSaveFile)

src/Accounts/Accounts/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Upgraded Azure.Identity to 1.4 and MSAL lib to 4.30.1
 * Removed obsolete parameters `ManagedServiceHostName`, `ManagedServicePort` and `ManagedServiceSecret` of cmdlet `Connect-AzAccount`, environment variables `MSI_ENDPOINT` and `MSI_SECRET` could be used instead
 * Customize display format of PSAzureRmAccount to hide secret of service principal [#14208]
 * Added optional parameter `AuthScope` to `Connect-AzAccount` to support enhanced authentication of data plane features

src/Accounts/Authentication.ResourceManager/AzureRmProfileConverter.cs

@@ -14,9 +14,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.IO;
-
-using Azure.Identity;
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
@@ -79,25 +76,13 @@ public override object ReadJson(JsonReader reader, Type objectType, object exist
             else if (objectType == typeof(IAzureTokenCache) || objectType == typeof(AzureTokenCache))
             {
                 var tempResult = serializer.Deserialize<CacheBuffer>(reader);
-                if (_serializeCache && tempResult != null && tempResult.CacheData != null && tempResult.CacheData.Length > 0)
+                if (_serializeCache && tempResult?.CacheData?.Length > 0)
                 {
-                    if(AzureSession.Instance.TryGetComponent(nameof(PowerShellTokenCache), out PowerShellTokenCache oldTokenCache))
+                    if (AzureSession.Instance.TryGetComponent(
+                        PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey,
+                        out PowerShellTokenCacheProvider tokenCacheProvider))
                     {
-                        if(!oldTokenCache.IsPersistentCache)
-                        {
-                            var stream = new MemoryStream(tempResult.CacheData);
-                            var tokenCache = new PowerShellTokenCache(stream);
-                            AzureSession.Instance.RegisterComponent(nameof(PowerShellTokenCache), () => tokenCache, true);
-                        }
-                        else
-                        {
-                            if (AzureSession.Instance.TryGetComponent(
-                                PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey,
-                                out PowerShellTokenCacheProvider tokenCacheProvider))
-                            {
-                                tokenCacheProvider.UpdateTokenDataWithoutFlush(tempResult.CacheData);
-                            }
-                        }
+                        tokenCacheProvider.UpdateTokenDataWithoutFlush(tempResult.CacheData);
                     }
                 }
                 // cache data is not for direct use, so we do not return anything
@@ -138,28 +123,12 @@ public override void WriteJson(JsonWriter writer, object value, JsonSerializer s
                 {
                     byte[] cacheData = null;
 
-
-                    if (AzureSession.Instance.TryGetComponent(nameof(PowerShellTokenCache), out PowerShellTokenCache tokenCache))
+                    if (AzureSession.Instance.TryGetComponent(
+                            PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey,
+                            out PowerShellTokenCacheProvider tokenCacheProvider))
                     {
-                        if (tokenCache.IsPersistentCache)
-                        {
-                            if (AzureSession.Instance.TryGetComponent(
-                                PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey,
-                                out PowerShellTokenCacheProvider tokenCacheProvider))
-                            {
-                                cacheData = tokenCacheProvider.ReadTokenData();
-                            }
-                        }
-                        else
-                        {
-                            using (var stream = new MemoryStream())
-                            {
-                                tokenCache.Serialize(stream);
-                                cacheData = stream.ToArray();
-                            }
-                        }
+                        cacheData = tokenCacheProvider.ReadTokenData();
                     }
-
                     value = new CacheBuffer { CacheData = cacheData };
                 }
                 else

src/Accounts/Authentication/Authentication.csproj

@@ -12,7 +12,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Azure.Identity" Version="1.4.0-beta.3" />
+    <PackageReference Include="Azure.Identity" Version="1.4.0" />
   </ItemGroup>
 
   <ItemGroup>

src/Accounts/Authentication/Authentication/TokenCache/InMemoryTokenCacheOptions.cs

@@ -0,0 +1,93 @@
+// ----------------------------------------------------------------------------------
+//
+// 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.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+using Azure.Identity;
+
+namespace Microsoft.Azure.Commands.Common.Authentication
+{
+    public class InMemoryTokenCacheOptions : UnsafeTokenCacheOptions
+    {
+        internal ReadOnlyMemory<byte> CachedToken { get; private set; }
+        private ReaderWriterLockSlim readerWriterLockSlim = new ReaderWriterLockSlim();
+
+        public InMemoryTokenCacheOptions()
+            : this(new ReadOnlyMemory<byte>())
+        {
+        }
+
+        public InMemoryTokenCacheOptions(ReadOnlyMemory<byte> token)
+        {
+            CachedToken = token;
+        }
+
+        protected override async Task<ReadOnlyMemory<byte>> RefreshCacheAsync()
+        {
+            readerWriterLockSlim.EnterReadLock();
+            try
+            {
+                return await Task.FromResult(CachedToken);
+            }
+            finally
+            {
+                readerWriterLockSlim.ExitReadLock();
+            }
+        }
+
+        protected override Task TokenCacheUpdatedAsync(TokenCacheUpdatedArgs tokenCacheUpdatedArgs)
+        {
+            readerWriterLockSlim.EnterWriteLock();
+            try
+            {
+                CachedToken = tokenCacheUpdatedArgs.UnsafeCacheData;
+            }
+            finally
+            {
+                readerWriterLockSlim.ExitWriteLock();
+            }
+            return Task.CompletedTask;
+        }
+
+        public void Serialize(Stream stream)
+        {
+            readerWriterLockSlim.EnterReadLock();
+            try
+            {
+                if (CachedToken.Length > 0)
+                {
+                    var bytes = CachedToken.ToArray();
+                    stream.Write(bytes, 0, bytes.Length);
+                }
+            }
+            finally
+            {
+                readerWriterLockSlim.ExitReadLock();
+            }
+        }
+
+        public static InMemoryTokenCacheOptions Deserialize(Stream stream)
+        {
+            using (MemoryStream memoryStream = new MemoryStream())
+            {
+                stream.CopyTo(memoryStream);
+                var token = memoryStream.ToArray();
+                return new InMemoryTokenCacheOptions(token);
+            }
+        }
+    }
+}

src/Accounts/Authentication/Authentication/TokenCache/InMemoryTokenCacheProvider.cs

@@ -12,39 +12,52 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
-using System.Threading.Tasks;
+using Azure.Identity;
 
 using Microsoft.Identity.Client;
 
 namespace Microsoft.Azure.Commands.Common.Authentication
 {
     public class InMemoryTokenCacheProvider : PowerShellTokenCacheProvider
     {
-        public InMemoryTokenCacheProvider()
+        private InMemoryTokenCacheOptions InMemoryTokenCacheOptions { get; set; }
+
+        public InMemoryTokenCacheProvider(byte[] tokenCache)
         {
+            InMemoryTokenCacheOptions = new InMemoryTokenCacheOptions(tokenCache);
+        }
+
+        public InMemoryTokenCacheProvider(InMemoryTokenCacheOptions options = null)
+        {
+            InMemoryTokenCacheOptions = options ?? new InMemoryTokenCacheOptions();
         }
 
         public override byte[] ReadTokenData()
         {
-            return null;
+            return InMemoryTokenCacheOptions.CachedToken.ToArray();
         }
 
         public override void FlushTokenData()
         {
+            if (_tokenCacheDataToFlush != null)
+            {
+                InMemoryTokenCacheOptions = new InMemoryTokenCacheOptions(_tokenCacheDataToFlush);
+                _tokenCacheDataToFlush = null;
+            }
         }
 
         public override void ClearCache()
         {
+            InMemoryTokenCacheOptions = new InMemoryTokenCacheOptions();
         }
 
         protected override void RegisterCache(IPublicClientApplication client)
         {
-
         }
 
-        public override PowerShellTokenCache GetTokenCache()
+        public override TokenCachePersistenceOptions GetTokenCachePersistenceOptions()
         {
-            return new PowerShellTokenCache(new global::Azure.Identity.TokenCache());
+            return InMemoryTokenCacheOptions;
         }
     }
 }