SERVER-44857 Allow SCRAM conversation to avoid empty exchange

8 files changed, 115 insertions, 10 deletions

diff --git a/src/mongo/client/sasl_client_authenticate_impl.cpp b/src/mongo/client/sasl_client_authenticate_impl.cpp
index 8a5794f95e0..3ff44444287 100644
--- a/src/mongo/client/sasl_client_authenticate_impl.cpp
+++ b/src/mongo/client/sasl_client_authenticate_impl.cpp
@@ -196,6 +196,12 @@ Future<void> asyncSaslConversation(auth::RunCommandHook runCommand,
 
     LOG(saslLogLevel) << "sasl client output: " << base64::encode(responsePayload) << endl;
 
+    // Handle a done from the server which comes before the client is complete.
+    const bool serverDone = inputObj[saslCommandDoneFieldName].trueValue();
+    if (serverDone && responsePayload.empty() && session->isSuccess()) {
+        return Status::OK();
+    }
+
     // Build command using our new payload and conversationId
     BSONObjBuilder commandBuilder;
     commandBuilder.appendElements(saslCommandPrefix);
@@ -268,9 +274,11 @@ Future<void> saslClientAuthenticateImpl(auth::RunCommandHook runCommand,
     if (!status.isOK())
         return status;
 
+    auto mechanismName = session->getParameter(SaslClientSession::parameterMechanism);
     BSONObj saslFirstCommandPrefix =
-        BSON(saslStartCommandName << 1 << saslCommandMechanismFieldName
-                                  << session->getParameter(SaslClientSession::parameterMechanism));
+        BSON(saslStartCommandName << 1 << saslCommandMechanismFieldName << mechanismName
+                                  << "options" << BSON(saslCommandOptionSkipEmptyExchange << true));
+
     BSONObj inputObj = BSON(saslCommandPayloadFieldName << "");
     return asyncSaslConversation(runCommand,
                                  session,
diff --git a/src/mongo/db/auth/authentication_session.h b/src/mongo/db/auth/authentication_session.h
index a7d5e779828..5a26c46efd4 100644
--- a/src/mongo/db/auth/authentication_session.h
+++ b/src/mongo/db/auth/authentication_session.h
@@ -67,6 +67,11 @@ public:
         return *_mech;
     }
 
+    Status setOptions(BSONObj options) {
+        invariant(_mech);
+        return _mech->setOptions(options);
+    }
+
 private:
     std::unique_ptr<ServerMechanismBase> _mech;
 };
diff --git a/src/mongo/db/auth/sasl_authentication_session_test.cpp b/src/mongo/db/auth/sasl_authentication_session_test.cpp
index 39ead833986..bdf8edc175d 100644
--- a/src/mongo/db/auth/sasl_authentication_session_test.cpp
+++ b/src/mongo/db/auth/sasl_authentication_session_test.cpp
@@ -42,6 +42,7 @@
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/auth/authz_manager_external_state_mock.h"
 #include "mongo/db/auth/authz_session_external_state_mock.h"
+#include "mongo/db/auth/sasl_command_constants.h"
 #include "mongo/db/auth/sasl_mechanism_registry.h"
 #include "mongo/db/auth/sasl_options.h"
 #include "mongo/db/auth/sasl_plain_server_conversation.h"
@@ -66,6 +67,7 @@ public:
     void testBadPassword();
     void testWrongClientMechanism();
     void testWrongServerMechanism();
+    void testSCRAMSkipEmptyExchange();
 
     ServiceContext::UniqueOperationContext opCtx;
     AuthzManagerExternalStateMock* authManagerExternalState;
@@ -230,6 +232,54 @@ void SaslConversation::testWrongServerMechanism() {
     assertConversationFailure();
 }
 
+void SaslConversation::testSCRAMSkipEmptyExchange() {
+    if ((mechanism != "SCRAM-SHA-1") && (mechanism != "SCRAM-SHA-256")) {
+        return;
+    }
+
+    for (bool enabled : {true, false}) {
+        client.reset(SaslClientSession::create(mechanism));
+        client->setParameter(SaslClientSession::parameterServiceName, mockServiceName);
+        client->setParameter(SaslClientSession::parameterServiceHostname, mockHostName);
+        client->setParameter(SaslClientSession::parameterMechanism, mechanism);
+        client->setParameter(SaslClientSession::parameterUser, "andy");
+        client->setParameter(SaslClientSession::parameterPassword, "frim");
+        ASSERT_OK(client->initialize());
+
+        auto swServer = registry.getServerMechanism(mechanism, "test");
+        ASSERT_OK(swServer.getStatus());
+        server = std::move(swServer.getValue());
+        ASSERT_OK(server->setOptions(BSON(saslCommandOptionSkipEmptyExchange << enabled)));
+
+        const std::size_t expected = enabled ? 2 : 3;
+        std::size_t step = 0;
+
+        std::string clientMsg = "";
+        StatusWith<std::string> serverMsg = "";
+        for (;;) {
+            ASSERT_OK(client->step(serverMsg.getValue(), &clientMsg));
+            if (client->isSuccess() && server->isSuccess()) {
+                break;
+            }
+
+            if (step > expected) {
+                break;
+            }
+            ++step;
+
+            serverMsg = server->step(opCtx.get(), clientMsg);
+            ASSERT_OK(serverMsg.getStatus());
+            if (client->isSuccess() && server->isSuccess()) {
+                break;
+            }
+        }
+
+        ASSERT_TRUE(client->isSuccess());
+        ASSERT_TRUE(server->isSuccess());
+        ASSERT_EQ(step, expected);
+    }
+}
+
 #define DEFINE_MECHANISM_FIXTURE(CLASS_SUFFIX, MECH_NAME)                   \
     class SaslConversation##CLASS_SUFFIX : public SaslConversation {        \
     public:                                                                 \
@@ -250,7 +300,8 @@ void SaslConversation::testWrongServerMechanism() {
     DEFINE_MECHANISM_TEST(FIXTURE_NAME, NoSuchUser)               \
     DEFINE_MECHANISM_TEST(FIXTURE_NAME, BadPassword)              \
     DEFINE_MECHANISM_TEST(FIXTURE_NAME, WrongClientMechanism)     \
-    DEFINE_MECHANISM_TEST(FIXTURE_NAME, WrongServerMechanism)
+    DEFINE_MECHANISM_TEST(FIXTURE_NAME, WrongServerMechanism)     \
+    DEFINE_MECHANISM_TEST(FIXTURE_NAME, SCRAMSkipEmptyExchange)
 
 #define TEST_MECHANISM(CLASS_SUFFIX, MECH_NAME)        \
     DEFINE_MECHANISM_FIXTURE(CLASS_SUFFIX, MECH_NAME); \
diff --git a/src/mongo/db/auth/sasl_command_constants.h b/src/mongo/db/auth/sasl_command_constants.h
index 03eb2e21d2d..182c55d9738 100644
--- a/src/mongo/db/auth/sasl_command_constants.h
+++ b/src/mongo/db/auth/sasl_command_constants.h
@@ -89,7 +89,10 @@ constexpr auto saslDefaultServiceName = "mongodb"_sd;
 // be digested.
 constexpr auto saslCommandDigestPasswordFieldName = "digestPassword"_sd;
 
-/// Field containing optional session token information for MONGODB-IAM sasl mechanism.
+// Field containing optional session token information for MONGODB-IAM sasl mechanism.
 constexpr auto saslCommandIamSessionToken = "awsIamSessionToken"_sd;
 
+// Field in saslStart.options for mechanisms which omit empty "OK" exchange.
+constexpr auto saslCommandOptionSkipEmptyExchange = "skipEmptyExchange"_sd;
+
 }  // namespace mongo
diff --git a/src/mongo/db/auth/sasl_commands.cpp b/src/mongo/db/auth/sasl_commands.cpp
index 8106f5930f3..1932b7e9a4e 100644
--- a/src/mongo/db/auth/sasl_commands.cpp
+++ b/src/mongo/db/auth/sasl_commands.cpp
@@ -241,6 +241,17 @@ StatusWith<std::unique_ptr<AuthenticationSession>> doSaslStart(OperationContext*
     }
 
     auto session = std::make_unique<AuthenticationSession>(std::move(swMech.getValue()));
+    auto options = cmdObj["options"];
+    if (!options.eoo()) {
+        if (options.type() != Object) {
+            return {ErrorCodes::BadValue, "saslStart.options must be an object"};
+        }
+        status = session->setOptions(options.Obj());
+        if (!status.isOK()) {
+            return status;
+        }
+    }
+
     Status statusStep = doSaslStep(opCtx, session.get(), cmdObj, result);
 
     if (!statusStep.isOK() || session->getMechanism().isSuccess()) {
diff --git a/src/mongo/db/auth/sasl_mechanism_registry.h b/src/mongo/db/auth/sasl_mechanism_registry.h
index 960a93b2db5..50b47c388a7 100644
--- a/src/mongo/db/auth/sasl_mechanism_registry.h
+++ b/src/mongo/db/auth/sasl_mechanism_registry.h
@@ -191,6 +191,14 @@ public:
         }
     }
 
