From 1346571e4f3e7c4a68e92d3c72797022320cc39e Mon Sep 17 00:00:00 2001 From: Francis Tsui Date: Thu, 7 May 2026 10:16:42 -0700 Subject: [PATCH] Fix UAF in WorkerQueue.Stop(). PiperOrigin-RevId: 912027304 --- sharing/BUILD | 1 + sharing/worker_queue.h | 50 ++++++++++++++++++++++++++++-------- sharing/worker_queue_test.cc | 43 ++++++++++++++++++++++++++----- 3 files changed, 76 insertions(+), 18 deletions(-) diff --git a/sharing/BUILD b/sharing/BUILD index d2c16191..ee2baf42 100644 --- a/sharing/BUILD +++ b/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", ], ) diff --git a/sharing/worker_queue.h b/sharing/worker_queue.h index 0fe8a928..b9c4d958 100644 --- a/sharing/worker_queue.h +++ b/sharing/worker_queue.h @@ -16,6 +16,7 @@ #define THIRD_PARTY_NEARBY_SHARING_WORKER_QUEUE_H_ #include +#include #include #include @@ -38,7 +39,11 @@ namespace nearby::sharing { template 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()) { + 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 is_stopped; + absl::AnyInvocable 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 callback_; + std::shared_ptr run_data_; // Tracks whether Start() has been called. std::atomic is_started_ = false; // Tracks whether Stop() has been called. - std::atomic is_stopped_ = false; absl::Mutex mutex_; std::queue queue_ ABSL_GUARDED_BY(mutex_); // This is used track whether the callback is already scheduled so as to avoid diff --git a/sharing/worker_queue_test.cc b/sharing/worker_queue_test.cc index e9b6e3b0..a3671d61 100644 --- a/sharing/worker_queue_test.cc +++ b/sharing/worker_queue_test.cc @@ -14,10 +14,12 @@ #include "sharing/worker_queue.h" +#include #include #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 queue(&task_runner); - queue.Queue(1); - queue.Queue(2); + auto queue = std::make_unique>(&task_runner); + queue->Queue(1); + queue->Queue(2); absl::Notification notification; - EXPECT_TRUE(queue.Start([&queue, ¬ification]() { - std::queue items = queue.ReadAll(); + auto queue_ptr = queue.get(); + EXPECT_TRUE(queue->Start([queue_ptr, ¬ification]() { + std::queue 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>(&task_runner); + absl::Notification notification1; + absl::Notification notification2; + auto queue_ptr = queue.get(); + EXPECT_TRUE(queue->Start([queue_ptr, ¬ification1, ¬ification2]() { + std::queue 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. }