From 50b10d6ba38d261c542176f4fb1387a0740512c2 Mon Sep 17 00:00:00 2001 From: Jin Lei <54836179+msJinLei@users.noreply.github.com> Date: Thu, 2 Feb 2023 16:23:37 +0800 Subject: [PATCH] Address review comments Co-authored-by: Yeming Liu <11371776+isra-fel@users.noreply.github.com> Address review comments --- src/Accounts/Accounts.Test/AutosaveTests.cs | 3 +- .../Accounts.Test/ContextCmdletTests.cs | 2 +- .../Accounts.Test/ProfileCmdletTests.cs | 2 +- .../Accounts/Account/ConnectAzureRmAccount.cs | 6 +- .../AutoSave/DisableAzureRmContextAutosave.cs | 5 + .../AutoSave/EnableAzureRmContextAutosave.cs | 5 + src/Accounts/Accounts/ChangeLog.md | 2 - .../Accounts/Context/ImportAzureRMContext.cs | 4 +- .../Accounts/Context/SetAzureRMContext.cs | 4 +- .../AzureRmProfile.cs | 4 +- .../Authentication.Test/AzKeyStorageTest.cs | 26 +-- .../Authentication.Test/StorageHelperTest.cs | 10 +- .../Factories/AuthenticationFactory.cs | 8 +- .../Identity/AsyncLockWithValue.cs | 18 -- .../KeyStore/AsyncLockWithValue.cs | 212 ++++++++++++++++++ .../Authentication/KeyStore/AzKeyStore.cs | 58 +++-- .../KeyStore/{IKeyStore.cs => IKeyCache.cs} | 2 +- .../Authentication/KeyStore/IKeyStoreKey.cs | 20 +- .../Authentication/KeyStore/IStorage.cs | 1 - .../Authentication/KeyStore/IStorageHelper.cs | 9 +- ...nMemoryKeyStore.cs => InMemoryKeyCache.cs} | 12 +- .../KeyStore/KeyStoreNotificationArgs.cs | 2 +- .../Authentication/KeyStore/StorageHelper.cs | 59 +---- .../Authentication/KeyStore/StorageWrapper.cs | 2 +- tools/Test/SmokeTest/InstallAzModules.ps1 | 3 +- tools/Test/SmokeTest/RmCoreSmokeTests.ps1 | 8 +- 26 files changed, 336 insertions(+), 151 deletions(-) create mode 100644 src/Accounts/Authentication/KeyStore/AsyncLockWithValue.cs rename src/Accounts/Authentication/KeyStore/{IKeyStore.cs => IKeyCache.cs} (97%) rename src/Accounts/Authentication/KeyStore/{InMemoryKeyStore.cs => InMemoryKeyCache.cs} (97%) diff --git a/src/Accounts/Accounts.Test/AutosaveTests.cs b/src/Accounts/Accounts.Test/AutosaveTests.cs index 50ea60b43fe5..0986b9808436 100644 --- a/src/Accounts/Accounts.Test/AutosaveTests.cs +++ b/src/Accounts/Accounts.Test/AutosaveTests.cs @@ -41,7 +41,6 @@ public AutosaveTests(ITestOutputHelper output) XunitTracingInterceptor.AddToContext(new XunitTracingInterceptor(output)); commandRuntimeMock = new MockCommandRuntime(); dataStore = new MemoryDataStore(); - ResetState(); keyStore = SetMockedAzKeyStore(); } @@ -51,7 +50,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())).Callback((byte[] s) => {}); - var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", storageMocker.Object); + 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 index 7b977a5a7482..9e4d44bb4cb5 100644 --- a/src/Accounts/Accounts.Test/ContextCmdletTests.cs +++ b/src/Accounts/Accounts.Test/ContextCmdletTests.cs @@ -64,7 +64,7 @@ public ContextCmdletTests(ITestOutputHelper output) Mock storageMocker = new Mock(); AzKeyStore azKeyStore = null; string profilePath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), Resources.AzureDirectoryName); - azKeyStore = new AzKeyStore(profilePath, AzureSession.Instance.KeyStoreFile, storageMocker.Object); + azKeyStore = new AzKeyStore(profilePath, AzureSession.Instance.KeyStoreFile, true, storageMocker.Object); AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => azKeyStore, true); } diff --git a/src/Accounts/Accounts.Test/ProfileCmdletTests.cs b/src/Accounts/Accounts.Test/ProfileCmdletTests.cs index 386998a78ffd..fdd16bf07357 100644 --- 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())).Callback((byte[] s) => { }); - var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", 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 index 218fc28abe8f..2e8208f0e560 100644 --- 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?.SaveCredential(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?.SaveCredential(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret + keyStore?.SaveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret ,azureAccount.Id, Tenant), password); if (GetContextModificationScope() == ContextModificationScope.CurrentUser && !keyStore.IsProtected) { @@ -713,7 +713,7 @@ public void OnImport() } AzKeyStore keyStore = null; - keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, AzureSession.Instance.KeyStoreFile); + keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, AzureSession.Instance.KeyStoreFile, autoSaveEnabled); AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => keyStore); if (!InitializeProfileProvider(autoSaveEnabled)) diff --git a/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs b/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs index 1d69aba78ce5..a2b9be38aa7a 100644 --- a/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs +++ b/src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs @@ -92,6 +92,11 @@ void DisableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextA builder.Reset(); } + if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keystore)) + { + keystore.DisableSyncToStorage(); + } + if (writeAutoSaveFile) { FileUtilities.EnsureDirectoryExists(session.ProfileDirectory); diff --git a/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs b/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs index 66325469ea71..830876d6463b 100644 --- a/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs +++ b/src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs @@ -102,6 +102,11 @@ void EnableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextAu AzureSession.Instance.RegisterComponent(PowerShellTokenCacheProvider.PowerShellTokenCacheProviderKey, () => newCacheProvider, true); } + if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keystore)) + { + keystore.EnableSyncToStorage(); + } + if (writeAutoSaveFile) { try diff --git a/src/Accounts/Accounts/ChangeLog.md b/src/Accounts/Accounts/ChangeLog.md index ebd8bc1d7230..0de24d71e9d9 100644 --- a/src/Accounts/Accounts/ChangeLog.md +++ b/src/Accounts/Accounts/ChangeLog.md @@ -22,8 +22,6 @@ * 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] -* Used Lazy load for AzKeyStore. -* Used update on change mechanism for AzKeyStore and remove `Flush` interface. ## 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 index 9304d88d13ba..f75d2e813fc5 100644 --- 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.SaveCredential(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.SaveCredential(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 index fab411685fea..9c0badcf8470 100644 --- 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.SaveCredential(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.SaveCredential(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 index 9c81702be477..a1a53393c7b2 100644 --- 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?.SaveCredential(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?.SaveCredential(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); } diff --git a/src/Accounts/Authentication.Test/AzKeyStorageTest.cs b/src/Accounts/Authentication.Test/AzKeyStorageTest.cs index 3e6165a547b3..ed4fb6b7f38b 100644 --- a/src/Accounts/Authentication.Test/AzKeyStorageTest.cs +++ b/src/Accounts/Authentication.Test/AzKeyStorageTest.cs @@ -56,15 +56,15 @@ private static bool CompareJsonObjects(string expected, string acutal) [Trait(Category.AcceptanceType, Category.CheckIn)] public void SaveKey() { - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("ServicePrincipalSecret", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); var secret = "secret".ConvertToSecureString(); - store.SaveCredential(servicePrincipalKey, secret); + store.SaveSecureString(servicePrincipalKey, secret); IKeyStoreKey certificatePassword = new ServicePrincipalKey("CertificatePassword", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); var passowrd = "password".ConvertToSecureString(); - store.SaveCredential(certificatePassword, passowrd); + store.SaveSecureString(certificatePassword, passowrd); var result = Encoding.UTF8.GetString(storageChecker.ToArray()); const string EXPECTEDSTRING = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""CertificatePassword\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""password\""""},{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secret\""""}]"; @@ -81,12 +81,12 @@ public void FindKey() { const string EXPECTED = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secret\""""}]"; storageChecker.AddRange(Encoding.UTF8.GetBytes(EXPECTED)); - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { storageMocker.Setup(f => f.ReadData()).Returns(storageChecker.ToArray()); IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("ServicePrincipalSecret", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); - var secret = store.GetCredential(servicePrincipalKey); + var secret = store.GetSecureString(servicePrincipalKey); Assert.Equal("secret", secret.ConvertToString()); store.Clear(); @@ -100,12 +100,12 @@ public void FindFallbackKey() { const string EXPECTED = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secretFallback\""""}]"; storageChecker.AddRange(Encoding.UTF8.GetBytes(EXPECTED)); - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { storageMocker.Setup(f => f.ReadData()).Returns(storageChecker.ToArray()); IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("ServicePrincipalSecret", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-0000-bad9-b7b93a3e9c5a"); - var secret = store.GetCredential(servicePrincipalKey); + var secret = store.GetSecureString(servicePrincipalKey); Assert.Equal("secretFallback", secret.ConvertToString()); store.Clear(); @@ -120,12 +120,12 @@ public void FindNoKey() { const string EXPECTED = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secret\""""}]"; storageChecker.AddRange(Encoding.UTF8.GetBytes(EXPECTED)); - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { storageMocker.Setup(f => f.ReadData()).Returns(storageChecker.ToArray()); IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("CertificatePassword", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); - Assert.Throws(() => store.GetCredential(servicePrincipalKey)); + Assert.Throws(() => store.GetSecureString(servicePrincipalKey)); store.Clear(); } @@ -138,12 +138,12 @@ public void RemoveKey() { const string EXPECTED = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secret\""""}]"; storageChecker.AddRange(Encoding.UTF8.GetBytes(EXPECTED)); - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { storageMocker.Setup(f => f.ReadData()).Returns(storageChecker.ToArray()); IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("ServicePrincipalSecret", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); - store.RemoveCredential(servicePrincipalKey); + store.RemoveSecureString(servicePrincipalKey); var result = Encoding.UTF8.GetString(storageChecker.ToArray()); var objects = JsonConvert.DeserializeObject>(result); @@ -160,12 +160,12 @@ public void RemoveNoKey() { const string EXPECTED = @"[{""keyType"":""ServicePrincipalKey"",""keyStoreKey"":""{\""appId\"":\""6c984d31-5b4f-4734-b548-e230a248e347\"",\""tenantId\"":\""54826b22-38d6-4fb2-bad9-b7b93a3e9c5a\"",\""name\"":\""ServicePrincipalSecret\""}"",""valueType"":""SecureString"",""keyStoreValue"":""\""secret\""""}]"; storageChecker.AddRange(Encoding.UTF8.GetBytes(EXPECTED)); - using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, storageMocker.Object)) + using (var store = new AzKeyStore(dummpyPath, keyStoreFileName, true, storageMocker.Object)) { storageMocker.Setup(f => f.ReadData()).Returns(storageChecker.ToArray()); IKeyStoreKey servicePrincipalKey = new ServicePrincipalKey("CertificatePassword", "6c984d31-5b4f-4734-b548-e230a248e347", "54826b22-38d6-4fb2-bad9-b7b93a3e9c5a"); - store.RemoveCredential(servicePrincipalKey); + store.RemoveSecureString(servicePrincipalKey); var result = Encoding.UTF8.GetString(storageChecker.ToArray()); var objects = JsonConvert.DeserializeObject>(result); diff --git a/src/Accounts/Authentication.Test/StorageHelperTest.cs b/src/Accounts/Authentication.Test/StorageHelperTest.cs index a261cae27470..7cde8dd554e6 100644 --- a/src/Accounts/Authentication.Test/StorageHelperTest.cs +++ b/src/Accounts/Authentication.Test/StorageHelperTest.cs @@ -27,7 +27,7 @@ namespace Common.Authenticators.Test public class StorageHelperTest { private Mock storageMocker = null; - private Mock keystoreMocker = null; + private Mock keystoreMocker = null; private string profilePath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), Resources.AzureDirectoryName); private string keyStoreFileName = "azkeystore"; @@ -35,7 +35,7 @@ public StorageHelperTest() { storageMocker = new Mock(); storageMocker.Setup(f => f.Create()).Returns(storageMocker.Object); - keystoreMocker = new Mock(); + keystoreMocker = new Mock(); } [Fact] @@ -72,11 +72,7 @@ public void SaveToStorageTest() var helper = StorageHelper.GetStorageHelperAsync(true, keyStoreFileName, profilePath , keystoreMocker.Object, storageMocker.Object).GetAwaiter().GetResult(); - var args = new KeyStoreNotificationArgs() - { - KeyStore = keystoreMocker.Object - }; - helper.WriteToCachedStorage(args); + helper.WriteToCachedStorage(keystoreMocker.Object); string actual = Encoding.UTF8.GetString(storageChecker.ToArray()); Assert.Equal(EXPECTED, actual); diff --git a/src/Accounts/Authentication/Factories/AuthenticationFactory.cs b/src/Accounts/Authentication/Factories/AuthenticationFactory.cs index 830d0520d6b4..34772928a5c3 100644 --- a/src/Accounts/Authentication/Factories/AuthenticationFactory.cs +++ b/src/Accounts/Authentication/Factories/AuthenticationFactory.cs @@ -433,9 +433,9 @@ public void RemoveUser(IAzureAccount account, IAzureTokenCache tokenCache) case AzureAccount.AccountType.ServicePrincipal: try { - KeyStore.RemoveCredential(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, + KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret, account.Id, account.GetTenants().FirstOrDefault())); - KeyStore.RemoveCredential(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, + KeyStore.RemoveSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, account.GetTenants().FirstOrDefault())); } catch @@ -577,7 +577,7 @@ private AuthenticationParameters GetAuthenticationParameters( { try { - password = KeyStore.GetCredential(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret + password = KeyStore.GetSecureString(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret , account.Id, tenant)); } catch @@ -591,7 +591,7 @@ private AuthenticationParameters GetAuthenticationParameters( { try { - certificatePassword = KeyStore.GetCredential(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword + certificatePassword = KeyStore.GetSecureString(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword , account.Id, tenant)); } catch diff --git a/src/Accounts/Authentication/Identity/AsyncLockWithValue.cs b/src/Accounts/Authentication/Identity/AsyncLockWithValue.cs index 6ac99a950115..e89ca5eb2394 100644 --- a/src/Accounts/Authentication/Identity/AsyncLockWithValue.cs +++ b/src/Accounts/Authentication/Identity/AsyncLockWithValue.cs @@ -126,24 +126,6 @@ private void SetValue(T value) } } - /// - /// Try to reset value and fail if value is locked. - /// - /// - public bool TryClearValue() - { - lock (_syncObj) - { - if (!_isLocked) - { - _value = default(T); - _hasValue = false; - return true; - } - } - return false; - } - /// /// Release the lock and allow next waiter acquire it /// diff --git a/src/Accounts/Authentication/KeyStore/AsyncLockWithValue.cs b/src/Accounts/Authentication/KeyStore/AsyncLockWithValue.cs new file mode 100644 index 000000000000..4e5d3d6b93c7 --- /dev/null +++ b/src/Accounts/Authentication/KeyStore/AsyncLockWithValue.cs @@ -0,0 +1,212 @@ +// ---------------------------------------------------------------------------------- +// +// 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 Microsoft.Azure.PowerShell.Authenticators.Identity; +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; + +namespace Microsoft.Azure.Commands.ResourceManager.Common +{ + /// + /// Primitive that combines async lock and value cache + /// + /// + internal sealed class AsyncLockWithValue + { + private readonly object _syncObj = new object(); + private Queue> _waiters; + private bool _isLocked; + private bool _hasValue; + private T _value; + + /// + /// 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. + /// + /// + /// + /// + public async ValueTask GetLockOrValueAsync(bool async, CancellationToken cancellationToken = default) + { + TaskCompletionSource 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>(); + } + // if async == false, valueTcs will be waited only in this thread and only synchronously, so RunContinuationsAsynchronously isn't needed. + valueTcs = new TaskCompletionSource(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; + } + } + + /// + /// Set value to the cache and to all the waiters + /// + /// + private void SetValue(T value) + { + Queue> 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)); + } + } + + /// + /// Try to reset value and fail if value is locked. + /// + /// + public bool TryClearValue() + { + lock (_syncObj) + { + if (!_isLocked) + { + _value = default(T); + _hasValue = false; + return true; + } + } + return false; + } + + /// + /// Release the lock and allow next waiter acquire it + /// + private void Reset() + { + TaskCompletionSource nextWaiter = UnlockOrGetNextWaiter(); + while (nextWaiter != default && !nextWaiter.TrySetResult(new Lock(this))) + { + nextWaiter = UnlockOrGetNextWaiter(); + } + } + + private TaskCompletionSource 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 _owner; + public bool HasValue => _owner == default; + public T Value { get; } + + public Lock(T value) + { + _owner = default; + Value = value; + } + + public Lock(AsyncLockWithValue owner) + { + _owner = owner; + Value = default; + } + + public void SetValue(T value) => _owner.SetValue(value); + + public void Dispose() => _owner?.Reset(); + } + } +} diff --git a/src/Accounts/Authentication/KeyStore/AzKeyStore.cs b/src/Accounts/Authentication/KeyStore/AzKeyStore.cs index 7c610f20fd6d..ade907b83f9e 100644 --- a/src/Accounts/Authentication/KeyStore/AzKeyStore.cs +++ b/src/Accounts/Authentication/KeyStore/AzKeyStore.cs @@ -23,11 +23,11 @@ public class AzKeyStore : IDisposable public string FileName { get; set; } public string Directory { get; set; } - private IKeyStore _inMemoryStore = null; - public IKeyStore InMemoryStore + private IKeyCache _inMemoryKeyCache = null; + public IKeyCache InMemoryKeyCache { - get => _inMemoryStore; - set => _inMemoryStore = value; + get => _inMemoryKeyCache; + set => _inMemoryKeyCache = value; } private IStorage inputStorage = null; @@ -38,31 +38,43 @@ public bool IsProtected private set; } - public AzKeyStore(string directory, string fileName, IStorage storage = null) + public AzKeyStore(string directory, string fileName, bool enableContextAutoSaving = true, IStorage storage = null) { - InMemoryStore = new InMemoryKeyStore(); - InMemoryStore.SetBeforeAccess(LoadStorage); + InMemoryKeyCache = new InMemoryKeyCache(); + + if (enableContextAutoSaving) + { + InMemoryKeyCache.SetBeforeAccess(LoadStorage); + InMemoryKeyCache.SetOnUpdate(UpdateStorage); + } FileName = fileName; Directory = directory; inputStorage = storage; - InMemoryKeyStore.RegisterJsonConverter(typeof(ServicePrincipalKey), typeof(ServicePrincipalKey).Name); - InMemoryKeyStore.RegisterJsonConverter(typeof(SecureString), typeof(SecureString).Name, new SecureStringConverter()); + Common.InMemoryKeyCache.RegisterJsonConverter(typeof(ServicePrincipalKey), typeof(ServicePrincipalKey).Name); + Common.InMemoryKeyCache.RegisterJsonConverter(typeof(SecureString), typeof(SecureString).Name, new SecureStringConverter()); } private void LoadStorage(KeyStoreNotificationArgs args) { - var asyncHelper = StorageHelper.GetStorageHelperAsync(true, FileName, Directory, args.KeyStore, inputStorage); + var asyncHelper = StorageHelper.GetStorageHelperAsync(true, FileName, Directory, args.KeyCache, inputStorage); var helper = asyncHelper.GetAwaiter().GetResult(); IsProtected = helper.IsProtected; } + private void UpdateStorage(KeyStoreNotificationArgs args) + { + var asyncHelper = StorageHelper.GetStorageHelperAsync(false, FileName, Directory, args.KeyCache, inputStorage); + var helper = asyncHelper.GetAwaiter().GetResult(); + helper.WriteToCachedStorage(args.KeyCache); + } + public void Clear() { - InMemoryStore.Clear(); + InMemoryKeyCache.Clear(); } public void Dispose() @@ -70,19 +82,31 @@ public void Dispose() StorageHelper.TryClearLockedStorageHelper(); } - public void SaveCredential(IKeyStoreKey key, SecureString value) + public void SaveSecureString(IKeyStoreKey key, SecureString value) + { + InMemoryKeyCache.SaveKey(key, value); + } + + public SecureString GetSecureString(IKeyStoreKey key) + { + return InMemoryKeyCache.GetKey(key); + } + + public bool RemoveSecureString(IKeyStoreKey key) { - InMemoryStore.SaveKey(key, value); + return InMemoryKeyCache.DeleteKey(key); } - public SecureString GetCredential(IKeyStoreKey key) + public void EnableSyncToStorage() { - return InMemoryStore.GetKey(key); + InMemoryKeyCache.SetBeforeAccess(LoadStorage); + InMemoryKeyCache.SetOnUpdate(UpdateStorage); } - public bool RemoveCredential(IKeyStoreKey key) + public void DisableSyncToStorage() { - return InMemoryStore.DeleteKey(key); + InMemoryKeyCache.SetBeforeAccess(null); + InMemoryKeyCache.SetOnUpdate(null); } } } diff --git a/src/Accounts/Authentication/KeyStore/IKeyStore.cs b/src/Accounts/Authentication/KeyStore/IKeyCache.cs similarity index 97% rename from src/Accounts/Authentication/KeyStore/IKeyStore.cs rename to src/Accounts/Authentication/KeyStore/IKeyCache.cs index e7c1f637bc4e..81d631b0dd35 100644 --- a/src/Accounts/Authentication/KeyStore/IKeyStore.cs +++ b/src/Accounts/Authentication/KeyStore/IKeyCache.cs @@ -14,7 +14,7 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common { - public interface IKeyStore + public interface IKeyCache { void SaveKey(IKeyStoreKey key, T value); diff --git a/src/Accounts/Authentication/KeyStore/IKeyStoreKey.cs b/src/Accounts/Authentication/KeyStore/IKeyStoreKey.cs index 4dbac64496d3..c126232d6193 100644 --- a/src/Accounts/Authentication/KeyStore/IKeyStoreKey.cs +++ b/src/Accounts/Authentication/KeyStore/IKeyStoreKey.cs @@ -16,14 +16,30 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common { public abstract class IKeyStoreKey { + /// + /// Create key from the data fields of KeyStoreKey. + /// protected abstract string CreateKey(); + /// + /// Convert key to string. + /// public override abstract string ToString(); - public override abstract bool Equals(object obj); - + /// + /// Generate hash code of KeyStoreKey. + /// public override abstract int GetHashCode(); + /// + /// Check whether the current key is exactly equal to another. + /// + public override abstract bool Equals(object obj); + + /// + /// Check whether the current key can be treated as equal to another even though they are not equal. + /// This method can be used as fuzzy search of the keys. + /// public abstract bool BeEquivalent(object obj); } } diff --git a/src/Accounts/Authentication/KeyStore/IStorage.cs b/src/Accounts/Authentication/KeyStore/IStorage.cs index ab8b6653973c..71b7055c6f06 100644 --- a/src/Accounts/Authentication/KeyStore/IStorage.cs +++ b/src/Accounts/Authentication/KeyStore/IStorage.cs @@ -11,7 +11,6 @@ // See the License for the specific language governing permissions and // limitations under the License. // ---------------------------------------------------------------------------------- -using System; namespace Microsoft.Azure.Commands.ResourceManager.Common { diff --git a/src/Accounts/Authentication/KeyStore/IStorageHelper.cs b/src/Accounts/Authentication/KeyStore/IStorageHelper.cs index eaa4c1619726..49be99e02563 100644 --- a/src/Accounts/Authentication/KeyStore/IStorageHelper.cs +++ b/src/Accounts/Authentication/KeyStore/IStorageHelper.cs @@ -11,7 +11,6 @@ // See the License for the specific language governing permissions and // limitations under the License. // ---------------------------------------------------------------------------------- -using System; namespace Microsoft.Azure.Commands.ResourceManager.Common { @@ -19,13 +18,9 @@ public interface IStorageHelper { void Clear(); - byte[] LoadUnencryptedTokenCache(); + void LoadFromCachedStorage(IKeyCache keycache); - void SaveUnencryptedTokenCache(byte[] tokenCache); - - void LoadFromCachedStorage(IKeyStore keystore); - - void WriteToCachedStorage(KeyStoreNotificationArgs args); + void WriteToCachedStorage(IKeyCache keycache); bool IsProtected { diff --git a/src/Accounts/Authentication/KeyStore/InMemoryKeyStore.cs b/src/Accounts/Authentication/KeyStore/InMemoryKeyCache.cs similarity index 97% rename from src/Accounts/Authentication/KeyStore/InMemoryKeyStore.cs rename to src/Accounts/Authentication/KeyStore/InMemoryKeyCache.cs index de4a3be78306..a2c8a78a5d73 100644 --- a/src/Accounts/Authentication/KeyStore/InMemoryKeyStore.cs +++ b/src/Accounts/Authentication/KeyStore/InMemoryKeyCache.cs @@ -11,18 +11,16 @@ // See the License for the specific language governing permissions and // limitations under the License. // ---------------------------------------------------------------------------------- -using Microsoft.Identity.Client; using Newtonsoft.Json; using System; using System.Collections.Concurrent; using System.Collections.Generic; using System.Linq; using System.Text; -using System.Threading.Tasks; namespace Microsoft.Azure.Commands.ResourceManager.Common { - internal class InMemoryKeyStore : IKeyStore + internal class InMemoryKeyCache : IKeyCache { internal class KeyStoreElement { @@ -48,7 +46,7 @@ public void SaveKey(IKeyStoreKey key, T value) { var args = new KeyStoreNotificationArgs() { - KeyStore = this + KeyCache = this }; BeforeAccess?.Invoke(args) ; if (!_typeNameMap.ContainsKey(key.GetType()) || !_typeNameMap.ContainsKey(value.GetType())) @@ -63,7 +61,7 @@ public T GetKey(IKeyStoreKey key) { var args = new KeyStoreNotificationArgs() { - KeyStore = this + KeyCache = this }; BeforeAccess?.Invoke(args); @@ -88,7 +86,7 @@ public bool DeleteKey(IKeyStoreKey key) { var args = new KeyStoreNotificationArgs() { - KeyStore = this + KeyCache = this }; BeforeAccess?.Invoke(args); bool ret = false; @@ -159,7 +157,7 @@ public void Clear() { var args = new KeyStoreNotificationArgs() { - KeyStore = this + KeyCache = this }; BeforeAccess?.Invoke(args); _credentials.Clear(); diff --git a/src/Accounts/Authentication/KeyStore/KeyStoreNotificationArgs.cs b/src/Accounts/Authentication/KeyStore/KeyStoreNotificationArgs.cs index 25cdcf22a50a..3c87ffe1ba44 100644 --- a/src/Accounts/Authentication/KeyStore/KeyStoreNotificationArgs.cs +++ b/src/Accounts/Authentication/KeyStore/KeyStoreNotificationArgs.cs @@ -16,7 +16,7 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common { public class KeyStoreNotificationArgs { - public IKeyStore KeyStore; + public IKeyCache KeyCache; } public delegate void KeyStoreCallbak(KeyStoreNotificationArgs args); diff --git a/src/Accounts/Authentication/KeyStore/StorageHelper.cs b/src/Accounts/Authentication/KeyStore/StorageHelper.cs index 9b97101744a9..85ed4675b180 100644 --- a/src/Accounts/Authentication/KeyStore/StorageHelper.cs +++ b/src/Accounts/Authentication/KeyStore/StorageHelper.cs @@ -11,7 +11,6 @@ // See the License for the specific language governing permissions and // limitations under the License. // ---------------------------------------------------------------------------------- -using Microsoft.Azure.PowerShell.Authenticators.Identity; using Microsoft.Identity.Client.Extensions.Msal; using Microsoft.IdentityModel.Abstractions; using System; @@ -21,7 +20,7 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common { - public class StorageHelper : IStorageHelper + internal class StorageHelper : IStorageHelper { private const string KeyChainServiceName = "Microsoft.Azure.PowerShell"; @@ -89,7 +88,7 @@ private static StorageHelper Create(StorageCreationProperties storageCreationPro } #region Public API - public static async Task GetStorageHelperAsync(bool async, string fileName, string directory, IKeyStore keystore, IStorage storage = null) + public static async Task GetStorageHelperAsync(bool async, string fileName, string directory, IKeyCache keycache, IStorage storage = null) { StorageHelper storageHelper = null; @@ -110,8 +109,7 @@ public static async Task GetStorageHelperAsync(bool async, strin storageHelper = GetFallbackStorageHelper(fileName, directory, storage); storageHelper.VerifyPersistence(); } - storageHelper.RegisterCache(keystore); - storageHelper.LoadFromCachedStorage(keystore); + storageHelper.LoadFromCachedStorage(keycache); asyncLock.SetValue(storageHelper); } return storageHelper; @@ -130,48 +128,9 @@ public void Clear() } } - public byte[] LoadUnencryptedTokenCache() + public void LoadFromCachedStorage(IKeyCache keycache) { - using (CreateCrossPlatLock(_storageCreationProperties)) - { - return PersistanceStore.ReadData(); - } - } - - public void SaveUnencryptedTokenCache(byte[] tokenCache) - { - using (CreateCrossPlatLock(_storageCreationProperties)) - { - PersistanceStore.WriteData(tokenCache); - } - } - - public void RegisterCache(IKeyStore keystore) - { - if (keystore == null) - { - throw new ArgumentNullException(nameof(keystore)); - } - - _logger.LogInformation($"Registering token cache with on disk storage"); - - keystore.SetOnUpdate(WriteToCachedStorage); - - _logger.LogInformation($"Done initializing"); - } - - public void UnregisterCache(IKeyStore keystore) - { - if (keystore == null) - { - throw new ArgumentNullException(nameof(keystore)); - } - keystore.SetOnUpdate(null); - } - - public void LoadFromCachedStorage(IKeyStore keystore) - { - LogMessage(EventLogLevel.Verbose, $"Before access\nAcquiring lock for token cache"); + LogMessage(EventLogLevel.Verbose, $"Before access\nAcquiring lock for keystore"); using (CreateCrossPlatLock(_storageCreationProperties)) { @@ -184,7 +143,7 @@ public void LoadFromCachedStorage(IKeyStore keystore) } catch (Exception ex) { - LogMessage(EventLogLevel.Error, $"Could not read the token cache. Ignoring. Exception: {ex}"); + LogMessage(EventLogLevel.Error, $"Could not read the keystore. Ignoring. Exception: {ex}"); return; } @@ -193,7 +152,7 @@ public void LoadFromCachedStorage(IKeyStore keystore) try { LogMessage(EventLogLevel.Verbose, $"Deserializing the store"); - keystore.Deserialize(cachedStoreData, false); + keycache.Deserialize(cachedStoreData, false); } catch (Exception e) { @@ -206,7 +165,7 @@ public void LoadFromCachedStorage(IKeyStore keystore) } } - public void WriteToCachedStorage(KeyStoreNotificationArgs args) + public void WriteToCachedStorage(IKeyCache keycache) { using (CreateCrossPlatLock(_storageCreationProperties)) { @@ -215,7 +174,7 @@ public void WriteToCachedStorage(KeyStoreNotificationArgs args) LogMessage(EventLogLevel.Verbose, $"After access, cache in memory HasChanged"); try { - data = args.KeyStore.Serialize(); + data = keycache.Serialize(); } catch (Exception e) { diff --git a/src/Accounts/Authentication/KeyStore/StorageWrapper.cs b/src/Accounts/Authentication/KeyStore/StorageWrapper.cs index aaa15a7600d2..647cabe8acee 100644 --- a/src/Accounts/Authentication/KeyStore/StorageWrapper.cs +++ b/src/Accounts/Authentication/KeyStore/StorageWrapper.cs @@ -16,7 +16,7 @@ namespace Microsoft.Azure.Commands.ResourceManager.Common { - class StorageWrapper : IStorage + internal class StorageWrapper : IStorage { public StorageCreationProperties StorageCreationProperties { get; set; } diff --git a/tools/Test/SmokeTest/InstallAzModules.ps1 b/tools/Test/SmokeTest/InstallAzModules.ps1 index ebacf1d52484..8982e959a0a6 100644 --- a/tools/Test/SmokeTest/InstallAzModules.ps1 +++ b/tools/Test/SmokeTest/InstallAzModules.ps1 @@ -44,7 +44,8 @@ Write-Host "Installing Az..." Install-Module -Name Az -Repository $gallery -Scope CurrentUser -AllowClobber -Force $file = Get-ChildItem $localRepoLocation | Where-Object {$_.Name -like "ThreadJob*"} -if ($file -ne $null) { +$installedModule = Get-Module -ListAVailable -Name ThreadJob +if ($file -ne $null -and $installedModule -ne $null) { Write-Host "Install ThreadJob..." Install-Module -Name ThreadJob -Repository $gallery -Scope CurrentUser -AllowClobber -Force } diff --git a/tools/Test/SmokeTest/RmCoreSmokeTests.ps1 b/tools/Test/SmokeTest/RmCoreSmokeTests.ps1 index b7b231f5139b..ea2341b0ddd2 100644 --- a/tools/Test/SmokeTest/RmCoreSmokeTests.ps1 +++ b/tools/Test/SmokeTest/RmCoreSmokeTests.ps1 @@ -164,7 +164,7 @@ $resourceTestCommands = @( $generalCommands = @( @{ - Name = "Import Az.Accounts in Parallel Job"; + Name = "Import Az.Accounts in Parallel (Process)"; Command = { if ($null -ne $env:SYSTEM_DEFINITIONID -or $null -ne $env:Release_DefinitionId -or $null -ne $env:AZUREPS_HOST_ENVIRONMENT) { Write-Warning "Skipping because 'Start-Job' is not supported by design in scenarios where PowerShell is being hosted in other applications." @@ -185,13 +185,9 @@ $generalCommands = @( Retry = 0; # no need to retry } @{ - Name = "Import Az.Accounts in Parallel Thread Job"; + Name = "Import Az.Accounts in Parallel (Thread)"; Command = { - #if ($null -ne $env:SYSTEM_DEFINITIONID -or $null -ne $env:Release_DefinitionId) { - #Write-Warning "Skipping because 'Start-ThreadJob' is not supported by design in scenarios where PowerShell is being hosted in other applications." - #return - #} $importJobs = @() 1..50 | ForEach-Object { $importJobs += Start-ThreadJob -name "import-no.$_" -ScriptBlock { Import-Module Az.Accounts; Get-AzTenant; }