CommandSender and CommandHandler reserves space for Finalize to succeed

src/app/CommandHandler.cpp

@@ -69,6 +69,10 @@ CHIP_ERROR CommandHandler::AllocateBuffer()
 
         mInvokeResponseBuilder.CreateInvokeResponses();
         ReturnErrorOnFailure(mInvokeResponseBuilder.GetError());
+
+        auto * tlvWriter = mInvokeResponseBuilder.GetWriter();
+        ReturnErrorOnFailure(tlvWriter->ReserveBuffer(mInvokeResponseBuilder.GetInvokeResponses().GetSizeToEndInvokeResponses()));
+        ReturnErrorOnFailure(tlvWriter->ReserveBuffer(mInvokeResponseBuilder.GetSizeToEndInvokeResponseMessage()));
         mBufferAllocated = true;
     }
 
@@ -756,6 +760,10 @@ CommandHandler::Handle::Handle(CommandHandler * handle)
 CHIP_ERROR CommandHandler::Finalize(System::PacketBufferHandle & commandPacket)
 {
     VerifyOrReturnError(mState == State::AddedCommand, CHIP_ERROR_INCORRECT_STATE);
+
+    auto * tlvWriter = mInvokeResponseBuilder.GetWriter();
+    ReturnErrorOnFailure(tlvWriter->UnreserveBuffer(mInvokeResponseBuilder.GetInvokeResponses().GetSizeToEndInvokeResponses()));
+    ReturnErrorOnFailure(tlvWriter->UnreserveBuffer(mInvokeResponseBuilder.GetSizeToEndInvokeResponseMessage()));
     ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().EndOfInvokeResponses());
     ReturnErrorOnFailure(mInvokeResponseBuilder.EndOfInvokeResponseMessage());
     return mCommandMessageWriter.Finalize(&commandPacket);

src/app/CommandSender.cpp

@@ -91,6 +91,10 @@ CHIP_ERROR CommandSender::AllocateBuffer()
         mInvokeRequestBuilder.CreateInvokeRequests();
         ReturnErrorOnFailure(mInvokeRequestBuilder.GetError());
 
+        auto * tlvWriter = mInvokeRequestBuilder.GetWriter();
+        ReturnErrorOnFailure(tlvWriter->ReserveBuffer(mInvokeRequestBuilder.GetInvokeRequests().GetSizeToEndInvokeRequests()));
+        ReturnErrorOnFailure(tlvWriter->ReserveBuffer(mInvokeRequestBuilder.GetSizeToEndInvokeRequestMessage()));
+
         mBufferAllocated = true;
     }
 
@@ -508,6 +512,11 @@ CHIP_ERROR CommandSender::FinishCommand(const Optional<uint16_t> & aTimedInvokeT
 CHIP_ERROR CommandSender::Finalize(System::PacketBufferHandle & commandPacket)
 {
     VerifyOrReturnError(mState == State::AddedCommand, CHIP_ERROR_INCORRECT_STATE);
+
+    auto * tlvWriter = mInvokeRequestBuilder.GetWriter();
+    ReturnErrorOnFailure(tlvWriter->UnreserveBuffer(mInvokeRequestBuilder.GetInvokeRequests().GetSizeToEndInvokeRequests()));
+    ReturnErrorOnFailure(tlvWriter->UnreserveBuffer(mInvokeRequestBuilder.GetSizeToEndInvokeRequestMessage()));
+
     ReturnErrorOnFailure(mInvokeRequestBuilder.GetInvokeRequests().EndOfInvokeRequests());
     ReturnErrorOnFailure(mInvokeRequestBuilder.EndOfInvokeRequestMessage());
     return mCommandMessageWriter.Finalize(&commandPacket);

src/app/MessageDef/InvokeRequestMessage.cpp

@@ -142,6 +142,9 @@ InvokeRequests::Builder & InvokeRequestMessage::Builder::CreateInvokeRequests()
 
 CHIP_ERROR InvokeRequestMessage::Builder::EndOfInvokeRequestMessage()
 {
+    // If any changes are made to how we end the invoke request message that involves how many
+    // bytes are needed, a corresponding change to GetSizeToEndInvokeRequestMessage indicating
+    // the new size that will be required.
     if (mError == CHIP_NO_ERROR)
     {
         mError = MessageBuilder::EncodeInteractionModelRevision();
@@ -152,5 +155,15 @@ CHIP_ERROR InvokeRequestMessage::Builder::EndOfInvokeRequestMessage()
     }
     return GetError();
 }
+
+uint32_t InvokeRequestMessage::Builder::GetSizeToEndInvokeRequestMessage()
+{
+    // EncodeInteractionModelRevision() encodes a uint8_t with context tag 0xFF. This means 1 control byte,
+    // 1 byte for the tag, 1 byte for the value.
+    uint32_t kEncodeInteractionModelSize = 1 + 1 + 1;
+    uint32_t kEndOfContainerSize         = 1;
+
+    return kEncodeInteractionModelSize + kEndOfContainerSize;
+}
 }; // namespace app
 }; // namespace chip

