From da77cb0c457e2c64feecf46f316890e6ed35e41b Mon Sep 17 00:00:00 2001
From: Song Guo <songguo@google.com>
Date: Mon, 7 Mar 2022 16:07:23 +0800
Subject: [PATCH] [IM] Make CommandHandler::AddResponseData atomic

---
 src/app/CommandHandler.cpp | 11 +++++++
 src/app/CommandHandler.h   | 67 ++++++++++++++++++++++++++++++++++----
 2 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/src/app/CommandHandler.cpp b/src/app/CommandHandler.cpp
index 56f60e1fbd9d1e..15374a93fe0135 100644
--- a/src/app/CommandHandler.cpp
+++ b/src/app/CommandHandler.cpp
@@ -540,6 +540,17 @@ CHIP_ERROR CommandHandler::FinishStatus()
     return CHIP_NO_ERROR;
 }
 
+CHIP_ERROR CommandHandler::ResetResponse()
+{
+    VerifyOrReturnError(mState == State::Idle || mState == State::AddedCommand || mState == State::AddingCommand,
+                        CHIP_ERROR_INCORRECT_STATE);
+    // Calling mCommandMessageWriter will release its underlying buffer, thus we can allocate another one when encode something.
+    mCommandMessageWriter.Reset();
+    mBufferAllocated = false;
+    MoveToState(State::Idle);
+    return CHIP_NO_ERROR;
+}
+
 TLV::TLVWriter * CommandHandler::GetCommandDataIBTLVWriter()
 {
     if (mState != State::AddingCommand)
diff --git a/src/app/CommandHandler.h b/src/app/CommandHandler.h
index 975fdf4b9ffbdd..1fa1397e6a8cff 100644
--- a/src/app/CommandHandler.h
+++ b/src/app/CommandHandler.h
@@ -162,6 +162,15 @@ class CommandHandler
     TLV::TLVWriter * GetCommandDataIBTLVWriter();
     FabricIndex GetAccessingFabricIndex() const;
 
+    /**
+     * API for resetting the internal response builder, useful when encoding the response manually.
+     * The TLVWriter got from GetCommandDataIBTLVWriter will be invalid after calling this.
+     *
+     * After calling this, users must call PrepareCommand or PrepareStatus before encoding something else. AddResponseData and
+     * AddStatus will handle this correctly.
+     */
+    CHIP_ERROR ResetResponse();
+
     /**
      * API for adding a data response.  The template parameter T is generally
      * expected to be a ClusterName::Commands::CommandName::Type struct, but any
@@ -175,13 +184,37 @@ class CommandHandler
     template <typename CommandData>
     CHIP_ERROR AddResponseData(const ConcreteCommandPath & aRequestCommandPath, const CommandData & aData)
     {
-        ConcreteCommandPath path = { aRequestCommandPath.mEndpointId, aRequestCommandPath.mClusterId, CommandData::GetCommandId() };
-        ReturnErrorOnFailure(PrepareCommand(path, false));
-        TLV::TLVWriter * writer = GetCommandDataIBTLVWriter();
-        VerifyOrReturnError(writer != nullptr, CHIP_ERROR_INCORRECT_STATE);
-        ReturnErrorOnFailure(DataModel::Encode(*writer, TLV::ContextTag(to_underlying(CommandDataIB::Tag::kData)), aData));
+        // We should not encode anything when we are not in the correct state, the user must be using CommandHandler incorrectly.
+        VerifyOrReturnError(mState == State::Idle, CHIP_ERROR_INCORRECT_STATE);
 
-        return FinishCommand(/* aEndDataStruct = */ false);
+        CHIP_ERROR err = TryAddResponseData(aRequestCommandPath, aData);
+        if (err != CHIP_NO_ERROR)
+        {
+            // We have verified the state above, so this call must success since we must be one of the state required by
+            // ResetResponse.
+            ResetResponse();
+        }
+        return err;
+    }
+
+    /**
+     * API for adding a data response. Will encode Protocols::InteractionModel::Status::Failure status code when it failed to encode
+     * the command data.  The template parameter T is generally expected to be a ClusterName::Commands::CommandName::Type struct,
+     * but any object that can be encoded using the DataModel::Encode machinery and exposes the right command id will work.
+     *
+     * @param [in] aRequestCommandPath the concrete path of the command we are
+     *             responding to.
+     * @param [in] aData the data for the response.
+     */
+    template <typename CommandData>
+    CHIP_ERROR AddResponseDataOrFailureStatus(const ConcreteCommandPath & aRequestCommandPath, const CommandData & aData)
+    {
+        CHIP_ERROR err = AddResponseData(aRequestCommandPath, aData);
+        if (err != CHIP_NO_ERROR)
+        {
+            return AddStatus(aRequestCommandPath, Protocols::InteractionModel::Status::Failure);
+        }
+        return err;
     }
 
     /**
@@ -270,6 +303,28 @@ class CommandHandler
     CHIP_ERROR AddStatusInternal(const ConcreteCommandPath & aCommandPath, const Protocols::InteractionModel::Status aStatus,
                                  const Optional<ClusterStatus> & aClusterStatus);
 
+    /**
+     * Adds a data response.  The template parameter T is generally
+     * expected to be a ClusterName::Commands::CommandName::Type struct, but any
+     * object that can be encoded using the DataModel::Encode machinery and
+     * exposes the right command id will work.
+     *
+     * @param [in] aRequestCommandPath the concrete path of the command we are
+     *             responding to.
+     * @param [in] aData the data for the response.
+     */
+    template <typename CommandData>
+    CHIP_ERROR TryAddResponseData(const ConcreteCommandPath & aRequestCommandPath, const CommandData & aData)
+    {
+        ConcreteCommandPath path = { aRequestCommandPath.mEndpointId, aRequestCommandPath.mClusterId, CommandData::GetCommandId() };
+        ReturnErrorOnFailure(PrepareCommand(path, false));
+        TLV::TLVWriter * writer = GetCommandDataIBTLVWriter();
+        VerifyOrReturnError(writer != nullptr, CHIP_ERROR_INCORRECT_STATE);
+        ReturnErrorOnFailure(DataModel::Encode(*writer, TLV::ContextTag(to_underlying(CommandDataIB::Tag::kData)), aData));
+
+        return FinishCommand(/* aEndDataStruct = */ false);
+    }
+
     Messaging::ExchangeContext * mpExchangeCtx = nullptr;
     Callback * mpCallback                      = nullptr;
     InvokeResponseMessage::Builder mInvokeResponseBuilder;