Implement popup grabs

2 files changed, 155 insertions, 7 deletions

diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 3a0caee9..29907e49 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -1,7 +1,9 @@
 use smithay::desktop::{
     find_popup_root_surface, get_popup_toplevel_coords, layer_map_for_output, LayerSurface,
-    PopupKind, PopupManager, Window, WindowSurfaceType,
+    PopupKeyboardGrab, PopupKind, PopupManager, PopupPointerGrab, PopupUngrabStrategy, Window,
+    WindowSurfaceType,
 };
+use smithay::input::pointer::Focus;
 use smithay::output::Output;
 use smithay::reexports::wayland_protocols::xdg::decoration::zv1::server::zxdg_toplevel_decoration_v1;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_positioner::ConstraintAdjustment;
@@ -12,6 +14,7 @@ use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::utils::{Logical, Rectangle, Serial};
 use smithay::wayland::compositor::{send_surface_state, with_states};
 use smithay::wayland::shell::kde::decoration::{KdeDecorationHandler, KdeDecorationState};
+use smithay::wayland::shell::wlr_layer::Layer;
 use smithay::wayland::shell::xdg::decoration::XdgDecorationHandler;
 use smithay::wayland::shell::xdg::{
     PopupSurface, PositionerState, ToplevelSurface, XdgPopupSurfaceData, XdgShellHandler,
@@ -19,7 +22,7 @@ use smithay::wayland::shell::xdg::{
 };
 use smithay::{delegate_kde_decoration, delegate_xdg_decoration, delegate_xdg_shell};
 
-use crate::niri::State;
+use crate::niri::{PopupGrabState, State};
 use crate::utils::clone2;
 
 impl XdgShellHandler for State {
@@ -90,8 +93,93 @@ impl XdgShellHandler for State {
         surface.send_repositioned(token);
     }
 
-    fn grab(&mut self, _surface: PopupSurface, _seat: WlSeat, _serial: Serial) {
-        // FIXME popup grabs
+    fn grab(&mut self, surface: PopupSurface, _seat: WlSeat, serial: Serial) {
+        let popup = PopupKind::Xdg(surface);
+        let Ok(root) = find_popup_root_surface(&popup) else {
+            return;
+        };
+
+        // We need to hand out the grab in a way consistent with what update_keyboard_focus()
+        // thinks the current focus is, otherwise it will desync and cause weird issues with
+        // keyboard focus being at the wrong place.
+        if self.niri.is_locked() {
+            if Some(&root) != self.niri.lock_surface_focus().as_ref() {
+                let _ = PopupManager::dismiss_popup(&root, &popup);
+                return;
+            }
+        } else if self.niri.screenshot_ui.is_open() {
+            let _ = PopupManager::dismiss_popup(&root, &popup);
+            return;
+        } else if let Some(output) = self.niri.layout.active_output() {
+            let layers = layer_map_for_output(output);
+
+            if let Some(layer_surface) =
+                layers.layer_for_surface(&root, WindowSurfaceType::TOPLEVEL)
+            {
+                if !matches!(layer_surface.layer(), Layer::Overlay | Layer::Top) {
+                    let _ = PopupManager::dismiss_popup(&root, &popup);
+                    return;
+                }
+            } else {
+                if layers
+                    .layers_on(Layer::Overlay)
+                    .any(|l| l.can_receive_keyboard_focus())
+                {
+                    let _ = PopupManager::dismiss_popup(&root, &popup);
+                    return;
+                }
+
+                let mon = self.niri.layout.monitor_for_output(output).unwrap();
+                if !mon.render_above_top_layer()
+                    && layers
+                        .layers_on(Layer::Top)
+                        .any(|l| l.can_receive_keyboard_focus())
+                {
+                    let _ = PopupManager::dismiss_popup(&root, &popup);
+                    return;
+                }
+
+                let layout_focus = self.niri.layout.focus();
+                if Some(&root) != layout_focus.map(|win| win.toplevel().wl_surface()) {
+                    let _ = PopupManager::dismiss_popup(&root, &popup);
+                    return;
+                }
+            }
+        } else {
+            let _ = PopupManager::dismiss_popup(&root, &popup);
+            return;
+        }
+
+        let seat = &self.niri.seat;
+        let Ok(mut grab) = self
+            .niri
+            .popups
+            .grab_popup(root.clone(), popup, seat, serial)
+        else {
+            return;
+        };
+
+        let keyboard = seat.get_keyboard().unwrap();
+        let pointer = seat.get_pointer().unwrap();
+
+        let keyboard_grab_mismatches = keyboard.is_grabbed()
+            && !(keyboard.has_grab(serial)
+                || grab
+                    .previous_serial()
+                    .map_or(true, |s| keyboard.has_grab(s)));
+        let pointer_grab_mismatches = pointer.is_grabbed()
+            && !(pointer.has_grab(serial)
+                || grab.previous_serial().map_or(true, |s| pointer.has_grab(s)));
+        if keyboard_grab_mismatches || pointer_grab_mismatches {
+            grab.ungrab(PopupUngrabStrategy::All);
+            return;
+        }
+
+        trace!("new grab for root {:?}", root);
+        keyboard.set_focus(self, grab.current_grab(), serial);
+        keyboard.set_grab(PopupKeyboardGrab::new(&grab), serial);
+        pointer.set_grab(self, PopupPointerGrab::new(&grab), serial, Focus::Keep);
+        self.niri.popup_grab = Some(PopupGrabState { root, grab });
     }
 
     fn maximize_request(&mut self, surface: ToplevelSurface) {
diff --git a/src/niri.rs b/src/niri.rs
index 78679146..a39530a4 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -35,7 +35,8 @@ use smithay::desktop::utils::{
     under_from_surface_tree, update_surface_primary_scanout_output, OutputPresentationFeedback,
 };
 use smithay::desktop::{
-    layer_map_for_output, LayerSurface, PopupManager, Space, Window, WindowSurfaceType,
+    layer_map_for_output, LayerSurface, PopupGrab, PopupManager, PopupUngrabStrategy, Space,
+    Window, WindowSurfaceType,
 };
 use smithay::input::keyboard::{Layout as KeyboardLayout, XkbContextHandler};
 use smithay::input::pointer::{CursorIcon, CursorImageAttributes, CursorImageStatus, MotionEvent};
@@ -157,6 +158,7 @@ pub struct Niri {
     pub primary_selection_state: PrimarySelectionState,
     pub data_control_state: DataControlState,
     pub popups: PopupManager,
+    pub popup_grab: Option<PopupGrabState>,
     pub presentation_state: PresentationState,
 
     pub seat: Seat<State>,
@@ -226,6 +228,11 @@ pub enum RedrawState {
     WaitingForEstimatedVBlankAndQueued((RegistrationToken, Idle<'static>)),
 }
 
+pub struct PopupGrabState {
+    pub root: WlSurface,
+    pub grab: PopupGrab<State>,
+}
+
 #[derive(Clone, PartialEq, Eq)]
 pub struct PointerFocus {
     pub output: Output,
@@ -303,6 +310,7 @@ impl State {
         self.niri.cursor_manager.check_cursor_image_surface_alive();
         self.niri.refresh_pointer_outputs();
         self.niri.popups.cleanup();
+        self.niri.refresh_popup_grab();
         self.update_keyboard_focus();
         self.refresh_pointer_focus();
 
@@ -405,6 +413,17 @@ impl State {
             let mon = self.niri.layout.monitor_for_output(output).unwrap();
             let layers = layer_map_for_output(output);
 
+            // Explicitly check for layer-shell popup grabs here, our keyboard focus will stay on
+            // the root layer surface while it has grabs.
+            let layer_grab = self.niri.popup_grab.as_ref().and_then(|g| {
+                layers
+                    .layer_for_surface(&g.root, WindowSurfaceType::TOPLEVEL)
+                    .map(|l| (&g.root, l.layer()))
+            });
+            let grab_on_layer = |layer: Layer| {
+                layer_grab.and_then(move |(s, l)| if l == layer { Some(s.clone()) } else { None })
+            };
+
             let layout_focus = || {
                 self.niri
                     .layout
@@ -417,7 +436,14 @@ impl State {
                     .then(|| surface.wl_surface().clone())
             };
 
-            let mut surface = layers.layers_on(Layer::Overlay).find_map(layer_focus);
+            let mut surface = grab_on_layer(Layer::Overlay);
+            // FIXME: we shouldn't prioritize the top layer grabs over regular overlay input or a
+            // fullscreen layout window. This will need tracking in grab() to avoid handing it out
+            // in the first place. Or a better way to structure this code.
+            surface = surface.or_else(|| grab_on_layer(Layer::Top));
+
+            surface = surface.or_else(|| layers.layers_on(Layer::Overlay).find_map(layer_focus));
+
             if mon.render_above_top_layer() {
                 surface = surface.or_else(layout_focus);
                 surface = surface.or_else(|| layers.layers_on(Layer::Top).find_map(layer_focus));
@@ -433,6 +459,31 @@ impl State {
 
         let keyboard = self.niri.seat.get_keyboard().unwrap();
         if self.niri.keyboard_focus != focus {
+            trace!(
+                "keyboard focus changed from {:?} to {:?}",
+                self.niri.keyboard_focus,
+                focus
+            );
+
+            if let Some(grab) = self.niri.popup_grab.as_mut() {
+                if Some(&grab.root) != focus.as_ref() {
+                    trace!(
+                        "grab root {:?} is not the new focus {:?}, ungrabbing",
+                        grab.root,
+                        focus
+                    );
+
+                    grab.grab.ungrab(PopupUngrabStrategy::All);
+                    keyboard.unset_grab();
+                    self.niri.seat.get_pointer().unwrap().unset_grab(
+                        self,
+                        SERIAL_COUNTER.next_serial(),
+                        get_monotonic_time().as_millis() as u32,
+                    );
+                    self.niri.popup_grab = None;
+                }
+            }
+
             if self.niri.config.borrow().input.keyboard.track_layout == TrackLayout::Window {
                 let current_layout =
                     keyboard.with_xkb_state(self, |context| context.active_layout());
@@ -783,6 +834,7 @@ impl Niri {
             primary_selection_state,
             data_control_state,
             popups: PopupManager::default(),
+            popup_grab: None,
             suppressed_keys: HashSet::new(),
             presentation_state,
 
@@ -1275,7 +1327,7 @@ impl Niri {
         layout_output.or_else(layer_shell_output)
     }
 
-    fn lock_surface_focus(&self) -> Option<WlSurface> {
+    pub fn lock_surface_focus(&self) -> Option<WlSurface> {
         let output_under_cursor = self.output_under_cursor();
         let output = output_under_cursor
             .as_ref()
@@ -1286,6 +1338,14 @@ impl Niri {
         state.lock_surface.as_ref().map(|s| s.wl_surface()).cloned()
     }
 
+    pub fn refresh_popup_grab(&mut self) {
+        if let Some(grab) = &self.popup_grab {
+            if grab.grab.has_ended() {
+                self.popup_grab = None;
+            }
+        }
+    }
+
     /// Schedules an immediate redraw on all outputs if one is not already scheduled.
     pub fn queue_redraw_all(&mut self) {
         let outputs: Vec<_> = self.output_state.keys().cloned().collect();