[PA] Purge Scheduler-Loop Quarantine on UI Thread Task Completion

This change introduces an observer to the UI thread's task queues. The
observer purges the Scheduler-Loop Quarantine after each task completes.
The Scheduler-Loop Quarantine is part of the PA/AC (PartitionAlloc with
Advanced Checks) experiment, which aims to mitigate use-after-free (UaF)
vulnerabilities by quarantining objects freed on the Browser UI thread.
By purging the quarantine after task completion, when stack memory is
likely empty, this patch should mitigate the memory regression caused by
the PA/AC experiment.

Bug: 329027914, 351974425
Change-Id: I0193f986bb8eb064914c7bc36ab3db1fd19c6f46
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5528149
Reviewed-by: Arthur Sonzogni <[email protected]>
Reviewed-by: Daniel Cheng <[email protected]>
Reviewed-by: Keishi Hattori <[email protected]>
Commit-Queue: Mikihito Matsuura <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1496114}
NOKEYCHECK=True
GitOrigin-RevId: 541ac7163f78ec14b1e0afc52145d70f2b44b783

Author: mikt

src/partition_alloc/scheduler_loop_quarantine.cc

@@ -241,7 +241,7 @@ SchedulerLoopQuarantineBranch<thread_bound>::PurgeInternal(
       }
     }
 
-    freed_count++;
+    ++freed_count;
     PA_DCHECK(slot_size > 0);
     freed_size_in_bytes += slot_size;
     branch_size_in_bytes_ -= slot_size;
@@ -254,6 +254,31 @@ SchedulerLoopQuarantineBranch<thread_bound>::PurgeInternal(
   root_->count_.fetch_sub(freed_count, std::memory_order_relaxed);
 }
 
