[IM] Make CommandHandler::AddResponseData atomic (#15906)

* Add encode failure tests

* [IM] Make CommandHandler::AddResponseData atomic

* Reset -> Rollback

* Rename: AddResponseData -> AddResponse

* Address comments

* Update cluster code

* 50 ms is too short on CI

* Fix

* Fix

examples/ota-provider-app/ota-provider-common/OTAProviderExample.cpp

@@ -282,8 +282,7 @@ CHIP_ERROR OTAProviderExample::SendQueryImageResponse(chip::app::CommandHandler
         response.metadataForRequestor.Emplace(chip::ByteSpan());
     }
 
-    ReturnErrorOnFailure(commandObj->AddResponseData(commandPath, response));
-
+    commandObj->AddResponse(commandPath, response);
     return CHIP_NO_ERROR;
 }
 
@@ -390,7 +389,7 @@ EmberAfStatus OTAProviderExample::HandleApplyUpdateRequest(chip::app::CommandHan
     // Reset back to success case for subsequent uses
     mUpdateAction = OTAApplyUpdateAction::kProceed;
 
-    VerifyOrReturnError(commandObj->AddResponseData(commandPath, response) == CHIP_NO_ERROR, EMBER_ZCL_STATUS_FAILURE);
+    commandObj->AddResponse(commandPath, response);
 
     return EMBER_ZCL_STATUS_SUCCESS;
 }

src/app/CommandHandler.cpp

@@ -481,10 +481,13 @@ CHIP_ERROR CommandHandler::AddClusterSpecificFailure(const ConcreteCommandPath &
 CHIP_ERROR CommandHandler::PrepareCommand(const ConcreteCommandPath & aCommandPath, bool aStartDataStruct)
 {
     ReturnErrorOnFailure(AllocateBuffer());
+
+    mInvokeResponseBuilder.Checkpoint(mBackupWriter);
     //
     // We must not be in the middle of preparing a command, or having prepared or sent one.
     //
     VerifyOrReturnError(mState == State::Idle, CHIP_ERROR_INCORRECT_STATE);
+    MoveToState(State::Preparing);
     InvokeResponseIBs::Builder & invokeResponses = mInvokeResponseBuilder.GetInvokeResponses();
     InvokeResponseIB::Builder & invokeResponse   = invokeResponses.CreateInvokeResponse();
     ReturnErrorOnFailure(invokeResponses.GetError());
@@ -526,6 +529,7 @@ CHIP_ERROR CommandHandler::PrepareStatus(const ConcreteCommandPath & aCommandPat
     // We must not be in the middle of preparing a command, or having prepared or sent one.
     //
     VerifyOrReturnError(mState == State::Idle, CHIP_ERROR_INCORRECT_STATE);
+    MoveToState(State::Preparing);
     InvokeResponseIBs::Builder & invokeResponses = mInvokeResponseBuilder.GetInvokeResponses();
     InvokeResponseIB::Builder & invokeResponse   = invokeResponses.CreateInvokeResponse();
     ReturnErrorOnFailure(invokeResponses.GetError());
@@ -550,6 +554,17 @@ CHIP_ERROR CommandHandler::FinishStatus()
     return CHIP_NO_ERROR;
 }
 
+CHIP_ERROR CommandHandler::RollbackResponse()
+{
+    VerifyOrReturnError(mState == State::Preparing || mState == State::AddingCommand, CHIP_ERROR_INCORRECT_STATE);
+    mInvokeResponseBuilder.Rollback(mBackupWriter);
+    mInvokeResponseBuilder.ResetError();
+    // Note: We only support one command per request, so we reset the state to Idle here, need to review the states when adding
+    // supports of having multiple requests in the same transaction.
+    MoveToState(State::Idle);
+    return CHIP_NO_ERROR;
+}
+
 TLV::TLVWriter * CommandHandler::GetCommandDataIBTLVWriter()
 {
     if (mState != State::AddingCommand)
@@ -607,6 +622,9 @@ const char * CommandHandler::GetStateStr() const
     case State::Idle:
         return "Idle";
 
+    case State::Preparing:
+        return "Preparing";
+
     case State::AddingCommand:
         return "AddingCommand";
 

src/app/CommandHandler.h

@@ -175,13 +175,42 @@ 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));
+        // TryAddResponseData will ensure we are in the correct state when calling AddResponseData.
+        CHIP_ERROR err = TryAddResponseData(aRequestCommandPath, aData);
+        if (err != CHIP_NO_ERROR)
+        {
+            // The state guarantees that either we can rollback or we don't have to rollback the buffer, so we don't care about the
+            // return value of RollbackResponse.
+            RollbackResponse();
+        }
+        return err;
+    }
 
