Don't automatically close exchanges if the underlying connection closes.

Someone might still have a reference to those exchanges, and we will
end up with an unbalanced release.  Instead just notify our consumer
in the specific case when we are waiting for a response so it knows
the response will never come (and presumably closes the exchange as
needed).

src/messaging/ExchangeContext.cpp

@@ -300,6 +300,21 @@ bool ExchangeContext::MatchExchange(SecureSessionHandle session, const PacketHea
         && (payloadHeader.IsInitiator() != IsInitiator());
 }
 
+void ExchangeContext::OnConnectionExpired()
+{
+    if (!IsResponseExpected())
+    {
+        // Nothing to do in this case
+        return;
+    }
+
+    // If we're waiting on a response, we now know it's never going to show up
+    // and we should notify our delegate accordingly.
+    CancelResponseTimer();
+    SetResponseExpected(false);
+    NotifyResponseTimeout();
+}
+
 CHIP_ERROR ExchangeContext::StartResponseTimer()
 {
     System::Layer * lSystemLayer = mExchangeMgr->GetSessionMgr()->SystemLayer();
@@ -334,11 +349,18 @@ void ExchangeContext::HandleResponseTimeout(System::Layer * aSystemLayer, void *
     // NOTE: we don't set mResponseExpected to false here because the response could still arrive. If the user
     // wants to never receive the response, they must close the exchange context.
 
-    ExchangeDelegate * delegate = ec->GetDelegate();
+    ec->NotifyResponseTimeout();
+}
+
+void ExchangeContext::NotifyResponseTimeout()
+{
+    ExchangeDelegate * delegate = GetDelegate();
 
     // Call the user's timeout handler.
     if (delegate != nullptr)
-        delegate->OnResponseTimeout(ec);
+    {
+        delegate->OnResponseTimeout(this);
+    }
 }
 
 CHIP_ERROR ExchangeContext::HandleMessage(const PacketHeader & packetHeader, const PayloadHeader & payloadHeader,

src/messaging/ExchangeContext.h

@@ -209,6 +209,17 @@ class DLL_EXPORT ExchangeContext : public ReliableMessageContext, public Referen
      */
     bool MatchExchange(SecureSessionHandle session, const PacketHeader & packetHeader, const PayloadHeader & payloadHeader);
 
+    /**
+     * Notify the exchange that its connection has expired.
+     */
+    void OnConnectionExpired();
+
+    /**
+     * Notify our delegate, if any, that we have timed out waiting for a
+     * response.
+     */
+    void NotifyResponseTimeout();
+
     CHIP_ERROR StartResponseTimer();
 
     void CancelResponseTimer();

src/messaging/ExchangeMgr.cpp

@@ -314,8 +314,9 @@ void ExchangeManager::OnConnectionExpired(SecureSessionHandle session, SecureSes
     mContextPool.ForEachActiveObject([&](auto * ec) {
         if (ec->mSecureSession == session)
         {
-            ec->Close();
-            // Continue iterate because there can be multiple contexts associated with the connection.
+            ec->OnConnectionExpired();
+            // Continue to iterate because there can be multiple exchanges
+            // associated with the connection.
         }
         return true;
     });

src/messaging/ExchangeMgr.h

@@ -249,9 +249,13 @@ class DLL_EXPORT ExchangeManager : public SecureSessionMgrDelegate, public Trans
                            SecureSessionMgr * msgLayer) override;
 
     void OnNewConnection(SecureSessionHandle session, SecureSessionMgr * mgr) override;
+#if CHIP_CONFIG_TEST
+public: // Allow OnConnectionExpired to be called directly from tests.
+#endif  // CHIP_CONFIG_TEST
     void OnConnectionExpired(SecureSessionHandle session, SecureSessionMgr * mgr) override;
 
     // TransportMgrDelegate interface for rendezvous sessions
+private:
     void OnMessageReceived(const Transport::PeerAddress & source, System::PacketBufferHandle && msgBuf) override;
 };
 

src/messaging/tests/TestExchangeMgr.cpp

@@ -89,6 +89,22 @@ class MockAppDelegate : public ExchangeDelegate
     bool IsOnMessageReceivedCalled = false;
 };
 
