[K32W] Fix apply action corner case in OTATlvProcessor interface (#33214)

* [k32w] Add mShouldNotApply flag in OTATlvProcessor interface

OTATlvProcessor::ApplyAction now has a default implementation,
but derived classes are still able to overwrite this.

The flag is used by the default ApplyAction implementation.
If something goes wrong during ExitAction of the TLV processor,
then mShouldNotApply should be set to true and the image processor
should abort. In this case, the BDX transfer was already finished
and calling CancelImageUpdate will not abort the transfer, hence
the device will reboot even though it should not have. If ApplyAction
fails during HandleApply, then the process will be aborted.

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>

* [k32w0] Use mShouldNotApply flag in ExitAction

During ExitAction, set mShouldNotApply to true if an error occurs.
This ensures that the OTA will be aborted and the device does not
reboot.

Also remove the ApplyAction override, since the default implementation
is enough.

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>

* [k32w1] Use mShouldNotApply during ExitAction

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>

* [k32w] Update OTA error naming

All OTA errors should be prefixed with CHIP_ERROR.

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>

* [k32w] Replace boolean mShouldNotApply with an enum class

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>

---------

Signed-off-by: marius-alex-tache <marius.tache@nxp.com>
diff --git a/src/platform/nxp/k32w/common/OTAImageProcessorImpl.cpp b/src/platform/nxp/k32w/common/OTAImageProcessorImpl.cpp
index 1e00ddc..e972550 100644
--- a/src/platform/nxp/k32w/common/OTAImageProcessorImpl.cpp
+++ b/src/platform/nxp/k32w/common/OTAImageProcessorImpl.cpp
@@ -142,7 +142,7 @@
         }
 
         status = mCurrentProcessor->Process(block);
-        if (status == CHIP_OTA_CHANGE_PROCESSOR)
+        if (status == CHIP_ERROR_OTA_CHANGE_PROCESSOR)
         {
             mAccumulator.Clear();
             mAccumulator.Init(sizeof(OTATlvHeader));
@@ -180,7 +180,7 @@
     if (pair == mProcessorMap.end())
     {
         ChipLogError(SoftwareUpdate, "There is no registered processor for tag: %" PRIu32, header.tag);
-        return CHIP_OTA_PROCESSOR_NOT_REGISTERED;
+        return CHIP_ERROR_OTA_PROCESSOR_NOT_REGISTERED;
     }
 
     ChipLogDetail(SoftwareUpdate, "Selected processor with tag: %ld", pair->first);
@@ -197,7 +197,7 @@
     if (pair != mProcessorMap.end())
     {
         ChipLogError(SoftwareUpdate, "A processor for tag %" PRIu32 " is already registered.", tag);
-        return CHIP_OTA_PROCESSOR_ALREADY_REGISTERED;
+        return CHIP_ERROR_OTA_PROCESSOR_ALREADY_REGISTERED;
     }
 
     mProcessorMap.insert({ tag, processor });
@@ -249,7 +249,7 @@
         mParams.downloadedBytes += mBlock.size();
         FetchNextData(0);
     }
-    else if (status == CHIP_OTA_FETCH_ALREADY_SCHEDULED)
+    else if (status == CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED)
     {
         mParams.downloadedBytes += mBlock.size();
     }
diff --git a/src/platform/nxp/k32w/common/OTATlvProcessor.cpp b/src/platform/nxp/k32w/common/OTATlvProcessor.cpp
index d9d2ccb..5dc0eae 100644
--- a/src/platform/nxp/k32w/common/OTATlvProcessor.cpp
+++ b/src/platform/nxp/k32w/common/OTATlvProcessor.cpp
@@ -31,6 +31,12 @@
 #if OTA_ENCRYPTION_ENABLE
 constexpr uint8_t au8Iv[] = { 0x00, 0x00, 0x00, 0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x00, 0x00, 0x00, 0x00 };
 #endif
+
+CHIP_ERROR OTATlvProcessor::ApplyAction()
+{
+    return mApplyState == ApplyState::kApply ? CHIP_NO_ERROR : CHIP_ERROR_OTA_PROCESSOR_DO_NOT_APPLY;
+}
+
 CHIP_ERROR OTATlvProcessor::Process(ByteSpan & block)
 {
     CHIP_ERROR status     = CHIP_NO_ERROR;
@@ -50,7 +56,7 @@
                 // If current block was processed fully and the block still contains data, it
                 // means that the block contains another TLV's data and the current processor
                 // should be changed by OTAImageProcessorImpl.
-                return CHIP_OTA_CHANGE_PROCESSOR;
+                return CHIP_ERROR_OTA_CHANGE_PROCESSOR;
             }
         }
     }
