CSHARP-3032: Support speculative SCRAM-SHA authentication attempts in isMaster

Author: vincentkam

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

@@ -35,6 +35,7 @@ public class DefaultAuthenticator : IAuthenticator
         // fields
         private readonly UsernamePasswordCredential _credential;
         private readonly IRandomStringGenerator _randomStringGenerator;
+        private IAuthenticator _speculativeAuthenticator;
 
         // constructors
         /// <summary>
@@ -79,9 +80,8 @@ public void Authenticate(IConnection connection, ConnectionDescription descripti
                     description.BuildInfoResult);
             }
 
-            var authenticator = CreateAuthenticator(connection, description);
+            var authenticator = GetOrCreateAuthenticator(connection, description);
             authenticator.Authenticate(connection, description, cancellationToken);
-            
         }
 
         /// <inheritdoc/>
@@ -106,15 +106,18 @@ public async Task AuthenticateAsync(IConnection connection, ConnectionDescriptio
                     description.BuildInfoResult);
             }  
 
-            var authenticator = CreateAuthenticator(connection, description);
+            var authenticator = GetOrCreateAuthenticator(connection, description);
             await authenticator.AuthenticateAsync(connection, description, cancellationToken).ConfigureAwait(false);
         }
 
 
         /// <inheritdoc/>
         public BsonDocument CustomizeInitialIsMasterCommand(BsonDocument isMasterCommand)
         {
-            return isMasterCommand.Merge(CreateSaslSupportedMechsRequest(_credential.Source, _credential.Username)); 
+            var saslSupportedMechs = CreateSaslSupportedMechsRequest(_credential.Source, _credential.Username);
+            isMasterCommand = isMasterCommand.Merge(saslSupportedMechs);
+            _speculativeAuthenticator = new ScramSha256Authenticator(_credential, _randomStringGenerator);
+            return _speculativeAuthenticator.CustomizeInitialIsMasterCommand(isMasterCommand);
         }
 
         private static BsonDocument CreateSaslSupportedMechsRequest(string authenticationDatabaseName, string userName)
@@ -143,5 +146,14 @@ private IAuthenticator CreateAuthenticator(IConnection connection, ConnectionDes
                     : new MongoDBCRAuthenticator(_credential);
 #pragma warning restore 618
         }
+
+        private IAuthenticator GetOrCreateAuthenticator(IConnection connection, ConnectionDescription description)
+        {
+            /* It is possible to have for IsMaster["SpeculativeAuthenticate"] != null and for
+             * _speculativeScramSha256Authenticator to be null in the case of multiple authenticators */
+            var speculativeAuthenticateResult = description.IsMasterResult.SpeculativeAuthenticate;
+            var canUseSpeculativeAuthenticator = _speculativeAuthenticator != null && speculativeAuthenticateResult != null;
+            return canUseSpeculativeAuthenticator ? _speculativeAuthenticator : CreateAuthenticator(connection, description);
+        }
     }
-}
+}

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