+template <bool thread_bound>
+void SchedulerLoopQuarantineBranch<thread_bound>::AllowScanlessPurge() {
+  PA_DCHECK(kThreadBound);
+  // Always thread-bound; no need to lock.
+  FakeScopedGuard guard(lock_);
+
+  PA_CHECK(disallow_scanless_purge_ > 0);
+  --disallow_scanless_purge_;
+  if (disallow_scanless_purge_ == 0) {
+    // Now scanless purge is allowed. Purging at this timing is more performance
+    // efficient.
+    PurgeInternal(0);
+  }
+}
+
+template <bool thread_bound>
+void SchedulerLoopQuarantineBranch<thread_bound>::DisallowScanlessPurge() {
+  PA_DCHECK(kThreadBound);
+  // Always thread-bound; no need to lock.
+  FakeScopedGuard guard(lock_);
+
+  ++disallow_scanless_purge_;
+  PA_CHECK(disallow_scanless_purge_ > 0);  // Overflow check.
+}
+
 template <bool thread_bound>
 const SchedulerLoopQuarantineConfig&
 SchedulerLoopQuarantineBranch<thread_bound>::GetConfigurationForTesting() {

src/partition_alloc/scheduler_loop_quarantine.h

@@ -156,6 +156,9 @@ class SchedulerLoopQuarantineBranch {
                   SlotSpanMetadata* slot_span,
                   uintptr_t slot_start) PA_LOCKS_EXCLUDED(lock_);
 
+  void AllowScanlessPurge();
+  void DisallowScanlessPurge();
+
   const SchedulerLoopQuarantineConfig& GetConfigurationForTesting();
 
   class ScopedQuarantineExclusion {
@@ -223,6 +226,15 @@ class SchedulerLoopQuarantineBranch {
   // Using `std::atomic` here so that other threads can update this value.
   std::atomic_size_t branch_capacity_in_bytes_ = 0;
 
+  // TODO(http://crbug.com/329027914): Implement stack scanning, to be performed
+  // when this value is non-zero.
+  //
+  // Currently, a scanless purge is always performed. However, this value is
+  // still used as a hint to determine safer purge timings for memory
+  // optimization.
+  uint32_t disallow_scanless_purge_ PA_GUARDED_BY(lock_) = 0;
+
+  // Debug and testing data.
 #if PA_BUILDFLAG(DCHECKS_ARE_ON)
   std::atomic_bool being_destructed_ = false;
 #endif  // PA_BUILDFLAG(DCHECKS_ARE_ON)

src/partition_alloc/scheduler_loop_quarantine_support.cc

@@ -21,6 +21,50 @@ ScopedSchedulerLoopQuarantineExclusion::
 ScopedSchedulerLoopQuarantineExclusion::
     ~ScopedSchedulerLoopQuarantineExclusion() {}
 
+SchedulerLoopQuarantineScanPolicyUpdater::
+    SchedulerLoopQuarantineScanPolicyUpdater() = default;
+
+SchedulerLoopQuarantineScanPolicyUpdater::
+    ~SchedulerLoopQuarantineScanPolicyUpdater() {
+  // Ensure all `DisallowScanlessPurge()` calls were followed by
+  // `AllowScanlessPurge()`.
+  PA_CHECK(disallow_scanless_purge_calls_ == 0);
+}
+
+void SchedulerLoopQuarantineScanPolicyUpdater::DisallowScanlessPurge() {
+  disallow_scanless_purge_calls_++;
+  PA_CHECK(0 < disallow_scanless_purge_calls_);  // Overflow check.
+
+  auto* branch = GetQuarantineBranch();
+  PA_CHECK(branch);
+  branch->DisallowScanlessPurge();
+}
+
+void SchedulerLoopQuarantineScanPolicyUpdater::AllowScanlessPurge() {
+  PA_CHECK(0 < disallow_scanless_purge_calls_);
+  disallow_scanless_purge_calls_--;
+
+  auto* branch = GetQuarantineBranch();
+  PA_CHECK(branch);
+  branch->AllowScanlessPurge();
+}
+
+internal::ThreadBoundSchedulerLoopQuarantineBranch*
+SchedulerLoopQuarantineScanPolicyUpdater::GetQuarantineBranch() {
+  ThreadCache* tcache = ThreadCache::EnsureAndGet();
+  if (!ThreadCache::IsValid(tcache)) {
+    return nullptr;
+  }
+
+  uintptr_t current_tcache_addr = reinterpret_cast<uintptr_t>(tcache);
+  if (tcache_address_ == 0) {
+    tcache_address_ = current_tcache_addr;
+  } else {
+    PA_CHECK(tcache_address_ == current_tcache_addr);
+  }
+  return &tcache->GetSchedulerLoopQuarantineBranch();
+}
+
 namespace internal {
 ScopedSchedulerLoopQuarantineBranchAccessorForTesting::
     ScopedSchedulerLoopQuarantineBranchAccessorForTesting(

src/partition_alloc/scheduler_loop_quarantine_support.h

@@ -10,11 +10,13 @@
 #ifndef PARTITION_ALLOC_SCHEDULER_LOOP_QUARANTINE_SUPPORT_H_
 #define PARTITION_ALLOC_SCHEDULER_LOOP_QUARANTINE_SUPPORT_H_
 
+#include <map>
 #include <optional>
 #include <variant>
 
 #include "partition_alloc/build_config.h"
 #include "partition_alloc/buildflags.h"
+#include "partition_alloc/partition_alloc_base/compiler_specific.h"
 #include "partition_alloc/partition_root.h"
 #include "partition_alloc/scheduler_loop_quarantine.h"
 #include "partition_alloc/thread_cache.h"
@@ -39,6 +41,49 @@ class PA_COMPONENT_EXPORT(PARTITION_ALLOC)
       instance_;
 };
 
+// An utility class to update Scheduler-Loop Quarantine's purging strategy for
+// the current thread. By default it uses "scanless" purge for best performance.
+// However, it also supports stack-scanning before purging to verify there is no
+// dangling pointer in stack memory. Stack-scanning comes with some performance
+// cost, but there is security benefit. This class can be used to switch between
+// these two strategies dynamically.
+// An example usage is to allow scanless purge only around "stack bottom".
+// We can safely assume there is no dangling pointer if stack memory is barely
+// used thus safe to purge quarantine.
+// At Chrome layer it is task execution and we expect
+// `DisallowScanlessPurge()` to be called before task execution and
+// `AllowScanlessPurge()` after. Since there is no unified way to hook
+// task execution in Chrome, we provide an abstract utility here.
+// This class is not thread-safe.
+//
+// TODO(http://crbug.com/329027914): stack-scanning is not implemented yet
+// and this class is effectively "disallow any purge unless really needed".
+// It still gives some hints on purging timing for memory efficiency.
+class PA_COMPONENT_EXPORT(PARTITION_ALLOC)
+    SchedulerLoopQuarantineScanPolicyUpdater {
+ public:
+  SchedulerLoopQuarantineScanPolicyUpdater();
+  ~SchedulerLoopQuarantineScanPolicyUpdater();
+
+  // Disallows scanless purge and performs stack-scanning when needed.
+  // Can be called multiple times, but each call to this function must be
+  // followed by `AllowScanlessPurge()`.
+  void DisallowScanlessPurge();
+
+  // Re-activate scanless purge. `DisallowScanlessPurge()` must be called prior
+  // to use of this function. This may trigger purge immediately.
+  void AllowScanlessPurge();
+
+ private:
+  PA_ALWAYS_INLINE internal::ThreadBoundSchedulerLoopQuarantineBranch*
+  GetQuarantineBranch();
+
+  uint32_t disallow_scanless_purge_calls_ = 0;
+
+  // An address of `ThreadCache` instance works as a thread ID.
+  uintptr_t tcache_address_ = 0;
+};
+
 namespace internal {
 
 class PA_COMPONENT_EXPORT(PARTITION_ALLOC)