Fix the bugs for AzKeyStore (Azure#20768)

* Fix AzKeyStore issues and optimize initialization and update * Enabled credential to be found only by applicationId while tenant was not matched when accquire token. [Azure#20484] * When Az.Accounts ran in parallel, the waiters were allowed to wait infinitely to avoid throw exception in automation enviroment. [Azure#20455] * Used Lazy load for AzKeyStore. * Used update on change mechanism for AzKeyStore and remove `Flush` interface. * Add parallel test case of Az.Accounts to smoke test * Address review comments Co-authored-by: Yeming Liu <11371776+isra-fel@users.noreply.github.com> Address review comments * Fix failed test cases. * Fix an issue of AzKeyStore when context autosaving switching
Xynoclafe · Feb 3, 2023 · 1115b94 · 1115b94
1 parent bc78f20
commit 1115b94
Show file tree

Hide file tree

Showing 28 changed files with 1,057 additions and 360 deletions.
diff --git a/.azure-pipelines/util/smoke-test-steps.yml b/.azure-pipelines/util/smoke-test-steps.yml
@@ -78,6 +78,13 @@ jobs:
        Write-Host "List artifacts..."
        Get-ChildItem "$(Pipeline.Workspace)\\LocalRepo\\"
 
+  - task: NuGetCommand@2
+    condition: and(succeeded(), eq('${{ parameters.psVersion }}', '5.1.14'))
+    displayName: 'Download ThreadJob .nupkg File for PowerShell 5.1.14'
+    inputs:
+      command: custom
+      arguments: 'install ThreadJob -directdownload -packagesavemode nupkg -source https://www.powershellgallery.com/api/v2 -OutputDirectory packages'
+
   - task: NuGetCommand@2
     condition: and(succeeded(), eq(variables['GalleryName'], 'LocalRepo'))
     displayName: 'Download Previous Az .nupkg Files'

diff --git a/src/Accounts/Accounts.Test/AutosaveTests.cs b/src/Accounts/Accounts.Test/AutosaveTests.cs
@@ -51,9 +51,7 @@ private AzKeyStore SetMockedAzKeyStore()
             storageMocker.Setup(f => f.Create()).Returns(storageMocker.Object);
             storageMocker.Setup(f => f.ReadData()).Returns(new byte[0]);
             storageMocker.Setup(f => f.WriteData(It.IsAny<byte[]>())).Callback((byte[] s) => {});
-            var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", false, false, storageMocker.Object);
-            AzKeyStore.RegisterJsonConverter(typeof(ServicePrincipalKey), typeof(ServicePrincipalKey).Name);
-            AzKeyStore.RegisterJsonConverter(typeof(SecureString), typeof(SecureString).Name, new SecureStringConverter());
+            var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", true, storageMocker.Object);
             return keyStore;
         }
 

diff --git a/src/Accounts/Accounts.Test/ContextCmdletTests.cs b/src/Accounts/Accounts.Test/ContextCmdletTests.cs
@@ -23,17 +23,21 @@
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 using Microsoft.WindowsAzure.Commands.Test.Utilities.Common;
 using Microsoft.WindowsAzure.Commands.Utilities.Common;
-using Xunit;
-using Xunit.Abstractions;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
-using System;
+using Microsoft.Azure.Commands.Common.Authentication.Properties;
 using Microsoft.Azure.Commands.Profile.Context;
-using System.Linq;
 using Microsoft.Azure.Commands.Common.Authentication.ResourceManager;
 using Microsoft.Azure.Commands.Profile.Common;
 using Microsoft.Azure.Commands.ScenarioTest.Mocks;
 using Microsoft.Azure.Commands.TestFx.Mocks;
 using Microsoft.Azure.Commands.TestFx;
+using Microsoft.Azure.Commands.ResourceManager.Common;
+using Moq;
+using System;
+using System.IO;
+using System.Linq;
+using Xunit;
+using Xunit.Abstractions;
 
 namespace Microsoft.Azure.Commands.Profile.Test
 {
@@ -56,6 +60,12 @@ public ContextCmdletTests(ITestOutputHelper output)
             tokenCacheProviderMock = new MockPowerShellTokenCacheProvider();
             AzureSession.Instance.RegisterComponent<PowerShellTokenCacheProvider>(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => tokenCacheProviderMock);
             Environment.SetEnvironmentVariable("Azure_PS_Data_Collection", "True");
+
+            Mock<IStorage> storageMocker = new Mock<IStorage>();
+            AzKeyStore azKeyStore = null;
+            string profilePath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), Resources.AzureDirectoryName);
+            azKeyStore = new AzKeyStore(profilePath, AzureSession.Instance.KeyStoreFile, true, storageMocker.Object);
+            AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => azKeyStore, true);
         }
 
         [Fact]

