Store a raw_ptr in SafeRef

This will allow it to be visible to DanglingPointerDetector, instead
of hiding the dangling pointer when moving from raw_ref to SafeRef.

This moves the pointer back out of WeakPtrBase so that SafeRef can
hold a raw_ptr<T> and WeakPtr can hold a T*. The latter is done to
avoid holding memory in quarantine for a dangling WeakPtr - since
WeakPtr guarantees no access to that memory beyond its lifetime
already. At this point WeakPtrBase is not providing anything of use,
so we remove it entirely and put the WeakReference in WeakPtr and
SafeRef. To avoid binary size growth, we out-of-line more stuff in
WeakReference. With that, binary size decreases another 4k with the
change.

The SafeRef and WeakPtr tests are improved to actually fail in the
case of virtual inheritance if the pointer held in WeakPtr is casted
incorrectly to the wrong base class. This won't be an issue with the
current implementation, but they were not testing the old uintptr_t
implementation correctly.

Bug: 1414420
Change-Id: I869cb04cf3711fab892c2223ab4e4cfc5f5f4fbe
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4245786
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Commit-Queue: danakj <danakj@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1108684}

base/memory/safe_ref.h

@@ -50,88 +50,113 @@ class SafeRef {
   // SafeRef if the pointer may not be present.
 
   // Copy construction and assignment.
-  SafeRef(const SafeRef& p) : w_(p.w_) {
+  SafeRef(const SafeRef& other) : ref_(other.ref_), ptr_(other.ptr_) {
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
   }
-  SafeRef& operator=(const SafeRef& p) {
-    w_ = p.w_;
+  SafeRef& operator=(const SafeRef& other) {
+    ref_ = other.ref_;
+    ptr_ = other.ptr_;
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
     return *this;
   }
 
   // Move construction and assignment.
-  SafeRef(SafeRef&& p) : w_(std::move(p.w_)) { CHECK(w_); }
-  SafeRef& operator=(SafeRef&& p) {
-    w_ = std::move(p.w_);
+  SafeRef(SafeRef&& other)
+      : ref_(std::move(other.ref_)), ptr_(std::move(other.ptr_)) {
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
+  }
+  SafeRef& operator=(SafeRef&& other) {
+    ref_ = std::move(other.ref_);
+    ptr_ = std::move(other.ptr_);
+    // Avoid use-after-move.
+    CHECK(ref_.IsValid());
     return *this;
   }
 
   // Copy conversion from SafeRef<U>.
-  template <typename U>
+  template <typename U,
+            typename = std::enable_if_t<std::is_convertible_v<U*, T*>>>
   // NOLINTNEXTLINE(google-explicit-constructor)
-  SafeRef(const SafeRef<U>& p) : w_(p.w_) {
+  SafeRef(const SafeRef<U>& other)
+      : ref_(other.ref_),
+        ptr_(other.ptr_)  // raw_ptr<U> converts to raw_ptr<T>.
+  {
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
   }
   template <typename U>
-  SafeRef& operator=(const SafeRef<U>& p) {
-    w_ = p.w_;
+  SafeRef& operator=(const SafeRef<U>& other) {
+    ref_ = other.ref_;
+    ptr_ = other.ptr_;  // raw_ptr<U> converts to raw_ptr<T>.
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
     return *this;
   }
 
   // Move conversion from SafeRef<U>.
   template <typename U>
   // NOLINTNEXTLINE(google-explicit-constructor)
-  SafeRef(SafeRef<U>&& p) : w_(std::move(p.w_)) {
+  SafeRef(SafeRef<U>&& other)
+      : ref_(std::move(other.ref_)),
+        ptr_(std::move(other.ptr_))  // raw_ptr<U> converts to raw_ptr<T>.
+  {
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
   }
   template <typename U>
-  SafeRef& operator=(SafeRef<U>&& p) {
-    w_ = std::move(p.w_);
+  SafeRef& operator=(SafeRef<U>&& other) {
+    ref_ = std::move(other.ref_);
+    ptr_ = std::move(other.ptr_);  // raw_ptr<U> converts to raw_ptr<T>.
     // Avoid use-after-move.
-    CHECK(w_);
+    CHECK(ref_.IsValid());
     return *this;
   }
 
-  // Call methods on the underlying T. Will CHECK() if the T pointee is no
-  // longer alive.
-  T* operator->() const {
-    // We rely on WeakPtr<T> to CHECK() on a bad deref; tests verify this.
-    return w_.operator->();
-  }
-
   // Provide access to the underlying T as a reference. Will CHECK() if the T
   // pointee is no longer alive.
-  T& operator*() const { return *operator->(); }
+  T& operator*() const {
+    CHECK(ref_.IsValid());
+    return *ptr_;
+  }
+
+  // Used to call methods on the underlying T. Will CHECK() if the T pointee is
+  // no longer alive.
+  T* operator->() const {
+    CHECK(ref_.IsValid());
+    return &*ptr_;
+  }
 
  private:
   template <typename U>
   friend class SafeRef;
   template <typename U>
   friend SafeRef<U> internal::MakeSafeRefFromWeakPtrInternals(
-      const internal::WeakReference& ref,
+      internal::WeakReference&& ref,
       U* ptr);
 
-  // Construction from a from WeakPtr. Will CHECK() if the WeakPtr is already
-  // invalid.
-  explicit SafeRef(WeakPtr<T> w) : w_(std::move(w)) { CHECK(w_); }
+  // Construction from a from a WeakPtr's internals. Will CHECK() if the WeakPtr
+  // is already invalid.
+  explicit SafeRef(internal::WeakReference&& ref, T* ptr)
+      : ref_(std::move(ref)), ptr_(ptr) {
+    CHECK(ref_.IsValid());
+  }
+
+  internal::WeakReference ref_;
 
-  WeakPtr<T> w_;
+  // This pointer is only valid when ref_.is_valid() is true.  Otherwise, its
+  // value is undefined (as opposed to nullptr). Unlike WeakPtr, this raw_ptr is
+  // not allowed to dangle.
+  raw_ptr<T> ptr_;
 };
 
 namespace internal {
 template <typename T>
-SafeRef<T> MakeSafeRefFromWeakPtrInternals(const internal::WeakReference& ref,
+SafeRef<T> MakeSafeRefFromWeakPtrInternals(internal::WeakReference&& ref,
                                            T* ptr) {
-  CHECK(ptr);
-  return SafeRef<T>(WeakPtr<T>(ref, ptr));
+  return SafeRef<T>(std::move(ref), ptr);
 }
 }  // namespace internal
 

base/memory/safe_ref_unittest.cc

@@ -11,24 +11,30 @@
 #include "base/memory/raw_ptr_exclusion.h"
 #include "base/memory/weak_ptr.h"
 #include "base/test/gtest_util.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/abseil-cpp/absl/types/optional.h"
 
 namespace base {
 namespace {
 
 struct ReallyBaseClass {};
-struct BaseClass : ReallyBaseClass {};
+struct BaseClass : ReallyBaseClass {
+  virtual ~BaseClass() = default;
+  void VirtualMethod() {}
+};
+struct OtherBaseClass {
+  virtual ~OtherBaseClass() = default;
+  virtual void VirtualMethod() {}
+};
 
-struct WithWeak : BaseClass {
-  ~WithWeak() { self = nullptr; }
+struct WithWeak final : BaseClass, OtherBaseClass {
+  ~WithWeak() final { self = nullptr; }
 
   void Method() {}
 
   int i = 1;
-  // This field is not a raw_ptr<> because it was filtered by the rewriter for:
-  // #union
-  RAW_PTR_EXCLUSION WithWeak* self{this};
+  raw_ptr<WithWeak> self{this};
   base::WeakPtrFactory<WithWeak> factory{this};
 };
 
@@ -118,18 +124,34 @@ TEST(SafeRefDeathTest, StarOperatorCrashIfBadPointer) {
   EXPECT_CHECK_DEATH(safe.operator*());  // Will crash since not live.
 }
 
+TEST(SafeRefTest, ConversionToBaseClassFromCopyConstruct) {
+  WithWeak with;
+  SafeRef<WithWeak> safe(with.factory.GetSafeRef());
+  SafeRef<OtherBaseClass> base_safe = safe;
+  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe), &with);
+}
+
+TEST(SafeRefTest, ConversionToBaseClassFromMoveConstruct) {
+  WithWeak with;
+  SafeRef<WithWeak> safe(with.factory.GetSafeRef());
+  SafeRef<OtherBaseClass> base_safe = std::move(safe);
+  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe), &with);
+}
+
 TEST(SafeRefTest, ConversionToBaseClassFromCopyAssign) {
   WithWeak with;
   SafeRef<WithWeak> safe(with.factory.GetSafeRef());
-  SafeRef<BaseClass> base_safe = safe;
-  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe)->self, &with);
+  SafeRef<OtherBaseClass> base_safe(with.factory.GetSafeRef());
+  base_safe = safe;
+  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe), &with);
 }
 
 TEST(SafeRefTest, ConversionToBaseClassFromMoveAssign) {
   WithWeak with;
   SafeRef<WithWeak> safe(with.factory.GetSafeRef());
-  SafeRef<BaseClass> base_safe = std::move(safe);
-  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe)->self, &with);
+  SafeRef<OtherBaseClass> base_safe(with.factory.GetSafeRef());
+  base_safe = std::move(safe);
+  EXPECT_EQ(static_cast<WithWeak*>(&*base_safe), &with);
 }
 
 TEST(SafeRefTest, CanDerefConst) {
@@ -266,5 +288,15 @@ TEST(SafeRefTest, Bind) {
   BindOnce(&WithWeak::Method, with.factory.GetSafeRef()).Run();
 }
 
+#if BUILDFLAG(ENABLE_DANGLING_RAW_PTR_CHECKS)
+// TODO(crbug.com/1416264): Test this when we are able to.
+TEST(SafeRefDeathTest, DISABLED_DanglingPointerDetector) {
+  auto with = std::make_unique<WithWeak>();
+  SafeRef<WithWeak> safe(with->factory.GetSafeRef());
+  BASE_EXPECT_DEATH({ with.reset(); },
+                    testing::HasSubstr("Detected dangling raw_ptr"));
+}
+#endif
+
 }  // namespace
 }  // namespace base

