Fix crashes when closing all exchanges for fabric. (#19780)

* Fix crashes when closing all exchanges for fabric.

Our "close all exchanges except this special one for this fabric" API
messes up exchange refcounting, leading to use-after-free.

The fix is to reuse, as much as possible, the normal "session is going
away" flow to notify exchanges, and other session consumers, that the
sessions are in fact going away.

Fixes #19747

* Address review comments.

* Updates to fix fallout from #19502.

We need to allow messages on inactive sessions to reach the exchange manager,
because such sessions need to be able to deliver an MRP ack to an exchange
waiting for one.

We also don't want to crash on an attempt to transition from Inactive to Defunct
state; the transition should just be ignored.  This way if we start trying to
transitionin to Defunct on MRP delivery failures we will not start crashing if
such a failure happens on an Inactive session.

* Address review comment

* Address review comments

* Address Jerry's review comments.

src/messaging/ExchangeContext.cpp

@@ -325,9 +325,6 @@ ExchangeContext::~ExchangeContext()
     VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() == 0);
     VerifyOrDie(!IsAckPending());
 
-    if (ReleaseSessionOnDestruction() && mSession)
-        mSession->AsSecureSession()->MarkForEviction();
-
 #if CONFIG_DEVICE_LAYER && CHIP_DEVICE_CONFIG_ENABLE_SED
     // Make sure that the exchange withdraws the request for Sleepy End Device active mode.
     UpdateSEDIntervalMode(false);
