Implement working screenshot capture with wlr-screencopy

- Fixed brace mismatch and duplicate code in screenshot.rs
- Increased SHM pool size to 256MB for high-res displays
- Corrected wlr-screencopy protocol usage: send copy request in Buffer event
- Fixed background handling: set_background now updates LockedSurface.background
- Removed dummy background creation; screenshots set when ready
- Added format conversion for Xbgr8888 and Xrgb8888 to ARGB32
- Proper Y-inversion handling based on flags

The lock screen now displays captured screenshots as background.
Screenshots are taken immediately after lock, before UI is shown.
Matches swaylock-effects behavior: lock appears first, then screenshot
applies when ready (no blocking).

Fixes: black/red screen issues, screenshot not displaying.

Cargo.lock

@@ -645,6 +645,21 @@ dependencies = [
  "miniz_oxide",
 ]
 
+[[package]]
+name = "futures"
+version = "0.3.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
 [[package]]
 name = "futures-channel"
 version = "0.3.32"
@@ -652,6 +667,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d"
 dependencies = [
  "futures-core",
+ "futures-sink",
 ]
 
 [[package]]
@@ -671,6 +687,12 @@ dependencies = [
  "futures-util",
 ]
 
+[[package]]
+name = "futures-io"
+version = "0.3.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
+
 [[package]]
 name = "futures-macro"
 version = "0.3.32"
@@ -682,6 +704,12 @@ dependencies = [
  "syn 2.0.117",
 ]
 
+[[package]]
+name = "futures-sink"
+version = "0.3.32"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
+
 [[package]]
 name = "futures-task"
 version = "0.3.32"
@@ -694,9 +722,13 @@ version = "0.3.32"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
 dependencies = [
+ "futures-channel",
  "futures-core",
+ "futures-io",
  "futures-macro",
+ "futures-sink",
  "futures-task",
+ "memchr",
  "pin-project-lite",
  "slab",
 ]
