Improved handling of XWayland grabs

A solution for https://github.com/Smithay/smithay/issues/1714.

With this, the lock screen is able to get keyboard focus normally, but
focus then reverts to the XWayland grab surface. This can be tested with
the example client from the issue.

Adding a new variant of `KeyboardFocusTarget` is annoying. Maybe it
could map to a different variant of the enum. But it presumably needs to
handle any `wl_surface` XWayland uses. (Override redirect surfaces?
Subsurfaces?) This seems as good as anything for now.

src/debug.rs

@@ -344,6 +344,7 @@ fn format_keyboard_focus(focus: Option<KeyboardFocusTarget>) -> String {
         Some(Popup(x)) => format!("Popup {}", x.wl_surface().id().protocol_id()),
         Some(Group(_)) => format!("Window Group"),
         Some(LockSurface(x)) => format!("LockSurface {}", x.wl_surface().id().protocol_id()),
+        Some(XWaylandGrab(x)) => format!("XWayland Grab {}", x.id().protocol_id()),
         None => format!("None"),
     }
 }

src/shell/focus/mod.rs

@@ -2,7 +2,7 @@ use crate::{
     shell::{element::CosmicMapped, Shell},
     state::Common,
     utils::prelude::*,
-    wayland::handlers::xdg_shell::PopupGrabData,
+    wayland::handlers::{xdg_shell::PopupGrabData, xwayland_keyboard_grab::XWaylandGrabSeatData},
 };
 use indexmap::IndexSet;
 use smithay::{
@@ -464,6 +464,18 @@ fn focus_target_is_valid(
         return matches!(target, KeyboardFocusTarget::LockSurface(_));
     }
 
+    let xwayland_grab = seat
+        .user_data()
+        .get_or_insert(XWaylandGrabSeatData::default)
+        .grab
+        .lock()
+        .unwrap();
+    if let Some((surface, grab)) = &*xwayland_grab {
+        if grab.grab().is_alive() {
+            return target == KeyboardFocusTarget::XWaylandGrab(surface.clone());
+        }
+    }
+
     // If an exclusive layer shell surface exists (on any output), only exclusive
     // shell surfaces can have focus, on the highest layer with exclusive surfaces.
     if let Some(layer) = exclusive_layer_surface_layer(shell) {
@@ -533,6 +545,7 @@ fn focus_target_is_valid(
         }
         KeyboardFocusTarget::Popup(_) => true,
         KeyboardFocusTarget::LockSurface(_) => false,
+        KeyboardFocusTarget::XWaylandGrab(_) => false,
     }
 }
 
@@ -541,12 +554,21 @@ fn update_focus_target(
     seat: &Seat<State>,
     output: &Output,
 ) -> Option<KeyboardFocusTarget> {
+    let mut xwayland_grab = seat
+        .user_data()
+        .get_or_insert(XWaylandGrabSeatData::default)
+        .grab
+        .lock()
+        .unwrap();
+    xwayland_grab.take_if(|(_, g)| !g.grab().is_alive());
     if let Some(session_lock) = &shell.session_lock {
         session_lock
             .surfaces
             .get(output)
             .cloned()
             .map(KeyboardFocusTarget::from)
+    } else if let Some((surface, _)) = &*xwayland_grab {
+        Some(KeyboardFocusTarget::XWaylandGrab(surface.clone()))
     } else if let Some(layer) = exclusive_layer_surface_layer(shell) {
         layer_map_for_output(output)
             .layers()

src/shell/focus/target.rs

@@ -74,6 +74,7 @@ pub enum KeyboardFocusTarget {
     LayerSurface(LayerSurface),
     Popup(PopupKind),
     LockSurface(LockSurface),
+    XWaylandGrab(WlSurface),
 }
 
 // TODO: This should be TryFrom, but PopupGrab needs to be able to convert. Fix this in smithay
@@ -234,6 +235,7 @@ impl IsAlive for KeyboardFocusTarget {
             KeyboardFocusTarget::LayerSurface(l) => l.alive(),
             KeyboardFocusTarget::Popup(p) => p.alive(),
             KeyboardFocusTarget::LockSurface(l) => l.alive(),
+            KeyboardFocusTarget::XWaylandGrab(g) => g.alive(),
         }
     }
 }
