Re-fix OTA exchange leaks (#20632)

* [ota] Fix exchange context leak in OTA requestor (#20304) * [ota] Fix exchange context leak It turns out that Exchange Context allocated for BDX transfer is not released on completion or connection abort. It is not seen in a happy path that results in applying and rebooting into a new firmware, but may lead to the exchange leak when the transfer is interrupted. Furthermore, if an exchange is never released, a Sleepy End Device never returns to the idle mode, needlessly draining the battery. Signed-off-by: Damian Krolik <[email protected]> * Restyled by clang-format Co-authored-by: Restyled.io <[email protected]> * Fix exchange lifetime management in BDXMessenger. * Address review comments. Co-authored-by: Damian Królik <[email protected]> Co-authored-by: Restyled.io <[email protected]>
project-chip · Oct 19, 2023 · 1433536 · 1433536
1 parent 15ca724
commit 1433536
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 17 deletions.
diff --git a/src/app/clusters/ota-requestor/BDXDownloader.cpp b/src/app/clusters/ota-requestor/BDXDownloader.cpp
@@ -201,10 +201,12 @@ void BDXDownloader::EndDownload(CHIP_ERROR reason)
         {
             mImageProcessor->Abort();
         }
-        SetState(State::kIdle, OTAChangeReasonEnum::kSuccess);
 
         // Because AbortTransfer() will generate a StatusReport to send.
         PollTransferSession();
+
+        // Now that we've sent our report, we're idle.
+        SetState(State::kIdle, OTAChangeReasonEnum::kSuccess);
     }
     else
     {

diff --git a/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp b/src/app/clusters/ota-requestor/DefaultOTARequestor.cpp
@@ -638,12 +638,14 @@ void DefaultOTARequestor::OnDownloadStateChanged(OTADownloader::State state, OTA
     {
     case OTADownloader::State::kComplete:
         mOtaRequestorDriver->UpdateDownloaded();
+        mBdxMessenger.Reset();
         break;
     case OTADownloader::State::kIdle:
         if (reason != OTAChangeReasonEnum::kSuccess)
         {
             RecordErrorUpdateState(CHIP_ERROR_CONNECTION_ABORTED, reason);
         }
+        mBdxMessenger.Reset();
         break;
     default:
         break;
@@ -823,15 +825,25 @@ CHIP_ERROR DefaultOTARequestor::StartDownload(OperationalDeviceProxy & devicePro
     Optional<SessionHandle> session = deviceProxy.GetSecureSession();
     VerifyOrReturnError(session.HasValue(), CHIP_ERROR_INCORRECT_STATE);
 
-    mExchangeCtx = exchangeMgr->NewContext(session.Value(), &mBdxMessenger);
-    VerifyOrReturnError(mExchangeCtx != nullptr, CHIP_ERROR_NO_MEMORY);
+    chip::Messaging::ExchangeContext * exchangeCtx = exchangeMgr->NewContext(session.Value(), &mBdxMessenger);
+    VerifyOrReturnError(exchangeCtx != nullptr, CHIP_ERROR_NO_MEMORY);
 
-    mBdxMessenger.Init(mBdxDownloader, mExchangeCtx);
+    mBdxMessenger.Init(mBdxDownloader, exchangeCtx);
     mBdxDownloader->SetMessageDelegate(&mBdxMessenger);
     mBdxDownloader->SetStateDelegate(this);
 
-    ReturnErrorOnFailure(mBdxDownloader->SetBDXParams(initOptions, kDownloadTimeoutSec));
-    return mBdxDownloader->BeginPrepareDownload();
+    CHIP_ERROR err = mBdxDownloader->SetBDXParams(initOptions, kDownloadTimeoutSec);
+    if (err == CHIP_NO_ERROR)
+    {
+        err = mBdxDownloader->BeginPrepareDownload();
+    }
+
+    if (err != CHIP_NO_ERROR)
+    {
+        mBdxMessenger.Reset();
+    }
+
+    return err;
 }
 
 CHIP_ERROR DefaultOTARequestor::SendApplyUpdateRequest(OperationalDeviceProxy & deviceProxy)

diff --git a/src/app/clusters/ota-requestor/DefaultOTARequestor.h b/src/app/clusters/ota-requestor/DefaultOTARequestor.h
@@ -136,9 +136,14 @@ class DefaultOTARequestor : public OTARequestorInterface, public BDXDownloader::
             {
                 sendFlags.Set(chip::Messaging::SendMessageFlags::kExpectResponse);
             }
-            ReturnErrorOnFailure(mExchangeCtx->SendMessage(event.msgTypeData.ProtocolId, event.msgTypeData.MessageType,
-                                                           event.MsgData.Retain(), sendFlags));
-            return CHIP_NO_ERROR;
+            CHIP_ERROR err = mExchangeCtx->SendMessage(event.msgTypeData.ProtocolId, event.msgTypeData.MessageType,
+                                                       event.MsgData.Retain(), sendFlags);
+            if (err != CHIP_NO_ERROR)
+            {
+                Reset();
+            }
+
+            return err;
         }
 
         CHIP_ERROR OnMessageReceived(chip::Messaging::ExchangeContext * ec, const chip::PayloadHeader & payloadHeader,
@@ -150,13 +155,14 @@ class DefaultOTARequestor : public OTARequestorInterface, public BDXDownloader::
                 return CHIP_NO_ERROR;
             }
 
-            mDownloader->OnMessageReceived(payloadHeader, payload.Retain());
+            mDownloader->OnMessageReceived(payloadHeader, std::move(payload));
 
             // For a receiver using BDX Protocol, all received messages will require a response except for a StatusReport
             if (!payloadHeader.HasMessageType(chip::Protocols::SecureChannel::MsgType::StatusReport))
             {
                 ec->WillSendMessage();
             }
+
             return CHIP_NO_ERROR;
         }
 
@@ -169,12 +175,21 @@ class DefaultOTARequestor : public OTARequestorInterface, public BDXDownloader::
             }
         }
 
