diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 2e7ee7f0..1ad6842c 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -39,7 +39,7 @@ use wayland_protocols_wlr::screencopy::v1::client::{ use crate::{ output::OutputInfo, region::{LogicalRegion, Position, Size}, - screencopy::FrameFormat, + screencopy::{DMAFrameFormat, FrameFormat}, }; #[derive(Debug)] @@ -180,14 +180,14 @@ pub enum FrameState { pub struct CaptureFrameState { pub formats: Vec, - pub dmabuf_formats: Vec<(u32, u32, u32)>, + pub dmabuf_formats: Vec, pub state: Option, pub buffer_done: AtomicBool, } impl Dispatch for CaptureFrameState { fn event( - _state: &mut Self, + _frame: &mut Self, _proxy: &ZwpLinuxDmabufV1, _event: zwp_linux_dmabuf_v1::Event, _data: &(), @@ -201,7 +201,7 @@ impl Dispatch for CaptureFrameState { fn event( _state: &mut Self, _proxy: &ZwpLinuxBufferParamsV1, - event: zwp_linux_buffer_params_v1::Event, + _event: zwp_linux_buffer_params_v1::Event, _data: &(), _conn: &Connection, _qhandle: &QueueHandle, @@ -227,6 +227,7 @@ impl Dispatch for CaptureFrameState { stride, } => { if let Value(f) = format { + tracing::debug!("Received Buffer event with format: {f:?}"); frame.formats.push(FrameFormat { format: f, size: Size { width, height }, @@ -250,7 +251,13 @@ impl Dispatch for CaptureFrameState { width, height, } => { - frame.dmabuf_formats.push((format, width, height)); + tracing::debug!( + "Received wlr-screencopy linux_dmabuf event with format: {format} and size {width}x{height}" + ); + frame.dmabuf_formats.push(DMAFrameFormat { + format, + size: Size { width, height }, + }); } zwlr_screencopy_frame_v1::Event::BufferDone => { frame.buffer_done.store(true, Ordering::SeqCst); diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index e4140426..13463a13 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -14,17 +14,16 @@ mod screencopy; use std::{ collections::HashSet, fs::File, - os::fd::{AsFd, BorrowedFd}, + os::fd::{AsFd, BorrowedFd, OwnedFd}, sync::atomic::{AtomicBool, Ordering}, thread, }; use dispatch::LayerShellState; -use drm::Device; use image::{imageops::replace, DynamicImage}; use memmap2::MmapMut; use region::{EmbeddedRegion, RegionCapturer}; -use screencopy::FrameGuard; +use screencopy::{DMAFrameFormat, DMAFrameGuard, FrameData, FrameGuard}; use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, @@ -36,7 +35,9 @@ use wayland_client::{ Connection, EventQueue, Proxy, }; use wayland_protocols::{ - wp::linux_dmabuf::zv1::client::zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1, + wp::linux_dmabuf::zv1::client::{ + zwp_linux_buffer_params_v1, zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1, + }, xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, }, @@ -66,7 +67,7 @@ pub mod reexport { use wayland_client::protocol::wl_output; pub use wl_output::{Transform, WlOutput}; } -use gbm::Device as GBMDevice; +use gbm::{BufferObject, BufferObjectFlags, Device as GBMDevice}; struct Card(std::fs::File); /// Implementing [`AsFd`] is a prerequisite to implementing the traits found @@ -137,7 +138,6 @@ impl WayshotConnection { let linux_dmabuf = globals.bind(&evq.handle(), 4..=ZwpLinuxDmabufV1::interface().version, ())?; let gpu = Card::open("/dev/dri/renderD128"); - println!("{:#?}", gpu.get_driver().unwrap()); // init a GBM device let gbm = GBMDevice::new(gpu).unwrap(); let mut initial_state = Self { @@ -222,14 +222,74 @@ impl WayshotConnection { capture_region: Option, ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = - self.capture_output_frame_get_state(cursor_overlay, output, capture_region)?; + 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)?; Ok((frame_format, frame_guard)) } - fn capture_output_frame_get_state( + fn capture_output_frame_shm_from_file( + &self, + cursor_overlay: bool, + output: &WlOutput, + file: &File, + capture_region: Option, + ) -> Result<(FrameFormat, FrameGuard)> { + let (state, event_queue, frame, frame_format) = + self.capture_output_frame_get_state_shm(cursor_overlay as i32, output, capture_region)?; + + file.set_len(frame_format.byte_size())?; + + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; + + Ok((frame_format, frame_guard)) + } + + pub fn capture_output_frame_dmabuf( + &self, + cursor_overlay: bool, + output: &WlOutput, + capture_region: Option, + ) -> Result<(DMAFrameFormat, DMAFrameGuard, BufferObject<()>)> { + match &self.dmabuf_state { + Some(dmabuf_state) => { + let (state, event_queue, frame, frame_format) = self + .capture_output_frame_get_state_dmabuf( + cursor_overlay as i32, + output, + capture_region, + )?; + let gbm = &dmabuf_state.gbmdev; + let bo = gbm + .create_buffer_object::<()>( + frame_format.size.width, + frame_format.size.height, + gbm::Format::try_from(frame_format.format).unwrap(), + BufferObjectFlags::RENDERING | BufferObjectFlags::LINEAR, + ) + .unwrap(); + let stride = bo.stride().unwrap(); + let modifier: u64 = bo.modifier().unwrap().into(); + + let frame_guard = self.capture_output_frame_inner_dmabuf( + state, + event_queue, + frame, + frame_format, + stride, + modifier, + bo.fd().unwrap(), + )?; + + Ok((frame_format, frame_guard, bo)) + } + None => todo!(), + } + } + + fn capture_output_frame_get_state_shm( &self, cursor_overlay: i32, output: &WlOutput, @@ -320,6 +380,144 @@ impl WayshotConnection { Ok((state, event_queue, frame, frame_format)) } + fn capture_output_frame_get_state_dmabuf( + &self, + cursor_overlay: i32, + output: &WlOutput, + capture_region: Option, + ) -> Result<( + CaptureFrameState, + EventQueue, + ZwlrScreencopyFrameV1, + DMAFrameFormat, + )> { + let mut state = CaptureFrameState { + formats: Vec::new(), + dmabuf_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) => { + tracing::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"); + tracing::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "ZwlrScreencopy Manager not found".to_string(), + )); + } + }; + + debug!("Capturing output..."); + let frame = if let Some(embedded_region) = capture_region { + screencopy_manager.capture_output_region( + cursor_overlay, + output, + embedded_region.inner.position.x, + embedded_region.inner.position.y, + embedded_region.inner.size.width as i32, + embedded_region.inner.size.height as i32, + &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)?; + } + + tracing::trace!( + "Received compositor frame buffer formats: {:#?}", + state.formats + ); + // TODO select appropriate format if there is more than one + let frame_format = state.dmabuf_formats[0]; + tracing::trace!("Selected frame buffer format: {:#?}", frame_format); + + Ok((state, event_queue, frame, frame_format)) + } + + fn capture_output_frame_inner_dmabuf( + &self, + mut state: CaptureFrameState, + mut event_queue: EventQueue, + frame: ZwlrScreencopyFrameV1, + frame_format: DMAFrameFormat, + stride: u32, + modifier: u64, + fd: OwnedFd, + ) -> Result { + match &self.dmabuf_state { + Some(dmabuf_state) => { + // Connecting to wayland environment. + let qh = event_queue.handle(); + + let linux_dmabuf = &dmabuf_state.linux_dmabuf; + let dma_width = frame_format.size.width; + let dma_height = frame_format.size.height; + + let dma_params = linux_dmabuf.create_params(&qh, ()); + dma_params.add( + fd.as_fd(), + 0, + 0, + stride, + (modifier >> 32) as u32, + (modifier & 0xffffffff) as u32, + ); + let dmabuf_wlbuf = dma_params.create_immed( + dma_width as i32, + dma_height as i32, + frame_format.format, + zwp_linux_buffer_params_v1::Flags::empty(), + &qh, + (), + ); + + // Copy the pixel data advertised by the compositor into the buffer we just created. + frame.copy(&dmabuf_wlbuf); + tracing::debug!("wlr-screencopy copy() with dmabuf complete"); + + //dma_params.destroy(); + //linux_dmabuf.destroy(); + // 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 => { + tracing::error!("Frame copy failed"); + return Err(Error::FramecopyFailed); + } + FrameState::Finished => { + tracing::trace!("Frame copy finished"); + + return Ok(DMAFrameGuard { + buffer: dmabuf_wlbuf, + }); + } + } + } + + event_queue.blocking_dispatch(&mut state)?; + } + } + None => todo!(), + } + } + fn capture_output_frame_inner( &self, mut state: CaptureFrameState, @@ -374,23 +572,47 @@ impl WayshotConnection { } } - fn capture_output_frame_shm_from_file( - &self, - cursor_overlay: bool, - output: &WlOutput, - file: &File, - capture_region: Option, - ) -> Result<(FrameFormat, FrameGuard)> { - let (state, event_queue, frame, frame_format) = - self.capture_output_frame_get_state(cursor_overlay as i32, output, capture_region)?; - - file.set_len(frame_format.byte_size())?; - - let frame_guard = - self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; - - Ok((frame_format, frame_guard)) - } + // /// Get a FrameCopy instance with screenshot pixel data for any wl_output object via dmabuf. + // #[tracing::instrument(skip_all, fields(output = format!("{output_info}"), region = capture_region.map(|r| format!("{:}", r)).unwrap_or("fullscreen".to_string())))] + // fn capture_frame_copy_dmabuf( + // &self, + // cursor_overlay: bool, + // output_info: &OutputInfo, + // capture_region: Option, + // ) -> Result<(FrameCopy, FrameGuard)> { + // match &self.dmabuf_state { + // Some(dmabuf_state) => { + // let (frame_format, frame_guard, bo) = self.capture_output_frame_dmabuf_from_file( + // cursor_overlay, + // &output_info.wl_output, + // capture_region, + // )?; + // let rotated_physical_size = match output_info.transform { + // Transform::_90 + // | Transform::_270 + // | Transform::Flipped90 + // | Transform::Flipped270 => Size { + // width: frame_format.size.height, + // height: frame_format.size.width, + // }, + // _ => frame_format.size, + // }; + // let frame_copy = FrameCopy { + // frame_format, + // frame_color_type, + // frame_data: FrameData::GBMBo(bo), + // transform: output_info.transform, + // logical_region: capture_region + // .map(|capture_region| capture_region.logical()) + // .unwrap_or(output_info.logical_region), + // physical_size: rotated_physical_size, + // }; + // tracing::debug!("Created frame copy: {:#?}", frame_copy); + // Ok((frame_copy, frame_guard)) + // } + // None => todo!(), + // } + // } /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. #[tracing::instrument(skip_all, fields(output = format!("{output_info}"), region = capture_region.map(|r| format!("{:}", r)).unwrap_or("fullscreen".to_string())))] @@ -433,7 +655,7 @@ impl WayshotConnection { let frame_copy = FrameCopy { frame_format, frame_color_type, - frame_mmap, + frame_data: FrameData::Mmap(frame_mmap), transform: output_info.transform, logical_region: capture_region .map(|capture_region| capture_region.logical()) diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index 8b01c7c6..1bebcd33 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -4,6 +4,7 @@ use std::{ time::{SystemTime, UNIX_EPOCH}, }; +use gbm::BufferObject; use image::{ColorType, DynamicImage, ImageBuffer, Pixel}; use memmap2::MmapMut; use nix::{ @@ -32,6 +33,15 @@ impl Drop for FrameGuard { } } +pub struct DMAFrameGuard { + pub buffer: WlBuffer, +} +impl Drop for DMAFrameGuard { + fn drop(&mut self) { + self.buffer.destroy(); + } +} + /// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and /// Xbgr8888. /// @@ -46,6 +56,17 @@ pub struct FrameFormat { pub stride: u32, } +/// Type of DMABUF frame supported by the compositor +/// +/// See `zwlr_screencopy_frame_v1::Event::linux_dmabuf` as it's retrieved from there. +#[derive(Debug, Copy, Clone, PartialEq)] +pub struct DMAFrameFormat { + pub format: u32, + /// Size of the frame in pixels. This will always be in "landscape" so a + /// portrait 1080x1920 frame will be 1920x1080 and will need to be rotated! + pub size: Size, +} + impl FrameFormat { /// Returns the size of the frame in bytes, which is the stride * height. pub fn byte_size(&self) -> u64 { @@ -53,30 +74,38 @@ impl FrameFormat { } } -#[tracing::instrument(skip(frame_mmap))] +#[tracing::instrument(skip(frame_data))] fn create_image_buffer

( frame_format: &FrameFormat, - frame_mmap: &MmapMut, + frame_data: &FrameData, ) -> Result>> where P: Pixel, { tracing::debug!("Creating image buffer"); - ImageBuffer::from_vec( - frame_format.size.width, - frame_format.size.height, - frame_mmap.to_vec(), - ) - .ok_or(Error::BufferTooSmall) + match frame_data { + FrameData::Mmap(frame_mmap) => ImageBuffer::from_vec( + frame_format.size.width, + frame_format.size.height, + frame_mmap.to_vec(), + ) + .ok_or(Error::BufferTooSmall), + FrameData::GBMBo(_) => todo!(), + } } +#[derive(Debug)] +pub enum FrameData { + Mmap(MmapMut), + GBMBo(BufferObject<()>), +} /// The copied frame comprising of the FrameFormat, ColorType (Rgba8), and a memory backed shm /// file that holds the image data in it. #[derive(Debug)] pub struct FrameCopy { pub frame_format: FrameFormat, pub frame_color_type: ColorType, - pub frame_mmap: MmapMut, + pub frame_data: FrameData, pub transform: wl_output::Transform, /// Logical region with the transform already applied. pub logical_region: LogicalRegion, @@ -89,10 +118,10 @@ impl TryFrom<&FrameCopy> for DynamicImage { fn try_from(value: &FrameCopy) -> Result { Ok(match value.frame_color_type { ColorType::Rgb8 => { - Self::ImageRgb8(create_image_buffer(&value.frame_format, &value.frame_mmap)?) + Self::ImageRgb8(create_image_buffer(&value.frame_format, &value.frame_data)?) } ColorType::Rgba8 => { - Self::ImageRgba8(create_image_buffer(&value.frame_format, &value.frame_mmap)?) + Self::ImageRgba8(create_image_buffer(&value.frame_format, &value.frame_data)?) } _ => return Err(Error::InvalidColor), })