src/app/MessageDef/InvokeRequestMessage.h

@@ -106,6 +106,13 @@ class Builder : public MessageBuilder
      */
     CHIP_ERROR EndOfInvokeRequestMessage();
 
+    /**
+     *  @brief Get number of bytes required in the buffer by EndOfInvokeRequestMessage()
+     *
+     *  @return Expected number of bytes required in the buffer by EndOfInvokeRequestMessage()
+     */
+    uint32_t GetSizeToEndInvokeRequestMessage();
+
 private:
     InvokeRequests::Builder mInvokeRequests;
 };

src/app/MessageDef/InvokeRequests.cpp

@@ -81,8 +81,17 @@ CommandDataIB::Builder & InvokeRequests::Builder::CreateCommandData()
 
 CHIP_ERROR InvokeRequests::Builder::EndOfInvokeRequests()
 {
+    // If any changes are made to how we end the invoke requests that involves how many bytes are
+    // needed, a corresponding change to GetSizeToEndInvokeRequests indicating the new size that
+    // will be required.
     EndOfContainer();
     return GetError();
 }
+
+uint32_t InvokeRequests::Builder::GetSizeToEndInvokeRequests()
+{
+    uint32_t kEndOfContainerSize = 1;
+    return kEndOfContainerSize;
+}
 } // namespace app
 } // namespace chip

src/app/MessageDef/InvokeRequests.h

@@ -61,6 +61,13 @@ class Builder : public ArrayBuilder
      */
     CHIP_ERROR EndOfInvokeRequests();
 
+    /**
+     *  @brief Get number of bytes required in the buffer by EndOfInvokeRequests()
+     *
+     *  @return Expected number of bytes required in the buffer by EndOfInvokeRequests()
+     */
+    uint32_t GetSizeToEndInvokeRequests();
+
 private:
     CommandDataIB::Builder mCommandData;
 };

src/app/MessageDef/InvokeResponseIBs.cpp

@@ -81,8 +81,17 @@ InvokeResponseIB::Builder & InvokeResponseIBs::Builder::CreateInvokeResponse()
 
 CHIP_ERROR InvokeResponseIBs::Builder::EndOfInvokeResponses()
 {
+    // If any changes are made to how we end the invoke responses that involves how many bytes are
+    // needed, a corresponding change to GetSizeToEndInvokeResponses indicating the new size that
+    // will be required.
     EndOfContainer();
     return GetError();
 }
+
+uint32_t InvokeResponseIBs::Builder::GetSizeToEndInvokeResponses()
+{
+    uint32_t kEndOfContainerSize = 1;
+    return kEndOfContainerSize;
+}
 } // namespace app
 } // namespace chip

src/app/MessageDef/InvokeResponseIBs.h

@@ -61,6 +61,13 @@ class Builder : public ArrayBuilder
      */
     CHIP_ERROR EndOfInvokeResponses();
 
+    /**
+     *  @brief Get number of bytes required in the buffer by EndOfInvokeResponses()
+     *
+     *  @return Expected number of bytes required in the buffer by EndOfInvokeResponses()
+     */
+    uint32_t GetSizeToEndInvokeResponses();
+
 private:
     InvokeResponseIB::Builder mInvokeResponse;
 };

src/app/MessageDef/InvokeResponseMessage.cpp

