Handle modifiers even when consumed by a11y

This is apparently required: Orca expects us to keep track of modifiers
sent in KeyboardMonitor even during an Orca modifier key combo.

2 files changed, 57 insertions, 15 deletions

diff --git a/src/dbus/freedesktop_a11y.rs b/src/dbus/freedesktop_a11y.rs
index 47ef621f..930f0f52 100644
--- a/src/dbus/freedesktop_a11y.rs
+++ b/src/dbus/freedesktop_a11y.rs
@@ -44,6 +44,17 @@ pub struct KeyboardMonitor {
     iface: Arc<OnceLock<InterfaceRef<Self>>>,
 }
 
+/// Keyboard monitor key block reason.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum KbMonBlock {
+    /// Not blocked.
+    Pass,
+    /// Blocked, and this is the first press/release of the a11y modifier.
+    ModifierFirstPress,
+    /// Blocked, and this is not the a11y modifier.
+    Block,
+}
+
 /// Interface for monitoring of keyboard input by assistive technologies.
 ///
 /// This interface is used by assistive technologies to monitor keyboard input of the compositor.
@@ -211,7 +222,7 @@ impl KeyboardMonitor {
         mods: u32,
         keysym: Keysym,
         unichar: u32,
-    ) -> bool {
+    ) -> KbMonBlock {
         let _span = tracy_client::span!("KeyboardMonitor::process_key");
 
         let mut ctxt = self.iface.get().unwrap().signal_emitter().clone();
@@ -250,7 +261,7 @@ impl KeyboardMonitor {
                 // second press that got handled normally.
                 if !data.suppressed_keys.contains(&keysym) {
                     trace!("handling release for second press of grabbed modifier: {keysym:?}");
-                    return false;
+                    return KbMonBlock::Pass;
                 }
             } else {
                 let last_press_entry = data
@@ -263,7 +274,7 @@ impl KeyboardMonitor {
                 // Modifier pressed twice; handle it as normal.
                 if time <= last_press.saturating_add(repeat_delay) {
                     trace!("handling second press of grabbed modifier: {keysym:?}");
-                    return false;
+                    return KbMonBlock::Pass;
                 }
             }
         }
@@ -293,7 +304,13 @@ impl KeyboardMonitor {
             }
         }
 
-        block
+        if !block {
+            KbMonBlock::Pass
+        } else if data.grabbed_mods.contains(&keysym) {
+            KbMonBlock::ModifierFirstPress
+        } else {
+            KbMonBlock::Block
+        }
     }
 }
 
@@ -430,9 +447,14 @@ impl Start for KeyboardMonitor {
 }
 
 impl State {
-    pub fn a11y_process_key(&mut self, time: Duration, keycode: Keycode, state: KeyState) -> bool {
+    pub fn a11y_process_key(
+        &mut self,
+        time: Duration,
+        keycode: Keycode,
+        state: KeyState,
+    ) -> KbMonBlock {
         if self.niri.a11y_keyboard_monitor.is_none() {
-            return false;
+            return KbMonBlock::Pass;
         }
 
         let keyboard = self.niri.seat.get_keyboard().unwrap();
@@ -454,7 +476,7 @@ impl State {
         let released = state == KeyState::Released;
 
         let Some(monitor) = &self.niri.a11y_keyboard_monitor else {
-            return false;
+            return KbMonBlock::Pass;
         };
         monitor.process_key(repeat_delay, time, keycode, released, mods, keysym, unichar)
     }
diff --git a/src/input/mod.rs b/src/input/mod.rs
index d584c940..8e640b7b 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -41,6 +41,8 @@ use touch_overview_grab::TouchOverviewGrab;
 use self::move_grab::MoveGrab;
 use self::resize_grab::ResizeGrab;
 use self::spatial_movement_grab::SpatialMovementGrab;
+#[cfg(feature = "dbus")]
+use crate::dbus::freedesktop_a11y::KbMonBlock;
 use crate::layout::scrolling::ScrollDirection;
 use crate::layout::{ActivateWindow, LayoutElement as _};
 use crate::niri::{CastTarget, PointerVisibility, State};
@@ -388,15 +390,26 @@ impl State {
 
         // Accessibility modifier grabs should override XKB state changes (e.g. Caps Lock), so we
         // need to process them before keyboard.input() below.
+        //
+        // Other accessibility-grabbed keys should still update our XKB state, but not cause any
+        // other changes.
         #[cfg(feature = "dbus")]
-        if self.a11y_process_key(
-            Duration::from_millis(u64::from(time)),
-            event.key_code(),
-            event.state(),
-        ) {
-            *consumed_by_a11y = true;
-            return;
-        }
+        let block = {
+            let block = self.a11y_process_key(
+                Duration::from_millis(u64::from(time)),
+                event.key_code(),
+                event.state(),
+            );
+            if block != KbMonBlock::Pass {
+                *consumed_by_a11y = true;
+            }
+            // The accessibility modifier first press must not change XKB state, so we return
+            // early here.
+            if block == KbMonBlock::ModifierFirstPress {
+                return;
+            }
+            block
+        };
 
         let Some(Some(bind)) = self.niri.seat.get_keyboard().unwrap().input(
             self,
@@ -405,6 +418,13 @@ impl State {
             serial,
             time,
             |this, mods, keysym| {
+                // After updating XKB state from accessibility-grabbed keys, return right away and
+                // don't handle them.
+                #[cfg(feature = "dbus")]
+                if block != KbMonBlock::Pass {
+                    return FilterResult::Intercept(None);
+                }
+
                 let key_code = event.key_code();
                 let modified = keysym.modified_sym();
                 let raw = keysym.raw_latin_sym_or_raw_current_sym();