@@ -63,6 +69,7 @@
     mLength          = 0;
     mProcessedLength = 0;
     mWasSelected     = false;
+    mApplyState      = ApplyState::kApply;
 #if OTA_ENCRYPTION_ENABLE
     mIVOffset = 0;
 #endif
@@ -70,7 +77,7 @@
 
 bool OTATlvProcessor::IsError(CHIP_ERROR & status)
 {
-    return status != CHIP_NO_ERROR && status != CHIP_ERROR_BUFFER_TOO_SMALL && status != CHIP_OTA_FETCH_ALREADY_SCHEDULED;
+    return status != CHIP_NO_ERROR && status != CHIP_ERROR_BUFFER_TOO_SMALL && status != CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED;
 }
 
 void OTADataAccumulator::Init(uint32_t threshold)
diff --git a/src/platform/nxp/k32w/common/OTATlvProcessor.h b/src/platform/nxp/k32w/common/OTATlvProcessor.h
index 65d1d68..f1faef7 100644
--- a/src/platform/nxp/k32w/common/OTATlvProcessor.h
+++ b/src/platform/nxp/k32w/common/OTATlvProcessor.h
@@ -27,20 +27,20 @@
 #define CHIP_ERROR_TLV_PROCESSOR(e)                                                                                                \
     ChipError(ChipError::Range::kLastRange, ((uint8_t) ChipError::Range::kLastRange << 3) | e, __FILE__, __LINE__)
 
-#define CHIP_OTA_TLV_CONTINUE_PROCESSING CHIP_ERROR_TLV_PROCESSOR(0x01)
-#define CHIP_OTA_CHANGE_PROCESSOR CHIP_ERROR_TLV_PROCESSOR(0x02)
-#define CHIP_OTA_PROCESSOR_NOT_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x03)
-#define CHIP_OTA_PROCESSOR_ALREADY_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x04)
-#define CHIP_OTA_PROCESSOR_CLIENT_INIT CHIP_ERROR_TLV_PROCESSOR(0x05)
-#define CHIP_OTA_PROCESSOR_MAKE_ROOM CHIP_ERROR_TLV_PROCESSOR(0x06)
-#define CHIP_OTA_PROCESSOR_PUSH_CHUNK CHIP_ERROR_TLV_PROCESSOR(0x07)
-#define CHIP_OTA_PROCESSOR_IMG_AUTH CHIP_ERROR_TLV_PROCESSOR(0x08)
-#define CHIP_OTA_FETCH_ALREADY_SCHEDULED CHIP_ERROR_TLV_PROCESSOR(0x09)
-#define CHIP_OTA_PROCESSOR_IMG_COMMIT CHIP_ERROR_TLV_PROCESSOR(0x0A)
-#define CHIP_OTA_PROCESSOR_CB_NOT_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x0B)
-#define CHIP_OTA_PROCESSOR_EEPROM_OFFSET CHIP_ERROR_TLV_PROCESSOR(0x0C)
-#define CHIP_OTA_PROCESSOR_EXTERNAL_STORAGE CHIP_ERROR_TLV_PROCESSOR(0x0D)
-#define CHIP_OTA_PROCESSOR_START_IMAGE CHIP_ERROR_TLV_PROCESSOR(0x0E)
+#define CHIP_ERROR_OTA_CHANGE_PROCESSOR CHIP_ERROR_TLV_PROCESSOR(0x02)
+#define CHIP_ERROR_OTA_PROCESSOR_NOT_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x03)
+#define CHIP_ERROR_OTA_PROCESSOR_ALREADY_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x04)
+#define CHIP_ERROR_OTA_PROCESSOR_CLIENT_INIT CHIP_ERROR_TLV_PROCESSOR(0x05)
+#define CHIP_ERROR_OTA_PROCESSOR_MAKE_ROOM CHIP_ERROR_TLV_PROCESSOR(0x06)
+#define CHIP_ERROR_OTA_PROCESSOR_PUSH_CHUNK CHIP_ERROR_TLV_PROCESSOR(0x07)
+#define CHIP_ERROR_OTA_PROCESSOR_IMG_AUTH CHIP_ERROR_TLV_PROCESSOR(0x08)
+#define CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED CHIP_ERROR_TLV_PROCESSOR(0x09)
+#define CHIP_ERROR_OTA_PROCESSOR_IMG_COMMIT CHIP_ERROR_TLV_PROCESSOR(0x0A)
+#define CHIP_ERROR_OTA_PROCESSOR_CB_NOT_REGISTERED CHIP_ERROR_TLV_PROCESSOR(0x0B)
+#define CHIP_ERROR_OTA_PROCESSOR_EEPROM_OFFSET CHIP_ERROR_TLV_PROCESSOR(0x0C)
+#define CHIP_ERROR_OTA_PROCESSOR_EXTERNAL_STORAGE CHIP_ERROR_TLV_PROCESSOR(0x0D)
+#define CHIP_ERROR_OTA_PROCESSOR_START_IMAGE CHIP_ERROR_TLV_PROCESSOR(0x0E)
+#define CHIP_ERROR_OTA_PROCESSOR_DO_NOT_APPLY CHIP_ERROR_TLV_PROCESSOR(0x0F)
 
 // Descriptor constants
 constexpr size_t kVersionStringSize = 64;
