Fix[BMQ]: handling early channel close when authenticating (#897)

Signed-off-by: dorjesinpo <[email protected]>

Author: dorjesinpo

src/groups/mqb/mqbnet/mqbnet_initialconnectioncontext.cpp

@@ -26,9 +26,13 @@ namespace mqbnet {
 // -------------------------------------
 
 InitialConnectionContext::InitialConnectionContext(bool isIncoming)
-: d_isIncoming(isIncoming)
-, d_resultState_p(0)
+: d_resultState_p(0)
 , d_userData_p(0)
+, d_channelSp()
+, d_initialConnectionCompleteCb()
+, d_negotiationCtxSp()
+, d_isIncoming(isIncoming)
+, d_isClosed(false)
 {
     // NOTHING
 }
@@ -71,6 +75,11 @@ InitialConnectionContext& InitialConnectionContext::setNegotiationContext(
     return *this;
 }
 
+void InitialConnectionContext::onClose()
+{
+    d_isClosed = true;
+}
+
 bool InitialConnectionContext::isIncoming() const
 {
     return d_isIncoming;
@@ -108,5 +117,10 @@ InitialConnectionContext::negotiationContext() const
     return d_negotiationCtxSp;
 }
 
+bool InitialConnectionContext::isClosed() const
+{
+    return d_isClosed;
+}
+
 }  // close package namespace
 }  // close enterprise namespace

src/groups/mqb/mqbnet/mqbnet_initialconnectioncontext.h

