CommandSender and CommandHandler reserves space for Finalize to succe…

…ed (#31077) * CommandSender and CommandHandler reserves space for Finalize to succeed * Address PR comments * Fix CI * Address PR comments * Restyled by clang-format * Update src/app/CommandHandler.cpp Co-authored-by: Boris Zbarsky <bzbarsky@apple.com> --------- Co-authored-by: Restyled.io <commits@restyled.io> Co-authored-by: Boris Zbarsky <bzbarsky@apple.com>
project-chip · Jan 3, 2024 · d2de5d5 · d2de5d5
1 parent c7f8ec9
commit d2de5d5
Show file tree

Hide file tree

Showing 11 changed files with 269 additions and 10 deletions.
diff --git a/src/app/CommandHandler.cpp b/src/app/CommandHandler.cpp
@@ -62,13 +62,14 @@ CHIP_ERROR CommandHandler::AllocateBuffer()
         VerifyOrReturnError(!commandPacket.IsNull(), CHIP_ERROR_NO_MEMORY);
 
         mCommandMessageWriter.Init(std::move(commandPacket));
-        ReturnErrorOnFailure(mInvokeResponseBuilder.Init(&mCommandMessageWriter));
+        ReturnErrorOnFailure(mInvokeResponseBuilder.InitWithEndBufferReserved(&mCommandMessageWriter));
 
         mInvokeResponseBuilder.SuppressResponse(mSuppressResponse);
         ReturnErrorOnFailure(mInvokeResponseBuilder.GetError());
 
-        mInvokeResponseBuilder.CreateInvokeResponses();
+        mInvokeResponseBuilder.CreateInvokeResponses(/* aReserveEndBuffer = */ true);
         ReturnErrorOnFailure(mInvokeResponseBuilder.GetError());
+
         mBufferAllocated = true;
     }
 

diff --git a/src/app/CommandSender.cpp b/src/app/CommandSender.cpp
@@ -83,12 +83,12 @@ CHIP_ERROR CommandSender::AllocateBuffer()
         VerifyOrReturnError(!commandPacket.IsNull(), CHIP_ERROR_NO_MEMORY);
 
         mCommandMessageWriter.Init(std::move(commandPacket));
-        ReturnErrorOnFailure(mInvokeRequestBuilder.Init(&mCommandMessageWriter));
+        ReturnErrorOnFailure(mInvokeRequestBuilder.InitWithEndBufferReserved(&mCommandMessageWriter));
 
         mInvokeRequestBuilder.SuppressResponse(mSuppressResponse).TimedRequest(mTimedRequest);
         ReturnErrorOnFailure(mInvokeRequestBuilder.GetError());
 
-        mInvokeRequestBuilder.CreateInvokeRequests();
+        mInvokeRequestBuilder.CreateInvokeRequests(/* aReserveEndBuffer = */ true);
         ReturnErrorOnFailure(mInvokeRequestBuilder.GetError());
 
         mBufferAllocated = true;

diff --git a/src/app/MessageDef/InvokeRequestMessage.cpp b/src/app/MessageDef/InvokeRequestMessage.cpp
@@ -113,6 +113,14 @@ CHIP_ERROR InvokeRequestMessage::Parser::GetInvokeRequests(InvokeRequests::Parse
     return apInvokeRequests->Init(reader);
 }
 