+        void OnExchangeClosing(Messaging::ExchangeContext * ec) override { mExchangeCtx = nullptr; }
+
         void Init(chip::BDXDownloader * downloader, chip::Messaging::ExchangeContext * ec)
         {
             mExchangeCtx = ec;
             mDownloader  = downloader;
         }
 
+        void Reset()
+        {
+            VerifyOrReturn(mExchangeCtx != nullptr);
+            mExchangeCtx->Close();
+            mExchangeCtx = nullptr;
+        }
+
     private:
         chip::Messaging::ExchangeContext * mExchangeCtx;
         chip::BDXDownloader * mDownloader;
@@ -292,13 +307,12 @@ class DefaultOTARequestor : public OTARequestorInterface, public BDXDownloader::
      */
     static void OnCommissioningCompleteRequestor(const DeviceLayer::ChipDeviceEvent * event, intptr_t arg);
 
-    OTARequestorStorage * mStorage            = nullptr;
-    OTARequestorDriver * mOtaRequestorDriver  = nullptr;
-    CASESessionManager * mCASESessionManager  = nullptr;
-    OnConnectedAction mOnConnectedAction      = kQueryImage;
-    Messaging::ExchangeContext * mExchangeCtx = nullptr;
-    BDXDownloader * mBdxDownloader            = nullptr; // TODO: this should be OTADownloader
-    BDXMessenger mBdxMessenger;                          // TODO: ideally this is held by the application
+    OTARequestorStorage * mStorage           = nullptr;
+    OTARequestorDriver * mOtaRequestorDriver = nullptr;
+    CASESessionManager * mCASESessionManager = nullptr;
+    OnConnectedAction mOnConnectedAction     = kQueryImage;
+    BDXDownloader * mBdxDownloader           = nullptr; // TODO: this should be OTADownloader
+    BDXMessenger mBdxMessenger;                         // TODO: ideally this is held by the application
     uint8_t mUpdateTokenBuffer[kMaxUpdateTokenLen];
     ByteSpan mUpdateToken;
     uint32_t mCurrentVersion = 0;