CSHARP-1977: Cache SCRAM ClientKey

Author: ksadoff

src/MongoDB.Driver.Core/Core/Authentication/ScramCache.cs

@@ -0,0 +1,163 @@
+/* Copyright 2019–present MongoDB Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+using System;
+using System.Linq;
+using System.Security;
+using MongoDB.Driver.Core.Misc;
+using MongoDB.Shared;
+
+namespace MongoDB.Driver.Core.Authentication
+{
+    /// <summary>
+    /// A cache for Client and Server keys, to be used during authentication.
+    /// </summary>
+    internal class ScramCache
+    {
+        private ScramCacheKey _cacheKey;
+        private ScramCacheEntry _cachedEntry;
+
+        /// <summary>
+        /// Try to get a cached entry.
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="entry"></param>
+        /// <returns></returns>
+        public bool TryGet(ScramCacheKey key, out ScramCacheEntry entry)
+        {
+            if (key.Equals(_cacheKey))
+            {
+                entry = _cachedEntry;
+                return true;
+            }
+            else
+            {
+                entry = null;
+                return false;
+            }
+        }
+
+        /// <summary>
+        /// Add a cached entry.
+        /// </summary>
+        /// <param name="key"></param>
+        /// <param name="entry"></param>
+        public void Add(ScramCacheKey key, ScramCacheEntry entry)
+        {
+            _cacheKey = key;
+            _cachedEntry = entry;
+        }
+    }
+
+    internal class ScramCacheKey
+    {
+        private int _iterationCount;
+        private SecureString _password;
+        private byte[] _salt;
+
+        internal ScramCacheKey(SecureString password, byte[] salt, int iterationCount)
+        {
+            _iterationCount = iterationCount;
+            _password = password;
+            _salt = salt;
+        }
+
+        public override bool Equals(object obj)
+        {
+            if (this == obj)
+            {
+                return true;
+            }
+
+            if (obj == null || obj.GetType() != obj.GetType())
+            {
+                return false;
+            }
+
+            ScramCacheKey other = (ScramCacheKey) obj;
+
+            return
+                Equals(_password,other._password) &&
+                _iterationCount == other._iterationCount &&
+                _salt.SequenceEqual(other._salt);
+        }
+
+        public override int GetHashCode()
+        {
+            // ignore _password when computing the hash code
+            return new Hasher()
+                .Hash(_iterationCount)
+                .Hash(_salt)
+                .GetHashCode();
+        }
+
+        // private methods
+        private bool Equals(SecureString x, SecureString y)
+        {
+            if (object.ReferenceEquals(x, y))
+            {
+                return true;
+            }
+
+            if (object.ReferenceEquals(x, null) || object.ReferenceEquals(y, null))
+            {
+                return false;
+
+            }
+            using (var dx = new DecryptedSecureString(x))
+            using (var dy = new DecryptedSecureString(y))
+            {
+                var xchars = dx.GetChars();
+                var ychars = dy.GetChars();
+                return Equals(xchars, ychars);
+            }
+        }
+
+        private bool Equals(char[] x, char[] y)
+        {
+            if (x.Length != y.Length)
+            {
+                return false;
+            }
+
+            for (var i = 0; i < x.Length; i++)
+            {
+                if (x[i] != y[i])
+                {
+                    return false;
+                }
+            }
+
+            return true;
+        }
+
+    }
+
+    internal class ScramCacheEntry
+    {
+        private byte[] _clientKey;
+        private byte[] _serverKey;
+
+        public ScramCacheEntry(byte[] clientKey, byte[] serverKey)
+        {
+            _clientKey = clientKey;
+            _serverKey = serverKey;
+        }
+
+        public byte[] ClientKey => _clientKey;
+
+        public byte[] ServerKey => _serverKey;
+    }
+}

src/MongoDB.Driver.Core/Core/Authentication/ScramSha1Authenticator.cs

@@ -55,7 +55,7 @@ public ScramSha1Authenticator(UsernamePasswordCredential credential)
         }
 
         internal ScramSha1Authenticator(UsernamePasswordCredential credential, IRandomStringGenerator randomStringGenerator)
-            : base(credential, HashAlgorithmName.SHA1, randomStringGenerator, H1, Hi1, Hmac1) 
+            : base(credential, HashAlgorithmName.SHA1, randomStringGenerator, H1, Hi1, Hmac1, new ScramCache()) 
         {
         }
 

src/MongoDB.Driver.Core/Core/Authentication/ScramSha256Authenticator.cs

@@ -60,7 +60,7 @@ public ScramSha256Authenticator(UsernamePasswordCredential credential)
         }
 
         internal ScramSha256Authenticator(UsernamePasswordCredential credential, IRandomStringGenerator randomStringGenerator)
-            : base(credential, HashAlgorithmName.SHA256, randomStringGenerator, H256, Hi256, Hmac256)
+            : base(credential, HashAlgorithmName.SHA256, randomStringGenerator, H256, Hi256, Hmac256, new ScramCache())
         {
         }
 

src/MongoDB.Driver.Core/Core/Authentication/ScramShaAuthenticator.cs

@@ -1,4 +1,4 @@
-/* Copyright 2018–present MongoDB Inc.
+/* Copyright 2018–present MongoDB Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -72,8 +72,8 @@ protected ScramShaAuthenticator(UsernamePasswordCredential credential,
             H h,
             Hi hi,
             Hmac hmac)
-            : this(credential, hashAlgorithmName, new DefaultRandomStringGenerator(), h, hi, hmac) { }
-        
+            : this(credential, hashAlgorithmName, new DefaultRandomStringGenerator(), h, hi, hmac, new ScramCache()) { }
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ScramShaAuthenticator"/> class.
         /// </summary>
@@ -83,14 +83,16 @@ protected ScramShaAuthenticator(UsernamePasswordCredential credential,
         /// <param name="h">The H function to use.</param>
         /// <param name="hi">The Hi function to use.</param>
         /// <param name="hmac">The Hmac function to use.</param>
+        /// <param name="cache">The cache to use.</param>
         internal ScramShaAuthenticator(
             UsernamePasswordCredential credential, 
             HashAlgorithmName hashAlgorithName,
             IRandomStringGenerator randomStringGenerator,
             H h,
             Hi hi,
-            Hmac hmac)
-            : base(new ScramShaMechanism(credential, hashAlgorithName, randomStringGenerator, h, hi, hmac))
+            Hmac hmac,
+            ScramCache cache)
+            : base(new ScramShaMechanism(credential, hashAlgorithName, randomStringGenerator, h, hi, hmac, cache))
         {
             _databaseName = credential.Source;
         }
@@ -108,14 +110,16 @@ private class ScramShaMechanism : ISaslMechanism
             private readonly Hi _hi;
             private readonly Hmac _hmac;
             private readonly string _name;
+            private ScramCache _cache;
 
             public ScramShaMechanism(
                 UsernamePasswordCredential credential, 
                 HashAlgorithmName hashAlgorithmName, 
                 IRandomStringGenerator randomStringGenerator,
                 H h,
                 Hi hi,
-                Hmac hmac)
+                Hmac hmac,
+                ScramCache cache)
             {
                 _credential = Ensure.IsNotNull(credential, nameof(credential));
                 _h = h;
@@ -127,6 +131,7 @@ public ScramShaMechanism(
                 }
                 _name = $"SCRAM-SHA-{hashAlgorithmName.ToString().Substring(3)}";
                 _randomStringGenerator = Ensure.IsNotNull(randomStringGenerator, nameof(randomStringGenerator));
+                _cache = cache;
             }
 
             public string Name => _name;
@@ -143,9 +148,9 @@ public ISaslStep Initialize(IConnection connection, SaslConversation conversatio
 
                 var clientFirstMessageBare = username + "," + nonce;
                 var clientFirstMessage = gs2Header + clientFirstMessageBare;
-                var clientFirstMessageBytes = Utf8Encodings.Strict.GetBytes(clientFirstMessage); 
+                var clientFirstMessageBytes = Utf8Encodings.Strict.GetBytes(clientFirstMessage);
 
-                return new ClientFirst(clientFirstMessageBytes, clientFirstMessageBare, _credential, r, _h, _hi, _hmac);
+                return new ClientFirst(clientFirstMessageBytes, clientFirstMessageBare, _credential, r, _h, _hi, _hmac, _cache);
             }
 
             private string GenerateRandomString()
@@ -163,7 +168,7 @@ private string PrepUsername(string username)
 
         private class ClientFirst : ISaslStep
         {
-            
+
             private readonly byte[] _bytesToSendToServer;
             private readonly string _clientFirstMessageBare;
             private readonly UsernamePasswordCredential _credential;
@@ -173,14 +178,17 @@ private class ClientFirst : ISaslStep
             private readonly Hi _hi;
             private readonly Hmac _hmac;
 
+            private ScramCache _cache;
+
             public ClientFirst(
                 byte[] bytesToSendToServer, 
                 string clientFirstMessageBare, 
                 UsernamePasswordCredential credential, 
                 string rPrefix,
                 H h,
                 Hi hi,
-                Hmac hmac)
+                Hmac hmac,
+                ScramCache cache)
             {
                 _bytesToSendToServer = bytesToSendToServer;
                 _clientFirstMessageBare = clientFirstMessageBare;
@@ -189,6 +197,7 @@ public ClientFirst(
                 _hi = hi;
                 _hmac = hmac;
                 _rPrefix = rPrefix;
+                _cache = cache;
             }
 
             public byte[] BytesToSendToServer => _bytesToSendToServer;
@@ -214,19 +223,31 @@ public ISaslStep Transition(SaslConversation conversation, byte[] bytesReceivedF
                 var nonce = "r=" + r;
                 var clientFinalMessageWithoutProof = channelBinding + "," + nonce;
 
-                var saltedPassword = _hi(
-                    _credential,
-                    Convert.FromBase64String(s),
-                    int.Parse(i));
+                var salt = Convert.FromBase64String(map['s']);
+                var iterations = int.Parse(map['i']);
+
+                byte[] clientKey;
+                byte[] serverKey;
+
+                var cacheKey = new ScramCacheKey(_credential.SaslPreppedPassword, salt, iterations);
+                if (_cache.TryGet(cacheKey, out var cacheEntry))
+                {
+                    clientKey = cacheEntry.ClientKey;
+                    serverKey = cacheEntry.ServerKey;
+                }
+                else
+                {
+                    var saltedPassword = _hi( _credential, salt, iterations);
+                    clientKey = _hmac(encoding, saltedPassword, "Client Key");
+                    serverKey = _hmac(encoding, saltedPassword, "Server Key");
+                    _cache.Add(cacheKey, new ScramCacheEntry(clientKey, serverKey));
+                }
 
-                var clientKey = _hmac(encoding, saltedPassword, "Client Key");
                 var storedKey = _h(clientKey);
                 var authMessage = _clientFirstMessageBare + "," + serverFirstMessage + "," + clientFinalMessageWithoutProof;
                 var clientSignature = _hmac(encoding, storedKey, authMessage);
                 var clientProof = XOR(clientKey, clientSignature);
-                var serverKey = _hmac(encoding, saltedPassword, "Server Key");
                 var serverSignature = _hmac(encoding, serverKey, authMessage);
-
                 var proof = "p=" + Convert.ToBase64String(clientProof);
                 var clientFinalMessage = clientFinalMessageWithoutProof + "," + proof;
 
@@ -243,7 +264,6 @@ private byte[] XOR(byte[] a, byte[] b)
 
                 return result;
             }
-
         }
 
         private class ClientLast : ISaslStep
@@ -287,4 +307,4 @@ private bool ConstantTimeEquals(byte[] a, byte[] b)
             }
         }
     }
-}
+}

src/MongoDB.Driver.Core/Core/Connections/CommandEventHelper.cs

@@ -1038,6 +1038,7 @@ private BsonDocument BuildFindCommandFromQuery(QueryMessage message)
                         default:
                             if (element.Name.StartsWith("$"))
                             {
+                                // should we actually remove the $ or not?
                                 command[element.Name.Substring(1)] = element.Value;
                             }
                             else