Skip to content

Commit b0c398d

Browse files
SNOW-715524: Add SSO token cache (#921)
Co-authored-by: Krzysztof Nozderko <[email protected]>
1 parent 44c7cb3 commit b0c398d

23 files changed

+1289
-218
lines changed

README.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,12 @@ Logging description and configuration:
143143
Method of validating the connection's certificates in the .NET driver differs from the rest of the Snowflake drivers.
144144
Read more in [certificate validation](doc/CertficateValidation.md) docs.
145145

146+
## Cache
147+
148+
Storing tokens in cache for SSO/MFA authentication.
149+
150+
Read more in [cache](doc/Cache.md) docs.
151+
146152
---------------
147153

148154
## Notice

Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs

Lines changed: 109 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
using NUnit.Framework;
1010
using Snowflake.Data.Client;
1111
using Snowflake.Data.Core;
12+
using Snowflake.Data.Core.CredentialManager;
13+
using Snowflake.Data.Core.CredentialManager.Infrastructure;
1214
using Snowflake.Data.Core.Session;
1315
using Snowflake.Data.Core.Tools;
1416
using Snowflake.Data.Log;
@@ -17,7 +19,6 @@
1719

1820
namespace Snowflake.Data.Tests.IntegrationTests
1921
{
20-
2122
[TestFixture]
2223
class SFConnectionIT : SFBaseTest
2324
{
@@ -1042,6 +1043,75 @@ public void TestSSOConnectionTimeoutAfter10s()
10421043
Assert.LessOrEqual(stopwatch.ElapsedMilliseconds, (waitSeconds + 5) * 1000);
10431044
}
10441045

1046+
[Test]
1047+
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
1048+
public void TestSSOConnectionWithTokenCaching()
1049+
{
1050+
/*
1051+
* This test checks that the connector successfully stores an SSO token and uses it for authentication if it exists
1052+
* 1. Login normally using external browser with CLIENT_STORE_TEMPORARY_CREDENTIAL enabled
1053+
* 2. Login again, this time without a browser, as the connector should be using the SSO token retrieved from step 1
1054+
*/
1055+
1056+
// Set the CLIENT_STORE_TEMPORARY_CREDENTIAL property to true to enable token caching
1057+
// The specified user should be configured for SSO
1058+
var externalBrowserConnectionString
1059+
= ConnectionStringWithoutAuth
1060+
+ $";authenticator=externalbrowser;user={testConfig.user};CLIENT_STORE_TEMPORARY_CREDENTIAL=true;poolingEnabled=false";
1061+
1062+
using (IDbConnection conn = new SnowflakeDbConnection())
1063+
{
1064+
conn.ConnectionString = externalBrowserConnectionString;
1065+
1066+
// Authenticate to retrieve and store the token if doesn't exist or invalid
1067+
conn.Open();
1068+
Assert.AreEqual(ConnectionState.Open, conn.State);
1069+
}
1070+
1071+
using (IDbConnection conn = new SnowflakeDbConnection())
1072+
{
1073+
conn.ConnectionString = externalBrowserConnectionString;
1074+
1075+
// Authenticate using the SSO token (the connector will automatically use the token and a browser should not pop-up in this step)
1076+
conn.Open();
1077+
Assert.AreEqual(ConnectionState.Open, conn.State);
1078+
}
1079+
}
1080+
1081+
[Test]
1082+
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
1083+
public void TestSSOConnectionWithInvalidCachedToken()
1084+
{
1085+
/*
1086+
* This test checks that the connector will attempt to re-authenticate using external browser if the token retrieved from the cache is invalid
1087+
* 1. Create a credential manager and save credentials for the user with a wrong token
1088+
* 2. Open a connection which initially should try to use the token and then switch to external browser when the token fails
1089+
*/
1090+
1091+
using (IDbConnection conn = new SnowflakeDbConnection())
1092+
{
1093+
// Set the CLIENT_STORE_TEMPORARY_CREDENTIAL property to true to enable token caching
1094+
conn.ConnectionString
1095+
= ConnectionStringWithoutAuth
1096+
+ $";authenticator=externalbrowser;user={testConfig.user};CLIENT_STORE_TEMPORARY_CREDENTIAL=true;";
1097+
1098+
// Create a credential manager and save a wrong token for the test user
1099+
var key = SnowflakeCredentialManagerFactory.GetSecureCredentialKey(testConfig.host, testConfig.user, TokenType.IdToken);
1100+
var credentialManager = SFCredentialManagerInMemoryImpl.Instance;
1101+
credentialManager.SaveCredentials(key, "wrongToken");
1102+
1103+
// Use the credential manager with the wrong token
1104+
SnowflakeCredentialManagerFactory.SetCredentialManager(credentialManager);
1105+
1106+
// Open a connection which should switch to external browser after trying to connect using the wrong token
1107+
conn.Open();
1108+
Assert.AreEqual(ConnectionState.Open, conn.State);
1109+
1110+
// Switch back to the default credential manager
1111+
SnowflakeCredentialManagerFactory.UseDefaultCredentialManager();
1112+
}
1113+
}
1114+
10451115
[Test]
10461116
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
10471117
public void TestSSOConnectionWithWrongUser()
@@ -2353,6 +2423,44 @@ public void TestOpenAsyncThrowExceptionWhenOperationIsCancelled()
23532423
}
23542424
}
23552425

