From 1077876eee2cff0aa126bb792d2484bd187cbfcf Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Thu, 30 Mar 2023 20:06:59 -0400
Subject: [PATCH] Gracefully handle an exchange delegate evicting the session
 on timeout. (#25908)

Otherwise the session eviction (which closes our exchange) would destroy it, and
then we would unwind into ExchangeContext::NotifyResponseTimeout and try to use
the exchange's members.
---
 src/messaging/ExchangeContext.cpp       |  5 ++++
 src/messaging/tests/TestExchangeMgr.cpp | 36 +++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/src/messaging/ExchangeContext.cpp b/src/messaging/ExchangeContext.cpp
index e8447527e8dc90..daee580971b404 100644
--- a/src/messaging/ExchangeContext.cpp
+++ b/src/messaging/ExchangeContext.cpp
@@ -480,6 +480,11 @@ void ExchangeContext::NotifyResponseTimeout(bool aCloseIfNeeded)
 {
     SetResponseExpected(false);
 
+    // Hold a ref to ourselves so we can make calls into our delegate that might
+    // decrease our refcount (e.g. by expiring out session) without worrying
+    // about use-after-free.
+    ExchangeHandle ref(*this);
+
     // mSession might be null if this timeout is due to the session being
     // evicted.
     if (mSession)
diff --git a/src/messaging/tests/TestExchangeMgr.cpp b/src/messaging/tests/TestExchangeMgr.cpp
index ae6fcca8c97570..38aec126549fe6 100644
--- a/src/messaging/tests/TestExchangeMgr.cpp
+++ b/src/messaging/tests/TestExchangeMgr.cpp
@@ -90,6 +90,15 @@ class WaitForTimeoutDelegate : public ExchangeDelegate
     bool IsOnResponseTimeoutCalled = false;
 };
 
+class ExpireSessionFromTimeoutDelegate : public WaitForTimeoutDelegate
+{
+    void OnResponseTimeout(ExchangeContext * ec) override
+    {
+        ec->GetSessionHandle()->AsSecureSession()->MarkForEviction();
+        WaitForTimeoutDelegate::OnResponseTimeout(ec);
+    }
+};
+
 void CheckNewContextTest(nlTestSuite * inSuite, void * inContext)
 {
     TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
@@ -162,6 +171,32 @@ void CheckSessionExpirationTimeout(nlTestSuite * inSuite, void * inContext)
     NL_TEST_ASSERT(inSuite, ctx.CreateSessionAliceToBob() == CHIP_NO_ERROR);
 }
 
+void CheckSessionExpirationDuringTimeout(nlTestSuite * inSuite, void * inContext)
+{
+    using namespace chip::System::Clock::Literals;
+
+    TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
+
+    ExpireSessionFromTimeoutDelegate sendDelegate;
+    ExchangeContext * ec1 = ctx.NewExchangeToBob(&sendDelegate);
+
+    ec1->SetResponseTimeout(System::Clock::Timeout(100));
+
+    ec1->SendMessage(Protocols::BDX::Id, kMsgType_TEST1, System::PacketBufferHandle::New(System::PacketBuffer::kMaxSize),
+                     SendFlags(Messaging::SendMessageFlags::kExpectResponse).Set(Messaging::SendMessageFlags::kNoAutoRequestAck));
+
+    ctx.DrainAndServiceIO();
+    NL_TEST_ASSERT(inSuite, !sendDelegate.IsOnResponseTimeoutCalled);
+
+    // Wait for our timeout to elapse. Give it an extra 100ms.
+    ctx.GetIOContext().DriveIOUntil(200_ms32, [&sendDelegate] { return sendDelegate.IsOnResponseTimeoutCalled; });
+
+    NL_TEST_ASSERT(inSuite, sendDelegate.IsOnResponseTimeoutCalled);
+
+    // recreate closed session.
+    NL_TEST_ASSERT(inSuite, ctx.CreateSessionAliceToBob() == CHIP_NO_ERROR);
+}
+
 void CheckUmhRegistrationTest(nlTestSuite * inSuite, void * inContext)
 {
     TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
@@ -237,6 +272,7 @@ const nlTest sTests[] =
     NL_TEST_DEF("Test ExchangeMgr::CheckExchangeMessages",    CheckExchangeMessages),
     NL_TEST_DEF("Test OnConnectionExpired basics",            CheckSessionExpirationBasics),
     NL_TEST_DEF("Test OnConnectionExpired timeout handling",  CheckSessionExpirationTimeout),
+    NL_TEST_DEF("Test session eviction in timeout handling",  CheckSessionExpirationDuringTimeout),
 
     NL_TEST_SENTINEL()
 };