Skip to content

Commit

Permalink
SNOW-955536: Multiple SAML Integration (#852)
Browse files Browse the repository at this point in the history 
### Description
Add multiple SAML Integration

### Checklist
- [x] Code compiles correctly
- [x] Code is formatted according to [Coding
Conventions](../CodingConventions.md)
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing (`dotnet test`)
- [x] Extended the README / documentation, if necessary
- [x] Provide JIRA issue id (if possible) or GitHub issue id in PR name
  • Loading branch information
@sfc-gh-ext-simba-lf
sfc-gh-ext-simba-lf authored Jan 23, 2024
1 parent 7d478f0 commit fa09941
Show file tree
Hide file tree
Showing 6 changed files with 149 additions and 17 deletions.
44 changes: 44 additions & 0 deletions Snowflake.Data.Tests/IntegrationTests/SFConnectionIT.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -894,6 +894,50 @@ public void TestSSOConnectionWithUserAsync()
}
}

[Test]
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
public void TestSSOConnectionWithUserAndDisableConsoleLogin()
{
// Use external browser to log in using proper password for [email protected]
using (IDbConnection conn = new SnowflakeDbConnection())
{
conn.ConnectionString
= ConnectionStringWithoutAuth
+ ";authenticator=externalbrowser;[email protected];disable_console_login=false;";
conn.Open();
Assert.AreEqual(ConnectionState.Open, conn.State);
using (IDbCommand command = conn.CreateCommand())
{
command.CommandText = "SELECT CURRENT_USER()";
Assert.AreEqual("QA", command.ExecuteScalar().ToString());
}
}
}

[Test]
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
public void TestSSOConnectionWithUserAsyncAndDisableConsoleLogin()
{
// Use external browser to log in using proper password for [email protected]
using (SnowflakeDbConnection conn = new SnowflakeDbConnection())
{
conn.ConnectionString
= ConnectionStringWithoutAuth
+ ";authenticator=externalbrowser;[email protected];disable_console_login=false;";

Task connectTask = conn.OpenAsync(CancellationToken.None);
connectTask.Wait();
Assert.AreEqual(ConnectionState.Open, conn.State);
using (DbCommand command = conn.CreateCommand())
{
command.CommandText = "SELECT CURRENT_USER()";
Task<object> task = command.ExecuteScalarAsync(CancellationToken.None);
task.Wait(CancellationToken.None);
Assert.AreEqual("QA", task.Result);
}
}
}

[Test]
[Ignore("This test requires manual interaction and therefore cannot be run in CI")]
public void TestSSOConnectionTimeoutAfter10s()
Expand Down
42 changes: 42 additions & 0 deletions Snowflake.Data.Tests/UnitTests/SFSessionPropertyTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
string defMaxHttpRetries = "7";
string defIncludeRetryReason = "true";
string defDisableQueryContextCache = "false";
string defDisableConsoleLogin = "true";
string defAllowUnderscoresInHost = "false";

