Address review comments

src/darwin/Framework/CHIP/MTROTAImageTransferHandler.h

@@ -71,6 +71,7 @@ class MTROTAImageTransferHandler : public chip::bdx::AsyncResponder
 
  // The OTA provider delegate queue used by the controller.
  dispatch_queue_t mDelegateNotificationQueue = nil;
+
 };
 
 NS_ASSUME_NONNULL_END

src/darwin/Framework/CHIP/MTROTAImageTransferHandler.mm

@@ -32,10 +32,29 @@
 
 constexpr bdx::TransferRole kBdxRole = bdx::TransferRole::kSender;
 
+@interface MTROTAImageTransferHandlerWrapper : NSObject
+- (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler;
+@property(nonatomic, nullable, readwrite, assign) MTROTAImageTransferHandler * otaImageTransferHandler;
+@end
+
+@implementation MTROTAImageTransferHandlerWrapper
+
+- (instancetype)initWithMTROTAImageTransferHandler:(MTROTAImageTransferHandler *)otaImageTransferHandler
+{
+ if (self = [super init])
+ {
+ _otaImageTransferHandler = otaImageTransferHandler;
+ }
+ return self;
+}
+@end
+
+MTROTAImageTransferHandlerWrapper * mOTAImageTransferHandlerWrapper;
+
 MTROTAImageTransferHandler::MTROTAImageTransferHandler()
 {
- // Increment the number of delegates handling BDX by 1.
- MTROTAUnsolicitedBDXMessageHandler::IncrementNumberOfDelegates();
+ MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnDelegateCreated(this);
+ mOTAImageTransferHandlerWrapper = [[MTROTAImageTransferHandlerWrapper alloc] initWithMTROTAImageTransferHandler:this];
 }
 
 CHIP_ERROR MTROTAImageTransferHandler::Init(System::Layer * layer,
@@ -55,14 +74,15 @@
 
 MTROTAImageTransferHandler::~MTROTAImageTransferHandler()
 {
+ assertChipStackLockedByCurrentThread();
  mFabricIndex.ClearValue();
  mNodeId.ClearValue();
  mDelegate = nil;
  mDelegateNotificationQueue = nil;
  mSystemLayer = nil;
 
- // Decrement the number of delegates handling BDX by 1.
- MTROTAUnsolicitedBDXMessageHandler::DecrementNumberOfDelegates();
+ MTROTAUnsolicitedBDXMessageHandler::GetInstance()->OnDelegateDestroyed(this);
+ mOTAImageTransferHandlerWrapper.otaImageTransferHandler = nullptr;
 }
 
 CHIP_ERROR MTROTAImageTransferHandler::OnTransferSessionBegin(TransferSession::OutputEvent & event)
@@ -86,7 +106,16 @@
  auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mFabricIndex.Value()];
  VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);
 
+ MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;
+
  auto completionHandler = ^(NSError * _Nullable error) {
+
+ // Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
+ MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
+ if (!strongWrapper || !strongWrapper.otaImageTransferHandler)
+ {
+ return;
+ }
  [controller
  asyncDispatchToMatterQueue:^() {
  assertChipStackLockedByCurrentThread();
@@ -202,7 +231,17 @@
  auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mFabricIndex.Value()];
  VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);
 
+ MTROTAImageTransferHandlerWrapper * __weak weakWrapper = mOTAImageTransferHandlerWrapper;
+
  auto completionHandler = ^(NSData * _Nullable data, BOOL isEOF) {
+
+ // Check if the OTA image transfer handler is still valid. If not, return from the completion handler.
+ MTROTAImageTransferHandlerWrapper * strongWrapper = weakWrapper;
+ if (!strongWrapper || !strongWrapper.otaImageTransferHandler)
+ {
+ return;
+ }
+
  [controller
  asyncDispatchToMatterQueue:^() {
  assertChipStackLockedByCurrentThread();
@@ -337,6 +376,8 @@
 CHIP_ERROR MTROTAImageTransferHandler::OnMessageReceived(
  Messaging::ExchangeContext * ec, const PayloadHeader & payloadHeader, System::PacketBufferHandle && payload)
 {
+ assertChipStackLockedByCurrentThread();
+
  ChipLogProgress(BDX, "MTROTAImageTransferHandler: OnMessageReceived: message " ChipLogFormatMessageType " protocol " ChipLogFormatProtocolId,
  payloadHeader.GetMessageType(), ChipLogValueProtocolId(payloadHeader.GetProtocolID()));
 

src/darwin/Framework/CHIP/MTROTAProviderDelegateBridge.mm

@@ -65,7 +65,6 @@
 MTROTAProviderDelegateBridge::~MTROTAProviderDelegateBridge()
 {
  Clusters::OTAProvider::SetDelegate(kOtaProviderEndpoint, nullptr);
- Shutdown();
 }
 
 CHIP_ERROR MTROTAProviderDelegateBridge::Init(System::Layer * systemLayer, Messaging::ExchangeManager * exchangeManager)

src/darwin/Framework/CHIP/MTROTAUnsolicitedBDXMessageHandler.h

@@ -16,6 +16,7 @@
  */
 
 #import <Foundation/Foundation.h>
+#import "MTROTAImageTransferHandler.h"
 
 #include <lib/core/CHIPError.h>
 #include <messaging/ExchangeMgr.h>
@@ -38,20 +39,28 @@ class MTROTAUnsolicitedBDXMessageHandler : public chip::Messaging::UnsolicitedMe
  MTROTAUnsolicitedBDXMessageHandler()
  : mExchangeMgr(nullptr)
  {
+ sInstance = this;
  }
- ~MTROTAUnsolicitedBDXMessageHandler() { mExchangeMgr = nullptr; }
+
+ ~MTROTAUnsolicitedBDXMessageHandler() { mExchangeMgr = nullptr;}
+
+ static MTROTAUnsolicitedBDXMessageHandler * GetInstance();
 
  CHIP_ERROR Init(chip::Messaging::ExchangeManager * exchangeManager);
 
  // Returns the number of delegates that are currently handling BDX transfers.
  static uint8_t GetNumberOfDelegates();
 
- // Increase the number of delegates handling BDX transfers by 1.
+
  static void IncrementNumberOfDelegates();
 
  // Decrease the number of delegates handling BDX transfers by 1.
  static void DecrementNumberOfDelegates();
 
+ void OnDelegateCreated(void * imageTransferHandler);
+
+ void OnDelegateDestroyed(void * imageTransferHandler);
+
  void Shutdown();
 
  void ControllerShuttingDown(MTRDeviceController * controller);
@@ -62,10 +71,15 @@ class MTROTAUnsolicitedBDXMessageHandler : public chip::Messaging::UnsolicitedMe
 
  void OnExchangeCreationFailed(chip::Messaging::ExchangeDelegate * _Nonnull delegate) override;
 
-protected:
+ // TODO: #36181 - Have a set of MTROTAImageTransferHandler objects.
+ MTROTAImageTransferHandler * mOTAImageTransferHandler = nullptr;
+
  chip::Messaging::ExchangeManager * mExchangeMgr;
 
  static inline uint8_t mNumberOfDelegates = 0;
+
+ static MTROTAUnsolicitedBDXMessageHandler * sInstance;
+
 };
 
 NS_ASSUME_NONNULL_END

src/darwin/Framework/CHIP/MTROTAUnsolicitedBDXMessageHandler.mm

@@ -16,14 +16,15 @@
  */
 
 #import "MTROTAUnsolicitedBDXMessageHandler.h"
-#import "MTROTAImageTransferHandler.h"
 
 #include <protocols/bdx/BdxMessages.h>
 
 using namespace chip;
 using namespace chip::Messaging;
 using namespace chip::bdx;
 
+MTROTAUnsolicitedBDXMessageHandler * MTROTAUnsolicitedBDXMessageHandler::sInstance = nullptr;
+
 CHIP_ERROR MTROTAUnsolicitedBDXMessageHandler::Init(ExchangeManager * exchangeManager)
 {
  assertChipStackLockedByCurrentThread();
@@ -35,23 +36,35 @@
  return mExchangeMgr->RegisterUnsolicitedMessageHandlerForProtocol(Protocols::BDX::Id, this);
 }
 
+MTROTAUnsolicitedBDXMessageHandler * MTROTAUnsolicitedBDXMessageHandler::GetInstance()
+{
+ return sInstance;
+}
+
 void MTROTAUnsolicitedBDXMessageHandler::Shutdown()
 {
  assertChipStackLockedByCurrentThread();
 
- MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates = 0;
- VerifyOrReturn(mExchangeMgr != nullptr, CHIP_ERROR_INCORRECT_STATE);
+ VerifyOrReturn(mExchangeMgr != nullptr);
 
  mExchangeMgr->UnregisterUnsolicitedMessageHandlerForProtocol(Protocols::BDX::Id);
+ MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates = 0;
+
+ if (mOTAImageTransferHandler != nullptr)
+ {
+ mOTAImageTransferHandler->DestroySelf();
+ delete mOTAImageTransferHandler;
+ }
+
 }
 
 void MTROTAUnsolicitedBDXMessageHandler::ControllerShuttingDown(MTRDeviceController * controller)
 {
  assertChipStackLockedByCurrentThread();
 
  // Since the OTA provider delegate only supports one BDX transfer at a time, calling ShutDown is fine for now.
- // TODO: This class needs to keep a list of all MTROTAImageTransferHandlers mapped to the fabric index and only
- // delete the MTROTAImageTransferHandler with a fabric index matching the MTRDeviceController's fabric index.
+ // TODO: #36181 - This class needs to keep a list of all MTROTAImageTransferHandlers mapped to the fabric index and only
+ // delete the MTROTAImageTransferHandler objects with a fabric index matching the MTRDeviceController's fabric index.
  Shutdown();
 }
 
@@ -71,11 +84,12 @@
 
  // Only proceed if there is a valid fabric index for the SessionHandle.
  if (session->IsSecureSession() && session->AsSecureSession() != nullptr && session->AsSecureSession()->GetFabricIndex() != kUndefinedFabricIndex) {
+
  // If we receive a ReceiveInit BDX message, create a new MTROTAImageTransferHandler and register it
- // as the handler for all BDX messages that will come over this exchange and increment the number of delegates.
+ // as the handler for all BDX messages that will come over this exchange.
  if (payloadHeader.HasMessageType(MessageType::ReceiveInit)) {
- MTROTAImageTransferHandler * otaImageTransferHandler = new MTROTAImageTransferHandler();
- newDelegate = otaImageTransferHandler;
+ mOTAImageTransferHandler = new MTROTAImageTransferHandler();
+ newDelegate = mOTAImageTransferHandler;
  return CHIP_NO_ERROR;
  }
  }
@@ -84,15 +98,48 @@
 
 void MTROTAUnsolicitedBDXMessageHandler::OnExchangeCreationFailed(ExchangeDelegate * delegate)
 {
+ assertChipStackLockedByCurrentThread();
  auto * otaTransferHandler = static_cast<MTROTAImageTransferHandler *>(delegate);
+ VerifyOrReturn(otaTransferHandler != nullptr);
+
  otaTransferHandler->DestroySelf();
 }
 
 uint8_t MTROTAUnsolicitedBDXMessageHandler::GetNumberOfDelegates()
 {
+ assertChipStackLockedByCurrentThread();
  return MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates;
 }
 
-void MTROTAUnsolicitedBDXMessageHandler::IncrementNumberOfDelegates() { MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates++; }
+void MTROTAUnsolicitedBDXMessageHandler::IncrementNumberOfDelegates()
+{
+ assertChipStackLockedByCurrentThread();
+ MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates++;
+}
+
+void MTROTAUnsolicitedBDXMessageHandler::DecrementNumberOfDelegates()
+{
+ assertChipStackLockedByCurrentThread();
+ MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates--;
+}
+
+void MTROTAUnsolicitedBDXMessageHandler::OnDelegateCreated(void * imageTransferHandler)
+{
+ assertChipStackLockedByCurrentThread();
 
-void MTROTAUnsolicitedBDXMessageHandler::DecrementNumberOfDelegates() { MTROTAUnsolicitedBDXMessageHandler::mNumberOfDelegates--; }
+ // TODO: #36181 - Store the imageTransferHandler in the set of MTROTAImageTransferHandler objects.
+ mOTAImageTransferHandler = static_cast<MTROTAImageTransferHandler *>(imageTransferHandler);
+ MTROTAUnsolicitedBDXMessageHandler::IncrementNumberOfDelegates();
+}
+
+void MTROTAUnsolicitedBDXMessageHandler::OnDelegateDestroyed(void * imageTransferHandler)
+{
+ assertChipStackLockedByCurrentThread();
+
+ // TODO: #36181 - Remove the object matching imageTransferHandler in the set of MTROTAImageTransferHandler objects.
+ if (mOTAImageTransferHandler == static_cast<MTROTAImageTransferHandler *>(imageTransferHandler))
+ {
+ mOTAImageTransferHandler = nullptr;
+ }
+ MTROTAUnsolicitedBDXMessageHandler::DecrementNumberOfDelegates();
+}