+CHIP_ERROR InvokeRequestMessage::Builder::InitWithEndBufferReserved(TLV::TLVWriter * const apWriter)
+{
+    ReturnErrorOnFailure(Init(apWriter));
+    ReturnErrorOnFailure(GetWriter()->ReserveBuffer(GetSizeToEndInvokeRequestMessage()));
+    mIsEndBufferReserved = true;
+    return CHIP_NO_ERROR;
+}
+
 InvokeRequestMessage::Builder & InvokeRequestMessage::Builder::SuppressResponse(const bool aSuppressResponse)
 {
     if (mError == CHIP_NO_ERROR)
@@ -131,17 +139,33 @@ InvokeRequestMessage::Builder & InvokeRequestMessage::Builder::TimedRequest(cons
     return *this;
 }
 
-InvokeRequests::Builder & InvokeRequestMessage::Builder::CreateInvokeRequests()
+InvokeRequests::Builder & InvokeRequestMessage::Builder::CreateInvokeRequests(const bool aReserveEndBuffer)
 {
     if (mError == CHIP_NO_ERROR)
     {
-        mError = mInvokeRequests.Init(mpWriter, to_underlying(Tag::kInvokeRequests));
+        if (aReserveEndBuffer)
+        {
+            mError = mInvokeRequests.InitWithEndBufferReserved(mpWriter, to_underlying(Tag::kInvokeRequests));
+        }
+        else
+        {
+            mError = mInvokeRequests.Init(mpWriter, to_underlying(Tag::kInvokeRequests));
+        }
     }
     return mInvokeRequests;
 }
 
 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.
+    ReturnErrorOnFailure(mError);
+    if (mIsEndBufferReserved)
+    {
+        ReturnErrorOnFailure(GetWriter()->UnreserveBuffer(GetSizeToEndInvokeRequestMessage()));
+        mIsEndBufferReserved = false;
+    }
     if (mError == CHIP_NO_ERROR)
     {
         mError = MessageBuilder::EncodeInteractionModelRevision();
@@ -152,5 +176,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
diff --git a/src/app/MessageDef/InvokeRequestMessage.h b/src/app/MessageDef/InvokeRequestMessage.h
@@ -75,6 +75,12 @@ class Parser : public MessageParser
 class Builder : public MessageBuilder
 {
 public:
+    /**
+     *  @brief Performs underlying StructBuilder::Init, but reserves memory need in
+     *  EndOfInvokeRequestMessage() with underlying TLVWriter.
+     */
+    CHIP_ERROR InitWithEndBufferReserved(TLV::TLVWriter * const apWriter);
+
     /**
      *  @brief when sets to true, it means do not send a response to this action
      */
@@ -90,7 +96,7 @@ class Builder : public MessageBuilder
      *
      *  @return A reference to InvokeRequests::Builder
      */
-    InvokeRequests::Builder & CreateInvokeRequests();
+    InvokeRequests::Builder & CreateInvokeRequests(const bool aReserveEndBuffer = false);
 
     /**
      *  @brief Get reference to InvokeRequests::Builder
@@ -106,8 +112,16 @@ 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;
+    bool mIsEndBufferReserved = false;
 };
 } // namespace InvokeRequestMessage
 } // namespace app

diff --git a/src/app/MessageDef/InvokeRequests.cpp b/src/app/MessageDef/InvokeRequests.cpp
@@ -70,6 +70,14 @@ CHIP_ERROR InvokeRequests::Parser::PrettyPrint() const
 }
 #endif // CHIP_CONFIG_IM_PRETTY_PRINT
 
+CHIP_ERROR InvokeRequests::Builder::InitWithEndBufferReserved(TLV::TLVWriter * const apWriter, const uint8_t aContextTagToUse)
+{
+    ReturnErrorOnFailure(Init(apWriter, aContextTagToUse));
+    ReturnErrorOnFailure(GetWriter()->ReserveBuffer(GetSizeToEndInvokeRequests()));
+    mIsEndBufferReserved = true;
+    return CHIP_NO_ERROR;
+}
+
 CommandDataIB::Builder & InvokeRequests::Builder::CreateCommandData()
 {
     if (mError == CHIP_NO_ERROR)
@@ -81,8 +89,22 @@ 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.
+    if (mIsEndBufferReserved)
+    {
+        ReturnErrorOnFailure(GetWriter()->UnreserveBuffer(GetSizeToEndInvokeRequests()));
+        mIsEndBufferReserved = false;
+    }
     EndOfContainer();
     return GetError();
 }
+
+uint32_t InvokeRequests::Builder::GetSizeToEndInvokeRequests()
+{
+    uint32_t kEndOfContainerSize = 1;
+    return kEndOfContainerSize;
+}
 } // namespace app
 } // namespace chip
diff --git a/src/app/MessageDef/InvokeRequests.h b/src/app/MessageDef/InvokeRequests.h
@@ -42,6 +42,12 @@ class Parser : public ArrayParser
 class Builder : public ArrayBuilder
 {
 public:
+    /**
+     *  @brief Performs underlying StructBuilder::Init, but reserves memory need in
+     *  EndOfInvokeRequests() with underlying TLVWriter.
+     */
+    CHIP_ERROR InitWithEndBufferReserved(TLV::TLVWriter * const apWriter, const uint8_t aContextTagToUse);
+
     /**
      *  @brief Initialize a CommandDataIB::Builder for writing into the TLV stream
      *
@@ -61,8 +67,16 @@ 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;
+    bool mIsEndBufferReserved = false;
 };
 } // namespace InvokeRequests
 } // namespace app

diff --git a/src/app/MessageDef/InvokeResponseIBs.cpp b/src/app/MessageDef/InvokeResponseIBs.cpp
@@ -70,6 +70,14 @@ CHIP_ERROR InvokeResponseIBs::Parser::PrettyPrint() const
 }
 #endif // CHIP_CONFIG_IM_PRETTY_PRINT
 
+CHIP_ERROR InvokeResponseIBs::Builder::InitWithEndBufferReserved(TLV::TLVWriter * const apWriter, const uint8_t aContextTagToUse)
+{
+    ReturnErrorOnFailure(Init(apWriter, aContextTagToUse));
+    ReturnErrorOnFailure(GetWriter()->ReserveBuffer(GetSizeToEndInvokeResponses()));
+    mIsEndBufferReserved = true;
+    return CHIP_NO_ERROR;
+}
+
 InvokeResponseIB::Builder & InvokeResponseIBs::Builder::CreateInvokeResponse()
 {
     if (mError == CHIP_NO_ERROR)
@@ -81,8 +89,22 @@ 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.
+    if (mIsEndBufferReserved)
+    {
+        ReturnErrorOnFailure(GetWriter()->UnreserveBuffer(GetSizeToEndInvokeResponses()));
+        mIsEndBufferReserved = false;
+    }
     EndOfContainer();
     return GetError();
 }
+
+uint32_t InvokeResponseIBs::Builder::GetSizeToEndInvokeResponses()
+{
+    uint32_t kEndOfContainerSize = 1;
+    return kEndOfContainerSize;
+}
 } // namespace app
 } // namespace chip
diff --git a/src/app/MessageDef/InvokeResponseIBs.h b/src/app/MessageDef/InvokeResponseIBs.h
@@ -42,6 +42,12 @@ class Parser : public ArrayParser
 class Builder : public ArrayBuilder
 {
 public:
+    /**
+     *  @brief Performs underlying StructBuilder::Init, but reserves memory need in
+     *  EndOfInvokeResponses() with underlying TLVWriter.
+     */
+    CHIP_ERROR InitWithEndBufferReserved(TLV::TLVWriter * const apWriter, const uint8_t aContextTagToUse);
+
     /**
      *  @brief Initialize a InvokeResponseIB::Builder for writing into the TLV stream
      *
@@ -61,8 +67,16 @@ 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;
+    bool mIsEndBufferReserved = false;
 };
 } // namespace InvokeResponseIBs
 } // namespace app

diff --git a/src/app/MessageDef/InvokeResponseMessage.cpp b/src/app/MessageDef/InvokeResponseMessage.cpp
@@ -112,6 +112,14 @@ CHIP_ERROR InvokeResponseMessage::Parser::GetMoreChunkedMessages(bool * const ap
     return GetSimpleValue(to_underlying(Tag::kMoreChunkedMessages), TLV::kTLVType_Boolean, apMoreChunkedMessages);
 }
 
+CHIP_ERROR InvokeResponseMessage::Builder::InitWithEndBufferReserved(TLV::TLVWriter * const apWriter)
+{
+    ReturnErrorOnFailure(Init(apWriter));
+    ReturnErrorOnFailure(GetWriter()->ReserveBuffer(GetSizeToEndInvokeResponseMessage()));
+    mIsEndBufferReserved = true;
+    return CHIP_NO_ERROR;
+}
+
 InvokeResponseMessage::Builder & InvokeResponseMessage::Builder::SuppressResponse(const bool aSuppressResponse)
 {
     if (mError == CHIP_NO_ERROR)
@@ -121,11 +129,18 @@ InvokeResponseMessage::Builder & InvokeResponseMessage::Builder::SuppressRespons
     return *this;
 }
 
-InvokeResponseIBs::Builder & InvokeResponseMessage::Builder::CreateInvokeResponses()
+InvokeResponseIBs::Builder & InvokeResponseMessage::Builder::CreateInvokeResponses(const bool aReserveEndBuffer)
 {
     if (mError == CHIP_NO_ERROR)
     {
-        mError = mInvokeResponses.Init(mpWriter, to_underlying(Tag::kInvokeResponses));
+        if (aReserveEndBuffer)
+        {
+            mError = mInvokeResponses.InitWithEndBufferReserved(mpWriter, to_underlying(Tag::kInvokeResponses));
+        }
+        else
+        {
+            mError = mInvokeResponses.Init(mpWriter, to_underlying(Tag::kInvokeResponses));
+        }
     }
     return mInvokeResponses;
 }
@@ -142,6 +157,15 @@ 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.
+    ReturnErrorOnFailure(mError);
+    if (mIsEndBufferReserved)
+    {
+        ReturnErrorOnFailure(GetWriter()->UnreserveBuffer(GetSizeToEndInvokeResponseMessage()));
+        mIsEndBufferReserved = false;
+    }
     if (mError == CHIP_NO_ERROR)
     {
         mError = MessageBuilder::EncodeInteractionModelRevision();
@@ -152,5 +176,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
diff --git a/src/app/MessageDef/InvokeResponseMessage.h b/src/app/MessageDef/InvokeResponseMessage.h
@@ -77,6 +77,12 @@ class Parser : public MessageParser
 class Builder : public MessageBuilder
 {
 public:
+    /**
+     *  @brief Performs underlying StructBuilder::Init, but reserves memory need in
+     *  EndOfInvokeResponseMessage() with underlying TLVWriter.
+     */
+    CHIP_ERROR InitWithEndBufferReserved(TLV::TLVWriter * const apWriter);
+
     /**
      *  @brief This is set to 'true' by the subscriber to indicate preservation of previous subscriptions. If omitted, it implies
      * 'false' as a value.
@@ -88,7 +94,7 @@ class Builder : public MessageBuilder
      *
      *  @return A reference to InvokeResponseIBs::Builder
      */
-    InvokeResponseIBs::Builder & CreateInvokeResponses();
+    InvokeResponseIBs::Builder & CreateInvokeResponses(const bool aReserveEndBuffer = false);
 
     /**
      *  @brief Get reference to InvokeResponseIBs::Builder
@@ -111,8 +117,16 @@ 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;
+    bool mIsEndBufferReserved = false;
 };
 } // namespace InvokeResponseMessage
 } // namespace app