From 6453535113a24bf89119ea72f8603acafb6ca452 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 b73540d7956f29..70d9ed3fdd440e 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 80531491f288a0..78381559a8a364 100644
--- a/src/transport/raw/tests/TestTCP.cpp
+++ b/src/transport/raw/tests/TestTCP.cpp
@@ -609,6 +609,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;