Skip to content

Commit 7554665

Browse files
cormacpaynecormacpayne
authored
Merge pull request #4146 from markcowl/cacheload
Fixing issues with Import- and Save Context cmdlets
2 parents 2683def + d0f5b7d commit 7554665

File tree

6 files changed

+141
-59
lines changed

6 files changed

+141
-59
lines changed

‎src/Common/Commands.Common.Authentication.Test/AzureSessionTests.cs

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
// ----------------------------------------------------------------------------------
2+
//
3+
// Copyright Microsoft Corporation
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
// ----------------------------------------------------------------------------------
14+
15+
using Microsoft.Azure.Commands.Common.Authentication;
16+
using Microsoft.Azure.Commands.Common.Authentication.Factories;
17+
using Microsoft.Azure.Commands.Common.Authentication.Models;
18+
using Microsoft.WindowsAzure.Commands.Test.Utilities.Common;
19+
using System;
20+
using System.Collections.Generic;
21+
using Microsoft.WindowsAzure.Commands.ScenarioTest;
22+
using Xunit;
23+
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
24+
using System.IO;
25+
26+
namespace Common.Authentication.Test
27+
{
28+
public class AzureSessionTests
29+
{
30+
[Fact]
31+
[Trait(Category.AcceptanceType, Category.CheckIn)]
32+
public void InitializerCreatesTokenCacheFile()
33+
{
34+
IAzureSession oldSession = null;
35+
try
36+
{
37+
oldSession = AzureSession.Instance;
38+
}
39+
catch { }
40+
try
41+
{
42+
var store = new MemoryDataStore();
43+
AzureSessionInitializer.CreateOrReplaceSession(store);
44+
var session = AzureSession.Instance;
45+
var tokenCacheFile = Path.Combine(session.ProfileDirectory, session.TokenCacheFile);
46+
Assert.True(store.FileExists(tokenCacheFile));
47+
48+
}
49+
finally
50+
{
51+
AzureSession.Initialize(() => oldSession, true);
52+
}
53+
}
54+
55+
[Fact]
56+
[Trait(Category.AcceptanceType, Category.CheckIn)]
57+
public void TokenCacheIgnoresInvalidData()
58+
{
59+
var store = new AzureTokenCache { CacheData = new byte[] { 3, 0, 0, 0, 0, 0, 0, 0 } };
60+
var cache = new AuthenticationStoreTokenCache(store);
61+
Assert.NotEqual(cache.CacheData, store.CacheData);
62+
}
63+
64+
[Fact]
65+
[Trait(Category.AcceptanceType, Category.CheckIn)]
66+
public void TokenCacheUsesValidData()
67+
{
68+
var store = new AzureTokenCache { CacheData = new byte[] { 2, 0, 0, 0, 0, 0, 0, 0 } };
69+
var cache = new AuthenticationStoreTokenCache(store);
70+
Assert.Equal(cache.CacheData, store.CacheData);
71+
}
72+
73+
}
74+
}

‎src/Common/Commands.Common.Authentication.Test/Commands.Common.Authentication.Test.csproj

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,7 @@
120120
</Reference>
121121
</ItemGroup>
122122
<ItemGroup>
123+
<Compile Include="AzureSessionTests.cs" />
123124
<Compile Include="ClientFactoryHandlerTests.cs" />
124125
<Compile Include="ClientFactoryTests.cs" />
125126
<Compile Include="AuthenticationFactoryTests.cs" />

‎src/Common/Commands.Common.Authentication/Authentication/AuthenticationStoreTokenCache.cs

Lines changed: 11 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -22,19 +22,17 @@ namespace Microsoft.Azure.Commands.Common.Authentication
2222
[Serializable]
2323
public class AuthenticationStoreTokenCache : TokenCache, IAzureTokenCache, IDisposable
2424
{
25-
AzureTokenCache _tokenStore;
26-
25+
IAzureTokenCache _store = new AzureTokenCache();
2726
public byte[] CacheData
2827
{
2928
get
3029
{
31-
return _tokenStore.CacheData;
30+
return Serialize();
3231
}
3332

3433
set
3534
{
36-
this.Clear();
37-
_tokenStore.CacheData = value;
35+
this.Deserialize(value);
3836
}
3937
}
4038

@@ -45,10 +43,9 @@ public AuthenticationStoreTokenCache(AzureTokenCache store) : base()
4543
throw new ArgumentNullException("store");
4644
}
4745

48-
_tokenStore = store;
49-
if (_tokenStore != null && _tokenStore.CacheData != null && _tokenStore.CacheData.Length > 0)
46+
if (store.CacheData != null && store.CacheData.Length > 0)
5047
{
51-
base.Deserialize(_tokenStore.CacheData);
48+
CacheData = store.CacheData;
5249
}
5350

