From 204d5904e6e5145db467063df041c997975b3203 Mon Sep 17 00:00:00 2001
From: Pradip De <pradipd@google.com>
Date: Tue, 17 Dec 2024 23:02:54 +0000
Subject: [PATCH] Fix for Bug #36731.

Add CloseActiveConnections() call in TCPBase::Close(), which
is called as part of Server::Shutdown().
Active connections should be closed as part of Server shutdown.
This allows the TCPConnectionState to also close the associated
TCPEndpoint object as part of this shutdown flow.

Previously, the CloseActiveConnections() call was present in the
TCPBase destructor alone.

Add test for Connection Close() and checking for TCPEndPoint.
---
 src/transport/raw/ActiveTCPConnectionState.h |  5 ++++-
 src/transport/raw/TCP.cpp                    | 13 +++++------
 src/transport/raw/tests/TestTCP.cpp          | 23 ++++++++++++++++++++
 3 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/transport/raw/ActiveTCPConnectionState.h b/src/transport/raw/ActiveTCPConnectionState.h
index 0f53d4479e9300..2176048cc25364 100644
--- a/src/transport/raw/ActiveTCPConnectionState.h
+++ b/src/transport/raw/ActiveTCPConnectionState.h
@@ -62,7 +62,10 @@ struct ActiveTCPConnectionState
 
     void Free()
     {
-        mEndPoint->Free();
+        if (mEndPoint)
+        {
+            mEndPoint->Free();
+        }
         mPeerAddr = PeerAddress::Uninitialized();
         mEndPoint = nullptr;
         mReceived = nullptr;
diff --git a/src/transport/raw/TCP.cpp b/src/transport/raw/TCP.cpp
index f12ee7892ad2d0..9976874d8dc6f1 100644
--- a/src/transport/raw/TCP.cpp
+++ b/src/transport/raw/TCP.cpp
@@ -55,14 +55,8 @@ constexpr int kListenBacklogSize = 2;
 
 TCPBase::~TCPBase()
 {
-    if (mListenSocket != nullptr)
-    {
-        // endpoint is only non null if it is initialized and listening
-        mListenSocket->Free();
-        mListenSocket = nullptr;
-    }
-
-    CloseActiveConnections();
+    // Call Close to free the listening socket and close all active connections.
+    Close();
 }
 
 void TCPBase::CloseActiveConnections()
@@ -125,6 +119,9 @@ void TCPBase::Close()
         mListenSocket->Free();
         mListenSocket = nullptr;
     }
+
+    CloseActiveConnections();
+
     mState = TCPState::kNotReady;
 }
 
diff --git a/src/transport/raw/tests/TestTCP.cpp b/src/transport/raw/tests/TestTCP.cpp
index 80d56707481c45..00eb0562c96172 100644
--- a/src/transport/raw/tests/TestTCP.cpp
+++ b/src/transport/raw/tests/TestTCP.cpp
@@ -610,6 +610,29 @@ TEST_F(TestTCP, HandleConnCloseCalledTest6)
     HandleConnCloseTest(addr);
 }
 
+TEST_F(TestTCP, CheckTCPEndpointAfterCloseTest)
+{
+    TCPImpl tcp;
+
+    IPAddress addr;
+    IPAddress::FromString("::1", addr);
+
+    MockTransportMgrDelegate gMockTransportMgrDelegate(mIOContext);
+    gMockTransportMgrDelegate.InitializeMessageTest(tcp, addr);
+    gMockTransportMgrDelegate.ConnectTest(tcp, addr);
+
+    Transport::PeerAddress lPeerAddress = Transport::PeerAddress::TCP(addr, gChipTCPPort);
+    void * state                        = TestAccess::FindActiveConnection(tcp, lPeerAddress);
+    ASSERT_NE(state, nullptr);
+    TCPEndPoint * lEndPoint = TestAccess::GetEndpoint(state);
+    ASSERT_NE(lEndPoint, nullptr);
+
+    // Call Close and check the TCPEndpoint
+    tcp.Close();
+    lEndPoint = TestAccess::GetEndpoint(state);
+    ASSERT_EQ(lEndPoint, nullptr);
+}
+
 TEST_F(TestTCP, CheckProcessReceivedBuffer)
 {
     TCPImpl tcp;