[Darwin] Enable taking assertion on MTRDeviceController

src/darwin/Framework/CHIP/MTRDeviceController.mm

@@ -46,6 +46,7 @@
 #import "MTRSetupPayload.h"
 #import "MTRTimeUtils.h"
 #import "MTRUnfairLock.h"
+#import "MTRUtilities.h"
 #import "NSDataSpanConversion.h"
 #import "NSStringSpanConversion.h"
 #import <setup_payload/ManualSetupPayloadGenerator.h>
@@ -129,6 +130,11 @@ @implementation MTRDeviceController {
     std::atomic<std::optional<uint64_t>> _storedCompressedFabricID;
     MTRP256KeypairBridge _signingKeypairBridge;
     MTRP256KeypairBridge _operationalKeypairBridge;
+
+    // Counters to track assertion status and access controlled by the _assertionLock
+    NSUInteger _keepRunningAssertionCounter;
+    BOOL _shutdownPending;
+    os_unfair_lock _assertionLock;
 }
 
 - (os_unfair_lock_t)deviceMapLock
@@ -142,6 +148,11 @@ - (instancetype)initForSubclasses
         // nothing, as superclass of MTRDeviceController is NSObject
     }
     _underlyingDeviceMapLock = OS_UNFAIR_LOCK_INIT;
+
+    // Setup assertion variables
+    _keepRunningAssertionCounter = 0;
+    _shutdownPending = NO;
+    _assertionLock = OS_UNFAIR_LOCK_INIT;
     return self;
 }
 
@@ -178,6 +189,12 @@ - (instancetype)initWithFactory:(MTRDeviceControllerFactory *)factory
         // Make sure our storage is all set up to work as early as possible,
         // before we start doing anything else with the controller.
         _uniqueIdentifier = uniqueIdentifier;
+
+        // Setup assertion variables
+        _keepRunningAssertionCounter = 0;
+        _shutdownPending = NO;
+        _assertionLock = OS_UNFAIR_LOCK_INIT;
+
         if (storageDelegate != nil) {
             if (storageDelegateQueue == nil) {
                 MTR_LOG_ERROR("storageDelegate provided without storageDelegateQueue");
@@ -295,6 +312,9 @@ - (instancetype)initWithFactory:(MTRDeviceControllerFactory *)factory
 
         _storedFabricIndex = chip::kUndefinedFabricIndex;
         _storedCompressedFabricID = std::nullopt;
+        self.nodeID = nil;
+        self.fabricID = nil;
+        self.rootPublicKey = nil;
 
         _storageBehaviorConfiguration = storageBehaviorConfiguration;
     }
@@ -311,8 +331,68 @@ - (BOOL)isRunning
     return _cppCommissioner != nullptr;
 }
 
