Fix CRMP resend null out retained buffer

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->PrepareMessage(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 & message) const
 {
-    return mSessionMgr->SendEncryptedMessage(session, std::move(message), retainedMessage);
+    return mSessionMgr->SendPreparedMessage(session, message);
 }
 
 bool ApplicationExchangeDispatch::MessagePermitted(uint16_t protocol, uint8_t type)

src/messaging/ApplicationExchangeDispatch.h

@@ -45,15 +45,13 @@ class ApplicationExchangeDispatch : public ExchangeMessageDispatch
         return ExchangeMessageDispatch::Init(reliableMessageMgr);
     }
 
-    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

@@ -75,24 +75,25 @@ CHIP_ERROR ExchangeMessageDispatch::SendMessage(SecureSessionHandle session, uin
         // Add to Table for subsequent sending
         ReturnErrorOnFailure(mReliableMessageMgr->AddToRetransTable(reliableMessageContext, &entry));
 
-        CHIP_ERROR err = SendMessageImpl(session, payloadHeader, std::move(message), &entry->retainedBuf);
+        CHIP_ERROR err = PrepareMessage(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));
             mReliableMessageMgr->ClearRetransTable(*entry);
-            ReturnErrorOnFailure(err);
-        }
-        else
-        {
-            mReliableMessageMgr->StartRetransmision(entry);
+            return err;
         }
+
+        ReturnErrorOnFailure(SendPreparedMessage(session, entry->retainedBuf));
+        mReliableMessageMgr->StartRetransmision(entry);
     }
     else
     {
         // If the channel itself is providing reliability, let's not request CRMP 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

@@ -48,28 +48,17 @@ class ExchangeMessageDispatch : public ReferenceCounted<ExchangeMessageDispatch>
                            ReliableMessageContext * reliableMessageContext, bool isReliableTransmission, Protocols::Id protocol,
                            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.
-     */
-    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 & message) 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; }
 
 private:
     ReliableMessageMgr * mReliableMessageMgr = nullptr;

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

@@ -113,37 +113,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 & message) 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(), message.Retain());
     }
 
+    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 & message) 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, message.Retain());
 }
 
 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(reliableMessageMgr);
     }
 
-    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 & message) 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);

src/transport/SecureSessionMgr.cpp

@@ -122,44 +122,49 @@ Transport::Type SecureSessionMgr::GetTransportType(NodeId peerNodeId)
     return Transport::Type::kUndefined;
 }
 
-CHIP_ERROR SecureSessionMgr::SendMessage(SecureSessionHandle session, PayloadHeader & payloadHeader,
-                                         System::PacketBufferHandle && msgBuf, EncryptedPacketBufferHandle * bufferRetainSlot)
+CHIP_ERROR SecureSessionMgr::PrepareMessage(SecureSessionHandle session, PayloadHeader & payloadHeader,
+                                            System::PacketBufferHandle && msgBuf, EncryptedPacketBufferHandle & preparedMessage)
 {
-    PacketHeader unusedPacketHeader;
-    return SendMessage(session, payloadHeader, unusedPacketHeader, std::move(msgBuf), bufferRetainSlot,
-                       EncryptionState::kPayloadIsUnencrypted);
-}
+    PacketHeader packetHeader;
+    if (IsControlMessage(payloadHeader))
+    {
+        packetHeader.SetSecureSessionControlMsg(true);
+    }
 