-        return FinishCommand(/* aEndDataStruct = */ false);
+    /**
+     * API for adding a response.  This will try to encode a data response (response command), and if that fails will encode a a
+     * Protocols::InteractionModel::Status::Failure status response instead.
+     *
+     * 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.
+     *
+     * Since the function will call AddStatus when it fails to encode the data, it cannot send any response when it fails to encode
+     * a status code since another AddStatus call will also fail. The error from AddStatus will just be logged.
+     *
+     * @param [in] aRequestCommandPath the concrete path of the command we are
+     *             responding to.
+     * @param [in] aData the data for the response.
+     */
+    template <typename CommandData>
+    void AddResponse(const ConcreteCommandPath & aRequestCommandPath, const CommandData & aData)
+    {
+        if (CHIP_NO_ERROR != AddResponseData(aRequestCommandPath, aData))
+        {
+            CHIP_ERROR err = AddStatus(aRequestCommandPath, Protocols::InteractionModel::Status::Failure);
+            if (err != CHIP_NO_ERROR)
+            {
+                ChipLogError(DataManagement, "Failed to encode status: %" CHIP_ERROR_FORMAT, err.Format());
+            }
+        }
     }
 
     /**
@@ -207,6 +236,7 @@ class CommandHandler
     enum class State
     {
         Idle,                ///< Default state that the object starts out in, where no work has commenced
+        Preparing,           ///< We are prepaing the command or status header.
         AddingCommand,       ///< In the process of adding a command.
         AddedCommand,        ///< A command has been completely encoded and is awaiting transmission.
         CommandSent,         ///< The command has been sent successfully.
@@ -216,6 +246,11 @@ class CommandHandler
     void MoveToState(const State aTargetState);
     const char * GetStateStr() const;
 
+    /**
+     * Rollback the state to before encoding the current ResponseData (before calling PrepareCommand / PrepareStatus)
+     */
+    CHIP_ERROR RollbackResponse();
+
     /*
      * This forcibly closes the exchange context if a valid one is pointed to. Such a situation does
      * not arise during normal message processing flows that all normally call Close() above. This can only
@@ -270,6 +305,26 @@ class CommandHandler
     CHIP_ERROR AddStatusInternal(const ConcreteCommandPath & aCommandPath, const Protocols::InteractionModel::Status aStatus,
                                  const Optional<ClusterStatus> & aClusterStatus);
 
+    /**
+     * If this function fails, it may leave our TLV buffer in an inconsistent state.  Callers should snapshot as needed before
+     * calling this function, and roll back as needed afterward.
+     *
+     * @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;
@@ -280,6 +335,7 @@ class CommandHandler
 
     State mState = State::Idle;
     chip::System::PacketBufferTLVWriter mCommandMessageWriter;
+    TLV::TLVWriter mBackupWriter;
     bool mBufferAllocated = false;
 };
 

src/app/CommandResponseHelper.h

@@ -36,12 +36,9 @@ class CommandResponseHelper
 
     CHIP_ERROR Success(const CommandData & aResponse)
     {
-        CHIP_ERROR err = mCommandHandler->AddResponseData(mCommandPath, aResponse);
-        if (err == CHIP_NO_ERROR)
-        {
-            mSentResponse = true;
-        }
-        return err;
+        mCommandHandler->AddResponse(mCommandPath, aResponse);
+        mSentResponse = true;
+        return CHIP_NO_ERROR;
     };
 
     CHIP_ERROR Success(ClusterStatus aClusterStatus)

src/app/clusters/diagnostic-logs-server/diagnostic-logs-server.cpp

@@ -45,7 +45,7 @@ void DiagnosticLogsCommandHandler::InvokeCommand(HandlerContext & handlerContext
                 if (mBuffer.IsEmpty())
                 {
                     response.status = chip::app::Clusters::DiagnosticLogs::LogsStatus::kNoLogs;
-                    handlerContext.mCommandHandler.AddResponseData(handlerContext.mRequestPath, response);
+                    handlerContext.mCommandHandler.AddResponse(handlerContext.mRequestPath, response);
                     break;
                 }
 
@@ -58,7 +58,7 @@ void DiagnosticLogsCommandHandler::InvokeCommand(HandlerContext & handlerContext
                 if (!buf)
                 {
                     response.status = chip::app::Clusters::DiagnosticLogs::LogsStatus::kBusy;
-                    handlerContext.mCommandHandler.AddResponseData(handlerContext.mRequestPath, response);
+                    handlerContext.mCommandHandler.AddResponse(handlerContext.mRequestPath, response);
                     break;
                 }
 
@@ -71,19 +71,19 @@ void DiagnosticLogsCommandHandler::InvokeCommand(HandlerContext & handlerContext
                 response.status    = chip::app::Clusters::DiagnosticLogs::LogsStatus::kSuccess;
                 response.content   = chip::ByteSpan(buf.get() + sizeof(timeMs), logSize - sizeof(timeMs));
                 response.timeStamp = timeMs;
-                handlerContext.mCommandHandler.AddResponseData(handlerContext.mRequestPath, response);
+                handlerContext.mCommandHandler.AddResponse(handlerContext.mRequestPath, response);
             }
             break;
             case chip::app::Clusters::DiagnosticLogs::LogsIntent::kNetworkDiag: {
                 chip::app::Clusters::DiagnosticLogs::Commands::RetrieveLogsResponse::Type response;
                 response.status = chip::app::Clusters::DiagnosticLogs::LogsStatus::kNoLogs;
-                handlerContext.mCommandHandler.AddResponseData(handlerContext.mRequestPath, response);
+                handlerContext.mCommandHandler.AddResponse(handlerContext.mRequestPath, response);
             }
             break;
             case chip::app::Clusters::DiagnosticLogs::LogsIntent::kCrashLogs: {
                 chip::app::Clusters::DiagnosticLogs::Commands::RetrieveLogsResponse::Type response;
                 response.status = chip::app::Clusters::DiagnosticLogs::LogsStatus::kNoLogs;
-                handlerContext.mCommandHandler.AddResponseData(handlerContext.mRequestPath, response);
+                handlerContext.mCommandHandler.AddResponse(handlerContext.mRequestPath, response);
             }
             break;
             }

src/app/clusters/door-lock-server/door-lock-server.cpp

@@ -2169,11 +2169,10 @@ DlStatus DoorLockServer::clearSchedules(chip::EndpointId endpointId, uint16_t us
     return DlStatus::kSuccess;
 }
 
-CHIP_ERROR DoorLockServer::sendGetWeekDayScheduleResponse(chip::app::CommandHandler * commandObj,
-                                                          const chip::app::ConcreteCommandPath & commandPath, uint8_t weekdayIndex,
-                                                          uint16_t userIndex, DlStatus status, DlDaysMaskMap daysMask,
-                                                          uint8_t startHour, uint8_t startMinute, uint8_t endHour,
-                                                          uint8_t endMinute)
+void DoorLockServer::sendGetWeekDayScheduleResponse(chip::app::CommandHandler * commandObj,
+                                                    const chip::app::ConcreteCommandPath & commandPath, uint8_t weekdayIndex,
+                                                    uint16_t userIndex, DlStatus status, DlDaysMaskMap daysMask, uint8_t startHour,
+                                                    uint8_t startMinute, uint8_t endHour, uint8_t endMinute)
 {
     VerifyOrDie(nullptr != commandObj);
 
@@ -2190,7 +2189,7 @@ CHIP_ERROR DoorLockServer::sendGetWeekDayScheduleResponse(chip::app::CommandHand
         response.endMinute   = Optional<uint8_t>(endMinute);
     }
 
-    return commandObj->AddResponseData(commandPath, response);
+    commandObj->AddResponse(commandPath, response);
 }
 
 bool DoorLockServer::yearDayIndexValid(chip::EndpointId endpointId, uint8_t yearDayIndex)
@@ -2242,10 +2241,10 @@ DlStatus DoorLockServer::clearYearDaySchedules(chip::EndpointId endpointId, uint
     return DlStatus::kSuccess;
 }
 
-CHIP_ERROR DoorLockServer::sendGetYearDayScheduleResponse(chip::app::CommandHandler * commandObj,
-                                                          const chip::app::ConcreteCommandPath & commandPath, uint8_t yearDayIndex,
-                                                          uint16_t userIndex, DlStatus status, uint32_t localStartTime,
-                                                          uint32_t localEndTime)
+void DoorLockServer::sendGetYearDayScheduleResponse(chip::app::CommandHandler * commandObj,
+                                                    const chip::app::ConcreteCommandPath & commandPath, uint8_t yearDayIndex,
+                                                    uint16_t userIndex, DlStatus status, uint32_t localStartTime,
+                                                    uint32_t localEndTime)
 {
     VerifyOrDie(nullptr != commandObj);
 
@@ -2259,7 +2258,7 @@ CHIP_ERROR DoorLockServer::sendGetYearDayScheduleResponse(chip::app::CommandHand
         response.localEndTime   = Optional<uint32_t>(localEndTime);
     }
 
-    return commandObj->AddResponseData(commandPath, response);
+    commandObj->AddResponse(commandPath, response);
 }
 
 EmberAfStatus DoorLockServer::clearCredential(chip::EndpointId endpointId, chip::FabricIndex modifier, chip::NodeId sourceNodeId,

src/app/clusters/door-lock-server/door-lock-server.h

@@ -238,21 +238,19 @@ class DoorLockServer
     DlStatus clearWeekDaySchedules(chip::EndpointId endpointId, uint16_t userIndex);
     DlStatus clearSchedules(chip::EndpointId endpointId, uint16_t userIndex);
 
-    CHIP_ERROR sendGetWeekDayScheduleResponse(chip::app::CommandHandler * commandObj,
-                                              const chip::app::ConcreteCommandPath & commandPath, uint8_t weekdayIndex,
-                                              uint16_t userIndex, DlStatus status, DlDaysMaskMap daysMask = DlDaysMaskMap(0),
-                                              uint8_t startHour = 0, uint8_t startMinute = 0, uint8_t endHour = 0,
-                                              uint8_t endMinute = 0);
+    void sendGetWeekDayScheduleResponse(chip::app::CommandHandler * commandObj, const chip::app::ConcreteCommandPath & commandPath,
+                                        uint8_t weekdayIndex, uint16_t userIndex, DlStatus status,
+                                        DlDaysMaskMap daysMask = DlDaysMaskMap(0), uint8_t startHour = 0, uint8_t startMinute = 0,
+                                        uint8_t endHour = 0, uint8_t endMinute = 0);
 
     bool yearDayIndexValid(chip::EndpointId endpointId, uint8_t yearDayIndex);
 
     DlStatus clearYearDaySchedule(chip::EndpointId endpointId, uint16_t userIndex, uint8_t weekDayIndex);
     DlStatus clearYearDaySchedules(chip::EndpointId endpointId, uint16_t userIndex);
 
-    CHIP_ERROR sendGetYearDayScheduleResponse(chip::app::CommandHandler * commandObj,
-                                              const chip::app::ConcreteCommandPath & commandPath, uint8_t yearDayIndex,
-                                              uint16_t userIndex, DlStatus status, uint32_t localStartTime = 0,
-                                              uint32_t localEndTime = 0);
+    void sendGetYearDayScheduleResponse(chip::app::CommandHandler * commandObj, const chip::app::ConcreteCommandPath & commandPath,
+                                        uint8_t yearDayIndex, uint16_t userIndex, DlStatus status, uint32_t localStartTime = 0,
+                                        uint32_t localEndTime = 0);
 
     bool sendRemoteLockUserChange(chip::EndpointId endpointId, DlLockDataType dataType, DlDataOperationType operation,
                                   chip::NodeId nodeId, chip::FabricIndex fabricIndex, uint16_t userIndex = 0,

src/app/clusters/general-commissioning-server/general-commissioning-server.cpp

@@ -151,12 +151,12 @@ bool emberAfGeneralCommissioningClusterArmFailSafeCallback(app::CommandHandler *
         CheckSuccess(failSafeContext.ArmFailSafe(accessingFabricIndex, System::Clock::Seconds16(commandData.expiryLengthSeconds)),
                      Failure);
         response.errorCode = CommissioningError::kOk;
-        CheckSuccess(commandObj->AddResponseData(commandPath, response), Failure);
+        commandObj->AddResponse(commandPath, response);
     }
     else
     {
         response.errorCode = CommissioningError::kBusyWithOtherAdmin;
-        CheckSuccess(commandObj->AddResponseData(commandPath, response), Failure);
+        commandObj->AddResponse(commandPath, response);
     }
 
     return true;
@@ -201,7 +201,7 @@ bool emberAfGeneralCommissioningClusterCommissioningCompleteCallback(
         }
     }
 
-    CheckSuccess(commandObj->AddResponseData(commandPath, response), Failure);
+    commandObj->AddResponse(commandPath, response);
 
     return true;
 }
@@ -218,7 +218,7 @@ bool emberAfGeneralCommissioningClusterSetRegulatoryConfigCallback(app::CommandH
 
     Commands::SetRegulatoryConfigResponse::Type response;
     response.errorCode = CommissioningError::kOk;
-    CheckSuccess(commandObj->AddResponseData(commandPath, response), Failure);
+    commandObj->AddResponse(commandPath, response);
 
     return true;
 }

src/app/clusters/group-key-mgmt-server/group-key-mgmt-server.cpp

@@ -443,12 +443,7 @@ bool emberAfGroupKeyManagementClusterKeySetReadCallback(
     }
     response.groupKeySet.epochKey2.SetNull();
 
-    CHIP_ERROR err = commandObj->AddResponseData(commandPath, response);
-    if (CHIP_NO_ERROR != err)
-    {
-        ChipLogDetail(Zcl, "GroupKeyManagementCluster: KeySetRead failed: %s", ErrorStr(err));
-        emberAfSendImmediateDefaultResponse(EMBER_ZCL_STATUS_FAILURE);
-    }
+    commandObj->AddResponse(commandPath, response);
     return true;
 }
 
@@ -535,11 +530,7 @@ bool emberAfGroupKeyManagementClusterKeySetReadAllIndicesCallback(
         return true;
     }
 
-    CHIP_ERROR err = commandObj->AddResponseData(commandPath, KeySetReadAllIndicesResponse(keysIt));
-    if (CHIP_NO_ERROR != err)
-    {
-        ChipLogDetail(Zcl, "GroupKeyManagementCluster: KeySetReadAllIndices failed: %s", ErrorStr(err));
-    }
+    commandObj->AddResponse(commandPath, KeySetReadAllIndicesResponse(keysIt));
     keysIt->Release();
     return true;
 }