5451
AfterAccess += HandleAfterAccess;
@@ -59,41 +56,33 @@ public AuthenticationStoreTokenCache(AzureTokenCache store) : base()
5956
/// </summary>
6057
/// <param name="cache">The cache to copy</param>
6158
/// <param name="store">The store to use for persisting state</param>
62-
public AuthenticationStoreTokenCache(TokenCache cache, AzureTokenCache store) : this(store)
59+
public AuthenticationStoreTokenCache(TokenCache cache) : base()
6360
{
6461
if (null == cache)
6562
{
6663
throw new ArgumentNullException("Cache");
6764
}
6865

69-
Deserialize(cache.Serialize());
70-
}
71-
72-
/// <summary>
73-
/// Create a token cache, copying any data from the given token cache
74-
/// </summary>
75-
/// <param name="cache">The cache to copy</param>
76-
public AuthenticationStoreTokenCache(TokenCache cache) : this(cache, new AzureTokenCache())
77-
{
66+
CacheData = cache.Serialize();
67+
AfterAccess += HandleAfterAccess;
7868
}
7969

80-
8170
public void HandleAfterAccess(TokenCacheNotificationArgs args)
8271
{
8372
if (HasStateChanged)
8473
{
85-
_tokenStore.CacheData = Serialize();
74+
_store.CacheData = Serialize();
8675
}
8776
}
8877

8978
protected virtual void Dispose(bool disposing)
9079
{
9180
if (disposing)
9281
{
93-
var cache = Interlocked.Exchange(ref _tokenStore, null);
82+
var cache = Interlocked.Exchange(ref _store, null);
9483
if (cache != null)
9584
{
96-
cache.CacheData = base.Serialize();
85+
cache.CacheData = Serialize();
9786
}
9887
}
9988
}

‎src/Common/Commands.Common.Authentication/Authentication/ProtectedFileTokenCache.cs

Lines changed: 34 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ public class ProtectedFileTokenCache : TokenCache, IAzureTokenCache
3737

3838
private static readonly Lazy<ProtectedFileTokenCache> instance = new Lazy<ProtectedFileTokenCache>(() => new ProtectedFileTokenCache());
3939

40+
IDataStore _store;
41+
4042
public byte[] CacheData
4143
{
4244
get
@@ -47,32 +49,37 @@ public byte[] CacheData
4749
set
4850
{
4951
Deserialize(value);
52+
HasStateChanged = true;
53+
EnsureStateSaved();
5054
}
5155
}
5256

5357
// Initializes the cache against a local file.
5458
// If the file is already present, it loads its content in the ADAL cache
5559
private ProtectedFileTokenCache()
5660
{
61+
_store = AzureSession.Instance.DataStore;
5762
Initialize(CacheFileName);
5863
}
5964

60-
public ProtectedFileTokenCache(byte[] inputData)
65+
public ProtectedFileTokenCache(byte[] inputData, IDataStore store = null) : this(CacheFileName, store)
6166
{
62-
AfterAccess = AfterAccessNotification;
63-
BeforeAccess = BeforeAccessNotification;
6467
CacheData = inputData;
6568
}
6669

70+
public ProtectedFileTokenCache(string cacheFile, IDataStore store = null)
71+
{
72+
_store = store ?? AzureSession.Instance.DataStore;
73+
Initialize(cacheFile);
74+
}
75+
6776
private void Initialize(string fileName)
6877
{
69-
AfterAccess = AfterAccessNotification;
70-
BeforeAccess = BeforeAccessNotification;
7178
lock (fileLock)
7279
{
73-
if (AzureSession.Instance.DataStore.FileExists(fileName))
80+
if (_store.FileExists(fileName))
7481
{
75-
var existingData = AzureSession.Instance.DataStore.ReadFileAsBytes(fileName);
82+
var existingData = _store.ReadFileAsBytes(fileName);
7683
if (existingData != null)
7784
{
7885
try
@@ -81,25 +88,26 @@ private void Initialize(string fileName)
8188
}
8289
catch (CryptographicException)
8390
{
84-
AzureSession.Instance.DataStore.DeleteFile(fileName);
91+
_store.DeleteFile(fileName);
8592
}
8693
}
8794
}
95+
96+
// Create the file to start with
97+
_store.WriteFile(CacheFileName, ProtectedData.Protect(Serialize(), null, DataProtectionScope.CurrentUser));
8898
}
89-
}
9099

91-
public ProtectedFileTokenCache(string cacheFile)
92-
{
93-
Initialize(cacheFile);
100+
AfterAccess = AfterAccessNotification;
101+
BeforeAccess = BeforeAccessNotification;
94102
}
95103

96104
// Empties the persistent store.
97105
public override void Clear()
98106
{
99107
base.Clear();
100-
if (AzureSession.Instance.DataStore.FileExists(CacheFileName))
108+
if (_store.FileExists(CacheFileName))
101109
{
102-
AzureSession.Instance.DataStore.DeleteFile(CacheFileName);
110+
_store.DeleteFile(CacheFileName);
103111
}
104112
}
105113