base/memory/weak_ptr.cc

@@ -10,8 +10,7 @@
 #include "base/debug/stack_trace.h"
 #endif
 
-namespace base {
-namespace internal {
+namespace base::internal {
 
 WeakReference::Flag::Flag() {
   // Flags only become bound when checked for validity, or invalidated,
@@ -53,14 +52,19 @@ void WeakReference::Flag::DetachFromSequence() {
 WeakReference::Flag::~Flag() = default;
 
 WeakReference::WeakReference() = default;
-
 WeakReference::WeakReference(const scoped_refptr<Flag>& flag) : flag_(flag) {}
-
 WeakReference::~WeakReference() = default;
 
+WeakReference::WeakReference(const WeakReference& other) = default;
+WeakReference& WeakReference::operator=(const WeakReference& other) = default;
+
 WeakReference::WeakReference(WeakReference&& other) noexcept = default;
+WeakReference& WeakReference::operator=(WeakReference&& other) noexcept =
+    default;
 
-WeakReference::WeakReference(const WeakReference& other) = default;
+void WeakReference::Reset() {
+  flag_ = nullptr;
+}
 
 bool WeakReference::IsValid() const {
   return flag_ && flag_->IsValid();
@@ -92,15 +96,6 @@ void WeakReferenceOwner::Invalidate() {
   flag_ = MakeRefCounted<WeakReference::Flag>();
 }
 
-WeakPtrBase::WeakPtrBase() : ptr_(0) {}
-
-WeakPtrBase::~WeakPtrBase() = default;
-
-WeakPtrBase::WeakPtrBase(const WeakReference& ref, uintptr_t ptr)
-    : ref_(ref), ptr_(ptr) {
-  DCHECK(ptr_);
-}
-
 WeakPtrFactoryBase::WeakPtrFactoryBase(uintptr_t ptr) : ptr_(ptr) {
   DCHECK(ptr_);
 }
@@ -109,5 +104,4 @@ WeakPtrFactoryBase::~WeakPtrFactoryBase() {
   ptr_ = 0;
 }
 
-}  // namespace internal
-}  // namespace base
+}  // namespace base::internal