Fix CRMP resend null out retained buffer (#7312)

* Fix CRMP resend null out retained buffer

* Resolve comments

* Apply suggestions from code review

Co-authored-by: Boris Zbarsky <bzbarsky@apple.com>

Co-authored-by: Boris Zbarsky <bzbarsky@apple.com>

src/messaging/ApplicationExchangeDispatch.cpp

@@ -26,17 +26,17 @@
 namespace chip {
 namespace Messaging {
 
-CHIP_ERROR ApplicationExchangeDispatch::SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader,
-                                                        System::PacketBufferHandle && message,
-                                                        EncryptedPacketBufferHandle * retainedMessage)
+CHIP_ERROR ApplicationExchangeDispatch::PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader,
+                                                       System::PacketBufferHandle && message,
+                                                       EncryptedPacketBufferHandle & preparedMessage)
 {
-    return mSessionMgr->SendMessage(session, payloadHeader, std::move(message), retainedMessage);
+    return mSessionMgr->BuildEncryptedMessagePayload(session, payloadHeader, std::move(message), preparedMessage);
 }
 
-CHIP_ERROR ApplicationExchangeDispatch::ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                                                      EncryptedPacketBufferHandle * retainedMessage) const
+CHIP_ERROR ApplicationExchangeDispatch::SendPreparedMessage(SecureSessionHandle session,
+                                                            const EncryptedPacketBufferHandle & preparedMessage) const
 {
-    return mSessionMgr->SendEncryptedMessage(session, std::move(message), retainedMessage);
+    return mSessionMgr->SendPreparedMessage(session, preparedMessage);
 }
 
 bool ApplicationExchangeDispatch::MessagePermitted(uint16_t protocol, uint8_t type)

src/messaging/ApplicationExchangeDispatch.h

@@ -45,15 +45,13 @@ class ApplicationExchangeDispatch : public ExchangeMessageDispatch
         return ExchangeMessageDispatch::Init();
     }
 
-    CHIP_ERROR ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                             EncryptedPacketBufferHandle * retainedMessage) const override;
+    CHIP_ERROR PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
+                              EncryptedPacketBufferHandle & preparedMessage) override;
+    CHIP_ERROR SendPreparedMessage(SecureSessionHandle session, const EncryptedPacketBufferHandle & message) const override;
 
     SecureSessionMgr * GetSessionMgr() const { return mSessionMgr; }
 
 protected:
-    CHIP_ERROR SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
-                               EncryptedPacketBufferHandle * retainedMessage) override;
-
     bool MessagePermitted(uint16_t protocol, uint8_t type) override;
 
 private:

src/messaging/ExchangeMessageDispatch.cpp

@@ -29,6 +29,7 @@
 #endif
 
 #include <inttypes.h>
+#include <memory>
 
 #include <messaging/ExchangeMessageDispatch.h>
 #include <messaging/ReliableMessageContext.h>
@@ -76,25 +77,22 @@ CHIP_ERROR ExchangeMessageDispatch::SendMessage(SecureSessionHandle session, uin
 
         // Add to Table for subsequent sending
         ReturnErrorOnFailure(reliableMessageMgr->AddToRetransTable(reliableMessageContext, &entry));
-
-        CHIP_ERROR err = SendMessageImpl(session, payloadHeader, std::move(message), &entry->retainedBuf);
-        if (err != CHIP_NO_ERROR)
-        {
-            // Remove from table
-            ChipLogError(ExchangeManager, "Failed to send message with err %s", ::chip::ErrorStr(err));
-            reliableMessageMgr->ClearRetransTable(*entry);
-            ReturnErrorOnFailure(err);
-        }
-        else
-        {
-            reliableMessageMgr->StartRetransmision(entry);
-        }
+        auto deleter = [reliableMessageMgr](ReliableMessageMgr::RetransTableEntry * e) {
+            reliableMessageMgr->ClearRetransTable(*e);
+        };
+        std::unique_ptr<ReliableMessageMgr::RetransTableEntry, decltype(deleter)> entryOwner(entry, deleter);
+
+        ReturnErrorOnFailure(PrepareMessage(session, payloadHeader, std::move(message), entryOwner->retainedBuf));
+        ReturnErrorOnFailure(SendPreparedMessage(session, entryOwner->retainedBuf));
+        reliableMessageMgr->StartRetransmision(entryOwner.release());
     }
     else
     {
         // If the channel itself is providing reliability, let's not request MRP acks
         payloadHeader.SetNeedsAck(false);
-        ReturnErrorOnFailure(SendMessageImpl(session, payloadHeader, std::move(message), nullptr));
+        EncryptedPacketBufferHandle preparedMessage;
+        ReturnErrorOnFailure(PrepareMessage(session, payloadHeader, std::move(message), preparedMessage));
+        ReturnErrorOnFailure(SendPreparedMessage(session, preparedMessage));
     }
 
     return CHIP_NO_ERROR;

