From 3ac709cf2797d9ccb67e14afe64f0729addc54e1 Mon Sep 17 00:00:00 2001
From: Song Guo <songguo@google.com>
Date: Tue, 15 Mar 2022 12:01:46 +0800
Subject: [PATCH] Reset -> Rollback

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

diff --git a/src/app/CommandHandler.cpp b/src/app/CommandHandler.cpp
index 15374a93fe0135..c3bad5c143f43c 100644
--- a/src/app/CommandHandler.cpp
+++ b/src/app/CommandHandler.cpp
@@ -471,6 +471,8 @@ 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.
     //
@@ -512,6 +514,7 @@ CHIP_ERROR CommandHandler::FinishCommand(bool aStartDataStruct)
 CHIP_ERROR CommandHandler::PrepareStatus(const ConcreteCommandPath & aCommandPath)
 {
     ReturnErrorOnFailure(AllocateBuffer());
+    mInvokeResponseBuilder.Checkpoint(mBackupWriter);
     //
     // We must not be in the middle of preparing a command, or having prepared or sent one.
     //
@@ -540,13 +543,14 @@ CHIP_ERROR CommandHandler::FinishStatus()
     return CHIP_NO_ERROR;
 }
 
-CHIP_ERROR CommandHandler::ResetResponse()
+CHIP_ERROR CommandHandler::RollbackResponse()
 {
     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;
+    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;
 }
diff --git a/src/app/CommandHandler.h b/src/app/CommandHandler.h
index 1fa1397e6a8cff..8aee802937de79 100644
--- a/src/app/CommandHandler.h
+++ b/src/app/CommandHandler.h
@@ -163,13 +163,9 @@ class CommandHandler
     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.
+     * Rollback the state to before encoding the current ResponseData (before calling PrepareCommand / PrepareStatus)
      */
-    CHIP_ERROR ResetResponse();
+    CHIP_ERROR RollbackResponse();
 
     /**
      * API for adding a data response.  The template parameter T is generally
@@ -192,7 +188,7 @@ class CommandHandler
         {
             // We have verified the state above, so this call must success since we must be one of the state required by
             // ResetResponse.
-            ResetResponse();
+            RollbackResponse();
         }
         return err;
     }
@@ -335,6 +331,7 @@ class CommandHandler
 
     State mState = State::Idle;
     chip::System::PacketBufferTLVWriter mCommandMessageWriter;
+    TLV::TLVWriter mBackupWriter;
     bool mBufferAllocated = false;
 };