Fix UAF in WorkerQueue.Stop().

PiperOrigin-RevId: 912027304

sharing/BUILD

@@ -1014,6 +1014,7 @@ cc_test(
         "//internal/test",
         "@com_github_protobuf_matchers//protobuf-matchers",
         "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
     ],
 )

sharing/worker_queue.h

@@ -16,6 +16,7 @@
 #define THIRD_PARTY_NEARBY_SHARING_WORKER_QUEUE_H_
 
 #include <atomic>
+#include <memory>
 #include <queue>
 #include <utility>
 
@@ -38,7 +39,11 @@ namespace nearby::sharing {
 template <typename T>
 class WorkerQueue {
  public:
-  explicit WorkerQueue(TaskRunner* task_runner) : task_runner_(task_runner) {}
+  explicit WorkerQueue(TaskRunner* task_runner)
+      : task_runner_(task_runner),
+        run_data_(std::make_shared<RunData>()) {
+    run_data_->is_stopped = false;
+  }
 
   ~WorkerQueue() { Stop(); }
 
@@ -52,11 +57,11 @@ class WorkerQueue {
       LOG(ERROR) << "WorkerQueue is already started.";
       return false;
     }
-    if (is_stopped_) {
+    if (run_data_->is_stopped) {
       LOG(ERROR) << "WorkerQueue is already stopped, cannot restart.";
       return false;
     }
-    callback_ = std::move(callback);
+    run_data_->callback = std::move(callback);
     {
       absl::MutexLock lock(mutex_);
       if (!queue_.empty()) {
@@ -67,12 +72,22 @@ class WorkerQueue {
   }
 
   // Stops the queue.  No new callback will be scheduled.
+  // This method will block until the callback finishes if it is currently
+  // running.
   void Stop() {
-    bool already_stopped = is_stopped_.exchange(true);
+    bool already_stopped = run_data_->is_stopped.exchange(true);
     if (already_stopped || !is_started_) {
       return;
     }
+    // Prevent new callbacks from being scheduled.
     is_scheduled_ = true;
+    // Wait for inflight callback to finish.
+    absl::MutexLock lock(run_data_->running_mutex);
+    auto stopped_running =
+        [this]() ABSL_EXCLUSIVE_LOCKS_REQUIRED(run_data_->running_mutex) {
+          return !run_data_->is_running;
+        };
+    run_data_->running_mutex.Await(absl::Condition(&stopped_running));
   }
 
   // Queues an item to be processed by the callback.
@@ -95,9 +110,16 @@ class WorkerQueue {
   }
 
  private:
+  struct RunData {
+    std::atomic<bool> is_stopped;
+    absl::AnyInvocable<void()> callback;
+    absl::Mutex running_mutex;
+    bool is_running ABSL_GUARDED_BY(running_mutex) = false;
+  };
+
   void ScheduleCallback() {
     // Skip if not started or stopped
-    if (!is_started_ || is_stopped_) {
+    if (!is_started_ || run_data_->is_stopped) {
       return;
     }
     if (is_scheduled_.exchange(true)) {
@@ -106,20 +128,26 @@ class WorkerQueue {
       return;
     }
     VLOG(1) << "Scheduling callback";
-    task_runner_->PostTask([this]() {
-      if (is_stopped_) {
-        return;
+    task_runner_->PostTask([run_data = run_data_]() {
+      {
+        absl::MutexLock lock(run_data->running_mutex);
+        run_data->is_running = true;
+      }
+      if (!run_data->is_stopped) {
+        run_data->callback();
+      }
+      {
+        absl::MutexLock lock(run_data->running_mutex);
+        run_data->is_running = false;
       }
-      callback_();
     });
   }
 
   TaskRunner* const task_runner_ = nullptr;
-  absl::AnyInvocable<void()> callback_;
+  std::shared_ptr<RunData> run_data_;
   // Tracks whether Start() has been called.
   std::atomic<bool> is_started_ = false;
   // Tracks whether Stop() has been called.
-  std::atomic<bool> is_stopped_ = false;
   absl::Mutex mutex_;
   std::queue<T> queue_ ABSL_GUARDED_BY(mutex_);
   // This is used track whether the callback is already scheduled so as to avoid

sharing/worker_queue_test.cc

@@ -14,10 +14,12 @@
 
 #include "sharing/worker_queue.h"
 
+#include <memory>
 #include <queue>
 
 #include "gtest/gtest.h"
 #include "absl/synchronization/notification.h"
+#include "absl/time/time.h"
 #include "internal/test/fake_clock.h"
 #include "internal/test/fake_task_runner.h"
 
@@ -91,12 +93,13 @@ TEST(WorkerQueueTest, QueueItemsWhileCallbackRunning) {
 TEST(WorkerQueueTest, StopStopsCallback) {
   FakeClock fake_clock;
   FakeTaskRunner task_runner(&fake_clock, 1);
-  WorkerQueue<int> queue(&task_runner);
-  queue.Queue(1);
-  queue.Queue(2);
+  auto queue = std::make_unique<WorkerQueue<int>>(&task_runner);
+  queue->Queue(1);
+  queue->Queue(2);
   absl::Notification notification;
-  EXPECT_TRUE(queue.Start([&queue, &notification]() {
-    std::queue<int> items = queue.ReadAll();
+  auto queue_ptr = queue.get();
+  EXPECT_TRUE(queue->Start([queue_ptr, &notification]() {
+    std::queue<int> items = queue_ptr->ReadAll();
     EXPECT_EQ(items.size(), 2);
     EXPECT_EQ(items.front(), 1);
     EXPECT_EQ(items.back(), 2);
@@ -104,8 +107,34 @@ TEST(WorkerQueueTest, StopStopsCallback) {
   }));
   // Wait for the callback to start.
   notification.WaitForNotification();
-  queue.Stop();
-  queue.Queue(3);
+  queue->Stop();
+  queue->Queue(3);
+  queue.reset();
+  task_runner.Sync();
+  // No more callbacks.
+}
+
+TEST(WorkerQueueTest, StopWaitsForInFlightCallback) {
+  FakeClock fake_clock;
+  FakeTaskRunner task_runner(&fake_clock, 1);
+  auto queue = std::make_unique<WorkerQueue<int>>(&task_runner);
+  absl::Notification notification1;
+  absl::Notification notification2;
+  auto queue_ptr = queue.get();
+  EXPECT_TRUE(queue->Start([queue_ptr, &notification1, &notification2]() {
+    std::queue<int> items = queue_ptr->ReadAll();
+    notification1.Notify();
+    notification2.WaitForNotificationWithTimeout(absl::Milliseconds(500));
+    EXPECT_EQ(items.size(), 2);
+    EXPECT_EQ(items.front(), 1);
+    EXPECT_EQ(items.back(), 2);
+  }));
+  queue->Queue(1);
+  queue->Queue(2);
+  // Wait for the callback to start.
+  notification1.WaitForNotification();
+  queue->Queue(3);
+  queue.reset();
   task_runner.Sync();
   // No more callbacks.
 }