Fix potential deadlock and use-after-free in BleServerSocket on Apple.

PiperOrigin-RevId: 895638821

connections/implementation/flags/nearby_connections_feature_flags.h

@@ -114,6 +114,9 @@ constexpr auto kSafeToDisconnectVersion =
 // Enable/Disable single copy read/write for input/output buffers.
 constexpr auto kEnableSingleCopy =
     flags::Flag<bool>(kConfigPackage, "45775979", true);
+// When true, fix the BleServerSocket deadlock/use-after-free (b/494335036).
+constexpr auto kFixBleServerSocketDeadlock =
+    flags::Flag<bool>(kConfigPackage, "45775192", true);
 
 }  // namespace nearby_connections_feature
 }  // namespace config_package_nearby

internal/platform/implementation/apple/Flags/GNCFeatureFlags.h

@@ -35,4 +35,7 @@
 /** Checks whether single copy read/write is enabled in the Nearby Connections SDK. */
 @property(nonatomic, class, readonly) BOOL singleCopyEnabled;
 
+/** Checks whether BLE server socket deadlock is fixed in the Nearby Connections SDK. */
+@property(nonatomic, class, readonly) BOOL fixBleServerSocketDeadlockEnabled;
+
 @end

internal/platform/implementation/apple/Flags/GNCFeatureFlags.mm

@@ -53,4 +53,10 @@
       nearby::connections::config_package_nearby::nearby_connections_feature::kEnableSingleCopy);
 }
 
++ (BOOL)fixBleServerSocketDeadlockEnabled {
+  return nearby::NearbyFlags::GetInstance().GetBoolFlag(
+      nearby::connections::config_package_nearby::nearby_connections_feature::
+          kFixBleServerSocketDeadlock);
+}
+
 @end

internal/platform/implementation/apple/Tests/ble_medium_test.mm

@@ -23,6 +23,7 @@
 #include <utility>
 #include <vector>
 
+#import "internal/platform/implementation/apple/Flags/GNCFeatureFlags.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/GNCBLEMedium.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/GNCPeripheral.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/Sockets/Source/Central/GNSCentralManager.h"
@@ -32,10 +33,11 @@
 #import "internal/platform/implementation/apple/Mediums/BLE/Sockets/Source/Shared/GNSSocket.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCBLEMedium+Testing.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakeBLEGATTServer.h"
-#import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakeCentralManager.h"
-#import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakePeripheralManager.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakeBLEMedium.h"
+#import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakeCentralManager.h"
 #import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakePeripheral.h"
+#import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakePeripheralManager.h"
+#import "internal/platform/implementation/apple/Mediums/BLE/Tests/GNCFakeSocket.h"
 #include "internal/platform/implementation/apple/ble_utils.h"
 #include "internal/platform/implementation/ble.h"
 #import "third_party/objective_c/ocmock/v3/Source/OCMock/OCMock.h"
@@ -51,6 +53,9 @@ class BleMediumPeer {
   static void SetSocketPeripheralManager(BleMedium *ble_medium, GNSPeripheralManager *manager) {
     ble_medium->socketPeripheralManager_ = manager;
   }
+  static GNSPeripheralServiceManager *GetSocketPeripheralServiceManager(BleMedium *ble_medium) {
+    return ble_medium->socketPeripheralServiceManager_;
+  }
 };
 
 }  // namespace apple
@@ -74,8 +79,8 @@ static const char *const kTestServiceID = "TestServiceID";
   GNCFakeCentralManager *fakeCentralManager = [[GNCFakeCentralManager alloc] init];
   GNCFakePeripheralManager *fakePeripheralManager = [[GNCFakePeripheralManager alloc] init];
   _fakeGNCBLEMedium = [[GNCFakeBLEMedium alloc] initWithCentralManager:fakeCentralManager
-                                                   peripheralManager:fakePeripheralManager
-                                                               queue:dispatch_get_main_queue()];
+                                                     peripheralManager:fakePeripheralManager
+                                                                 queue:dispatch_get_main_queue()];
   _medium = std::make_unique<nearby::apple::BleMedium>((GNCBLEMedium *)_fakeGNCBLEMedium);
 }
 
