1 files changed, 171 insertions, 0 deletions
diff --git a/crates/wayland/src/lock.rs b/crates/wayland/src/lock.rs
new file mode 100644
index 0000000..51b6158
--- /dev/null
+++ b/crates/wayland/src/lock.rs
@@ -0,0 +1,171 @@
+//! Session lock (`ext-session-lock-v1`).
+//!
+//! Kept as its own module because the security-relevant invariant spans
+//! state, protocol handling *and* rendering: while [`SessionLock::locked`]
+//! is set, no client surface other than the lock surface may be drawn or
+//! receive input. The render helpers here are the "drawn" half; the input
+//! half is the locked-session branches in [`crate::input`], and the focus
+//! half is the guard in [`crate::state::CompState::set_keyboard_focus`].
+
+use smithay::backend::renderer::element::surface::{render_elements_from_surface_tree, WaylandSurfaceRenderElement};
+use smithay::backend::renderer::element::Kind;
+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::session_lock::{LockSurface, SessionLockHandler, SessionLockManagerState, SessionLocker};
+
+use crate::input::dwindow_wl_surface;
+use std::collections::{HashMap, HashSet};
+
+use crate::state::CompState;
+
+/// Session-lock runtime state (`ext-session-lock-v1`).
+///
+/// The security-relevant invariant is that `locked` gates *both* rendering
+/// and input: while it is set, no client surface other than a lock surface
+/// is drawn or receives events. `pending_confirm` exists because the
+/// protocol requires the compositor to confirm the lock only *after* a frame
+/// with no client content has actually been presented - confirming earlier
+/// would tell the locker "the screen is safe" while the user's windows were
+/// still on screen. So `lock()` only stashes the confirmation here, and the
+/// render paths call `confirm_lock_if_presented()` once such a frame is out.
+///
+/// With multiple monitors the locker creates **one lock surface per
+/// output**, and the confirmation must wait for *every* output to have both
+/// a surface and a presented frame - otherwise a second monitor could still
+/// be showing the desktop at the moment the locker is told the session is
+/// locked.
+#[derive(Default)]
+pub(crate) struct SessionLock {
+    pub(crate) locked: bool,
+    /// Lock surface per output, keyed by `Output::name()`.
+    pub(crate) surfaces: HashMap<String, LockSurface>,
+    pub(crate) pending_confirm: Option<SessionLocker>,
+    /// Outputs that have presented a client-content-free frame since the
+    /// lock request, keyed the same way.
+    pub(crate) presented: HashSet<String>,
+}
+
+impl CompState {
+    /// Records that `output` has presented a lock frame, and confirms the
+    /// lock once *every* output has done so. Called from both backends right
+    /// after a frame is presented. See `SessionLock::pending_confirm`.
+    pub(crate) fn confirm_lock_if_presented(&mut self, output: &Output) {
+        if !self.lock.locked || self.lock.pending_confirm.is_none() {
+            return;
+        }
+        self.lock.presented.insert(output.name());
+
+        // Every output must be both covered by a lock surface and have shown
+        // a frame of it. An output still missing its surface means the locker
+        // hasn't got to it yet - keep waiting rather than confirm early.
+        let all_covered = self
+            .outputs()
+            .all(|o| self.lock.surfaces.contains_key(&o.name()) && self.lock.presented.contains(&o.name()));
+        if !all_covered {
+            return;
+        }
+        if let Some(confirm) = self.lock.pending_confirm.take() {
+            log::info!("session lock: all {} output(s) presented a cleared frame, confirming lock", self.outputs.len());
+            confirm.lock();
+        }
+    }
+
+    /// The lock surface covering `output`, if we are locked and the locker
+    /// has created one for it. Used by the render paths and input routing.
+    pub(crate) fn lock_surface_for(&self, output: &Output) -> Option<&WlSurface> {
+        if !self.lock.locked {
+            return None;
+        }
+        self.lock.surfaces.get(&output.name()).map(|s| s.wl_surface())
+    }
+
+    /// Any lock surface at all - used by input routing, which needs a
+    /// keyboard-focus target rather than a per-output one.
+    pub(crate) fn any_lock_surface(&self) -> Option<&WlSurface> {
+        if !self.lock.locked {
+            return None;
+        }
+        self.lock.surfaces.values().next().map(|s| s.wl_surface())
+    }
+}
+
+impl SessionLockHandler for CompState {
+    fn lock_state(&mut self) -> &mut SessionLockManagerState {
+        &mut self.session_lock_state
+    }
+
+    fn lock(&mut self, confirmation: SessionLocker) {
+        log::info!("session lock: locking");
+        self.lock.locked = true;
+        self.lock.pending_confirm = Some(confirmation);
+        // Drop keyboard focus off whatever client had it immediately, so no
+        // keystroke can reach a normal client in the window between the lock
+        // request and the lock surface being mapped. Focus moves to the lock
+        // surface in `new_surface` once it exists.
+        self.set_keyboard_focus(None);
+    }
+
+    fn unlock(&mut self) {
+        log::info!("session lock: unlocking");
+        self.lock.locked = false;
+        self.lock.surfaces.clear();
+        self.lock.presented.clear();
+        self.lock.pending_confirm = None;
+        // Hand focus back to whatever srdwm considers the focused window.
+        let surface = self
+            .wm
+            .borrow()
+            .focused_id()
+            .and_then(|id| self.id_to_window.get(&id).cloned())
+            .and_then(|w| dwindow_wl_surface(&w));
+        self.set_keyboard_focus(surface);
+    }
+
+    fn new_surface(&mut self, surface: LockSurface, wl_output: WlOutput) {
+        // The locker creates one surface per output and names which; size it
+        // to that output specifically, since monitors differ in resolution.
+        let Some(entry) = self.output_for_wl(&wl_output) else {
+            log::warn!("session lock: lock surface for an output we don't drive; ignoring");
+            return;
+        };
+        let name = entry.output.name();
+        let size = entry.size();
+        surface.with_pending_state(|state| {
+            state.size = Some((size.w as u32, size.h as u32).into());
+        });
+        surface.send_configure();
+        // Store before focusing: `set_keyboard_focus`'s locked-session guard
+        // checks the focus target *against* the stored lock surfaces, so
+        // setting it afterwards would make the guard reject this very
+        // surface.
+        let wl_surface = surface.wl_surface().clone();
+        self.lock.surfaces.insert(name, surface);
+        self.set_keyboard_focus(Some(wl_surface));
+    }
+}
+
+/// Render elements for a locked session: the lock surface alone, at the
+/// output origin. Deliberately *not* built from the `Space` or the
+/// `LayerMap` - while locked, nothing else may reach the screen, so the
+/// caller pairs this with an opaque black clear color. Shared by both
+/// backends (winit and udev), which differ only in their renderer type.
+/// Takes the surface rather than `&CompState` so the udev backend can call
+/// it while already holding a `&mut` borrow of its own `UdevOutput`.
+pub(crate) fn lock_render_elements<R>(lock_surface: Option<&WlSurface>, renderer: &mut R) -> Vec<WaylandSurfaceRenderElement<R>>
+where
+    R: smithay::backend::renderer::Renderer + smithay::backend::renderer::ImportAll,
+    R::TextureId: Clone + 'static,
+{
+    let Some(surface) = lock_surface else { return Vec::new() };
+    render_elements_from_surface_tree(renderer, surface, (0, 0), 1.0, 1.0, Kind::Unspecified)
+}
+
+/// Frame callbacks for the lock surface. Without these the locker never
+/// gets to draw a second frame - no cursor blink, no password-dot
+/// feedback, no failed-attempt shake.
+pub(crate) fn send_lock_frame(lock_surface: Option<&WlSurface>, output: &Output, time: std::time::Duration) {
+    if let Some(surface) = lock_surface {
+        smithay::desktop::utils::send_frames_surface_tree(surface, output, time, None, |_, _| Some(output.clone()));
+    }
+}