Skip to content

Commit

Permalink
Address review comments
Browse files Browse the repository at this point in the history
Co-authored-by: Yeming Liu <[email protected]>

Address review comments
  • Loading branch information
msJinLei committed Feb 2, 2023
1 parent a72f3d2 commit 50b10d6
Show file tree
Hide file tree
Showing 26 changed files with 336 additions and 151 deletions.
3 changes: 1 addition & 2 deletions src/Accounts/Accounts.Test/AutosaveTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@ public AutosaveTests(ITestOutputHelper output)
XunitTracingInterceptor.AddToContext(new XunitTracingInterceptor(output));
commandRuntimeMock = new MockCommandRuntime();
dataStore = new MemoryDataStore();
ResetState();
keyStore = SetMockedAzKeyStore();
}

Expand All @@ -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<byte[]>())).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;
}

Expand Down
2 changes: 1 addition & 1 deletion src/Accounts/Accounts.Test/ContextCmdletTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ public ContextCmdletTests(ITestOutputHelper output)
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);
azKeyStore = new AzKeyStore(profilePath, AzureSession.Instance.KeyStoreFile, true, storageMocker.Object);
AzureSession.Instance.RegisterComponent(AzKeyStore.Name, () => azKeyStore, true);
}

Expand Down
2 changes: 1 addition & 1 deletion src/Accounts/Accounts.Test/ProfileCmdletTests.cs
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", storageMocker.Object);
var keyStore = new AzKeyStore(AzureSession.Instance.ARMProfileDirectory, "azkeystore", false, storageMocker.Object);
return keyStore;
}

Expand Down
6 changes: 3 additions & 3 deletions src/Accounts/Accounts/Account/ConnectAzureRmAccount.cs
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?.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));
Expand All @@ -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)
{
Expand Down Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 0 additions & 2 deletions src/Accounts/Accounts/ChangeLog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
4 changes: 2 additions & 2 deletions src/Accounts/Accounts/Context/ImportAzureRMContext.cs
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.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());
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/Accounts/Accounts/Context/SetAzureRMContext.cs
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.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());
}
}
Expand Down
4 changes: 2 additions & 2 deletions src/Accounts/Authentication.ResourceManager/AzureRmProfile.cs
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?.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);
}
Expand Down
26 changes: 13 additions & 13 deletions src/Accounts/Authentication.Test/AzKeyStorageTest.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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\""""}]";
Expand All @@ -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();
Expand All @@ -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();
Expand All @@ -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<ArgumentException>(() => store.GetCredential(servicePrincipalKey));
Assert.Throws<ArgumentException>(() => store.GetSecureString(servicePrincipalKey));

store.Clear();
}
Expand All @@ -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<List<Object>>(result);
Expand All @@ -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<List<Object>>(result);
Expand Down
10 changes: 3 additions & 7 deletions src/Accounts/Authentication.Test/StorageHelperTest.cs
Original file line number Diff line number Diff line change
Expand Up @@ -27,15 +27,15 @@ namespace Common.Authenticators.Test
public class StorageHelperTest
{
private Mock<IStorage> storageMocker = null;
private Mock<IKeyStore> keystoreMocker = null;
private Mock<IKeyCache> keystoreMocker = null;
private string profilePath = Path.Combine(Environment.GetFolderPath(Environment.SpecialFolder.UserProfile), Resources.AzureDirectoryName);
private string keyStoreFileName = "azkeystore";

public StorageHelperTest()
{
storageMocker = new Mock<IStorage>();
storageMocker.Setup(f => f.Create()).Returns(storageMocker.Object);
keystoreMocker = new Mock<IKeyStore>();
keystoreMocker = new Mock<IKeyCache>();
}

[Fact]
Expand Down Expand Up @@ -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);
Expand Down
Loading

0 comments on commit 50b10d6

Please sign in to comment.