Darwin: Ensure we tell the OTA delegate when a transfer ends (#34190)

* Darwin: Tweak OTA provider test to make it easier to extend

Use AssertIdentical for obejct identity tests
Fulfill expectations at the start of callback blocks before calling more methods

* Darwin: Ensure we tell the OTA delegate when a transfer ends

Ensure we call handleBDXTransferSessionEndForNodeID:... if we have called
handleBDXTransferSessionBeginForNodeID:..., even if the session is reset via an
error code path.

* Remove TODO clarified in review

* restyle

Author: ksperling-apple

src/darwin/Framework/CHIP/MTROTAProviderDelegateBridge.mm

@@ -104,7 +104,7 @@ CHIP_ERROR Shutdown()
         VerifyOrReturnError(mExchangeMgr != nullptr, CHIP_ERROR_INCORRECT_STATE);
 
         mExchangeMgr->UnregisterUnsolicitedMessageHandlerForProtocol(Protocols::BDX::Id);
-        ResetState();
+        ResetState(CHIP_ERROR_CANCELLED);
 
         mExchangeMgr = nullptr;
         mSystemLayer = nullptr;
@@ -117,18 +117,35 @@ void ControllerShuttingDown(MTRDeviceController * controller)
         assertChipStackLockedByCurrentThread();
 
         if (mInitialized && mFabricIndex.Value() == controller.fabricIndex) {
-            ResetState();
+            ResetState(CHIP_ERROR_CANCELLED);
         }
     }
 
-    void ResetState()
+    void ResetState(CHIP_ERROR error)
     {
         assertChipStackLockedByCurrentThread();
         if (mNodeId.HasValue() && mFabricIndex.HasValue()) {
             ChipLogProgress(Controller,
                 "Resetting state for OTA Provider; no longer providing an update for node id 0x" ChipLogFormatX64
                 ", fabric index %u",
                 ChipLogValueX64(mNodeId.Value()), mFabricIndex.Value());
+
+            if (mTransferStarted) {
+                auto controller = [MTRDeviceControllerFactory.sharedInstance runningControllerForFabricIndex:mFabricIndex.Value()];
+                if (controller) {
+                    auto nodeId = @(mNodeId.Value());
+                    auto strongDelegate = mDelegate;
+                    if ([strongDelegate respondsToSelector:@selector(handleBDXTransferSessionEndForNodeID:controller:error:)]) {
+                        dispatch_async(mDelegateNotificationQueue, ^{
+                            [strongDelegate handleBDXTransferSessionEndForNodeID:nodeId
+                                                                      controller:controller
+                                                                           error:[MTRError errorForCHIPErrorCode:error]];
+                        });
+                    }
+                } else {
+                    ChipLogError(Controller, "Not notifying delegate of BDX Transfer Session End, controller is not running");
+                }
+            }
         } else {
             ChipLogProgress(Controller, "Resetting state for OTA Provider");
         }
@@ -153,6 +170,7 @@ void ResetState()
         mDelegateNotificationQueue = nil;
 
         mInitialized = false;
+        mTransferStarted = false;
     }
 
 private:
@@ -162,7 +180,7 @@ void ResetState()
     static void HandleBdxInitReceivedTimeoutExpired(chip::System::Layer * systemLayer, void * state)
     {
         VerifyOrReturn(state != nullptr);
-        static_cast<BdxOTASender *>(state)->ResetState();
+        static_cast<BdxOTASender *>(state)->ResetState(CHIP_ERROR_TIMEOUT);
     }
 
     CHIP_ERROR OnMessageToSend(TransferSession::OutputEvent & event)
@@ -188,12 +206,12 @@ CHIP_ERROR OnMessageToSend(TransferSession::OutputEvent & event)
         if (err != CHIP_NO_ERROR) {
             mExchangeCtx->Close();
             mExchangeCtx = nullptr;
-            ResetState();
-        } else if (event.msgTypeData.HasMessageType(Protocols::SecureChannel::MsgType::StatusReport)) {
+            ResetState(err);
+        } else if (msgTypeData.HasMessageType(Protocols::SecureChannel::MsgType::StatusReport)) {
             // If the send was successful for a status report, since we are not expecting a response the exchange context is
             // already closed. We need to null out the reference to avoid having a dangling pointer.
             mExchangeCtx = nullptr;
-            ResetState();
+            ResetState(CHIP_ERROR_INTERNAL);
         }
         return err;
     }
@@ -258,8 +276,8 @@ CHIP_ERROR OnTransferSessionBegin(TransferSession::OutputEvent & event)
                               }];
         };
 