src/messaging/ExchangeMessageDispatch.h

@@ -45,27 +45,26 @@ class ExchangeMessageDispatch : public ReferenceCounted<ExchangeMessageDispatch>
                            uint8_t type, System::PacketBufferHandle && message);
 
     /**
-     * The 'message' and 'retainedMessage' arguments may point to the same
-     * handle.  Therefore, callees _must_ ensure that any moving out of
-     * 'message' happens before writing to *retainedMessage.
+     * @brief
+     *   This interface takes the payload and returns the prepared message which can be send multiple times.
+     *
+     * @param session         Peer node to which the payload to be sent
+     * @param payloadHeader   The payloadHeader to be encoded into the packet
+     * @param message         The payload to be sent
+     * @param preparedMessage The handle to hold the prepared message
      */
-    virtual CHIP_ERROR ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                                     EncryptedPacketBufferHandle * retainedMessage) const
-    {
-        return CHIP_ERROR_NOT_IMPLEMENTED;
-    }
+    virtual CHIP_ERROR PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader,
+                                      System::PacketBufferHandle && message, EncryptedPacketBufferHandle & preparedMessage) = 0;
+    virtual CHIP_ERROR SendPreparedMessage(SecureSessionHandle session,
+                                           const EncryptedPacketBufferHandle & preparedMessage) const                       = 0;
 
     virtual CHIP_ERROR OnMessageReceived(const PayloadHeader & payloadHeader, uint32_t messageId,
                                          const Transport::PeerAddress & peerAddress,
                                          ReliableMessageContext * reliableMessageContext);
 
 protected:
     virtual bool MessagePermitted(uint16_t protocol, uint8_t type) = 0;
-
-    virtual CHIP_ERROR SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader,
-                                       System::PacketBufferHandle && message, EncryptedPacketBufferHandle * retainedMessage) = 0;
-
-    virtual bool IsReliableTransmissionAllowed() { return true; }
+    virtual bool IsReliableTransmissionAllowed() const { return true; }
 };
 
 } // namespace Messaging

src/messaging/ReliableMessageMgr.cpp

@@ -366,8 +366,7 @@ CHIP_ERROR ReliableMessageMgr::SendFromRetransTable(RetransTableEntry * entry)
     const ExchangeMessageDispatch * dispatcher = rc->GetExchangeContext()->GetMessageDispatch();
     VerifyOrExit(dispatcher != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
 
-    err =
-        dispatcher->ResendMessage(rc->GetExchangeContext()->GetSecureSession(), std::move(entry->retainedBuf), &entry->retainedBuf);
+    err = dispatcher->SendPreparedMessage(rc->GetExchangeContext()->GetSecureSession(), entry->retainedBuf);
     SuccessOrExit(err);
 
     // Update the counters

src/messaging/tests/TestReliableMessageProtocol.cpp

@@ -117,37 +117,25 @@ class MockAppDelegate : public ExchangeDelegate
 class MockSessionEstablishmentExchangeDispatch : public Messaging::ExchangeMessageDispatch
 {
 public:
-    CHIP_ERROR SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
-                               EncryptedPacketBufferHandle * retainedMessage) override
+    CHIP_ERROR PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
+                              EncryptedPacketBufferHandle & preparedMessage) override
     {
         PacketHeader packetHeader;
 
         ReturnErrorOnFailure(payloadHeader.EncodeBeforeData(message));
         ReturnErrorOnFailure(packetHeader.EncodeBeforeData(message));
 
-        if (retainedMessage != nullptr && mRetainMessageOnSend)
-        {
-            *retainedMessage = EncryptedPacketBufferHandle::MarkEncrypted(message.Retain());
-        }
-        return gTransportMgr.SendMessage(Transport::PeerAddress(), std::move(message));
+        preparedMessage = EncryptedPacketBufferHandle::MarkEncrypted(std::move(message));
+        return CHIP_NO_ERROR;
     }
 
-    CHIP_ERROR ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                             EncryptedPacketBufferHandle * retainedMessage) const override
+    CHIP_ERROR SendPreparedMessage(SecureSessionHandle session, const EncryptedPacketBufferHandle & preparedMessage) const override
     {
-        // Our send path needs a (writable) PacketBuffer, so get that from the
-        // EncryptedPacketBufferHandle.  Note that we have to do this before we
-        // set *retainedMessage, because 'message' and '*retainedMessage' might
-        // be the same memory location and we have to guarantee that we move out
-        // of 'message' before we write to *retainedMessage.
-        System::PacketBufferHandle writableBuf(std::move(message).CastToWritable());
-        if (retainedMessage != nullptr && mRetainMessageOnSend)
-        {
-            *retainedMessage = EncryptedPacketBufferHandle::MarkEncrypted(writableBuf.Retain());
-        }
-        return gTransportMgr.SendMessage(Transport::PeerAddress(), std::move(writableBuf));
+        return gTransportMgr.SendMessage(Transport::PeerAddress(), preparedMessage.CastToWritable());
     }
 