2426+
[Test]
2427+
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
2428+
public void TestSSOConnectionWithTokenCachingAsync()
2429+
{
2430+
/*
2431+
* This test checks that the connector successfully stores an SSO token and uses it for authentication if it exists
2432+
* 1. Login normally using external browser with CLIENT_STORE_TEMPORARY_CREDENTIAL enabled
2433+
* 2. Login again, this time without a browser, as the connector should be using the SSO token retrieved from step 1
2434+
*/
2435+
2436+
// Set the CLIENT_STORE_TEMPORARY_CREDENTIAL property to true to enable token caching
2437+
// The specified user should be configured for SSO
2438+
var externalBrowserConnectionString
2439+
= ConnectionStringWithoutAuth
2440+
+ $";authenticator=externalbrowser;user={testConfig.user};CLIENT_STORE_TEMPORARY_CREDENTIAL=true;poolingEnabled=false";
2441+
2442+
using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
2443+
{
2444+
conn.ConnectionString = externalBrowserConnectionString;
2445+
2446+
// Authenticate to retrieve and store the token if doesn't exist or invalid
2447+
Task connectTask = conn.OpenAsync(CancellationToken.None);
2448+
connectTask.Wait();
2449+
Assert.AreEqual(ConnectionState.Open, conn.State);
2450+
}
2451+
2452+
using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
2453+
{
2454+
conn.ConnectionString = externalBrowserConnectionString;
2455+
2456+
// Authenticate using the SSO token (the connector will automatically use the token and a browser should not pop-up in this step)
2457+
Task connectTask = conn.OpenAsync(CancellationToken.None);
2458+
connectTask.Wait();
2459+
Assert.AreEqual(ConnectionState.Open, conn.State);
2460+
}
2461+
2462+
}
2463+
23562464
[Test]
23572465
public void TestCloseSessionWhenGarbageCollectorFinalizesConnection()
23582466
{
Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
using Snowflake.Data.Core;
2+
using System.Net.Http;
3+
using System.Threading;
4+
using System.Threading.Tasks;
5+
6+
namespace Snowflake.Data.Tests.Mock
7+
{
8+
9+
class MockExternalBrowserRestRequester : IMockRestRequester
10+
{
11+
public string ProofKey { get; set; }
12+
13+
public string SSOUrl { get; set; }
14+
15+
public string IdToken { get; set; }
16+
17+
public bool ThrowInvalidIdToken { get; set; } = false;
18+
19+
public bool ThrowNonInvalidIdToken { get; set; } = false;
20+
21+
public T Get<T>(IRestRequest request)
22+
{
23+
throw new System.NotImplementedException();
24+
}
25+
26+
public Task<T> GetAsync<T>(IRestRequest request, CancellationToken cancellationToken)
27+
{
28+
throw new System.NotImplementedException();
29+
}
30+
31+
public T Post<T>(IRestRequest postRequest)
32+
{
33+
return Task.Run(async () => await (PostAsync<T>(postRequest, CancellationToken.None)).ConfigureAwait(false)).Result;
34+
}
35+
36+
public Task<T> PostAsync<T>(IRestRequest postRequest, CancellationToken cancellationToken)
37+
{
38+
SFRestRequest sfRequest = (SFRestRequest)postRequest;
39+
if (sfRequest.jsonBody is AuthenticatorRequest)
40+
{
41+
if (string.IsNullOrEmpty(SSOUrl))
42+
{
43+
var body = (AuthenticatorRequest)sfRequest.jsonBody;
44+
var port = body.Data.BrowserModeRedirectPort;
45+
SSOUrl = $"http://localhost:{port}/?token=mockToken";
46+
}
47+
48+
// authenticator
49+
var authnResponse = new AuthenticatorResponse
50+
{
51+
success = true,
52+
data = new AuthenticatorResponseData
53+
{
54+
proofKey = ProofKey,
55+
ssoUrl = SSOUrl,
56+
}
57+
};
58+
59+
return Task.FromResult<T>((T)(object)authnResponse);
60+
}
61+
else
62+
{
63+
// login
64+
LoginResponse loginResponse;
65+
if (ThrowInvalidIdToken)
66+
{
67+
ThrowInvalidIdToken = false;
68+
loginResponse = new LoginResponse
69+
{
70+
success = false,
71+
code = SFError.ID_TOKEN_INVALID.GetAttribute<SFErrorAttr>().errorCode,
72+
message = "",
73+
};
74+
}
75+
else if (ThrowNonInvalidIdToken)
76+
{
77+
ThrowInvalidIdToken = false;
78+
loginResponse = new LoginResponse
79+
{
80+
success = false,
81+
code = SFError.INTERNAL_ERROR.GetAttribute<SFErrorAttr>().errorCode,
82+
message = "",
83+
};
84+
}
85+
else
86+
{
87+
loginResponse = new LoginResponse
88+
{
89+
success = true,
90+
data = new LoginResponseData
91+
{
92+
idToken = IdToken,
93+
sessionId = "",
94+
token = "",
95+
masterToken = "",
96+
masterValidityInSeconds = 0,
97+
authResponseSessionInfo = new SessionInfo
98+
{
99+
databaseName = "",
100+
schemaName = "",
101+
roleName = "",
102+
warehouseName = "",
103+
}
104+
}
105+
};
106+
}
107+
108+
return Task.FromResult<T>((T)(object)loginResponse);
109+
}
110+
}
111+
112+
public HttpResponseMessage Get(IRestRequest request)
113+
{
114+
throw new System.NotImplementedException();
115+
}
116+
117+
public Task<HttpResponseMessage> GetAsync(IRestRequest request, CancellationToken cancellationToken)
118+
{
119+
throw new System.NotImplementedException();
120+
}
121+
122+
public void setHttpClient(HttpClient httpClient)
123+
{
124+
// Nothing to do
125+
}
126+
}
127+
}

Snowflake.Data.Tests/UnitTests/CredentialManager/SFCredentialManagerFileImplTest.cs

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -258,9 +258,10 @@ public void TestReadingAndWritingAreUnavailableIfDirLockExists()
258258
}
259259