@@ -77,13 +77,19 @@
 class OTATlvProcessor
 {
 public:
+    enum class ApplyState : uint8_t
+    {
+        kApply = 0,
+        kDoNotApply
+    };
+
     virtual ~OTATlvProcessor() {}
 
     virtual CHIP_ERROR Init()        = 0;
     virtual CHIP_ERROR Clear()       = 0;
-    virtual CHIP_ERROR ApplyAction() = 0;
     virtual CHIP_ERROR AbortAction() = 0;
     virtual CHIP_ERROR ExitAction() { return CHIP_NO_ERROR; }
+    virtual CHIP_ERROR ApplyAction();
 
     CHIP_ERROR Process(ByteSpan & block);
     void RegisterDescriptorCallback(ProcessDescriptor callback) { mCallbackProcessDescriptor = callback; }
@@ -102,7 +108,7 @@
      * If more image chunks are needed, CHIP_ERROR_BUFFER_TOO_SMALL error is returned.
      * Other error codes indicate that an error occurred during processing. Fetching
      * next data is scheduled automatically by OTAImageProcessorImpl if the return value
-     * is neither an error code, nor CHIP_OTA_FETCH_ALREADY_SCHEDULED (which implies the
+     * is neither an error code, nor CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED (which implies the
      * scheduling is done inside ProcessInternal or will be done in the future, through a
      * callback).
      *
@@ -110,13 +116,13 @@
      *              returns CHIP_NO_ERROR, the byte span is used to return a remaining part
      *              of the chunk, not used by current TLV processor.
      *
-     * @retval CHIP_NO_ERROR                    Block was processed successfully.
-     * @retval CHIP_ERROR_BUFFER_TOO_SMALL      Provided buffers are insufficient to decode some
-     *                                          metadata (e.g. a descriptor).
-     * @retval CHIP_OTA_FETCH_ALREADY_SCHEDULED Should be returned if ProcessInternal schedules
-     *                                          fetching next data (e.g. through a callback).
-     * @retval Error code                       Something went wrong. Current OTA process will be
-     *                                          canceled.
+     * @retval CHIP_NO_ERROR                          Block was processed successfully.
+     * @retval CHIP_ERROR_BUFFER_TOO_SMALL            Provided buffers are insufficient to decode some
+     *                                                metadata (e.g. a descriptor).
+     * @retval CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED Should be returned if ProcessInternal schedules
+     *                                                fetching next data (e.g. through a callback).
+     * @retval Error code                             Something went wrong. Current OTA process will be
+     *                                                canceled.
      */
     virtual CHIP_ERROR ProcessInternal(ByteSpan & block) = 0;
 
@@ -130,9 +136,23 @@
     /* Expected byte size of the OTAEncryptionKeyLength */
     static constexpr size_t kOTAEncryptionKeyLength = 16;
 #endif
-    uint32_t mLength                             = 0;
-    uint32_t mProcessedLength                    = 0;
-    bool mWasSelected                            = false;
+    uint32_t mLength          = 0;
+    uint32_t mProcessedLength = 0;
+    bool mWasSelected         = false;
+
+    /**
+     * @brief A flag to account for corner cases during OTA apply
+     *
+     * Used by the default ApplyAction implementation.
+     *
+     * If something goes wrong during ExitAction of the TLV processor,
+     * then mApplyState should be set to kDoNotApply and the image processor
+     * should abort. In this case, the BDX transfer was already finished
+     * and calling CancelImageUpdate will not abort the transfer, hence
+     * the device will reboot even though it should not have. If ApplyAction
+     * fails during HandleApply, then the process will be aborted.
+     */
+    ApplyState mApplyState                       = ApplyState::kApply;
     ProcessDescriptor mCallbackProcessDescriptor = nullptr;
 };
 
diff --git a/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.cpp b/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.cpp
index 760c9ef..f63aa74 100644
--- a/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.cpp
+++ b/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.cpp
@@ -28,12 +28,12 @@
 
 CHIP_ERROR OTAFirmwareProcessor::Init()
 {
-    ReturnErrorCodeIf(mCallbackProcessDescriptor == nullptr, CHIP_OTA_PROCESSOR_CB_NOT_REGISTERED);
+    ReturnErrorCodeIf(mCallbackProcessDescriptor == nullptr, CHIP_ERROR_OTA_PROCESSOR_CB_NOT_REGISTERED);
     mAccumulator.Init(sizeof(Descriptor));
 #if OTA_ENCRYPTION_ENABLE
     mUnalignmentNum = 0;
 #endif
-    ReturnErrorCodeIf(gOtaSuccess_c != OTA_ClientInit(), CHIP_OTA_PROCESSOR_CLIENT_INIT);
+    ReturnErrorCodeIf(gOtaSuccess_c != OTA_ClientInit(), CHIP_ERROR_OTA_PROCESSOR_CLIENT_INIT);
 
     auto offset = OTA_GetCurrentEepromAddressOffset();
     if (offset != 0)
@@ -41,8 +41,8 @@
         offset += 1;
     }
 
-    ReturnErrorCodeIf(OTA_UTILS_IMAGE_INVALID_ADDR == OTA_SetStartEepromOffset(offset), CHIP_OTA_PROCESSOR_EEPROM_OFFSET);
-    ReturnErrorCodeIf(gOtaSuccess_c != OTA_StartImage(mLength - sizeof(Descriptor)), CHIP_OTA_PROCESSOR_START_IMAGE);
+    ReturnErrorCodeIf(OTA_UTILS_IMAGE_INVALID_ADDR == OTA_SetStartEepromOffset(offset), CHIP_ERROR_OTA_PROCESSOR_EEPROM_OFFSET);
+    ReturnErrorCodeIf(gOtaSuccess_c != OTA_StartImage(mLength - sizeof(Descriptor)), CHIP_ERROR_OTA_PROCESSOR_START_IMAGE);
 
     return CHIP_NO_ERROR;
 }
