fix issue: fails to migrate ADAL token for MSA account (#13292)

Author: erich-wang

src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs

@@ -21,6 +21,7 @@
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Core;
 using Microsoft.Azure.Commands.Common.Authentication.Factories;
 using Microsoft.Azure.Commands.Common.Authentication.Models;
 using Microsoft.Azure.Commands.Profile.Common;
@@ -463,6 +464,12 @@ private void SetContextWithOverwritePrompt(Action<AzureRmProfile, RMProfileClien
             }
         }
 
+        private static IAzureContextContainer GetAzureContextContainer()
+        {
+            var provider = new ProtectedProfileProvider();
+            return provider.Profile;
+        }
+
         /// <summary>
         /// Load global aliases for ARM
         /// </summary>
@@ -472,7 +479,8 @@ public void OnImport()
             try
             {
 #endif
-                 AzureSessionInitializer.InitializeAzureSession();
+                AzureSessionInitializer.InitializeAzureSession();
+                AzureSessionInitializer.MigrateAdalCache(AzureSession.Instance, GetAzureContextContainer);
 #if DEBUG
                 if (!TestMockSupport.RunningMocked)
                 {

src/Accounts/Authentication/Authentication/TokenCache/AdalTokenMigrator.cs

@@ -16,6 +16,7 @@
 using System.Linq;
 
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Core;
 using Microsoft.Identity.Client;
 using Microsoft.Identity.Client.Extensions.Msal;
 
@@ -24,13 +25,18 @@ namespace Microsoft.Azure.Commands.Common.Authentication.Authentication.TokenCac
     public class AdalTokenMigrator
     {
         protected const string PowerShellClientId = "1950a258-227b-4e31-a9cf-717495945fc2";
+        private const string TenantsString = "Tenants";
+
         private byte[] AdalToken { get; set; }
 
         private bool HasRegistered { get; set; }
 
-        public AdalTokenMigrator(byte[] adalToken)
+        private Lazy<IAzureContextContainer> ContextContainerInitializer { get; set; }
+
+        public AdalTokenMigrator(byte[] adalToken, Func<IAzureContextContainer> getContextContainer)
         {
             AdalToken = adalToken;
+            ContextContainerInitializer = new Lazy<IAzureContextContainer>(getContextContainer);
         }
 
         public void MigrateFromAdalToMsal()
@@ -85,7 +91,19 @@ public void MigrateFromAdalToMsal()
                     }
 
                     var scopes = new string[] { string.Format("{0}{1}", environment.ActiveDirectoryServiceEndpointResourceId, ".default") };
-                    var token = clientApplication.AcquireTokenSilent(scopes, account).ExecuteAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+
+                    try
+                    {
+                        clientApplication.AcquireTokenSilent(scopes, account).ExecuteAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+                    }
+                    catch //For MSA account, real AAD tenant must be specified, otherwise MSAL library will request token against its home tenant
+                    {
+                        var tenantId = GetTenantId(account.Username);
+                        if(!string.IsNullOrEmpty(tenantId))
+                        {
+                            clientApplication.AcquireTokenSilent(scopes, account).WithAuthority(environment.ActiveDirectoryAuthority, tenantId).ExecuteAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+                        }
+                    }
                     //TODO: Set HomeAccountId for migration
                 }
                 catch
@@ -96,5 +114,20 @@ public void MigrateFromAdalToMsal()
             }
             cacheHelper?.UnregisterCache(clientApplication.UserTokenCache);
         }
+
+        private string GetTenantId(string accountId)
+        {
+            var contextContainer = ContextContainerInitializer.Value;
+            string tenantId = null;
+            if (contextContainer != null)
+            {
+                var matchedAccount = contextContainer?.Accounts?.FirstOrDefault(account => string.Equals(account.Id, accountId, StringComparison.InvariantCultureIgnoreCase));
+                if (matchedAccount != null && matchedAccount.ExtendedProperties.ContainsKey(TenantsString))
+                {
+                    tenantId = matchedAccount.ExtendedProperties[TenantsString]?.Split(',')?.FirstOrDefault();
+                }
+            }
+            return tenantId;
+        }
     }
 }

src/Accounts/Authentication/AzureSessionInitializer.cs

@@ -20,6 +20,7 @@
 using Hyak.Common;
 
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Core;
 using Microsoft.Azure.Commands.Common.Authentication.Authentication.TokenCache;
 using Microsoft.Azure.Commands.Common.Authentication.Factories;
 using Microsoft.Azure.Commands.Common.Authentication.Properties;
@@ -98,14 +99,17 @@ static bool MigrateSettings(IDataStore store, string oldProfileDirectory, string
             return false;
         }
 
-        static void MigrateAdalCache(AzureSession session, IDataStore store, string adalCachePath, string msalCachePath)
+        public static void MigrateAdalCache(IAzureSession session, Func<IAzureContextContainer> getContextContainer)
         {
             if (session.ARMContextSaveMode == ContextSaveMode.Process)
             {
                 // Don't attempt to migrate if context autosave is disabled
                 return;
             }
 
+            var adalCachePath = Path.Combine(session.ProfileDirectory, "TokenCache.dat");
+            var msalCachePath = Path.Combine(session.TokenCacheDirectory, "msal.cache");
+            var store = session.DataStore;
             if (!store.FileExists(adalCachePath) || store.FileExists(msalCachePath))
             {
                 // Return if
@@ -127,7 +131,7 @@ static void MigrateAdalCache(AzureSession session, IDataStore store, string adal
 
             if(adalData != null && adalData.Length > 0)
             {
-                new AdalTokenMigrator(adalData).MigrateFromAdalToMsal();
+                new AdalTokenMigrator(adalData, getContextContainer).MigrateFromAdalToMsal();
             }
         }
 
@@ -232,7 +236,7 @@ static IAzureSession CreateInstance(IDataStore dataStore = null)
             session.ARMProfileFile = autoSave.ContextFile;
             session.TokenCacheDirectory = autoSave.CacheDirectory;
             session.TokenCacheFile = autoSave.CacheFile;
-            MigrateAdalCache(session, dataStore, oldCachePath, Path.Combine(cachePath, "msal.cache"));
+
             InitializeDataCollection(session);
             session.RegisterComponent(HttpClientOperationsFactory.Name, () => HttpClientOperationsFactory.Create());
             session.TokenCache = session.TokenCache ?? new AzureTokenCache();

src/Accounts/Authentication/Utilities/CustomAssemblyResolver.cs

@@ -55,7 +55,8 @@ public static Assembly CurrentDomain_AssemblyResolve(object sender, ResolveEvent
                 AssemblyName name = new AssemblyName(args.Name);
                 if (NetFxPreloadAssemblies.TryGetValue(name.Name, out Version version))
                 {
-                    if (version >= name.Version && version.Major == name.Version.Major)
+                    //For Newtonsoft.Json, allow to use bigger version to replace smaller version
+                    if (version >= name.Version && (version.Major == name.Version.Major || string.Equals(name.Name, "Newtonsoft.Json", StringComparison.OrdinalIgnoreCase)))
                     {
                         string requiredAssembly = Path.Combine(PreloadAssemblyFolder, $"{name.Name}.dll");
                         return Assembly.LoadFrom(requiredAssembly);