[base] Protect thread-safe refcounts against overflow.

Implementing tests for this also tickled a surprising edge case in the
way RefCounted/RefCountedThreadSafe signalled whether the refcount
should begin at 0 or 1. Previously, this was signalled using a static
data member; however, local classes (which the updated tests use) cannot
contain static data members.

As a result, this CL "minimally" updates the tagging mechanism to use a
type alias while leaving as much of the existing tagging mechanism
intact as possible.

Bug: 1238642
Change-Id: I16f3ab243745e11bbde755f826d430e85cc33a93
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3131643
Reviewed-by: Lei Zhang <thestig@chromium.org>
Owners-Override: Lei Zhang <thestig@chromium.org>
Commit-Queue: Daniel Cheng <dcheng@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1071489}

base/atomic_ref_count.h

@@ -12,6 +12,10 @@
 
 namespace base {
 
+namespace subtle {
+class RefCountedOverflowTest;
+}  // namespace subtle
+
 class AtomicRefCount {
  public:
   constexpr AtomicRefCount() : ref_count_(0) {}
@@ -61,6 +65,8 @@ class AtomicRefCount {
   }
 
  private:
+  friend subtle::RefCountedOverflowTest;
+
   std::atomic_int ref_count_;
 };
 

base/memory/ref_counted.h

@@ -15,6 +15,7 @@
 #include "base/check_op.h"
 #include "base/compiler_specific.h"
 #include "base/dcheck_is_on.h"
+// TODO(dcheng): Remove this separately.
 #include "base/gtest_prod_util.h"
 #include "base/memory/scoped_refptr.h"
 #include "base/sequence_checker.h"
@@ -113,7 +114,7 @@ class BASE_EXPORT RefCountedBase {
   template <typename U>
   friend scoped_refptr<U> base::AdoptRef(U*);
 
-  FRIEND_TEST_ALL_PREFIXES(RefCountedDeathTest, TestOverflowCheck);
+  friend class RefCountedOverflowTest;
 
   void Adopted() const {
 #if DCHECK_IS_ON()
@@ -195,6 +196,8 @@ class BASE_EXPORT RefCountedThreadSafeBase {
   template <typename U>
   friend scoped_refptr<U> base::AdoptRef(U*);
 
+  friend class RefCountedOverflowTest;
+
   void Adopted() const {
 #if DCHECK_IS_ON()
     DCHECK(needs_adopt_ref_);
@@ -210,7 +213,7 @@ class BASE_EXPORT RefCountedThreadSafeBase {
     // MakeRefCounted.
     DCHECK(!needs_adopt_ref_);
 #endif
-    ref_count_.Increment();
+    CHECK_NE(ref_count_.Increment(), std::numeric_limits<int>::max());
   }
 
   ALWAYS_INLINE void AddRefWithCheckImpl() const {
@@ -221,7 +224,9 @@ class BASE_EXPORT RefCountedThreadSafeBase {
     // MakeRefCounted.
     DCHECK(!needs_adopt_ref_);
 #endif
-    CHECK_GT(ref_count_.Increment(), 0);
+    int pre_increment_count = ref_count_.Increment();
+    CHECK_GT(pre_increment_count, 0);
+    CHECK_NE(pre_increment_count, std::numeric_limits<int>::max());
   }
 
   ALWAYS_INLINE bool ReleaseImpl() const {
@@ -314,9 +319,8 @@ class BASE_EXPORT ScopedAllowCrossThreadRefCountAccess final {
 //    And start-from-one ref count is a step to merge WTF::RefCounted into
 //    base::RefCounted.
 //
-#define REQUIRE_ADOPTION_FOR_REFCOUNTED_TYPE()             \
-  static constexpr ::base::subtle::StartRefCountFromOneTag \
-      kRefCountPreference = ::base::subtle::kStartRefCountFromOneTag
+#define REQUIRE_ADOPTION_FOR_REFCOUNTED_TYPE() \
+  using RefCountPreferenceTag = ::base::subtle::StartRefCountFromOneTag
 
 template <class T, typename Traits>
 class RefCounted;
@@ -331,10 +335,9 @@ struct DefaultRefCountedTraits {
 template <class T, typename Traits = DefaultRefCountedTraits<T>>
 class RefCounted : public subtle::RefCountedBase {
  public:
-  static constexpr subtle::StartRefCountFromZeroTag kRefCountPreference =
-      subtle::kStartRefCountFromZeroTag;
+  using RefCountPreferenceTag = subtle::StartRefCountFromZeroTag;
 
-  RefCounted() : subtle::RefCountedBase(T::kRefCountPreference) {}
+  RefCounted() : subtle::RefCountedBase(subtle::GetRefCountPreference<T>()) {}
 
   RefCounted(const RefCounted&) = delete;
   RefCounted& operator=(const RefCounted&) = delete;
@@ -399,16 +402,15 @@ struct DefaultRefCountedThreadSafeTraits {
 template <class T, typename Traits = DefaultRefCountedThreadSafeTraits<T> >
 class RefCountedThreadSafe : public subtle::RefCountedThreadSafeBase {
  public:
-  static constexpr subtle::StartRefCountFromZeroTag kRefCountPreference =
-      subtle::kStartRefCountFromZeroTag;
+  using RefCountPreferenceTag = subtle::StartRefCountFromZeroTag;
 
   explicit RefCountedThreadSafe()
-      : subtle::RefCountedThreadSafeBase(T::kRefCountPreference) {}
+      : subtle::RefCountedThreadSafeBase(subtle::GetRefCountPreference<T>()) {}
 
   RefCountedThreadSafe(const RefCountedThreadSafe&) = delete;
   RefCountedThreadSafe& operator=(const RefCountedThreadSafe&) = delete;
 
-  void AddRef() const { AddRefImpl(T::kRefCountPreference); }
+  void AddRef() const { AddRefImpl(subtle::GetRefCountPreference<T>()); }
 
   void Release() const {
     if (subtle::RefCountedThreadSafeBase::Release()) {

base/memory/ref_counted_delete_on_sequence.h

@@ -32,14 +32,13 @@ namespace base {
 template <class T>
 class RefCountedDeleteOnSequence : public subtle::RefCountedThreadSafeBase {
  public:
-  static constexpr subtle::StartRefCountFromZeroTag kRefCountPreference =
-      subtle::kStartRefCountFromZeroTag;
+  using RefCountPreferenceTag = subtle::StartRefCountFromZeroTag;
 
   // A SequencedTaskRunner for the current sequence can be acquired by calling
   // SequencedTaskRunnerHandle::Get().
   explicit RefCountedDeleteOnSequence(
       scoped_refptr<SequencedTaskRunner> owning_task_runner)
-      : subtle::RefCountedThreadSafeBase(T::kRefCountPreference),
+      : subtle::RefCountedThreadSafeBase(subtle::GetRefCountPreference<T>()),
         owning_task_runner_(std::move(owning_task_runner)) {
     DCHECK(owning_task_runner_);
   }
@@ -48,7 +47,7 @@ class RefCountedDeleteOnSequence : public subtle::RefCountedThreadSafeBase {
   RefCountedDeleteOnSequence& operator=(const RefCountedDeleteOnSequence&) =
       delete;
 
-  void AddRef() const { AddRefImpl(T::kRefCountPreference); }
+  void AddRef() const { AddRefImpl(subtle::GetRefCountPreference<T>()); }
 
   void Release() const {
     if (subtle::RefCountedThreadSafeBase::Release())

base/memory/ref_counted_unittest.cc

@@ -177,15 +177,6 @@ class CheckRefptrNull : public base::RefCounted<CheckRefptrNull> {
   raw_ptr<scoped_refptr<CheckRefptrNull>> ptr_ = nullptr;
 };
 
-class Overflow : public base::RefCounted<Overflow> {
- public:
-  Overflow() = default;
-
- private:
-  friend class base::RefCounted<Overflow>;
-  ~Overflow() = default;
-};
-
 }  // namespace
 
 TEST(RefCountedUnitTest, TestSelfAssignment) {
@@ -693,12 +684,88 @@ TEST(RefCountedDeathTest, TestAdoptRef) {
 }
 
 #if defined(ARCH_CPU_64_BITS)
-TEST(RefCountedDeathTest, TestOverflowCheck) {
+class RefCountedOverflowTest : public ::testing::Test {
+ public:
+  static uint32_t& GetMutableRefCount(RefCountedBase* ref_counted) {
+    return ref_counted->ref_count_;
+  }
+
+  static std::atomic_int& GetMutableRefCount(
+      RefCountedThreadSafeBase* ref_counted) {
+    return ref_counted->ref_count_.ref_count_;
+  }
+};
+
+TEST_F(RefCountedOverflowTest, NonThreadSafeStartFromZero) {
+  class Overflow : public base::RefCounted<Overflow> {
+   public:
+    Overflow() { EXPECT_FALSE(HasOneRef()); }
+
+   private:
+    friend class base::RefCounted<Overflow>;
+    ~Overflow() = default;
+  };
+
+  auto p = base::MakeRefCounted<Overflow>();
+  GetMutableRefCount(p.get()) = std::numeric_limits<uint32_t>::max();
+  EXPECT_CHECK_DEATH(p->AddRef());
+  // Ensure `p` doesn't leak and fail lsan builds.
+  GetMutableRefCount(p.get()) = 1;
+}
+
+TEST_F(RefCountedOverflowTest, NonThreadSafeStartFromOne) {
+  class Overflow : public base::RefCounted<Overflow> {
+   public:
+    REQUIRE_ADOPTION_FOR_REFCOUNTED_TYPE();
+
+    Overflow() { EXPECT_TRUE(HasOneRef()); }
+
+   private:
+    friend class base::RefCounted<Overflow>;
+    ~Overflow() = default;
+  };
+
+  auto p = base::MakeRefCounted<Overflow>();
+  GetMutableRefCount(p.get()) = std::numeric_limits<uint32_t>::max();
+  EXPECT_CHECK_DEATH(p->AddRef());
+  // Ensure `p` doesn't leak and fail lsan builds.
+  GetMutableRefCount(p.get()) = 1;
+}
+
+TEST_F(RefCountedOverflowTest, ThreadSafeStartFromZero) {
+  class Overflow : public base::RefCountedThreadSafe<Overflow> {
+   public:
+    Overflow() { EXPECT_FALSE(HasOneRef()); }
+
+   private:
+    friend class base::RefCountedThreadSafe<Overflow>;
+    ~Overflow() = default;
+  };
+
+  auto p = base::MakeRefCounted<Overflow>();
+  GetMutableRefCount(p.get()) = std::numeric_limits<int>::max();
+  EXPECT_CHECK_DEATH(p->AddRef());
+  // Ensure `p` doesn't leak and fail lsan builds.
+  GetMutableRefCount(p.get()) = 1;
+}
+
+TEST_F(RefCountedOverflowTest, ThreadSafeStartFromOne) {
+  class Overflow : public base::RefCountedThreadSafe<Overflow> {
+   public:
+    REQUIRE_ADOPTION_FOR_REFCOUNTED_TYPE();
+
+    Overflow() { EXPECT_TRUE(HasOneRef()); }
+
+   private:
+    friend class base::RefCountedThreadSafe<Overflow>;
+    ~Overflow() = default;
+  };
+
   auto p = base::MakeRefCounted<Overflow>();
-  p->ref_count_ = std::numeric_limits<uint32_t>::max();
+  GetMutableRefCount(p.get()) = std::numeric_limits<int>::max();
   EXPECT_CHECK_DEATH(p->AddRef());
   // Ensure `p` doesn't leak and fail lsan builds.
-  p->ref_count_ = 1;
+  GetMutableRefCount(p.get()) = 1;
 }
 #endif
 

base/memory/scoped_refptr.h

@@ -43,31 +43,49 @@ enum AdoptRefTag { kAdoptRefTag };
 enum StartRefCountFromZeroTag { kStartRefCountFromZeroTag };
 enum StartRefCountFromOneTag { kStartRefCountFromOneTag };
 
+template <typename TagType>
+struct RefCountPreferenceTagTraits;
+
+template <>
+struct RefCountPreferenceTagTraits<StartRefCountFromZeroTag> {
+  static constexpr StartRefCountFromZeroTag kTag = kStartRefCountFromZeroTag;
+};
+
+template <>
+struct RefCountPreferenceTagTraits<StartRefCountFromOneTag> {
+  static constexpr StartRefCountFromOneTag kTag = kStartRefCountFromOneTag;
+};
+
+template <typename T, typename Tag = typename T::RefCountPreferenceTag>
+constexpr Tag GetRefCountPreference() {
+  return RefCountPreferenceTagTraits<Tag>::kTag;
+}
+
 // scoped_refptr<T> is typically used with one of several RefCounted<T> base
 // classes or with custom AddRef and Release methods. These overloads dispatch
 // on which was used.
 
 template <typename T, typename U, typename V>
 constexpr bool IsRefCountPreferenceOverridden(const T*,
                                               const RefCounted<U, V>*) {
-  return !std::is_same<std::decay_t<decltype(T::kRefCountPreference)>,
-                       std::decay_t<decltype(U::kRefCountPreference)>>::value;
+  return !std::is_same_v<std::decay_t<decltype(GetRefCountPreference<T>())>,
+                         std::decay_t<decltype(GetRefCountPreference<U>())>>;
 }
 
 template <typename T, typename U, typename V>
 constexpr bool IsRefCountPreferenceOverridden(
     const T*,
     const RefCountedThreadSafe<U, V>*) {
-  return !std::is_same<std::decay_t<decltype(T::kRefCountPreference)>,
-                       std::decay_t<decltype(U::kRefCountPreference)>>::value;
+  return !std::is_same_v<std::decay_t<decltype(GetRefCountPreference<T>())>,
+                         std::decay_t<decltype(GetRefCountPreference<U>())>>;
 }
 
 template <typename T, typename U>
 constexpr bool IsRefCountPreferenceOverridden(
     const T*,
     const RefCountedDeleteOnSequence<U>*) {
-  return !std::is_same<std::decay_t<decltype(T::kRefCountPreference)>,
-                       std::decay_t<decltype(U::kRefCountPreference)>>::value;
+  return !std::is_same_v<std::decay_t<decltype(GetRefCountPreference<T>())>,
+                         std::decay_t<decltype(GetRefCountPreference<U>())>>;
 }
 
 constexpr bool IsRefCountPreferenceOverridden(...) {
@@ -105,7 +123,7 @@ constexpr void AssertRefCountBaseMatches(...) {}
 // from 1 instead of 0.
 template <typename T>
 scoped_refptr<T> AdoptRef(T* obj) {
-  using Tag = std::decay_t<decltype(T::kRefCountPreference)>;
+  using Tag = std::decay_t<decltype(subtle::GetRefCountPreference<T>())>;
   static_assert(std::is_same<subtle::StartRefCountFromOneTag, Tag>::value,
                 "Use AdoptRef only if the reference count starts from one.");
 
@@ -134,7 +152,7 @@ scoped_refptr<T> AdoptRefIfNeeded(T* obj, StartRefCountFromOneTag) {
 template <typename T, typename... Args>
 scoped_refptr<T> MakeRefCounted(Args&&... args) {
   T* obj = new T(std::forward<Args>(args)...);
-  return subtle::AdoptRefIfNeeded(obj, T::kRefCountPreference);
+  return subtle::AdoptRefIfNeeded(obj, subtle::GetRefCountPreference<T>());
 }
 
 // Takes an instance of T, which is a ref counted type, and wraps the object

gpu/command_buffer/service/dxgi_shared_handle_manager.cc

@@ -49,7 +49,8 @@ DXGISharedHandleState::DXGISharedHandleState(
     gfx::DXGIHandleToken token,
     base::win::ScopedHandle shared_handle,
     Microsoft::WRL::ComPtr<ID3D11Texture2D> d3d11_texture)
-    : base::subtle::RefCountedThreadSafeBase(kRefCountPreference),
+    : base::subtle::RefCountedThreadSafeBase(
+          base::subtle::GetRefCountPreference<DXGISharedHandleState>()),
       manager_(std::move(manager)),
       token_(std::move(token)),
       shared_handle_(std::move(shared_handle)),