From e1ff0309c7889857b94f78181c667dd4ab210141 Mon Sep 17 00:00:00 2001
From: Victoria Brekenfeld <victoria@system76.com>
Date: Thu, 19 Mar 2026 13:41:52 +0100
Subject: [PATCH] wayland/image-copy-capture: Break toplevel reference cycle

---
 src/wayland/handlers/image_capture_source.rs  |  2 +-
 .../handlers/image_copy_capture/mod.rs        | 35 +++++++++++++++----
 src/wayland/protocols/image_capture_source.rs |  8 ++---
 3 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/wayland/handlers/image_capture_source.rs b/src/wayland/handlers/image_capture_source.rs
index 6c705c34..32858c2d 100644
--- a/src/wayland/handlers/image_capture_source.rs
+++ b/src/wayland/handlers/image_capture_source.rs
@@ -45,7 +45,7 @@ impl ToplevelCaptureSourceHandler for State {
         toplevel: ForeignToplevelHandle,
     ) {
         let data = match window_from_ext(self, toplevel) {
-            Some(toplevel) => ImageCaptureSourceKind::Toplevel(toplevel.clone()),
+            Some(toplevel) => ImageCaptureSourceKind::Toplevel(toplevel.downgrade()),
             None => ImageCaptureSourceKind::Destroyed,
         };
         source.user_data().insert_if_missing(|| data);
diff --git a/src/wayland/handlers/image_copy_capture/mod.rs b/src/wayland/handlers/image_copy_capture/mod.rs
index b9d7626b..c6b6c9af 100644
--- a/src/wayland/handlers/image_copy_capture/mod.rs
+++ b/src/wayland/handlers/image_copy_capture/mod.rs
@@ -61,7 +61,11 @@ impl ImageCopyCaptureHandler for State {
                 constraints_for_output(output, &mut self.backend)
             }
             ImageCaptureSourceKind::Toplevel(window) => {
-                constraints_for_toplevel(window, &mut self.backend)
+                if let Some(window) = window.upgrade() {
+                    constraints_for_toplevel(&window, &mut self.backend)
+                } else {
+                    None
+                }
             }
             _ => None,
         }
@@ -128,7 +132,12 @@ impl ImageCopyCaptureHandler for State {
                 });
                 workspace.add_session(session);
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
+            ImageCaptureSourceKind::Toplevel(toplevel) => {
+                let Some(mut toplevel) = toplevel.upgrade() else {
+                    session.stop();
+                    return;
+                };
+
                 let size = toplevel.geometry().size.to_physical(1);
                 session.user_data().insert_if_missing_threadsafe(|| {
                     Mutex::new(SessionUserData::new(OutputDamageTracker::new(
@@ -238,7 +247,11 @@ impl ImageCopyCaptureHandler for State {
 
                 workspace.add_cursor_session(session);
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
+            ImageCaptureSourceKind::Toplevel(toplevel) => {
+                let Some(mut toplevel) = toplevel.upgrade() else {
+                    return;
+                };
+
                 let shell = self.common.shell.read();
                 if let Some(element) = shell.element_for_surface(&toplevel)
                     && element.has_active_window(&toplevel)
@@ -287,6 +300,10 @@ impl ImageCopyCaptureHandler for State {
                 render_workspace_to_buffer(self, session, frame, handle)
             }
             ImageCaptureSourceKind::Toplevel(toplevel) => {
+                let Some(toplevel) = toplevel.upgrade() else {
+                    return;
+                };
+
                 render_window_to_buffer(self, session, frame, &toplevel)
             }
             ImageCaptureSourceKind::Destroyed => {
@@ -336,7 +353,11 @@ impl ImageCopyCaptureHandler for State {
                     workspace.remove_session(&session)
                 }
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => toplevel.remove_session(&session),
+            ImageCaptureSourceKind::Toplevel(toplevel) => {
+                if let Some(mut toplevel) = toplevel.upgrade() {
+                    toplevel.remove_session(&session);
+                }
+            }
             ImageCaptureSourceKind::Destroyed => {}
         }
     }
@@ -365,8 +386,10 @@ impl ImageCopyCaptureHandler for State {
                     workspace.remove_cursor_session(&session)
                 }
             }
-            ImageCaptureSourceKind::Toplevel(mut toplevel) => {
-                toplevel.remove_cursor_session(&session)
+            ImageCaptureSourceKind::Toplevel(toplevel) => {
+                if let Some(mut toplevel) = toplevel.upgrade() {
+                    toplevel.remove_cursor_session(&session)
+                }
             }
             ImageCaptureSourceKind::Destroyed => {}
         }
diff --git a/src/wayland/protocols/image_capture_source.rs b/src/wayland/protocols/image_capture_source.rs
index 1214cbe0..bf130da6 100644
--- a/src/wayland/protocols/image_capture_source.rs
+++ b/src/wayland/protocols/image_capture_source.rs
@@ -1,9 +1,7 @@
 use super::{
     workspace::{WorkspaceHandle, WorkspaceHandler},
 };
-use crate::{
-    shell::CosmicSurface,
-};
+use crate::shell::element::surface::WeakCosmicSurface;
 use cosmic_protocols::image_capture_source::v1::server::{
     zcosmic_workspace_image_capture_source_manager_v1::{
         Request as CosmicWorkspaceSourceRequest, ZcosmicWorkspaceImageCaptureSourceManagerV1,
@@ -30,11 +28,11 @@ pub struct WorkspaceImageCaptureSourceManagerGlobalData {
     filter: Box<dyn for<'a> Fn(&'a Client) -> bool + Send + Sync>,
 }
 
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone)]
 pub enum ImageCaptureSourceKind {
     Output(WeakOutput),
     Workspace(WorkspaceHandle),
-    Toplevel(CosmicSurface),
+    Toplevel(WeakCosmicSurface),
     Destroyed,
 }