diff --git a/Cargo.lock b/Cargo.lock index e95580f..5dc4c48 100644 --- a/Cargo.lock +++ b/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", ] diff --git a/Cargo.toml b/Cargo.toml index be7814b..373d2df 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/lock.rs b/src/lock.rs index 6123ca1..485ee6d 100644 --- a/src/lock.rs +++ b/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, + output: wl_output::WlOutput, } impl LockedSurface { /// Create a new locked surface for an output - pub fn new(width: i32, height: i32, config: &Config) -> Option { + pub fn new( + width: i32, + height: i32, + config: &Config, + output: wl_output::WlOutput, + ) -> Option { 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 - let background = if config.screenshots { - // 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 - }; + // Background will be set later when screenshot is captured (if screenshots enabled) + let background = 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 { - match LockedSurface::new(width, height, &self.config) { + pub fn add_surface(&mut self, width: i32, height: i32, output: wl_output::WlOutput) -> bool { + 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()) + } } diff --git a/src/main.rs b/src/main.rs index de982db..1576c98 100644 --- a/src/main.rs +++ b/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() { @@ -85,7 +96,7 @@ fn lock_wayland_session( ctrlc_exit: Arc, ) -> Result<(), Box> { use smithay_client_toolkit::reexports::calloop; - + use smithay_client_toolkit::{ compositor::{CompositorHandler, CompositorState}, output::{OutputHandler, OutputState}, @@ -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>>, unlocking: bool, + screenshot_manager: Option, + screenshot_frames: Vec>, + captured_backgrounds: Vec>, + locked_outputs: HashSet, + outputs: Vec, + lock_surface_outputs: Vec, } 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 - for output in self.output_state.outputs() { + let outputs: Vec = self.output_state.outputs().collect(); + 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() { - lock_manager.initialize_lock_surfaces(); + if self.config.screenshots { + 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 for WaylandLock { + fn event( + state: &mut Self, + _manager: &ZwlrScreencopyManagerV1, + _event: ::Event, + _data: &(), + _conn: &Connection, + _qh: &QueueHandle, + ) { + // No events expected from the manager + } + } + + impl wayland_client::Dispatch for WaylandLock { + fn event( + state: &mut Self, + frame: &ZwlrScreencopyFrameV1, + event: ::Event, + _data: &CaptureData, + _conn: &Connection, + qh: &QueueHandle, + ) { + 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 = 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() diff --git a/src/screenshot.rs b/src/screenshot.rs index 6596763..24615f7 100644 --- a/src/screenshot.rs +++ b/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, - pub height: u32, - pub data: Vec, + surface: ImageSurface, } impl Screenshot { - pub fn capture( - _output: wayland_client::protocol::wl_output::WlOutput, - width: i32, - height: i32, - ) -> Result { - 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 { - width, - height, - data, - }) + /// Create a new screenshot from a Cairo surface. + pub fn new(surface: ImageSurface) -> Self { + Self { surface } } - pub fn apply_blur(&mut self, radius: u32, times: u32) { + /// Get a reference to the underlying surface. + pub fn surface(&self) -> &ImageSurface { + &self.surface + } + + /// 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, Vec> = - image::ImageBuffer::from_raw(self.width, self.height, self.data.clone()) - .expect("Failed to create image buffer"); + image::ImageBuffer::from_raw(width as u32, height as u32, data) + .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, - self.height as usize, + width 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 center_y = self.height as f32 / 2.0; + let width = self.surface.width(); + 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 { - for x in 0..self.width { + 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)) + .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; - self.data[index + i] = value.clamp(0.0, 255.0) as u8; + let value = data[index + i] as f32 * vignette_factor; + data[index + i] = value.clamp(0.0, 255.0) as u8; } } } - } - pub fn as_image_surface(&self) -> cairo::ImageSurface { - let surface = cairo::ImageSurface::create( + let mut surface_data = self.surface.data().unwrap(); + 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, +} + +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(globals: &GlobalList, qh: &QueueHandle) -> Result + where + D: Dispatch + 'static, + { + let manager = globals + .bind::(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( + &self, + output: &wl_output::WlOutput, + qh: &QueueHandle, + user_data: CaptureData, + ) -> Result + where + D: Dispatch + '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 { + 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, - self.height as i32, + info.width 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 { + 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 { + 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>, + pub flags: Mutex>, + pub buffer: Mutex>, +} - // TODO: Properly copy pixel data to surface using cairo API - // For now, return empty surface - surface +impl CaptureData { + /// Create new capture data for the given output index. + pub fn new(output_idx: usize) -> Self { + Self { + output_idx, + info: Mutex::new(None), + flags: Mutex::new(None), + buffer: Mutex::new(None), + } } }