@@ -70,11 +70,6 @@ class InitialConnectionContext {
   private:
     // DATA
 
-    /// True if the session being negotiated originates
-    /// from a remote peer (i.e., a 'listen'); false if
-    /// it originates from us (i.e., a 'connect).
-    bool d_isIncoming;
-
     /// Raw pointer, held not owned, to some user data
     /// the session factory will pass back to the
     /// 'resultCb' method (used to inform of the
@@ -119,11 +114,19 @@ class InitialConnectionContext {
     /// The NegotiationContext updated upon receiving a negotiation message.
     bsl::shared_ptr<NegotiationContext> d_negotiationCtxSp;
 
+    /// True if the session being negotiated originates
+    /// from a remote peer (i.e., a 'listen'); false if
+    /// it originates from us (i.e., a 'connect).
+    bool d_isIncoming;
+
+    /// True if the associated channel is closed (with `onClose`).
+    bool d_isClosed;
+
   public:
     // CREATORS
 
     /// Create a new object having the specified `isIncoming` value.
-    InitialConnectionContext(bool isIncoming);
+    explicit InitialConnectionContext(bool isIncoming);
 
     ~InitialConnectionContext();
 
@@ -140,6 +143,9 @@ class InitialConnectionContext {
     InitialConnectionContext&
     setNegotiationContext(const bsl::shared_ptr<NegotiationContext>& value);
 
+    /// Called by the IO upon `onCLose` signal
+    void onClose();
+
     // ACCESSORS
 
     /// Return the value of the corresponding field.
@@ -148,6 +154,7 @@ class InitialConnectionContext {
     void*                                      resultState() const;
     const bsl::shared_ptr<bmqio::Channel>&     channel() const;
     const bsl::shared_ptr<NegotiationContext>& negotiationContext() const;
+    bool                                       isClosed() const;
 
     void complete(int                                     rc,
                   const bsl::string&                      error,

src/groups/mqb/mqbnet/mqbnet_tcpsessionfactory.cpp

@@ -366,6 +366,14 @@ void TCPSessionFactory::handleInitialConnection(
             bdlf::PlaceHolders::_5,  // initialConnectionContext
             context));
 
+    // Register as observer of the channel to get the 'onClose'
+    channel->onClose(
+        bdlf::BindUtil::bindS(d_allocator_p,
+                              &TCPSessionFactory::onClose,
+                              this,
+                              initialConnectionContext,
+                              bdlf::PlaceHolders::_1 /* bmqio::Status */));
+
     // NOTE: we must ensure the 'initialConnectionCompleteCb' can be invoked
     // from the
     //       'handleInitialConnection()' call as specified on the
@@ -558,6 +566,24 @@ void TCPSessionFactory::negotiationComplete(
 
         ++d_nbSessions;
 
+        if (isClientOrProxy(session.get())) {
+            ++d_nbOpenClients;
+        }
+
+        // check if the channel is not closed (we can be in authentication
+        // thread)
+
+        if (initialConnectionContext_p->isClosed()) {
+            BALL_LOG_WARN
+                << "#TCP_UNEXPECTED_STATE TCPSessionFactory '"
+                << d_config.name()
+                << "' got an already closed channel after negotiation.";
+
+            // Since the 'session' has missed 'onClose', call 'tearDown' here.
+            session->tearDown(session, false);
+            return;  // RETURN
+        }
+
         info.createInplace(d_allocator_p,
                            channel,
                            monitoredSession,
@@ -573,10 +599,6 @@ void TCPSessionFactory::negotiationComplete(
         inserted = d_channels.insert(toInsert);
         info     = inserted.first->second;
 
-        if (isClientOrProxy(info->d_session_sp.get())) {
-            ++d_nbOpenClients;
-        }
-
         if (info->d_monitor.isHearbeatEnabled() &&
             d_heartbeatSchedulerActive) {
             // Enable/Disable heartbeating under the lock
@@ -711,14 +733,6 @@ void TCPSessionFactory::channelStateCallback(
             // Keep track of active channels, for logging purposes
             ++d_nbActiveChannels;
 
-            // Register as observer of the channel to get the 'onClose'
-            channel->onClose(bdlf::BindUtil::bindS(
-                d_allocator_p,
-                &TCPSessionFactory::onClose,
-                this,
-                channel,
-                bdlf::PlaceHolders::_1 /* bmqio::Status */));
-
             handleInitialConnection(channel, context);
         }
     } break;
@@ -743,11 +757,17 @@ void TCPSessionFactory::channelStateCallback(
     }
 }
 
-void TCPSessionFactory::onClose(const bsl::shared_ptr<bmqio::Channel>& channel,
-                                const bmqio::Status&                   status)
+void TCPSessionFactory::onClose(
+    const bsl::shared_ptr<InitialConnectionContext>& initialConnectionContext,
+    const bmqio::Status&                             status)
 {
+    // Executed by one of the IO threads.
+
     --d_nbActiveChannels;
 
+    const bsl::shared_ptr<bmqio::Channel>& channel =
+        initialConnectionContext->channel();
+
     int port;
     channel->properties().load(
         &port,
@@ -758,6 +778,10 @@ void TCPSessionFactory::onClose(const bsl::shared_ptr<bmqio::Channel>& channel,
         // Lookup the session and remove it from internal map
         bslmt::LockGuard<bslmt::Mutex> guard(&d_mutex);  // LOCK
 
+        // set the 'isClosed' flag under lock to be checked under lock in
+        // 'negotiationComplete'.
+        initialConnectionContext->onClose();
+
         ChannelMap::const_iterator it = d_channels.find(channel.get());
         if (it != d_channels.end()) {
             channelInfo = it->second;

src/groups/mqb/mqbnet/mqbnet_tcpsessionfactory.h

@@ -457,8 +457,9 @@ class TCPSessionFactory {
     /// the channel's status, `userData` corresponding to the one provided
     /// when calling `addObserver` to register this object as observer of
     /// the channel.
-    virtual void onClose(const bsl::shared_ptr<bmqio::Channel>& channel,
-                         const bmqio::Status&                   status);
+    virtual void onClose(const bsl::shared_ptr<InitialConnectionContext>&
+                                              initialConnectionContext,
+                         const bmqio::Status& status);
 
     /// Reccuring scheduler event to check for all `heartbeat-enabled`
     /// channels : this will send a heartbeat if no data has been received