From 872a1a682ada0eb928e76694bb342d19765b1f64 Mon Sep 17 00:00:00 2001
From: Carol Yang <clyang@apple.com>
Date: Thu, 3 Mar 2022 21:32:21 -0800
Subject: [PATCH] [OTA] Set the current state on every attempt for CASE session
 establishment

---
 .../clusters/ota-requestor/OTARequestor.cpp   | 23 +++++++------------
 src/app/clusters/ota-requestor/OTARequestor.h |  6 ++---
 src/include/platform/OTARequestorInterface.h  |  4 ++--
 3 files changed, 13 insertions(+), 20 deletions(-)

diff --git a/src/app/clusters/ota-requestor/OTARequestor.cpp b/src/app/clusters/ota-requestor/OTARequestor.cpp
index e66052813d5363..91f62a1276a1f7 100644
--- a/src/app/clusters/ota-requestor/OTARequestor.cpp
+++ b/src/app/clusters/ota-requestor/OTARequestor.cpp
@@ -385,11 +385,9 @@ void OTARequestor::OnConnected(void * context, OperationalDeviceProxy * devicePr
             requestorCore->RecordErrorUpdateState(UpdateFailureState::kQuerying, err);
             return;
         }
-
-        // The kQuerying state is set in TriggerImmediateQueryInternal(), do not set it here
         break;
     }
-    case kStartBDX: {
+    case kDownload: {
         CHIP_ERROR err = requestorCore->StartDownload(*deviceProxy);
 
         if (err != CHIP_NO_ERROR)
@@ -398,8 +396,6 @@ void OTARequestor::OnConnected(void * context, OperationalDeviceProxy * devicePr
             requestorCore->RecordErrorUpdateState(UpdateFailureState::kDownloading, err);
             return;
         }
-
-        requestorCore->RecordNewUpdateState(OTAUpdateStateEnum::kDownloading, OTAChangeReasonEnum::kSuccess);
         break;
     }
     case kApplyUpdate: {
@@ -411,8 +407,6 @@ void OTARequestor::OnConnected(void * context, OperationalDeviceProxy * devicePr
             requestorCore->RecordErrorUpdateState(UpdateFailureState::kApplying, err);
             return;
         }
-
-        requestorCore->RecordNewUpdateState(OTAUpdateStateEnum::kApplying, OTAChangeReasonEnum::kSuccess);
         break;
     }
     case kNotifyUpdateApplied: {
@@ -424,8 +418,6 @@ void OTARequestor::OnConnected(void * context, OperationalDeviceProxy * devicePr
             requestorCore->RecordErrorUpdateState(UpdateFailureState::kNotifying, err);
             return;
         }
-
-        requestorCore->RecordNewUpdateState(OTAUpdateStateEnum::kIdle, OTAChangeReasonEnum::kSuccess);
         break;
     }
     default:
@@ -447,7 +439,7 @@ void OTARequestor::OnConnectionFailure(void * context, PeerId peerId, CHIP_ERROR
     case kQueryImage:
         requestorCore->RecordErrorUpdateState(UpdateFailureState::kQuerying, error);
         break;
-    case kStartBDX:
+    case kDownload:
         requestorCore->RecordErrorUpdateState(UpdateFailureState::kDownloading, error);
         break;
     case kApplyUpdate:
@@ -483,10 +475,6 @@ OTARequestorInterface::OTATriggerResult OTARequestor::TriggerImmediateQuery()
 
     SetCurrentProviderLocation(providerLocation);
 
-    // We are now querying a provider, leave the kIdle state.
-    // No state matches this one fully but we can't be in kIdle.
-    RecordNewUpdateState(OTAUpdateStateEnum::kQuerying, OTAChangeReasonEnum::kSuccess);
-
     // Go through the driver as it has additional logic to execute
     mOtaRequestorDriver->SendQueryImage();
 
@@ -495,11 +483,13 @@ OTARequestorInterface::OTATriggerResult OTARequestor::TriggerImmediateQuery()
 
 void OTARequestor::DownloadUpdate()
 {
-    ConnectToProvider(kStartBDX);
+    RecordNewUpdateState(OTAUpdateStateEnum::kDownloading, OTAChangeReasonEnum::kSuccess);
+    ConnectToProvider(kDownload);
 }
 
 void OTARequestor::ApplyUpdate()
 {
+    RecordNewUpdateState(OTAUpdateStateEnum::kApplying, OTAChangeReasonEnum::kSuccess);
     ConnectToProvider(kApplyUpdate);
 }
 
@@ -519,6 +509,9 @@ void OTARequestor::NotifyUpdateApplied(uint32_t version)
 
     OtaRequestorServerOnVersionApplied(version, productId);
 
+    // There is no response for a notify so consider this OTA complete
+    RecordNewUpdateState(OTAUpdateStateEnum::kIdle, OTAChangeReasonEnum::kSuccess);
+
     ConnectToProvider(kNotifyUpdateApplied);
 }
 
diff --git a/src/app/clusters/ota-requestor/OTARequestor.h b/src/app/clusters/ota-requestor/OTARequestor.h
index 13de1b351a3494..dad5af7d008162 100644
--- a/src/app/clusters/ota-requestor/OTARequestor.h
+++ b/src/app/clusters/ota-requestor/OTARequestor.h
@@ -54,10 +54,10 @@ class OTARequestor : public OTARequestorInterface, public BDXDownloader::StateDe
     // Initiate download of the new image
     void DownloadUpdate() override;
 
-    // Send ApplyImage
+    // Initiate the session to send ApplyUpdateRequest command
     void ApplyUpdate() override;
 
-    // Send NotifyUpdateApplied, update Basic cluster SoftwareVersion attribute, log the VersionApplied event
+    // Initiate the session to send NotifyUpdateApplied command
     void NotifyUpdateApplied(uint32_t version) override;
 
     // Get image update progress in percents unit
@@ -230,7 +230,7 @@ class OTARequestor : public OTARequestorInterface, public BDXDownloader::StateDe
     enum OnConnectedAction
     {
         kQueryImage = 0,
-        kStartBDX,
+        kDownload,
         kApplyUpdate,
         kNotifyUpdateApplied,
     };
diff --git a/src/include/platform/OTARequestorInterface.h b/src/include/platform/OTARequestorInterface.h
index 17c20ddb84957b..6fb5853e124677 100644
--- a/src/include/platform/OTARequestorInterface.h
+++ b/src/include/platform/OTARequestorInterface.h
@@ -182,10 +182,10 @@ class OTARequestorInterface
     // Download image
     virtual void DownloadUpdate() = 0;
 
-    // Send ApplyImage command
+    // Initiate the session to send ApplyUpdateRequest command
     virtual void ApplyUpdate() = 0;
 
-    // Send NotifyUpdateApplied command
+    // Initiate the session to send NotifyUpdateApplied command
     virtual void NotifyUpdateApplied(uint32_t version) = 0;
 
     // Get image update progress in percents unit