Added grace implementation

src/auth.rs

@@ -7,7 +7,7 @@ use smithay_client_toolkit::reexports::{calloop::channel, calloop::EventLoop};
 use users::get_current_username;
 use zeroize::Zeroizing;
 
-const SERVICE_NAME: &str = "wayrustlock";
+const SERVICE_NAME: &str = "rustlock";
 
 pub struct LockConversation {
     pub password: Option<Zeroizing<String>>,
@@ -36,12 +36,6 @@ impl pam_client::ConversationHandler for LockConversation {
 }
 
 pub fn create_and_run_auth_loop() -> (channel::Sender<Zeroizing<String>>, channel::Channel<bool>) {
-    struct AuthLoopState {
-        auth_res_send: channel::Sender<bool>,
-        main_closed: bool,
-        context: pam_client::Context<LockConversation>,
-    }
-
     let username = get_current_username()
         .expect("Failed to get username")
         .to_str()
@@ -49,7 +43,7 @@ pub fn create_and_run_auth_loop() -> (channel::Sender<Zeroizing<String>>, channe
         .to_string();
 
     let conversation = LockConversation { password: None };
-    let context = Context::new(SERVICE_NAME, Some(username.as_str()), conversation)
+    let _context = Context::new(SERVICE_NAME, Some(username.as_str()), conversation)
         .expect("Failed to initialize PAM context");
     debug!("Prepared to authenticate user '{}'", username);
 
@@ -57,35 +51,33 @@ pub fn create_and_run_auth_loop() -> (channel::Sender<Zeroizing<String>>, channe
     let (auth_res_send, auth_res_recv) = channel::channel::<bool>();
 
     thread::spawn(move || {
-        let mut event_loop: EventLoop<AuthLoopState> = EventLoop::try_new().unwrap();
+        let mut event_loop: EventLoop<()> = EventLoop::try_new().unwrap();
         event_loop
             .handle()
-            .insert_source(auth_req_recv, |evt, _metadata, state| match evt {
+            .insert_source(auth_req_recv, |evt, _metadata, _state| match evt {
                 channel::Event::Msg(password) => {
-                    state.context.conversation_mut().password = Some(password);
-                    let status = match state.context.authenticate(Flag::NONE) {
-                        Ok(()) => true,
+                    let conversation = LockConversation {
+                        password: Some(password),
+                    };
+                    let mut context =
+                        Context::new(SERVICE_NAME, Some(username.as_str()), conversation)
+                            .expect("Failed to initialize PAM context");
+                    match context.authenticate(Flag::NONE) {
+                        Ok(()) => {
+                            auth_res_send.send(true).unwrap();
+                        }
                         Err(err) => {
                             error!("Pam authenticate failed with {:?}", err);
-                            false
+                            auth_res_send.send(false).unwrap();
                         }
-                    };
-                    state.auth_res_send.send(status).unwrap();
+                    }
                 }
-                channel::Event::Closed => state.main_closed = true,
+                channel::Event::Closed => {}
             })
             .unwrap();
 
-        let mut state = AuthLoopState {
-            auth_res_send,
-            main_closed: false,
-            context,
-        };
-
-        while !state.main_closed {
-            event_loop
-                .dispatch(None, &mut state)
-                .expect("Failed to run");
+        loop {
+            event_loop.dispatch(None, &mut ()).expect("Failed to run");
         }
     });
 

src/config.rs

@@ -57,7 +57,7 @@ pub struct Config {
     #[arg(long)]
     pub debug: bool,
 
-    /// Write verbose logs to ~/.wayrustlock.log
+    /// Write verbose logs to ~/.rustlock.log
     #[arg(long)]
     pub log_file: bool,
 
@@ -69,18 +69,18 @@ pub struct Config {
 impl Config {
     pub fn load() -> Self {
         let cli_config = Config::parse();
-        
+
         // Use path from CLI if provided, otherwise default
         let config_path = cli_config.config.clone().unwrap_or_else(|| {
             let mut path = std::path::PathBuf::from(std::env::var("HOME").unwrap_or_default());
-            path.push(".config/wayrustlock/config.toml");
+            path.push(".config/rustlock/config.toml");
             path
         });
 
         if config_path.exists() {
             if let Ok(file_content) = std::fs::read_to_string(&config_path) {
                 if let Ok(_file_config) = toml::from_str::<Config>(&file_content) {
-                    return cli_config; 
+                    return cli_config;
                 }
             }
         }

src/main.rs

@@ -9,16 +9,18 @@ mod util;
 
 use config::Config;
 use lock::LockManager;
-use screenshot::{CaptureData, ScreenshotManager, Screenshot};
+use screenshot::{CaptureData, Screenshot, ScreenshotManager};
 use std::error::Error;
 use std::fs::OpenOptions;
 use std::io::Write;
 use std::sync::{Arc, Mutex};
-use std::time::Duration;
+use std::time::{Duration, Instant};
 use wayland_client::protocol::wl_output::WlOutput;
 use wayland_client::protocol::wl_surface::WlSurface;
-use wayland_client::{Connection, Dispatch, QueueHandle, Proxy};
-use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{ZwlrScreencopyFrameV1, Flags};
+use wayland_client::{Connection, Dispatch, Proxy, QueueHandle};
+use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{
+    Flags, ZwlrScreencopyFrameV1,
+};
 use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1;
 use zeroize::Zeroizing;
 
@@ -45,9 +47,7 @@ static FILE_LOGGER: std::sync::LazyLock<std::sync::Mutex<Option<std::fs::File>>>
     std::sync::LazyLock::new(|| std::sync::Mutex::new(None));
 
 fn setup_file_logging(_config: &Config) {
-    let log_path = std::env::var("HOME")
-        .unwrap_or_else(|_| "/tmp".to_string())
-        + "/.wayrustlock.log";
+    let log_path = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()) + "/.rustlock.log";
 
     match OpenOptions::new()
         .create(true)
@@ -124,10 +124,14 @@ struct WaylandLock {
     exit: bool,
     unlocking: bool,
     screenshot_manager: Option<ScreenshotManager>,
+    grace_until: Option<Instant>,
 }
 
 impl WaylandLock {
     fn handle_auth_result(&mut self, success: bool) {
+        // Clear grace period on any auth result
+        self.grace_until = None;
+
         if success {
             log::info!("✅ Authentication successful - unlocking session");
             if let Some(session_lock) = &self.session_lock {
@@ -159,6 +163,16 @@ impl WaylandLock {
     }
 
     fn handle_key_event(&mut self, event: KeyEvent) {
+        // Check if we're in the grace period (any key unlocks without password)
+        if let Some(grace_until) = self.grace_until {
+            if Instant::now() < grace_until {
+                self.handle_auth_result(true);
+                return;
+            } else {
+                self.grace_until = None;
+            }
+        }
+
         use crate::input::InputAction;
         use smithay_client_toolkit::seat::keyboard::Keysym;
 
@@ -167,7 +181,9 @@ impl WaylandLock {
             if let Ok(mut lock_manager) = self.lock_manager.lock() {
                 let mut password = Zeroizing::new(String::new());
                 for surface in &mut lock_manager.surfaces {
-                    if let Some(InputAction::SubmitPassword(p)) = surface.handle_key_event(event.clone()) {
+                    if let Some(InputAction::SubmitPassword(p)) =
+                        surface.handle_key_event(event.clone())
+                    {
                         password = p;
                     }
                 }
@@ -194,32 +210,75 @@ impl WaylandLock {
 }
 
 impl ProvidesRegistryState for WaylandLock {
-    fn registry(&mut self) -> &mut RegistryState { &mut self.registry_state }
+    fn registry(&mut self) -> &mut RegistryState {
+        &mut self.registry_state
+    }
     registry_handlers![OutputState, SeatState];
 }
 
 impl CompositorHandler for WaylandLock {
-    fn scale_factor_changed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _surface: &WlSurface, _new_factor: i32) {}
-    fn transform_changed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _surface: &WlSurface, _new_transform: wayland_client::protocol::wl_output::Transform) {}
-    fn frame(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _surface: &WlSurface, _time: u32) {}
-    fn surface_enter(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _surface: &WlSurface, _output: &WlOutput) {}
-    fn surface_leave(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _surface: &WlSurface, _output: &WlOutput) {}
+    fn scale_factor_changed(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _surface: &WlSurface,
+        _new_factor: i32,
+    ) {
+    }
+    fn transform_changed(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _surface: &WlSurface,
+        _new_transform: wayland_client::protocol::wl_output::Transform,
+    ) {
+    }
+    fn frame(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _surface: &WlSurface,
+        _time: u32,
+    ) {
+    }
+    fn surface_enter(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _surface: &WlSurface,
+        _output: &WlOutput,
+    ) {
+    }
+    fn surface_leave(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _surface: &WlSurface,
+        _output: &WlOutput,
+    ) {
+    }
 }
 
 impl OutputHandler for WaylandLock {
-    fn output_state(&mut self) -> &mut OutputState { &mut self.output_state }
+    fn output_state(&mut self) -> &mut OutputState {
+        &mut self.output_state
+    }
     fn new_output(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {}
     fn update_output(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {}
-    fn output_destroyed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {}
+    fn output_destroyed(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _output: WlOutput) {
+    }
 }
 
 impl ShmHandler for WaylandLock {
-    fn shm_state(&mut self) -> &mut Shm { &mut self.shm_state }
+    fn shm_state(&mut self) -> &mut Shm {
+        &mut self.shm_state
+    }
 }
 
 impl SessionLockHandler for WaylandLock {
     fn locked(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _session_lock: SessionLock) {
         log::info!("Session LOCKED confirmed by compositor");
+        self.grace_until = Some(Instant::now() + Duration::from_secs_f32(self.config.grace));
     }
 
     fn finished(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _lock: SessionLock) {
@@ -239,21 +298,27 @@ impl SessionLockHandler for WaylandLock {
         log::debug!("CONFIGURE callback: {}x{}", width, height);
 
         if let Ok(mut lock_manager) = self.lock_manager.lock() {
-            if let Some(locked_surface) = lock_manager.find_surface_by_wayland_surface(session_lock_surface.wl_surface()) {
+            if let Some(locked_surface) =
+                lock_manager.find_surface_by_wayland_surface(session_lock_surface.wl_surface())
+            {
                 locked_surface.resize(width as i32, height as i32);
                 locked_surface.set_configured();
-                
+
                 let output = locked_surface.output();
                 let output_id = Proxy::id(output);
                 let output_idx = self.outputs.iter().position(|o| Proxy::id(o) == output_id);
-                
+
                 if let Some(idx) = output_idx {
                     if let Some(bg) = self.captured_backgrounds.get(idx).and_then(|b| b.as_ref()) {
-                        log::info!("Applying background for output {} (ID: {:?})", idx, output_id);
+                        log::info!(
+                            "Applying background for output {} (ID: {:?})",
+                            idx,
+                            output_id
+                        );
                         locked_surface.set_background(bg.clone());
                     }
                 }
-                
+
                 locked_surface.update();
                 let _ = locked_surface.commit(&mut self.pool);
             }
@@ -262,40 +327,143 @@ impl SessionLockHandler for WaylandLock {
 }
 
 impl KeyboardHandler for WaylandLock {
-    fn enter(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard, _surface: &WlSurface, _serial: u32, _raw: &[u32], _keysyms: &[smithay_client_toolkit::seat::keyboard::Keysym]) {}
-    fn leave(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard, _surface: &WlSurface, _serial: u32) {}
-    fn press_key(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard, _serial: u32, event: KeyEvent) {
+    fn enter(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
+        _surface: &WlSurface,
+        _serial: u32,
+        _raw: &[u32],
+        _keysyms: &[smithay_client_toolkit::seat::keyboard::Keysym],
+    ) {
+    }
+    fn leave(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
+        _surface: &WlSurface,
+        _serial: u32,
+    ) {
+    }
+    fn press_key(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
+        _serial: u32,
+        event: KeyEvent,
+    ) {
         self.handle_key_event(event);
     }
-    fn release_key(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard, _serial: u32, _event: KeyEvent) {}
-    fn update_modifiers(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard, _serial: u32, _modifiers: Modifiers, _layout: u32) {}
+    fn release_key(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
+        _serial: u32,
+        _event: KeyEvent,
+    ) {
+    }
+    fn update_modifiers(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _keyboard: &wayland_client::protocol::wl_keyboard::WlKeyboard,
+        _serial: u32,
+        _modifiers: Modifiers,
+        _layout: u32,
+    ) {
+    }
 }
 
 impl SeatHandler for WaylandLock {
-    fn seat_state(&mut self) -> &mut SeatState { &mut self.seat_state }
-    fn new_seat(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _seat: wayland_client::protocol::wl_seat::WlSeat) {}
-    fn new_capability(&mut self, _conn: &Connection, qh: &QueueHandle<Self>, seat: wayland_client::protocol::wl_seat::WlSeat, capability: smithay_client_toolkit::seat::Capability) {
+    fn seat_state(&mut self) -> &mut SeatState {
+        &mut self.seat_state
+    }
+    fn new_seat(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _seat: wayland_client::protocol::wl_seat::WlSeat,
+    ) {
+    }
+    fn new_capability(
+        &mut self,
+        _conn: &Connection,
+        qh: &QueueHandle<Self>,
+        seat: wayland_client::protocol::wl_seat::WlSeat,
+        capability: smithay_client_toolkit::seat::Capability,
+    ) {
         if capability == smithay_client_toolkit::seat::Capability::Keyboard {
-            let _ = self.seat_state.get_keyboard_with_repeat(qh, &seat, None, self.loop_handle.clone(), Box::new(|_state, _kbd, _event| {}));
+            let _ = self.seat_state.get_keyboard_with_repeat(
+                qh,
+                &seat,
+                None,
+                self.loop_handle.clone(),
+                Box::new(|_state, _kbd, _event| {}),
+            );
         }
     }
-    fn remove_capability(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _seat: wayland_client::protocol::wl_seat::WlSeat, _capability: smithay_client_toolkit::seat::Capability) {}
-    fn remove_seat(&mut self, _conn: &Connection, _qh: &QueueHandle<Self>, _seat: wayland_client::protocol::wl_seat::WlSeat) {}
+    fn remove_capability(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _seat: wayland_client::protocol::wl_seat::WlSeat,
+        _capability: smithay_client_toolkit::seat::Capability,
+    ) {
+    }
+    fn remove_seat(
+        &mut self,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+        _seat: wayland_client::protocol::wl_seat::WlSeat,
+    ) {
+    }
 }
 
 impl Dispatch<ZwlrScreencopyManagerV1, ()> for WaylandLock {
-    fn event(_state: &mut Self, _proxy: &ZwlrScreencopyManagerV1, _event: <ZwlrScreencopyManagerV1 as Proxy>::Event, _data: &(), _conn: &Connection, _qh: &QueueHandle<Self>) {}
+    fn event(
+        _state: &mut Self,
+        _proxy: &ZwlrScreencopyManagerV1,
+        _event: <ZwlrScreencopyManagerV1 as Proxy>::Event,
+        _data: &(),
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+    ) {
+    }
 }
 
 impl Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
-    fn event(state: &mut Self, frame: &ZwlrScreencopyFrameV1, event: <ZwlrScreencopyFrameV1 as Proxy>::Event, data: &CaptureData, _conn: &Connection, _qh: &QueueHandle<Self>) {
+    fn event(
+        state: &mut Self,
+        frame: &ZwlrScreencopyFrameV1,
+        event: <ZwlrScreencopyFrameV1 as Proxy>::Event,
+        data: &CaptureData,
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+    ) {
         use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event;
         match event {
-            Event::Buffer { format, width, height, stride } => {
+            Event::Buffer {
+                format,
+                width,
+                height,
+                stride,
+            } => {
                 let format = format.into_result().unwrap();
                 let mut info = data.info.lock().unwrap();
-                *info = Some(screenshot::BufferInfo { width, height, stride, format });
-                match state.pool.create_buffer(width as i32, height as i32, stride as i32, format) {
+                *info = Some(screenshot::BufferInfo {
+                    width,
+                    height,
+                    stride,
+                    format,
+                });
+                match state
+                    .pool
+                    .create_buffer(width as i32, height as i32, stride as i32, format)
+                {
                     Ok((buffer, _canvas)) => {
                         frame.copy(buffer.wl_buffer());
                         *data.buffer.lock().unwrap() = Some(buffer);
@@ -303,7 +471,9 @@ impl Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
                     Err(e) => log::error!("Screencopy: Buffer creation failed: {:?}", e),
                 }
             }
-            Event::Flags { flags } => { *data.flags.lock().unwrap() = Some(flags.into_result().unwrap()); }
+            Event::Flags { flags } => {
+                *data.flags.lock().unwrap() = Some(flags.into_result().unwrap());
+            }
             Event::Ready { .. } => {
                 log::info!("Screencopy: Ready for output {}", data.output_idx);
                 if let Some(mgr) = &state.screenshot_manager {
@@ -311,7 +481,11 @@ impl Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
                     let info = data.info.lock().unwrap().take();
                     let flags = data.flags.lock().unwrap().take();
                     if let (Some(buffer), Some(info), Some(flags)) = (buffer, info, flags) {
-                        let handle = screenshot::ScreencopyBufferHandle { buffer, info, y_invert: flags.contains(Flags::YInvert) };
+                        let handle = screenshot::ScreencopyBufferHandle {
+                            buffer,
+                            info,
+                            y_invert: flags.contains(Flags::YInvert),
+                        };
                         if let Ok(surface) = mgr.buffer_to_surface(handle, &mut state.pool) {
                             let mut ss = Screenshot::new(surface);
                             let _ = ss.apply_effects(&state.config);
@@ -334,8 +508,18 @@ impl Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
     }
 }
 
-impl wayland_client::Dispatch<wayland_client::protocol::wl_registry::WlRegistry, ()> for WaylandLock {
-    fn event(_state: &mut Self, _proxy: &wayland_client::protocol::wl_registry::WlRegistry, _event: wayland_client::protocol::wl_registry::Event, _data: &(), _conn: &Connection, _qh: &QueueHandle<Self>) {}
+impl wayland_client::Dispatch<wayland_client::protocol::wl_registry::WlRegistry, ()>
+    for WaylandLock
+{
+    fn event(
+        _state: &mut Self,
+        _proxy: &wayland_client::protocol::wl_registry::WlRegistry,
+        _event: wayland_client::protocol::wl_registry::Event,
+        _data: &(),
+        _conn: &Connection,
+        _qh: &QueueHandle<Self>,
+    ) {
+    }
 }
 
 smithay_client_toolkit::delegate_compositor!(WaylandLock);
@@ -353,12 +537,13 @@ fn main() -> Result<(), Box<dyn Error>> {
     static LOGGER: DualLogger = DualLogger;
     log::set_logger(&LOGGER).map(|()| log::set_max_level(log::LevelFilter::Debug))?;
 
-    log::info!("Starting wayrustlock v{}", env!("CARGO_PKG_VERSION"));
+    log::info!("Starting rustlock v{}", env!("CARGO_PKG_VERSION"));
     let lock_manager = Arc::new(Mutex::new(LockManager::new(config.clone())));
     let ctrlc_exit = Arc::new(std::sync::atomic::AtomicBool::new(false));
 
     let conn = Connection::connect_to_env()?;
-    let (globals, mut event_queue) = wayland_client::globals::registry_queue_init::<WaylandLock>(&conn)?;
+    let (globals, mut event_queue) =
+        wayland_client::globals::registry_queue_init::<WaylandLock>(&conn)?;
     let qh: QueueHandle<WaylandLock> = event_queue.handle();
 
     let shm_state = Shm::bind(&globals, &qh).map_err(|_| "wl_shm not supported")?;
@@ -389,35 +574,57 @@ fn main() -> Result<(), Box<dyn Error>> {
         exit: false,
         unlocking: false,
         screenshot_manager: ScreenshotManager::new(&globals, &qh).ok(),
+        grace_until: None,
     };
 
     event_queue.blocking_dispatch(&mut state)?;
 
-    let _wayland_source = WaylandSource::new(conn.clone(), event_queue).insert(event_loop.handle())?;
+    let _wayland_source =
+        WaylandSource::new(conn.clone(), event_queue).insert(event_loop.handle())?;
 
-    event_loop.handle().insert_source(auth_feedback_rx_actual, |event, _, state| {
-        if let calloop::channel::Event::Msg(success) = event {
-            state.handle_auth_result(success);
-            if success { state.lock_surfaces.clear(); }
-        }
-    })?;
+    event_loop
+        .handle()
+        .insert_source(auth_feedback_rx_actual, |event, _, state| {
+            if let calloop::channel::Event::Msg(success) = event {
+                state.handle_auth_result(success);
+                if success {
+                    state.lock_surfaces.clear();
+                }
+            }
+        })?;
 
     let timer = calloop::timer::Timer::from_duration(Duration::from_millis(16));
     event_loop.handle().insert_source(timer, |_, _, state| {
+        // Clear expired grace period
+        if let Some(grace_until) = state.grace_until {
+            if Instant::now() >= grace_until {
+                state.grace_until = None;
+            }
+        }
+
         if state.unlocking {
             return calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(100));
         }
 
         if let Ok(mut lm) = state.lock_manager.lock() {
             lm.update();
-            for surface in &mut lm.surfaces { let _ = surface.commit(&mut state.pool); }
+            for surface in &mut lm.surfaces {
+                let _ = surface.commit(&mut state.pool);
+            }
+        }
+        if state.exit {
+            calloop::timer::TimeoutAction::Drop
+        } else {
+            calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(16))
         }
-        if state.exit { calloop::timer::TimeoutAction::Drop } else { calloop::timer::TimeoutAction::ToDuration(Duration::from_millis(16)) }
     })?;
 
     if state.config.screenshots {
         state.outputs = state.output_state.outputs().collect();
-        log::info!("Capturing screenshots for {} outputs...", state.outputs.len());
+        log::info!(
+            "Capturing screenshots for {} outputs...",
+            state.outputs.len()
+        );
         state.screenshot_frames = vec![None; state.outputs.len()];
         state.captured_backgrounds = vec![None; state.outputs.len()];
         state.pending_screenshots = state.outputs.len();
@@ -427,11 +634,16 @@ fn main() -> Result<(), Box<dyn Error>> {
                 let data = CaptureData::new(i);
                 match mgr.capture_output(output, &qh, data) {
                     Ok(frame) => state.screenshot_frames[i] = Some(frame),
-                    Err(e) => { log::error!("Screencopy failed for output {}: {:?}", i, e); state.pending_screenshots -= 1; }
+                    Err(e) => {
+                        log::error!("Screencopy failed for output {}: {:?}", i, e);
+                        state.pending_screenshots -= 1;
+                    }
                 }
-            } else { state.pending_screenshots -= 1; }
+            } else {
+                state.pending_screenshots -= 1;
+            }
         }
-        
+
         let start = std::time::Instant::now();
         while state.pending_screenshots > 0 && start.elapsed() < Duration::from_secs(2) {
             event_loop.dispatch(Duration::from_millis(50), &mut state)?;
@@ -440,7 +652,10 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
 
     log::info!("Attempting to lock Wayland session...");
-    let session_lock = state.session_lock_state.lock(&qh).map_err(|e| { log::error!("Lock failed: {}", e); e })?;
+    let session_lock = state.session_lock_state.lock(&qh).map_err(|e| {
+        log::error!("Lock failed: {}", e);
+        e
+    })?;
 
     let outputs_to_lock: Vec<WlOutput> = state.output_state.outputs().collect();
     for output in outputs_to_lock {
@@ -448,11 +663,15 @@ fn main() -> Result<(), Box<dyn Error>> {
         let (width, height) = state.get_output_dimensions(&output);
         let lock_surface = session_lock.create_lock_surface(surface.clone(), &output, &qh);
         state.lock_surfaces.push(lock_surface);
-        if !state.outputs.contains(&output) { state.outputs.push(output.clone()); }
+        if !state.outputs.contains(&output) {
+            state.outputs.push(output.clone());
+        }
         if let Ok(mut lm) = state.lock_manager.lock() {
             lm.add_surface(width, height, output.clone());
             let count = lm.surface_count();
-            if let Some(ls) = lm.get_surface_mut(count - 1) { ls.set_wayland_surface(surface); }
+            if let Some(ls) = lm.get_surface_mut(count - 1) {
+                ls.set_wayland_surface(surface);
+            }
         }
     }
     state.session_lock = Some(session_lock);
@@ -461,6 +680,6 @@ fn main() -> Result<(), Box<dyn Error>> {
         event_loop.dispatch(Duration::from_millis(16), &mut state)?;
     }
 
-    log::info!("Exiting wayrustlock");
+    log::info!("Exiting rustlock");
     Ok(())
 }