delay the creation of AuthenticationContext

emelialei88 · emelialei88 · commit 95f948051d53 · 2025-05-20T17:44:45.000-04:00
Signed-off-by: Emelia Lei &lt;wlei29@bloomberg.net&gt;
diff --git a/src/groups/mqb/mqba/mqba_authenticator.cpp b/src/groups/mqb/mqba/mqba_authenticator.cpp
@@ -28,6 +28,7 @@
 // MQB
 #include <mqbblp_clustercatalog.h>
 #include <mqbnet_authenticationcontext.h>
+#include <mqbnet_initialconnectioncontext.h>
 
 // BMQ
 #include <bmqio_status.h>
@@ -57,26 +58,42 @@ const int k_AUTHENTICATION_READTIMEOUT = 3 * 60;  // 3 minutes
 // -------------------
 
 int Authenticator::onAuthenticationRequest(
-    bsl::ostream&                  errorDescription,
-    const AuthenticationContextSp& context)
+    bsl::ostream&                              errorDescription,
+    const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+    const InitialConnectionContextSp&          context)
 {
     // PRECONDITIONS
-    BSLS_ASSERT_SAFE(
-        context->authenticationMessage().isAuthenticateRequestValue());
-    BSLS_ASSERT_SAFE(context->initialConnectionContext()->isIncoming());
-    BSLS_ASSERT_SAFE(!context->isReversed());  // not supported for now
+    BSLS_ASSERT_SAFE(authenticationMsg.isAuthenticateRequestValue());
+    BSLS_ASSERT_SAFE(context->isIncoming());
 
     const bmqp_ctrlmsg::AuthenticateRequest& authenticateRequest =
-        context->authenticationMessage().authenticateRequest();
+        authenticationMsg.authenticateRequest();
 
     BALL_LOG_DEBUG << "Received authentication message from '"
-                   << context->initialConnectionContext()->channel()->peerUri()
+                   << context->channel()->peerUri()
                    << "': " << authenticateRequest;
 
     bmqp_ctrlmsg::AuthenticationMessage authenticationResponse;
     bmqp_ctrlmsg::AuthenticateResponse& response =
         authenticationResponse.makeAuthenticateResponse();
 
+    if (!context->authenticationContext()) {
+        // Create an AuthenticationContext for that connection
+        bsl::shared_ptr<mqbnet::AuthenticationContext> authenticationContext =
+            bsl::allocate_shared<mqbnet::AuthenticationContext>(
+                d_allocator_p,
+                context.get(),           // initialConnectionContext
+                false,                   // isReversed
+                State::e_AUTHENTICATING  // state
+            );
+        context->setAuthenticationContext(authenticationContext);
+    }
+    else {
+        context->authenticationContext()->testAndSwapState(
+            State::e_AUTHENTICATED,
+            State::e_AUTHENTICATING);
+    }
+
     // Always succeeds for now
     // TODO: For later implementation, plugins will perform authentication,
     // taking the `AuthenticationContext` and updates it with the
@@ -85,20 +102,20 @@ int Authenticator::onAuthenticationRequest(
     response.status().code()     = 0;
     response.lifetimeMs()        = 10 * 60 * 1000;
 
-    context->testAndSwapState(
-        mqbnet::AuthenticationContext::State::e_AUTHENTICATING,
-        mqbnet::AuthenticationContext::State::e_AUTHENTICATED);
+    context->authenticationContext()->testAndSwapState(State::e_AUTHENTICATING,
+                                                       State::e_AUTHENTICATED);
 
     int rc = sendAuthenticationMessage(errorDescription,
                                        authenticationResponse,
-                                       context);
+                                       context->authenticationContext());
 
     return rc;
 }
 
 int Authenticator::onAuthenticationResponse(
-    bsl::ostream&                  errorDescription,
-    const AuthenticationContextSp& context)
+    bsl::ostream&                              errorDescription,
+    const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+    const InitialConnectionContextSp&          context)
 {
     BALL_LOG_ERROR << "Not Implemented";
 
@@ -167,9 +184,11 @@ Authenticator::~Authenticator()
     // NOTHING: (required because of inheritance)
 }
 
