Merge branch 'google3' to master, roll forward up to cl/355302206.

cpp/core/internal/mediums/ble.cc

@@ -335,7 +335,12 @@ BleSocket Ble::Connect(BlePeripheral& peripheral, const std::string& service_id,
     return socket;
   }
 
-  socket = medium_.Connect(peripheral, service_id);
+  if (cancellation_flag->Cancelled()) {
+    NEARBY_LOGS(INFO) << "Can't create client BLE socket due to cancel.";
+    return socket;
+  }
+
+  socket = medium_.Connect(peripheral, service_id, cancellation_flag);
   if (!socket.IsValid()) {
     NEARBY_LOGS(INFO) << "Failed to Connect via BLE [service=" << service_id
                       << "]";

cpp/core/internal/mediums/ble_test.cc

@@ -29,12 +29,23 @@ namespace nearby {
 namespace connections {
 namespace {
 
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
 constexpr absl::Duration kWaitDuration = absl::Milliseconds(1000);
 constexpr absl::string_view kServiceID{"com.google.location.nearby.apps.test"};
 constexpr absl::string_view kAdvertisementString{"\x0a\x0b\x0c\x0d"};
 constexpr absl::string_view kFastAdvertisementServiceUuid{"\xf3\xfe"};
 
-class BleTest : public ::testing::Test {
+class BleTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveredPeripheralCallback = BleMedium::DiscoveredPeripheralCallback;
 
@@ -43,6 +54,123 @@ class BleTest : public ::testing::Test {
   MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
+TEST_P(BleTest, CanStartAcceptingConnectionsAndConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  BluetoothRadio radio_a;
+  BluetoothRadio radio_b;
+  Ble ble_a{radio_a};
+  Ble ble_b{radio_b};
+  radio_a.Enable();
+  radio_b.Enable();
+  std::string service_id(kServiceID);
+  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
+  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
+  CountDownLatch found_latch(1);
+  CountDownLatch accept_latch(1);
+
+  ble_a.StartAdvertising(service_id, advertisement_bytes,
+                         fast_advertisement_service_uuid);
+  ble_a.StartAcceptingConnections(
+      service_id,
+      {
+          .accepted_cb = [&accept_latch](
+                             BleSocket socket,
+                             const std::string&) { accept_latch.CountDown(); },
+      });
+  BlePeripheral discovered_peripheral;
+  ble_b.StartScanning(
+      service_id, fast_advertisement_service_uuid,
+      {
+          .peripheral_discovered_cb =
+              [&found_latch, &discovered_peripheral](
+                  BlePeripheral& peripheral, const std::string& service_id,
+                  const ByteArray& advertisement_bytes,
+                  bool fast_advertisement) {
+                discovered_peripheral = peripheral;
+                NEARBY_LOG(
+                    INFO,
+                    "Discovered peripheral=%p [impl=%p], fast advertisement=%d",
+                    &peripheral, &peripheral.GetImpl(), fast_advertisement);
+                found_latch.CountDown();
+              },
+      });
+
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  ASSERT_TRUE(discovered_peripheral.IsValid());
+  CancellationFlag flag;
+  BleSocket socket = ble_b.Connect(discovered_peripheral, service_id, &flag);
+  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(socket.IsValid());
+  ble_b.StopScanning(service_id);
+  ble_a.StopAdvertising(service_id);
+  env_.Stop();
+}
+
+TEST_P(BleTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  BluetoothRadio radio_a;
+  BluetoothRadio radio_b;
+  Ble ble_a{radio_a};
+  Ble ble_b{radio_b};
+  radio_a.Enable();
+  radio_b.Enable();
+  std::string service_id(kServiceID);
+  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
+  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
+  CountDownLatch found_latch(1);
+  CountDownLatch accept_latch(1);
+
+  ble_a.StartAdvertising(service_id, advertisement_bytes,
+                         fast_advertisement_service_uuid);
+  ble_a.StartAcceptingConnections(
+      service_id,
+      {
+          .accepted_cb = [&accept_latch](
+                             BleSocket socket,
+                             const std::string&) { accept_latch.CountDown(); },
+      });
+  BlePeripheral discovered_peripheral;
+  ble_b.StartScanning(
+      service_id, fast_advertisement_service_uuid,
+      {
+          .peripheral_discovered_cb =
+              [&found_latch, &discovered_peripheral](
+                  BlePeripheral& peripheral, const std::string& service_id,
+                  const ByteArray& advertisement_bytes,
+                  bool fast_advertisement) {
+                discovered_peripheral = peripheral;
+                NEARBY_LOG(
+                    INFO,
+                    "Discovered peripheral=%p [impl=%p], fast advertisement=%d",
+                    &peripheral, &peripheral.GetImpl(), fast_advertisement);
+                found_latch.CountDown();
+              },
+      });
+
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  ASSERT_TRUE(discovered_peripheral.IsValid());
+  CancellationFlag flag(true);
+  BleSocket socket = ble_b.Connect(discovered_peripheral, service_id, &flag);
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_TRUE(socket.IsValid());
+  } else {
+    EXPECT_FALSE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_FALSE(socket.IsValid());
+  }
+  ble_b.StopScanning(service_id);
+  ble_a.StopAdvertising(service_id);
+  env_.Stop();
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedBleTest, BleTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(BleTest, CanConstructValidObject) {
   env_.Start();
   BluetoothRadio radio_a;
@@ -129,58 +257,6 @@ TEST_F(BleTest, CanStartDiscovery) {
   env_.Stop();
 }
 
-TEST_F(BleTest, CanStartAcceptingConnectionsAndConnect) {
-  env_.Start();
-  BluetoothRadio radio_a;
-  BluetoothRadio radio_b;
-  Ble ble_a{radio_a};
-  Ble ble_b{radio_b};
-  radio_a.Enable();
-  radio_b.Enable();
-  std::string service_id(kServiceID);
-  ByteArray advertisement_bytes{std::string(kAdvertisementString)};
-  std::string fast_advertisement_service_uuid(kFastAdvertisementServiceUuid);
-  CountDownLatch found_latch(1);
-  CountDownLatch accept_latch(1);
-
-  ble_a.StartAdvertising(service_id, advertisement_bytes,
-                         fast_advertisement_service_uuid);
-  ble_a.StartAcceptingConnections(
-      service_id,
-      {
-          .accepted_cb = [&accept_latch](
-                             BleSocket socket,
-                             const std::string&) { accept_latch.CountDown(); },
-      });
-  BlePeripheral discovered_peripheral;
-  ble_b.StartScanning(
-      service_id, fast_advertisement_service_uuid,
-      {
-          .peripheral_discovered_cb =
-              [&found_latch, &discovered_peripheral](
-                  BlePeripheral& peripheral, const std::string& service_id,
-                  const ByteArray& advertisement_bytes,
-                  bool fast_advertisement) {
-                discovered_peripheral = peripheral;
-                NEARBY_LOG(
-                    INFO,
-                    "Discovered peripheral=%p [impl=%p], fast advertisement=%d",
-                    &peripheral, &peripheral.GetImpl(), fast_advertisement);
-                found_latch.CountDown();
-              },
-      });
-
-  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
-  ASSERT_TRUE(discovered_peripheral.IsValid());
-  CancellationFlag flag;
-  BleSocket socket = ble_b.Connect(discovered_peripheral, service_id, &flag);
-  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
-  EXPECT_TRUE(socket.IsValid());
-  ble_b.StopScanning(service_id);
-  ble_a.StopAdvertising(service_id);
-  env_.Stop();
-}
-
 }  // namespace
 }  // namespace connections
 }  // namespace nearby

cpp/core/internal/mediums/bluetooth_classic.cc

@@ -350,7 +350,8 @@ BluetoothSocket BluetoothClassic::Connect(BluetoothDevice& bluetooth_device,
                                           CancellationFlag* cancellation_flag) {
   for (int attempts_count = 0; attempts_count < kConnectAttemptsLimit;
        attempts_count++) {
-    auto wrapper_result = AttemptToConnect(bluetooth_device, service_name);
+    auto wrapper_result =
+        AttemptToConnect(bluetooth_device, service_name, cancellation_flag);
     if (wrapper_result.IsValid()) {
       return wrapper_result;
     }
@@ -359,7 +360,8 @@ BluetoothSocket BluetoothClassic::Connect(BluetoothDevice& bluetooth_device,
 }
 
 BluetoothSocket BluetoothClassic::AttemptToConnect(
-    BluetoothDevice& bluetooth_device, const std::string& service_name) {
+    BluetoothDevice& bluetooth_device, const std::string& service_name,
+    CancellationFlag* cancellation_flag) {
   MutexLock lock(&mutex_);
   NEARBY_LOG(INFO, "BluetoothClassic::Connect: device=%p", &bluetooth_device);
   // Socket to return. To allow for NRVO to work, it has to be a single object.
@@ -386,8 +388,14 @@ BluetoothSocket BluetoothClassic::AttemptToConnect(
     return socket;
   }
 
+  if (cancellation_flag->Cancelled()) {
+    NEARBY_LOGS(INFO) << "Can't create client BT socket due to cancel.";
+    return socket;
+  }
+
   socket = medium_.ConnectToService(bluetooth_device,
-                                    GenerateUuidFromString(service_name));
+                                    GenerateUuidFromString(service_name),
+                                    cancellation_flag);
   if (!socket.IsValid()) {
     NEARBY_LOG(INFO, "Failed to Connect via BT [service=%s]",
                service_name.c_str());

cpp/core/internal/mediums/bluetooth_classic.h

@@ -170,7 +170,8 @@ class BluetoothClassic {
   // Returns socket instance. On success, BluetoothSocket.IsValid() return true.
   // Called by client.
   BluetoothSocket AttemptToConnect(BluetoothDevice& bluetooth_device,
-                                   const std::string& service_name);
+                                   const std::string& service_name,
+                                   CancellationFlag* cancellation_flag);
 
   mutable Mutex mutex_;
   BluetoothRadio& radio_ ABSL_GUARDED_BY(mutex_);

cpp/core/internal/mediums/bluetooth_classic_test.cc

@@ -31,9 +31,20 @@ namespace nearby {
 namespace connections {
 namespace {
 
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
 constexpr absl::Duration kWaitDuration = absl::Milliseconds(1000);
 
-class BluetoothClassicTest : public ::testing::Test {
+class BluetoothClassicTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveryCallback = BluetoothClassicMedium::DiscoveryCallback;
 
@@ -72,6 +83,125 @@ class BluetoothClassicTest : public ::testing::Test {
   std::unique_ptr<BluetoothClassic> bt_b_;
 };
 
+TEST_P(BluetoothClassicTest, CanConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+
+  constexpr absl::string_view kDeviceName{"Simulated BT device #1"};
+  constexpr absl::string_view kServiceName{"service name"};
+
+  BluetoothRadio& radio_for_client = *radio_a_;
+  BluetoothRadio& radio_for_server = *radio_b_;
+  BluetoothClassic& bt_client = *bt_a_;
+  BluetoothClassic& bt_server = *bt_b_;
+
+  EXPECT_TRUE(radio_for_client.IsEnabled());
+  EXPECT_TRUE(radio_for_server.IsEnabled());
+
+  EXPECT_TRUE(bt_server.TurnOnDiscoverability(std::string(kDeviceName)));
+  EXPECT_EQ(radio_for_server.GetBluetoothAdapter().GetName(),
+            std::string(kDeviceName));
+  CountDownLatch latch(1);
+  BluetoothDevice discovered_device;
+  EXPECT_TRUE(bt_client.StartDiscovery({
+      .device_discovered_cb =
+          [&latch, &discovered_device](BluetoothDevice& device) {
+            discovered_device = device;
+            NEARBY_LOG(INFO, "Discovered device=%p [impl=%p]", &device,
+                       &device.GetImpl());
+            latch.CountDown();
+          },
+  }));
+  EXPECT_TRUE(latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(bt_server.TurnOffDiscoverability());
+  ASSERT_TRUE(discovered_device.IsValid());
+  BluetoothSocket socket_for_server;
+  CountDownLatch accept_latch(1);
+  EXPECT_TRUE(bt_server.StartAcceptingConnections(
+      std::string(kServiceName),
+      {
+          .accepted_cb =
+              [&socket_for_server, &accept_latch](BluetoothSocket socket) {
+                socket_for_server = std::move(socket);
+                accept_latch.CountDown();
+              },
+      }));
+  CancellationFlag flag;
+  BluetoothSocket socket_for_client =
+      bt_client.Connect(discovered_device, std::string(kServiceName), &flag);
+  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(bt_server.StopAcceptingConnections(std::string(kServiceName)));
+  EXPECT_TRUE(socket_for_server.IsValid());
+  EXPECT_TRUE(socket_for_client.IsValid());
+  EXPECT_TRUE(socket_for_server.GetRemoteDevice().IsValid());
+  EXPECT_TRUE(socket_for_client.GetRemoteDevice().IsValid());
+}
+
+TEST_P(BluetoothClassicTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+
+  constexpr absl::string_view kDeviceName{"Simulated BT device #1"};
+  constexpr absl::string_view kServiceName{"service name"};
+
+  BluetoothRadio& radio_for_client = *radio_a_;
+  BluetoothRadio& radio_for_server = *radio_b_;
+  BluetoothClassic& bt_client = *bt_a_;
+  BluetoothClassic& bt_server = *bt_b_;
+
+  EXPECT_TRUE(radio_for_client.IsEnabled());
+  EXPECT_TRUE(radio_for_server.IsEnabled());
+
+  EXPECT_TRUE(bt_server.TurnOnDiscoverability(std::string(kDeviceName)));
+  EXPECT_EQ(radio_for_server.GetBluetoothAdapter().GetName(),
+            std::string(kDeviceName));
+  CountDownLatch latch(1);
+  BluetoothDevice discovered_device;
+  EXPECT_TRUE(bt_client.StartDiscovery({
+      .device_discovered_cb =
+          [&latch, &discovered_device](BluetoothDevice& device) {
+            discovered_device = device;
+            NEARBY_LOG(INFO, "Discovered device=%p [impl=%p]", &device,
+                       &device.GetImpl());
+            latch.CountDown();
+          },
+  }));
+  EXPECT_TRUE(latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(bt_server.TurnOffDiscoverability());
+  ASSERT_TRUE(discovered_device.IsValid());
+  BluetoothSocket socket_for_server;
+  CountDownLatch accept_latch(1);
+  EXPECT_TRUE(bt_server.StartAcceptingConnections(
+      std::string(kServiceName),
+      {
+          .accepted_cb =
+              [&socket_for_server, &accept_latch](BluetoothSocket socket) {
+                socket_for_server = std::move(socket);
+                accept_latch.CountDown();
+              },
+      }));
+  CancellationFlag flag(true);
+  BluetoothSocket socket_for_client =
+      bt_client.Connect(discovered_device, std::string(kServiceName), &flag);
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_TRUE(bt_server.StopAcceptingConnections(std::string(kServiceName)));
+    EXPECT_TRUE(socket_for_server.IsValid());
+    EXPECT_TRUE(socket_for_client.IsValid());
+    EXPECT_TRUE(socket_for_server.GetRemoteDevice().IsValid());
+    EXPECT_TRUE(socket_for_client.GetRemoteDevice().IsValid());
+  } else {
+    EXPECT_FALSE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_TRUE(bt_server.StopAcceptingConnections(std::string(kServiceName)));
+    EXPECT_FALSE(socket_for_server.IsValid());
+    EXPECT_FALSE(socket_for_client.IsValid());
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedBluetoothClassicTest, BluetoothClassicTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(BluetoothClassicTest, CanConstructValidObject) {
   EXPECT_TRUE(bt_a_->IsMediumValid());
   EXPECT_TRUE(bt_a_->IsAdapterValid());
@@ -154,57 +284,6 @@ TEST_F(BluetoothClassicTest, CanStartAcceptingConnections) {
   EXPECT_TRUE(bt_server.StopAcceptingConnections(std::string(kServiceName)));
 }
 
-TEST_F(BluetoothClassicTest, CanConnect) {
-  constexpr absl::string_view kDeviceName{"Simulated BT device #1"};
-  constexpr absl::string_view kServiceName{"service name"};
-
-  BluetoothRadio& radio_for_client = *radio_a_;
-  BluetoothRadio& radio_for_server = *radio_b_;
-  BluetoothClassic& bt_client = *bt_a_;
-  BluetoothClassic& bt_server = *bt_b_;
-
-  EXPECT_TRUE(radio_for_client.IsEnabled());
-  EXPECT_TRUE(radio_for_server.IsEnabled());
-
-  EXPECT_TRUE(bt_server.TurnOnDiscoverability(std::string(kDeviceName)));
-  EXPECT_EQ(radio_for_server.GetBluetoothAdapter().GetName(),
-            std::string(kDeviceName));
-  CountDownLatch latch(1);
-  BluetoothDevice discovered_device;
-  EXPECT_TRUE(bt_client.StartDiscovery({
-      .device_discovered_cb =
-          [&latch, &discovered_device](BluetoothDevice& device) {
-            discovered_device = device;
-            NEARBY_LOG(INFO, "Discovered device=%p [impl=%p]", &device,
-                       &device.GetImpl());
-            latch.CountDown();
-          },
-  }));
-  EXPECT_TRUE(latch.Await(kWaitDuration).result());
-  EXPECT_TRUE(bt_server.TurnOffDiscoverability());
-  ASSERT_TRUE(discovered_device.IsValid());
-  BluetoothSocket socket_for_server;
-  CountDownLatch accept_latch(1);
-  EXPECT_TRUE(bt_server.StartAcceptingConnections(
-      std::string(kServiceName),
-      {
-          .accepted_cb =
-              [&socket_for_server, &accept_latch](BluetoothSocket socket) {
-                socket_for_server = std::move(socket);
-                accept_latch.CountDown();
-              },
-      }));
-  CancellationFlag flag;
-  BluetoothSocket socket_for_client =
-      bt_client.Connect(discovered_device, std::string(kServiceName), &flag);
-  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
-  EXPECT_TRUE(bt_server.StopAcceptingConnections(std::string(kServiceName)));
-  EXPECT_TRUE(socket_for_server.IsValid());
-  EXPECT_TRUE(socket_for_client.IsValid());
-  EXPECT_TRUE(socket_for_server.GetRemoteDevice().IsValid());
-  EXPECT_TRUE(socket_for_client.GetRemoteDevice().IsValid());
-}
-
 }  // namespace
 }  // namespace connections
 }  // namespace nearby

cpp/core/internal/mediums/webrtc.cc

@@ -212,8 +212,8 @@ WebRtcSocketWrapper WebRtc::Connect(const std::string& service_id,
                                     CancellationFlag* cancellation_flag) {
   for (int attempts_count = 0; attempts_count < kConnectAttemptsLimit;
        attempts_count++) {
-    auto wrapper_result =
-        AttemptToConnect(service_id, remote_peer_id, location_hint);
+    auto wrapper_result = AttemptToConnect(service_id, remote_peer_id,
+                                           location_hint, cancellation_flag);
     if (wrapper_result.IsValid()) {
       return wrapper_result;
     }
@@ -223,7 +223,7 @@ WebRtcSocketWrapper WebRtc::Connect(const std::string& service_id,
 
 WebRtcSocketWrapper WebRtc::AttemptToConnect(
     const std::string& service_id, const PeerId& remote_peer_id,
-    const LocationHint& location_hint) {
+    const LocationHint& location_hint, CancellationFlag* cancellation_flag) {
   ConnectionRequestInfo info = ConnectionRequestInfo();
   info.self_peer_id = PeerId::FromRandom();
   Future<WebRtcSocketWrapper> socket_future = info.socket_future;
@@ -238,6 +238,11 @@ WebRtcSocketWrapper WebRtc::AttemptToConnect(
       return WebRtcSocketWrapper();
     }
 
+    if (cancellation_flag->Cancelled()) {
+      NEARBY_LOGS(INFO) << "Cannot connect with WebRtc due to cancel.";
+      return WebRtcSocketWrapper();
+    }
+
     // Create a new ConnectionFlow for this connection attempt.
     std::unique_ptr<ConnectionFlow> connection_flow =
         CreateConnectionFlow(service_id, remote_peer_id);

cpp/core/internal/mediums/webrtc.h

@@ -143,7 +143,8 @@ class WebRtc {
   // Runs on @MainThread.
   WebRtcSocketWrapper AttemptToConnect(const std::string& service_id,
                                        const PeerId& peer_id,
-                                       const LocationHint& location_hint)
+                                       const LocationHint& location_hint,
+                                       CancellationFlag* cancellation_flag)
       ABSL_LOCKS_EXCLUDED(mutex_);
 
   // Returns if the device is accepting connection with specific service id.

cpp/core/internal/mediums/webrtc_test.cc

@@ -28,7 +28,18 @@ namespace mediums {
 
 namespace {
 
-class WebRtcTest : public ::testing::Test {
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
+class WebRtcTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   WebRtcTest() {
     MediumEnvironment::Instance().Stop();
@@ -36,6 +47,87 @@ class WebRtcTest : public ::testing::Test {
   }
 };
 
+// Tests the flow when the two devices exchange SDP messages and connect to each
+// other but the signaling channel is closed before sending the data.
+TEST_P(WebRtcTest, ConnectBothDevices_ShutdownSignaling_SendData) {
+  FeatureFlags feature_flags = GetParam();
+  MediumEnvironment::Instance().SetFeatureFlags(feature_flags);
+  WebRtc receiver, sender;
+  WebRtcSocketWrapper receiver_socket, sender_socket;
+  const PeerId self_id("self_id");
+  const std::string service_id("NearbySharing");
+  LocationHint location_hint;
+  Future<bool> connected;
+  ByteArray message("message xyz");
+
+  receiver.StartAcceptingConnections(
+      service_id, self_id, location_hint,
+      {[&receiver_socket, connected](WebRtcSocketWrapper wrapper) mutable {
+        receiver_socket = wrapper;
+        connected.Set(receiver_socket.IsValid());
+      }});
+
+  CancellationFlag flag;
+  sender_socket = sender.Connect(service_id, self_id, location_hint, &flag);
+  EXPECT_TRUE(sender_socket.IsValid());
+
+  ExceptionOr<bool> devices_connected = connected.Get();
+  ASSERT_TRUE(devices_connected.ok());
+  EXPECT_TRUE(devices_connected.result());
+
+  // Only shuts down signaling channel.
+  receiver.StopAcceptingConnections(service_id);
+
+  sender_socket.GetOutputStream().Write(message);
+  ExceptionOr<ByteArray> received_msg =
+      receiver_socket.GetInputStream().Read(/*size=*/32);
+  ASSERT_TRUE(received_msg.ok());
+  EXPECT_EQ(message, received_msg.result());
+}
+
+TEST_P(WebRtcTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  MediumEnvironment::Instance().SetFeatureFlags(feature_flags);
+  WebRtc receiver, sender;
+  WebRtcSocketWrapper receiver_socket, sender_socket;
+  const PeerId self_id("self_id");
+  const std::string service_id("NearbySharing");
+  LocationHint location_hint;
+  Future<bool> connected;
+  ByteArray message("message");
+
+  receiver.StartAcceptingConnections(
+      service_id, self_id, location_hint,
+      {[&receiver_socket, connected](WebRtcSocketWrapper wrapper) mutable {
+        receiver_socket = wrapper;
+        connected.Set(receiver_socket.IsValid());
+      }});
+
+  CancellationFlag flag(true);
+  sender_socket = sender.Connect(service_id, self_id, location_hint, &flag);
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(sender_socket.IsValid());
+
+    ExceptionOr<bool> devices_connected = connected.Get();
+    ASSERT_TRUE(devices_connected.ok());
+    EXPECT_TRUE(devices_connected.result());
+
+    sender_socket.GetOutputStream().Write(message);
+    ExceptionOr<ByteArray> received_msg =
+        receiver_socket.GetInputStream().Read(/*size=*/32);
+    ASSERT_TRUE(received_msg.ok());
+    EXPECT_EQ(message, received_msg.result());
+
+    receiver_socket.Close();
+  } else {
+    EXPECT_FALSE(sender_socket.IsValid());
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedWebRtcTest, WebRtcTest,
+                         ::testing::ValuesIn(kTestCases));
+
 // Basic test to check that device is accepting connections when initialized.
 TEST_F(WebRtcTest, NotAcceptingConnections) {
   WebRtc webrtc;
@@ -241,42 +333,6 @@ TEST_F(WebRtcTest, ConnectBothDevicesAndSendData) {
   receiver_socket.Close();
 }
 
-// Tests the flow when the two devices exchange SDP messages and connect to each
-// other but the signaling channel is closed before sending the data.
-TEST_F(WebRtcTest, ConnectBothDevices_ShutdownSignaling_SendData) {
-  WebRtc receiver, sender;
-  WebRtcSocketWrapper receiver_socket, sender_socket;
-  const PeerId self_id("self_id");
-  const std::string service_id("NearbySharing");
-  LocationHint location_hint;
-  Future<bool> connected;
-  ByteArray message("message xyz");
-
-  receiver.StartAcceptingConnections(
-      service_id, self_id, location_hint,
-      {[&receiver_socket, connected](WebRtcSocketWrapper wrapper) mutable {
-        receiver_socket = wrapper;
-        connected.Set(receiver_socket.IsValid());
-      }});
-
-  CancellationFlag flag;
-  sender_socket = sender.Connect(service_id, self_id, location_hint, &flag);
-  EXPECT_TRUE(sender_socket.IsValid());
-
-  ExceptionOr<bool> devices_connected = connected.Get();
-  ASSERT_TRUE(devices_connected.ok());
-  EXPECT_TRUE(devices_connected.result());
-
-  // Only shuts down signaling channel.
-  receiver.StopAcceptingConnections(service_id);
-
-  sender_socket.GetOutputStream().Write(message);
-  ExceptionOr<ByteArray> received_msg =
-      receiver_socket.GetInputStream().Read(/*size=*/32);
-  ASSERT_TRUE(received_msg.ok());
-  EXPECT_EQ(message, received_msg.result());
-}
-
 TEST_F(WebRtcTest, Connect_NullPeerConnection) {
   using MockAcceptedCallback =
       testing::MockFunction<void(WebRtcSocketWrapper socket)>;

cpp/core/internal/mediums/wifi_lan.cc

@@ -248,7 +248,12 @@ WifiLanSocket WifiLan::Connect(WifiLanService& wifi_lan_service,
     return socket;
   }
 
-  socket = medium_.Connect(wifi_lan_service, service_id);
+  if (cancellation_flag->Cancelled()) {
+    NEARBY_LOGS(INFO) << "Can't create client WifiLan socket due to cancel.";
+    return socket;
+  }
+
+  socket = medium_.Connect(wifi_lan_service, service_id, cancellation_flag);
   if (!socket.IsValid()) {
     NEARBY_LOG(INFO, "Failed to Connect via WifiLan [service_id=%s]",
                service_id.c_str());

cpp/core/internal/mediums/wifi_lan_test.cc

@@ -29,6 +29,17 @@ namespace nearby {
 namespace connections {
 namespace {
 
+using FeatureFlags = FeatureFlags::Flags;
+
+constexpr FeatureFlags kTestCases[] = {
+    FeatureFlags{
+        .enable_cancellation_flag = true,
+    },
+    FeatureFlags{
+        .enable_cancellation_flag = false,
+    },
+};
+
 constexpr absl::Duration kWaitDuration = absl::Milliseconds(1000);
 constexpr absl::string_view kServiceID{"com.google.location.nearby.apps.test"};
 constexpr absl::string_view kServiceInfoName{
@@ -36,7 +47,7 @@ constexpr absl::string_view kServiceInfoName{
 constexpr absl::string_view kEndpointName{"Simulated endpoint name"};
 constexpr absl::string_view kEndpointInfoKey{"n"};
 
-class WifiLanTest : public ::testing::Test {
+class WifiLanTest : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   using DiscoveredServiceCallback = WifiLanMedium::DiscoveredServiceCallback;
 
@@ -45,6 +56,117 @@ class WifiLanTest : public ::testing::Test {
   MediumEnvironment& env_{MediumEnvironment::Instance()};
 };
 
+TEST_P(WifiLanTest, CanStartAcceptingConnectionsAndConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  WifiLan wifi_lan_a;
+  WifiLan wifi_lan_b;
+  std::string service_id(kServiceID);
+  std::string service_info_name{kServiceInfoName};
+  std::string endpoint_info_name{kEndpointName};
+  CountDownLatch found_latch(1);
+  CountDownLatch accept_latch(1);
+
+  NsdServiceInfo nsd_service_info;
+  nsd_service_info.SetServiceInfoName(service_info_name);
+  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
+                                endpoint_info_name);
+  wifi_lan_a.StartAdvertising(service_id, nsd_service_info);
+  wifi_lan_a.StartAcceptingConnections(
+      service_id,
+      {
+          .accepted_cb = [&accept_latch](
+                             WifiLanSocket socket,
+                             absl::string_view) { accept_latch.CountDown(); },
+      });
+  WifiLanService discovered_service;
+  wifi_lan_b.StartDiscovery(
+      service_id,
+      {
+          .service_discovered_cb =
+              [&found_latch, &discovered_service](
+                  WifiLanService& service, absl::string_view service_id) {
+                discovered_service = service;
+                NEARBY_LOG(INFO, "Discovered service=%p [impl=%p]", &service,
+                           &service.GetImpl());
+                found_latch.CountDown();
+              },
+      });
+
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  ASSERT_TRUE(discovered_service.IsValid());
+  CancellationFlag flag;
+  WifiLanSocket socket =
+      wifi_lan_b.Connect(discovered_service, service_id, &flag);
+  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+  EXPECT_TRUE(socket.IsValid());
+  wifi_lan_b.StopDiscovery(service_id);
+  wifi_lan_a.StopAcceptingConnections(service_id);
+  wifi_lan_a.StopAdvertising(service_id);
+  env_.Stop();
+}
+
+TEST_P(WifiLanTest, CanCancelConnect) {
+  FeatureFlags feature_flags = GetParam();
+  env_.SetFeatureFlags(feature_flags);
+  env_.Start();
+  WifiLan wifi_lan_a;
+  WifiLan wifi_lan_b;
+  std::string service_id(kServiceID);
+  std::string service_info_name{kServiceInfoName};
+  std::string endpoint_info_name{kEndpointName};
+  CountDownLatch found_latch(1);
+  CountDownLatch accept_latch(1);
+
+  NsdServiceInfo nsd_service_info;
+  nsd_service_info.SetServiceInfoName(service_info_name);
+  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
+                                endpoint_info_name);
+  wifi_lan_a.StartAdvertising(service_id, nsd_service_info);
+  wifi_lan_a.StartAcceptingConnections(
+      service_id,
+      {
+          .accepted_cb = [&accept_latch](
+                             WifiLanSocket socket,
+                             absl::string_view) { accept_latch.CountDown(); },
+      });
+  WifiLanService discovered_service;
+  wifi_lan_b.StartDiscovery(
+      service_id,
+      {
+          .service_discovered_cb =
+              [&found_latch, &discovered_service](
+                  WifiLanService& service, absl::string_view service_id) {
+                discovered_service = service;
+                NEARBY_LOG(INFO, "Discovered service=%p [impl=%p]", &service,
+                           &service.GetImpl());
+                found_latch.CountDown();
+              },
+      });
+
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  ASSERT_TRUE(discovered_service.IsValid());
+  CancellationFlag flag(true);
+  WifiLanSocket socket =
+      wifi_lan_b.Connect(discovered_service, service_id, &flag);
+  // If FeatureFlag is disabled, Cancelled is false as no-op.
+  if (!feature_flags.enable_cancellation_flag) {
+    EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_TRUE(socket.IsValid());
+  } else {
+    EXPECT_FALSE(accept_latch.Await(kWaitDuration).result());
+    EXPECT_FALSE(socket.IsValid());
+  }
+  wifi_lan_b.StopDiscovery(service_id);
+  wifi_lan_a.StopAcceptingConnections(service_id);
+  wifi_lan_a.StopAdvertising(service_id);
+  env_.Stop();
+}
+
+INSTANTIATE_TEST_SUITE_P(ParametrisedWifiLanTest, WifiLanTest,
+                         ::testing::ValuesIn(kTestCases));
+
 TEST_F(WifiLanTest, CanConstructValidObject) {
   env_.Start();
   WifiLan wifi_lan_a;
@@ -121,55 +243,6 @@ TEST_F(WifiLanTest, CanStartDiscovery) {
   env_.Stop();
 }
 
-TEST_F(WifiLanTest, CanStartAcceptingConnectionsAndConnect) {
-  env_.Start();
-  WifiLan wifi_lan_a;
-  WifiLan wifi_lan_b;
-  std::string service_id(kServiceID);
-  std::string service_info_name{kServiceInfoName};
-  std::string endpoint_info_name{kEndpointName};
-  CountDownLatch found_latch(1);
-  CountDownLatch accept_latch(1);
-
-  NsdServiceInfo nsd_service_info;
-  nsd_service_info.SetServiceInfoName(service_info_name);
-  nsd_service_info.SetTxtRecord(std::string(kEndpointInfoKey),
-                                endpoint_info_name);
-  wifi_lan_a.StartAdvertising(service_id, nsd_service_info);
-  wifi_lan_a.StartAcceptingConnections(
-      service_id,
-      {
-          .accepted_cb = [&accept_latch](
-                             WifiLanSocket socket,
-                             absl::string_view) { accept_latch.CountDown(); },
-      });
-  WifiLanService discovered_service;
-  wifi_lan_b.StartDiscovery(
-      service_id,
-      {
-          .service_discovered_cb =
-              [&found_latch, &discovered_service](
-                  WifiLanService& service, absl::string_view service_id) {
-                discovered_service = service;
-                NEARBY_LOG(INFO, "Discovered service=%p [impl=%p]", &service,
-                           &service.GetImpl());
-                found_latch.CountDown();
-              },
-      });
-
-  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
-  ASSERT_TRUE(discovered_service.IsValid());
-  CancellationFlag flag;
-  WifiLanSocket socket =
-      wifi_lan_b.Connect(discovered_service, service_id, &flag);
-  EXPECT_TRUE(accept_latch.Await(kWaitDuration).result());
-  EXPECT_TRUE(socket.IsValid());
-  wifi_lan_b.StopDiscovery(service_id);
-  wifi_lan_a.StopAcceptingConnections(service_id);
-  wifi_lan_a.StopAdvertising(service_id);
-  env_.Stop();
-}
-
 }  // namespace
 }  // namespace connections
 }  // namespace nearby