diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 2d2a4200..7507ae9a 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -13,10 +13,7 @@ mod screencopy; use std::{ cmp, fs::File, - os::{ - fd::{AsFd, AsRawFd}, - unix::prelude::FromRawFd, - }, + os::fd::AsFd, process::exit, sync::atomic::{AtomicBool, Ordering}, }; @@ -175,17 +172,132 @@ impl WayshotConnection { fd: T, capture_region: Option, ) -> Result { - self.capture_output_frame_shm_fd_inner(cursor_overlay, output, fd, None, capture_region) + // Connecting to wayland environment. + let mut state = CaptureFrameState { + formats: Vec::new(), + state: None, + buffer_done: AtomicBool::new(false), + }; + let mut event_queue = self.conn.new_event_queue::(); + let qh = event_queue.handle(); + + // Instantiating screencopy manager. + let screencopy_manager = match self.globals.bind::( + &qh, + 3..=3, + (), + ) { + Ok(x) => x, + Err(e) => { + log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"); + log::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "ZwlrScreencopy Manager not found".to_string(), + )); + } + }; + + // Capture output. + let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { + screencopy_manager.capture_output_region( + cursor_overlay, + output, + region.x_coordinate, + region.y_coordinate, + region.width, + region.height, + &qh, + (), + ) + } else { + screencopy_manager.capture_output(cursor_overlay, output, &qh, ()) + }; + + // Empty internal event buffer until buffer_done is set to true which is when the Buffer done + // event is fired, aka the capture from the compositor is succesful. + while !state.buffer_done.load(Ordering::SeqCst) { + event_queue.blocking_dispatch(&mut state)?; + } + + log::debug!( + "Received compositor frame buffer formats: {:#?}", + state.formats + ); + // Filter advertised wl_shm formats and select the first one that matches. + let frame_format = state + .formats + .iter() + .find(|frame| { + matches!( + frame.format, + wl_shm::Format::Xbgr2101010 + | wl_shm::Format::Abgr2101010 + | wl_shm::Format::Argb8888 + | wl_shm::Format::Xrgb8888 + | wl_shm::Format::Xbgr8888 + ) + }) + .copied(); + log::debug!("Selected frame buffer format: {:#?}", frame_format); + + // Check if frame format exists. + let frame_format = match frame_format { + Some(format) => format, + None => { + log::error!("No suitable frame format found"); + return Err(Error::NoSupportedBufferFormat); + } + }; + + // Bytes of data in the frame = stride * height. + let frame_bytes = frame_format.stride * frame_format.height; + + // Create an in memory file and return it's file descriptor. + + // Instantiate shm global. + let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); + let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ()); + let buffer = shm_pool.create_buffer( + 0, + frame_format.width as i32, + frame_format.height as i32, + frame_format.stride as i32, + frame_format.format, + &qh, + (), + ); + + // Copy the pixel data advertised by the compositor into the buffer we just created. + frame.copy(&buffer); + // On copy the Ready / Failed events are fired by the frame object, so here we check for them. + loop { + // Basically reads, if frame state is not None then... + if let Some(state) = state.state { + match state { + FrameState::Failed => { + log::error!("Frame copy failed"); + return Err(Error::FramecopyFailed); + } + FrameState::Finished => { + buffer.destroy(); + shm_pool.destroy(); + return Ok(frame_format); + } + } + } + + event_queue.blocking_dispatch(&mut state)?; + } } - fn capture_output_frame_shm_fd_inner( + fn capture_output_frame_shm_from_file( &self, cursor_overlay: i32, output: &WlOutput, - fd: T, - file: Option<&File>, + file: &File, capture_region: Option, ) -> Result { + let fd = file.as_fd(); // Connecting to wayland environment. let mut state = CaptureFrameState { formats: Vec::new(), @@ -265,9 +377,7 @@ impl WayshotConnection { // Bytes of data in the frame = stride * height. let frame_bytes = frame_format.stride * frame_format.height; - if let Some(file) = file { - file.set_len(frame_bytes as u64)?; - } + file.set_len(frame_bytes as u64)?; // Create an in memory file and return it's file descriptor. // Instantiate shm global. @@ -317,13 +427,12 @@ impl WayshotConnection { // Create an in memory file and return it's file descriptor. let fd = create_shm_fd()?; // Create a writeable memory map backed by a mem_file. - let mem_file = unsafe { File::from_raw_fd(fd.as_raw_fd()) }; + let mem_file = File::from(fd); - let frame_format = self.capture_output_frame_shm_fd_inner( + let frame_format = self.capture_output_frame_shm_from_file( cursor_overlay, output, - fd, - Some(&mem_file), + &mem_file, capture_region, )?;