Add API to invalid sessions/exchanges for UpdateNOC command (#19328)

* Add API to invalid sessions/exchanges for UpdateNOC command

* Complete the idea

* Add purge all fabric API

* Resolve comments: adjust function name

* Restyle

* Fix rebase conflict

* Apply suggestions from code review

Co-authored-by: Boris Zbarsky <[email protected]>

* Resolve comments

Co-authored-by: Boris Zbarsky <[email protected]>

src/messaging/ExchangeContext.cpp

@@ -325,6 +325,9 @@ ExchangeContext::~ExchangeContext()
     VerifyOrDie(mExchangeMgr != nullptr && GetReferenceCount() == 0);
     VerifyOrDie(!IsAckPending());
 
+    if (ReleaseSessionOnDestruction() && mSession)
+        mSession->AsSecureSession()->MarkForRemoval();
+
 #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);

src/messaging/ExchangeMgr.cpp

@@ -376,5 +376,23 @@ 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,6 +183,11 @@ 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,6 +124,11 @@ 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.
@@ -164,6 +169,9 @@ 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),
     };
 
     BitFlags<Flags> mFlags; // Internal state flags
@@ -245,5 +253,15 @@ 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/transport/SecureSession.cpp

@@ -39,6 +39,7 @@ void SecureSession::MarkForRemoval()
         NotifySessionReleased();
         return;
     case State::kActive:
+    case State::kInactive:
         Release(); // Decrease the ref which is retained at Activate
         mState = State::kPendingRemoval;
         NotifySessionReleased();
@@ -49,6 +50,27 @@ void SecureSession::MarkForRemoval()
     }
 }
 
+void SecureSession::MarkInactive()
+{
+    ChipLogDetail(Inet, "SecureSession[%p]: MarkInactive Type:%d LSID:%d", this, to_underlying(mSecureSessionType),
+                  mLocalSessionId);
+    ReferenceCountedHandle<Transport::Session> ref(*this);
+    switch (mState)
+    {
+    case State::kPairing:
+        VerifyOrDie(false);
+        return;
+    case State::kActive:
+        // By setting this state, IsActiveSession() will return false, which prevents creating new exchanges.
+        mState = State::kInactive;
+        return;
+    case State::kInactive:
+    case State::kPendingRemoval:
+        // Do nothing
+        return;
+    }
+}
+
 Access::SubjectDescriptor SecureSession::GetSubjectDescriptor() const
 {
     Access::SubjectDescriptor subjectDescriptor;

src/transport/SecureSession.h

@@ -144,6 +144,9 @@ class SecureSession : public Session, public ReferenceCounted<SecureSession, Sec
     /// @brief Mark as pending removal, all holders to this session will be cleared, and disallow future grab
     void MarkForRemoval();
 
+    // Used to prevent any new exchange created on the session while the existing exchanges finish their work.
+    void MarkInactive();
+
     Session::SessionType GetSessionType() const override { return Session::SessionType::kSecure; }
 #if CHIP_PROGRESS_LOGGING
     const char * GetSessionTypeString() const override { return "secure"; };
@@ -249,6 +252,13 @@ class SecureSession : public Session, public ReferenceCounted<SecureSession, Sec
         // grab this session, when all SessionHandles go out of scope, the
         // session object will be released automatically.
         kPendingRemoval = 3,
+
+        // The session is still functional, but it can't yield any new exchanges,
+        // This is meant to be used in conjunction with
+        // ExchangeManager::AbortExchangesForFabricExceptOne, with the one
+        // exceptional exchange handling moving this session out of this state when
+        // it finishes whatever it needs the session for.
+        kInactive = 4,
     };
 
     friend class SecureSessionDeleter;

src/transport/SessionManager.cpp

@@ -387,6 +387,23 @@ void SessionManager::ExpireAllPASEPairings()
     });
 }
 
+void SessionManager::ReleaseSessionsForFabricExceptOne(FabricIndex fabricIndex, const SessionHandle & deferred)
+{
+    VerifyOrDie(deferred->IsSecureSession());
+    SecureSession * deferredSecureSession = deferred->AsSecureSession();
+
+    mSecureSessions.ForEachSession([&](auto session) {
+        if (session->GetPeer().GetFabricIndex() == fabricIndex)
+        {
+            if (session == deferredSecureSession)
+                session->MarkInactive();
+            else
+                session->MarkForRemoval();
+        }
+        return Loop::Continue;
+    });
+}
+
 Optional<SessionHandle> SessionManager::AllocateSession(SecureSession::Type secureSessionType,
                                                         const ScopedNodeId & sessionEvictionHint)
 {

src/transport/SessionManager.h

@@ -176,6 +176,11 @@ class DLL_EXPORT SessionManager : public TransportMgrDelegate
     void ExpireAllPairingsForFabric(FabricIndex fabric);
     void ExpireAllPASEPairings();
 
+    // This API is used by commands that need to release all existing sessions 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 the release of the session used by the
+    // exchange is deferred until the exchange is done.
+    void ReleaseSessionsForFabricExceptOne(FabricIndex fabricIndex, const SessionHandle & deferred);
+
     /**
      * @brief
      *   Return the System Layer pointer used by current SessionManager.