diff --git a/src/Accounts/Accounts.Test/ProfileCmdletTests.cs b/src/Accounts/Accounts.Test/ProfileCmdletTests.cs
@@ -55,7 +55,7 @@ private AzKeyStore SetMockedAzKeyStore()
             storageMocker.Setup(f => f.Create()).Returns(storageMocker.Object);
             storageMocker.Setup(f => f.ReadData()).Returns(new byte[0]);
             storageMocker.Setup(f => f.WriteData(It.IsAny<byte[]>())).Callback((byte[] s) => { });
-            var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", false, false, storageMocker.Object);
+            var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", false, storageMocker.Object);
             return keyStore;
         }
 

diff --git a/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs b/src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
@@ -425,7 +425,7 @@ public override void ExecuteCmdlet()
                 azureAccount.SetProperty(AzureAccount.Property.CertificatePath, resolvedPath);
                 if (CertificatePassword != null)
                 {
-                    keyStore?.SaveKey(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, azureAccount.Id, Tenant), CertificatePassword);
+                    keyStore?.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, azureAccount.Id, Tenant), CertificatePassword);
                     if (GetContextModificationScope() == ContextModificationScope.CurrentUser && !keyStore.IsProtected)
                     {
                         WriteWarning(string.Format(Resources.ServicePrincipalWarning, AzureSession.Instance.KeyStoreFile, AzureSession.Instance.ARMProfileDirectory));
@@ -451,7 +451,7 @@ public override void ExecuteCmdlet()
 
             if (azureAccount.Type == AzureAccount.AccountType.ServicePrincipal && password != null)
             {
-                keyStore?.SaveKey(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret
+                keyStore?.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret
                     ,azureAccount.Id, Tenant), password);
                 if (GetContextModificationScope() == ContextModificationScope.CurrentUser && !keyStore.IsProtected)
                 {
@@ -713,9 +713,7 @@ public void OnImport()
                 }
 
                 AzKeyStore keyStore = null;
-                keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, AzureSession.Instance.KeyStoreFile, false, autoSaveEnabled);
-                AzKeyStore.RegisterJsonConverter(typeof(ServicePrincipalKey), typeof(ServicePrincipalKey).Name);
-                AzKeyStore.RegisterJsonConverter(typeof(SecureString), typeof(SecureString).Name, new SecureStringConverter());
+                keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, AzureSession.Instance.KeyStoreFile, autoSaveEnabled);
                 AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => keyStore);
 
                 if (!InitializeProfileProvider(autoSaveEnabled))
@@ -724,11 +722,6 @@ public void OnImport()
                     autoSaveEnabled = false;
                 }
 
-                if (!keyStore.LoadStorage())
-                {
-                    WriteInitializationWarnings(Resources.KeyStoreLoadingError);
-                }
-
                 IAuthenticatorBuilder builder = null;
                 if (!AzureSession.Instance.TryGetComponent(AuthenticatorBuilder.AuthenticatorBuilderKey, out builder))
                 {

diff --git a/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs b/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs
@@ -94,7 +94,7 @@ void DisableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextA
 
             if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keystore))
             {
-                keystore.DisableAutoSaving();
+                keystore.DisableSyncToStorage();
             }
 
             if (writeAutoSaveFile)

diff --git a/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs b/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs
@@ -104,11 +104,9 @@ void EnableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextAu
 
             if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keystore))
             {
-                keystore.Flush();
-                keystore.DisableAutoSaving();
+                keystore.EnableSyncToStorage();
             }
 