@@ -31,7 +31,8 @@ namespace MongoDB.Driver.Core.Authentication
     public abstract class SaslAuthenticator : IAuthenticator
     {
         // fields
-        private readonly ISaslMechanism _mechanism;
+        private protected readonly ISaslMechanism _mechanism;
+        private protected ISaslStep _speculativeFirstStep;
 
         // constructors
         /// <summary>
@@ -67,10 +68,20 @@ public void Authenticate(IConnection connection, ConnectionDescription descripti
 
             using (var conversation = new SaslConversation(description.ConnectionId))
             {
-                var currentStep = _mechanism.Initialize( connection, conversation, description);
+                ISaslStep currentStep;
+                BsonDocument command;
+                var speculativeAuthenticateResult = description.IsMasterResult.SpeculativeAuthenticate;
+                if (_speculativeFirstStep != null && speculativeAuthenticateResult != null)
+                {
+                    currentStep = Transition(conversation, _speculativeFirstStep, speculativeAuthenticateResult, out command);
+                }
+                else
+                {
+                    currentStep = _mechanism.Initialize(connection, conversation, description);
+                    command = CreateStartCommand(currentStep);
+                }
 
-                var command = CreateStartCommand(currentStep);
-                while (true)
+                while (currentStep != null)
                 {
                     BsonDocument result;
                     try
@@ -83,13 +94,7 @@ public void Authenticate(IConnection connection, ConnectionDescription descripti
                         throw CreateException(connection, ex);
                     }
 
-                    currentStep = Transition(conversation, currentStep, result);
-                    if (currentStep == null)
-                    {
-                        return;
-                    }
-
-                    command = CreateContinueCommand(currentStep, result);
+                    currentStep = Transition(conversation, currentStep, result , out command);
                 }
             }
         }
@@ -102,10 +107,20 @@ public async Task AuthenticateAsync(IConnection connection, ConnectionDescriptio
 
             using (var conversation = new SaslConversation(description.ConnectionId))
             {
-                var currentStep = _mechanism.Initialize(connection, conversation, description);
+                ISaslStep currentStep;
+                BsonDocument command;
+                var speculativeAuthenticateResult = description.IsMasterResult.SpeculativeAuthenticate;
+                if (_speculativeFirstStep != null && speculativeAuthenticateResult != null)
+                {
+                    currentStep = Transition(conversation, _speculativeFirstStep, speculativeAuthenticateResult, out command);
+                }
+                else
+                {
+                    currentStep = _mechanism.Initialize(connection, conversation, description);
+                    command = CreateStartCommand(currentStep);
+                }
 
-                var command = CreateStartCommand(currentStep);
-                while (true)
+                while (currentStep != null)
                 {
                     BsonDocument result;
                     try
@@ -118,23 +133,29 @@ public async Task AuthenticateAsync(IConnection connection, ConnectionDescriptio
                         throw CreateException(connection, ex);
                     }
 
-                    currentStep = Transition(conversation, currentStep, result);
-                    if (currentStep == null)
-                    {
-                        return;
-                    }
-
-                    command = CreateContinueCommand(currentStep, result);
+                    currentStep = Transition(conversation, currentStep, result, out command);
                 }
             }
         }
 
         /// <inheritdoc/>
-        public BsonDocument CustomizeInitialIsMasterCommand(BsonDocument isMasterCommand)
+        public virtual BsonDocument CustomizeInitialIsMasterCommand(BsonDocument isMasterCommand)
         {
             return isMasterCommand;
         }
 
+        private protected virtual BsonDocument CreateStartCommand(ISaslStep currentStep)
+        {
+            var startCommand = new BsonDocument
+            {
+                { "saslStart", 1 },
+                { "mechanism", _mechanism.Name },
+                { "payload", currentStep.BytesToSendToServer }
+            };
+
+            return startCommand;
+        }
+
         private CommandWireProtocol<BsonDocument> CreateCommandProtocol(BsonDocument command)
         {
             return new CommandWireProtocol<BsonDocument>(
@@ -161,28 +182,16 @@ private MongoAuthenticationException CreateException(IConnection connection, Exc
             return new MongoAuthenticationException(connection.ConnectionId, message, ex);
         }
 
-        private BsonDocument CreateStartCommand(ISaslStep currentStep)
-        {
-            var startCommand = new BsonDocument
-            {
-                { "saslStart", 1 },
-                { "mechanism", _mechanism.Name },
-                { "payload", currentStep.BytesToSendToServer }
-            };
-
-            if (_mechanism.Name.StartsWith("SCRAM", StringComparison.OrdinalIgnoreCase))
-            {
-                startCommand.Add("options", new BsonDocument("skipEmptyExchange", true));
-            }
-
-            return startCommand;
-        }
-
-        private ISaslStep Transition(SaslConversation conversation, ISaslStep currentStep, BsonDocument result)
+        private ISaslStep Transition(
+            SaslConversation conversation,
+            ISaslStep currentStep,
+            BsonDocument result,
+            out BsonDocument command)
         {
             // we might be done here if the client is not expecting a reply from the server
             if (result.GetValue("done", false).ToBoolean() && currentStep.IsComplete)
             {
+                command = null;
                 return null;
             }
 
@@ -191,12 +200,15 @@ private ISaslStep Transition(SaslConversation conversation, ISaslStep currentSte
             // we might be done here if the client had some final verification it needed to do
             if (result.GetValue("done", false).ToBoolean() && currentStep.IsComplete)
             {
+                command = null;
                 return null;
             }
 
+            command = CreateContinueCommand(currentStep, result);
             return currentStep;
         }
 
+
         // nested classes
         /// <summary>
         /// Represents a SASL conversation.

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

@@ -16,6 +16,7 @@
 using System;
 using System.Security.Cryptography;
 using System.Text;
+using MongoDB.Bson;
 using MongoDB.Bson.IO;
 #if NET452
 using MongoDB.Driver.Core.Authentication.Vendored;
@@ -35,7 +36,7 @@ public abstract class ScramShaAuthenticator : SaslAuthenticator
         /// </summary>
         /// <param name="data">The data to hash. Also called "str" in RFC5802.</param>
         protected internal delegate byte[] H(byte[] data);
-        
+
         /// <summary>
         /// A Hi function used to compute the SaltedPassword as defined in RFC5802, except with "str" parameter replaced
         /// with a UsernamePassword credential so that the password can be optionally digested/prepped in a secure fashion
@@ -46,15 +47,15 @@ public abstract class ScramShaAuthenticator : SaslAuthenticator
         /// <param name="salt">The salt.</param>
         /// <param name="iterations">The iteration count.</param>
         protected internal delegate byte[] Hi(UsernamePasswordCredential credentials, byte[] salt, int iterations);
-        
+
         /// <summary>
         /// An HMAC function as defined in RFC5802, plus the encoding of the data.
         /// </summary>
         /// <param name="encoding">The encoding of the data.</param>
         /// <param name="data">The data. Also called "str" in RFC5802.</param>
         /// <param name="key">The key.</param>
         protected internal delegate byte[] Hmac(UTF8Encoding encoding, byte[] data, string key);
-        
+
         // fields
         private readonly string _databaseName;
 
@@ -67,7 +68,7 @@ public abstract class ScramShaAuthenticator : SaslAuthenticator
         /// <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>
-        protected ScramShaAuthenticator(UsernamePasswordCredential credential, 
+        protected ScramShaAuthenticator(UsernamePasswordCredential credential,
             HashAlgorithmName hashAlgorithmName,
             H h,
             Hi hi,
@@ -85,7 +86,7 @@ protected ScramShaAuthenticator(UsernamePasswordCredential credential,
         /// <param name="hmac">The Hmac function to use.</param>
         /// <param name="cache">The cache to use.</param>
         internal ScramShaAuthenticator(
-            UsernamePasswordCredential credential, 
+            UsernamePasswordCredential credential,
             HashAlgorithmName hashAlgorithName,
             IRandomStringGenerator randomStringGenerator,
             H h,
@@ -101,6 +102,24 @@ internal ScramShaAuthenticator(
         /// <inheritdoc/>
         public override string DatabaseName => _databaseName;
 
+        /// <inheritdoc/>
+        public override BsonDocument CustomizeInitialIsMasterCommand(BsonDocument isMasterCommand)
+        {
+            isMasterCommand = base.CustomizeInitialIsMasterCommand(isMasterCommand);
+            _speculativeFirstStep = _mechanism.Initialize(connection: null, conversation: null, description: null);
+            var firstCommand = CreateStartCommand(_speculativeFirstStep);
+            firstCommand.Add("db", DatabaseName);
+            isMasterCommand.Add("speculativeAuthenticate", firstCommand);
+            return isMasterCommand;
+        }
+
+        private protected override BsonDocument CreateStartCommand(ISaslStep currentStep)
+        {
+            var startCommand = base.CreateStartCommand(currentStep);
+            startCommand.Add("options", new BsonDocument("skipEmptyExchange", true));
+            return startCommand;
+        }
+
         // nested classes
         private class ScramShaMechanism : ISaslMechanism
         {
@@ -113,8 +132,8 @@ private class ScramShaMechanism : ISaslMechanism
             private ScramCache _cache;
 
             public ScramShaMechanism(
-                UsernamePasswordCredential credential, 
-                HashAlgorithmName hashAlgorithmName, 
+                UsernamePasswordCredential credential,
+                HashAlgorithmName hashAlgorithmName,
                 IRandomStringGenerator randomStringGenerator,
                 H h,
                 Hi hi,
@@ -138,9 +157,6 @@ public ScramShaMechanism(
 
             public ISaslStep Initialize(IConnection connection, SaslConversation conversation, ConnectionDescription description)
             {
-                Ensure.IsNotNull(connection, nameof(connection));
-                Ensure.IsNotNull(description, nameof(description));
-
                 const string gs2Header = "n,,";
                 var username = "n=" + PrepUsername(_credential.Username);
                 var r = GenerateRandomString();
@@ -165,13 +181,13 @@ private string PrepUsername(string username)
                 return username.Replace("=", "=3D").Replace(",", "=2C");
             }
         }
-        
+
         private class ClientFirst : ISaslStep
         {
             private readonly byte[] _bytesToSendToServer;
             private readonly string _clientFirstMessageBare;
             private readonly UsernamePasswordCredential _credential;
-           
+
             private readonly string _rPrefix;
             private readonly H _h;
             private readonly Hi _hi;
@@ -180,9 +196,9 @@ private class ClientFirst : ISaslStep
             private ScramCache _cache;
 
             public ClientFirst(
-                byte[] bytesToSendToServer, 
-                string clientFirstMessageBare, 
-                UsernamePasswordCredential credential, 
+                byte[] bytesToSendToServer,
+                string clientFirstMessageBare,
+                UsernamePasswordCredential credential,
                 string rPrefix,
                 H h,
                 Hi hi,
@@ -266,7 +282,7 @@ private byte[] XOR(byte[] a, byte[] b)
         }
 
         private class ClientLast : ISaslStep
-        {            
+        {
             private readonly byte[] _bytesToSendToServer;
             private readonly byte[] _serverSignature64;
 
@@ -293,7 +309,7 @@ public ISaslStep Transition(SaslConversation conversation, byte[] bytesReceivedF
 
                 return new CompletedStep();
             }
-            
+
             private bool ConstantTimeEquals(byte[] a, byte[] b)
             {
                 var diff = a.Length ^ b.Length;
@@ -306,4 +322,4 @@ private bool ConstantTimeEquals(byte[] a, byte[] b)
             }
         }
     }
-}
+}