From 62892d636112230dc16778b91331fbe97576b005 Mon Sep 17 00:00:00 2001
From: Ivan Molodetskikh <yalterz@gmail.com>
Date: Sat, 17 Feb 2024 07:47:06 +0400
Subject: [PATCH] Prevent locking while another lock client is already active

Fixes double swaylock from manual + swayidle.
---
 src/niri.rs | 25 ++++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/src/niri.rs b/src/niri.rs
index 97511b6ec..353376538 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -44,6 +44,7 @@ use smithay::reexports::calloop::{
     self, Idle, Interest, LoopHandle, LoopSignal, Mode, PostAction, RegistrationToken,
 };
 use smithay::reexports::input;
+use smithay::reexports::wayland_protocols::ext::session_lock::v1::server::ext_session_lock_v1::ExtSessionLockV1;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel::WmCapabilities;
 use smithay::reexports::wayland_protocols_misc::server_decoration as _server_decoration;
 use smithay::reexports::wayland_server::backend::{
@@ -263,7 +264,7 @@ pub enum LockState {
     #[default]
     Unlocked,
     Locking(SessionLocker),
-    Locked,
+    Locked(ExtSessionLockV1),
 }
 
 #[derive(PartialEq, Eq)]
@@ -1300,8 +1301,9 @@ impl Niri {
                     .all(|state| state.lock_render_state == LockRenderState::Locked);
 
                 if all_locked {
+                    let lock = confirmation.ext_session_lock().clone();
                     confirmation.lock();
-                    self.lock_state = LockState::Locked;
+                    self.lock_state = LockState::Locked(lock);
                 } else {
                     // Still waiting.
                     self.lock_state = LockState::Locking(confirmation);
@@ -2108,8 +2110,9 @@ impl Niri {
                         .all(|state| state.lock_render_state == LockRenderState::Locked);
 
                     if all_locked {
+                        let lock = confirmation.ext_session_lock().clone();
                         confirmation.lock();
-                        self.lock_state = LockState::Locked;
+                        self.lock_state = LockState::Locked(lock);
                     } else {
                         // Still waiting.
                         self.lock_state = LockState::Locking(confirmation);
@@ -2833,6 +2836,22 @@ impl Niri {
     }
 
     pub fn lock(&mut self, confirmation: SessionLocker) {
+        // Check if another client is in the process of locking.
+        if matches!(self.lock_state, LockState::Locking(_)) {
+            info!("refusing lock as another client is currently locking");
+            return;
+        }
+
+        // Check if we're already locked with an active client.
+        if let LockState::Locked(lock) = &self.lock_state {
+            if lock.is_alive() {
+                info!("refusing lock as already locked with an active client");
+                return;
+            }
+
+            // If the client had died, continue with the new lock.
+        }
+
         info!("locking session");
 
         self.screenshot_ui.close();