var simpleTestCase = new TestCase()
Expand Down Expand Up @@ -105,6 +106,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -132,6 +134,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -162,6 +165,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
},
ConnectionString =
Expand Down Expand Up @@ -194,6 +198,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
},
ConnectionString =
Expand Down Expand Up @@ -225,6 +230,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.FILE_TRANSFER_MEMORY_THRESHOLD, "25" },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -253,6 +259,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, "false" },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -280,11 +287,42 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, "true" },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
},
ConnectionString =
$"ACCOUNT={defAccount};USER={defUser};PASSWORD={defPassword};DISABLEQUERYCONTEXTCACHE=true"
};
var testCaseWithDisableConsoleLogin = new TestCase()
{
ExpectedProperties = new SFSessionProperties()
{
{ SFSessionProperty.ACCOUNT, defAccount },
{ SFSessionProperty.USER, defUser },
{ SFSessionProperty.HOST, defHost },
{ SFSessionProperty.AUTHENTICATOR, defAuthenticator },
{ SFSessionProperty.SCHEME, defScheme },
{ SFSessionProperty.CONNECTION_TIMEOUT, defConnectionTimeout },
{ SFSessionProperty.PASSWORD, defPassword },
{ SFSessionProperty.PORT, defPort },
{ SFSessionProperty.VALIDATE_DEFAULT_PARAMETERS, "true" },
{ SFSessionProperty.USEPROXY, "false" },
{ SFSessionProperty.INSECUREMODE, "false" },
{ SFSessionProperty.DISABLERETRY, "false" },
{ SFSessionProperty.FORCERETRYON404, "false" },
{ SFSessionProperty.CLIENT_SESSION_KEEP_ALIVE, "false" },
{ SFSessionProperty.FORCEPARSEERROR, "false" },
{ SFSessionProperty.BROWSER_RESPONSE_TIMEOUT, defBrowserResponseTime },
{ SFSessionProperty.RETRY_TIMEOUT, defRetryTimeout },
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, "false" },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
},
ConnectionString =
$"ACCOUNT={defAccount};USER={defUser};PASSWORD={defPassword};DISABLE_CONSOLE_LOGIN=false"
};
var complicatedAccount = $"{defAccount}.region-name.host-name";
var testCaseComplicatedAccountName = new TestCase()
{
Expand All @@ -311,6 +349,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -339,6 +378,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, defAllowUnderscoresInHost }
}
};
Expand Down Expand Up @@ -367,6 +407,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
{ SFSessionProperty.MAXHTTPRETRIES, defMaxHttpRetries },
{ SFSessionProperty.INCLUDERETRYREASON, defIncludeRetryReason },
{ SFSessionProperty.DISABLEQUERYCONTEXTCACHE, defDisableQueryContextCache },
{ SFSessionProperty.DISABLE_CONSOLE_LOGIN, defDisableConsoleLogin },
{ SFSessionProperty.ALLOWUNDERSCORESINHOST, "true" }
}
};
Expand All @@ -379,6 +420,7 @@ public static IEnumerable<TestCase> ConnectionStringTestCases()
testCaseWithFileTransferMaxBytesInMemory,
testCaseWithIncludeRetryReason,
testCaseWithDisableQueryContextCache,
testCaseWithDisableConsoleLogin,
testCaseComplicatedAccountName,
testCaseUnderscoredAccountName,
testCaseUnderscoredAccountNameWithEnabledAllowUnderscores
Expand Down
73 changes: 56 additions & 17 deletions Snowflake.Data/Core/Authenticator/ExternalBrowserAuthenticator.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
using Snowflake.Data.Log;
using Snowflake.Data.Client;
using System.Text.RegularExpressions;
using System.Collections.Generic;