@@ -371,6 +368,11 @@ bool ExchangeContext::MatchExchange(const SessionHandle & session, const PacketH
 
 void ExchangeContext::OnSessionReleased()
 {
+    if (ShouldIgnoreSessionRelease())
+    {
+        return;
+    }
+
     if (mFlags.Has(Flags::kFlagClosed))
     {
         // Exchange is already being closed. It may occur when closing an exchange after sending
@@ -573,5 +575,31 @@ ExchangeMessageDispatch & ExchangeContext::GetMessageDispatch(bool isEphemeralEx
     return ApplicationExchangeDispatch::Instance();
 }
 
+void ExchangeContext::AbortAllOtherCommunicationOnFabric()
+{
+    if (!mSession || !mSession->IsSecureSession())
+    {
+        ChipLogError(ExchangeManager, "AbortAllOtherCommunicationOnFabric called when we don't have a PASE/CASE session");
+        return;
+    }
+
+    // Save our session so it does not actually go away.
+    Optional<SessionHandle> session = mSession.Get();
+
+    SetIgnoreSessionRelease(true);
+
+    GetExchangeMgr()->GetSessionManager()->ExpireAllPairingsForFabric(mSession->GetFabricIndex());
+
+    mSession.GrabExpiredSession(session.Value());
+
+    SetIgnoreSessionRelease(false);
+}
+
+void ExchangeContext::ExchangeSessionHolder::GrabExpiredSession(const SessionHandle & session)
+{
+    VerifyOrDie(session->AsSecureSession()->IsPendingEviction());
+    GrabUnchecked(session);
+}
+
 } // namespace Messaging
 } // namespace chip

src/messaging/ExchangeContext.h

@@ -184,15 +184,33 @@ class DLL_EXPORT ExchangeContext : public ReliableMessageContext,
     // using this function is recommended.
     void SetResponseTimeout(Timeout timeout);
 
+    // This API is used by commands that need to shut down all existing
+    // sessions/exchanges on a fabric but need to make sure the response to the
+    // command still goes out on the exchange the command came in on.  This API
+    // will ensure that all secure sessions for the fabric this exchanges is on
+    // are released except the one this exchange is using, and will release
+    // that session once this exchange is done sending the response.
+    //
+    // This API is a no-op if called on an exchange that is not using a
+    // SecureSession.
+    void AbortAllOtherCommunicationOnFabric();
+
 private:
+    class ExchangeSessionHolder : public SessionHolderWithDelegate
+    {
+    public:
+        ExchangeSessionHolder(ExchangeContext & exchange) : SessionHolderWithDelegate(exchange) {}
+        void GrabExpiredSession(const SessionHandle & session);
+    };
+
     Timeout mResponseTimeout{ 0 }; // Maximum time to wait for response (in milliseconds); 0 disables response timeout.
     ExchangeDelegate * mDelegate   = nullptr;
     ExchangeManager * mExchangeMgr = nullptr;
 
     ExchangeMessageDispatch & mDispatch;
 
-    SessionHolderWithDelegate mSession; // The connection state
-    uint16_t mExchangeId;               // Assigned exchange ID.
+    ExchangeSessionHolder mSession; // The connection state
+    uint16_t mExchangeId;           // Assigned exchange ID.
 
     /**
      *  Determine whether a response is currently expected for a message that was sent over
@@ -274,7 +292,22 @@ class DLL_EXPORT ExchangeContext : public ReliableMessageContext,
     void UpdateSEDIntervalMode(bool activeMode);
 
     static ExchangeMessageDispatch & GetMessageDispatch(bool isEphemeralExchange, ExchangeDelegate * delegate);
+
+    // If SetAutoReleaseSession() is called, this exchange must be using a SecureSession, and should
+    // evict it when the exchange is done with all its work (including any MRP traffic).
+    inline void SetIgnoreSessionRelease(bool ignore);
+    inline bool ShouldIgnoreSessionRelease();
 };
 
+inline void ExchangeContext::SetIgnoreSessionRelease(bool ignore)
+{
+    mFlags.Set(Flags::kFlagIgnoreSessionRelease, ignore);
+}
+
+inline bool ExchangeContext::ShouldIgnoreSessionRelease()
+{
+    return mFlags.Has(Flags::kFlagIgnoreSessionRelease);
+}
+
 } // namespace Messaging
 } // namespace chip

src/messaging/ExchangeMgr.cpp

@@ -236,7 +236,14 @@ void ExchangeManager::OnMessageReceived(const PacketHeader & packetHeader, const
                         packetHeader.GetDestinationGroupId().Value());
     }
 
-    // Do not handle unsolicited messages on a inactive session.
+    // Do not handle messages that don't match an existing exchange on an
+    // inactive session, since we should not be creating new exchanges there.
+    if (!session->IsActiveSession())
+    {
+        ChipLogProgress(ExchangeManager, "Dropping message on inactive session that does not match an existing exchange");
+        return;
+    }
+
     // If it's not a duplicate message, search for an unsolicited message handler if it is marked as being sent by an initiator.
     // Since we didn't find an existing exchange that matches the message, it must be an unsolicited message. However all
     // unsolicited messages must be marked as being from an initiator.
@@ -376,23 +383,5 @@ void ExchangeManager::CloseAllContextsForDelegate(const ExchangeDelegate * deleg
     });
 }
 
-void ExchangeManager::AbortExchangesForFabricExceptOne(FabricIndex fabricIndex, ExchangeContext * deferred)
-{
-    VerifyOrDie(deferred->HasSessionHandle() && deferred->GetSessionHandle()->IsSecureSession());
-
-    mContextPool.ForEachActiveObject([&](auto * ec) {
-        if (ec->HasSessionHandle() && ec->GetSessionHandle()->GetFabricIndex() == fabricIndex)
-        {
-            if (ec == deferred)
-                ec->SetAutoReleaseSession();
-            else
-                ec->Abort();
-        }
-        return Loop::Continue;
-    });
-
-    mSessionManager->ReleaseSessionsForFabricExceptOne(fabricIndex, deferred->GetSessionHandle());
-}
-
 } // namespace Messaging
 } // namespace chip

src/messaging/ExchangeMgr.h

@@ -183,11 +183,6 @@ class DLL_EXPORT ExchangeManager : public SessionMessageDelegate
      */
     void CloseAllContextsForDelegate(const ExchangeDelegate * delegate);
 
-    // This API is used by commands that need to shut down all existing exchanges on a fabric but need to make sure the response to
-    // the command still goes out on the exchange the command came in on.  This API flags that one exchange to shut down its session
-    // when it's done.
-    void AbortExchangesForFabricExceptOne(FabricIndex fabricIndex, ExchangeContext * deferred);
-
     SessionManager * GetSessionManager() const { return mSessionManager; }
 
     ReliableMessageMgr * GetReliableMessageMgr() { return &mReliableMessageMgr; };

src/messaging/ReliableMessageContext.h

@@ -124,11 +124,6 @@ class ReliableMessageContext
     /// Determine whether this exchange is a EphemeralExchange for replying a StandaloneAck
     bool IsEphemeralExchange() const;
 
-    // If SetAutoReleaseSession() is called, this exchange must be using a SecureSession, and should
-    // evict it when the exchange is done with all its work (including any MRP traffic).
-    void SetAutoReleaseSession();
-    bool ReleaseSessionOnDestruction();
-
     /**
      * Get the reliable message manager that corresponds to this reliable
      * message context.
@@ -170,8 +165,8 @@ class ReliableMessageContext
         /// When set, signifies that the exchange created sorely for replying a StandaloneAck
         kFlagEphemeralExchange = (1u << 9),
 
-        /// When set, automatically release the session when this exchange is destroyed.
-        kFlagAutoReleaseSession = (1u << 10),
+        /// When set, ignore session being released, because we are releasing it ourselves.
+        kFlagIgnoreSessionRelease = (1u << 10),
     };
 
     BitFlags<Flags> mFlags; // Internal state flags
@@ -253,15 +248,5 @@ inline bool ReliableMessageContext::IsEphemeralExchange() const
     return mFlags.Has(Flags::kFlagEphemeralExchange);
 }
 
-inline void ReliableMessageContext::SetAutoReleaseSession()
-{
-    mFlags.Set(Flags::kFlagAutoReleaseSession, true);
-}
-
-inline bool ReliableMessageContext::ReleaseSessionOnDestruction()
-{
-    return mFlags.Has(Flags::kFlagAutoReleaseSession);
-}
-
 } // namespace Messaging
 } // namespace chip

src/messaging/tests/BUILD.gn

@@ -43,11 +43,13 @@ static_library("helpers") {
 chip_test_suite("tests") {
   output_name = "libMessagingLayerTests"
 
-  sources = [ "TestMessagingLayer.h" ]
+  test_sources = []
 
   if (chip_device_platform != "efr32") {
     # TODO(#10447): ReliableMessage Test has HF, and ExchangeMgr hangs on EFR32.
-    sources += [
+    # And TestAbortExchangesForFabric does not link on EFR32 for some reason.
+    test_sources += [
+      "TestAbortExchangesForFabric.cpp",
       "TestExchangeMgr.cpp",
       "TestReliableMessageProtocol.cpp",
     ]
@@ -67,9 +69,4 @@ chip_test_suite("tests") {
     "${nlio_root}:nlio",
     "${nlunit_test_root}:nlunit-test",
   ]
-
-  tests = [
-    "TestExchangeMgr",
-    "TestReliableMessageProtocol",
-  ]
 }

src/messaging/tests/MessagingContext.h

@@ -103,8 +103,6 @@ class MessagingContext : public PlatformMemoryUser
     static const uint16_t kAliceKeyId   = 2;
     static const uint16_t kCharlieKeyId = 3;
     static const uint16_t kDavidKeyId   = 4;
-    NodeId GetBobNodeId() const;
-    NodeId GetAliceNodeId() const;
     GroupId GetFriendsGroupId() const { return mFriendsGroupId; }
 
     SessionManager & GetSecureSessionManager() { return mSessionManager; }