+    bool IsReliableTransmissionAllowed() const override { return mRetainMessageOnSend; }
+
     bool MessagePermitted(uint16_t protocol, uint8_t type) override { return true; }
 
     bool mRetainMessageOnSend = true;

src/protocols/secure_channel/SessionEstablishmentExchangeDispatch.cpp

@@ -28,40 +28,23 @@ namespace chip {
 
 using namespace Messaging;
 
-CHIP_ERROR SessionEstablishmentExchangeDispatch::SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader,
-                                                                 System::PacketBufferHandle && message,
-                                                                 EncryptedPacketBufferHandle * retainedMessage)
+CHIP_ERROR SessionEstablishmentExchangeDispatch::PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader,
+                                                                System::PacketBufferHandle && message,
+                                                                EncryptedPacketBufferHandle & preparedMessage)
 {
-    ReturnErrorCodeIf(mTransportMgr == nullptr, CHIP_ERROR_INCORRECT_STATE);
     PacketHeader packetHeader;
-
     ReturnErrorOnFailure(payloadHeader.EncodeBeforeData(message));
     ReturnErrorOnFailure(packetHeader.EncodeBeforeData(message));
 
-    if (retainedMessage != nullptr)
-    {
-        *retainedMessage = EncryptedPacketBufferHandle::MarkEncrypted(message.Retain());
-        ChipLogError(Inet, "RETAINED IN SESS: %p %d", retainedMessage, (*retainedMessage).IsNull());
-    }
-    return mTransportMgr->SendMessage(mPeerAddress, std::move(message));
+    preparedMessage = EncryptedPacketBufferHandle::MarkEncrypted(std::move(message));
+    return CHIP_NO_ERROR;
 }
 
-CHIP_ERROR SessionEstablishmentExchangeDispatch::ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                                                               EncryptedPacketBufferHandle * retainedMessage) const
+CHIP_ERROR SessionEstablishmentExchangeDispatch::SendPreparedMessage(SecureSessionHandle session,
+                                                                     const EncryptedPacketBufferHandle & preparedMessage) const
 {
     ReturnErrorCodeIf(mTransportMgr == nullptr, CHIP_ERROR_INCORRECT_STATE);
-
-    // Our send path needs a (writable) PacketBuffer, so get that from the
-    // EncryptedPacketBufferHandle.  Note that we have to do this before we set
-    // *retainedMessage, because 'message' and '*retainedMessage' might be the
-    // same memory location and we have to guarantee that we move out of
-    // 'message' before we write to *retainedMessage.
-    System::PacketBufferHandle writableBuf(std::move(message).CastToWritable());
-    if (retainedMessage != nullptr)
-    {
-        *retainedMessage = EncryptedPacketBufferHandle::MarkEncrypted(writableBuf.Retain());
-    }
-    return mTransportMgr->SendMessage(mPeerAddress, std::move(writableBuf));
+    return mTransportMgr->SendMessage(mPeerAddress, preparedMessage.CastToWritable());
 }
 
 CHIP_ERROR SessionEstablishmentExchangeDispatch::OnMessageReceived(const PayloadHeader & payloadHeader, uint32_t messageId,

src/protocols/secure_channel/SessionEstablishmentExchangeDispatch.h

@@ -43,8 +43,9 @@ class SessionEstablishmentExchangeDispatch : public Messaging::ExchangeMessageDi
         return ExchangeMessageDispatch::Init();
     }
 
-    CHIP_ERROR ResendMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && message,
-                             EncryptedPacketBufferHandle * retainedMessage) const override;
+    CHIP_ERROR PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
+                              EncryptedPacketBufferHandle & out) override;
+    CHIP_ERROR SendPreparedMessage(SecureSessionHandle session, const EncryptedPacketBufferHandle & preparedMessage) const override;
 
     CHIP_ERROR OnMessageReceived(const PayloadHeader & payloadHeader, uint32_t messageId,
                                  const Transport::PeerAddress & peerAddress,
@@ -55,12 +56,9 @@ class SessionEstablishmentExchangeDispatch : public Messaging::ExchangeMessageDi
     void SetPeerAddress(const Transport::PeerAddress & address) { mPeerAddress = address; }
 
 protected:
-    CHIP_ERROR SendMessageImpl(SecureSessionHandle session, PayloadHeader & payloadHeader, System::PacketBufferHandle && message,
-                               EncryptedPacketBufferHandle * retainedMessage) override;
-
     bool MessagePermitted(uint16_t protocol, uint8_t type) override;
 
-    bool IsReliableTransmissionAllowed() override
+    bool IsReliableTransmissionAllowed() const override
     {
         // If the underlying transport is UDP.
         return (mPeerAddress.GetTransportType() == Transport::Type::kUdp);