WindowServer: Fix animation crash

If an Animation's on_stop handler cleared itself inside the on_stop
event handler, it would remove itself from the animation map that was
still being iterated, leading to a crash.

To solve this, we'll iterate over the animations using the
remove_all_matching function, which enables us to delete it by simply
returning true when the animation finished. In the event that the
Animation is kept alive elsewhere and the on_stop event clears its own
reference, we need to temporarily bump the reference count. Another
advantage is that we only need to bump the reference count when an
animation is finished, whereas before this we bumped it
unconditionally.

3 files changed, 31 insertions, 8 deletions

diff --git a/Userland/Services/WindowServer/Animation.cpp b/Userland/Services/WindowServer/Animation.cpp
index 633331c156..37067a9c43 100644
--- a/Userland/Services/WindowServer/Animation.cpp
+++ b/Userland/Services/WindowServer/Animation.cpp
@@ -17,7 +17,8 @@ Animation::Animation()
 
 Animation::~Animation()
 {
-    Compositor::the().unregister_animation({}, *this);
+    if (!m_was_removed)
+        Compositor::the().unregister_animation({}, *this);
 }
 
 void Animation::set_duration(int duration_in_ms)
@@ -39,7 +40,12 @@ void Animation::stop()
         on_stop();
 }
 
-void Animation::update(Badge<Compositor>, Gfx::Painter& painter, Screen& screen, Gfx::DisjointRectSet& flush_rects)
+void Animation::was_removed(Badge<Compositor>)
+{
+    m_was_removed = true;
+}
+
+bool Animation::update(Badge<Compositor>, Gfx::Painter& painter, Screen& screen, Gfx::DisjointRectSet& flush_rects)
 {
     int elapsed_ms = m_timer.elapsed();
     float progress = min((float)elapsed_ms / (float)m_duration, 1.0f);
@@ -47,8 +53,7 @@ void Animation::update(Badge<Compositor>, Gfx::Painter& painter, Screen& screen,
     if (on_update)
         on_update(progress, painter, screen, flush_rects);
 
-    if (progress >= 1.0f)
-        stop();
+    return progress < 1.0f;
 }
 
 }
diff --git a/Userland/Services/WindowServer/Animation.h b/Userland/Services/WindowServer/Animation.h
index b8786d8734..89f5e7787d 100644
--- a/Userland/Services/WindowServer/Animation.h
+++ b/Userland/Services/WindowServer/Animation.h
@@ -28,11 +28,12 @@ public:
 
     void start();
     void stop();
+    void was_removed(Badge<Compositor>);
 
     void set_duration(int duration_in_ms);
     int duration() const { return m_duration; }
 
-    void update(Badge<Compositor>, Gfx::Painter&, Screen&, Gfx::DisjointRectSet& flush_rects);
+    bool update(Badge<Compositor>, Gfx::Painter&, Screen&, Gfx::DisjointRectSet& flush_rects);
 
     Function<void(float progress, Gfx::Painter&, Screen&, Gfx::DisjointRectSet& flush_rects)> on_update;
     Function<void()> on_stop;
@@ -43,6 +44,7 @@ private:
     Core::ElapsedTimer m_timer;
     int m_duration { 0 };
     bool m_running { false };
+    bool m_was_removed { false };
 };
 
 }
diff --git a/Userland/Services/WindowServer/Compositor.cpp b/Userland/Services/WindowServer/Compositor.cpp
index a173a888ae..2ca7f3ef52 100644
--- a/Userland/Services/WindowServer/Compositor.cpp
+++ b/Userland/Services/WindowServer/Compositor.cpp
@@ -1460,9 +1460,25 @@ void Compositor::unregister_animation(Badge<Animation>, Animation& animation)
 void Compositor::update_animations(Screen& screen, Gfx::DisjointRectSet& flush_rects)
 {
     auto& painter = *screen.compositor_screen_data().m_back_painter;
-    for (RefPtr<Animation> animation : m_animations) {
-        animation->update({}, painter, screen, flush_rects);
-    }
+    // Iterating over the animations using remove_all_matching we can iterate
+    // and immediately remove finished animations without having to keep track
+    // of them in a separate container.
+    m_animations.remove_all_matching([&](auto* animation) {
+        if (!animation->update({}, painter, screen, flush_rects)) {
+            // Mark it as removed so that the Animation::on_stop handler doesn't
+            // trigger the Animation object from being destroyed, causing it to
+            // unregister while we still loop over them.
+            animation->was_removed({});
+
+            // Temporarily bump the ref count so that if the Animation::on_stop
+            // handler clears its own reference, it doesn't immediately destroy
+            // itself while we're still in the Function<> call
+            NonnullRefPtr<Animation> protect_animation(*animation);
+            animation->stop();
+            return true;
+        }
+        return false;
+    });
 }
 
 void Compositor::create_window_stack_switch_overlay(WindowStack& target_stack)