@@ -95,7 +95,7 @@
     if (gOtaSuccess_c != status)
     {
         ChipLogError(SoftwareUpdate, "Failed to make room for next block. Status: %d", status);
-        return CHIP_OTA_PROCESSOR_MAKE_ROOM;
+        return CHIP_ERROR_OTA_PROCESSOR_MAKE_ROOM;
     }
 #if OTA_ENCRYPTION_ENABLE
     status = OTA_PushImageChunk((uint8_t *) mBlock.data(), (uint16_t) mBlock.size(), NULL, NULL);
@@ -105,10 +105,10 @@
     if (gOtaSuccess_c != status)
     {
         ChipLogError(SoftwareUpdate, "Failed to write image block. Status: %d", status);
-        return CHIP_OTA_PROCESSOR_PUSH_CHUNK;
+        return CHIP_ERROR_OTA_PROCESSOR_PUSH_CHUNK;
     }
 
-    return CHIP_OTA_FETCH_ALREADY_SCHEDULED;
+    return CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED;
 }
 
 CHIP_ERROR OTAFirmwareProcessor::ProcessDescriptor(ByteSpan & block)
@@ -122,12 +122,6 @@
     return CHIP_NO_ERROR;
 }
 
-CHIP_ERROR OTAFirmwareProcessor::ApplyAction()
-{
-
-    return CHIP_NO_ERROR;
-}
-
 CHIP_ERROR OTAFirmwareProcessor::AbortAction()
 {
     OTA_CancelImage();
@@ -144,13 +138,15 @@
     if (OTA_CommitImage(NULL) != gOtaSuccess_c)
     {
         ChipLogError(SoftwareUpdate, "Failed to commit firmware image.");
-        return CHIP_OTA_PROCESSOR_IMG_COMMIT;
+        mApplyState = ApplyState::kDoNotApply;
+        return CHIP_ERROR_OTA_PROCESSOR_IMG_COMMIT;
     }
 
     if (OTA_ImageAuthenticate() != gOtaImageAuthPass_c)
     {
         ChipLogError(SoftwareUpdate, "Failed to authenticate firmware image.");
-        return CHIP_OTA_PROCESSOR_IMG_AUTH;
+        mApplyState = ApplyState::kDoNotApply;
+        return CHIP_ERROR_OTA_PROCESSOR_IMG_AUTH;
     }
 
     OTA_AddNewImageFlag();
diff --git a/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.h b/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.h
index 885dc23..5a1092e 100644
--- a/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.h
+++ b/src/platform/nxp/k32w/k32w0/OTAFirmwareProcessor.h
@@ -35,7 +35,6 @@
 
     CHIP_ERROR Init() override;
     CHIP_ERROR Clear() override;
-    CHIP_ERROR ApplyAction() override;
     CHIP_ERROR AbortAction() override;
     CHIP_ERROR ExitAction() override;
 
diff --git a/src/platform/nxp/k32w/k32w1/OTAFirmwareProcessor.cpp b/src/platform/nxp/k32w/k32w1/OTAFirmwareProcessor.cpp
index 528261b..ee3ff9b 100644
--- a/src/platform/nxp/k32w/k32w1/OTAFirmwareProcessor.cpp
+++ b/src/platform/nxp/k32w/k32w1/OTAFirmwareProcessor.cpp
@@ -26,16 +26,16 @@
 
 CHIP_ERROR OTAFirmwareProcessor::Init()
 {
-    ReturnErrorCodeIf(mCallbackProcessDescriptor == nullptr, CHIP_OTA_PROCESSOR_CB_NOT_REGISTERED);
+    ReturnErrorCodeIf(mCallbackProcessDescriptor == nullptr, CHIP_ERROR_OTA_PROCESSOR_CB_NOT_REGISTERED);
     mAccumulator.Init(sizeof(Descriptor));
 
-    ReturnErrorCodeIf(gOtaSuccess_c != OTA_SelectExternalStoragePartition(), CHIP_OTA_PROCESSOR_EXTERNAL_STORAGE);
+    ReturnErrorCodeIf(gOtaSuccess_c != OTA_SelectExternalStoragePartition(), CHIP_ERROR_OTA_PROCESSOR_EXTERNAL_STORAGE);
 
     otaResult_t ota_status;
     ota_status = OTA_ServiceInit(&mPostedOperationsStorage[0], NB_PENDING_TRANSACTIONS * TRANSACTION_SZ);
 
-    ReturnErrorCodeIf(ota_status != gOtaSuccess_c, CHIP_OTA_PROCESSOR_CLIENT_INIT);
-    ReturnErrorCodeIf(gOtaSuccess_c != OTA_StartImage(mLength - sizeof(Descriptor)), CHIP_OTA_PROCESSOR_START_IMAGE);
+    ReturnErrorCodeIf(ota_status != gOtaSuccess_c, CHIP_ERROR_OTA_PROCESSOR_CLIENT_INIT);
+    ReturnErrorCodeIf(gOtaSuccess_c != OTA_StartImage(mLength - sizeof(Descriptor)), CHIP_ERROR_OTA_PROCESSOR_START_IMAGE);
 
     return CHIP_NO_ERROR;
 }
@@ -70,7 +70,7 @@
         if (gOtaSuccess_c != status)
         {
             ChipLogError(SoftwareUpdate, "Failed to make room for next block. Status: %d", status);
-            return CHIP_OTA_PROCESSOR_MAKE_ROOM;
+            return CHIP_ERROR_OTA_PROCESSOR_MAKE_ROOM;
         }
     }
     else
@@ -82,10 +82,10 @@
     if (gOtaSuccess_c != status)
     {
         ChipLogError(SoftwareUpdate, "Failed to write image block. Status: %d", status);
-        return CHIP_OTA_PROCESSOR_PUSH_CHUNK;
+        return CHIP_ERROR_OTA_PROCESSOR_PUSH_CHUNK;
     }
 
-    return CHIP_OTA_FETCH_ALREADY_SCHEDULED;
+    return CHIP_ERROR_OTA_FETCH_ALREADY_SCHEDULED;
 }
 
 CHIP_ERROR OTAFirmwareProcessor::ProcessDescriptor(ByteSpan & block)
@@ -117,7 +117,8 @@
     if (OTA_CommitImage(NULL) != gOtaSuccess_c)
     {
         ChipLogError(SoftwareUpdate, "Failed to commit firmware image.");
-        return CHIP_OTA_PROCESSOR_IMG_COMMIT;
+        mApplyState = ApplyState::kDoNotApply;
+        return CHIP_ERROR_OTA_PROCESSOR_IMG_COMMIT;
     }
 
     return CHIP_NO_ERROR;