-CHIP_ERROR SecureSessionMgr::SendEncryptedMessage(SecureSessionHandle session, EncryptedPacketBufferHandle && msgBuf,
-                                                  EncryptedPacketBufferHandle * bufferRetainSlot)
-{
-    VerifyOrReturnError(!msgBuf.IsNull(), CHIP_ERROR_INVALID_ARGUMENT);
-    VerifyOrReturnError(!msgBuf.HasChainedBuffer(), CHIP_ERROR_INVALID_MESSAGE_LENGTH);
+    PeerConnectionState * state = GetPeerConnectionState(session);
+    if (state == nullptr)
+    {
+        return CHIP_ERROR_NOT_CONNECTED;
+    };
 
-    // Our send path needs a (writable) PacketBuffer (e.g. so it can encode a
-    // PacketHeader into it), so get that from the EncryptedPacketBufferHandle.
-    System::PacketBufferHandle mutableBuf(std::move(msgBuf).CastToWritable());
+    Transport::AdminPairingInfo * admin = mAdmins->FindAdminWithId(state->GetAdminId());
+    if (admin == nullptr)
+    {
+        return CHIP_ERROR_INCORRECT_STATE;
+    }
 
-    // Advancing the start to encrypted header, since SendMessage will attach the packet header on top of it.
-    PacketHeader packetHeader;
-    ReturnErrorOnFailure(packetHeader.DecodeAndConsume(mutableBuf));
+    NodeId localNodeId       = admin->GetNodeId();
+    MessageCounter & counter = GetSendCounterForPacket(payloadHeader, *state);
+    ReturnErrorOnFailure(SecureMessageCodec::Encode(localNodeId, state, payloadHeader, packetHeader, msgBuf, counter));
 
-    PayloadHeader payloadHeader;
-    return SendMessage(session, payloadHeader, packetHeader, std::move(mutableBuf), bufferRetainSlot,
-                       EncryptionState::kPayloadIsEncrypted);
+    ReturnErrorOnFailure(packetHeader.EncodeBeforeData(msgBuf));
+
+    preparedMessage = EncryptedPacketBufferHandle::MarkEncrypted(std::move(msgBuf));
+    ChipLogProgress(Inet, "Prepared msg %p from 0x" ChipLogFormatX64 " to 0x" ChipLogFormatX64 ".", &preparedMessage,
+                    ChipLogValueX64(localNodeId), ChipLogValueX64(state->GetPeerNodeId()));
+
+    return CHIP_NO_ERROR;
 }
 
-CHIP_ERROR SecureSessionMgr::SendMessage(SecureSessionHandle session, PayloadHeader & payloadHeader, PacketHeader & packetHeader,
-                                         System::PacketBufferHandle && msgBuf, EncryptedPacketBufferHandle * bufferRetainSlot,
-                                         EncryptionState encryptionState)
+CHIP_ERROR SecureSessionMgr::SendPreparedMessage(SecureSessionHandle session, const EncryptedPacketBufferHandle & preparedMessage)
 {
     CHIP_ERROR err              = CHIP_NO_ERROR;
     PeerConnectionState * state = nullptr;
-    NodeId localNodeId          = mLocalNodeId;
-
-    Transport::AdminPairingInfo * admin = nullptr;
+    PacketBufferHandle msgBuf;
 
     VerifyOrExit(mState == State::kInitialized, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(!preparedMessage.IsNull(), err = CHIP_ERROR_INVALID_ARGUMENT);
+    msgBuf = preparedMessage.Retain();
     VerifyOrExit(!msgBuf.IsNull(), err = CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrExit(!msgBuf->HasChainedBuffer(), err = CHIP_ERROR_INVALID_MESSAGE_LENGTH);
 
@@ -169,33 +174,9 @@ CHIP_ERROR SecureSessionMgr::SendMessage(SecureSessionHandle session, PayloadHea
 
     // This marks any connection where we send data to as 'active'
     mPeerConnections.MarkConnectionActive(state);
-    admin = mAdmins->FindAdminWithId(state->GetAdminId());
-    VerifyOrExit(admin != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
-    localNodeId = admin->GetNodeId();
-
-    if (IsControlMessage(payloadHeader))
-    {
-        packetHeader.SetSecureSessionControlMsg(true);
-    }
-
-    if (encryptionState == EncryptionState::kPayloadIsUnencrypted)
-    {
-        MessageCounter & counter = GetSendCounterForPacket(payloadHeader, *state);
-        err                      = SecureMessageCodec::Encode(localNodeId, state, payloadHeader, packetHeader, msgBuf, counter);
-        SuccessOrExit(err);
-    }
-
-    err = packetHeader.EncodeBeforeData(msgBuf);
-    SuccessOrExit(err);
-
-    // Retain the packet buffer in case it's needed for retransmissions.
-    if (bufferRetainSlot != nullptr)
-    {
-        *bufferRetainSlot = EncryptedPacketBufferHandle::MarkEncrypted(msgBuf.Retain());
-    }
 
-    ChipLogProgress(Inet, "Sending msg from 0x" ChipLogFormatX64 " to 0x" ChipLogFormatX64 " at utc time: %" PRId64 " msec",
-                    ChipLogValueX64(localNodeId), ChipLogValueX64(state->GetPeerNodeId()), System::Layer::GetClock_MonotonicMS());
+    ChipLogProgress(Inet, "Sending msg %p to 0x" ChipLogFormatX64 " at utc time: %" PRId64 " msec", &preparedMessage,
+                    ChipLogValueX64(state->GetPeerNodeId()), System::Layer::GetClock_MonotonicMS());
 
     if (state->GetTransport() != nullptr)
     {