+class WaitForTimeoutDelegate : public ExchangeDelegate
+{
+public:
+    void OnMessageReceived(ExchangeContext * ec, const PacketHeader & packetHeader, const PayloadHeader & payloadHeader,
+                           System::PacketBufferHandle && buffer) override
+    {}
+
+    void OnResponseTimeout(ExchangeContext * ec) override
+    {
+        IsOnResponseTimeoutCalled = true;
+        ec->Close();
+    }
+
+    bool IsOnResponseTimeoutCalled = false;
+};
+
 void CheckNewContextTest(nlTestSuite * inSuite, void * inContext)
 {
     TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
@@ -114,6 +130,47 @@ void CheckNewContextTest(nlTestSuite * inSuite, void * inContext)
     ec2->Close();
 }
 
+void CheckSessionExpirationBasics(nlTestSuite * inSuite, void * inContext)
+{
+    TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
+
+    MockAppDelegate sendDelegate;
+    ExchangeContext * ec1 = ctx.NewExchangeToLocal(&sendDelegate);
+
+    // Expire the session this exchange is supposedly on.
+    ctx.GetExchangeManager().OnConnectionExpired(ec1->GetSecureSession(), &ctx.GetSecureSessionManager());
+
+    MockAppDelegate receiveDelegate;
+    CHIP_ERROR err =
+        ctx.GetExchangeManager().RegisterUnsolicitedMessageHandlerForType(Protocols::BDX::Id, kMsgType_TEST1, &receiveDelegate);
+    NL_TEST_ASSERT(inSuite, err == CHIP_NO_ERROR);
+
+    ec1->SendMessage(Protocols::BDX::Id, kMsgType_TEST1, System::PacketBufferHandle::New(System::PacketBuffer::kMaxSize),
+                     SendFlags(Messaging::SendMessageFlags::kNoAutoRequestAck));
+    NL_TEST_ASSERT(inSuite, receiveDelegate.IsOnMessageReceivedCalled);
+
+    ec1->Close();
+    err = ctx.GetExchangeManager().UnregisterUnsolicitedMessageHandlerForType(Protocols::BDX::Id, kMsgType_TEST1);
+    NL_TEST_ASSERT(inSuite, err == CHIP_NO_ERROR);
+}
+
+void CheckSessionExpirationTimeout(nlTestSuite * inSuite, void * inContext)
+{
+    TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
+
+    WaitForTimeoutDelegate sendDelegate;
+    ExchangeContext * ec1 = ctx.NewExchangeToLocal(&sendDelegate);
+
+    ec1->SendMessage(Protocols::BDX::Id, kMsgType_TEST1, System::PacketBufferHandle::New(System::PacketBuffer::kMaxSize),
+                     SendFlags(Messaging::SendMessageFlags::kExpectResponse).Set(Messaging::SendMessageFlags::kNoAutoRequestAck));
+    NL_TEST_ASSERT(inSuite, !sendDelegate.IsOnResponseTimeoutCalled);
+
+    // Expire the session this exchange is supposedly on.  This should close the
+    // exchange.
+    ctx.GetExchangeManager().OnConnectionExpired(ec1->GetSecureSession(), &ctx.GetSecureSessionManager());
+    NL_TEST_ASSERT(inSuite, sendDelegate.IsOnResponseTimeoutCalled);
+}
+
 void CheckUmhRegistrationTest(nlTestSuite * inSuite, void * inContext)
 {
     TestContext & ctx = *reinterpret_cast<TestContext *>(inContext);
@@ -167,6 +224,8 @@ void CheckExchangeMessages(nlTestSuite * inSuite, void * inContext)
     NL_TEST_ASSERT(inSuite, mockUnsolicitedAppDelegate.IsOnMessageReceivedCalled);
 
     ec1->Close();
+    err = ctx.GetExchangeManager().UnregisterUnsolicitedMessageHandlerForType(Protocols::BDX::Id, kMsgType_TEST1);
+    NL_TEST_ASSERT(inSuite, err == CHIP_NO_ERROR);
 }
 
 // Test Suite
@@ -180,6 +239,8 @@ const nlTest sTests[] =
     NL_TEST_DEF("Test ExchangeMgr::NewContext",               CheckNewContextTest),
     NL_TEST_DEF("Test ExchangeMgr::CheckUmhRegistrationTest", CheckUmhRegistrationTest),
     NL_TEST_DEF("Test ExchangeMgr::CheckExchangeMessages",    CheckExchangeMessages),
+    NL_TEST_DEF("Test OnConnectionExpired basics",            CheckSessionExpirationBasics),
+    NL_TEST_DEF("Test OnConnectionExpired timeout handling",  CheckSessionExpirationTimeout),
 
     NL_TEST_SENTINEL()
 };