From 6db879ee663249041a6a731a8a6f31bd40ca8089 Mon Sep 17 00:00:00 2001
From: Andreas Kling <awesomekling@gmail.com>
Date: Fri, 2 Aug 2019 11:35:05 +0200
Subject: [PATCH] AK: Fix ref leak in NonnullRefPtr::operator=(T&).

We would leak a ref when assigning a T& to a NonnullRefPtr that already
contains that same T.
---
 AK/NonnullRefPtr.h             |  7 ++++---
 AK/Tests/Makefile              |  5 ++++-
 AK/Tests/TestNonnullRefPtr.cpp | 36 ++++++++++++++++++++++++++++++++++
 3 files changed, 44 insertions(+), 4 deletions(-)
 create mode 100644 AK/Tests/TestNonnullRefPtr.cpp

diff --git a/AK/NonnullRefPtr.h b/AK/NonnullRefPtr.h
index 2a6d9b1d29..56772dea4c 100644
--- a/AK/NonnullRefPtr.h
+++ b/AK/NonnullRefPtr.h
@@ -131,10 +131,11 @@ public:
 
     NonnullRefPtr& operator=(T& object)
     {
-        if (m_ptr != &object)
+        if (m_ptr != &object) {
             deref_if_not_null(m_ptr);
-        m_ptr = &object;
-        m_ptr->ref();
+            m_ptr = &object;
+            m_ptr->ref();
+        }
         return *this;
     }
 
diff --git a/AK/Tests/Makefile b/AK/Tests/Makefile
index 3847c4834e..f1a28322a9 100644
--- a/AK/Tests/Makefile
+++ b/AK/Tests/Makefile
@@ -1,4 +1,4 @@
-PROGRAMS = TestString TestQueue TestVector TestHashMap TestJSON TestWeakPtr
+PROGRAMS = TestString TestQueue TestVector TestHashMap TestJSON TestWeakPtr TestNonnullRefPtr
 
 CXXFLAGS = -std=c++17 -Wall -Wextra -ggdb3 -O2 -I../ -I../../
 
@@ -42,6 +42,9 @@ TestJSON: TestJSON.o $(SHARED_TEST_OBJS)
 TestWeakPtr: TestWeakPtr.o $(SHARED_TEST_OBJS) 
 	$(PRE_CXX) $(CXX) $(CXXFLAGS) -o $@ TestWeakPtr.o $(SHARED_TEST_OBJS)
 
+TestNonnullRefPtr: TestNonnullRefPtr.o $(SHARED_TEST_OBJS)
+	$(PRE_CXX) $(CXX) $(CXXFLAGS) -o $@ TestNonnullRefPtr.o $(SHARED_TEST_OBJS)
+
 clean:
 	rm -f $(SHARED_TEST_OBJS)
 	rm -f $(PROGRAMS)
diff --git a/AK/Tests/TestNonnullRefPtr.cpp b/AK/Tests/TestNonnullRefPtr.cpp
new file mode 100644
index 0000000000..9012206586
--- /dev/null
+++ b/AK/Tests/TestNonnullRefPtr.cpp
@@ -0,0 +1,36 @@
+#include <AK/TestSuite.h>
+
+#include <AK/NonnullRefPtr.h>
+#include <AK/AKString.h>
+
+struct Object : public RefCounted<Object> {
+    int x;
+};
+
+TEST_CASE(basics)
+{
+    auto object = adopt(*new Object);
+    EXPECT(object.ptr() != nullptr);
+    EXPECT_EQ(object->ref_count(), 1);
+    object->ref();
+    EXPECT_EQ(object->ref_count(), 2);
+    object->deref();
+    EXPECT_EQ(object->ref_count(), 1);
+
+    {
+        NonnullRefPtr another = object;
+        EXPECT_EQ(object->ref_count(), 2);
+    }
+
+    EXPECT_EQ(object->ref_count(), 1);
+}
+
+TEST_CASE(assign_reference)
+{
+    auto object = adopt(*new Object);
+    EXPECT_EQ(object->ref_count(), 1);
+    object = *object;
+    EXPECT_EQ(object->ref_count(), 1);
+}
+
+TEST_MAIN(String)