-
             if (writeAutoSaveFile)
             {
                 try

diff --git a/src/Accounts/Accounts/ChangeLog.md b/src/Accounts/Accounts/ChangeLog.md
@@ -20,6 +20,8 @@
 
 ## Upcoming Release
 * Supported Web Account Manager on ARM64-based Windows systems. Fixed an issue where `Connect-AzAccount` failed with error "Unable to load DLL 'msalruntime_arm64'". [#20700]
+* Enabled credential to be found only by applicationId while tenant was not matched when accquire token. [#20484]
+* When Az.Accounts ran in parallel, the waiters were allowed to wait infinitely to avoid throw exception in automation enviroment. [#20455]
 
 ## Version 2.11.1
 * Fixed an issue where Az.Accounts cannot be imported correctly. [#20615]

diff --git a/src/Accounts/Accounts/Context/ImportAzureRMContext.cs b/src/Accounts/Accounts/Context/ImportAzureRMContext.cs
@@ -78,13 +78,13 @@ void CopyProfile(AzureRmProfile source, IProfileOperations target)
                         var secret = account.GetProperty(AzureAccount.Property.ServicePrincipalSecret);
                         if (!string.IsNullOrEmpty(secret))
                         {
-                            keyStore.SaveKey(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, context.Value.Tenant?.Id)
+                            keyStore.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, context.Value.Tenant?.Id)
                                 , secret.ConvertToSecureString());
                         }
                         var password = account.GetProperty(AzureAccount.Property.CertificatePassword);
                         if (!string.IsNullOrEmpty(password))
                         {
-                            keyStore.SaveKey(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, context.Value.Tenant?.Id)
+                            keyStore.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, context.Value.Tenant?.Id)
                                 ,password.ConvertToSecureString());
                         }
                     }

diff --git a/src/Accounts/Accounts/Context/SetAzureRMContext.cs b/src/Accounts/Accounts/Context/SetAzureRMContext.cs
@@ -97,13 +97,13 @@ public override void ExecuteCmdlet()
                                     var secret = account.GetProperty(AzureAccount.Property.ServicePrincipalSecret);
                                     if (!string.IsNullOrEmpty(secret))
                                     {
-                                        keyStore.SaveKey(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, Context.Tenant?.Id)
+                                        keyStore.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, Context.Tenant?.Id)
                                             , secret.ConvertToSecureString());
                                     }
                                     var password = account.GetProperty(AzureAccount.Property.CertificatePassword);
                                     if (!string.IsNullOrEmpty(password))
                                     {
-                                        keyStore.SaveKey(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, Context.Tenant?.Id)
+                                        keyStore.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, Context.Tenant?.Id)
                                             , password.ConvertToSecureString());
                                     }
                                 }

diff --git a/src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs b/src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs
@@ -225,13 +225,13 @@ private IAzureContext MigrateSecretToKeyStore(IAzureContext context, AzKeyStore
                 var account = context.Account;
                 if (account.IsPropertySet(AzureAccount.Property.ServicePrincipalSecret))
                 {
-                    keystore?.SaveKey(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, account.GetTenants().First())
+                    keystore?.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, account.GetTenants().First())
                         , account.ExtendedProperties.GetProperty(AzureAccount.Property.ServicePrincipalSecret).ConvertToSecureString());
                     account.ExtendedProperties.Remove(AzureAccount.Property.ServicePrincipalSecret);
                 }
                 if (account.IsPropertySet(AzureAccount.Property.CertificatePassword))
                 {
-                    keystore?.SaveKey(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, account.GetTenants().First())
+                    keystore?.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, account.GetTenants().First())
     , account.ExtendedProperties.GetProperty(AzureAccount.Property.CertificatePassword).ConvertToSecureString());
                     account.ExtendedProperties.Remove(AzureAccount.Property.CertificatePassword);
                 }
@@ -336,10 +336,6 @@ public void Save(IFileProvider provider, bool serializeCache = true)
                     // so that previous data is overwritten
                     provider.Stream.SetLength(provider.Stream.Position);
                 }
-
-                AzKeyStore keystore = null;
-                AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out keystore);
-                keystore?.Flush();
             }
             finally
             {