From 70df3565bc6d01f302cd4dfaf0360a21780fb2b2 Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Fri, 26 May 2023 15:50:35 -0400
Subject: [PATCH] Fix thread race in MTRXPCListenerSampleTests. (#26862)

WARNING: ThreadSanitizer: race on NSMutableDictionary (pid=72050)
      Modifying access of NSMutableDictionary at 0x7b0800149020 by thread T7:
        #0 -[__NSDictionaryM setObject:forKey:] <null>:2 (CoreFoundation:x86_64+0x40d4e)
        #1 -[MTRXPCListenerSample listener:shouldAcceptNewConnection:] MTRXPCListenerSampleTests.m:116 (MatterTests:x86_64+0x48437b)
        #2 service_connection_handler_make_connection <null>:2 (Foundation:x86_64+0xdea93)

      Previous modifying access of NSMutableDictionary at 0x7b0800149020 by thread T4:
        #0 -[__NSDictionaryM removeObjectForKey:] <null>:2 (CoreFoundation:x86_64+0x6cd6d)
        #1 __59-[MTRXPCListenerSample listener:shouldAcceptNewConnection:]_block_invoke MTRXPCListenerSampleTests.m:119 (MatterTests:x86_64+0x484630)
---
 .../Framework/CHIPTests/MTRXPCListenerSampleTests.m    | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/darwin/Framework/CHIPTests/MTRXPCListenerSampleTests.m b/src/darwin/Framework/CHIPTests/MTRXPCListenerSampleTests.m
index c22bc7a72125ea..392f0d1f1a3513 100644
--- a/src/darwin/Framework/CHIPTests/MTRXPCListenerSampleTests.m
+++ b/src/darwin/Framework/CHIPTests/MTRXPCListenerSampleTests.m
@@ -33,6 +33,7 @@
 
 // system dependencies
 #import <XCTest/XCTest.h>
+#import <os/lock.h>
 
 static uint16_t kTestVendorId = 0xFFF1u;
 
@@ -73,6 +74,8 @@ @interface MTRXPCListenerSample ()
 @property (nonatomic, readonly, strong)
     NSMutableDictionary<NSNumber *, MTRClusterStateCacheContainer *> * clusterStateCacheDictionary;
 
+// serversLock controls access to _servers.
+@property (nonatomic, readonly) os_unfair_lock serversLock;
 @end
 
 @implementation MTRXPCListenerSample
@@ -86,6 +89,7 @@ - (instancetype)init
         _clusterStateCacheDictionary = [NSMutableDictionary dictionary];
         _xpcListener = [NSXPCListener anonymousListener];
         [_xpcListener setDelegate:(id<NSXPCListenerDelegate>) self];
+        _serversLock = OS_UNFAIR_LOCK_INIT;
     }
     return self;
 }
@@ -113,10 +117,16 @@ - (BOOL)listener:(NSXPCListener *)listener shouldAcceptNewConnection:(NSXPCConne
     __auto_type newServer = [[MTRDeviceControllerServerSample alloc] initWithClientProxy:[newConnection remoteObjectProxy]
                                                              clusterStateCacheDictionary:_clusterStateCacheDictionary];
     newConnection.exportedObject = newServer;
+
+    os_unfair_lock_lock(&_serversLock);
     [_servers setObject:newServer forKey:newServer.identifier];
+    os_unfair_lock_unlock(&_serversLock);
+
     newConnection.invalidationHandler = ^{
         NSLog(@"XPC connection disconnected");
+        os_unfair_lock_lock(&self->_serversLock);
         [self.servers removeObjectForKey:newServer.identifier];
+        os_unfair_lock_unlock(&self->_serversLock);
     };
     [newConnection resume];
     return YES;