@@ -2057,11 +2089,13 @@ dependencies = [
 name = "wayrustlock"
 version = "0.1.0"
 dependencies = [
+ "anyhow",
  "cairo-rs",
  "chrono",
  "clap",
  "env_logger",
  "fastblur",
+ "futures",
  "image",
  "log",
  "pam-client",
@@ -2072,6 +2106,7 @@ dependencies = [
  "users",
  "wayland-client",
  "wayland-protocols",
+ "wayland-protocols-wlr",
  "xkbcommon",
  "zeroize",
 ]

Cargo.toml

@@ -7,6 +7,9 @@ edition = "2021"
 smithay-client-toolkit = { version = "0.19", features = ["calloop"] }
 wayland-client = "0.31"
 wayland-protocols = { version = "0.32", features = ["client", "unstable"] }
+wayland-protocols-wlr = { version = "0.3", features = ["client"] }
+anyhow = "1.0"
+futures = "0.3"
 cairo-rs = { version = "0.20", features = ["png"] }
 image = "0.25"
 fastblur = "0.1"

src/lock.rs

@@ -1,7 +1,8 @@
 use cairo::ImageSurface;
 use std::error::Error;
 use std::time::Instant;
-use wayland_client::protocol::{wl_shm, wl_surface};
+use wayland_client::protocol::{wl_output, wl_shm, wl_surface};
+use wayland_client::Proxy;
 
 use crate::config::Config;
 use crate::input::{InputAction, InputHandler};
@@ -23,11 +24,17 @@ pub struct LockedSurface {
     temp_screenshot_shown: bool,
     last_update: Instant,
     wayland_surface: Option<wl_surface::WlSurface>,
+    output: wl_output::WlOutput,
 }
 
 impl LockedSurface {
     /// Create a new locked surface for an output
-    pub fn new(width: i32, height: i32, config: &Config) -> Option<Self> {
+    pub fn new(
+        width: i32,
+        height: i32,
+        config: &Config,
+        output: wl_output::WlOutput,
+    ) -> Option<Self> {
         if width <= 0 || height <= 0 {
             return None;
         }
@@ -35,38 +42,8 @@ impl LockedSurface {
         let renderer = Renderer::new(width, height, config.clone());
         let input_handler = InputHandler::new(config.clone());
 
-        // Create background if screenshots are enabled
+        // Background will be set later when screenshot is captured (if screenshots enabled)
-        let background = if config.screenshots {
+        let background = None;
-            // For now, create a dummy screenshot with the output dimensions
-            // In a real implementation, this would capture actual screenshots via Wayland
-            let mut screenshot = Screenshot {
-                width: width as u32,
-                height: height as u32,
-                data: vec![0u8; (width * height * 4) as usize],
-            };
-
-            // Fill with a dark gray color (similar to swaylock default)
-            for i in 0..(screenshot.width * screenshot.height) as usize {
-                let offset = i * 4;
-                screenshot.data[offset] = 40; // R
-                screenshot.data[offset + 1] = 44; // G
-                screenshot.data[offset + 2] = 52; // B
-                screenshot.data[offset + 3] = 255; // A
-            }
-
-            // Apply effects if configured
-            if let Some((blur_radius, blur_times)) = config.effect_blur {
-                screenshot.apply_blur(blur_radius, blur_times);
-            }
-
-            if let Some((vignette_base, vignette_factor)) = config.effect_vignette {
-                screenshot.apply_vignette(vignette_base, vignette_factor);
-            }
-
-            Some(screenshot.as_image_surface())
-        } else {
-            None
-        };
 
         Some(Self {
             width,
@@ -81,6 +58,7 @@ impl LockedSurface {
             temp_screenshot_shown: false,
             last_update: Instant::now(),
             wayland_surface: None,
+            output,
         })
     }
 
@@ -334,10 +312,20 @@ impl LockedSurface {
         self.wayland_surface.as_ref()
     }
 
+    /// Get the output associated with this locked surface
+    pub fn output(&self) -> &wl_output::WlOutput {
+        &self.output
+    }
+
     /// Check if this surface has a Wayland surface attached
     pub fn has_wayland_surface(&self) -> bool {
         self.wayland_surface.is_some()
     }
+
+    /// Set the background image for this locked surface
+    pub fn set_background(&mut self, surface: ImageSurface) {
+        self.background = Some(surface);
+    }
 }
 
 /// Manager for all locked surfaces (multiple outputs)
@@ -358,8 +346,8 @@ impl LockManager {
     }
 
     /// Add a locked surface for an output
-    pub fn add_surface(&mut self, width: i32, height: i32) -> bool {
+    pub fn add_surface(&mut self, width: i32, height: i32, output: wl_output::WlOutput) -> bool {
-        match LockedSurface::new(width, height, &self.config) {
+        match LockedSurface::new(width, height, &self.config, output) {
             Some(surface) => {
                 self.surfaces.push(surface);
                 true
@@ -423,16 +411,6 @@ impl LockManager {
         self.surfaces.get_mut(index)
     }
 
-    /// Initialize lock surfaces for all outputs (called after session is locked)
-    pub fn initialize_lock_surfaces(&mut self) {
-        // In a real implementation, this would create Wayland surfaces for each output
-        // For now, we'll create dummy surfaces with default dimensions
-        if self.surfaces.is_empty() {
-            // Add a default surface (single monitor)
-            self.add_surface(1920, 1080);
-        }
-    }
-
     /// Toggle temp screenshot peek mode
     pub fn toggle_peek(&mut self) {
         for surface in &mut self.surfaces {
@@ -449,4 +427,14 @@ impl LockManager {
             .iter_mut()
             .find(|surface| surface.matches_surface(wayland_surface))
     }
+
+    /// Find a locked surface by output
+    pub fn find_surface_by_output(
+        &mut self,
+        output: &wl_output::WlOutput,
+    ) -> Option<&mut LockedSurface> {
+        self.surfaces
+            .iter_mut()
+            .find(|surface| surface.output().id() == output.id())
+    }
 }

src/main.rs

@@ -9,10 +9,21 @@ mod util;
 
 use config::Config;
 use lock::LockManager;
+use screenshot::{CaptureData, Screenshot, ScreenshotManager};
+use std::collections::HashSet;
 use std::error::Error;
 use std::fs::OpenOptions;
 use std::io::Write;
 use std::sync::{Arc, Mutex};
+use wayland_client::globals::GlobalList;
+use wayland_client::protocol::wl_output::WlOutput;
+use wayland_client::protocol::wl_surface::WlSurface;
+use wayland_client::Proxy;
+use wayland_client::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;
 
 fn setup_file_logging() {
@@ -119,6 +130,7 @@ fn lock_wayland_session(
         compositor_state: CompositorState,
         output_state: OutputState,
         registry_state: RegistryState,
+        globals: GlobalList,
         session_lock_state: SessionLockState,
         seat_state: SeatState,
         shm_state: Shm,
@@ -131,6 +143,12 @@ fn lock_wayland_session(
         exit: bool,
         auth_tx: Option<channel::Sender<Zeroizing<String>>>,
         unlocking: bool,
+        screenshot_manager: Option<ScreenshotManager>,
+        screenshot_frames: Vec<Option<ZwlrScreencopyFrameV1>>,
+        captured_backgrounds: Vec<Option<cairo::ImageSurface>>,
+        locked_outputs: HashSet<WlOutput>,
+        outputs: Vec<WlOutput>,
+        lock_surface_outputs: Vec<WlOutput>,
     }
 
     impl SessionLockHandler for WaylandLock {
@@ -146,15 +164,22 @@ fn lock_wayland_session(
             log_to_file("Session LOCKED - creating surfaces");
             eprintln!("SESSION LOCKED - creating lock surfaces");
 
-            // Take ownership of session_lock
             let lock_ref = &mut session_lock;
 
-            // Create lock surfaces for all outputs
+            let outputs: Vec<WlOutput> = self.output_state.outputs().collect();
-            for output in self.output_state.outputs() {
+            let output_count = outputs.len();
+
+            self.outputs = outputs.clone();
+            self.screenshot_frames = vec![None; output_count];
+            self.captured_backgrounds = vec![None; output_count];
+
+            for output in &outputs {
                 log_to_file(&format!("Creating lock surface for output"));
                 let surface = self.compositor_state.create_surface(qh);
-                let lock_surface = lock_ref.create_lock_surface(surface, &output, qh);
+                let lock_surface = lock_ref.create_lock_surface(surface, output, qh);
                 self.lock_surfaces.push(lock_surface);
+                self.lock_surface_outputs.push(output.clone());
+                self.locked_outputs.insert(output.clone());
             }
 
             log_to_file(&format!(
@@ -164,8 +189,35 @@ fn lock_wayland_session(
 
             self.session_lock = Some(session_lock);
 
-            if let Ok(mut lock_manager) = self.lock_manager.lock() {
+            if self.config.screenshots {
-                lock_manager.initialize_lock_surfaces();
+                log::info!("Initializing screenshot capture...");
+                match ScreenshotManager::new(&self.globals, qh) {
+                    Ok(manager) => {
+                        self.screenshot_manager = Some(manager);
+                        log::info!("Screenshot manager initialized");
+
+                        for (idx, output) in self.outputs.iter().enumerate() {
+                            log::info!("Starting screenshot capture for output {}...", idx);
+                            match self.screenshot_manager.as_mut().unwrap().capture_output(
+                                output,
+                                qh,
+                                CaptureData::new(idx),
+                            ) {
+                                Ok(frame) => {
+                                    self.screenshot_frames[idx] = Some(frame);
+                                    log::debug!("Screenshot capture initiated for output {}", idx);
+                                }
+                                Err(e) => {
+                                    log::error!("Failed to capture output {}: {}", idx, e);
+                                }
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        log::warn!("Screenshot capture not available: {}", e);
+                        self.screenshot_manager = None;
+                    }
+                }
             }
         }
 
@@ -204,7 +256,29 @@ fn lock_wayland_session(
             eprintln!("CONFIGURE: {}x{}", width, height);
 
             let surface_added = if let Ok(mut lock_manager) = self.lock_manager.lock() {
-                lock_manager.add_surface(width as i32, height as i32)
+                let surface_wl_id = session_lock_surface.wl_surface().id();
+                let output_idx = self
+                    .lock_surfaces
+                    .iter()
+                    .position(|s| s.wl_surface().id() == surface_wl_id);
+
+                if let Some(idx) = output_idx {
+                    if idx < self.lock_surface_outputs.len() {
+                        let output = &self.lock_surface_outputs[idx];
+                        let (width, height) = configure.new_size;
+                        lock_manager.add_surface(width as i32, height as i32, output.clone())
+                    } else {
+                        log::error!(
+                            "lock_surface_outputs index out of bounds: {} vs {}",
+                            idx,
+                            self.lock_surface_outputs.len()
+                        );
+                        false
+                    }
+                } else {
+                    log::error!("Could not find lock surface in self.lock_surfaces");
+                    false
+                }
             } else {
                 false
             };
@@ -216,7 +290,6 @@ fn lock_wayland_session(
             }
 
             log::debug!("Surface added to lock manager");
-            self.lock_surfaces.push(session_lock_surface.clone());
 
             if let Ok(mut lock_manager) = self.lock_manager.lock() {
                 let surface_count = lock_manager.surface_count();
@@ -229,7 +302,36 @@ fn lock_wayland_session(
                 if let Some(locked_surface) = lock_manager.get_surface_mut(surface_count - 1) {
                     locked_surface.set_wayland_surface(session_lock_surface.wl_surface().clone());
 
-                    // Render the surface first
+                    if self.config.screenshots {
+                        let output = locked_surface.output();
+                        let output_idx = self.outputs.iter().position(|o| o.id() == output.id());
+
+                        if let Some(idx) = output_idx {
+                            if let Some(background) = self
+                                .captured_backgrounds
+                                .get_mut(idx)
+                                .and_then(|b| b.take())
+                            {
+                                log::info!(
+                                    "Using captured screenshot background for output {}",
+                                    idx
+                                );
+                                locked_surface.set_background(background);
+                            } else {
+                                log::warn!("No screenshot available for output {}, will render without background", idx);
+                            }
+                        } else {
+                            log::warn!(
+                                "Output not found in self.outputs, cannot assign screenshot"
+                            );
+                        }
+                    }
+                }
+
+                if let Some(locked_surface) = lock_manager.get_surface_mut(surface_count - 1) {
+                    locked_surface.set_wayland_surface(session_lock_surface.wl_surface().clone());
+
+                    // Render the surface
                     locked_surface.renderer.render();
 
                     match locked_surface.renderer.get_pixel_data() {
@@ -364,9 +466,15 @@ fn lock_wayland_session(
             output: wl_output::WlOutput,
         ) {
             log::info!("New output detected: creating lock surface");
+            if self.locked_outputs.contains(&output) {
+                log::debug!("Output already has lock surface, ignoring");
+                return;
+            }
             if let Some(ref session_lock) = self.session_lock {
                 let surface = self.compositor_state.create_surface(qh);
                 let _lock_surface = session_lock.create_lock_surface(surface, &output, qh);
+                self.locked_outputs.insert(output.clone());
+                self.lock_surface_outputs.push(output.clone());
             }
         }
 
@@ -634,6 +742,184 @@ fn lock_wayland_session(
     smithay_client_toolkit::delegate_keyboard!(WaylandLock);
     smithay_client_toolkit::delegate_pointer!(WaylandLock);
 
+    impl wayland_client::Dispatch<ZwlrScreencopyManagerV1, ()> for WaylandLock {
+        fn event(
+            state: &mut Self,
+            _manager: &ZwlrScreencopyManagerV1,
+            _event: <ZwlrScreencopyManagerV1 as wayland_client::Proxy>::Event,
+            _data: &(),
+            _conn: &Connection,
+            _qh: &QueueHandle<Self>,
+        ) {
+            // No events expected from the manager
+        }
+    }
+
+    impl wayland_client::Dispatch<ZwlrScreencopyFrameV1, CaptureData> for WaylandLock {
+        fn event(
+            state: &mut Self,
+            frame: &ZwlrScreencopyFrameV1,
+            event: <ZwlrScreencopyFrameV1 as wayland_client::Proxy>::Event,
+            _data: &CaptureData,
+            _conn: &Connection,
+            qh: &QueueHandle<Self>,
+        ) {
+            let output_idx = _data.output_idx;
+            log::debug!("Screencopy event for output {}: {:?}", output_idx, event);
+
+            match event {
+                wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event::Buffer {
+                    format,
+                    width,
+                    height,
+                    stride,
+                } => {
+                    log::debug!("Buffer event: {}x{}, stride={}, format={:?}", width, height, stride, format);
+
+                    // Extract format value from the compositor's Buffer event
+                    let format_value = match format {
+                        wayland_client::WEnum::Value(f) => f,
+                        wayland_client::WEnum::Unknown(_) => {
+                            log::error!("Unknown format value");
+                            return;
+                        }
+                    };
+
+                    // Store buffer info including the actual format
+                    if let Ok(mut info_guard) = _data.info.lock() {
+                        *info_guard = Some(screenshot::BufferInfo {
+                            width,
+                            height,
+                            stride,
+                            format: format_value,
+                        });
+                    }
+
+                    // Create SHM buffer with the EXACT format the compositor specified
+                    if let Ok(mut buffer_guard) = _data.buffer.lock() {
+                        if let Ok((buf, _canvas)) = state.pool.create_buffer(
+                            i32::try_from(width).unwrap(),
+                            i32::try_from(height).unwrap(),
+                            i32::try_from(stride).unwrap(),
+                            format_value,
+                        ) {
+                            // Store the buffer for later use in Ready event
+                            *buffer_guard = Some(buf);
+                            log::debug!("Created SHM buffer for screencopy with format {:?}", format_value);
+                            // Immediately send copy request (protocol requires this after buffer creation)
+                            if let Some(buffer) = buffer_guard.as_ref() {
+                                frame.copy(buffer.wl_buffer());
+                                log::debug!("Sent copy request for screenshot");
+                            }
+                        } else {
+                            log::error!("Failed to create SHM buffer");
+                        }
+                    }
+                }
+                wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event::Flags { flags } => {
+                    log::debug!("Flags: {:?}", flags);
+                    if let Ok(mut flags_guard) = _data.flags.lock() {
+                        let flags_value = match flags {
+                            wayland_client::WEnum::Value(f) => f,
+                            wayland_client::WEnum::Unknown(_) => {
+                                log::warn!("Unknown flags value");
+                                return;
+                            }
+                        };
+                        *flags_guard = Some(flags_value);
+                    }
+                }
+                wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event::Ready { 
+                    tv_sec_hi: _, 
+                    tv_sec_lo: _, 
+                    tv_nsec: _ 
+                } => {
+                    log::info!("Ready event: processing screenshot for output {}", output_idx);
+
+                    // Extract all data from CaptureData
+                    let info_opt = _data.info.lock().ok().and_then(|i| i.clone());
+                    let buffer_opt = _data.buffer.lock().ok().and_then(|mut b| b.take());
+                    let flags_opt = _data.flags.lock().ok().and_then(|f| f.clone());
+
+                    if let (Some(info), Some(buffer), Some(flags)) = (info_opt, buffer_opt, flags_opt) {
+                        // Determine if Y-inversion is needed (Y_INVERT flag is bit 0)
+                        let y_invert = flags.bits() & 1 != 0;
+
+                        let handle = screenshot::ScreencopyBufferHandle {
+                            buffer,
+                            info: screenshot::BufferInfo {
+                                width: info.width,
+                                height: info.height,
+                                stride: info.stride,
+                                format: info.format,
+                            },
+                            y_invert,
+                        };
+
+                        match state.screenshot_manager.as_ref().unwrap().buffer_to_surface(handle, &mut state.pool) {
+                            Ok(surface) => {
+                                // Apply effects if configured
+                                let mut screenshot = Screenshot::new(surface);
+                                if let Err(e) = screenshot.apply_effects(&state.config) {
+                                    log::error!("Failed to apply effects: {}", e);
+                                }
+                                let surface = screenshot.into_inner();
+
+                                // Save screenshot to file for debugging if debug mode enabled
+                                if state.config.debug {
+                                    let path = format!("/tmp/wayrustlock_output{}.png", output_idx);
+                                            match std::fs::File::create(&path) {
+                                                Ok(mut file) => {
+                                                    if let Err(e) = surface.write_to_png(&mut file) {
+                                                        log::error!("Failed to write PNG to {}: {}", path, e);
+                                                    } else {
+                                                        log::info!("Saved screenshot to {}", path);
+                                                    }
+                                                }
+                                                Err(e) => {
+                                                    log::error!("Failed to create file {}: {}", path, e);
+                                                }
+                                            }
+                                }
+
+                                if output_idx < state.outputs.len() {
+                                    let output = &state.outputs[output_idx];
+                                    if let Ok(mut lock_manager) = state.lock_manager.lock() {
+                                        if let Some(locked_surface) = lock_manager.find_surface_by_output(output) {
+                                            log::debug!("Setting background directly on lock surface for output {}", output_idx);
+                                            locked_surface.set_background(surface);
+                                        } else {
+                                            log::debug!("Storing background in captured_backgrounds for output {}", output_idx);
+                                            state.captured_backgrounds[output_idx] = Some(surface);
+                                        }
+                                    } else {
+                                        log::debug!("Failed to lock manager, storing background");
+                                        state.captured_backgrounds[output_idx] = Some(surface);
+                                    }
+                                } else {
+                                    log::error!("output_idx {} out of bounds (outputs len {})", output_idx, state.outputs.len());
+                                }
+                            }
+                            Err(e) => {
+                                log::error!("Failed to convert buffer to surface: {}", e);
+                            }
+                        }
+                    } else {
+                        log::error!("Missing buffer data for Ready event");
+                    }
+
+                    // Clean up frame
+                    frame.destroy();
+                }
+                wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::Event::Failed => {
+                    log::error!("Screenshot capture failed for output {}", output_idx);
+                    frame.destroy();
+                }
+                _ => {}
+            }
+        }
+    }
+
     let conn = Connection::connect_to_env()?;
     let (globals, event_queue) = registry_queue_init(&conn)?;
     let qh: QueueHandle<WaylandLock> = event_queue.handle();
@@ -645,26 +931,32 @@ fn lock_wayland_session(
     let mut state = WaylandLock {
         loop_handle: event_loop.handle(),
         conn: conn.clone(),
+        lock_manager,
+        config,
+        ctrlc_exit: ctrlc_exit.clone(),
+        auth_tx: Some(auth_tx),
         compositor_state: CompositorState::bind(&globals, &qh)?,
         output_state: OutputState::new(&globals, &qh),
         registry_state: RegistryState::new(&globals),
         session_lock_state: SessionLockState::new(&globals, &qh),
         seat_state: SeatState::new(&globals, &qh),
-
         shm_state: Shm::bind(&globals, &qh).map_err(|_| "wl_shm protocol not supported")?,
         pool: SlotPool::new(
-            1920 * 1080 * 4 * 3, // triple buffering for smooth rendering
+            256 * 1024 * 1024, // 256 MB pool to support high-resolution displays and screenshots
             &Shm::bind(&globals, &qh).map_err(|_| "wl_shm protocol not supported")?,
         )
         .map_err(|e| format!("Failed to create slot pool: {:?}", e))?,
+        globals,
         session_lock: None,
         lock_surfaces: Vec::new(),
-        lock_manager,
-        config,
-        ctrlc_exit: ctrlc_exit.clone(),
         exit: false,
-        auth_tx: Some(auth_tx),
         unlocking: false,
+        screenshot_manager: None,
+        screenshot_frames: Vec::new(),
+        captured_backgrounds: Vec::new(),
+        locked_outputs: HashSet::new(),
+        outputs: Vec::new(),
+        lock_surface_outputs: Vec::new(),
     };
 
     state.session_lock = Some(
@@ -747,7 +1039,6 @@ fn run_demonstration_mode(
 
     {
         let mut lock_manager = lock_manager.lock().unwrap();
-        lock_manager.initialize_lock_surfaces();
         log::info!(
             "Initialized {} lock surface(s)",
             lock_manager.surface_count()

src/screenshot.rs

@@ -1,55 +1,85 @@
+//!
+//! This module provides functionality to capture the current screen contents
+//! and apply visual effects like blur and vignette, similar to swaylock-effects.
+
+use anyhow::{Context, Result};
+use cairo::ImageSurface;
+use log::debug;
+use smithay_client_toolkit::shm::{slot::Buffer, slot::SlotPool};
+use std::sync::Mutex;
+use wayland_client::globals::GlobalList;
+use wayland_client::protocol::{wl_output, wl_shm};
+use wayland_client::{Dispatch, 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 crate::config::Config;
+
+/// A captured screenshot with optional visual effects applied.
 pub struct Screenshot {
-    pub width: u32,
+    surface: ImageSurface,
-    pub height: u32,
-    pub data: Vec<u8>,
 }
 
 impl Screenshot {
-    pub fn capture(
+    /// Create a new screenshot from a Cairo surface.
-        _output: wayland_client::protocol::wl_output::WlOutput,
+    pub fn new(surface: ImageSurface) -> Self {
-        width: i32,
+        Self { surface }
-        height: i32,
-    ) -> Result<Self, String> {
-        let width = width as u32;
-        let height = height as u32;
-        let size = (width * height * 4) as usize;
-        let mut data = vec![0u8; size];
-
-        for i in 0..(width * height) as usize {
-            let offset = i * 4;
-            data[offset] = 40;
-            data[offset + 1] = 44;
-            data[offset + 2] = 52;
-            data[offset + 3] = 255;
     }
 
-        Ok(Self {
+    /// Get a reference to the underlying surface.
-            width,
+    pub fn surface(&self) -> &ImageSurface {
-            height,
+        &self.surface
-            data,
-        })
     }
 
-    pub fn apply_blur(&mut self, radius: u32, times: u32) {
+    /// Consume the screenshot and return the underlying Cairo surface.
+    pub fn into_inner(self) -> ImageSurface {
+        self.surface
+    }
+
+    /// Apply configured visual effects to the screenshot.
+    pub fn apply_effects(&mut self, config: &Config) -> Result<()> {
+        if let Some((radius, times)) = config.effect_blur {
+            self.apply_blur(radius, times)?;
+        }
+        if let Some((base, factor)) = config.effect_vignette {
+            self.apply_vignette(base, factor);
+        }
+        Ok(())
+    }
+
+    /// Apply a Gaussian blur effect.
+    pub fn apply_blur(&mut self, radius: u32, times: u32) -> Result<()> {
         if radius == 0 || times == 0 {
-            return;
+            return Ok(());
         }
 
+        let width = self.surface.width();
+        let height = self.surface.height();
+        let stride = self.surface.stride() as usize;
+        let mut data = vec![0u8; stride * height as usize];
+
+        self.surface
+            .with_data(|src| data.copy_from_slice(src))
+            .context("Failed to get surface data")?;
+
+        // Convert to image::RgbaImage for processing
         let mut img: image::ImageBuffer<image::Rgba<u8>, Vec<u8>> =
-            image::ImageBuffer::from_raw(self.width, self.height, self.data.clone())
+            image::ImageBuffer::from_raw(width as u32, height as u32, data)
-                .expect("Failed to create image buffer");
+                .context("Failed to create image buffer")?;
 
         for _ in 0..times {
             let mut rgb_data: Vec<[u8; 3]> =
-                Vec::with_capacity((self.width * self.height) as usize);
+                Vec::with_capacity((width as usize) * (height as usize));
             for pixel in img.pixels() {
                 rgb_data.push([pixel[0], pixel[1], pixel[2]]);
             }
 
             fastblur::gaussian_blur(
                 &mut rgb_data,
-                self.width as usize,
+                width as usize,
-                self.height as usize,
+                height as usize,
                 radius as f32,
             );
 
@@ -60,40 +90,240 @@ impl Screenshot {
             }
         }
 
-        self.data = img.into_raw();
+        // Copy back to surface
+        let new_data = img.into_raw();
+        let mut surface_data = self.surface.data()?;
+        surface_data.copy_from_slice(&new_data);
+        Ok(())
     }
 
+    /// Apply a vignette effect (darken edges).
     pub fn apply_vignette(&mut self, base: f32, factor: f32) {
-        let center_x = self.width as f32 / 2.0;
+        let width = self.surface.width();
-        let center_y = self.height as f32 / 2.0;
+        let height = self.surface.height();
+        let center_x = width as f32 / 2.0;
+        let center_y = height as f32 / 2.0;
         let max_distance = (center_x * center_x + center_y * center_y).sqrt();
 
-        for y in 0..self.height {
+        let stride = self.surface.stride() as usize;
-            for x in 0..self.width {
+        let mut data = vec![0u8; stride * height as usize];
+        self.surface
+            .with_data(|src| data.copy_from_slice(src))
+            .unwrap();
+
+        for y in 0..height {
+            for x in 0..width {
                 let dx = x as f32 - center_x;
                 let dy = y as f32 - center_y;
                 let distance = (dx * dx + dy * dy).sqrt();
                 let vignette_factor = base + (1.0 - base) * (distance / max_distance).powf(factor);
 
-                let index = ((y * self.width + x) * 4) as usize;
+                let index = ((y * width + x) * 4) as usize;
                 for i in 0..3 {
-                    let value = self.data[index + i] as f32 * vignette_factor;
+                    let value = data[index + i] as f32 * vignette_factor;
-                    self.data[index + i] = value.clamp(0.0, 255.0) as u8;
+                    data[index + i] = value.clamp(0.0, 255.0) as u8;
-                }
                 }
             }
         }
 
-    pub fn as_image_surface(&self) -> cairo::ImageSurface {
+        let mut surface_data = self.surface.data().unwrap();
-        let surface = cairo::ImageSurface::create(
+        surface_data.copy_from_slice(&data);
+    }
+}
+
+#[derive(Clone)]
+/// Information about a buffer from the screencopy protocol.
+pub struct BufferInfo {
+    pub width: u32,
+    pub height: u32,
+    pub stride: u32,
+    pub format: wl_shm::Format,
+}
+
+/// Handle to a captured buffer that can be converted to a Cairo surface.
+pub struct ScreencopyBufferHandle {
+    pub buffer: Buffer,
+    pub info: BufferInfo,
+    pub y_invert: bool,
+}
+
+/// Manager for the wlr-screencopy protocol.
+pub struct ScreenshotManager {
+    manager: Option<ZwlrScreencopyManagerV1>,
+}
+
+impl ScreenshotManager {
+    /// Bind to the wlr-screencopy global and create a new manager.
+    ///
+    /// Returns `Ok(Self)` if the protocol is available, otherwise `Err`.
+    pub fn new<D>(globals: &GlobalList, qh: &QueueHandle<D>) -> Result<Self>
+    where
+        D: Dispatch<ZwlrScreencopyManagerV1, ()> + 'static,
+    {
+        let manager = globals
+            .bind::<ZwlrScreencopyManagerV1, _, _>(qh, 1..=3, ())
+            .ok();
+
+        if manager.is_none() {
+            debug!("zwlr_screencopy_manager_v1 not available");
+        }
+
+        Ok(Self { manager })
+    }
+
+    /// Returns `true` if the wlr-screencopy protocol is available.
+    pub fn is_available(&self) -> bool {
+        self.manager.is_some()
+    }
+
+    /// Initiate a screencopy operation for the given output.
+    ///
+    /// This method sends a screencopy request and returns the frame object.
+    /// The frame events will be dispatched to the provided queue's dispatcher
+    /// with the given user data.
+    pub fn capture_output<D>(
+        &self,
+        output: &wl_output::WlOutput,
+        qh: &QueueHandle<D>,
+        user_data: CaptureData,
+    ) -> Result<ZwlrScreencopyFrameV1>
+    where
+        D: Dispatch<ZwlrScreencopyFrameV1, CaptureData> + 'static,
+    {
+        let manager = self.manager.as_ref().context("Screencopy not available")?;
+        let frame = manager.capture_output(0, output, qh, user_data);
+        Ok(frame)
+    }
+
+    /// Convert a captured buffer to a Cairo ImageSurface.
+    pub fn buffer_to_surface(
+        &self,
+        handle: ScreencopyBufferHandle,
+        pool: &mut SlotPool,
+    ) -> Result<ImageSurface> {
+        let info = handle.info;
+        let y_invert = handle.y_invert;
+
+        let canvas = handle
+            .buffer
+            .canvas(pool)
+            .context("Failed to get buffer canvas")?;
+
+        let pixel_width = (info.width * 4) as usize;
+        let stride = info.stride as usize;
+        let height = info.height as usize;
+
+        if stride < pixel_width {
+            anyhow::bail!("Stride smaller than pixel width");
+        }
+
+        let raw_data = {
+            let mut data = vec![0u8; (info.width * info.height * 4) as usize];
+            for row in 0..height {
+                let src_offset = row * stride;
+                let dst_offset = row * pixel_width;
+                data[dst_offset..dst_offset + pixel_width]
+                    .copy_from_slice(&canvas[src_offset..src_offset + pixel_width]);
+            }
+            data
+        };
+
+        let converted_data = match info.format {
+            wayland_client::protocol::wl_shm::Format::Argb8888 => raw_data,
+            wayland_client::protocol::wl_shm::Format::Xbgr8888 => {
+                convert_xbgr8888_to_argb32(&raw_data, info.width as usize, info.height as usize)
+            }
+            wayland_client::protocol::wl_shm::Format::Xrgb8888 => {
+                convert_xrgb8888_to_argb32(&raw_data, info.width as usize, info.height as usize)
+            }
+            _ => {
+                log::warn!("Unsupported format {:?}, using raw data as-is", info.format);
+                raw_data
+            }
+        };
+
+        if y_invert {
+            let mut flipped = vec![0u8; (info.width * info.height * 4) as usize];
+            let src_stride = (info.width * 4) as usize;
+            for row in 0..height {
+                let src_row = height - 1 - row;
+                let src_offset = src_row * src_stride;
+                let dst_offset = row * src_stride;
+                flipped[dst_offset..dst_offset + src_stride]
+                    .copy_from_slice(&converted_data[src_offset..src_offset + src_stride]);
+            }
+            return Ok(ImageSurface::create_for_data(
+                flipped,
+                cairo::Format::ARgb32,
+                info.width as i32,
+                info.height as i32,
+                src_stride as i32,
+            )
+            .context("Failed to create flipped Cairo surface")?);
+        }
+
+        Ok(ImageSurface::create_for_data(
+            converted_data,
             cairo::Format::ARgb32,
-            self.width as i32,
+            info.width as i32,
-            self.height as i32,
+            info.height as i32,
+            pixel_width as i32,
         )
-        .expect("Failed to create image surface");
+        .context("Failed to create Cairo surface")?)
+    }
+}
+
+/// Convert Xbgr8888 buffer data to ARGB32 format.
+fn convert_xbgr8888_to_argb32(data: &[u8], width: usize, height: usize) -> Vec<u8> {
+    let mut result = vec![0u8; width * height * 4];
+    for y in 0..height {
+        for x in 0..width {
+            let src_idx = (y * width + x) * 4;
+            let dst_idx = (y * width + x) * 4;
+            result[dst_idx] = 255;
+            result[dst_idx + 1] = data[src_idx + 2];
+            result[dst_idx + 2] = data[src_idx + 1];
+            result[dst_idx + 3] = data[src_idx];
+        }
+    }
+    result
+}
+
+/// Convert Xrgb8888 buffer data to ARGB32 format.
+fn convert_xrgb8888_to_argb32(data: &[u8], width: usize, height: usize) -> Vec<u8> {
+    let mut result = vec![0u8; width * height * 4];
+    for y in 0..height {
+        for x in 0..width {
+            let src_idx = (y * width + x) * 4;
+            let dst_idx = (y * width + x) * 4;
+            result[dst_idx] = 255;
+            result[dst_idx + 1] = data[src_idx + 1];
+            result[dst_idx + 2] = data[src_idx + 2];
+            result[dst_idx + 3] = data[src_idx + 3];
+        }
+    }
+    result
+}
+
+/// User data associated with a screencopy frame request.
+///
+/// Stores intermediate data needed to assemble the final screenshot once
+/// all frame events are received.
+pub struct CaptureData {
+    pub output_idx: usize,
+    pub info: Mutex<Option<BufferInfo>>,
+    pub flags: Mutex<Option<Flags>>,
+    pub buffer: Mutex<Option<Buffer>>,
+}
 
-        // TODO: Properly copy pixel data to surface using cairo API
+impl CaptureData {
-        // For now, return empty surface
+    /// Create new capture data for the given output index.
-        surface
+    pub fn new(output_idx: usize) -> Self {
+        Self {
+            output_idx,
+            info: Mutex::new(None),
+            flags: Mutex::new(None),
+            buffer: Mutex::new(None),
+        }
     }
 }