Make exchanges handle closing on message receipt more automatically.

The idea is that the cases that want to keep the exchange open should
do so explicitly and in all other cases the exchange should close.
This helps avoid exchange leaks and makes it much clearer when an
exchange is being kept open.

examples/shell/shell_common/cmd_send.cpp

@@ -110,7 +110,6 @@ class MockAppDelegate : public Messaging::ExchangeDelegate
         streamer_printf(sout, "Response received: len=%u time=%.3fms\n", buffer->DataLength(),
                         static_cast<double>(transitTime) / 1000);
 
-        gExchangeCtx->Close();
         gExchangeCtx = nullptr;
         return CHIP_NO_ERROR;
     }

src/app/CommandHandler.cpp

@@ -71,6 +71,8 @@ CHIP_ERROR CommandHandler::SendCommandResponse()
     MoveToState(CommandState::Sending);
 
 exit:
+    // Keep Shutdown() from double-closing our exchange.
+    mpExchangeCtx = nullptr;
     Shutdown();
     ChipLogFunctError(err);
     return err;

src/app/CommandSender.cpp

@@ -92,9 +92,8 @@ CHIP_ERROR CommandSender::OnMessageReceived(Messaging::ExchangeContext * apExcha
 exit:
     ChipLogFunctError(err);
 
-    // Close the exchange cleanly so that the ExchangeManager will send an ack for the message we just received.
-    // This needs to be done before the Reset() call, because Reset() aborts mpExchangeCtx if its not null.
-    mpExchangeCtx->Close();
+    // Null out mpExchangeCtx, so our Shutdown() call below won't try to abort
+    // it and fail to send an ack for the message we just received.
     mpExchangeCtx = nullptr;
 
     if (mpDelegate != nullptr)

src/app/InteractionModelEngine.cpp

@@ -190,7 +190,6 @@ CHIP_ERROR InteractionModelEngine::OnUnknownMsgType(Messaging::ExchangeContext *
     // err = SendStatusReport(ec, kChipProfile_Common, kStatus_UnsupportedMessage);
     // SuccessOrExit(err);
 
-    apExchangeContext->Close();
     apExchangeContext = nullptr;
 
     ChipLogFunctError(err);

src/app/ReadClient.cpp

@@ -217,8 +217,8 @@ CHIP_ERROR ReadClient::OnMessageReceived(Messaging::ExchangeContext * apExchange
 exit:
     ChipLogFunctError(err);
 
-    // Close the exchange cleanly so that the ExchangeManager will send an ack for the message we just received.
-    mpExchangeCtx->Close();
+    // Null out mpExchangeCtx, so our Shutdown() call below won't try to abort
+    // it and fail to send an ack for the message we just received.
     mpExchangeCtx = nullptr;
     MoveToState(ClientState::Initialized);
 

src/app/ReadHandler.cpp

@@ -82,6 +82,8 @@ CHIP_ERROR ReadHandler::OnReadRequest(Messaging::ExchangeContext * apExchangeCon
     if (err != CHIP_NO_ERROR)
     {
         ChipLogFunctError(err);
+        // Keep Shutdown() from double-closing our exchange.
+        mpExchangeCtx = nullptr;
         Shutdown();
     }
 
@@ -153,6 +155,14 @@ CHIP_ERROR ReadHandler::ProcessReadRequest(System::PacketBufferHandle && aPayloa
     MoveToState(HandlerState::Reportable);
 
     err = InteractionModelEngine::GetInstance()->GetReportingEngine().ScheduleRun();
+    SuccessOrExit(err);
+
+    // mpExchangeCtx can be null here due to
+    // https://github.com/project-chip/connectedhomeip/issues/8031
+    if (mpExchangeCtx)
+    {
+        mpExchangeCtx->WillSendMessage();
+    }
 
 exit:
     ChipLogFunctError(err);

src/app/WriteClient.cpp

@@ -294,20 +294,16 @@ CHIP_ERROR WriteClient::OnMessageReceived(Messaging::ExchangeContext * apExchang
 
     VerifyOrDie(apExchangeContext == mpExchangeCtx);
 
+    // We are done with this exchange, and it will be closing itself.
+    mpExchangeCtx = nullptr;
+
     // Verify that the message is an Invoke Command Response.
     // If not, close the exchange and free the payload.
     if (!aPayloadHeader.HasMessageType(Protocols::InteractionModel::MsgType::WriteResponse))
     {
-        apExchangeContext->Close();
-        mpExchangeCtx = nullptr;
         ExitNow();
     }
 
-    // Close the current exchange after receiving the response since the response message marks the
-    // end of conversation represented by the exchange. We should create an new exchange for a new
-    // conversation defined in Interaction Model protocol.
-    ClearExistingExchangeContext();
-
     err = ProcessWriteResponseMessage(std::move(aPayload));
 
 exit:

src/app/WriteHandler.cpp

@@ -74,6 +74,8 @@ CHIP_ERROR WriteHandler::OnWriteRequest(Messaging::ExchangeContext * apExchangeC
 
 exit:
     ChipLogFunctError(err);
+    // Keep Shutdown() from double-closing our exchange.
+    mpExchangeCtx = nullptr;
     Shutdown();
     return err;
 }

src/app/server/Server.cpp

@@ -387,7 +387,6 @@ class ServerCallback : public ExchangeDelegate
         }
 
     exit:
-        exchangeContext->Close();
         return err;
     }
 

src/controller/CHIPDevice.cpp

@@ -338,7 +338,6 @@ CHIP_ERROR Device::OnMessageReceived(Messaging::ExchangeContext * exchange, cons
             HandleDataModelMessage(exchange, std::move(msgBuf));
         }
     }
-    exchange->Close();
     return CHIP_NO_ERROR;
 }
 

src/controller/CHIPDeviceController.cpp

@@ -574,7 +574,6 @@ CHIP_ERROR DeviceController::OnMessageReceived(Messaging::ExchangeContext * ec,
                                                const PayloadHeader & payloadHeader, System::PacketBufferHandle && msgBuf)
 {
     uint16_t index;
-    bool needClose = true;
 
     VerifyOrExit(mState == State::Initialized, ChipLogError(Controller, "OnMessageReceived was called in incorrect state"));
 
@@ -584,14 +583,9 @@ CHIP_ERROR DeviceController::OnMessageReceived(Messaging::ExchangeContext * ec,
     index = FindDeviceIndex(packetHeader.GetSourceNodeId().Value());
     VerifyOrExit(index < kNumMaxActiveDevices, ChipLogError(Controller, "OnMessageReceived was called for unknown device object"));
 
-    needClose = false; // Device will handle it
     mActiveDevices[index].OnMessageReceived(ec, packetHeader, payloadHeader, std::move(msgBuf));
 
 exit:
-    if (needClose)
-    {
-        ec->Close();
-    }
     return CHIP_NO_ERROR;
 }
 

src/messaging/ExchangeContext.cpp

@@ -79,6 +79,9 @@ void ExchangeContext::SetResponseTimeout(Timeout timeout)
 CHIP_ERROR ExchangeContext::SendMessage(Protocols::Id protocolId, uint8_t msgType, PacketBufferHandle && msgBuf,
                                         const SendFlags & sendFlags)
 {
+    // If we were waiting for a message send, this is it.
+    mFlags.Clear(Flags::kFlagWillSendMessage);
+
     CHIP_ERROR err                         = CHIP_NO_ERROR;
     Transport::PeerConnectionState * state = nullptr;
 
@@ -378,18 +381,26 @@ CHIP_ERROR ExchangeContext::HandleMessage(const PacketHeader & packetHeader, con
     // layer has completed its work on the ExchangeContext.
     Retain();
 
+    // Keep track of whether we're nested under an outer HandleMessage
+    // invocation.
+    bool alreadyHandlingMessage = mFlags.Has(Flags::kFlagHandlingMessage);
+    mFlags.Set(Flags::kFlagHandlingMessage);
+
+    bool isStandaloneAck = payloadHeader.HasMessageType(Protocols::SecureChannel::MsgType::StandaloneAck);
+    bool isDuplicate     = msgFlags.Has(MessageFlagValues::kDuplicateMessage);
+
     CHIP_ERROR err = mDispatch->OnMessageReceived(payloadHeader, packetHeader.GetMessageId(), peerAddress, msgFlags,
                                                   GetReliableMessageContext());
     SuccessOrExit(err);
 
     // The SecureChannel::StandaloneAck message type is only used for MRP; do not pass such messages to the application layer.
-    if (payloadHeader.HasMessageType(Protocols::SecureChannel::MsgType::StandaloneAck))
+    if (isStandaloneAck)
     {
         ExitNow(err = CHIP_NO_ERROR);
     }
 
     // Since the message is duplicate, let's not forward it up the stack
-    if (msgFlags.Has(MessageFlagValues::kDuplicateMessage))
+    if (isDuplicate)
     {
         ExitNow(err = CHIP_NO_ERROR);
     }
@@ -412,6 +423,29 @@ CHIP_ERROR ExchangeContext::HandleMessage(const PacketHeader & packetHeader, con
     }
 
 exit:
+    // Don't close ourselves if this message is a standalone ack. Either we're
+    // still not closed and getting an ack should not affect that, or we're
+    // already closed and then closing again would be a bad idea.
+    //
+    // Don't close for duplicates for similar reasons, except when mDelegate is
+    // null: that case means we created the exchange just to send an ack for
+    // this duplicate, and we don't need to exist anymore.
+    //
+    // Also don't close if there's an outer HandleMessage invocation.  It'll
+    // deal with the closing.
+    if (!mFlags.Has(Flags::kFlagWillSendMessage) && !IsResponseExpected() && !isStandaloneAck &&
+        (!isDuplicate || (mDelegate == nullptr)) && !alreadyHandlingMessage)
+    {
+        Close();
+    }
+
+    if (!alreadyHandlingMessage)
+    {
+        // We are the outermost HandleMessage invocation.  We're not handling a
+        // message anymore.
+        mFlags.Clear(Flags::kFlagHandlingMessage);
+    }
+
     // Release the reference to the ExchangeContext that was held at the beginning of this function.
     // This call should also do the needful of closing the ExchangeContext if the protocol has
     // already made a prior call to Close().

src/messaging/ExchangeContext.h

@@ -109,6 +109,12 @@ class DLL_EXPORT ExchangeContext : public ReliableMessageContext, public Referen
                            std::move(msgPayload), sendFlags);
     }
 
+    /**
+     * A notification that we will have SendMessage called on us in the future
+     * (and should stay open until that happens).
+     */
+    void WillSendMessage() { mFlags.Set(Flags::kFlagWillSendMessage); }
+
     /**
      *  Handle a received CHIP message on this exchange.
      *

src/messaging/ExchangeDelegate.h

@@ -49,7 +49,18 @@ class DLL_EXPORT ExchangeDelegate
 
     /**
      * @brief
-     *   This function is the protocol callback for handling a received CHIP message.
+     *   This function is the protocol callback for handling a received CHIP
+     *   message.
+     *
+     *   After calling this method an exchange will close itself unless one of
+     *   two things happens:
+     *
+     *   1) A call to SendMessage on the exchange with the kExpectResponse flag
+     *      set.
+     *   2) A call to WillSendMessage on the exchange.
+     *
+     *   Consumers that don't do one of those things MUST NOT retain a pointer
+     *   to the exchange.
      *
      *  @param[in]    ec            A pointer to the ExchangeContext object.
      *  @param[in]    packetHeader  A reference to the PacketHeader object.

src/messaging/ExchangeMgr.cpp

@@ -201,7 +201,6 @@ void ExchangeManager::OnMessageReceived(const PacketHeader & packetHeader, const
 {
     CHIP_ERROR err                          = CHIP_NO_ERROR;
     UnsolicitedMessageHandler * matchingUMH = nullptr;
-    bool sendAckAndCloseExchange            = false;
 
     ChipLogProgress(ExchangeManager, "Received message of type %d and protocolId %" PRIu32 " on exchange %d",
                     payloadHeader.GetMessageType(), payloadHeader.GetProtocolID().ToFullyQualifiedSpecForm(),
@@ -269,36 +268,23 @@ void ExchangeManager::OnMessageReceived(const PacketHeader & packetHeader, const
         ExitNow(err = CHIP_ERROR_UNSOLICITED_MSG_NO_ORIGINATOR);
     }
 
-    // If we didn't find an existing exchange that matches the message, and no unsolicited message handler registered
-    // to hand this message, we need to create a temporary exchange to send an ack for this message and then close this exchange.
-    sendAckAndCloseExchange = payloadHeader.NeedsAck() && (matchingUMH == nullptr);
-
-    // If we found a handler or we need to create a new exchange context (EC).
-    if (matchingUMH != nullptr || sendAckAndCloseExchange)
+    // If we found a handler or we need to send an ack, create an exchange to
+    // handle the message.
+    if (matchingUMH != nullptr || payloadHeader.NeedsAck())
     {
-        ExchangeContext * ec = nullptr;
-
-        if (sendAckAndCloseExchange)
-        {
-            // If rcvd msg is from initiator then this exchange is created as not Initiator.
-            // If rcvd msg is not from initiator then this exchange is created as Initiator.
-            // TODO: Figure out which channel to use for the received message
-            ec = mContextPool.CreateObject(this, payloadHeader.GetExchangeID(), session, !payloadHeader.IsInitiator(), nullptr);
-        }
-        else
-        {
-            ec = mContextPool.CreateObject(this, payloadHeader.GetExchangeID(), session, false, matchingUMH->Delegate);
-        }
+        ExchangeDelegate * delegate = matchingUMH ? matchingUMH->Delegate : nullptr;
+        // If rcvd msg is from initiator then this exchange is created as not Initiator.
+        // If rcvd msg is not from initiator then this exchange is created as Initiator.
+        // Note that if matchingUMH is not null then rcvd msg if from initiator.
+        // TODO: Figure out which channel to use for the received message
+        ExchangeContext * ec =
+            mContextPool.CreateObject(this, payloadHeader.GetExchangeID(), session, !payloadHeader.IsInitiator(), delegate);
 
         VerifyOrExit(ec != nullptr, err = CHIP_ERROR_NO_MEMORY);
 
         ChipLogDetail(ExchangeManager, "ec id: %d, Delegate: 0x%p", ec->GetExchangeId(), ec->GetDelegate());
 
         ec->HandleMessage(packetHeader, payloadHeader, source, msgFlags, std::move(msgBuf));
-
-        // Close exchange if it was created only to send ack for a duplicate message.
-        if (sendAckAndCloseExchange)
-            ec->Close();
     }
 
 exit:

src/messaging/ReliableMessageContext.h

@@ -202,6 +202,12 @@ class ReliableMessageContext
 
         /// When set, signifies that at least one message has been received from peer on this exchange context.
         kFlagMsgRcvdFromPeer = 0x0040,
+
+        /// When set, signifies that this exchange is waiting for a call to SendMessage.
+        kFlagWillSendMessage = 0x0080,
+
+        /// When set, signifies that we are currently in the middle of HandleMessage.
+        kFlagHandlingMessage = 0x0100,
     };
 
     BitFlags<Flags> mFlags; // Internal state flags

src/messaging/tests/TestExchangeMgr.cpp

@@ -81,7 +81,6 @@ class MockAppDelegate : public ExchangeDelegate
                                  System::PacketBufferHandle && buffer) override
     {
         IsOnMessageReceivedCalled = true;
-        ec->Close();
         return CHIP_NO_ERROR;
     }
 

src/messaging/tests/TestReliableMessageProtocol.cpp

@@ -115,9 +115,12 @@ class MockAppDelegate : public ExchangeDelegate
 
         if (!mRetainExchange)
         {
-            ec->Close();
             ec = nullptr;
         }
+        else
+        {
+            ec->WillSendMessage();
+        }
         mExchange = ec;
 
         if (mTestSuite != nullptr)
@@ -181,7 +184,6 @@ class MockSessionEstablishmentDelegate : public ExchangeDelegate
                                  System::PacketBufferHandle && buffer) override
     {
         IsOnMessageReceivedCalled = true;
-        ec->Close();
         if (mTestSuite != nullptr)
         {
             NL_TEST_ASSERT(mTestSuite, buffer->TotalLength() == sizeof(PAYLOAD));

src/protocols/echo/EchoClient.cpp

@@ -95,22 +95,14 @@ CHIP_ERROR EchoClient::OnMessageReceived(Messaging::ExchangeContext * ec, const
     // which clears the OnMessageReceived callback.
     VerifyOrDie(ec == mExchangeCtx);
 
+    mExchangeCtx = nullptr;
+
     // Verify that the message is an Echo Response.
-    // If not, close the exchange and free the payload.
     if (!payloadHeader.HasMessageType(MsgType::EchoResponse))
     {
-        ec->Close();
-        mExchangeCtx = nullptr;
         return CHIP_ERROR_INVALID_ARGUMENT;
     }
 
-    // Remove the EC from the app state now. OnEchoResponseReceived can call
-    // SendEchoRequest and install a new one. We abort rather than close
-    // because we no longer care whether the echo request message has been
-    // acknowledged at the transport layer.
-    mExchangeCtx->Abort();
-    mExchangeCtx = nullptr;
-
     // Call the registered OnEchoResponseReceived handler, if any.
     if (OnEchoResponseReceived != nullptr)
     {