Skip to content

Commit

Permalink
Fix AzKeyStore issues and optimize initialization and update
Browse files Browse the repository at this point in the history 
* 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.
  • Loading branch information
@msJinLei
msJinLei committed Jan 30, 2023
1 parent df75309 commit 525af60
Show file tree
Hide file tree
Showing 25 changed files with 843 additions and 374 deletions.
4 changes: 1 addition & 3 deletions src/Accounts/Accounts.Test/AutosaveTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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", storageMocker.Object);
return keyStore;
}

Expand Down
18 changes: 14 additions & 4 deletions src/Accounts/Accounts.Test/ContextCmdletTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand All @@ -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, storageMocker.Object);
AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => azKeyStore, true);
}

[Fact]
Expand Down
2 changes: 1 addition & 1 deletion src/Accounts/Accounts.Test/ProfileCmdletTests.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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", storageMocker.Object);
return keyStore;
}

Expand Down
13 changes: 3 additions & 10 deletions src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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?.SaveCredential(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));
Expand All @@ -451,7 +451,7 @@ public override void ExecuteCmdlet()

if (azureAccount.Type == AzureAccount.AccountType.ServicePrincipal && password != null)
{
keyStore?.SaveKey(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret
keyStore?.SaveCredential(new ServicePrincipalKey(AzureAccount.Property.ServicePrincipalSecret
,azureAccount.Id, Tenant), password);
if (GetContextModificationScope() == ContextModificationScope.CurrentUser && !keyStore.IsProtected)
{
Expand Down Expand Up @@ -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);
AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => keyStore);

if (!InitializeProfileProvider(autoSaveEnabled))
Expand All @@ -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))
{
Expand Down
5 changes: 0 additions & 5 deletions src/Accounts/Accounts/AutoSave/DisableAzureRmContextAutosave.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -92,11 +92,6 @@ void DisableAutosave(IAzureSession session, bool writeAutoSaveFile, out ContextA
builder.Reset();
}

if (AzureSession.Instance.TryGetComponent(AzKeyStore.Name, out AzKeyStore keystore))
{
keystore.DisableAutoSaving();
}

if (writeAutoSaveFile)
{
FileUtilities.EnsureDirectoryExists(session.ProfileDirectory);
Expand Down
7 changes: 0 additions & 7 deletions src/Accounts/Accounts/AutoSave/EnableAzureRmContextAutosave.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,13 +102,6 @@ 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.Flush();
keystore.DisableAutoSaving();
}


if (writeAutoSaveFile)
{
try
Expand Down
4 changes: 4 additions & 0 deletions src/Accounts/Accounts/ChangeLog.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@
-->

## Upcoming Release
* 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]
Expand Down
4 changes: 2 additions & 2 deletions src/Accounts/Accounts/Context/ImportAzureRMContext.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.SaveCredential(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.SaveCredential(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, context.Value.Tenant?.Id)
,password.ConvertToSecureString());
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/Accounts/Accounts/Context/SetAzureRMContext.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.SaveCredential(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.SaveCredential(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, Context.Tenant?.Id)
, password.ConvertToSecureString());
}
}
Expand Down
8 changes: 2 additions & 6 deletions src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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?.SaveCredential(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?.SaveCredential(new ServicePrincipalKey(AzureAccount.Property.CertificatePassword, account.Id, account.GetTenants().First())
, account.ExtendedProperties.GetProperty(AzureAccount.Property.CertificatePassword).ConvertToSecureString());
account.ExtendedProperties.Remove(AzureAccount.Property.CertificatePassword);
}
Expand Down Expand Up @@ -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
{
Expand Down
Loading

0 comments on commit 525af60

Please sign in to comment.