From 138fb4f0d23bb5353e8b17b32de1a3893fa5a25d Mon Sep 17 00:00:00 2001
From: Pradip De <pradipd@google.com>
Date: Tue, 23 Jul 2024 17:47:15 -0700
Subject: [PATCH] Mark session for eviction on TCP Client side when
 disconnecting with peer. (#34451)

The TCP server on receiving a TCPDisconnect from the client marks
the corresponding secure session for eviction.
The TCP client should also mark its session for eviction when proactively
closing the connection with peer.
---
 src/transport/SessionManager.cpp | 51 +++++++++++++++++++-------------
 src/transport/SessionManager.h   |  4 +++
 2 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/transport/SessionManager.cpp b/src/transport/SessionManager.cpp
index e22f0722d7cb55..17065e319cf3bd 100644
--- a/src/transport/SessionManager.cpp
+++ b/src/transport/SessionManager.cpp
@@ -714,27 +714,7 @@ void SessionManager::HandleConnectionClosed(Transport::ActiveTCPConnectionState
 {
     VerifyOrReturn(conn != nullptr);
 
-    // Mark the corresponding secure sessions as defunct
-    mSecureSessions.ForEachSession([&](auto session) {
-        if (session->IsActiveSession() && session->GetTCPConnection() == conn)
-        {
-            SessionHandle handle(*session);
-            // Notify the SessionConnection delegate of the connection
-            // closure.
-            if (mConnDelegate != nullptr)
-            {
-                mConnDelegate->OnTCPConnectionClosed(handle, conErr);
-            }
-
-            // Dis-associate the connection from session by setting it to a
-            // nullptr.
-            session->SetTCPConnection(nullptr);
-            // Mark session as defunct
-            session->MarkAsDefunct();
-        }
-
-        return Loop::Continue;
-    });
+    MarkSecureSessionOverTCPForEviction(conn, conErr);
 
     // TODO: A mechanism to mark an unauthenticated session as unusable when
     // the underlying connection is broken. Issue #32323
@@ -785,6 +765,8 @@ void SessionManager::TCPDisconnect(Transport::ActiveTCPConnectionState * conn, b
         conn->mPeerAddr.ToString(peerAddrBuf);
         ChipLogProgress(Inet, "Disconnecting TCP connection from peer at %s.", peerAddrBuf);
         mTransportMgr->TCPDisconnect(conn, shouldAbort);
+
+        MarkSecureSessionOverTCPForEviction(conn, CHIP_NO_ERROR);
     }
 }
 #endif // INET_CONFIG_ENABLE_TCP_ENDPOINT
@@ -1336,6 +1318,33 @@ Optional<SessionHandle> SessionManager::FindSecureSessionForNode(ScopedNodeId pe
     return mrpSession != nullptr ? MakeOptional<SessionHandle>(*mrpSession) : Optional<SessionHandle>::Missing();
 }
 
+#if INET_CONFIG_ENABLE_TCP_ENDPOINT
+void SessionManager::MarkSecureSessionOverTCPForEviction(Transport::ActiveTCPConnectionState * conn, CHIP_ERROR conErr)
+{
+    // Mark the corresponding secure sessions for eviction
+    mSecureSessions.ForEachSession([&](auto session) {
+        if (session->IsActiveSession() && session->GetTCPConnection() == conn)
+        {
+            SessionHandle handle(*session);
+            // Notify the SessionConnection delegate of the connection
+            // closure.
+            if (mConnDelegate != nullptr)
+            {
+                mConnDelegate->OnTCPConnectionClosed(handle, conErr);
+            }
+
+            // Dis-associate the connection from session by setting it to a
+            // nullptr.
+            session->SetTCPConnection(nullptr);
+            // Mark session for eviction.
+            session->MarkForEviction();
+        }
+
+        return Loop::Continue;
+    });
+}
+#endif // INET_CONFIG_ENABLE_TCP_ENDPOINT
+
 /**
  * Provides a means to get diagnostic information such as number of sessions.
  */
diff --git a/src/transport/SessionManager.h b/src/transport/SessionManager.h
index 7d0c5645365f97..67871d4ee83b2a 100644
--- a/src/transport/SessionManager.h
+++ b/src/transport/SessionManager.h
@@ -616,6 +616,10 @@ class DLL_EXPORT SessionManager : public TransportMgrDelegate, public FabricTabl
 
     void OnReceiveError(CHIP_ERROR error, const Transport::PeerAddress & source);
 
+#if INET_CONFIG_ENABLE_TCP_ENDPOINT
+    void MarkSecureSessionOverTCPForEviction(Transport::ActiveTCPConnectionState * conn, CHIP_ERROR conErr);
+#endif // INET_CONFIG_ENABLE_TCP_ENDPOINT
+
     static bool IsControlMessage(PayloadHeader & payloadHeader)
     {
         return payloadHeader.HasMessageType(Protocols::SecureChannel::MsgType::MsgCounterSyncReq) ||