+- (BOOL)matchesPendingShutdownWithParams:(MTRDeviceControllerParameters *)parameters
+{
+    if (!parameters.operationalCertificate || !parameters.rootCertificate) {
+        return FALSE;
+    }
+    NSNumber * nodeID = [MTRDeviceControllerParameters nodeIDFromNOC:parameters.operationalCertificate];
+    NSNumber * fabricID = [MTRDeviceControllerParameters fabricIDFromNOC:parameters.operationalCertificate];
+    NSData * publicKey = [MTRDeviceControllerParameters publicKeyFromCertificate:parameters.rootCertificate];
+
+    std::lock_guard lock(_assertionLock);
+
+    return _keepRunningAssertionCounter > 0 && _shutdownPending && MTREqualObjects(nodeID, self.nodeID) && MTREqualObjects(fabricID, self.fabricID) && MTREqualObjects(publicKey, self.rootPublicKey);
+}
+
+- (void)addRunAssertion
+{
+    std::lock_guard lock(_assertionLock);
+
+    // Only take an assertion if running
+    if ([self isRunning]) {
+        ++_keepRunningAssertionCounter;
+        MTR_LOG("%@ Adding keep running assertion, total %lu", self, (unsigned long) _keepRunningAssertionCounter);
+    }
+}
+
+- (void)removeRunAssertion;
+{
+    std::lock_guard lock(_assertionLock);
+
+    if (_keepRunningAssertionCounter > 0) {
+        --_keepRunningAssertionCounter;
+        MTR_LOG("%@ Removing keep running assertion, total %lu", self, (unsigned long) _keepRunningAssertionCounter);
+
+        if ([self isRunning] && _keepRunningAssertionCounter == 0 && _shutdownPending) {
+            MTR_LOG("%@ All assertions removed and shutdown is pending, shutting down", self);
+            [self finalShutdown];
+        }
+    }
+}
+
+- (void)clearPendingShutdown
+{
+    std::lock_guard lock(_assertionLock);
+    _shutdownPending = NO;
+}
+
 - (void)shutdown
 {
+    std::lock_guard lock(_assertionLock);
+
+    if (_keepRunningAssertionCounter > 0) {
+        MTR_LOG("%@ Pending shutdown since %lu assertions are present", self, (unsigned long) _keepRunningAssertionCounter);
+        _shutdownPending = YES;
+        return;
+    }
+    [self finalShutdown];
+}
+
+- (void)finalShutdown
+{
+    os_unfair_lock_assert_owner(&_assertionLock);
+
     MTR_LOG("%@ shutdown called", self);
     if (_cppCommissioner == nullptr) {
         // Already shut down.
@@ -369,11 +449,16 @@ - (void)shutDownCppController
         // shuts down.
         _storedFabricIndex = chip::kUndefinedFabricIndex;
         _storedCompressedFabricID = std::nullopt;
+        self.nodeID = nil;
+        self.fabricID = nil;
+        self.rootPublicKey = nil;
+
         delete commissionerToShutDown;
         if (_operationalCredentialsDelegate != nil) {
             _operationalCredentialsDelegate->SetDeviceCommissioner(nullptr);
         }
     }
+    _shutdownPending = NO;
 }
 
 - (void)deinitFromFactory
@@ -609,6 +694,14 @@ - (BOOL)startup:(MTRDeviceControllerStartupParamsInternal *)startupParams
 
         self->_storedFabricIndex = fabricIdx;
         self->_storedCompressedFabricID = _cppCommissioner->GetCompressedFabricId();
+
+        chip::Crypto::P256PublicKey rootPublicKey;
+        if (_cppCommissioner->GetRootPublicKey(rootPublicKey) == CHIP_NO_ERROR) {
+            self.rootPublicKey = [NSData dataWithBytes:rootPublicKey.Bytes() length:rootPublicKey.Length()];
+            self.nodeID = @(_cppCommissioner->GetNodeId());
+            self.fabricID = @(_cppCommissioner->GetFabricId());
+        }
+
         commissionerInitialized = YES;
 
         MTR_LOG("%@ startup succeeded for nodeID 0x%016llX", self, self->_cppCommissioner->GetNodeId());

src/darwin/Framework/CHIP/MTRDeviceControllerFactory.mm

@@ -1133,12 +1133,31 @@ - (void)operationalInstanceAdded:(chip::PeerId &)operationalID
     }
 }
 
