LibWeb: Protect ourselves during ResourceClient iteration

Notifying a Resource's clients may lead to arbitrary JS execution,
so we can't rely on the ResourceClient pointers remaining valid.
Use WeakPtr to avoid this problem.

2 files changed, 17 insertions, 12 deletions

diff --git a/Libraries/LibWeb/Loader/Resource.cpp b/Libraries/LibWeb/Loader/Resource.cpp
index e261220b82..2b4e0d61a7 100644
--- a/Libraries/LibWeb/Loader/Resource.cpp
+++ b/Libraries/LibWeb/Loader/Resource.cpp
@@ -24,6 +24,7 @@
  * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
 
+#include <AK/Function.h>
 #include <LibWeb/DOM/HTMLImageElement.h>
 #include <LibWeb/Loader/Resource.h>
 
@@ -43,6 +44,18 @@ Resource::~Resource()
 {
 }
 
+void Resource::for_each_client(Function<void(ResourceClient&)> callback)
+{
+    Vector<WeakPtr<ResourceClient>, 16> clients_copy;
+    clients_copy.ensure_capacity(m_clients.size());
+    for (auto* client : m_clients)
+        clients_copy.append(client->make_weak_ptr());
+    for (auto client : clients_copy) {
+        if (client)
+            callback(*client);
+    }
+}
+
 void Resource::did_load(Badge<ResourceLoader>, const ByteBuffer& data, const HashMap<String, String, CaseInsensitiveStringTraits>& headers)
 {
     ASSERT(!m_loaded);
diff --git a/Libraries/LibWeb/Loader/Resource.h b/Libraries/LibWeb/Loader/Resource.h
index 9af4ecb406..932e58c124 100644
--- a/Libraries/LibWeb/Loader/Resource.h
+++ b/Libraries/LibWeb/Loader/Resource.h
@@ -32,6 +32,8 @@
 #include <AK/Noncopyable.h>
 #include <AK/RefCounted.h>
 #include <AK/URL.h>
+#include <AK/WeakPtr.h>
+#include <AK/Weakable.h>
 #include <LibWeb/Forward.h>
 #include <LibWeb/Loader/LoadRequest.h>
 
@@ -60,17 +62,7 @@ public:
     void register_client(Badge<ResourceClient>, ResourceClient&);
     void unregister_client(Badge<ResourceClient>, ResourceClient&);
 
-    template<typename Callback>
-    void for_each_client(Callback callback)
-    {
-        // FIXME: This should use some kind of smart pointer to ResourceClient!
-        Vector<ResourceClient*, 16> clients_copy;
-        clients_copy.ensure_capacity(m_clients.size());
-        for (auto* client : m_clients)
-            clients_copy.append(client);
-        for (auto* client : clients_copy)
-            callback(*client);
-    }
+    void for_each_client(Function<void(ResourceClient&)>);
 
     void did_load(Badge<ResourceLoader>, const ByteBuffer& data, const HashMap<String, String, CaseInsensitiveStringTraits>& headers);
     void did_fail(Badge<ResourceLoader>, const String& error);
@@ -87,7 +79,7 @@ private:
     HashTable<ResourceClient*> m_clients;
 };
 
-class ResourceClient {
+class ResourceClient : public Weakable<ResourceClient> {
 public:
     virtual ~ResourceClient();