Fix crash when closing windows and popups

src/client/server_handlers.rs

@@ -253,7 +253,13 @@ impl WprsClientState {
         if let Some(Role::XdgToplevel(toplevel)) = &surface.role {
             match request.payload {
                 ToplevelRequestPayload::Destroyed => {
-                    surface.role = None;
+                    let toplevel = surface
+                        .role
+                        .take()
+                        .location(loc!())?
+                        .into_xdg_toplevel()
+                        .unwrap();
+                    surface.local_surface = Some(toplevel.destroy())
                 },
                 ToplevelRequestPayload::SetMaximized => {
                     toplevel.local_window.set_maximized();
@@ -304,7 +310,13 @@ impl WprsClientState {
         let surface = client.surface(&request.surface).location(loc!())?;
         match request.payload {
             PopupRequestPayload::Destroyed => {
-                surface.role = None;
+                let popup = surface
+                    .role
+                    .take()
+                    .location(loc!())?
+                    .into_xdg_popup()
+                    .unwrap();
+                surface.local_surface = Some(popup.destroy())
             },
         }
         Ok(())

src/client/xdg_shell.rs

@@ -14,6 +14,7 @@
 
 use std::collections::HashMap;
 
+use smithay_client_toolkit::compositor::Surface;
 use smithay_client_toolkit::reexports::client::Proxy;
 use smithay_client_toolkit::reexports::client::QueueHandle;
 use smithay_client_toolkit::reexports::protocols::xdg::shell::client::xdg_positioner;
@@ -243,6 +244,10 @@ impl RemoteXdgToplevel {
         let surface = surfaces.get_mut(&surface_id).location(loc!())?;
         Self::update(surface_state, surface)
     }
+
+    pub fn destroy(self) -> Surface {
+        self.local_window.destroy()
+    }
 }
 
 #[derive(Debug)]
@@ -428,4 +433,8 @@ impl RemoteXdgPopup {
         let surface = surfaces.get_mut(&surface_id).location(loc!())?;
         Self::update(surface_state, surface, xdg_shell_state)
     }
+
+    pub fn destroy(self) -> Surface {
+        self.local_popup.destroy()
+    }
 }

src/server/smithay_handlers.rs

@@ -151,7 +151,7 @@ impl WprsServerState {
                 client: serialization::ClientId::new(&surface.client().unwrap()),
                 surface: (&surface.id()).into(),
                 payload,
-            })))
+            })));
     }
 
     #[allow(clippy::missing_panics_doc)]
@@ -212,12 +212,23 @@ impl XdgShellHandler for WprsServerState {
     }
 
     #[instrument(skip(self), level = "debug")]
-    fn toplevel_destroyed(&mut self, surface: ToplevelSurface) {
+    fn toplevel_destroyed(&mut self, toplevel: ToplevelSurface) {
         // If client() returns None, the surface was already destroyed and an
         // appropriate message would have been sent to the client, so we don't
         // need to worry about destroying the toplevel,
-        if surface.wl_surface().client().is_some() {
-            self.send_toplevel_request(&surface, ToplevelRequestPayload::Destroyed);
+        if toplevel.wl_surface().client().is_some() {
+            compositor::with_states(toplevel.wl_surface(), |surface_data| {
+                let surface_state = &mut surface_data
+                    .data_map
+                    .get::<LockedSurfaceState>()
+                    .unwrap()
+                    .0
+                    .lock()
+                    .unwrap();
+                surface_state.role = None;
+            });
+
+            self.send_toplevel_request(&toplevel, ToplevelRequestPayload::Destroyed);
         }
     }
 
@@ -246,16 +257,28 @@ impl XdgShellHandler for WprsServerState {
     }
 
     #[instrument(skip(self), level = "debug")]
-    fn popup_destroyed(&mut self, surface: PopupSurface) {
+    fn popup_destroyed(&mut self, popup: PopupSurface) {
         // If client() returns None, the surface was already destroyed and an
         // appropriate message would have been sent to the client, so we don't
         // need to worry about destroying the popup,
-        if let Some(client) = surface.wl_surface().client() {
+        if let Some(client) = popup.wl_surface().client() {
+            // Uses with_states internally and with_states is not reentrant.
+            compositor::with_states(popup.wl_surface(), |surface_data| {
+                let surface_state = &mut surface_data
+                    .data_map
+                    .get::<LockedSurfaceState>()
+                    .unwrap()
+                    .0
+                    .lock()
+                    .unwrap();
+                surface_state.role = None;
+            });
+
             self.serializer
                 .writer()
                 .send(SendType::Object(Request::Popup(PopupRequest {
                     client: serialization::ClientId::new(&client),
-                    surface: (&surface.wl_surface().id()).into(),
+                    surface: (&popup.wl_surface().id()).into(),
                     payload: PopupRequestPayload::Destroyed,
                 })));
         };