Handle session lock surfaces in visible_outputs_for_surface

src/backend/render/mod.rs

@@ -13,9 +13,10 @@ use crate::debug::{fps_ui, profiler_ui};
 use crate::{
     shell::{
         focus::target::WindowGroup, grabs::SeatMoveGrabState, layout::tiling::ANIMATION_DURATION,
-        CosmicMapped, CosmicMappedRenderElement, OverviewMode, Trigger, WorkspaceRenderElement,
+        CosmicMapped, CosmicMappedRenderElement, OverviewMode, SessionLock, Trigger,
+        WorkspaceRenderElement,
     },
-    state::{Common, Fps, SessionLock},
+    state::{Common, Fps},
     utils::prelude::*,
     wayland::{
         handlers::{
@@ -487,7 +488,7 @@ where
     }
 
     // If session locked, only show session lock surfaces
-    if let Some(session_lock) = &state.session_lock {
+    if let Some(session_lock) = &state.shell.session_lock {
         elements.extend(
             session_lock_elements(renderer, output, session_lock)
                 .into_iter()

src/input/mod.rs

@@ -10,10 +10,10 @@ use crate::{
             floating::ResizeGrabMarker,
             tiling::{SwapWindowGrab, TilingLayout},
         },
-        Direction, FocusResult, MoveResult, OverviewMode, ResizeDirection, ResizeMode, Trigger,
-        Workspace,
+        Direction, FocusResult, MoveResult, OverviewMode, ResizeDirection, ResizeMode, SessionLock,
+        Trigger, Workspace,
     },
-    state::{Common, SessionLock},
+    state::Common,
     utils::prelude::*,
     wayland::{
         handlers::{screencopy::ScreencopySessions, xdg_activation::ActivationContext},
@@ -569,7 +569,7 @@ impl State {
                         &self.common.shell.override_redirect_windows,
                         overview.0.clone(),
                         workspace,
-                        self.common.session_lock.as_ref(),
+                        self.common.shell.session_lock.as_ref(),
                     )
                     .map(|(target, pos)| (target, pos.as_logical()));
 
@@ -651,7 +651,7 @@ impl State {
                         &self.common.shell.override_redirect_windows,
                         overview.0,
                         workspace,
-                        self.common.session_lock.as_ref(),
+                        self.common.shell.session_lock.as_ref(),
                     )
                     .map(|(target, pos)| (target, pos.as_logical()));
 
@@ -777,7 +777,7 @@ impl State {
                         &self.common.shell.override_redirect_windows,
                         overview.0,
                         workspace,
-                        self.common.session_lock.as_ref(),
+                        self.common.shell.session_lock.as_ref(),
                     )
                     .map(|(target, pos)| (target, pos.as_logical()));
 
@@ -843,18 +843,20 @@ impl State {
                             && !seat.get_keyboard().map(|k| k.is_grabbed()).unwrap_or(false)
                         {
                             let output = seat.active_output();
+
                             let pos = seat.get_pointer().unwrap().current_location().as_global();
                             let relative_pos = pos.to_local(&output);
                             let overview = self.common.shell.overview_mode();
-                            let workspace = self.common.shell.active_space_mut(&output);
                             let mut under = None;
 
-                            if let Some(session_lock) = self.common.session_lock.as_ref() {
+                            if let Some(session_lock) = self.common.shell.session_lock.as_ref() {
                                 under = session_lock
                                     .surfaces
                                     .get(&output)
                                     .map(|lock| lock.clone().into());
-                            } else if let Some(window) = workspace.get_fullscreen() {
+                            } else if let Some(window) =
+                                self.common.shell.active_space(&output).get_fullscreen()
+                            {
                                 let layers = layer_map_for_output(&output);
                                 if let Some(layer) =
                                     layers.layer_under(WlrLayer::Overlay, relative_pos.as_logical())
@@ -874,6 +876,7 @@ impl State {
                                     under = Some(window.clone().into());
                                 }
                             } else {
+                                let workspace = self.common.shell.active_space_mut(&output);
                                 let done = {
                                     let layers = layer_map_for_output(&output);
                                     if let Some(layer) = layers
@@ -1157,7 +1160,7 @@ impl State {
     ) {
         // TODO: Detect if started from login manager or tty, and only allow
         // `Terminate` if it will return to login manager.
-        if self.common.session_lock.is_some()
+        if self.common.shell.session_lock.is_some()
             && !matches!(action, Action::Terminate | Action::Debug)
         {
             return;

src/shell/focus/mod.rs

@@ -308,7 +308,7 @@ fn focus_target_is_valid(
     target: KeyboardFocusTarget,
 ) -> bool {
     // If a session lock is active, only lock surfaces can be focused
-    if state.common.session_lock.is_some() {
+    if state.common.shell.session_lock.is_some() {
         return matches!(target, KeyboardFocusTarget::LockSurface(_));
     }
 
@@ -361,7 +361,7 @@ fn update_focus_target(
     seat: &Seat<State>,
     output: &Output,
 ) -> Option<KeyboardFocusTarget> {
-    if let Some(session_lock) = &state.common.session_lock {
+    if let Some(session_lock) = &state.common.shell.session_lock {
         session_lock
             .surfaces
             .get(output)

src/shell/mod.rs

@@ -19,11 +19,15 @@ use smithay::{
         Seat,
     },
     output::Output,
-    reexports::wayland_server::{protocol::wl_surface::WlSurface, Client, DisplayHandle},
+    reexports::{
+        wayland_protocols::ext::session_lock::v1::server::ext_session_lock_v1::ExtSessionLockV1,
+        wayland_server::{protocol::wl_surface::WlSurface, Client, DisplayHandle},
+    },
     utils::{Point, Rectangle, Serial, SERIAL_COUNTER},
     wayland::{
         compositor::with_states,
         seat::WaylandFocus,
+        session_lock::LockSurface,
         shell::{
             wlr_layer::{
                 KeyboardInteractivity, Layer, LayerSurfaceCachedState, WlrLayerShellState,
@@ -175,6 +179,7 @@ pub struct Shell {
     pub pending_layers: Vec<(LayerSurface, Output, Seat<State>)>,
     pub pending_activations: HashMap<ActivationKey, ActivationContext>,
     pub override_redirect_windows: Vec<X11Surface>,
+    pub session_lock: Option<SessionLock>,
 
     // wayland_state
     pub layer_shell_state: WlrLayerShellState,
@@ -199,6 +204,12 @@ pub struct Shell {
     resize_indicator: Option<ResizeIndicator>,
 }
 
+#[derive(Debug)]
+pub struct SessionLock {
+    pub ext_session_lock: ExtSessionLockV1,
+    pub surfaces: HashMap<Output, LockSurface>,
+}
+
 #[derive(Debug)]
 pub struct WorkspaceSet {
     previously_active: Option<(usize, Instant)>,
@@ -916,6 +927,7 @@ impl Shell {
             pending_layers: Vec::new(),
             pending_activations: HashMap::new(),
             override_redirect_windows: Vec::new(),
+            session_lock: None,
 
             layer_shell_state,
             toplevel_info_state,
@@ -1015,6 +1027,15 @@ impl Shell {
         &'a self,
         surface: &'a WlSurface,
     ) -> impl Iterator<Item = Output> + 'a {
+        if let Some(session_lock) = &self.session_lock {
+            let output = session_lock
+                .surfaces
+                .iter()
+                .find(|(_, v)| v.wl_surface() == surface)
+                .map(|(k, _)| k.clone());
+            return Box::new(output.into_iter()) as Box<dyn Iterator<Item = Output>>;
+        }
+
         match self
             .outputs()
             .find(|o| {

src/state.rs

@@ -50,7 +50,6 @@ use smithay::{
     output::{Mode as OutputMode, Output, Scale},
     reexports::{
         calloop::{LoopHandle, LoopSignal},
-        wayland_protocols::ext::session_lock::v1::server::ext_session_lock_v1::ExtSessionLockV1,
         wayland_protocols_misc::server_decoration::server::org_kde_kwin_server_decoration_manager::Mode,
         wayland_server::{
             backend::{ClientData, ClientId, DisconnectReason},
@@ -75,7 +74,7 @@ use smithay::{
             data_device::DataDeviceState, primary_selection::PrimarySelectionState,
             wlr_data_control::DataControlState,
         },
-        session_lock::{LockSurface, SessionLockManagerState},
+        session_lock::SessionLockManagerState,
         shell::{kde::decoration::KdeDecorationState, xdg::decoration::XdgDecorationState},
         shm::ShmState,
         text_input::TextInputManagerState,
@@ -87,10 +86,7 @@ use smithay::{
 use tracing::error;
 
 use std::{cell::RefCell, ffi::OsString, time::Duration};
-use std::{
-    collections::{HashMap, VecDeque},
-    time::Instant,
-};
+use std::{collections::VecDeque, time::Instant};
 
 #[derive(RustEmbed)]
 #[folder = "resources/i18n"]
@@ -173,8 +169,6 @@ pub struct Common {
     pub kde_decoration_state: KdeDecorationState,
     pub xdg_decoration_state: XdgDecorationState,
 
-    pub session_lock: Option<SessionLock>,
-
     // xwayland state
     pub xwayland_state: Option<XWaylandState>,
 }
@@ -195,12 +189,6 @@ pub struct SurfaceDmabufFeedback {
     pub scanout_feedback: DmabufFeedback,
 }
 
-#[derive(Debug)]
-pub struct SessionLock {
-    pub ext_session_lock: ExtSessionLockV1,
-    pub surfaces: HashMap<Output, LockSurface>,
-}
-
 impl BackendData {
     pub fn kms(&mut self) -> &mut KmsState {
         match self {
@@ -414,8 +402,6 @@ impl State {
                 kde_decoration_state,
                 xdg_decoration_state,
 
-                session_lock: None,
-
                 xwayland_state: None,
             },
             backend: BackendData::Unset,
@@ -515,7 +501,7 @@ impl Common {
         let time = self.clock.now();
         let throttle = Some(Duration::from_secs(1));
 
-        if let Some(session_lock) = self.session_lock.as_ref() {
+        if let Some(session_lock) = self.shell.session_lock.as_ref() {
             if let Some(lock_surface) = session_lock.surfaces.get(output) {
                 with_surfaces_surface_tree(lock_surface.wl_surface(), |_surface, states| {
                     with_fractional_scale(states, |fraction_scale| {

src/wayland/handlers/session_lock.rs

@@ -1,9 +1,6 @@
 // SPDX-License-Identifier: GPL-3.0-only
 
-use crate::{
-    state::{SessionLock, State},
-    utils::prelude::*,
-};
+use crate::{shell::SessionLock, state::State, utils::prelude::*};
 use smithay::{
     delegate_session_lock,
     output::Output,
@@ -22,7 +19,7 @@ impl SessionLockHandler for State {
 
     fn lock(&mut self, locker: SessionLocker) {
         // Reject lock if sesion lock exists and is still valid
-        if let Some(session_lock) = self.common.session_lock.as_ref() {
+        if let Some(session_lock) = self.common.shell.session_lock.as_ref() {
             if self
                 .common
                 .display_handle
@@ -35,18 +32,28 @@ impl SessionLockHandler for State {
 
         let ext_session_lock = locker.ext_session_lock().clone();
         locker.lock();
-        self.common.session_lock = Some(SessionLock {
+        self.common.shell.session_lock = Some(SessionLock {
             ext_session_lock,
             surfaces: HashMap::new(),
         });
+
+        for output in self.common.shell.outputs() {
+            self.backend
+                .schedule_render(&self.common.event_loop_handle, &output, None);
+        }
     }
 
     fn unlock(&mut self) {
-        self.common.session_lock = None;
+        self.common.shell.session_lock = None;
+
+        for output in self.common.shell.outputs() {
+            self.backend
+                .schedule_render(&self.common.event_loop_handle, &output, None);
+        }
     }
 
     fn new_surface(&mut self, lock_surface: LockSurface, wl_output: WlOutput) {
-        if let Some(session_lock) = &mut self.common.session_lock {
+        if let Some(session_lock) = &mut self.common.shell.session_lock {
             if let Some(output) = Output::from_resource(&wl_output) {
                 lock_surface.with_pending_state(|states| {
                     let size = output.geometry().size;