Moved logic to secure credential manager key when created instead of …

…each credential manager implementation.
snowflakedb · Nov 16, 2024 · fa7c5c3 · fa7c5c3
1 parent 66227de
commit fa7c5c3
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 23 deletions.
diff --git a/Snowflake.Data/Client/SnowflakeCredentialManagerFactory.cs b/Snowflake.Data/Client/SnowflakeCredentialManagerFactory.cs
@@ -6,6 +6,7 @@
 using Snowflake.Data.Core;
 using Snowflake.Data.Core.CredentialManager;
 using Snowflake.Data.Core.CredentialManager.Infrastructure;
+using Snowflake.Data.Core.Tools;
 using Snowflake.Data.Log;
 
 namespace Snowflake.Data.Client
@@ -19,11 +20,17 @@ public class SnowflakeCredentialManagerFactory
 
         private static ISnowflakeCredentialManager s_credentialManager;
 
-        internal static string BuildCredentialKey(string host, string user, TokenType tokenType, string authenticator = null)
+        private static string BuildCredentialKey(string host, string user, TokenType tokenType, string authenticator = null)
         {
             return $"{host.ToUpper()}:{user.ToUpper()}:{SFEnvironment.DriverName}:{tokenType.ToString().ToUpper()}:{authenticator?.ToUpper() ?? string.Empty}";
         }
 
