From 360d53c3eb8b1141131e60bacc56dcc4c63707f6 Mon Sep 17 00:00:00 2001
From: Alvar Hansen <alvar.hansen@bolt.eu>
Date: Thu, 16 May 2024 10:27:23 +0300
Subject: [PATCH] Fix threading issue in IPC connection

This change addresses the possibility of
`DTXIPCConnection._slaveDidConnectWithName` being called from non-main
thread by dispatching the work to main thread when it is being called
from non main thread. This is necessary so that `_otherConnection`
gets assigned on main thread.

`SBTUITestTunnelClient.serverDidConnect` is updated to wait for
`ipcConnection` to become valid. This is needed because
`DTXIPCConnection._slaveDidConnectWithName` and `SBTUITestTunnelClient.serverDidConnect`
can be called from different threads upon which we dispatch both to
main. But `SBTUITestTunnelClient.serverDidConnect` work can start
first at which point it can lead to assertion of `ipcConnection` not
being valid as `ipcConnection._otherConnection` has not yet been
assigned.
---
 .../SBTUITestTunnelClient.m                   | 21 ++++++++++++++-----
 .../DetoxIPC/DTXIPCConnection.m               | 12 +++++++++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/Sources/SBTUITestTunnelClient/SBTUITestTunnelClient.m b/Sources/SBTUITestTunnelClient/SBTUITestTunnelClient.m
index 2834c75f..90f05050 100755
--- a/Sources/SBTUITestTunnelClient/SBTUITestTunnelClient.m
+++ b/Sources/SBTUITestTunnelClient/SBTUITestTunnelClient.m
@@ -255,15 +255,26 @@ - (void)serverDidConnect:(id)sender
 
         NSLog(@"[SBTUITestTunnel] IPC tunnel did connect after, %fs", CFAbsoluteTimeGetCurrent() - weakSelf.launchStart);
 
-        if (weakSelf.startupBlock) {
-            weakSelf.startupBlock();
+        [weakSelf checkConnectionAndProceed];
+    });
+}
+
+- (void)checkConnectionAndProceed
+{
+    if (self.ipcConnection.isValid) {
+        if (_startupBlock) {
+            _startupBlock();
             NSLog(@"[SBTUITestTunnel] Did perform startupBlock");
         }
         
-        weakSelf.startupCompleted = [[weakSelf sendSynchronousRequestWithPath:SBTUITunneledApplicationCommandStartupCommandsCompleted params:@{}] isEqualToString:@"YES"];
+        _startupCompleted = [[self sendSynchronousRequestWithPath:SBTUITunneledApplicationCommandStartupCommandsCompleted params:@{}] isEqualToString:@"YES"];
 
-        NSLog(@"[SBTUITestTunnel] Tunnel ready after %fs", CFAbsoluteTimeGetCurrent() - weakSelf.launchStart);
-    });
+        NSLog(@"[SBTUITestTunnel] Tunnel ready after %fs", CFAbsoluteTimeGetCurrent() - _launchStart);
+    } else {
+        dispatch_async(dispatch_get_main_queue(), ^{
+            [self checkConnectionAndProceed];
+        });
+    }
 }
 
 - (void)performCommandWithParameters:(NSDictionary *)parameters block:(void (^)(NSDictionary *))block {}
diff --git a/Sources/SBTUITestTunnelCommon/DetoxIPC/DTXIPCConnection.m b/Sources/SBTUITestTunnelCommon/DetoxIPC/DTXIPCConnection.m
index 9e666ebd..cab77ec3 100644
--- a/Sources/SBTUITestTunnelCommon/DetoxIPC/DTXIPCConnection.m
+++ b/Sources/SBTUITestTunnelCommon/DetoxIPC/DTXIPCConnection.m
@@ -323,8 +323,16 @@ - (void)setExportedObject:(id)exportedObject
 
 - (oneway void)_slaveDidConnectWithName:(NSString*)slaveServiceName
 {
-	_otherConnection = [NSConnection connectionWithRegisteredName:slaveServiceName host:nil];
-	[NSNotificationCenter.defaultCenter addObserver:self selector:@selector(_otherConnectionDidDie:) name:NSConnectionDidDieNotification object:_otherConnection];
+    dispatch_block_t block = ^{
+        self->_otherConnection = [NSConnection connectionWithRegisteredName:slaveServiceName host:nil];
+        [NSNotificationCenter.defaultCenter addObserver:self selector:@selector(_otherConnectionDidDie:) name:NSConnectionDidDieNotification object:self->_otherConnection];
+    };
+
+    if ([NSThread isMainThread]) {
+        block();
+    } else {
+        dispatch_async(dispatch_get_main_queue(), block);
+    }
 }
 
 - (oneway void)_invokeFromRemote:(NSDictionary*)serializedInvocation