@@ -109,9 +117,9 @@ void BeforeAccessNotification(TokenCacheNotificationArgs args)
109117
{
110118
lock (fileLock)
111119
{
112-
if (AzureSession.Instance.DataStore.FileExists(CacheFileName))
120+
if (_store.FileExists(CacheFileName))
113121
{
114-
var existingData = AzureSession.Instance.DataStore.ReadFileAsBytes(CacheFileName);
122+
var existingData = _store.ReadFileAsBytes(CacheFileName);
115123
if (existingData != null)
116124
{
117125
try
@@ -120,7 +128,7 @@ void BeforeAccessNotification(TokenCacheNotificationArgs args)
120128
}
121129
catch (CryptographicException)
122130
{
123-
AzureSession.Instance.DataStore.DeleteFile(CacheFileName);
131+
_store.DeleteFile(CacheFileName);
124132
}
125133
}
126134
}
@@ -131,13 +139,18 @@ void BeforeAccessNotification(TokenCacheNotificationArgs args)
131139
void AfterAccessNotification(TokenCacheNotificationArgs args)
132140
{
133141
// if the access operation resulted in a cache update
134-
if (HasStateChanged)
142+
EnsureStateSaved();
143+
}
144+
145+
void EnsureStateSaved()
146+
{
147+
lock (fileLock)
135148
{
136-
lock (fileLock)
149+
if (HasStateChanged)
137150
{
138151
// reflect changes in the persistent store
139-
AzureSession.Instance.DataStore.WriteFile(CacheFileName,
140-
ProtectedData.Protect(Serialize(), null, DataProtectionScope.CurrentUser));
152+
_store.WriteFile(CacheFileName,
153+
ProtectedData.Protect(Serialize(), null, DataProtectionScope.CurrentUser));
141154
// once the write operation took place, restore the HasStateChanged bit to false
142155
HasStateChanged = false;
143156
}

‎src/Common/Commands.Common.Authentication/AzureSessionInitializer.cs

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -35,16 +35,32 @@ public static class AzureSessionInitializer
3535
/// </summary>
3636
public static void InitializeAzureSession()
3737
{
38-
AzureSession.Initialize(CreateInstance);
38+
AzureSession.Initialize(() => CreateInstance());
3939
}
4040

41-
static IAzureSession CreateInstance()
41+
/// <summary>
42+
/// Create a new session and replace any existing session
43+
/// </summary>
44+
public static void CreateOrReplaceSession()
45+
{
46+
CreateOrReplaceSession(new DiskDataStore());
47+
}
48+
49+
/// <summary>
50+
/// Create a new session and replace any existing session
51+
/// </summary>
52+
public static void CreateOrReplaceSession(IDataStore dataStore)
53+
{
54+
AzureSession.Initialize(() => CreateInstance(dataStore), true);
55+
}
56+
57+
static IAzureSession CreateInstance(IDataStore dataStore = null)
4258
{
4359
var session = new AdalSession
4460
{
4561
ClientFactory = new ClientFactory(),
4662
AuthenticationFactory = new AuthenticationFactory(),
47-
DataStore = new DiskDataStore(),
63+
DataStore = dataStore?? new DiskDataStore(),
4864
OldProfileFile = "WindowsAzureProfile.xml",
4965
OldProfileFileBackup = "WindowsAzureProfile.xml.bak",
5066
ProfileDirectory = Path.Combine(
@@ -58,18 +74,7 @@ static IAzureSession CreateInstance()
5874
FileUtilities.DataStore = session.DataStore;
5975
FileUtilities.EnsureDirectoryExists(session.ProfileDirectory);
6076
var cacheFile = Path.Combine(session.ProfileDirectory, session.TokenCacheFile);
61-
var contents = new byte[0];
62-
if (session.DataStore.FileExists(cacheFile))
63-
{
64-
contents = session.DataStore.ReadFileAsBytes(cacheFile);
65-
}
66-
67-
if (contents != null && contents.Length > 0)
68-
{
69-
contents = ProtectedData.Unprotect(contents, null, DataProtectionScope.CurrentUser);
70-
}
71-
72-
session.TokenCache = new ProtectedFileTokenCache(contents);
77+
session.TokenCache = new ProtectedFileTokenCache(cacheFile, session.DataStore);
7378
}
7479
catch
7580
{

‎src/ResourceManager/Profile/Commands.Profile.Test/AzureRMProfileTests.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -751,7 +751,7 @@ public void SavingProfileWorks()
751751
},
752752
""VersionProfile"": null,
753753
""TokenCache"": {
754-
""CacheData"": ""AQIDBAUGCAkA""
754+
""CacheData"": ""AgAAAAAAAAA=""
755755
},
756756
""ExtendedProperties"": {}
757757
}

0 commit comments

Comments
 (0)