260260
[Test]
261-
public void TestChangeCacheDirPermissionsWhenInsecure()
261+
public void TestThatDoesNotChangeCacheDirPermissionsWhenInsecure()
262262
{
263263
// arrange
264+
var insecurePermissions = FileAccessPermissions.UserReadWriteExecute | FileAccessPermissions.GroupRead;
264265
var tempDirectory = Path.Combine(Path.GetTempPath(), Path.GetRandomFileName());
265266
t_environmentOperations
266267
.Setup(e => e.GetEnvironmentVariable(SFCredentialManagerFileStorage.CredentialCacheDirectoryEnvironmentName))
@@ -269,15 +270,15 @@ public void TestChangeCacheDirPermissionsWhenInsecure()
269270
try
270271
{
271272
DirectoryOperations.Instance.CreateDirectory(tempDirectory);
272-
UnixOperations.Instance.ChangePermissions(tempDirectory, FileAccessPermissions.UserReadWriteExecute | FileAccessPermissions.GroupRead);
273+
UnixOperations.Instance.ChangePermissions(tempDirectory, insecurePermissions);
273274

274275
// act
275276
_credentialManager.SaveCredentials("key", "token");
276277
var result = _credentialManager.GetCredentials("key");
277278

278279
// assert
279280
Assert.AreEqual("token", result);
280-
Assert.AreEqual(FileAccessPermissions.UserReadWriteExecute, UnixOperations.Instance.GetDirectoryInfo(tempDirectory).Permissions);
281+
Assert.AreEqual(insecurePermissions, UnixOperations.Instance.GetDirectoryInfo(tempDirectory).Permissions);
281282
}
282283
finally
283284
{

0 commit comments

Comments
 (0)