Fixed a potential crash in BLE scanning

PiperOrigin-RevId: 663900592
2026-09-14 14:46:12 -04:00 · 2024-08-16 15:59:28 -07:00
parent 4088e5904f
commit f1a1c39177
3 changed files with 238 additions and 53 deletions
@@ -637,51 +637,24 @@ void DiscoveredPeripheralTracker::HandleAdvertisementHeader(
                           << absl::BytesToHexString(
                                  ByteArray(advertisement_header).data())
                           << " in thread";
-      ByteArray advertisement_data{advertisement_header};
-      if (fetching_advertisements_.contains(advertisement_data)) {
+
+      if (!fetching_advertisements_.insert(advertisement_header).second) {
        NEARBY_LOGS(VERBOSE) << ": Ignore the advertisement header due to it "
                                "is already in fetching.";
        return;
      }

-      fetching_advertisements_.insert(advertisement_data);
-
      if (executor_ == nullptr) {
        // The situation happens when flag value changed
        executor_ = std::make_unique<MultiThreadExecutor>(kGattThreadCount);
      }
-      executor_->Execute([this, peripheral, advertisement_header,
-                          advertisement_fetcher =
-                              std::move(advertisement_fetcher),
-                          advertisement_data =
-                              std::move(advertisement_data)]() mutable {
-        {
-          MutexLock lock(&mutex_);
-          if (!IsInterestingAdvertisementHeader(advertisement_header)) {
-            NEARBY_LOGS(INFO)
-                << ": Ignore to read raw advertisement from server due to it "
-                   "is not interesting header now.";
-            fetching_advertisements_.erase(advertisement_data);
-            return;
-          }
-        }
-
-        std::vector<const ByteArray*> gatt_advertisement_bytes_list =
+      executor_->Execute(
+          [this, peripheral, advertisement_header,
+           advertisement_fetcher = std::move(advertisement_fetcher),
+           advertisement_data = std::move(advertisement_data)]() mutable {
            FetchRawAdvertisementsInThread(peripheral, advertisement_header,
                                           std::move(advertisement_fetcher));
-        {
-          MutexLock lock(&mutex_);
-          HandleRawGattAdvertisements(peripheral, advertisement_header,
-                                      gatt_advertisement_bytes_list,
-                                      /*service_uuid=*/{});
-          UpdateCommonStateForFoundBleAdvertisement(advertisement_header);
-          fetching_advertisements_.erase(advertisement_data);
-          NEARBY_LOGS(VERBOSE)
-              << ": Completed to handle GATT advertisement "
-              << absl::BytesToHexString(ByteArray(advertisement_header).data())
-              << " in thread";
-        }
-      });
+          });
      return;
    } else {
      std::vector<const ByteArray*> gatt_advertisement_bytes_list =
@@ -796,8 +769,7 @@ DiscoveredPeripheralTracker::FetchRawAdvertisements(
  std::transform(service_id_infos_.begin(), service_id_infos_.end(),
                 std::back_inserter(service_ids),
                 [](auto& kv) { return kv.first; });
-  advertisement_fetcher(std::move(peripheral),
-                        advertisement_header.GetNumSlots(),
+  advertisement_fetcher(peripheral, advertisement_header.GetNumSlots(),
                        advertisement_header.GetPsm(), service_ids, *result);

  // Take those results and return all the advertisements we were able to
@@ -805,33 +777,63 @@ DiscoveredPeripheralTracker::FetchRawAdvertisements(
  return result->GetAdvertisements();
 }

-std::vector<const ByteArray*>
-DiscoveredPeripheralTracker::FetchRawAdvertisementsInThread(
+void DiscoveredPeripheralTracker::FetchRawAdvertisementsInThread(
    BleV2Peripheral peripheral,
    const BleAdvertisementHeader& advertisement_header,
    AdvertisementFetcher advertisement_fetcher) {
  std::vector<std::string> service_ids;
-  AdvertisementReadResult* result = nullptr;
+
  {
    MutexLock lock(&mutex_);
-    // Fetch the raw GATT advertisements and store the results.
-    auto& read_result = advertisement_read_results_[advertisement_header];
-    if (read_result == nullptr) {
-      read_result = std::make_unique<mediums::AdvertisementReadResult>();
+    if (!IsInterestingAdvertisementHeader(advertisement_header)) {
+      NEARBY_LOGS(INFO)
+          << ": Ignore to read raw advertisement from server due to it "
+             "is not interesting header now.";
+      fetching_advertisements_.erase(advertisement_header);
+      return;
    }

-    result = read_result.get();
    std::transform(service_id_infos_.begin(), service_id_infos_.end(),
                   std::back_inserter(service_ids),
                   [](auto& kv) { return kv.first; });
  }
-  advertisement_fetcher(std::move(peripheral),
-                        advertisement_header.GetNumSlots(),
-                        advertisement_header.GetPsm(), service_ids, *result);

-  // Take those results and return all the advertisements we were able to
-  // read.
-  return result->GetAdvertisements();
+  auto result = std::make_unique<mediums::AdvertisementReadResult>();
+  advertisement_fetcher(peripheral, advertisement_header.GetNumSlots(),
+                        advertisement_header.GetPsm(), service_ids, *result);
+  {
+    MutexLock lock(&mutex_);
+    // The fetching process might take a few seconds, and tracking settings
+    // could change during that time. We need to double-check if the result
+    // is still valid afterward.
+    if (!IsInterestingAdvertisementHeader(advertisement_header)) {
+      NEARBY_LOGS(WARNING)
+          << ": Ignore the fetched GATT advertisement from server due to it "
+             "is not interesting header now.";
+      return;
+    }
+
+    auto it = advertisement_read_results_.insert_or_assign(advertisement_header,
+                                                           std::move(result));
+    std::vector<const ByteArray*> gatt_advertisement_bytes_list =
+        it.first->second->GetAdvertisements();
+
+    if (gatt_advertisement_bytes_list.empty() ||
+        !IsInterestingAdvertisementHeader(advertisement_header)) {
+      fetching_advertisements_.erase(advertisement_header);
+      return;
+    }
+
+    HandleRawGattAdvertisements(peripheral, advertisement_header,
+                                gatt_advertisement_bytes_list,
+                                /*service_uuid=*/{});
+    UpdateCommonStateForFoundBleAdvertisement(advertisement_header);
+    fetching_advertisements_.erase(advertisement_header);
+    NEARBY_LOGS(VERBOSE) << ": Completed to handle GATT advertisement "
+                         << absl::BytesToHexString(
+                                ByteArray(advertisement_header).data())
+                         << " in thread";
+  }
 }

 void DiscoveredPeripheralTracker::UpdateCommonStateForFoundBleAdvertisement(
@@ -241,7 +241,7 @@ class DiscoveredPeripheralTracker {
      AdvertisementFetcher advertisement_fetcher)
      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);

-  std::vector<const ByteArray*> FetchRawAdvertisementsInThread(
+  void FetchRawAdvertisementsInThread(
      BleV2Peripheral peripheral,
      const BleAdvertisementHeader& advertisement_header,
      AdvertisementFetcher advertisement_fetcher);
@@ -319,7 +319,7 @@ class DiscoveredPeripheralTracker {
      gatt_advertisement_infos_ ABSL_GUARDED_BY(mutex_);

  // Tracks the advertisements in GATT fetching.
-  absl::flat_hash_set<ByteArray> fetching_advertisements_
+  absl::flat_hash_set<BleAdvertisementHeader> fetching_advertisements_
      ABSL_GUARDED_BY(mutex_);

  std::unique_ptr<MultiThreadExecutor> executor_ ABSL_GUARDED_BY(mutex_) =
@@ -24,6 +24,7 @@
 #include "gtest/gtest.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/strings/string_view.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "connections/implementation/flags/nearby_connections_feature_flags.h"
 #include "connections/implementation/mediums/ble_v2/advertisement_read_result.h"
@@ -165,7 +166,10 @@ class DiscoveredPeripheralTrackerTest : public testing::Test {
    ble_central_ = std::make_unique<BleV2Medium>(*adapter_central_);
  }

-  void TearDown() override { MediumEnvironment::Instance().Stop(); }
+  void TearDown() override {
+    MediumEnvironment::Instance().Stop();
+    NearbyFlags::GetInstance().ResetOverridedValues();
+  }

  BleV2Peripheral CreateBlePeripheral() {
    return ble_central_->GetRemotePeripheral(
@@ -196,6 +200,17 @@ class DiscoveredPeripheralTrackerTest : public testing::Test {
        GetAdvertisementFetcher(fetch_latch, advertisement_bytes_list));
  }

+  void FindAdvertisementWithSlowFetcher(
+      const api::ble_v2::BleAdvertisementData& advertisement_data,
+      const std::vector<ByteArray>& advertisement_bytes_list,
+      CountDownLatch& fetch_latch) {
+    BleV2Peripheral peripheral = CreateBlePeripheral();
+
+    discovered_peripheral_tracker_.ProcessFoundBleAdvertisement(
+        peripheral, advertisement_data,
+        GetSlowAdvertisementFetcher(fetch_latch, advertisement_bytes_list));
+  }
+
  int GetFetchAdvertisementCallbackCount() const {
    MutexLock lock(&mutex_);
    return fetch_count_;
@@ -214,6 +229,13 @@ class DiscoveredPeripheralTrackerTest : public testing::Test {
        true);
  }

+  void EnableFetchGattAdvertisementInThread() {
+    NearbyFlags::GetInstance().OverrideBoolFlagValue(
+        config_package_nearby::nearby_connections_feature::
+            kEnableGattQueryInThread,
+        true);
+  }
+
 protected:
  // A stub Advertisement fetcher.
  DiscoveredPeripheralTracker::AdvertisementFetcher GetAdvertisementFetcher(
@@ -235,6 +257,27 @@ class DiscoveredPeripheralTrackerTest : public testing::Test {
    };
  }

+  DiscoveredPeripheralTracker::AdvertisementFetcher GetSlowAdvertisementFetcher(
+      CountDownLatch& fetch_latch,
+      const std::vector<ByteArray>& advertisement_bytes_list) {
+    return [this, &fetch_latch, advertisement_bytes_list](
+               BleV2Peripheral peripheral, int num_slots, int psm,
+               const std::vector<std::string>& interesting_service_ids,
+               mediums::AdvertisementReadResult& advertisement_read_result) {
+      MutexLock lock(&mutex_);
+      fetch_count_++;
+      int slot = 0;
+      // In real environment, the GATT fetch may run about 3-5 seconds.
+      absl::SleepFor(absl::Milliseconds(200));
+      for (const auto& advertisement_bytes : advertisement_bytes_list) {
+        advertisement_read_result.AddAdvertisement(slot++, advertisement_bytes);
+      }
+      advertisement_read_result.RecordLastReadStatus(
+          /*is_success=*/true);
+      fetch_latch.CountDown();
+    };
+  }
+
  std::unique_ptr<BluetoothAdapter> adapter_peripheral_;
  std::unique_ptr<BluetoothAdapter> adapter_central_;
  std::unique_ptr<BleV2Medium> ble_peripheral_;
@@ -1401,6 +1444,146 @@ TEST_F(DiscoveredPeripheralTrackerTest, SkipDummyAdvertisement) {
  EXPECT_EQ(GetFetchAdvertisementCallbackCount(), 0);
 }

+TEST_F(DiscoveredPeripheralTrackerTest, FetchGattAdvertisementInThread) {
+  EnableFetchGattAdvertisementInThread();
+  std::vector<std::string> service_ids = {std::string(kServiceIdA)};
+  ByteArray advertisement_header_bytes = CreateBleAdvertisementHeader(
+      GenerateRandomAdvertisementHash(), service_ids);
+  ByteArray advertisement_bytes = CreateBleAdvertisement(
+      std::string(kServiceIdA), ByteArray(std::string(kData)),
+      ByteArray(std::string(kDeviceToken)));
+  CountDownLatch found_latch(1);
+  CountDownLatch fetch_latch(1);
+
+  discovered_peripheral_tracker_.StartTracking(
+      std::string(kServiceIdA),
+      {
+          .peripheral_discovered_cb =
+              [&found_latch](BleV2Peripheral peripheral,
+                             const std::string& service_id,
+                             const ByteArray& advertisement_bytes,
+                             bool fast_advertisement) {
+                EXPECT_EQ(advertisement_bytes, ByteArray(std::string(kData)));
+                EXPECT_FALSE(fast_advertisement);
+                found_latch.CountDown();
+              },
+      },
+      {});
+
+  api::ble_v2::BleAdvertisementData advertisement_data{};
+  if (!advertisement_header_bytes.Empty()) {
+    advertisement_data.service_data.insert(
+        {bleutils::kCopresenceServiceUuid, advertisement_header_bytes});
+  }
+
+  FindAdvertisementWithSlowFetcher(advertisement_data, {advertisement_bytes},
+                                   fetch_latch);
+
+  // We should receive a client callback of a peripheral discovery.
+  fetch_latch.Await(kWaitDuration);
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  EXPECT_EQ(GetFetchAdvertisementCallbackCount(), 1);
+}
+
+TEST_F(DiscoveredPeripheralTrackerTest,
+       IgnoreGattAdvertisementResultWhentrackingStoppedInThread) {
+  EnableFetchGattAdvertisementInThread();
+  std::vector<std::string> service_ids = {std::string(kServiceIdA)};
+  ByteArray advertisement_header_bytes = CreateBleAdvertisementHeader(
+      GenerateRandomAdvertisementHash(), service_ids);
+  ByteArray advertisement_bytes = CreateBleAdvertisement(
+      std::string(kServiceIdA), ByteArray(std::string(kData)),
+      ByteArray(std::string(kDeviceToken)));
+  CountDownLatch found_latch(1);
+  CountDownLatch fetch_latch(1);
+
+  discovered_peripheral_tracker_.StartTracking(
+      std::string(kServiceIdA),
+      {
+          .peripheral_discovered_cb =
+              [&found_latch](BleV2Peripheral peripheral,
+                             const std::string& service_id,
+                             const ByteArray& advertisement_bytes,
+                             bool fast_advertisement) {
+                EXPECT_EQ(advertisement_bytes, ByteArray(std::string(kData)));
+                EXPECT_FALSE(fast_advertisement);
+                found_latch.CountDown();
+              },
+      },
+      {});
+
+  api::ble_v2::BleAdvertisementData advertisement_data{};
+  if (!advertisement_header_bytes.Empty()) {
+    advertisement_data.service_data.insert(
+        {bleutils::kCopresenceServiceUuid, advertisement_header_bytes});
+  }
+
+  FindAdvertisementWithSlowFetcher(advertisement_data, {advertisement_bytes},
+                                   fetch_latch);
+
+  // We should receive a client callback of a peripheral discovery.
+  absl::SleepFor(absl::Milliseconds(20));
+  discovered_peripheral_tracker_.StopTracking(std::string(kServiceIdA));
+  fetch_latch.Await(kWaitDuration);
+  EXPECT_FALSE(found_latch.Await(kWaitDuration).result());
+  EXPECT_EQ(GetFetchAdvertisementCallbackCount(), 1);
+}
+
+TEST_F(DiscoveredPeripheralTrackerTest,
+       FetchMultipleGattAdvertisementResultsInThread) {
+  EnableFetchGattAdvertisementInThread();
+  std::vector<std::string> service_ids = {std::string(kServiceIdA)};
+  ByteArray advertisement_header_bytes = CreateBleAdvertisementHeader(
+      GenerateRandomAdvertisementHash(), service_ids);
+  ByteArray advertisement_header_bytes_2 = CreateBleAdvertisementHeader(
+      GenerateRandomAdvertisementHash(), service_ids);
+  ByteArray advertisement_bytes = CreateBleAdvertisement(
+      std::string(kServiceIdA), ByteArray(std::string(kData)),
+      ByteArray(std::string(kDeviceToken)));
+  ByteArray advertisement_bytes_2 = CreateBleAdvertisement(
+      std::string(kServiceIdA), ByteArray(std::string(kData2)),
+      ByteArray(std::string(kDeviceToken)));
+  CountDownLatch found_latch(2);
+  CountDownLatch fetch_latch(2);
+
+  discovered_peripheral_tracker_.StartTracking(
+      std::string(kServiceIdA),
+      {
+          .peripheral_discovered_cb =
+              [&found_latch](BleV2Peripheral peripheral,
+                             const std::string& service_id,
+                             const ByteArray& advertisement_bytes,
+                             bool fast_advertisement) {
+                EXPECT_FALSE(fast_advertisement);
+                found_latch.CountDown();
+              },
+      },
+      {});
+
+  api::ble_v2::BleAdvertisementData advertisement_data{};
+  if (!advertisement_header_bytes.Empty()) {
+    advertisement_data.service_data.insert(
+        {bleutils::kCopresenceServiceUuid, advertisement_header_bytes});
+  }
+
+  FindAdvertisementWithSlowFetcher(advertisement_data, {advertisement_bytes},
+                                   fetch_latch);
+
+  api::ble_v2::BleAdvertisementData advertisement_data_2{};
+  if (!advertisement_header_bytes_2.Empty()) {
+    advertisement_data_2.service_data.insert(
+        {bleutils::kCopresenceServiceUuid, advertisement_header_bytes_2});
+  }
+
+  FindAdvertisementWithSlowFetcher(advertisement_data_2,
+                                   {advertisement_bytes_2}, fetch_latch);
+
+  // We should receive a client callback of a peripheral discovery.
+  fetch_latch.Await(kWaitDuration);
+  EXPECT_TRUE(found_latch.Await(kWaitDuration).result());
+  EXPECT_EQ(GetFetchAdvertisementCallbackCount(), 2);
+}
+
 }  // namespace

 }  // namespace mediums