SERVER-52863 Instantiate AuthenticationSession during SASL mechanism negotiation

13 files changed, 424 insertions, 110 deletions

diff --git a/src/mongo/base/error_codes.yml b/src/mongo/base/error_codes.yml
index 50936d9cb53..1b9cc12f920 100644
--- a/src/mongo/base/error_codes.yml
+++ b/src/mongo/base/error_codes.yml
@@ -416,6 +416,8 @@ error_codes:
 
     - {code: 336, name: TimeseriesBucketCleared, categories: [InternalOnly]}
 
+    - {code: 337, name: AuthenticationAbandoned, categories: [InternalOnly]}
+
     # Error codes 4000-8999 are reserved.
 
     # Non-sequential error codes for compatibility only)
diff --git a/src/mongo/db/auth/README.md b/src/mongo/db/auth/README.md
index 4487fcbec2c..fce577a7b78 100644
--- a/src/mongo/db/auth/README.md
+++ b/src/mongo/db/auth/README.md
@@ -74,10 +74,12 @@ at runtime. `SASL` mechanisms define a method of communication between a client
 does not, however, define where the user credentials can be stored. With some `SASL` mechanisms,
 `PLAIN` for example, the credentials can be stored in the database itself or in `LDAP`.
 
-Before running authentication, the server sets an empty
-[`AuthenticationSession`](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/authentication_session.h)
-on the `Client`. During the first step of authentication, the client invokes `{saslStart: ...}`,
-which reaches
+Before running authentication, the server initializes an
+[`AuthenticationSession`](https://github.com/mongodb/mongo/blob/master/src/mongo/db/auth/authentication_session.h)
+on the `Client`. This session persists information between authentications steps and is released
+when authentication concludes, either successfully or unsuccessfully.
+
+During the first step of authentication, the client invokes `{saslStart: ...}`, which reaches
 [`doSaslStart`](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/sasl_commands.cpp#L237-L242)
 which gets the mechanism used and performs the actual authentication by calling the step function
 (inherited from
@@ -91,12 +93,10 @@ closes the session.
 
 If, after the first SASL step, there is more work to be done, the client sends a
 [`CMDSaslContinue`](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/sasl_commands.cpp#L98)
-to the server with whatever extra information the server requested. The server then retrieves the
-former
-[`AuthenticationSession`](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/authentication_session.h)
-from the current client and performs another SASL step. The server then sends the client a similar
-reply as it did from the `SASLStart` command. The `SASLContinue` phase repeats until the client is
-either authenticated or an error is encountered.
+to the server with whatever extra information the server requested. The server then performs another
+SASL step. The server then sends the client a similar reply as it did from the `SASLStart` command.
+The `SASLContinue` phase repeats until the client is either authenticated or an error is
+encountered.
 
 #### Speculative Authentication
 
@@ -639,6 +639,7 @@ Refer to the following links for definitions of the Classes referenced in this d
 | Class | File | Description |
 | --- | --- | --- |
 | `ActionType` | [mongo/db/auth/action\_type.h](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/action_type.h) | High level categories of actions which may be performed against a given resource (e.g. `find`, `insert`, `update`, etc...) |
+| `AuthenticationSession` | [mongo/db/auth/authentication\_session.h](https://github.com/mongodb/mongo/blob/master/src/mongo/db/auth/authentication_session.h) | Session object to persist Authentication state |
 | `AuthorizationManager` | [mongo/db/auth/authorization\_manager.h](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/authorization_manager.h) | Interface to external state providers |
 | `AuthorizationSession` | [mongo/db/auth/authorization\_session.h](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/authorization_session.h) | Representation of currently authenticated and authorized users on the `Client` connection |
 | `AuthzManagerExternalStateLocal` | [.../authz\_manager\_external\_state\_local.h](https://github.com/mongodb/mongo/blob/r4.4.0/src/mongo/db/auth/authz_manager_external_state_local.h) | `Local` implementation of user/role provider |
diff --git a/src/mongo/db/auth/SConscript b/src/mongo/db/auth/SConscript
index 5d1890e92d9..e4efc6d1935 100644
--- a/src/mongo/db/auth/SConscript
+++ b/src/mongo/db/auth/SConscript
@@ -38,6 +38,20 @@ env.Library(
 )
 
 env.Library(
+    target='authentication_session',
+    source=[
+        'authentication_session.cpp',
+    ],
+    LIBDEPS=[
+        'auth',
+        'saslauth',
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/stats/counters',
+    ],
+)
+
+env.Library(
     target='auth_op_observer',
     source=[
         "auth_op_observer.cpp",
@@ -258,8 +272,10 @@ env.Library(
         '$BUILD_DIR/mongo/db/audit',
         '$BUILD_DIR/mongo/db/commands',
         '$BUILD_DIR/mongo/db/commands/test_commands_enabled',
+        '$BUILD_DIR/mongo/db/stats/counters',
         'auth',
         'auth_impl_internal',
+        'authentication_session',
         'authorization_manager_global',
         'saslauth',
     ],
diff --git a/src/mongo/db/auth/auth_decorations.cpp b/src/mongo/db/auth/auth_decorations.cpp
index 233692329a6..c8cf55fd765 100644
--- a/src/mongo/db/auth/auth_decorations.cpp
+++ b/src/mongo/db/auth/auth_decorations.cpp
@@ -46,7 +46,7 @@ namespace mongo {
 namespace {
 
 const auto getAuthenticationSession =
-    Client::declareDecoration<std::unique_ptr<AuthenticationSession>>();
+    Client::declareDecoration<boost::optional<AuthenticationSession>>();
 
 const auto getAuthorizationManager =
     ServiceContext::declareDecoration<std::unique_ptr<AuthorizationManager>>();
@@ -96,13 +96,8 @@ ServiceContext::ConstructorActionRegisterer destroyAuthorizationManagerRegistere
 
 }  // namespace
 
-void AuthenticationSession::set(Client* client, std::unique_ptr<AuthenticationSession> newSession) {
-    getAuthenticationSession(client) = std::move(newSession);
-}
-
-void AuthenticationSession::swap(Client* client, std::unique_ptr<AuthenticationSession>& other) {
-    using std::swap;
-    swap(getAuthenticationSession(client), other);
+boost::optional<AuthenticationSession>& AuthenticationSession::_get(Client* client) {
+    return getAuthenticationSession(client);
 }
 
 AuthorizationManager* AuthorizationManager::get(ServiceContext* service) {
diff --git a/src/mongo/db/auth/authentication_session.cpp b/src/mongo/db/auth/authentication_session.cpp
new file mode 100644
index 00000000000..227709de932
--- /dev/null
+++ b/src/mongo/db/auth/authentication_session.cpp
@@ -0,0 +1,149 @@
+/**
+ *    Copyright (C) 2018-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kAccessControl
+
+#include "mongo/db/auth/authentication_session.h"
+
+#include "mongo/db/audit.h"
+#include "mongo/logv2/log.h"
+
+namespace mongo {
+namespace {
+constexpr auto kDiagnosticLogLevel = 0;
+}
+
+AuthenticationSession::StepGuard::StepGuard(OperationContext* opCtx, StepType currentStep)
+    : _opCtx(opCtx), _currentStep(currentStep) {
+    auto client = _opCtx->getClient();
+
+    LOGV2_DEBUG(
+        5286300, kDiagnosticLogLevel, "Starting authentication step", "step"_attr = _currentStep);
+
+    auto& maybeSession = _get(client);
+    ON_BLOCK_EXIT([&] {
+        if (maybeSession) {
+            // If we successfully made/kept a session, update it and track it.
+            auto& session = *maybeSession;
+            _session = &session;
+        }
+    });
+
+    auto createSession = [&] {
+        if (maybeSession) {
+            maybeSession->markFailed(
+                {ErrorCodes::AuthenticationAbandoned, "Overridden by new authentication session"});
+        }
+        maybeSession.emplace(client);
+    };
+
+    switch (_currentStep) {
+        case StepType::kSaslSupportedMechanisms: {
+            createSession();
+            authCounter.incSaslSupportedMechanismsReceived();
+        } break;
+        case StepType::kSpeculativeAuthenticate:
+        case StepType::kSpeculativeSaslStart: {
+            createSession();
+            maybeSession->_isSpeculative = true;
+        } break;
+        case StepType::kAuthenticate:
+        case StepType::kSaslStart: {
+            createSession();
+        } break;
+        case StepType::kSaslContinue: {
+            uassert(ErrorCodes::ProtocolError, "No SASL session state found", maybeSession);
+        } break;
+    }
+}
+
+AuthenticationSession::StepGuard::~StepGuard() {
+    auto& maybeSession = _get(_opCtx->getClient());
+    if (maybeSession) {
+        LOGV2_DEBUG(5286301,
+                    kDiagnosticLogLevel,
+                    "Finished authentication step",
+                    "step"_attr = _currentStep);
+        if (maybeSession->isFinished()) {
+            // We're done with this session, reset it.
+            maybeSession.reset();
+        }
+    }
+}
+
+AuthenticationSession* AuthenticationSession::get(Client* client) {
+    auto& maybeSession = _get(client);
+    tassert(5286302, "Unable to retrieve authentication session", static_cast<bool>(maybeSession));
+    return &(*maybeSession);
+}
+
+void AuthenticationSession::setMechanism(std::unique_ptr<ServerMechanismBase> mech,
+                                         boost::optional<BSONObj> options) {
+    tassert(5286303, "Attempted to override previous authentication mechanism", !_mech);
+
+    _mech = std::move(mech);
+    if (options) {
+        uassertStatusOK(_mech->setOptions(options->getOwned()));
+    }
+
+    if (_mech && _mech->isClusterMember()) {
+        setAsClusterMember();
+    }
+
+    LOGV2_DEBUG(5286304, kDiagnosticLogLevel, "Determined mechanism for authentication");
+}
+
+void AuthenticationSession::setAsClusterMember() {
+    _isClusterMember = true;
+
+    LOGV2_DEBUG(5286305, kDiagnosticLogLevel, "Marking as cluster member");
+}
+
+void AuthenticationSession::markSuccessful() {
+    // Log success.
+    _isFinished = true;
+    LOGV2_DEBUG(5286306,
+                kDiagnosticLogLevel,
+                "Successfully authenticated",
+                "isSpeculative"_attr = isSpeculative(),
+                "isClusterMember"_attr = isClusterMember());
+}
+
+void AuthenticationSession::markFailed(const Status& status) {
+    // Log the error.
+    _isFinished = true;
+    LOGV2_DEBUG(5286307,
+                kDiagnosticLogLevel,
+                "Failed to authenticate",
+                "isSpeculative"_attr = isSpeculative(),
+                "isClusterMember"_attr = isClusterMember(),
+                "error"_attr = status);
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/auth/authentication_session.h b/src/mongo/db/auth/authentication_session.h
index ff86ea7ca08..38807455086 100644
--- a/src/mongo/db/auth/authentication_session.h
+++ b/src/mongo/db/auth/authentication_session.h
@@ -32,6 +32,7 @@
 #include <memory>
 
 #include "mongo/db/auth/sasl_mechanism_registry.h"
+#include "mongo/db/stats/counters.h"
 
 namespace mongo {
 
@@ -45,40 +46,169 @@ class AuthenticationSession {
     AuthenticationSession& operator=(const AuthenticationSession&) = delete;
 
 public:
-    explicit AuthenticationSession(std::unique_ptr<ServerMechanismBase> mech, bool speculative)
-        : _mech(std::move(mech)), _speculative(speculative) {}
+    /**
+     * This enum enumerates the various steps that need access to the AuthenticationSession.
+     */
+    enum class StepType {
+        kSaslSupportedMechanisms,
+        kSaslStart,
+        kSaslContinue,
+        kAuthenticate,
+        kSpeculativeSaslStart,
+        kSpeculativeAuthenticate,
+    };
+
+    AuthenticationSession(Client* client) : _client(client) {}
 
     /**
-     * Sets the authentication session for the given "client" to "newSession".
+     * This guard creates and destroys the session as appropriate for the currentStep.
      */
-    static void set(Client* client, std::unique_ptr<AuthenticationSession> newSession);
+    class StepGuard {
+    public:
+        StepGuard(OperationContext* opCtx, StepType currentStep);
+        ~StepGuard();
+
+        AuthenticationSession* getSession() {
+            return _session;
+        }
+
+    private:
+        OperationContext* const _opCtx;
+        const StepType _currentStep;
+
+        AuthenticationSession* _session;
+    };
 
     /**
-     * Swaps "client"'s current authentication session with "other".
+     * Gets the authentication session for the given "client".
+     *
+     * This function always returns a valid pointer.
      */
-    static void swap(Client* client, std::unique_ptr<AuthenticationSession>& other);
+    static AuthenticationSession* get(Client* client);
+    static AuthenticationSession* get(OperationContext* opCtx) {
+        return get(opCtx->getClient());
+    }
 
     /**
      * Return an identifer of the type of session, so that a caller can safely cast it and
      * extract the type-specific data stored within.
+     *
+     * If a mechanism has not already been set, this may return nullptr.
      */
-    ServerMechanismBase& getMechanism() const {
-        invariant(_mech);
-        return *_mech;
+    ServerMechanismBase* getMechanism() const {
+        return _mech.get();
     }
 
-    Status setOptions(BSONObj options) {
-        invariant(_mech);
-        return _mech->setOptions(options);
+    /**
+     * This returns true if the session started with StepType::kSpeculativeSaslStart or
+     * StepType::kSpeculativeAuthenticate.
+     */
+    bool isSpeculative() const {
+        return _isSpeculative;
     }
 
-    bool isSpeculative() const {
-        return _speculative;
+    /**
+     * This returns true if the session currently believes itself to be a cluster member.
+     */
+    bool isClusterMember() const {
+        if (_mech && _mech->isClusterMember()) {
+            // If we're doing sasl and we have a mechanism, then we know.
+            return true;
+        }
+
+        // Otherwise, rely on what the implementation has told us directly.
+        return _isClusterMember;
+    }
+
+    /**
+     * This returns true once either markFailed or markSuccessful is invoked.
+     */
+    bool isFinished() const {
+        return _isFinished;
+    }
+
+    /**
+     * Mark the session as a cluster member.
+     *
+     * This is used for x509 authentication since it lacks a mechanism in the traditional sense.
+     */
+    void setAsClusterMember();
+
+    /**
+     * Set the mechanism for the session.
+     *
+     * This function is only valid to invoke when there is no current mechanism.
+     */
+    void setMechanism(std::unique_ptr<ServerMechanismBase> mech, boost::optional<BSONObj> options);
+
+    /**
+     * Mark the session as succssfully authenticated.
+     *
+     * TODO(SERVER-52862) This should increment counters and log.
+     */
+    void markSuccessful();
+
+    /**
+     * Mark the session as unable to authenticate.
+     *
+     * TODO(SERVER-52862) This should increment counters and log.
+     */
+    void markFailed(const Status& status);
+
+    /**
+     * This function invokes a functor with a StepGuard on the stack and observes any exceptions
+     * emitted.
+     */
+    template <typename F>
+    static auto doStep(OperationContext* opCtx, StepType state, F&& f) {
+        auto guard = StepGuard(opCtx, state);
+        auto session = guard.getSession();
+
+        try {
+            return std::forward<F>(f)(session);
+        } catch (const DBException& ex) {
+            session->markFailed(ex.toStatus());
+            throw;
+        } catch (...) {
+            // Swallow other errors.
+            session->markFailed(
+                Status(ErrorCodes::InternalError, "Encountered an unhandleable error"));
+            throw;
+        }
+    }
+
+    /**
+     * Convert a StepType to a constant string.
+     */
+    friend constexpr StringData toString(StepType step) {
+        switch (step) {
+            case StepType::kSaslSupportedMechanisms:
+                return "SaslSupportedMechanisms"_sd;
+            case StepType::kSaslStart:
+                return "SaslStart"_sd;
+            case StepType::kSaslContinue:
+                return "SaslContinue"_sd;
+            case StepType::kAuthenticate:
+                return "Authenticate"_sd;
+            case StepType::kSpeculativeSaslStart:
+                return "SpeculativeSaslStart"_sd;
+            case StepType::kSpeculativeAuthenticate:
+                return "SpeculativeAuthenticate"_sd;
+        }
+
+        return "Unknown"_sd;
     }
 
 private:
+    static boost::optional<AuthenticationSession>& _get(Client* client);
+
+    Client* const _client;
+
+    bool _isSpeculative = false;
+    bool _isClusterMember = false;
+    bool _isFinished = false;
+
     std::unique_ptr<ServerMechanismBase> _mech;
-    bool _speculative{false};
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/auth/sasl_commands.cpp b/src/mongo/db/auth/sasl_commands.cpp
index 25f09916a51..0a14aeba223 100644
--- a/src/mongo/db/auth/sasl_commands.cpp
+++ b/src/mongo/db/auth/sasl_commands.cpp
@@ -138,7 +138,9 @@ public:
 StatusWith<SaslReply> doSaslStep(OperationContext* opCtx,
                                  const SaslPayload& payload,
                                  AuthenticationSession* session) try {
-    auto& mechanism = session->getMechanism();
+    auto mechanismPtr = session->getMechanism();
+    invariant(mechanismPtr);
+    auto& mechanism = *mechanismPtr;
 
     // Passing in a payload and extracting a responsePayload
     StatusWith<std::string> swResponse = mechanism.step(opCtx, payload.get());
@@ -182,6 +184,8 @@ StatusWith<SaslReply> doSaslStep(OperationContext* opCtx,
                   "authenticationDatabase"_attr = mechanism.getAuthenticationDatabase(),
                   "remote"_attr = opCtx->getClient()->session()->remote());
         }
+
+        session->markSuccessful();
     }
 
     SaslReply reply;
@@ -198,41 +202,36 @@ StatusWith<SaslReply> doSaslStep(OperationContext* opCtx,
 }
 
 SaslReply doSaslStart(OperationContext* opCtx,
+                      AuthenticationSession* session,
                       const SaslStartCommand& request,
-                      bool speculative,
-                      std::string* principalName,
-                      std::unique_ptr<AuthenticationSession>* session) {
+                      std::string* principalName) {
     auto mechanism = uassertStatusOK(
         SASLServerMechanismRegistry::get(opCtx->getServiceContext())
             .getServerMechanism(request.getMechanism(), request.getDbName().toString()));
 
     uassert(ErrorCodes::BadValue,
             "Plaintext mechanisms may not be used with speculativeSaslStart",
-            !speculative ||
+            !session->isSpeculative() ||
                 mechanism->properties().hasAllProperties(
                     SecurityPropertySet({SecurityProperty::kNoPlainText})));
 
-    auto newSession = std::make_unique<AuthenticationSession>(std::move(mechanism), speculative);
-
-    if (auto options = request.getOptions()) {
-        uassertStatusOK(newSession->setOptions(options->getOwned()));
-    }
+    session->setMechanism(std::move(mechanism), request.getOptions());
 
-    auto swReply = doSaslStep(opCtx, request.getPayload(), newSession.get());
-    if (!swReply.isOK() || newSession->getMechanism().isSuccess()) {
+    auto swReply = doSaslStep(opCtx, request.getPayload(), session);
+    if (!swReply.isOK() || session->getMechanism()->isSuccess()) {
         // Only attempt to populate principal name if we're done (successfully or not).
-        *principalName = newSession->getMechanism().getPrincipalName().toString();
+        *principalName = session->getMechanism()->getPrincipalName().toString();
     }
 
     auto reply = uassertStatusOK(swReply);
-    session->reset(newSession.release());
     return reply;
 }
 
-SaslReply runSaslStart(OperationContext* opCtx, const SaslStartCommand& request, bool speculative) {
+SaslReply runSaslStart(OperationContext* opCtx,
+                       AuthenticationSession* session,
+                       const SaslStartCommand& request) {
     opCtx->markKillOnClientDisconnect();
     auto client = opCtx->getClient();
-    AuthenticationSession::set(client, std::unique_ptr<AuthenticationSession>());
 
     auto db = request.getDbName();
     auto mechanismName = request.getMechanism().toString();
@@ -240,32 +239,27 @@ SaslReply runSaslStart(OperationContext* opCtx, const SaslStartCommand& request,
     SaslReply reply;
     std::string principalName;
     try {
-        std::unique_ptr<AuthenticationSession> session;
-
         auto mechCounter = authCounter.getMechanismCounter(mechanismName);
         mechCounter.incAuthenticateReceived();
-        if (speculative) {
+        if (session->isSpeculative()) {
             mechCounter.incSpeculativeAuthenticateReceived();
         }
 
-        reply = doSaslStart(opCtx, request, speculative, &principalName, &session);
+        reply = doSaslStart(opCtx, session, request, &principalName);
 
-        const bool isClusterMember = session->getMechanism().isClusterMember();
-        if (isClusterMember) {
+        if (session->isClusterMember()) {
             mechCounter.incClusterAuthenticateReceived();
         }
-        if (session->getMechanism().isSuccess()) {
+        if (session->getMechanism()->isSuccess()) {
             mechCounter.incAuthenticateSuccessful();
-            if (isClusterMember) {
+            if (session->isClusterMember()) {
                 mechCounter.incClusterAuthenticateSuccessful();
             }
-            if (speculative) {
+            if (session->isSpeculative()) {
                 mechCounter.incSpeculativeAuthenticateSuccessful();
             }
             audit::logAuthentication(
                 client, mechanismName, UserName(principalName, db), ErrorCodes::OK);
-        } else {
-            AuthenticationSession::swap(client, session);
         }
     } catch (const AssertionException& ex) {
         audit::logAuthentication(client, mechanismName, UserName(principalName, db), ex.code());
@@ -275,25 +269,34 @@ SaslReply runSaslStart(OperationContext* opCtx, const SaslStartCommand& request,
     return reply;
 }
 
+SaslReply runSaslContinue(OperationContext* opCtx,
+                          AuthenticationSession* session,
+                          const SaslContinueCommand& request);
+
 SaslReply CmdSaslStart::Invocation::typedRun(OperationContext* opCtx) {
-    return runSaslStart(opCtx, request(), false);
+    return AuthenticationSession::doStep(
+        opCtx, AuthenticationSession::StepType::kSaslStart, [&](auto session) {
+            return runSaslStart(opCtx, session, request());
+        });
 }
 
 SaslReply CmdSaslContinue::Invocation::typedRun(OperationContext* opCtx) {
-    auto cmd = request();
+    return AuthenticationSession::doStep(
+        opCtx, AuthenticationSession::StepType::kSaslContinue, [&](auto session) {
+            return runSaslContinue(opCtx, session, request());
+        });
+}
 
+SaslReply runSaslContinue(OperationContext* opCtx,
+                          AuthenticationSession* session,
+                          const SaslContinueCommand& cmd) {
     opCtx->markKillOnClientDisconnect();
     auto* client = Client::getCurrent();
-    std::unique_ptr<AuthenticationSession> sessionGuard;
-    AuthenticationSession::swap(client, sessionGuard);
-
-    if (!sessionGuard) {
-        uasserted(ErrorCodes::ProtocolError, "No SASL session state found");
-    }
 
-    auto* session = static_cast<AuthenticationSession*>(sessionGuard.get());
+    auto mechanismPtr = session->getMechanism();
+    invariant(mechanismPtr);
+    auto& mechanism = *mechanismPtr;
 
-    auto& mechanism = session->getMechanism();
     // Authenticating the __system@local user to the admin database on mongos is required
     // by the auth passthrough test suite.
     if (mechanism.getAuthenticationDatabase() != cmd.getDbName() && !getTestCommandsEnabled()) {
@@ -324,8 +327,6 @@ SaslReply CmdSaslContinue::Invocation::typedRun(OperationContext* opCtx) {
                 mechCounter.incSpeculativeAuthenticateSuccessful();
             }
         }
-    } else {
-        AuthenticationSession::swap(client, sessionGuard);
     }
 
     return uassertStatusOK(swReply);
@@ -353,11 +354,13 @@ void doSpeculativeSaslStart(OperationContext* opCtx, BSONObj cmdObj, BSONObjBuil
         return;
     }
 
-    auto reply = auth::runSaslStart(
-        opCtx,
-        auth::SaslStartCommand::parse(IDLParserErrorContext("speculative saslStart"), cmd.obj()),
-        true);
-    result->append(auth::kSpeculativeAuthenticate, reply.toBSON());
+    AuthenticationSession::doStep(
+        opCtx, AuthenticationSession::StepType::kSpeculativeSaslStart, [&](auto session) {
+            auto request = auth::SaslStartCommand::parse(
+                IDLParserErrorContext("speculative saslStart"), cmd.obj());
+            auto reply = auth::runSaslStart(opCtx, session, request);
+            result->append(auth::kSpeculativeAuthenticate, reply.toBSON());
+        });
 } catch (...) {
     // Treat failure like we never even got a speculative start.
 }
diff --git a/src/mongo/db/auth/sasl_mechanism_registry.cpp b/src/mongo/db/auth/sasl_mechanism_registry.cpp
index 09ef7502056..d4ca26cf809 100644
--- a/src/mongo/db/auth/sasl_mechanism_registry.cpp
+++ b/src/mongo/db/auth/sasl_mechanism_registry.cpp
@@ -162,6 +162,18 @@ std::vector<std::string> SASLServerMechanismRegistry::getMechanismNames() const
     return names;
 }
 
+StringData ServerMechanismBase::getAuthenticationDatabase() const {
+    if (getTestCommandsEnabled() && _authenticationDatabase == "admin" &&
+        getPrincipalName() == internalSecurity.user->getName().getUser()) {
+        // Allows authenticating as the internal user against the admin database.  This is to
+        // support the auth passthrough test framework on mongos (since you can't use the local
+        // database on a mongos, so you can't auth as the internal user without this).
+        return internalSecurity.user->getName().getDB();
+    } else {
+        return _authenticationDatabase;
+    }
+}
+
 namespace {
 ServiceContext::ConstructorActionRegisterer SASLServerMechanismRegistryInitializer{
     "CreateSASLServerMechanismRegistry", {"EndStartupOptionStorage"}, [](ServiceContext* service) {
diff --git a/src/mongo/db/auth/sasl_mechanism_registry.h b/src/mongo/db/auth/sasl_mechanism_registry.h
index 3ac3a92dd4f..f23b0c5ac83 100644
--- a/src/mongo/db/auth/sasl_mechanism_registry.h
+++ b/src/mongo/db/auth/sasl_mechanism_registry.h
@@ -188,17 +188,7 @@ public:
     }
 
     /** Returns which database contains the user which authentication is being performed against. */
-    StringData getAuthenticationDatabase() const {
-        if (getTestCommandsEnabled() && _authenticationDatabase == "admin" &&
-            getPrincipalName() == internalSecurity.user->getName().getUser()) {
-            // Allows authenticating as the internal user against the admin database.  This is to
-            // support the auth passthrough test framework on mongos (since you can't use the local
-            // database on a mongos, so you can't auth as the internal user without this).
-            return internalSecurity.user->getName().getDB();
-        } else {
-            return _authenticationDatabase;
-        }
-    }
+    StringData getAuthenticationDatabase() const;
 
     /**
      * Flexible bag of options for a saslStart command.
diff --git a/src/mongo/db/commands/SConscript b/src/mongo/db/commands/SConscript
index a1e50e001e8..95c3134673e 100644
--- a/src/mongo/db/commands/SConscript
+++ b/src/mongo/db/commands/SConscript
@@ -203,6 +203,7 @@ env.Library(
         '$BUILD_DIR/mongo/bson/mutable/mutable_bson',
         '$BUILD_DIR/mongo/db/audit',
         '$BUILD_DIR/mongo/db/auth/auth',
+        '$BUILD_DIR/mongo/db/auth/authentication_session',
         '$BUILD_DIR/mongo/db/commands',
         '$BUILD_DIR/mongo/db/stats/counters',
         '$BUILD_DIR/mongo/rpc/client_metadata',
diff --git a/src/mongo/db/commands/authentication_commands.cpp b/src/mongo/db/commands/authentication_commands.cpp
index 485b292c7a2..68f27357867 100644
--- a/src/mongo/db/commands/authentication_commands.cpp
+++ b/src/mongo/db/commands/authentication_commands.cpp
@@ -187,7 +187,10 @@ constexpr auto kX509AuthenticationDisabledMessage = "x.509 authentication is dis
  * mechanism, and ProtocolError, indicating an error in the use of the authentication
  * protocol.
  */
-void _authenticateX509(OperationContext* opCtx, UserName& user, StringData dbname) {
+void _authenticateX509(OperationContext* opCtx,
+                       AuthenticationSession* session,
+                       UserName& user,
+                       StringData dbname) {
     if (user.getUser().empty()) {
         auto& sslPeerInfo = SSLPeerInfo::forSession(opCtx->getClient()->session());
         user = UserName(sslPeerInfo.subjectName.toString(), dbname);
@@ -259,18 +262,21 @@ void _authenticateX509(OperationContext* opCtx, UserName& user, StringData dbnam
 #endif  // MONGO_CONFIG_SSL
 
 void _authenticate(OperationContext* opCtx,
+                   AuthenticationSession* session,
                    StringData mechanism,
                    UserName& user,
                    StringData dbname) {
 #ifdef MONGO_CONFIG_SSL
     if (mechanism == kX509AuthMechanism) {
-        return _authenticateX509(opCtx, user, dbname);
+        return _authenticateX509(opCtx, session, user, dbname);
     }
 #endif
     uasserted(ErrorCodes::BadValue, "Unsupported mechanism: " + mechanism);
 }
 
-AuthenticateReply authCommand(OperationContext* opCtx, const AuthenticateCommand& cmd) {
+AuthenticateReply authCommand(OperationContext* opCtx,
+                              AuthenticationSession* session,
+                              const AuthenticateCommand& cmd) {
     auto dbname = cmd.getDbName();
     UserName user(cmd.getUser().value_or(""), dbname);
     const std::string mechanism(cmd.getMechanism());
@@ -300,7 +306,7 @@ AuthenticateReply authCommand(OperationContext* opCtx, const AuthenticateCommand
         auto mechCounter = authCounter.getMechanismCounter("MONGODB-X509");
         mechCounter.incAuthenticateReceived();
 
-        _authenticate(opCtx, mechanism, user, dbname);
+        _authenticate(opCtx, session, mechanism, user, dbname);
         audit::logAuthentication(opCtx->getClient(), mechanism, user, ErrorCodes::OK);
 
         if (!serverGlobalParams.quiet.load()) {
@@ -312,6 +318,7 @@ AuthenticateReply authCommand(OperationContext* opCtx, const AuthenticateCommand
                   "remote"_attr = opCtx->getClient()->session()->remote());
         }
 
+        session->markSuccessful();
         mechCounter.incAuthenticateSuccessful();
 
         return AuthenticateReply(user.getUser().toString(), user.getDB().toString());
@@ -354,8 +361,11 @@ public:
         void doCheckAuthorization(OperationContext*) const final {}
 
         Reply typedRun(OperationContext* opCtx) final {
-            CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
-            return authCommand(opCtx, request());
+            return AuthenticationSession::doStep(
+                opCtx, AuthenticationSession::StepType::kAuthenticate, [&](auto session) {
+                    CommandHelpers::handleMarkKillOnClientDisconnect(opCtx);
+                    return authCommand(opCtx, session, request());
+                });
         }
     };
 
@@ -386,18 +396,20 @@ void doSpeculativeAuthenticate(OperationContext* opCtx,
     auto authCmdObj = AuthenticateCommand::parse(
         IDLParserErrorContext("speculative X509 Authenticate"), cmd.obj());
 
+    AuthenticationSession::doStep(
+        opCtx, AuthenticationSession::StepType::kSpeculativeAuthenticate, [&](auto session) {
+            const auto mechanism = authCmdObj.getMechanism().toString();
+            try {
+                authCounter.getMechanismCounter(mechanism).incSpeculativeAuthenticateReceived();
+            } catch (...) {
+                // Run will make sure an audit entry happens. Let it reach that point.
+            }
 
-    const auto mechanism = authCmdObj.getMechanism().toString();
-    try {
-        authCounter.getMechanismCounter(mechanism).incSpeculativeAuthenticateReceived();
-    } catch (...) {
-        // Run will make sure an audit entry happens. Let it reach that point.
-    }
-
-    auto authReply = authCommand(opCtx, authCmdObj);
+            auto authReply = authCommand(opCtx, session, authCmdObj);
 
-    authCounter.getMechanismCounter(mechanism).incSpeculativeAuthenticateSuccessful();
-    result->append(auth::kSpeculativeAuthenticate, authReply.toBSON());
+            authCounter.getMechanismCounter(mechanism).incSpeculativeAuthenticateSuccessful();
+            result->append(auth::kSpeculativeAuthenticate, authReply.toBSON());
+        });
 } catch (...) {
     // Treat failure like we never even got a speculative start.
 }
diff --git a/src/mongo/db/repl/SConscript b/src/mongo/db/repl/SConscript
index bae691f94ef..b2747bf9a24 100644
--- a/src/mongo/db/repl/SConscript
+++ b/src/mongo/db/repl/SConscript
@@ -1742,8 +1742,8 @@ env.Library(
     ],
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/auth/authentication_session',
         '$BUILD_DIR/mongo/db/auth/authservercommon',
-        '$BUILD_DIR/mongo/db/stats/counters',
     ],
 )
 
diff --git a/src/mongo/db/repl/hello_auth.cpp b/src/mongo/db/repl/hello_auth.cpp
index 8c0c7ab8787..b88ea5b1325 100644
--- a/src/mongo/db/repl/hello_auth.cpp
+++ b/src/mongo/db/repl/hello_auth.cpp
@@ -30,6 +30,7 @@
 #include "mongo/db/repl/hello_auth.h"
 
 #include "mongo/client/authenticate.h"
+#include "mongo/db/auth/authentication_session.h"
 #include "mongo/db/auth/sasl_command_constants.h"
 #include "mongo/db/auth/sasl_commands.h"
 #include "mongo/db/auth/sasl_mechanism_registry.h"
@@ -41,12 +42,14 @@ namespace mongo {
 void handleHelloAuth(OperationContext* opCtx, BSONObj cmdObj, BSONObjBuilder* result) {
     auto ssme = cmdObj[auth::kSaslSupportedMechanisms];
     if (ssme.type() == BSONType::String) {
-        UserName userName = uassertStatusOK(UserName::parse(ssme.String()));
+        AuthenticationSession::doStep(
+            opCtx, AuthenticationSession::StepType::kSaslSupportedMechanisms, [&](auto session) {
+                UserName userName = uassertStatusOK(UserName::parse(ssme.String()));
 
-        authCounter.incSaslSupportedMechanismsReceived();
-
-        auto& saslMechanismRegistry = SASLServerMechanismRegistry::get(opCtx->getServiceContext());
-        saslMechanismRegistry.advertiseMechanismNamesForUser(opCtx, userName, result);
+                auto& saslMechanismRegistry =
+                    SASLServerMechanismRegistry::get(opCtx->getServiceContext());
+                saslMechanismRegistry.advertiseMechanismNamesForUser(opCtx, userName, result);
+            });
     }
 
     auto sae = cmdObj[auth::kSpeculativeAuthenticate];