From 2581009397cd43233df4b8af332584948c473d34 Mon Sep 17 00:00:00 2001
From: Zang MingJie <mingjiez@google.com>
Date: Sat, 18 Jun 2022 10:42:58 +0800
Subject: [PATCH] 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 <bzbarsky@apple.com>

* Resolve comments

Co-authored-by: Boris Zbarsky <bzbarsky@apple.com>
---
 src/messaging/ExchangeContext.cpp      |  3 +++
 src/messaging/ExchangeMgr.cpp          | 18 ++++++++++++++++++
 src/messaging/ExchangeMgr.h            |  5 +++++
 src/messaging/ReliableMessageContext.h | 18 ++++++++++++++++++
 src/transport/SecureSession.cpp        | 22 ++++++++++++++++++++++
 src/transport/SecureSession.h          | 10 ++++++++++
 src/transport/SessionManager.cpp       | 17 +++++++++++++++++
 src/transport/SessionManager.h         |  5 +++++
 8 files changed, 98 insertions(+)

diff --git a/src/messaging/ExchangeContext.cpp b/src/messaging/ExchangeContext.cpp
index a2aa221a7cc78e..422a2f88e34ce4 100644
--- a/src/messaging/ExchangeContext.cpp
+++ b/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);
diff --git a/src/messaging/ExchangeMgr.cpp b/src/messaging/ExchangeMgr.cpp
index 2c1242aff55c1c..b01d5b4fe4cb5c 100644
--- a/src/messaging/ExchangeMgr.cpp
+++ b/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
diff --git a/src/messaging/ExchangeMgr.h b/src/messaging/ExchangeMgr.h
index d6b775436799f4..112cddd82c7ecb 100644
--- a/src/messaging/ExchangeMgr.h
+++ b/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; };
diff --git a/src/messaging/ReliableMessageContext.h b/src/messaging/ReliableMessageContext.h
index e67b2d596e963e..9b2f8d67ab0539 100644
--- a/src/messaging/ReliableMessageContext.h
+++ b/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
diff --git a/src/transport/SecureSession.cpp b/src/transport/SecureSession.cpp
index f2ec416c0ba9d6..9363d92b323f21 100644
--- a/src/transport/SecureSession.cpp
+++ b/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;
diff --git a/src/transport/SecureSession.h b/src/transport/SecureSession.h
index a505a3237f00eb..829b179965a8b2 100644
--- a/src/transport/SecureSession.h
+++ b/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;
diff --git a/src/transport/SessionManager.cpp b/src/transport/SessionManager.cpp
index a9e95411e778aa..193f0d222eb88c 100644
--- a/src/transport/SessionManager.cpp
+++ b/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)
 {
diff --git a/src/transport/SessionManager.h b/src/transport/SessionManager.h
index 1b01f1d5bcb5e7..d88031f670f4ac 100644
--- a/src/transport/SessionManager.h
+++ b/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.