+    /**
+     * Flexible bag of options for a saslStart command.
+     */
+    virtual Status setOptions(BSONObj options) {
+        // Be default, ignore any options provided.
+        return Status::OK();
+    }
+
 protected:
     /**
      * Mechanism provided step implementation.
diff --git a/src/mongo/db/auth/sasl_scram_server_conversation.cpp b/src/mongo/db/auth/sasl_scram_server_conversation.cpp
index 82e6df36a36..cb401e0b45d 100644
--- a/src/mongo/db/auth/sasl_scram_server_conversation.cpp
+++ b/src/mongo/db/auth/sasl_scram_server_conversation.cpp
@@ -41,6 +41,7 @@
 #include "mongo/base/string_data.h"
 #include "mongo/crypto/mechanism_scram.h"
 #include "mongo/crypto/sha1_block.h"
+#include "mongo/db/auth/sasl_command_constants.h"
 #include "mongo/db/auth/sasl_mechanism_policies.h"
 #include "mongo/db/auth/sasl_mechanism_registry.h"
 #include "mongo/db/auth/sasl_options.h"
@@ -58,7 +59,8 @@ StatusWith<std::tuple<bool, std::string>> SaslSCRAMServerMechanism<Policy>::step
     OperationContext* opCtx, StringData inputData) {
     _step++;
 
-    if (_step > 3 || _step <= 0) {
+    const int numSteps = (_skipEmptyExchange ? 2 : 3);
+    if (_step > numSteps || _step <= 0) {
         return Status(ErrorCodes::AuthenticationFailed,
                       str::stream() << "Invalid SCRAM authentication step: " << _step);
     }
@@ -249,7 +251,7 @@ StatusWith<std::tuple<bool, std::string>> SaslSCRAMServerMechanism<Policy>::_fir
     std::string outputData = sb.str();
 
     // add client-first-message-bare and server-first-message to _authMessage
-    _authMessage = client_first_message_bare.toString() + "," + outputData;
+    _authMessage = str::stream() << client_first_message_bare.toString() << "," << outputData;
 
     return std::make_tuple(false, std::move(outputData));
 }
@@ -346,12 +348,24 @@ StatusWith<std::tuple<bool, std::string>> SaslSCRAMServerMechanism<Policy>::_sec
                       "SCRAM authentication failed, storedKey mismatch");
     }
 
-    invariant(!serverSignature.empty());
-    StringBuilder sb;
     // ServerSignature := HMAC(ServerKey, AuthMessage)
-    sb << "v=" << serverSignature;
+    invariant(!serverSignature.empty());
+    std::string reply("v=" + serverSignature);
+    return std::make_tuple(_skipEmptyExchange, reply);
+}
+
+template <typename Policy>
+Status SaslSCRAMServerMechanism<Policy>::setOptions(BSONObj options) {
+    auto see = options[saslCommandOptionSkipEmptyExchange];
+    if (!see.eoo()) {
+        if (!see.isBoolean()) {
+            return {ErrorCodes::BadValue,
+                    str::stream() << saslCommandOptionSkipEmptyExchange << " must be boolean"};
+        }
+        _skipEmptyExchange = see.boolean();
+    }
 
-    return std::make_tuple(false, sb.str());
+    return Status::OK();
 }
 
 template class SaslSCRAMServerMechanism<SCRAMSHA1Policy>;
diff --git a/src/mongo/db/auth/sasl_scram_server_conversation.h b/src/mongo/db/auth/sasl_scram_server_conversation.h
index 68b21da2d3d..4e596c0c8a4 100644
--- a/src/mongo/db/auth/sasl_scram_server_conversation.h
+++ b/src/mongo/db/auth/sasl_scram_server_conversation.h
@@ -67,6 +67,8 @@ public:
         }
     }
 
+    Status setOptions(BSONObj options) final;
+
 private:
     /**
      * Parse client-first-message and generate server-first-message
@@ -88,6 +90,9 @@ private:
 
     // client and server nonce concatenated
     std::string _nonce;
+
+    // Do not send empty 3rd reply in scram conversation.
+    bool _skipEmptyExchange{false};
 };
 
 extern template class SaslSCRAMServerMechanism<SCRAMSHA1Policy>;