Update CommandHandler to support batch commands

src/app/CommandHandler.cpp

@@ -32,6 +32,7 @@
 #include <app/RequiredPrivilege.h>
 #include <app/util/MatterCallbacks.h>
 #include <credentials/GroupDataProvider.h>
+#include <lib/core/CHIPConfig.h>
 #include <lib/core/TLVData.h>
 #include <lib/core/TLVUtilities.h>
 #include <lib/support/TypeTraits.h>
@@ -44,6 +45,13 @@ using Status = Protocols::InteractionModel::Status;
 
 CommandHandler::CommandHandler(Callback * apCallback) : mExchangeCtx(*this), mpCallback(apCallback), mSuppressResponse(false) {}
 
+CommandHandler::CommandHandler(TestOnlyMarker aTestMarker, Callback * apCallback, CommandPathRegistry * apCommandPathRegistry) :
+    CommandHandler(apCallback)
+{
+    mMaxPathsPerInvoke   = apCommandPathRegistry->MaxSize();
+    mCommandPathRegistry = apCommandPathRegistry;
+}
+
 CHIP_ERROR CommandHandler::AllocateBuffer()
 {
     if (!mBufferAllocated)
@@ -97,40 +105,108 @@ void CommandHandler::OnInvokeCommandRequest(Messaging::ExchangeContext * ec, con
     mGoneAsync = true;
 }
 
+CHIP_ERROR CommandHandler::ValidateInvokeRequestsAndBuildRegistry(TLV::TLVReader & invokeRequestsReader)
+{
+    CHIP_ERROR err          = CHIP_NO_ERROR;
+    size_t commandCount     = 0;
+    bool commandRefExpected = false;
+
+    ReturnErrorOnFailure(TLV::Utilities::Count(invokeRequestsReader, commandCount, false /* recurse */));
+
+    // If this is a GroupRequest the only thing to check is that there is only one
+    // CommandDataIB.
+    if (IsGroupRequest())
+    {
+        VerifyOrReturnError(commandCount == 1, CHIP_ERROR_INVALID_ARGUMENT);
+        return CHIP_NO_ERROR;
+    }
+    // While technically any commandCount == 1 should already be unique and does not need
+    // any further validation, we do need to read and populate the registry to help
+    // in building the InvokeResponse.
+
+    VerifyOrReturnError(commandCount <= MaxPathsPerInvoke(), CHIP_ERROR_INVALID_ARGUMENT);
+
+    // If there is more than one CommandDataIB, spec states that CommandRef must be provided.
+    commandRefExpected = commandCount > 1;
+
+    while (CHIP_NO_ERROR == (err = invokeRequestsReader.Next()))
+    {
+        VerifyOrReturnError(TLV::AnonymousTag() == invokeRequestsReader.GetTag(), CHIP_ERROR_INVALID_ARGUMENT);
+        CommandDataIB::Parser commandData;
+        ReturnErrorOnFailure(commandData.Init(invokeRequestsReader));
+
+        // First validate that we can get a ConcreteCommandPath.
+        CommandPathIB::Parser commandPath;
+        ConcreteCommandPath concretePath(0, 0, 0);
+        ReturnErrorOnFailure(commandData.GetPath(&commandPath));
+        ReturnErrorOnFailure(commandPath.GetConcreteCommandPath(concretePath));
+
+        // Grab the CommandRef if there is one, and validate that it's there when it
+        // has to be.
+        Optional<uint16_t> commandRef;
+        uint16_t ref;
+        err = commandData.GetRef(&ref);
+        VerifyOrReturnError(err == CHIP_NO_ERROR || err == CHIP_END_OF_TLV, err);
+        if (err == CHIP_END_OF_TLV && commandRefExpected)
+        {
+            return CHIP_ERROR_INVALID_ARGUMENT;
+        }
+        if (err == CHIP_NO_ERROR)
+        {
+            commandRef.SetValue(ref);
+        }
+
+        // Adding can fail if concretePath is not unique, or if commandRef is a value
+        // and is not unique, or if we have already added more paths than we support.
+        ReturnErrorOnFailure(GetCommandPathRegistry().Add(concretePath, commandRef));
+    }
+
+    // It's OK/expected to have reached the end of the container without failure.
+    if (CHIP_END_OF_TLV == err)
+    {
+        err = CHIP_NO_ERROR;
+    }
+    return err;
+}
+
 Status CommandHandler::ProcessInvokeRequest(System::PacketBufferHandle && payload, bool isTimedInvoke)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
     System::PacketBufferTLVReader reader;
-    TLV::TLVReader invokeRequestsReader;
     InvokeRequestMessage::Parser invokeRequestMessage;
     InvokeRequests::Parser invokeRequests;
     reader.Init(std::move(payload));
     VerifyOrReturnError(invokeRequestMessage.Init(reader) == CHIP_NO_ERROR, Status::InvalidAction);
 #if CHIP_CONFIG_IM_PRETTY_PRINT
     invokeRequestMessage.PrettyPrint();
 #endif
+    if (mExchangeCtx->IsGroupExchangeContext())
+    {
+        SetGroupRequest(true);
+    }
 
     VerifyOrReturnError(invokeRequestMessage.GetSuppressResponse(&mSuppressResponse) == CHIP_NO_ERROR, Status::InvalidAction);
     VerifyOrReturnError(invokeRequestMessage.GetTimedRequest(&mTimedRequest) == CHIP_NO_ERROR, Status::InvalidAction);
     VerifyOrReturnError(invokeRequestMessage.GetInvokeRequests(&invokeRequests) == CHIP_NO_ERROR, Status::InvalidAction);
     VerifyOrReturnError(mTimedRequest == isTimedInvoke, Status::UnsupportedAccess);
-    invokeRequests.GetReader(&invokeRequestsReader);
 
     {
-        // We don't support handling multiple commands but the protocol is ready to support it in the future, reject all of them and
-        // IM Engine will send a status response.
-        size_t commandCount = 0;
-        TLV::Utilities::Count(invokeRequestsReader, commandCount, false /* recurse */);
-        VerifyOrReturnError(commandCount == 1, Status::InvalidAction);
+        TLV::TLVReader validationInvokeRequestsReader;
+        invokeRequests.GetReader(&validationInvokeRequestsReader);
+        VerifyOrReturnError(ValidateInvokeRequestsAndBuildRegistry(validationInvokeRequestsReader) == CHIP_NO_ERROR,
+                            Status::InvalidAction);
     }
 
+    TLV::TLVReader invokeRequestsReader;
+    invokeRequests.GetReader(&invokeRequestsReader);
+
     while (CHIP_NO_ERROR == (err = invokeRequestsReader.Next()))
     {
         VerifyOrReturnError(TLV::AnonymousTag() == invokeRequestsReader.GetTag(), Status::InvalidAction);
         CommandDataIB::Parser commandData;
         VerifyOrReturnError(commandData.Init(invokeRequestsReader) == CHIP_NO_ERROR, Status::InvalidAction);
         Status status = Status::Success;
-        if (mExchangeCtx->IsGroupExchangeContext())
+        if (IsGroupRequest())
         {
             status = ProcessGroupCommandDataIB(commandData);
         }
@@ -200,7 +276,7 @@ void CommandHandler::DecrementHoldOff()
         {
             ChipLogProgress(DataManagement, "Skipping command response: exchange context is null");
         }
-        else if (!mExchangeCtx->IsGroupExchangeContext())
+        else if (!IsGroupRequest())
         {
             CHIP_ERROR err = SendCommandResponse();
             if (err != CHIP_NO_ERROR)
@@ -251,8 +327,6 @@ Status CommandHandler::ProcessCommandDataIB(CommandDataIB::Parser & aCommandElem
     ConcreteCommandPath concretePath(0, 0, 0);
     TLV::TLVReader commandDataReader;
 
-    SetGroupRequest(false);
-
     // NOTE: errors may occur before the concrete command path is even fully decoded.
 
     err = aCommandElement.GetPath(&commandPath);
@@ -354,7 +428,6 @@ Status CommandHandler::ProcessGroupCommandDataIB(CommandDataIB::Parser & aComman
     Credentials::GroupDataProvider::GroupEndpoint mapping;
     Credentials::GroupDataProvider * groupDataProvider = Credentials::GetGroupDataProvider();
     Credentials::GroupDataProvider::EndpointIterator * iterator;
-    SetGroupRequest(true);
 
     err = aCommandElement.GetPath(&commandPath);
     VerifyOrReturnError(err == CHIP_NO_ERROR, Status::InvalidAction);
@@ -465,7 +538,7 @@ CHIP_ERROR CommandHandler::AddStatusInternal(const ConcreteCommandPath & aComman
 void CommandHandler::AddStatus(const ConcreteCommandPath & aCommandPath, const Protocols::InteractionModel::Status aStatus,
                                const char * context)
 {
-    // Return prematurely in case of requests targeted to a group that should not add the status for response purposes.
+    // Return early in case of requests targeted to a group, since they should not add a response.
     VerifyOrReturn(!IsGroupRequest());
     VerifyOrDie(FallibleAddStatus(aCommandPath, aStatus, context) == CHIP_NO_ERROR);
 }
@@ -500,15 +573,49 @@ CHIP_ERROR CommandHandler::AddClusterSpecificFailure(const ConcreteCommandPath &
     return AddStatusInternal(aCommandPath, StatusIB(Status::Failure, aClusterStatus));
 }
 
-CHIP_ERROR CommandHandler::PrepareCommand(const ConcreteCommandPath & aCommandPath, bool aStartDataStruct)
+CHIP_ERROR CommandHandler::PrepareInvokeResponseCommand(const ConcreteCommandPath & aResponseCommandPath,
+                                                        const CommandHandler::InvokeResponseParameters & aPrepareParameters)
+{
+    auto commandPathRegistryEntry = GetCommandPathRegistry().Find(aPrepareParameters.mRequestCommandPath);
+    VerifyOrReturnValue(commandPathRegistryEntry.HasValue(), CHIP_ERROR_INCORRECT_STATE);
+
+    return PrepareInvokeResponseCommand(commandPathRegistryEntry.Value(), aResponseCommandPath,
+                                        aPrepareParameters.mStartOrEndDataStruct);
+}
+
+CHIP_ERROR CommandHandler::PrepareCommand(const ConcreteCommandPath & aResponseCommandPath, bool aStartDataStruct)
+{
+    // Legacy code is calling the deprecated version of PrepareCommand. If we are in a case where
+    // there was a single command in the request, we can just assume this response is triggered by
+    // the single command.
+    size_t countOfPathRegistryEntries = GetCommandPathRegistry().Count();
+
+    // At this point application supports Batch Invoke Commands since CommandPathRegistry has more than 1 entry,
+    // but application is calling the deprecated PrepareCommand. We have no way to determine the associated CommandRef
+    // to put into the InvokeResponse.
+    VerifyOrDieWithMsg(countOfPathRegistryEntries > 1, DataManagement,
+                       "Seemingly device supports batch commands, but is calling the deprecated PrepareCommand API");
+
+    auto commandPathRegistryEntry = GetCommandPathRegistry().GetFirstEntry();
+    VerifyOrReturnValue(commandPathRegistryEntry.HasValue(), CHIP_ERROR_INCORRECT_STATE);
+
+    return PrepareInvokeResponseCommand(commandPathRegistryEntry.Value(), aResponseCommandPath, aStartDataStruct);
+}
+
+CHIP_ERROR CommandHandler::PrepareInvokeResponseCommand(const CommandPathRegistryEntry & apCommandPathRegistryEntry,
+                                                        const ConcreteCommandPath & aCommandPath, bool aStartDataStruct)
 {
     ReturnErrorOnFailure(AllocateBuffer());
 
     mInvokeResponseBuilder.Checkpoint(mBackupWriter);
+    mBackupState = mState;
     //
     // 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);
+    VerifyOrReturnError(mState == State::Idle || mState == State::AddedCommand, CHIP_ERROR_INCORRECT_STATE);
+
+    mRefForResponse = apCommandPathRegistryEntry.ref;
+
     MoveToState(State::Preparing);
     InvokeResponseIBs::Builder & invokeResponses = mInvokeResponseBuilder.GetInvokeResponses();
     InvokeResponseIB::Builder & invokeResponse   = invokeResponses.CreateInvokeResponse();
@@ -536,10 +643,14 @@ CHIP_ERROR CommandHandler::FinishCommand(bool aStartDataStruct)
     {
         ReturnErrorOnFailure(commandData.GetWriter()->EndContainer(mDataElementContainerType));
     }
+
+    if (mRefForResponse.HasValue())
+    {
+        ReturnErrorOnFailure(commandData.Ref(mRefForResponse.Value()));
+    }
+
     ReturnErrorOnFailure(commandData.EndOfCommandDataIB());
     ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().GetInvokeResponse().EndOfInvokeResponseIB());
-    ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().EndOfInvokeResponses());
-    ReturnErrorOnFailure(mInvokeResponseBuilder.EndOfInvokeResponseMessage());
     MoveToState(State::AddedCommand);
     return CHIP_NO_ERROR;
 }
@@ -550,7 +661,12 @@ 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);
+    VerifyOrReturnError(mState == State::Idle || mState == State::AddedCommand, CHIP_ERROR_INCORRECT_STATE);
+
+    auto commandPathRegistryEntry = GetCommandPathRegistry().Find(aCommandPath);
+    VerifyOrReturnError(commandPathRegistryEntry.HasValue(), CHIP_ERROR_INCORRECT_STATE);
+    mRefForResponse = commandPathRegistryEntry.Value().ref;
+
     MoveToState(State::Preparing);
     InvokeResponseIBs::Builder & invokeResponses = mInvokeResponseBuilder.GetInvokeResponses();
     InvokeResponseIB::Builder & invokeResponse   = invokeResponses.CreateInvokeResponse();
@@ -567,10 +683,15 @@ CHIP_ERROR CommandHandler::PrepareStatus(const ConcreteCommandPath & aCommandPat
 CHIP_ERROR CommandHandler::FinishStatus()
 {
     VerifyOrReturnError(mState == State::AddingCommand, CHIP_ERROR_INCORRECT_STATE);
+
+    CommandStatusIB::Builder & commandStatus = mInvokeResponseBuilder.GetInvokeResponses().GetInvokeResponse().GetStatus();
+    if (mRefForResponse.HasValue())
+    {
+        ReturnErrorOnFailure(commandStatus.Ref(mRefForResponse.Value()));
+    }
+
     ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().GetInvokeResponse().GetStatus().EndOfCommandStatusIB());
     ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().GetInvokeResponse().EndOfInvokeResponseIB());
-    ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().EndOfInvokeResponses());
-    ReturnErrorOnFailure(mInvokeResponseBuilder.EndOfInvokeResponseMessage());
     MoveToState(State::AddedCommand);
     return CHIP_NO_ERROR;
 }
@@ -579,9 +700,7 @@ CHIP_ERROR CommandHandler::RollbackResponse()
 {
     VerifyOrReturnError(mState == State::Preparing || mState == State::AddingCommand, CHIP_ERROR_INCORRECT_STATE);
     mInvokeResponseBuilder.Rollback(mBackupWriter);
-    // 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);
+    MoveToState(mBackupState);
     return CHIP_NO_ERROR;
 }
 
@@ -635,6 +754,8 @@ CommandHandler::Handle::Handle(CommandHandler * handle)
 CHIP_ERROR CommandHandler::Finalize(System::PacketBufferHandle & commandPacket)
 {
     VerifyOrReturnError(mState == State::AddedCommand, CHIP_ERROR_INCORRECT_STATE);
+    ReturnErrorOnFailure(mInvokeResponseBuilder.GetInvokeResponses().EndOfInvokeResponses());
+    ReturnErrorOnFailure(mInvokeResponseBuilder.EndOfInvokeResponseMessage());
     return mCommandMessageWriter.Finalize(&commandPacket);
 }