+- (nullable MTRDeviceController *)_findControllerMatchingParams:(MTRDeviceControllerParameters *)parameters
+{
+    std::lock_guard lock(_controllersLock);
+    for (MTRDeviceController * controller in _controllers) {
+        if ([controller matchesPendingShutdownWithParams:parameters]) {
+            MTR_LOG("%@ Found existing controller %@ that is pending shutdown and matching parameters, re-using it", self, controller);
+            [controller clearPendingShutdown];
+            return controller;
+        }
+    }
+    return nil;
+}
+
 - (nullable MTRDeviceController *)initializeController:(MTRDeviceController *)controller
                                         withParameters:(MTRDeviceControllerParameters *)parameters
                                                  error:(NSError * __autoreleasing *)error
 {
     [self _assertCurrentQueueIsNotMatterQueue];
 
+    // If there is a controller already running with matching parameters, re-use it
+    MTRDeviceController * existingController = [self _findControllerMatchingParams:parameters];
+    if (existingController) {
+        return existingController;
+    }
+
     return [self _startDeviceController:controller
                           startupParams:parameters
                           fabricChecker:^MTRDeviceControllerStartupParamsInternal *(

src/darwin/Framework/CHIP/MTRDeviceControllerStartupParams.mm

@@ -306,6 +306,29 @@ - (void)setOTAProviderDelegate:(id<MTROTAProviderDelegate>)otaProviderDelegate q
     _otaProviderDelegateQueue = queue;
 }
 
++ (nullable NSNumber *)nodeIDFromNOC:(MTRCertificateDERBytes)noc
+{
+    NSNumber * nodeID = nil;
+    ExtractNodeIDFromNOC(noc, &nodeID);
+    return nodeID;
+}
+
++ (nullable NSNumber *)fabricIDFromNOC:(MTRCertificateDERBytes)noc
+{
+    NSNumber * fabricID = nil;
+    ExtractFabricIDFromNOC(noc, &fabricID);
+    return fabricID;
+}
+
++ (NSData *)publicKeyFromCertificate:(MTRCertificateDERBytes)certificate
+{
+    Crypto::P256PublicKey pubKey;
+    if (ExtractPubkeyFromX509Cert(AsByteSpan(certificate), pubKey) != CHIP_NO_ERROR) {
+        return nil;
+    }
+    return [NSData dataWithBytes:pubKey.Bytes() length:pubKey.Length()];
+}
+
 @end
 
 @implementation MTRDeviceControllerExternalCertificateParameters

src/darwin/Framework/CHIP/MTRDeviceControllerStartupParams_Internal.h

@@ -85,6 +85,10 @@ NS_ASSUME_NONNULL_BEGIN
 @property (nonatomic, strong, readonly, nullable) id<MTROTAProviderDelegate> otaProviderDelegate;
 @property (nonatomic, strong, readonly, nullable) dispatch_queue_t otaProviderDelegateQueue;
 
++ (nullable NSNumber *)nodeIDFromNOC:(MTRCertificateDERBytes)noc;
++ (nullable NSNumber *)fabricIDFromNOC:(MTRCertificateDERBytes)noc;
++ (NSData *)publicKeyFromCertificate:(MTRCertificateDERBytes)certificate;
+
 @end
 
 @interface MTRDeviceControllerStartupParamsInternal : MTRDeviceControllerStartupParams

src/darwin/Framework/CHIP/MTRDeviceController_Internal.h

@@ -45,6 +45,7 @@
 #import <Matter/MTRDiagnosticLogsType.h>
 #import <Matter/MTROTAProviderDelegate.h>
 
+@class MTRDeviceControllerParameters;
 @class MTRDeviceControllerStartupParamsInternal;
 @class MTRDeviceControllerFactory;
 @class MTRDevice;
@@ -117,6 +118,21 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @property (nonatomic, readonly) MTRAsyncWorkQueue<MTRDeviceController *> * concurrentSubscriptionPool;
 
+/**
+ * Fabric ID tied to controller
+ */
+@property (nonatomic, retain, nullable) NSNumber * fabricID;
+
+/**
+ * Node ID tied to controller
+ */
+@property (nonatomic, retain, nullable) NSNumber * nodeID;
+
+/**
+ * Root Public Key tied to controller
+ */
+@property (nonatomic, retain, nullable) NSData * rootPublicKey;
+
 /**
  * Init a newly created controller.
  *
@@ -289,6 +305,30 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (void)directlyGetSessionForNode:(chip::NodeId)nodeID completion:(MTRInternalDeviceConnectionCallback)completion;
 
+/**
+ * Takes an assertion to keep the controller running. If `-[MTRDeviceController shutdown]` is called while an assertion
+ * is held, the shutdown will be honored only after all assertions are released. Invoking this method multiple times increases
+ * the number of assertions and needs to be matched with equal amount of '-[MTRDeviceController removeRunAssertion]` to release
+ * the assertion.
+ */
+- (void)addRunAssertion;
+
+/**
+ * Removes an assertion to allow the controller to shutdown once all assertions have been released.
+ * Invoking this method once all assertions have been released in a noop.
+ */
+- (void)removeRunAssertion;
+
+/**
+ * This method returns TRUE if this controller matches the fabric reference and node ID as listed in the parameters.
+ */
+- (BOOL)matchesPendingShutdownWithParams:(MTRDeviceControllerParameters *)parameters;
+
+/**
+ * Clear any pending shutdown request.
+ */
+- (void)clearPendingShutdown;
+
 @end
 
 /**