Guard against removed outputs in several places

Output::from_resource() succeeds even after the global has been disabled
and removed from niri. Clients operating on these disabled outputs could
cause panics in several places because niri assumed the output existed.

src/handlers/layer_shell.rs

@@ -1,6 +1,5 @@
 use smithay::delegate_layer_shell;
 use smithay::desktop::{layer_map_for_output, LayerSurface, PopupKind, WindowSurfaceType};
-use smithay::output::Output;
 use smithay::reexports::wayland_server::protocol::wl_output::WlOutput;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::wayland::compositor::{add_pre_commit_hook, get_parent, with_states, HookId};
@@ -27,7 +26,7 @@ impl WlrLayerShellHandler for State {
         namespace: String,
     ) {
         let output = if let Some(wl_output) = &wl_output {
-            Output::from_resource(wl_output)
+            self.niri.output_from_resource(wl_output)
         } else {
             self.niri.layout.active_output().cloned()
         };

src/handlers/mod.rs

@@ -466,7 +466,7 @@ impl SessionLockHandler for State {
     }
 
     fn new_surface(&mut self, surface: LockSurface, output: WlOutput) {
-        let Some(output) = Output::from_resource(&output) else {
+        let Some(output) = self.niri.output_from_resource(&output) else {
             warn!("no Output matching WlOutput");
             return;
         };
@@ -551,7 +551,9 @@ impl ForeignToplevelHandler for State {
         {
             let window = mapped.window.clone();
 
-            if let Some(requested_output) = wl_output.as_ref().and_then(Output::from_resource) {
+            if let Some(requested_output) =
+                wl_output.and_then(|o| self.niri.output_from_resource(&o))
+            {
                 if Some(&requested_output) != current_output {
                     self.niri.layout.move_to_output(
                         Some(&window),
@@ -627,6 +629,12 @@ delegate_ext_workspace!(State);
 
 impl ScreencopyHandler for State {
     fn frame(&mut self, manager: &ZwlrScreencopyManagerV1, screencopy: Screencopy) {
+        // This can happen if the output was removed before this was called.
+        if !self.niri.output_exists(screencopy.output()) {
+            trace!("screencopy output no longer exists");
+            return;
+        }
+
         // If with_damage then push it onto the queue for redraw of the output,
         // otherwise render it immediately.
         if screencopy.with_damage() {

src/handlers/xdg_shell.rs

@@ -612,7 +612,7 @@ impl XdgShellHandler for State {
         toplevel: ToplevelSurface,
         wl_output: Option<wl_output::WlOutput>,
     ) {
-        let requested_output = wl_output.as_ref().and_then(Output::from_resource);
+        let requested_output = wl_output.and_then(|o| self.niri.output_from_resource(&o));
 
         if let Some((mapped, current_output)) = self
             .niri

src/input/mod.rs

@@ -295,6 +295,7 @@ impl State {
         I::Device: 'static,
     {
         let device_output = event.device().output(self);
+        let device_output = device_output.filter(|output| self.niri.output_exists(output));
         let device_output = device_output.as_ref();
         let (target_geo, keep_ratio, px, transform) =
             if let Some(output) = device_output.or_else(|| self.niri.output_for_tablet()) {
@@ -4039,6 +4040,7 @@ impl State {
         fallback_output: Option<&Output>,
     ) -> Option<Point<f64, Logical>> {
         let output = evt.device().output(self);
+        let output = output.filter(|output| self.niri.output_exists(output));
         let output = output.as_ref().or(fallback_output)?;
         let output_geo = self.niri.global_space.output_geometry(output).unwrap();
         let transform = output.current_transform();

src/niri.rs

@@ -110,6 +110,7 @@ use smithay::wayland::viewporter::ViewporterState;
 use smithay::wayland::virtual_keyboard::VirtualKeyboardManagerState;
 use smithay::wayland::xdg_activation::XdgActivationState;
 use smithay::wayland::xdg_foreign::XdgForeignState;
+use wayland_server::protocol::wl_output::WlOutput;
 
 #[cfg(feature = "dbus")]
 use crate::a11y::A11y;
@@ -2874,6 +2875,20 @@ impl Niri {
         self.reposition_outputs(Some(&output));
     }
 
+    pub fn output_exists(&self, output: &Output) -> bool {
+        self.output_state.contains_key(output)
+    }
+
+    /// Converts a `WlOutput` to a corresponding `Output` if it exists.
+    ///
+    /// Compared to raw `Output::from_resource`, this method also verifies that the output still
+    /// exists in niri. Right after the output global is disabled, but before it is removed for
+    /// good, `Output::from_resource` will succeed, but since niri already forgot the output,
+    /// accessing it can cause logic bugs.
+    pub fn output_from_resource(&self, wl_output: &WlOutput) -> Option<Output> {
+        Output::from_resource(wl_output).filter(|output| self.output_exists(output))
+    }
+
     pub fn remove_output(&mut self, output: &Output) {
         for layer in layer_map_for_output(output).layers() {
             layer.layer_surface().send_close();