chore: allow pass row fd to capture

libwayshot/src/error.rs

@@ -21,4 +21,10 @@ pub enum Error {
  Global(#[from] GlobalError),
  #[error("connect error: {0}")]
  Connect(#[from] ConnectError),
+ #[error("framecopy failed")]
+ FramecopyFailed,
+ #[error("No supported buffer format")]
+ NoSupportedBufferFormat,
+ #[error("Cannot find Wayland protocols")]
+ ProtocolNotFound(String),
 }

libwayshot/src/lib.rs

@@ -13,7 +13,10 @@ mod screencopy;
 use std::{
  cmp,
  fs::File,
- os::unix::prelude::FromRawFd,
+ os::{
+ fd::{AsFd, AsRawFd},
+ unix::prelude::FromRawFd,
+ },
  process::exit,
  sync::atomic::{AtomicBool, Ordering},
 };
@@ -45,6 +48,11 @@ use crate::{
 
 pub use crate::error::{Error, Result};
 
+pub mod reexport {
+ use wayland_client::protocol::wl_output;
+ pub use wl_output::{Transform, WlOutput};
+}
+
 type Frame = (Vec<FrameCopy>, Option<(i32, i32)>);
 
 /// Struct to store region capture details.
@@ -158,6 +166,150 @@ impl WayshotConnection {
  Ok(())
  }
 
+ /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
+ pub fn capture_output_frame_shm_fd<T: AsFd>(
+ &self,
+ cursor_overlay: i32,
+ output: &WlOutput,
+ transform: Transform,
+ fd: T,
+ capture_region: Option<CaptureRegion>,
+ ) -> Result<FrameCopy> {
+ // 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::<CaptureFrameState>();
+ let qh = event_queue.handle();
+
+ // Instantiating screencopy manager.
+ let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
+ &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.
+ let mem_file = unsafe { File::from_raw_fd(fd.as_fd().as_raw_fd()) };
+ mem_file.set_len(frame_bytes as u64)?;
+
+ // Instantiate shm global.
+ let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ()).unwrap();
+ let shm_pool = shm.create_pool(fd.as_fd().as_raw_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 => {
+ // Create a writeable memory map backed by a mem_file.
+ let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? };
+ let data = &mut *frame_mmap;
+ let frame_color_type = if let Some(converter) =
+ create_converter(frame_format.format)
+ {
+ converter.convert_inplace(data)
+ } else {
+ log::error!("Unsupported buffer format: {:?}", frame_format.format);
+ log::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot.");
+ return Err(Error::NoSupportedBufferFormat);
+ };
+ return Ok(FrameCopy {
+ frame_format,
+ frame_color_type,
+ frame_mmap,
+ transform,
+ });
+ }
+ }
+ }
+
+ event_queue.blocking_dispatch(&mut state)?;
+ }
+ }
  /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
  fn capture_output_frame(
  &self,