diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs index 384f8dc2..4142aef9 100644 --- a/libwayshot/src/error.rs +++ b/libwayshot/src/error.rs @@ -80,7 +80,7 @@ mod tests { #[test] fn test_from_io_error() { use std::io; - let io_error = io::Error::new(io::ErrorKind::Other, "test error"); + let io_error = io::Error::other("test error"); let wayshot_error: Error = io_error.into(); match wayshot_error { diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index f92a1db3..d402dc8a 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -102,6 +102,7 @@ pub struct WayshotConnection { toplevel_infos: Vec, dmabuf_state: Option, toplevel_capture_support: bool, + image_copy_support: bool, } pub enum WayshotFrame { @@ -109,6 +110,12 @@ pub enum WayshotFrame { ExtImageCopy(ExtImageCopyCaptureFrameV1), } +#[derive(Debug, Clone)] +pub enum FormatCheckTarget { + Screen(WlOutput), + Toplevel(TopLevel), +} + fn check_toplevel_protocols(globals: &GlobalList, conn: &Connection) -> Result<()> { let event_queue = conn.new_event_queue::(); let qh = event_queue.handle(); @@ -119,6 +126,15 @@ fn check_toplevel_protocols(globals: &GlobalList, conn: &Connection) -> Result<( Ok(()) } +fn check_ext_image_copy_protocols(globals: &GlobalList, conn: &Connection) -> Result<()> { + let event_queue = conn.new_event_queue::(); + let qh = event_queue.handle(); + + // Bind managers + globals.bind::(&qh, 1..=1, ())?; + Ok(()) +} + impl WayshotConnection { pub fn new() -> Result { let conn = Connection::connect_to_env()?; @@ -130,6 +146,7 @@ impl WayshotConnection { pub fn from_connection(conn: Connection) -> Result { let (globals, _) = registry_queue_init::(&conn)?; + let image_copy_support = check_ext_image_copy_protocols(&globals, &conn).is_ok(); let toplevel_capture_support = check_toplevel_protocols(&globals, &conn).is_ok(); let mut initial_state = Self { conn, @@ -138,6 +155,7 @@ impl WayshotConnection { toplevel_infos: Vec::new(), dmabuf_state: None, toplevel_capture_support, + image_copy_support, }; initial_state.refresh_outputs()?; @@ -158,6 +176,7 @@ impl WayshotConnection { let gpu = dispatch::Card::open(device_path); // init a GBM device let gbm = GBMDevice::new(gpu).unwrap(); + let image_copy_support = check_ext_image_copy_protocols(&globals, &conn).is_ok(); let toplevel_capture_support = check_toplevel_protocols(&globals, &conn).is_ok(); let mut initial_state = Self { conn, @@ -169,6 +188,7 @@ impl WayshotConnection { gbmdev: gbm, }), toplevel_capture_support, + image_copy_support, }; initial_state.refresh_outputs()?; @@ -181,6 +201,11 @@ impl WayshotConnection { self.toplevel_capture_support } + /// Check if the image copy is supported + pub fn image_copy_support(&self) -> bool { + self.image_copy_support + } + /// Fetch all accessible wayland outputs. pub fn get_all_outputs(&self) -> &[OutputInfo] { self.output_infos.as_slice() @@ -301,8 +326,21 @@ impl WayshotConnection { /// # Returns /// - A vector of [`FrameFormat`] if screen capture succeeds. /// - [`Error::ProtocolNotFound`] if wlr-screencopy protocol is not found. - pub fn get_available_frame_formats(&self, output: &WlOutput) -> Result> { - let (state, _event_queue, _frame) = self.capture_output_frame_get_state(0, output, None)?; + pub fn get_available_frame_formats( + &self, + target: &FormatCheckTarget, + ) -> Result> { + let state = match target { + FormatCheckTarget::Screen(output) => { + let (state, _, _) = self.capture_output_frame_get_state(0, output, None)?; + state + } + FormatCheckTarget::Toplevel(toplevel) => { + let (state, _, _) = self.capture_toplevel_frame_get_state(toplevel, false)?; + state + } + }; + Ok(state.formats) } @@ -335,7 +373,7 @@ impl WayshotConnection { .copied() { let frame_guard: FrameGuard = - self.capture_output_frame_inner(state, event_queue, frame, format, fd)?; + self.image_copy_frame_inner(state, event_queue, frame, format, fd)?; Ok(frame_guard) } else { Err(Error::NoSupportedBufferFormat) @@ -353,7 +391,7 @@ impl WayshotConnection { let (state, event_queue, frame, frame_format) = self.capture_output_frame_get_state_shm(cursor_overlay, output, capture_region)?; let frame_guard = - self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd)?; + self.image_copy_frame_inner(state, event_queue, frame, frame_format, fd)?; Ok((frame_format, frame_guard)) } @@ -371,7 +409,7 @@ impl WayshotConnection { file.set_len(frame_format.byte_size())?; let frame_guard = - self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; + self.image_copy_frame_inner(state, event_queue, frame, frame_format, file)?; Ok((frame_format, frame_guard)) } @@ -917,7 +955,7 @@ impl WayshotConnection { } } - fn capture_output_frame_inner( + fn image_copy_frame_inner( &self, state: CaptureFrameState, event_queue: EventQueue, @@ -927,15 +965,15 @@ impl WayshotConnection { ) -> Result { match frame { WayshotFrame::WlrScreenshot(frame) => { - self.capture_output_frame_inner_wlr(state, event_queue, frame, frame_format, fd) + self.wlr_screencopy_inner(state, event_queue, frame, frame_format, fd) } WayshotFrame::ExtImageCopy(frame) => { - self.capture_output_frame_inner_ext(state, event_queue, frame, frame_format, fd) + self.ext_image_copy_frame_inner(state, event_queue, frame, frame_format, fd) } } } - fn capture_output_frame_inner_wlr( + fn wlr_screencopy_inner( &self, mut state: CaptureFrameState, mut event_queue: EventQueue, @@ -980,7 +1018,11 @@ impl WayshotConnection { } FrameState::Finished => { tracing::trace!("Frame copy finished"); - return Ok(FrameGuard { buffer, shm_pool }); + return Ok(FrameGuard { + buffer, + shm_pool, + size: frame_format.size, + }); } } } @@ -988,7 +1030,7 @@ impl WayshotConnection { event_queue.blocking_dispatch(&mut state)?; } } - fn capture_output_frame_inner_ext( + fn ext_image_copy_frame_inner( &self, mut state: CaptureFrameState, mut event_queue: EventQueue, @@ -1034,7 +1076,11 @@ impl WayshotConnection { } FrameState::Finished => { tracing::trace!("Frame copy finished"); - return Ok(FrameGuard { buffer, shm_pool }); + return Ok(FrameGuard { + buffer, + shm_pool, + size: frame_format.size, + }); } } } @@ -1232,6 +1278,11 @@ impl WayshotConnection { .get_all_outputs() .iter() .filter_map(|output_info| { + // NOTE: ext-image-copy do not need to compute the location + // So we keep none here + if self.image_copy_support { + return None; + } tracing::span!( tracing::Level::DEBUG, "filter_map", @@ -1256,7 +1307,7 @@ impl WayshotConnection { }) .collect(), RegionCapturer::TopLevel(ref toplevel) => { - return self.capture_toplevel_using_ext_protocol(toplevel, cursor_overlay); + return self.capture_toplevel(toplevel, cursor_overlay); } RegionCapturer::Freeze(_) => self .get_all_outputs() @@ -1354,11 +1405,6 @@ impl WayshotConnection { }) } - /// Take a screenshot from the specified region. - /// Because ext-image-copy is very different from wlr-screencopy, and it caused bug here - /// So deprecate it now - /// Maybe we will fix it in the further - #[deprecated(since = "0.4.0", note = "use screenshot_freeze instead")] pub fn screenshot( &self, capture_region: LogicalRegion, @@ -1416,20 +1462,24 @@ impl WayshotConnection { ) } - fn capture_toplevel_using_ext_protocol( + fn capture_toplevel_frame_get_state( &self, toplevel: &TopLevel, cursor_overlay: bool, - ) -> Result { + ) -> Result<( + CaptureFrameState, + EventQueue, + ExtImageCopyCaptureFrameV1, + )> { // Create state and event queue similar to other ext-image flows - let state = CaptureFrameState { + let mut state = CaptureFrameState { formats: Vec::new(), dmabuf_formats: Vec::new(), state: None, buffer_done: AtomicBool::new(false), toplevels: Vec::new(), }; - let event_queue = self.conn.new_event_queue::(); + let mut event_queue = self.conn.new_event_queue::(); let qh = event_queue.handle(); // Bind managers @@ -1450,19 +1500,71 @@ impl WayshotConnection { }; let session = manager.create_session(&source, options, &qh, ()); let frame = session.create_frame(&qh, ()); + event_queue.blocking_dispatch(&mut state)?; - // Determine a suitable shm FrameFormat for this frame + Ok((state, event_queue, frame)) + } + + pub fn capture_toplevel_frame_shm_fd( + &self, + cursor_overlay: bool, + toplevel: &TopLevel, + fd: T, + ) -> Result<(FrameFormat, FrameGuard)> { + let (state, event_queue, frame, frame_format) = + self.capture_toplevel_frame_get_state_shm(toplevel, cursor_overlay)?; + let frame_guard = + self.ext_image_copy_frame_inner(state, event_queue, frame, frame_format, fd)?; + Ok((frame_format, frame_guard)) + } + + pub fn capture_toplevel_frame_shm_fd_with_format( + &self, + cursor_overlay: bool, + toplevel: &TopLevel, + fd: T, + frame_format: wl_shm::Format, + ) -> Result { + let (state, event_queue, frame) = + self.capture_toplevel_frame_get_state(toplevel, cursor_overlay)?; + if let Some(format) = state + .formats + .iter() + .find(|f| f.format == frame_format) + .copied() + { + let frame_guard: FrameGuard = + self.ext_image_copy_frame_inner(state, event_queue, frame, format, fd)?; + Ok(frame_guard) + } else { + Err(Error::NoSupportedBufferFormat) + } + } + + fn capture_toplevel_frame_shm_from_file( + &self, + cursor_overlay: bool, + toplevel: &TopLevel, + file: &File, + ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = - self.capture_output_frame_get_state_shm_for_toplevel(state, event_queue, frame)?; + self.capture_toplevel_frame_get_state_shm(toplevel, cursor_overlay)?; + + file.set_len(frame_format.byte_size())?; + + let frame_guard = + self.ext_image_copy_frame_inner(state, event_queue, frame, frame_format, file)?; + + Ok((frame_format, frame_guard)) + } + fn capture_toplevel(&self, toplevel: &TopLevel, cursor_overlay: bool) -> Result { // Back the buffer with a shm file of the required size let fd = create_shm_fd()?; let memfile = File::from(fd); - memfile.set_len(frame_format.byte_size())?; - - // Perform the copy using the existing ext-image path - let _guard = - self.capture_output_frame_inner_ext(state, event_queue, frame, frame_format, &memfile)?; + // Determine a suitable shm FrameFormat for this frame + let (frame_format, _) = + self.capture_toplevel_frame_shm_from_file(cursor_overlay, toplevel, &memfile)?; // Map and convert to image let frame_mmap = unsafe { MmapMut::map_mut(&memfile)? }; @@ -1484,11 +1586,10 @@ impl WayshotConnection { } // Helper method to get frame format for toplevel using ext-image session events - fn capture_output_frame_get_state_shm_for_toplevel( + fn capture_toplevel_frame_get_state_shm( &self, - mut state: CaptureFrameState, - mut event_queue: EventQueue, - frame: ExtImageCopyCaptureFrameV1, + toplevel: &TopLevel, + cursor_overlay: bool, ) -> Result<( CaptureFrameState, EventQueue, @@ -1496,8 +1597,8 @@ impl WayshotConnection { FrameFormat, )> { // Drain events at least once to populate formats from the session - event_queue.blocking_dispatch(&mut state)?; - + let (state, event_queue, frame) = + self.capture_toplevel_frame_get_state(toplevel, cursor_overlay)?; let frame_format = state .formats .iter() diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index e7505449..9cbe9c81 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -24,6 +24,7 @@ use crate::{ pub struct FrameGuard { pub buffer: WlBuffer, pub shm_pool: WlShmPool, + pub size: Size, } impl Drop for FrameGuard { diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs index 8f9866c5..7cf6ebb5 100644 --- a/wayshot/src/cli.rs +++ b/wayshot/src/cli.rs @@ -1,7 +1,7 @@ use std::path::PathBuf; use clap::{ - Parser, arg, + Parser, builder::{ Styles, styling::{AnsiColor, Effects},