-int Authenticator::handleAuthentication(bsl::ostream& errorDescription,
-                                        bool*         isContinueRead,
-                                        const AuthenticationContextSp& context)
+int Authenticator::handleAuthentication(
+    bsl::ostream&                              errorDescription,
+    const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+    bool*                                      isContinueRead,
+    const InitialConnectionContextSp&          context)
 {
     enum RcEnum {
         // Value for the various RC error categories
@@ -180,19 +199,19 @@ int Authenticator::handleAuthentication(bsl::ostream& errorDescription,
     bmqu::MemOutStream errStream;
     int                rc = rc_SUCCESS;
 
-    switch (context->authenticationMessage().selectionId()) {
+    switch (authenticationMsg.selectionId()) {
     case bmqp_ctrlmsg::AuthenticationMessage::
         SELECTION_ID_AUTHENTICATE_REQUEST: {
-        rc = onAuthenticationRequest(errStream, context);
+        rc = onAuthenticationRequest(errStream, authenticationMsg, context);
     } break;  // BREAK
     case bmqp_ctrlmsg::AuthenticationMessage::
         SELECTION_ID_AUTHENTICATE_RESPONSE: {
-        rc = onAuthenticationResponse(errStream, context);
+        rc = onAuthenticationResponse(errStream, authenticationMsg, context);
     } break;  // BREAK
     default: {
         errorDescription
             << "Invalid authentication message received (unknown type): "
-            << context->authenticationMessage();
+            << authenticationMsg;
         return rc_ERROR;  // RETURN
     }
     }
diff --git a/src/groups/mqb/mqba/mqba_authenticator.h b/src/groups/mqb/mqba/mqba_authenticator.h
@@ -35,6 +35,7 @@
 #include <mqbconfm_messages.h>
 #include <mqbnet_authenticationcontext.h>
 #include <mqbnet_authenticator.h>
+#include <mqbnet_initialconnectioncontext.h>
 
 // BMQ
 #include <bmqio_channel.h>
@@ -78,10 +79,15 @@ class Authenticator : public mqbnet::Authenticator {
         bdlcc::ObjectPoolFunctors::RemoveAll<bdlbb::Blob> >
         BlobSpPool;
 
+    typedef mqbnet::AuthenticationContext::State State;
+
   private:
     typedef bsl::shared_ptr<mqbnet::AuthenticationContext>
         AuthenticationContextSp;
 
+    typedef bsl::shared_ptr<mqbnet::InitialConnectionContext>
+        InitialConnectionContextSp;
+
   private:
     // DATA
 
@@ -102,20 +108,23 @@ class Authenticator : public mqbnet::Authenticator {
 
     /// Invoked when received a `AuthenticationRequest` authentication message
     /// with the specified `context`.  The behavior of this function is
-    /// undefined unless `d_authenticationMessage` in the `context` is an
-    /// `AuthenticateRequest` and this request is incoming or reversed
-    /// connection.  Returns 0 on success, or return a non-zero code and
-    /// populate the specified `errorDescription` with a description of the
-    /// error on failure.
-    int onAuthenticationRequest(bsl::ostream& errorDescription,
-                                const AuthenticationContextSp& context);
+    /// undefined unless `authenticationMsg` is an `AuthenticateRequest` and
+    /// this request is an incoming connection.  Returns 0 on success,
+    /// or return a non-zero code and populate the specified `errorDescription`
+    /// with a description of the error on failure.
+    int onAuthenticationRequest(
+        bsl::ostream&                              errorDescription,
+        const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+        const InitialConnectionContextSp&          context);
 
     /// Invoked when received a `AuthenticationResponse` authentication message
     /// with the specified `context`.  Returns 0 on success, or return a
     /// non-zero code and populate the specified `errorDescription` with
     /// a description of the error on failure.
-    int onAuthenticationResponse(bsl::ostream& errorDescription,
-                                 const AuthenticationContextSp& context);
+    int onAuthenticationResponse(
+        bsl::ostream&                              errorDescription,
+        const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+        const InitialConnectionContextSp&          context);
 
     /// Send the specified `message` to the peer associated with the
     /// specified `context` and return 0 on success, or return a non-zero
@@ -152,13 +161,21 @@ class Authenticator : public mqbnet::Authenticator {
     // MANIPULATORS
     //   (virtual: mqbnet::Authenticator)
 
-    int handleAuthentication(bsl::ostream&                  errorDescription,
-                             bool*                          isContinueRead,
-                             const AuthenticationContextSp& context)
-        BSLS_KEYWORD_OVERRIDE;
+    /// Authenticate the connection based on the type of AuthenticationMessage
+    /// `authenticationMsg`.  Set `isContinueRead` to true if we want to
+    /// continue reading instead of finishing authentication.
+    /// Return 0 on success, or a non-zero error code and populate the
+    /// specified `errorDescription` with a description of the error otherwise.
+    int handleAuthentication(
+        bsl::ostream&                              errorDescription,
+        const bmqp_ctrlmsg::AuthenticationMessage& authenticationMsg,
+        bool*                                      isContinueRead,
+        const InitialConnectionContextSp& context) BSLS_KEYWORD_OVERRIDE;
 
     /// Send out outbound authentication message or reverse connection request
     /// with the specified `context`.
+    /// Return 0 on success, or a non-zero error code and populate the
+    /// specified `errorDescription` with a description of the error otherwise.
     int authenticationOutboundOrReverse(const AuthenticationContextSp& context)
         BSLS_KEYWORD_OVERRIDE;
 };
diff --git a/src/groups/mqb/mqba/mqba_initialconnectionhandler.cpp b/src/groups/mqb/mqba/mqba_initialconnectionhandler.cpp
@@ -191,13 +191,11 @@ int InitialConnectionHandler::processBlob(
 
     // Authentication or Negotiation based on the type of message received.
     if (authenticationMsg.has_value()) {
-        context->authenticationContext()->setAuthenticationMessage(
-            authenticationMsg.value());
-
         rc = d_authenticator_mp->handleAuthentication(
             errorDescription,
+            authenticationMsg.value(),
             isContinueRead,
-            context->authenticationContext());
+            context);
     }
     else if (negotiationMsg.has_value()) {
         context->negotiationContext()->d_negotiationMessage =
@@ -344,15 +342,6 @@ InitialConnectionHandler::~InitialConnectionHandler()
 void InitialConnectionHandler::setupContext(
     const InitialConnectionContextSp& context)
 {
-    // Create an AuthenticationContext for that connection
-    bsl::shared_ptr<mqbnet::AuthenticationContext> authenticationContext =
-        bsl::allocate_shared<mqbnet::AuthenticationContext>(
-            d_allocator_p,
-            context.get(),  // initialConnectionContext
-            false           // isReversed
-        );
-    context->setAuthenticationContext(authenticationContext);
-
     // Create an NegotiationContext for that connection
     bsl::shared_ptr<mqbnet::NegotiationContext> negotiationContext;
     negotiationContext.createInplace(d_allocator_p);
diff --git a/src/groups/mqb/mqbnet/mqbnet_authenticationcontext.h b/src/groups/mqb/mqbnet/mqbnet_authenticationcontext.h
@@ -43,8 +43,7 @@ class AuthenticationContext {
   public:
     // TYPES
     enum State {
-        e_IDLE = 0,
-        e_AUTHENTICATING,
+        e_AUTHENTICATING = 0,
         e_AUTHENTICATED,
     };
 
@@ -72,7 +71,7 @@ class AuthenticationContext {
     AuthenticationContext(
         InitialConnectionContext* initialConnectionContext,
         bool                      isReversed,
-        State                     state          = State::e_IDLE,
+        State                     state,
         ConnectionType::Enum      connectionType = ConnectionType::e_UNKNOWN,
         bslma::Allocator*         basicAllocator = 0);
 
diff --git a/src/groups/mqb/mqbnet/mqbnet_authenticator.h b/src/groups/mqb/mqbnet/mqbnet_authenticator.h
@@ -25,6 +25,7 @@
 /// (re)authenticates a connection with a BlazingMQ client or another bmqbrkr.
 
 // MQB
+#include "mqbnet_initialconnectioncontext.h"
 #include <mqbnet_authenticationcontext.h>
 
 // BDE
@@ -50,14 +51,15 @@ class Authenticator {
     // MANIPULATORS
 
     /// Authenticate the connection based on the type of AuthenticationMessage
-    /// in the specified `context`.  Set `isContinueRead` to true if we want to
+    /// `authenticationMsg`.  Set `isContinueRead` to true if we want to
     /// continue reading instead of finishing authentication.
     /// Return 0 on success, or a non-zero error code and populate the
     /// specified `errorDescription` with a description of the error otherwise.
     virtual int handleAuthentication(
-        bsl::ostream&                                 errorDescription,
-        bool*                                         isContinueRead,
-        const bsl::shared_ptr<AuthenticationContext>& context) = 0;
+        bsl::ostream&                                    errorDescription,
+        const bmqp_ctrlmsg::AuthenticationMessage&       authenticationMsg,
+        bool*                                            isContinueRead,
+        const bsl::shared_ptr<InitialConnectionContext>& context) = 0;
 
     /// Send out outbound authentication message or reverse connection request
     /// with the specified `context`.