namespace Snowflake.Data.Core.Authenticator
{
Expand Down Expand Up @@ -54,19 +55,28 @@ async Task IAuthenticator.AuthenticateAsync(CancellationToken cancellationToken)
httpListener.Start();

logger.Debug("Get IdpUrl and ProofKey");
var authenticatorRestRequest = BuildAuthenticatorRestRequest(localPort);
var authenticatorRestResponse =
await session.restRequester.PostAsync<AuthenticatorResponse>(
authenticatorRestRequest,
cancellationToken
).ConfigureAwait(false);
authenticatorRestResponse.FilterFailedResponse();
string loginUrl;
if (session._disableConsoleLogin)
{
var authenticatorRestRequest = BuildAuthenticatorRestRequest(localPort);
var authenticatorRestResponse =
await session.restRequester.PostAsync<AuthenticatorResponse>(
authenticatorRestRequest,
cancellationToken
).ConfigureAwait(false);
authenticatorRestResponse.FilterFailedResponse();

var idpUrl = authenticatorRestResponse.data.ssoUrl;
_proofKey = authenticatorRestResponse.data.proofKey;
loginUrl = authenticatorRestResponse.data.ssoUrl;
_proofKey = authenticatorRestResponse.data.proofKey;
}
else
{
_proofKey = GenerateProofKey();
loginUrl = GetLoginUrl(_proofKey, localPort);
}

logger.Debug("Open browser");
StartBrowser(idpUrl);
StartBrowser(loginUrl);

logger.Debug("Get the redirect SAML request");
_successEvent = new ManualResetEvent(false);
Expand Down Expand Up @@ -96,15 +106,24 @@ void IAuthenticator.Authenticate()
httpListener.Start();

logger.Debug("Get IdpUrl and ProofKey");
var authenticatorRestRequest = BuildAuthenticatorRestRequest(localPort);
var authenticatorRestResponse = session.restRequester.Post<AuthenticatorResponse>(authenticatorRestRequest);
authenticatorRestResponse.FilterFailedResponse();
string loginUrl;
if (session._disableConsoleLogin)
{
var authenticatorRestRequest = BuildAuthenticatorRestRequest(localPort);
var authenticatorRestResponse = session.restRequester.Post<AuthenticatorResponse>(authenticatorRestRequest);
authenticatorRestResponse.FilterFailedResponse();

var idpUrl = authenticatorRestResponse.data.ssoUrl;
_proofKey = authenticatorRestResponse.data.proofKey;
loginUrl = authenticatorRestResponse.data.ssoUrl;
_proofKey = authenticatorRestResponse.data.proofKey;
}
else
{
_proofKey = GenerateProofKey();
loginUrl = GetLoginUrl(_proofKey, localPort);
}

logger.Debug("Open browser");
StartBrowser(idpUrl);
StartBrowser(loginUrl);

logger.Debug("Get the redirect SAML request");
_successEvent = new ManualResetEvent(false);
Expand Down Expand Up @@ -187,7 +206,7 @@ private static void StartBrowser(string url)
// The following code is learnt from https://brockallen.com/2016/09/24/process-start-for-urls-on-net-core/
#if NETFRAMEWORK
// .net standard would pass here
Process.Start(url);
Process.Start(new ProcessStartInfo(url) { UseShellExecute = true });
#else
// hack because of this: https://github.com/dotnet/corefx/issues/10361
if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
Expand Down Expand Up @@ -247,5 +266,25 @@ protected override void SetSpecializedAuthenticatorData(ref LoginRequestData dat
data.Token = _samlResponseToken;
data.ProofKey = _proofKey;
}

private string GetLoginUrl(string proofKey, int localPort)
{
Dictionary<string, string> parameters = new Dictionary<string, string>()
{
{ "login_name", session.properties[SFSessionProperty.USER]},
{ "proof_key", proofKey },
{ "browser_mode_redirect_port", localPort.ToString() }
};
Uri loginUrl = session.BuildUri(RestPath.SF_CONSOLE_LOGIN, parameters);
return loginUrl.ToString();
}

private string GenerateProofKey()
{
Random rnd = new Random();
Byte[] randomness = new Byte[32];
rnd.NextBytes(randomness);
return Convert.ToBase64String(randomness);
}
}
}
2 changes: 2 additions & 0 deletions Snowflake.Data/Core/RestParams.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,8 @@ internal static class RestPath
internal const string SF_QUERY_PATH = "/queries/v1/query-request";

internal const string SF_SESSION_HEARTBEAT_PATH = SF_SESSION_PATH + "/heartbeat";

internal const string SF_CONSOLE_LOGIN = "/console/login";
}

internal class SFEnvironment
Expand Down
3 changes: 3 additions & 0 deletions Snowflake.Data/Core/Session/SFSession.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,8 @@ public class SFSession

private bool _disableQueryContextCache = false;

internal bool _disableConsoleLogin;

internal void ProcessLoginResponse(LoginResponse authnResponse)
{
if (authnResponse.success)
Expand Down Expand Up @@ -148,6 +150,7 @@ internal SFSession(
connStr = connectionString;
properties = SFSessionProperties.parseConnectionString(connectionString, password);
_disableQueryContextCache = bool.Parse(properties[SFSessionProperty.DISABLEQUERYCONTEXTCACHE]);
_disableConsoleLogin = bool.Parse(properties[SFSessionProperty.DISABLE_CONSOLE_LOGIN]);
ValidateApplicationName(properties);
try
{
Expand Down
2 changes: 2 additions & 0 deletions Snowflake.Data/Core/Session/SFSessionProperty.cs
100755 → 100644
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,8 @@ internal enum SFSessionProperty
DISABLEQUERYCONTEXTCACHE,
[SFSessionPropertyAttr(required = false)]
CLIENT_CONFIG_FILE,
[SFSessionPropertyAttr(required = false, defaultValue = "true")]
DISABLE_CONSOLE_LOGIN,
[SFSessionPropertyAttr(required = false, defaultValue = "false")]
ALLOWUNDERSCORESINHOST
}
Expand Down

0 comments on commit fa09941

Please sign in to comment.