@@ -142,6 +142,9 @@ InvokeResponseMessage::Builder & InvokeResponseMessage::Builder::MoreChunkedMess
 
 CHIP_ERROR InvokeResponseMessage::Builder::EndOfInvokeResponseMessage()
 {
+    // If any changes are made to how we end the invoke response message that involves how many
+    // bytes are needed, a corresponding change to GetSizeToEndInvokeResponseMessage indicating
+    // the new size that will be required.
     if (mError == CHIP_NO_ERROR)
     {
         mError = MessageBuilder::EncodeInteractionModelRevision();
@@ -152,5 +155,15 @@ CHIP_ERROR InvokeResponseMessage::Builder::EndOfInvokeResponseMessage()
     }
     return GetError();
 }
+
+uint32_t InvokeResponseMessage::Builder::GetSizeToEndInvokeResponseMessage()
+{
+    // EncodeInteractionModelRevision() encodes a uint8_t with context tag 0xFF. This means 1 control byte,
+    // 1 byte for the tag, 1 byte for the value.
+    uint32_t kEncodeInteractionModelSize = 1 + 1 + 1;
+    uint32_t kEndOfContainerSize         = 1;
+
+    return kEncodeInteractionModelSize + kEndOfContainerSize;
+}
 } // namespace app
 } // namespace chip

src/app/MessageDef/InvokeResponseMessage.h

@@ -111,6 +111,13 @@ class Builder : public MessageBuilder
      */
     CHIP_ERROR EndOfInvokeResponseMessage();
 
+    /**
+     *  @brief Get number of bytes required in the buffer by EndOfInvokeResponseMessage()
+     *
+     *  @return Expected number of bytes required in the buffer by EndOfInvokeResponseMessage()
+     */
+    uint32_t GetSizeToEndInvokeResponseMessage();
+
 private:
     InvokeResponseIBs::Builder mInvokeResponses;
 };

src/app/tests/TestMessageDef.cpp

@@ -2057,6 +2057,54 @@ void InvokeInvokeRequestMessageTest(nlTestSuite * apSuite, void * apContext)
     ParseInvokeRequestMessage(apSuite, reader);
 }
 
+void InvokeRequestMessageEndOfMessageReservationTest(nlTestSuite * apSuite, void * apContext)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+    chip::System::PacketBufferTLVWriter writer;
+    InvokeRequestMessage::Builder invokeRequestMessageBuilder;
+    const uint32_t kSmallBufferSize = 100;
+    writer.Init(chip::System::PacketBufferHandle::New(kSmallBufferSize, /* aReservedSize = */ 0), /* useChainedBuffers = */ false);
+    err = invokeRequestMessageBuilder.Init(&writer);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t sizeToReserve = invokeRequestMessageBuilder.GetSizeToEndInvokeRequestMessage();
+    err                    = writer.ReserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    uint32_t remianingLengthAfterReservation = writer.GetRemainingFreeLength();
+
+    err = writer.UnreserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    err = invokeRequestMessageBuilder.EndOfInvokeRequestMessage();
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t remainingLengthAfterEndingInvokeRequestMessage = writer.GetRemainingFreeLength();
+    NL_TEST_ASSERT(apSuite, remianingLengthAfterReservation == remainingLengthAfterEndingInvokeRequestMessage);
+}
+
+void InvokeRequestsEndOfRequestReservationTest(nlTestSuite * apSuite, void * apContext)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+    chip::System::PacketBufferTLVWriter writer;
+    InvokeRequests::Builder invokeRequestsBuilder;
+    const uint32_t kSmallBufferSize = 100;
+    writer.Init(chip::System::PacketBufferHandle::New(kSmallBufferSize, /* aReservedSize = */ 0), /* useChainedBuffers = */ false);
+    err = invokeRequestsBuilder.Init(&writer);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t sizeToReserve = invokeRequestsBuilder.GetSizeToEndInvokeRequests();
+    err                    = writer.ReserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    uint32_t remianingLengthAfterReservation = writer.GetRemainingFreeLength();
+
+    err = writer.UnreserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    err = invokeRequestsBuilder.EndOfInvokeRequests();
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t remianingLengthAfterEndingInvokeRequests = writer.GetRemainingFreeLength();
+    NL_TEST_ASSERT(apSuite, remianingLengthAfterReservation == remianingLengthAfterEndingInvokeRequests);
+}
+
 void InvokeInvokeResponseMessageTest(nlTestSuite * apSuite, void * apContext)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