+        mTransferStarted = true;
         auto nodeId = @(mNodeId.Value());
-
         auto strongDelegate = mDelegate;
         dispatch_async(mDelegateNotificationQueue, ^{
             if ([strongDelegate respondsToSelector:@selector(handleBDXTransferSessionBeginForNodeID:controller:fileDesignator:offset:completion:)]) {
@@ -294,20 +312,7 @@ CHIP_ERROR OnTransferSessionEnd(TransferSession::OutputEvent & event)
             error = CHIP_ERROR_INTERNAL;
         }
 
-        auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:mFabricIndex.Value()];
-        VerifyOrReturnError(controller != nil, CHIP_ERROR_INCORRECT_STATE);
-        auto nodeId = @(mNodeId.Value());
-
-        auto strongDelegate = mDelegate;
-        if ([strongDelegate respondsToSelector:@selector(handleBDXTransferSessionEndForNodeID:controller:error:)]) {
-            dispatch_async(mDelegateNotificationQueue, ^{
-                [strongDelegate handleBDXTransferSessionEndForNodeID:nodeId
-                                                          controller:controller
-                                                               error:[MTRError errorForCHIPErrorCode:error]];
-            });
-        }
-
-        ResetState();
+        ResetState(error); // will notify the delegate
         return CHIP_NO_ERROR;
     }
 
@@ -438,7 +443,7 @@ CHIP_ERROR ConfigureState(chip::FabricIndex fabricIndex, chip::NodeId nodeId)
             VerifyOrReturnError(mFabricIndex.Value() == fabricIndex && mNodeId.Value() == nodeId, CHIP_ERROR_BUSY);
 
             // Reset stale connection from the same Node if exists.
-            ResetState();
+            ResetState(CHIP_ERROR_CANCELLED);
         }
 
         auto * controller = [[MTRDeviceControllerFactory sharedInstance] runningControllerForFabricIndex:fabricIndex];
@@ -466,6 +471,7 @@ CHIP_ERROR ConfigureState(chip::FabricIndex fabricIndex, chip::NodeId nodeId)
     }
 
     bool mInitialized = false;
+    bool mTransferStarted = false;
     Optional<FabricIndex> mFabricIndex;
     Optional<NodeId> mNodeId;
     id<MTROTAProviderDelegate> mDelegate = nil;
@@ -489,7 +495,7 @@ CHIP_ERROR ConfigureState(chip::FabricIndex fabricIndex, chip::NodeId nodeId)
 
 MTROTAProviderDelegateBridge::~MTROTAProviderDelegateBridge()
 {
-    gOtaSender->ResetState();
+    gOtaSender->ResetState(CHIP_ERROR_CANCELLED);
     Clusters::OTAProvider::SetDelegate(kOtaProviderEndpoint, nullptr);
 }
 
@@ -685,7 +691,7 @@ bool GetPeerNodeInfo(CommandHandler * commandHandler, const ConcreteCommandPath
                     handle.Release();
                     // We need to reset state here to clean up any initialization we might have done including starting the BDX
                     // timeout timer while preparing for transfer if any failure occurs afterwards.
-                    gOtaSender->ResetState();
+                    gOtaSender->ResetState(err);
                     return;
                 }
 
@@ -696,7 +702,7 @@ bool GetPeerNodeInfo(CommandHandler * commandHandler, const ConcreteCommandPath
                     LogErrorOnFailure(err);
                     handler->AddStatus(cachedCommandPath, StatusIB(err).mStatus);
                     handle.Release();
-                    gOtaSender->ResetState();
+                    gOtaSender->ResetState(err);
                     return;
                 }
                 delegateResponse.imageURI.SetValue(uri);

src/darwin/Framework/CHIPTests/MTROTAProviderTests.m

@@ -818,13 +818,15 @@ - (void)test003_ReceiveQueryImageRequestWhileHandlingBDX_RespondImplicitBusy
     XCTestExpectation * queryExpectation1 = [self expectationWithDescription:@"handleQueryImageForNodeID called first time"];
     XCTestExpectation * queryExpectation2 = [self expectationWithDescription:@"handleQueryImageForNodeID called second time"];
     XCTestExpectation * queryExpectation3 = [self expectationWithDescription:@"handleQueryImageForNodeID called third time"];