@@ -680,6 +682,9 @@ impl KeyboardTarget<State> for KeyboardFocusTarget {
             KeyboardFocusTarget::LockSurface(l) => {
                 KeyboardTarget::enter(l.wl_surface(), seat, data, keys, serial)
             }
+            KeyboardFocusTarget::XWaylandGrab(g) => {
+                KeyboardTarget::enter(g, seat, data, keys, serial)
+            }
         }
     }
     fn leave(&self, seat: &Seat<State>, data: &mut State, serial: Serial) {
@@ -696,6 +701,7 @@ impl KeyboardTarget<State> for KeyboardFocusTarget {
             KeyboardFocusTarget::LockSurface(l) => {
                 KeyboardTarget::leave(l.wl_surface(), seat, data, serial)
             }
+            KeyboardFocusTarget::XWaylandGrab(g) => KeyboardTarget::leave(g, seat, data, serial),
         }
     }
     fn key(
@@ -724,6 +730,9 @@ impl KeyboardTarget<State> for KeyboardFocusTarget {
             KeyboardFocusTarget::LockSurface(l) => {
                 KeyboardTarget::key(l.wl_surface(), seat, data, key, state, serial, time)
             }
+            KeyboardFocusTarget::XWaylandGrab(g) => {
+                KeyboardTarget::key(g, seat, data, key, state, serial, time)
+            }
         }
     }
     fn modifiers(
@@ -750,6 +759,9 @@ impl KeyboardTarget<State> for KeyboardFocusTarget {
             KeyboardFocusTarget::LockSurface(l) => {
                 KeyboardTarget::modifiers(l.wl_surface(), seat, data, modifiers, serial)
             }
+            KeyboardFocusTarget::XWaylandGrab(g) => {
+                KeyboardTarget::modifiers(g, seat, data, modifiers, serial)
+            }
         }
     }
     fn replace(
@@ -781,6 +793,7 @@ impl WaylandFocus for KeyboardFocusTarget {
             KeyboardFocusTarget::LayerSurface(l) => Some(Cow::Borrowed(l.wl_surface())),
             KeyboardFocusTarget::Popup(p) => Some(Cow::Borrowed(p.wl_surface())),
             KeyboardFocusTarget::LockSurface(l) => Some(Cow::Borrowed(l.wl_surface())),
+            KeyboardFocusTarget::XWaylandGrab(g) => Some(Cow::Borrowed(g)),
         }
     }
     fn same_client_as(&self, object_id: &ObjectId) -> bool {
@@ -791,6 +804,7 @@ impl WaylandFocus for KeyboardFocusTarget {
             KeyboardFocusTarget::LayerSurface(l) => l.wl_surface().id().same_client_as(object_id),
             KeyboardFocusTarget::Popup(p) => p.wl_surface().id().same_client_as(object_id),
             KeyboardFocusTarget::LockSurface(l) => l.wl_surface().id().same_client_as(object_id),
+            KeyboardFocusTarget::XWaylandGrab(g) => g.id().same_client_as(object_id),
         }
     }
 }

src/shell/mod.rs

@@ -53,6 +53,7 @@ use smithay::{
         session_lock::LockSurface,
         shell::wlr_layer::{KeyboardInteractivity, Layer, LayerSurfaceCachedState},
         xdg_activation::XdgActivationState,
+        xwayland_keyboard_grab::XWaylandKeyboardGrab,
     },
     xwayland::X11Surface,
 };
@@ -266,6 +267,7 @@ pub struct Shell {
     pub session_lock: Option<SessionLock>,
     pub seats: Seats,
     pub previous_workspace_idx: Option<(Serial, WeakOutput, usize)>,
+    pub xwayland_keyboard_grab: Option<XWaylandKeyboardGrab<State>>,
 
     theme: cosmic::Theme,
     pub active_hint: bool,
@@ -1474,6 +1476,7 @@ impl Shell {
             override_redirect_windows: Vec::new(),
             session_lock: None,
             previous_workspace_idx: None,
+            xwayland_keyboard_grab: None,
 
             theme,
             active_hint: config.cosmic_conf.active_hint,
@@ -1661,7 +1664,25 @@ impl Shell {
             }?;
 
             focus_target = KeyboardFocusTarget::Element(new_target);
-        };
+        } else if let KeyboardFocusTarget::XWaylandGrab(surface) = &focus_target {
+            if let Some(new_target) = self.element_for_surface(surface) {
+                focus_target = KeyboardFocusTarget::Element(new_target.clone());
+            } else if let Some(or) = self
+                .override_redirect_windows
+                .iter()
+                .find(|w| w.wl_surface().as_ref() == Some(surface))
+            {
+                // Find output the override redirect window overlaps the most with
+                let or_geo = or.geometry().as_global();
+                let (output, _) = self
+                    .outputs()
+                    .filter_map(|o| Some((o, o.geometry().intersection(or_geo)?)))
+                    .max_by_key(|(_, intersection)| intersection.size.w * intersection.size.h)?;
+                return Some(output.clone());
+            } else {
+                return None;
+            }
+        }
 
         match focus_target {
             KeyboardFocusTarget::Element(elem) => {
@@ -1726,6 +1747,7 @@ impl Shell {
                 .iter()
                 .find_map(|(output, s)| (s == &surface).then_some(output))
                 .cloned(),
+            KeyboardFocusTarget::XWaylandGrab(_) => unreachable!(),
             KeyboardFocusTarget::Popup(_) => unreachable!(),
         }
     }

src/wayland/handlers/xwayland_keyboard_grab.rs

@@ -2,11 +2,26 @@
 
 use crate::{shell::focus::target::KeyboardFocusTarget, state::State};
 use smithay::{
-    delegate_xwayland_keyboard_grab, reexports::wayland_server::protocol::wl_surface::WlSurface,
-    wayland::xwayland_keyboard_grab::XWaylandKeyboardGrabHandler,
+    delegate_xwayland_keyboard_grab,
+    input::Seat,
+    reexports::wayland_server::protocol::wl_surface::WlSurface,
+    wayland::xwayland_keyboard_grab::{XWaylandKeyboardGrab, XWaylandKeyboardGrabHandler},
 };
+use std::sync::Mutex;
+
+#[derive(Default)]
+pub struct XWaylandGrabSeatData {
+    pub grab: Mutex<Option<(WlSurface, XWaylandKeyboardGrab<State>)>>,
+}
 
 impl XWaylandKeyboardGrabHandler for State {
+    fn grab(&mut self, surface: WlSurface, seat: Seat<Self>, grab: XWaylandKeyboardGrab<Self>) {
+        let data = seat
+            .user_data()
+            .get_or_insert(XWaylandGrabSeatData::default);
+        *data.grab.lock().unwrap() = Some((surface, grab));
+    }
+
     fn keyboard_focus_for_xsurface(&self, surface: &WlSurface) -> Option<KeyboardFocusTarget> {
         let element = self
             .common