@@ -2074,6 +2122,54 @@ void InvokeInvokeResponseMessageTest(nlTestSuite * apSuite, void * apContext)
     ParseInvokeResponseMessage(apSuite, reader);
 }
 
+void InvokeResponseMessageEndOfMessageReservationTest(nlTestSuite * apSuite, void * apContext)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+    chip::System::PacketBufferTLVWriter writer;
+    InvokeResponseMessage::Builder invokeResponseMessageBuilder;
+    const uint32_t kSmallBufferSize = 100;
+    writer.Init(chip::System::PacketBufferHandle::New(kSmallBufferSize, /* aReservedSize = */ 0), /* useChainedBuffers = */ false);
+    err = invokeResponseMessageBuilder.Init(&writer);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t sizeToReserve = invokeResponseMessageBuilder.GetSizeToEndInvokeResponseMessage();
+    err                    = writer.ReserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    uint32_t remianingLengthAfterReservation = writer.GetRemainingFreeLength();
+
+    err = writer.UnreserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    err = invokeResponseMessageBuilder.EndOfInvokeResponseMessage();
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t remianingLengthAfterEndingInvokeResponseMessage = writer.GetRemainingFreeLength();
+    NL_TEST_ASSERT(apSuite, remianingLengthAfterReservation == remianingLengthAfterEndingInvokeResponseMessage);
+}
+
+void InvokeResponsesEndOfResponseReservationTest(nlTestSuite * apSuite, void * apContext)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+    chip::System::PacketBufferTLVWriter writer;
+    InvokeResponseIBs::Builder invokeResponsesBuilder;
+    const uint32_t kSmallBufferSize = 100;
+    writer.Init(chip::System::PacketBufferHandle::New(kSmallBufferSize, /* aReservedSize = */ 0), /* useChainedBuffers = */ false);
+    err = invokeResponsesBuilder.Init(&writer);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t sizeToReserve = invokeResponsesBuilder.GetSizeToEndInvokeResponses();
+    err                    = writer.ReserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    uint32_t remianingLengthAfterReservation = writer.GetRemainingFreeLength();
+
+    err = writer.UnreserveBuffer(sizeToReserve);
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+    err = invokeResponsesBuilder.EndOfInvokeResponses();
+    NL_TEST_ASSERT(apSuite, err == CHIP_NO_ERROR);
+
+    uint32_t remianingLengthAfterEndingInvokeResponses = writer.GetRemainingFreeLength();
+    NL_TEST_ASSERT(apSuite, remianingLengthAfterReservation == remianingLengthAfterEndingInvokeResponses);
+}
+
 void ReportDataMessageTest(nlTestSuite * apSuite, void * apContext)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
@@ -2283,7 +2379,11 @@ const nlTest sTests[] =
                 NL_TEST_DEF("InvokeRequestsTest", InvokeRequestsTest),
                 NL_TEST_DEF("InvokeResponsesTest", InvokeResponsesTest),
                 NL_TEST_DEF("InvokeInvokeRequestMessageTest", InvokeInvokeRequestMessageTest),
+                NL_TEST_DEF("InvokeRequestMessageEndOfMessageReservationTest", InvokeRequestMessageEndOfMessageReservationTest),
+                NL_TEST_DEF("InvokeRequestsEndOfRequestReservationTest", InvokeRequestsEndOfRequestReservationTest),
                 NL_TEST_DEF("InvokeInvokeResponseMessageTest", InvokeInvokeResponseMessageTest),
+                NL_TEST_DEF("InvokeResponseMessageEndOfMessageReservationTest", InvokeResponseMessageEndOfMessageReservationTest),
+                NL_TEST_DEF("InvokeResponsesEndOfResponseReservationTest", InvokeResponsesEndOfResponseReservationTest),
                 NL_TEST_DEF("ReportDataMessageTest", ReportDataMessageTest),
                 NL_TEST_DEF("ReadRequestMessageTest", ReadRequestMessageTest),
                 NL_TEST_DEF("WriteRequestMessageTest", WriteRequestMessageTest),