+    XCTestExpectation * transferEndExpectation = [self expectationWithDescription:@"handleBDXTransferSessionEndForNodeID called"];
 
     const uint16_t busyDelay = 1; // 1 second
     NSString * fakeImageURI = @"No such image, really";
 
     __block QueryImageHandler handleThirdQuery;
     sOTAProviderDelegate.queryImageHandler = ^(NSNumber * nodeID, MTRDeviceController * controller,
         MTROTASoftwareUpdateProviderClusterQueryImageParams * params, QueryImageCompletion completion) {
+        [queryExpectation1 fulfill];
         XCTAssertEqualObjects(nodeID, @(kDeviceId1));
         XCTAssertEqual(controller, sController);
         [sOTAProviderDelegate respondAvailableWithDelay:@(0)
@@ -833,23 +835,32 @@ - (void)test003_ReceiveQueryImageRequestWhileHandlingBDX_RespondImplicitBusy
                                         softwareVersion:kUpdatedSoftwareVersion_5
                                   softwareVersionString:kUpdatedSoftwareVersionString_5
                                              completion:completion];
-        [queryExpectation1 fulfill];
     };
     sOTAProviderDelegate.transferBeginHandler = ^(NSNumber * nodeID, MTRDeviceController * controller, NSString * fileDesignator,
         NSNumber * offset, MTRStatusCompletion outerCompletion) {
+        sOTAProviderDelegate.transferBeginHandler = nil;
+        // Now that we've begun a transfer, we expect to be told when it ends, even if it's due to an error
+        sOTAProviderDelegate.transferEndHandler = ^(NSNumber * nodeID, MTRDeviceController * controller, NSError * _Nullable error) {
+            [transferEndExpectation fulfill];
+            sOTAProviderDelegate.transferEndHandler = nil;
+            XCTAssertEqualObjects(nodeID, @(kDeviceId1));
+            XCTAssertIdentical(controller, sController);
+            XCTAssertNotNil(error); // we cancelled the transfer, so there should be an error
+        };
+
         XCTAssertEqualObjects(nodeID, @(kDeviceId1));
-        XCTAssertEqual(controller, sController);
+        XCTAssertIdentical(controller, sController);
 
         // Don't actually respond until the second requestor has queried us for
         // an image.  We need to reset queryImageHandler here, so we can close
         // over outerCompletion.
         sOTAProviderDelegate.queryImageHandler = ^(NSNumber * nodeID, MTRDeviceController * controller,
             MTROTASoftwareUpdateProviderClusterQueryImageParams * params, QueryImageCompletion innerCompletion) {
+            [queryExpectation2 fulfill];
             sOTAProviderDelegate.queryImageHandler = handleThirdQuery;
-            sOTAProviderDelegate.transferBeginHandler = nil;
 
             XCTAssertEqualObjects(nodeID, @(kDeviceId2));
-            XCTAssertEqual(controller, sController);
+            XCTAssertIdentical(controller, sController);
 
             // We respond UpdateAvailable, but since we are in the middle of
             // handling OTA for device1 we expect the requestor to get Busy and
@@ -860,28 +871,29 @@ - (void)test003_ReceiveQueryImageRequestWhileHandlingBDX_RespondImplicitBusy
                                             softwareVersion:kUpdatedSoftwareVersion_5
                                       softwareVersionString:kUpdatedSoftwareVersionString_5
                                                  completion:innerCompletion];
+
+            // Cancel the transfer with device1
             [sOTAProviderDelegate respondErrorWithCompletion:outerCompletion];
-            [queryExpectation2 fulfill];
         };
 
         announceResponseExpectation2 = [self announceProviderToDevice:device2];
     };
 
     handleThirdQuery = ^(NSNumber * nodeID, MTRDeviceController * controller,
         MTROTASoftwareUpdateProviderClusterQueryImageParams * params, QueryImageCompletion completion) {
+        [queryExpectation3 fulfill];
         XCTAssertEqualObjects(nodeID, @(kDeviceId2));
-        XCTAssertEqual(controller, sController);
+        XCTAssertIdentical(controller, sController);
 
         [sOTAProviderDelegate respondNotAvailableWithCompletion:completion];
-        [queryExpectation3 fulfill];
     };
 
     // Advertise ourselves as an OTA provider.
     XCTestExpectation * announceResponseExpectation1 = [self announceProviderToDevice:device1];
 
     // Make sure we get our queries in order.  Give it a bit more time, because
     // there will be a delay between the two queries.
-    [self waitForExpectations:@[ queryExpectation1, queryExpectation2, queryExpectation3 ]
+    [self waitForExpectations:@[ queryExpectation1, queryExpectation2, transferEndExpectation, queryExpectation3 ]
                       timeout:(kTimeoutInSeconds + busyDelay * 3)
                  enforceOrder:YES];