+        internal static string GetSecureCredentialKey(string host, string user, TokenType tokenType, string authenticator = null)
+        {
+            return BuildCredentialKey(host, user, tokenType, authenticator).ToSha256Hash();
+        }
+
+
         public static void UseDefaultCredentialManager()
         {
             SetCredentialManager(GetDefaultCredentialManager());

diff --git a/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerFileImpl.cs b/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerFileImpl.cs
@@ -113,8 +113,7 @@ public string GetCredentials(string key)
             if (_fileOperations.Exists(_jsonCacheFilePath))
             {
                 var keyTokenPairs = ReadJsonFile();
-                var hashKey = key.ToSha256Hash();
-                if (keyTokenPairs.TryGetValue(hashKey, out string token))
+                if (keyTokenPairs.TryGetValue(key, out string token))
                 {
                     return token;
                 }
@@ -130,18 +129,16 @@ public void RemoveCredentials(string key)
             if (_fileOperations.Exists(_jsonCacheFilePath))
             {
                 var keyTokenPairs = ReadJsonFile();
-                var hashKey = key.ToSha256Hash();
-                keyTokenPairs.Remove(hashKey);
+                keyTokenPairs.Remove(key);
                 WriteToJsonFile(JsonConvert.SerializeObject(keyTokenPairs));
             }
         }
 
         public void SaveCredentials(string key, string token)
         {
             s_logger.Debug($"Saving credentials to json file in {_jsonCacheFilePath} for key: {key}");
-            var hashKey = key.ToSha256Hash();
             KeyTokenDict keyTokenPairs = _fileOperations.Exists(_jsonCacheFilePath) ? ReadJsonFile() : new KeyTokenDict();
-            keyTokenPairs[hashKey] = token;
+            keyTokenPairs[key] = token;
 
             string jsonString = JsonConvert.SerializeObject(keyTokenPairs);
             WriteToJsonFile(jsonString);

diff --git a/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerInMemoryImpl.cs b/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerInMemoryImpl.cs
@@ -22,8 +22,7 @@ internal class SFCredentialManagerInMemoryImpl : ISnowflakeCredentialManager
         public string GetCredentials(string key)
         {
             s_logger.Debug($"Getting credentials from memory for key: {key}");
-            var hashKey = key.ToSha256Hash();
-            if (s_credentials.TryGetValue(hashKey, out var secureToken))
+            if (s_credentials.TryGetValue(key, out var secureToken))
             {
                 return SecureStringHelper.Decode(secureToken);
             }
@@ -36,16 +35,14 @@ public string GetCredentials(string key)
 
         public void RemoveCredentials(string key)
         {
-            var hashKey = key.ToSha256Hash();
             s_logger.Debug($"Removing credentials from memory for key: {key}");
-            s_credentials.Remove(hashKey);
+            s_credentials.Remove(key);
         }
 
         public void SaveCredentials(string key, string token)
         {
-            var hashKey = key.ToSha256Hash();
-            s_logger.Debug($"Saving credentials into memory for key: {hashKey}");
-            s_credentials[hashKey] = SecureStringHelper.Encode(token);
+            s_logger.Debug($"Saving credentials into memory for key: {key}");
+            s_credentials[key] = SecureStringHelper.Encode(token);
         }
     }
 }
diff --git a/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerWindowsNativeImpl.cs b/Snowflake.Data/Core/CredentialManager/Infrastructure/SFCredentialManagerWindowsNativeImpl.cs
@@ -22,9 +22,8 @@ internal class SFCredentialManagerWindowsNativeImpl : ISnowflakeCredentialManage
         public string GetCredentials(string key)
         {
             s_logger.Debug($"Getting the credentials for key: {key}");
-            var hashKey = key.ToSha256Hash();
             IntPtr nCredPtr;
-            if (!CredRead(hashKey, 1 /* Generic */, 0, out nCredPtr))
+            if (!CredRead(key, 1 /* Generic */, 0, out nCredPtr))
             {
                 s_logger.Info($"Unable to get credentials for key: {key}");
                 return "";
@@ -41,8 +40,7 @@ public void RemoveCredentials(string key)
         {
             s_logger.Debug($"Removing the credentials for key: {key}");
 
-            var hashKey = key.ToSha256Hash();
-            if (!CredDelete(hashKey, 1 /* Generic */, 0))
+            if (!CredDelete(key, 1 /* Generic */, 0))
             {
                 s_logger.Info($"Unable to remove credentials because the specified key did not exist: {key}");
             }
@@ -51,7 +49,6 @@ public void RemoveCredentials(string key)
         public void SaveCredentials(string key, string token)
         {
             s_logger.Debug($"Saving the credentials for key: {key}");
-            var hashKey = key.ToSha256Hash();
             byte[] byteArray = Encoding.Unicode.GetBytes(token);
             Credential credential = new Credential();
             credential.AttributeCount = 0;
@@ -61,7 +58,7 @@ public void SaveCredentials(string key, string token)
             credential.Type = 1; // Generic
             credential.Persist = 2; // Local Machine
             credential.CredentialBlobSize = (uint)(byteArray == null ? 0 : byteArray.Length);
-            credential.TargetName = hashKey;
+            credential.TargetName = key;
             credential.CredentialBlob = token;
             credential.UserName = Environment.UserName;
 

diff --git a/Snowflake.Data/Core/Session/SFSession.cs b/Snowflake.Data/Core/Session/SFSession.cs
@@ -124,7 +124,7 @@ internal void ProcessLoginResponse(LoginResponse authnResponse)
                 if (!string.IsNullOrEmpty(authnResponse.data.mfaToken))
                 {
                     _mfaToken = SecureStringHelper.Encode(authnResponse.data.mfaToken);
-                    var key = SnowflakeCredentialManagerFactory.BuildCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, properties[SFSessionProperty.AUTHENTICATOR]);
+                    var key = SnowflakeCredentialManagerFactory.GetSecureCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, properties[SFSessionProperty.AUTHENTICATOR]);
                     SnowflakeCredentialManagerFactory.GetCredentialManager().SaveCredentials(key, authnResponse.data.mfaToken);
                 }
                 logger.Debug($"Session opened: {sessionId}");
@@ -143,7 +143,7 @@ internal void ProcessLoginResponse(LoginResponse authnResponse)
                 {
                     logger.Info($"Unable to use cached MFA token is expired or invalid. Fails with the {e.Message}. ", e);
                     _mfaToken = null;
-                    var mfaKey = SnowflakeCredentialManagerFactory.BuildCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, properties[SFSessionProperty.AUTHENTICATOR]);
+                    var mfaKey = SnowflakeCredentialManagerFactory.GetSecureCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, properties[SFSessionProperty.AUTHENTICATOR]);
                     SnowflakeCredentialManagerFactory.GetCredentialManager().RemoveCredentials(mfaKey);
                 }
 
@@ -215,7 +215,7 @@ internal SFSession(
 
                 if (properties.TryGetValue(SFSessionProperty.AUTHENTICATOR, out var _authenticatorType) &&  _authenticatorType == "username_password_mfa")
                 {
-                    var mfaKey = SnowflakeCredentialManagerFactory.BuildCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, _authenticatorType);
+                    var mfaKey = SnowflakeCredentialManagerFactory.GetSecureCredentialKey(properties[SFSessionProperty.HOST], properties[SFSessionProperty.USER], TokenType.MFAToken, _authenticatorType);
                     _mfaToken = SecureStringHelper.Encode(SnowflakeCredentialManagerFactory.GetCredentialManager().GetCredentials(mfaKey));
                 }
             }