@@ -229,8 +234,8 @@ static const char *const kTestServiceID = "TestServiceID";
 #pragma mark - GATT Server Tests
 
 - (void)testStartGattServer_Success {
-  _fakeGNCBLEMedium.fakeGATTServer =
-      [[GNCFakeBLEGATTServer alloc] initWithPeripheralManager:nil queue:nil];
+  _fakeGNCBLEMedium.fakeGATTServer = [[GNCFakeBLEGATTServer alloc] initWithPeripheralManager:nil
+                                                                                       queue:nil];
   auto gatt_server = _medium->StartGattServer({});
 
   XCTAssertNotEqual(gatt_server.get(), nullptr);
@@ -474,7 +479,10 @@ static const char *const kTestServiceID = "TestServiceID";
 
 #pragma mark - Server Socket Tests
 
-- (void)testOpenServerSocket_Success {
+- (void)testOpenServerSocket_Success_LegacyPath {
+  id mockFeatureFlags = OCMClassMock([GNCFeatureFlags class]);
+  OCMStub([mockFeatureFlags fixBleServerSocketDeadlockEnabled]).andReturn(NO);
+
   id mockPeripheralManager = OCMClassMock([GNSPeripheralManager class]);
   OCMStub([mockPeripheralManager addPeripheralServiceManager:[OCMArg any]
                                    bleServiceAddedCompletion:[OCMArg any]])
@@ -487,6 +495,85 @@ static const char *const kTestServiceID = "TestServiceID";
   auto server_socket = _medium->OpenServerSocket(kTestServiceID);
 
   XCTAssertNotEqual(server_socket.get(), nullptr);
+
+  GNSPeripheralServiceManager *serviceManager =
+      nearby::apple::BleMediumPeer::GetSocketPeripheralServiceManager(_medium.get());
+  XCTAssertNotNil(serviceManager);
+}
+
+- (void)testOpenServerSocket_Success_OptimizedPath {
+  id mockFeatureFlags = OCMClassMock([GNCFeatureFlags class]);
+  OCMStub([mockFeatureFlags fixBleServerSocketDeadlockEnabled]).andReturn(YES);
+
+  id mockPeripheralManager = OCMClassMock([GNSPeripheralManager class]);
+  OCMStub([mockPeripheralManager addPeripheralServiceManager:[OCMArg any]
+                                   bleServiceAddedCompletion:[OCMArg any]])
+      .andDo(^(GNSPeripheralManager *localSelf, GNSPeripheralServiceManager *manager,
+               void (^completion)(NSError *error)) {
+        completion(nil);
+      });
+  nearby::apple::BleMediumPeer::SetSocketPeripheralManager(_medium.get(), mockPeripheralManager);
+
+  auto server_socket = _medium->OpenServerSocket(kTestServiceID);
+
+  XCTAssertNotEqual(server_socket.get(), nullptr);
+
+  GNSPeripheralServiceManager *serviceManager =
+      nearby::apple::BleMediumPeer::GetSocketPeripheralServiceManager(_medium.get());
+  XCTAssertNotNil(serviceManager);
+}
+
+- (void)testOpenServerSocket_OptimizedPath_AcceptSocketAfterClose {
+  id mockFeatureFlags = OCMClassMock([GNCFeatureFlags class]);
+  OCMStub([mockFeatureFlags fixBleServerSocketDeadlockEnabled]).andReturn(YES);
+
+  id mockPeripheralManager = OCMClassMock([GNSPeripheralManager class]);
+  OCMStub([mockPeripheralManager addPeripheralServiceManager:[OCMArg any]
+                                   bleServiceAddedCompletion:[OCMArg any]])
+      .andDo(^(GNSPeripheralManager *localSelf, GNSPeripheralServiceManager *manager,
+               void (^completion)(NSError *error)) {
+        completion(nil);
+      });
+  nearby::apple::BleMediumPeer::SetSocketPeripheralManager(_medium.get(), mockPeripheralManager);
+
+  auto server_socket = _medium->OpenServerSocket(kTestServiceID);
+  XCTAssertNotEqual(server_socket.get(), nullptr);
+  __block BOOL (^capturedHandler)(GNSSocket *) = nil;
+  id mockServiceManagerClass = OCMClassMock([GNSPeripheralServiceManager class]);
+  OCMStub([mockServiceManagerClass alloc]).andReturn(mockServiceManagerClass);
+  OCMStub([mockServiceManagerClass initWithBleServiceUUID:[OCMArg any]
+                                 addPairingCharacteristic:NO
+                                shouldAcceptSocketHandler:[OCMArg any]])
+      .andDo(^(NSInvocation *invocation) {
+        BOOL (^handler)(GNSSocket *);
+        [invocation getArgument:&handler atIndex:4];
+        capturedHandler = handler;
+      })
+      .andReturn(mockServiceManagerClass);
+
+  auto server_socket_for_handler_capture = _medium->OpenServerSocket(kTestServiceID);
+  XCTAssertNotEqual(server_socket_for_handler_capture.get(), nullptr);
+  XCTAssertNotNil(capturedHandler);
+
+  // Invoke the shouldAcceptSocketHandler with a fake socket.
+  GNCFakeSocket *fakeSocket = [[GNCFakeSocket alloc] init];
+  BOOL result = capturedHandler((GNSSocket *)fakeSocket);
+  XCTAssertTrue(result);
+
+  // Close the server_socket. This triggers the close notifier, setting server_socket_ptr_ to null.
+  server_socket_for_handler_capture->Close();
+
+  // Now simulate the connection completing. It should safely ignore the connection because
+  // server_socket_ptr_ is null, preventing use-after-free or deadlocks.
+  [fakeSocket simulateSocketDidConnect];
+
+  // Since we use dispatch_async internally for connection callback, give it a small amount of time
+  // to process so we know it didn't crash.
+  XCTestExpectation *expectation2 = [self expectationWithDescription:@"Wait for async execution"];
+  dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.1 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
+    [expectation2 fulfill];
+  });
+  [self waitForExpectations:@[ expectation2 ] timeout:1.0];
 }
 
 - (void)testOpenServerSocket_Failure {

internal/platform/implementation/apple/ble_medium.h

@@ -207,6 +207,13 @@ class BleMedium : public api::ble::BleMedium {
                                   NSDictionary<CBUUID *, NSData *> *service_data);
   NSDate *GetLastTimestampToCleanExpiredAdvertisementPackets();
 
+  // Opens a BLE server socket based on service ID with deadlock safety.
+  std::unique_ptr<api::ble::BleServerSocket> OpenServerSocketWithDeadlockSafety(
+      const std::string &service_id);
+
+  // Opens a BLE server socket based on service ID using the legacy implementation.
+  std::unique_ptr<api::ble::BleServerSocket> OpenServerSocketLegacy(const std::string &service_id);
+
   // The executor for handling callbacks.
   apple::SingleThreadExecutor callback_executor_;
 
@@ -237,6 +244,11 @@ class BleMedium : public api::ble::BleMedium {
   // callback.
   api::ble::BleMedium::ScanningCallback scanning_cb_;
 
+  // Used for the BleServerSocket.
+  absl::Mutex server_socket_mutex_;
+  BleServerSocket *server_socket_ptr_ ABSL_GUARDED_BY(server_socket_mutex_) = nullptr;
+
+  // Used for the L2CAP server socket.
   absl::Mutex l2cap_server_socket_mutex_;
   BleL2capServerSocket *l2cap_server_socket_ptr_ = nullptr;
 

internal/platform/implementation/apple/ble_medium.mm

@@ -448,23 +448,117 @@ std::unique_ptr<api::ble::GattClient> BleMedium::ConnectToGattServer(
 // TODO(b/293336684): Old Weave code that need to be deleted once shared Weave is complete.
 std::unique_ptr<api::ble::BleServerSocket> BleMedium::OpenServerSocket(
     const std::string &service_id) {
+  if (GNCFeatureFlags.fixBleServerSocketDeadlockEnabled) {
+    return OpenServerSocketWithDeadlockSafety(service_id);
+  } else {
+    return OpenServerSocketLegacy(service_id);
+  }
+}
+
+std::unique_ptr<api::ble::BleServerSocket> BleMedium::OpenServerSocketWithDeadlockSafety(
+    const std::string &service_id) {
   auto server_socket = std::make_unique<BleServerSocket>();
-  __block auto server_socket_ptr = server_socket.get();
 
   if (socketPeripheralManager_ == nil) {
     socketPeripheralManager_ = [[GNSPeripheralManager alloc] initWithAdvertisedName:nil
                                                                   restoreIdentifier:nil];
   }
-
   if (socketPeripheralManager_ == nil) {
     GNCLoggerError(@"Failed to create peripheral manager.");
     return nullptr;
   }
 
+  // Fix for b/494335036 (Registry + Background Queue)
+  {
+    absl::MutexLock lock(server_socket_mutex_);
+    server_socket_ptr_ = server_socket.get();
+  }
+  server_socket->SetCloseNotifier([this]() {
+    absl::MutexLock lock(server_socket_mutex_);
+    server_socket_ptr_ = nullptr;
+  });
+
   socketPeripheralServiceManager_ = [[GNSPeripheralServiceManager alloc]
          initWithBleServiceUUID:[CBUUID UUIDWithString:kWeaveServiceUUID]
        addPairingCharacteristic:NO
       shouldAcceptSocketHandler:^BOOL(GNSSocket *socket) {
+        // Optimized Path: Use background queue and registry validation.
+        GNCMWaitForConnection(socket, connection_callback_queue_, ^(BOOL didConnect) {
+          GNCMBleConnection *connection =
+              [GNCMBleConnection connectionWithSocket:socket
+                                            serviceID:nil
+                                  expectedIntroPacket:YES
+                                        callbackQueue:connection_callback_queue_];
+
+          auto socket_wrapper = std::make_unique<BleSocket>(connection);
+          socket_wrapper->SetCloseNotifier(
+              [socketPeripheralManager = socketPeripheralManager_,
+               serviceUUID = socketPeripheralServiceManager_.serviceUUID]() {
+                [socketPeripheralManager
+                    removePeripheralServiceManagerForServiceUUID:serviceUUID
+                                     bleServiceRemovedCompletion:^(NSError *_Nullable error) {
+                                       GNCLoggerInfo(@"BleSocket is removed peripheral manager.");
+                                     }];
+              });
+
+          connection.connectionHandlers = socket_wrapper->GetInputStream().GetConnectionHandlers();
+
+          // Fix: Verify the BleServerSocket still exists before calling Connect().
+          // This prevents the use-after-free/deadlock reported in b/494335036.
+          absl::MutexLock lock(server_socket_mutex_);
+          if (server_socket_ptr_) {
+            server_socket_ptr_->Connect(std::move(socket_wrapper));
+            GNCLoggerInfo(@"BleServerSocket is created with connection");
+          } else {
+            GNCLoggerWarning(@"BleServerSocket was destroyed; ignoring connection.");
+          }
+        });
+        return YES;
+      }];
+
+  dispatch_semaphore_t semaphore = dispatch_semaphore_create(0);
+  __block NSError *blockError = nil;
+  [socketPeripheralManager_ addPeripheralServiceManager:socketPeripheralServiceManager_
+                              bleServiceAddedCompletion:^(NSError *error) {
+                                if (error != nil) {
+                                  GNCLoggerError(@"Failed to add Weave service: %@", error);
+                                  blockError = error;
+                                }
+                                dispatch_semaphore_signal(semaphore);
+                              }];
+  [socketPeripheralManager_ start];
+  dispatch_time_t timeout = dispatch_time(DISPATCH_TIME_NOW, kApiTimeoutInSeconds * NSEC_PER_SEC);
+  if (dispatch_semaphore_wait(semaphore, timeout) != 0) {
+    GNCLoggerError(@"OpenServerSocket operation timed out.");
+    return nullptr;
+  }
+  if (blockError != nil) {
+    return nullptr;
+  }
+  return std::move(server_socket);
+}
+
+std::unique_ptr<api::ble::BleServerSocket> BleMedium::OpenServerSocketLegacy(
+    const std::string &service_id) {
+  auto server_socket = std::make_unique<BleServerSocket>();
+
+  if (socketPeripheralManager_ == nil) {
+    socketPeripheralManager_ = [[GNSPeripheralManager alloc] initWithAdvertisedName:nil
+                                                                  restoreIdentifier:nil];
+  }
+  if (socketPeripheralManager_ == nil) {
+    GNCLoggerError(@"Failed to create peripheral manager.");
+    return nullptr;
+  }
+
+  // Raw pointer for closure capture in the legacy path (risks use-after-free).
+  BleServerSocket *server_socket_ptr = server_socket.get();
+
+  socketPeripheralServiceManager_ = [[GNSPeripheralServiceManager alloc]
+         initWithBleServiceUUID:[CBUUID UUIDWithString:kWeaveServiceUUID]
+       addPairingCharacteristic:NO
+      shouldAcceptSocketHandler:^BOOL(GNSSocket *socket) {
+        // Legacy Path: Verbatim copy of original code (blocks Main Thread).
         GNCMWaitForConnection(socket, nil, ^(BOOL didConnect) {
           GNCMBleConnection *connection =
               [GNCMBleConnection connectionWithSocket:socket
@@ -603,7 +697,14 @@ std::unique_ptr<api::ble::BleSocket> BleMedium::Connect(
                                dispatch_semaphore_signal(semaphore);
                                return;
                              }
-                             GNCMWaitForConnection(nssocket, nil, ^(BOOL didConnect) {
+
+                             // Suggestion: Use the connection callback queue instead of nil
+                             dispatch_queue_t targetQueue =
+                                 GNCFeatureFlags.fixBleServerSocketDeadlockEnabled
+                                     ? connection_callback_queue_
+                                     : nil;
+
+                             GNCMWaitForConnection(nssocket, targetQueue, ^(BOOL didConnect) {
                                if (!didConnect) {
                                  dispatch_semaphore_signal(semaphore);
                                  return;