libwayshot_xcap/
lib.rs

1//! `libwayshot` is a convenient wrapper over the wlroots screenshot protocol
2//! that provides a simple API to take screenshots with.
3//!
4//! To get started, look at [`WayshotConnection`].
5
6mod convert;
7mod dispatch;
8mod error;
9mod image_util;
10pub mod output;
11pub mod region;
12mod screencopy;
13
14use std::{
15    collections::HashSet,
16    ffi::c_void,
17    fs::File,
18    os::fd::{AsFd, IntoRawFd, OwnedFd},
19    sync::atomic::{AtomicBool, Ordering},
20    thread,
21};
22
23use dispatch::{DMABUFState, LayerShellState};
24use image::{DynamicImage, imageops::replace};
25use khronos_egl::{self as egl, Instance};
26use memmap2::MmapMut;
27use region::{EmbeddedRegion, RegionCapturer};
28use screencopy::{DMAFrameFormat, DMAFrameGuard, EGLImageGuard, FrameData, FrameGuard};
29use tracing::debug;
30use wayland_client::{
31    Connection, EventQueue, Proxy,
32    globals::{GlobalList, registry_queue_init},
33    protocol::{
34        wl_compositor::WlCompositor,
35        wl_output::{Transform, WlOutput},
36        wl_shm::{self, WlShm},
37    },
38};
39use wayland_protocols::{
40    wp::{
41        linux_dmabuf::zv1::client::{
42            zwp_linux_buffer_params_v1, zwp_linux_dmabuf_v1::ZwpLinuxDmabufV1,
43        },
44        viewporter::client::wp_viewporter::WpViewporter,
45    },
46    xdg::xdg_output::zv1::client::{
47        zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1,
48    },
49};
50use wayland_protocols_wlr::{
51    layer_shell::v1::client::{
52        zwlr_layer_shell_v1::{Layer, ZwlrLayerShellV1},
53        zwlr_layer_surface_v1::Anchor,
54    },
55    screencopy::v1::client::{
56        zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1,
57        zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
58    },
59};
60
61use crate::{
62    convert::create_converter,
63    dispatch::{CaptureFrameState, FrameState, OutputCaptureState, WayshotState},
64    output::OutputInfo,
65    region::{LogicalRegion, Size},
66    screencopy::{FrameCopy, FrameFormat, create_shm_fd},
67};
68
69pub use crate::error::{Error, Result};
70
71pub mod reexport {
72    use wayland_client::protocol::wl_output;
73    pub use wl_output::{Transform, WlOutput};
74}
75use gbm::{BufferObject, BufferObjectFlags, Device as GBMDevice};
76
77/// Struct to store wayland connection and globals list.
78/// # Example usage
79///
80/// ```ignore
81/// use libwayshot::WayshotConnection;
82/// let wayshot_connection = WayshotConnection::new()?;
83/// let image_buffer = wayshot_connection.screenshot_all()?;
84/// ```
85#[derive(Debug)]
86pub struct WayshotConnection {
87    pub conn: Connection,
88    pub globals: GlobalList,
89    output_infos: Vec<OutputInfo>,
90    dmabuf_state: Option<DMABUFState>,
91}
92
93impl WayshotConnection {
94    pub fn new() -> Result<Self> {
95        let conn = Connection::connect_to_env()?;
96
97        Self::from_connection(conn)
98    }
99
100    /// Recommended if you already have a [`wayland_client::Connection`].
101    pub fn from_connection(conn: Connection) -> Result<Self> {
102        let (globals, _) = registry_queue_init::<WayshotState>(&conn)?;
103
104        let mut initial_state = Self {
105            conn,
106            globals,
107            output_infos: Vec::new(),
108            dmabuf_state: None,
109        };
110
111        initial_state.refresh_outputs()?;
112
113        Ok(initial_state)
114    }
115
116    ///Create a WayshotConnection struct having DMA-BUF support
117    /// Using this connection is required to make use of the dmabuf functions
118    ///# Parameters
119    /// - conn: a Wayland connection
120    /// - device_path: string pointing to the DRI device that is to be used for creating the DMA-BUFs on. For example: "/dev/dri/renderD128"
121    pub fn from_connection_with_dmabuf(conn: Connection, device_path: &str) -> Result<Self> {
122        let (globals, evq) = registry_queue_init::<WayshotState>(&conn)?;
123        let linux_dmabuf =
124            globals.bind(&evq.handle(), 4..=ZwpLinuxDmabufV1::interface().version, ())?;
125        let gpu = dispatch::Card::open(device_path);
126        // init a GBM device
127        let gbm = GBMDevice::new(gpu).unwrap();
128        let mut initial_state = Self {
129            conn,
130            globals,
131            output_infos: Vec::new(),
132            dmabuf_state: Some(DMABUFState {
133                linux_dmabuf,
134                gbmdev: gbm,
135            }),
136        };
137
138        initial_state.refresh_outputs()?;
139
140        Ok(initial_state)
141    }
142
143    /// Fetch all accessible wayland outputs.
144    pub fn get_all_outputs(&self) -> &[OutputInfo] {
145        self.output_infos.as_slice()
146    }
147
148    /// refresh the outputs, to get new outputs
149    pub fn refresh_outputs(&mut self) -> Result<()> {
150        // Connecting to wayland environment.
151        let mut state = OutputCaptureState {
152            outputs: Vec::new(),
153        };
154        let mut event_queue = self.conn.new_event_queue::<OutputCaptureState>();
155        let qh = event_queue.handle();
156
157        // Bind to xdg_output global.
158        let zxdg_output_manager = match self.globals.bind::<ZxdgOutputManagerV1, _, _>(
159            &qh,
160            3..=3,
161            (),
162        ) {
163            Ok(x) => x,
164            Err(e) => {
165                tracing::error!(
166                    "Failed to create ZxdgOutputManagerV1 version 3. Does your compositor implement ZxdgOutputManagerV1?"
167                );
168                panic!("{:#?}", e);
169            }
170        };
171
172        // Fetch all outputs; when their names arrive, add them to the list
173        let _ = self.conn.display().get_registry(&qh, ());
174        event_queue.roundtrip(&mut state)?;
175
176        // We loop over each output and request its position data.
177        let xdg_outputs: Vec<ZxdgOutputV1> = state
178            .outputs
179            .iter()
180            .enumerate()
181            .map(|(index, output)| {
182                zxdg_output_manager.get_xdg_output(&output.wl_output, &qh, index)
183            })
184            .collect();
185
186        event_queue.roundtrip(&mut state)?;
187
188        for xdg_output in xdg_outputs {
189            xdg_output.destroy();
190        }
191
192        if state.outputs.is_empty() {
193            tracing::error!("Compositor did not advertise any wl_output devices!");
194            return Err(Error::NoOutputs);
195        }
196        tracing::trace!("Outputs detected: {:#?}", state.outputs);
197        self.output_infos = state.outputs;
198
199        Ok(())
200    }
201
202    /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
203    ///  Data will be written to fd.
204    pub fn capture_output_frame_shm_fd<T: AsFd>(
205        &self,
206        cursor_overlay: i32,
207        output: &WlOutput,
208        fd: T,
209        capture_region: Option<EmbeddedRegion>,
210    ) -> Result<(FrameFormat, FrameGuard)> {
211        let (state, event_queue, frame, frame_format) =
212            self.capture_output_frame_get_state_shm(cursor_overlay, output, capture_region)?;
213        let frame_guard =
214            self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd)?;
215
216        Ok((frame_format, frame_guard))
217    }
218
219    fn capture_output_frame_shm_from_file(
220        &self,
221        cursor_overlay: bool,
222        output: &WlOutput,
223        file: &File,
224        capture_region: Option<EmbeddedRegion>,
225    ) -> Result<(FrameFormat, FrameGuard)> {
226        let (state, event_queue, frame, frame_format) =
227            self.capture_output_frame_get_state_shm(cursor_overlay as i32, output, capture_region)?;
228
229        file.set_len(frame_format.byte_size())?;
230
231        let frame_guard =
232            self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?;
233
234        Ok((frame_format, frame_guard))
235    }
236    /// # Safety
237    ///
238    /// Helper function/wrapper that uses the OpenGL extension OES_EGL_image to convert the EGLImage obtained from [`WayshotConnection::capture_output_frame_eglimage`]
239    /// into a OpenGL texture.
240    /// - The caller is supposed to setup everything required for the texture binding. An example call may look like:
241    /// ```no_run, ignore
242    /// gl::BindTexture(gl::TEXTURE_2D, self.gl_texture);
243    /// gl::TexParameteri(gl::TEXTURE_2D, gl::TEXTURE_MIN_FILTER, gl::LINEAR as i32);
244    /// wayshot_conn
245    ///     .bind_output_frame_to_gl_texture(
246    ///         true,
247    ///        &wayshot_conn.get_all_outputs()[0].wl_output,
248    ///        None)
249    ///```
250    /// # Parameters
251    /// - `cursor_overlay`: A boolean flag indicating whether the cursor should be included in the capture.
252    /// - `output`: Reference to the `WlOutput` from which the frame is to be captured.
253    /// - `capture_region`: Optional region specifying a sub-area of the output to capture. If `None`, the entire output is captured.
254    /// # Returns
255    /// - If the function was found and called, an OK(()), note that this does not necessarily mean that binding was successful, only that the function was called.
256    ///   The caller may check for any OpenGL errors using the standard routes.
257    /// - If the function was not found, [`Error::EGLImageToTexProcNotFoundError`] is returned
258    pub unsafe fn bind_output_frame_to_gl_texture(
259        &self,
260        cursor_overlay: bool,
261        output: &WlOutput,
262        capture_region: Option<EmbeddedRegion>,
263    ) -> Result<()> {
264        let egl = khronos_egl::Instance::new(egl::Static);
265        let eglimage_guard =
266            self.capture_output_frame_eglimage(&egl, cursor_overlay, output, capture_region)?;
267        unsafe {
268            let gl_egl_image_texture_target_2d_oes: unsafe extern "system" fn(
269                target: gl::types::GLenum,
270                image: gl::types::GLeglImageOES,
271            ) -> () =
272                std::mem::transmute(match egl.get_proc_address("glEGLImageTargetTexture2DOES") {
273                    Some(f) => {
274                        tracing::debug!("glEGLImageTargetTexture2DOES found at address {:#?}", f);
275                        f
276                    }
277                    None => {
278                        tracing::error!("glEGLImageTargetTexture2DOES not found");
279                        return Err(Error::EGLImageToTexProcNotFoundError);
280                    }
281                });
282
283            gl_egl_image_texture_target_2d_oes(gl::TEXTURE_2D, eglimage_guard.image.as_ptr());
284            tracing::trace!("glEGLImageTargetTexture2DOES called");
285            Ok(())
286        }
287    }
288
289    /// Obtain a screencapture in the form of a EGLImage.
290    /// The display on which this image is created is obtained from the Wayland Connection.
291    /// Uses the dma-buf provisions of the wlr-screencopy copy protocol to avoid VRAM->RAM copies
292    /// It returns the captured frame as an `EGLImage`, wrapped in an `EGLImageGuard`
293    /// for safe handling and cleanup.
294    /// # Parameters
295    /// - `egl_instance`: Reference to an egl API instance obtained from the khronos_egl crate, which is used to create the `EGLImage`.
296    /// - `cursor_overlay`: A boolean flag indicating whether the cursor should be included in the capture.
297    /// - `output`: Reference to the `WlOutput` from which the frame is to be captured.
298    /// - `capture_region`: Optional region specifying a sub-area of the output to capture. If `None`, the entire output is captured.
299    ///
300    /// # Returns
301    /// If successful, an EGLImageGuard which contains a pointer 'image' to the created EGLImage
302    /// On error, the EGL [error code](https://registry.khronos.org/EGL/sdk/docs/man/html/eglGetError.xhtml) is returned via this crates Error type
303    pub fn capture_output_frame_eglimage<'a, T: khronos_egl::api::EGL1_5>(
304        &self,
305        egl_instance: &'a Instance<T>,
306        cursor_overlay: bool,
307        output: &WlOutput,
308        capture_region: Option<EmbeddedRegion>,
309    ) -> Result<EGLImageGuard<'a, T>> {
310        let egl_display = unsafe {
311            match egl_instance.get_display(self.conn.display().id().as_ptr() as *mut c_void) {
312                Some(disp) => disp,
313                None => return Err(egl_instance.get_error().unwrap().into()),
314            }
315        };
316        tracing::trace!("eglDisplay obtained from Wayland connection's display");
317
318        egl_instance.initialize(egl_display)?;
319        self.capture_output_frame_eglimage_on_display(
320            egl_instance,
321            egl_display,
322            cursor_overlay,
323            output,
324            capture_region,
325        )
326    }
327
328    /// Obtain a screencapture in the form of a EGLImage on the given EGLDisplay.
329    ///
330    /// Uses the dma-buf provisions of the wlr-screencopy copy protocol to avoid VRAM->RAM copies
331    /// It returns the captured frame as an `EGLImage`, wrapped in an `EGLImageGuard`
332    /// for safe handling and cleanup.
333    /// # Parameters
334    /// - `egl_instance`: Reference to an `EGL1_5` instance, which is used to create the `EGLImage`.
335    /// - `egl_display`: The `EGLDisplay` on which the image should be created.
336    /// - `cursor_overlay`: A boolean flag indicating whether the cursor should be included in the capture.
337    /// - `output`: Reference to the `WlOutput` from which the frame is to be captured.
338    /// - `capture_region`: Optional region specifying a sub-area of the output to capture. If `None`, the entire output is captured.
339    ///
340    /// # Returns
341    /// If successful, an EGLImageGuard which contains a pointer 'image' to the created EGLImage
342    /// On error, the EGL [error code](https://registry.khronos.org/EGL/sdk/docs/man/html/eglGetError.xhtml) is returned via this crates Error type
343    pub fn capture_output_frame_eglimage_on_display<'a, T: khronos_egl::api::EGL1_5>(
344        &self,
345        egl_instance: &'a Instance<T>,
346        egl_display: egl::Display,
347        cursor_overlay: bool,
348        output: &WlOutput,
349        capture_region: Option<EmbeddedRegion>,
350    ) -> Result<EGLImageGuard<'a, T>> {
351        type Attrib = egl::Attrib;
352        let (frame_format, _guard, bo) =
353            self.capture_output_frame_dmabuf(cursor_overlay, output, capture_region)?;
354        let modifier: u64 = bo.modifier().into();
355        let image_attribs = [
356            egl::WIDTH as Attrib,
357            frame_format.size.width as Attrib,
358            egl::HEIGHT as Attrib,
359            frame_format.size.height as Attrib,
360            0x3271, //EGL_LINUX_DRM_FOURCC_EXT
361            bo.format() as Attrib,
362            0x3272, //EGL_DMA_BUF_PLANE0_FD_EXT
363            bo.fd_for_plane(0).unwrap().into_raw_fd() as Attrib,
364            0x3273, //EGL_DMA_BUF_PLANE0_OFFSET_EXT
365            bo.offset(0) as Attrib,
366            0x3274, //EGL_DMA_BUF_PLANE0_PITCH_EXT
367            bo.stride_for_plane(0) as Attrib,
368            0x3443, //EGL_DMA_BUF_PLANE0_MODIFIER_LO_EXT
369            (modifier as u32) as Attrib,
370            0x3444, //EGL_DMA_BUF_PLANE0_MODIFIER_HI_EXT
371            (modifier >> 32) as Attrib,
372            egl::ATTRIB_NONE as Attrib,
373        ];
374        tracing::debug!(
375            "Calling eglCreateImage with attributes: {:#?}",
376            image_attribs
377        );
378        unsafe {
379            match egl_instance.create_image(
380                egl_display,
381                khronos_egl::Context::from_ptr(egl::NO_CONTEXT),
382                0x3270, // EGL_LINUX_DMA_BUF_EXT
383                khronos_egl::ClientBuffer::from_ptr(std::ptr::null_mut()), //NULL
384                &image_attribs,
385            ) {
386                Ok(image) => Ok(EGLImageGuard {
387                    image,
388                    egl_instance,
389                    egl_display,
390                }),
391                Err(e) => {
392                    tracing::error!("eglCreateImage call failed with error {e}");
393                    Err(e.into())
394                }
395            }
396        }
397    }
398
399    /// Obtain a screencapture in the form of a WlBuffer backed by a GBM Bufferobject on the GPU.
400    /// Uses the dma-buf provisions of the wlr-screencopy copy protocol to avoid VRAM->RAM copies
401    /// The captured frame is returned as a tuple containing the frame format, a guard to manage
402    /// the WlBuffer's cleanup on drop, and the underlying `BufferObject`.
403    /// - `cursor_overlay`: A boolean flag indicating whether the cursor should be included in the capture.
404    /// - `output`: Reference to the `WlOutput` from which the frame is to be captured.
405    /// - `capture_region`: Optional region specifying a sub-area of the output to capture. If `None`, the entire output is captured.
406    ///# Returns
407    /// On success, returns a tuple containing the frame format,
408    ///   a guard to manage the frame's lifecycle, and the GPU-backed `BufferObject`.
409    /// # Errors
410    /// - Returns `NoDMAStateError` if the DMA-BUF state is not initialized a the time of initialization of this struct.
411    pub fn capture_output_frame_dmabuf(
412        &self,
413        cursor_overlay: bool,
414        output: &WlOutput,
415        capture_region: Option<EmbeddedRegion>,
416    ) -> Result<(DMAFrameFormat, DMAFrameGuard, BufferObject<()>)> {
417        match &self.dmabuf_state {
418            Some(dmabuf_state) => {
419                let (state, event_queue, frame, frame_format) = self
420                    .capture_output_frame_get_state_dmabuf(
421                        cursor_overlay as i32,
422                        output,
423                        capture_region,
424                    )?;
425                let gbm = &dmabuf_state.gbmdev;
426                let bo = gbm.create_buffer_object::<()>(
427                    frame_format.size.width,
428                    frame_format.size.height,
429                    gbm::Format::try_from(frame_format.format)?,
430                    BufferObjectFlags::RENDERING | BufferObjectFlags::LINEAR,
431                )?;
432
433                let stride = bo.stride();
434                let modifier: u64 = bo.modifier().into();
435                tracing::debug!(
436                    "Created GBM Buffer object with input frame format {:#?}, stride {:#?} and modifier {:#?} ",
437                    frame_format,
438                    stride,
439                    modifier
440                );
441                let frame_guard = self.capture_output_frame_inner_dmabuf(
442                    state,
443                    event_queue,
444                    frame,
445                    frame_format,
446                    stride,
447                    modifier,
448                    bo.fd_for_plane(0).unwrap(),
449                )?;
450
451                Ok((frame_format, frame_guard, bo))
452            }
453            None => Err(Error::NoDMAStateError),
454        }
455    }
456
457    // This API is exposed to provide users with access to window manager (WM)
458    // information. For instance, enabling Vulkan in wlroots alters the display
459    // format. Consequently, using PipeWire to capture streams without knowing
460    // the current format can lead to color distortion. This function attempts
461    // a trial screenshot to determine the screen's properties.
462    pub fn capture_output_frame_get_state_shm(
463        &self,
464        cursor_overlay: i32,
465        output: &WlOutput,
466        capture_region: Option<EmbeddedRegion>,
467    ) -> Result<(
468        CaptureFrameState,
469        EventQueue<CaptureFrameState>,
470        ZwlrScreencopyFrameV1,
471        FrameFormat,
472    )> {
473        let mut state = CaptureFrameState {
474            formats: Vec::new(),
475            dmabuf_formats: Vec::new(),
476            state: None,
477            buffer_done: AtomicBool::new(false),
478        };
479        let mut event_queue = self.conn.new_event_queue::<CaptureFrameState>();
480        let qh = event_queue.handle();
481
482        // Instantiating screencopy manager.
483        let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
484            &qh,
485            3..=3,
486            (),
487        ) {
488            Ok(x) => x,
489            Err(e) => {
490                tracing::error!(
491                    "Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"
492                );
493                tracing::error!("err: {e}");
494                return Err(Error::ProtocolNotFound(
495                    "ZwlrScreencopy Manager not found".to_string(),
496                ));
497            }
498        };
499
500        tracing::debug!("Capturing output(shm buffer)...");
501        let frame = if let Some(embedded_region) = capture_region {
502            screencopy_manager.capture_output_region(
503                cursor_overlay,
504                output,
505                embedded_region.inner.position.x,
506                embedded_region.inner.position.y,
507                embedded_region.inner.size.width as i32,
508                embedded_region.inner.size.height as i32,
509                &qh,
510                (),
511            )
512        } else {
513            screencopy_manager.capture_output(cursor_overlay, output, &qh, ())
514        };
515
516        // Empty internal event buffer until buffer_done is set to true which is when the Buffer done
517        // event is fired, aka the capture from the compositor is successful.
518        while !state.buffer_done.load(Ordering::SeqCst) {
519            event_queue.blocking_dispatch(&mut state)?;
520        }
521
522        tracing::trace!(
523            "Received compositor frame buffer formats: {:#?}",
524            state.formats
525        );
526        // Filter advertised wl_shm formats and select the first one that matches.
527        let frame_format = state
528            .formats
529            .iter()
530            .find(|frame| {
531                matches!(
532                    frame.format,
533                    wl_shm::Format::Xbgr2101010
534                        | wl_shm::Format::Abgr2101010
535                        | wl_shm::Format::Argb8888
536                        | wl_shm::Format::Xrgb8888
537                        | wl_shm::Format::Xbgr8888
538                        | wl_shm::Format::Bgr888
539                )
540            })
541            .copied()
542            // Check if frame format exists.
543            .ok_or_else(|| {
544                tracing::error!("No suitable frame format found");
545                Error::NoSupportedBufferFormat
546            })?;
547        tracing::trace!("Selected frame buffer format: {:#?}", frame_format);
548
549        Ok((state, event_queue, frame, frame_format))
550    }
551
552    fn capture_output_frame_get_state_dmabuf(
553        &self,
554        cursor_overlay: i32,
555        output: &WlOutput,
556        capture_region: Option<EmbeddedRegion>,
557    ) -> Result<(
558        CaptureFrameState,
559        EventQueue<CaptureFrameState>,
560        ZwlrScreencopyFrameV1,
561        DMAFrameFormat,
562    )> {
563        let mut state = CaptureFrameState {
564            formats: Vec::new(),
565            dmabuf_formats: Vec::new(),
566            state: None,
567            buffer_done: AtomicBool::new(false),
568        };
569        let mut event_queue = self.conn.new_event_queue::<CaptureFrameState>();
570        let qh = event_queue.handle();
571
572        // Instantiating screencopy manager.
573        let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
574            &qh,
575            3..=3,
576            (),
577        ) {
578            Ok(x) => x,
579            Err(e) => {
580                tracing::error!(
581                    "Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"
582                );
583                tracing::error!("err: {e}");
584                return Err(Error::ProtocolNotFound(
585                    "ZwlrScreencopy Manager not found".to_string(),
586                ));
587            }
588        };
589
590        tracing::debug!("Capturing output for DMA-BUF API...");
591        let frame = if let Some(embedded_region) = capture_region {
592            screencopy_manager.capture_output_region(
593                cursor_overlay,
594                output,
595                embedded_region.inner.position.x,
596                embedded_region.inner.position.y,
597                embedded_region.inner.size.width as i32,
598                embedded_region.inner.size.height as i32,
599                &qh,
600                (),
601            )
602        } else {
603            screencopy_manager.capture_output(cursor_overlay, output, &qh, ())
604        };
605
606        // Empty internal event buffer until buffer_done is set to true which is when the Buffer done
607        // event is fired, aka the capture from the compositor is successful.
608        while !state.buffer_done.load(Ordering::SeqCst) {
609            event_queue.blocking_dispatch(&mut state)?;
610        }
611
612        tracing::trace!(
613            "Received compositor frame buffer formats: {:#?}",
614            state.formats
615        );
616        // TODO select appropriate format if there is more than one
617        let frame_format = state.dmabuf_formats[0];
618        tracing::trace!("Selected frame buffer format: {:#?}", frame_format);
619
620        Ok((state, event_queue, frame, frame_format))
621    }
622
623    #[allow(clippy::too_many_arguments)]
624    fn capture_output_frame_inner_dmabuf(
625        &self,
626        mut state: CaptureFrameState,
627        mut event_queue: EventQueue<CaptureFrameState>,
628        frame: ZwlrScreencopyFrameV1,
629        frame_format: DMAFrameFormat,
630        stride: u32,
631        modifier: u64,
632        fd: OwnedFd,
633    ) -> Result<DMAFrameGuard> {
634        match &self.dmabuf_state {
635            Some(dmabuf_state) => {
636                // Connecting to wayland environment.
637                let qh = event_queue.handle();
638
639                let linux_dmabuf = &dmabuf_state.linux_dmabuf;
640                let dma_width = frame_format.size.width;
641                let dma_height = frame_format.size.height;
642
643                let dma_params = linux_dmabuf.create_params(&qh, ());
644
645                dma_params.add(
646                    fd.as_fd(),
647                    0,
648                    0,
649                    stride,
650                    (modifier >> 32) as u32,
651                    (modifier & 0xffffffff) as u32,
652                );
653                tracing::trace!("Called  ZwpLinuxBufferParamsV1::create_params ");
654                let dmabuf_wlbuf = dma_params.create_immed(
655                    dma_width as i32,
656                    dma_height as i32,
657                    frame_format.format,
658                    zwp_linux_buffer_params_v1::Flags::empty(),
659                    &qh,
660                    (),
661                );
662                tracing::trace!("Called  ZwpLinuxBufferParamsV1::create_immed to create WlBuffer ");
663                // Copy the pixel data advertised by the compositor into the buffer we just created.
664                frame.copy(&dmabuf_wlbuf);
665                tracing::debug!("wlr-screencopy copy() with dmabuf complete");
666
667                // On copy the Ready / Failed events are fired by the frame object, so here we check for them.
668                loop {
669                    // Basically reads, if frame state is not None then...
670                    if let Some(state) = state.state {
671                        match state {
672                            FrameState::Failed => {
673                                tracing::error!("Frame copy failed");
674                                return Err(Error::FramecopyFailed);
675                            }
676                            FrameState::Finished => {
677                                tracing::trace!("Frame copy finished");
678
679                                return Ok(DMAFrameGuard {
680                                    buffer: dmabuf_wlbuf,
681                                });
682                            }
683                        }
684                    }
685
686                    event_queue.blocking_dispatch(&mut state)?;
687                }
688            }
689            None => Err(Error::NoDMAStateError),
690        }
691    }
692
693    fn capture_output_frame_inner<T: AsFd>(
694        &self,
695        mut state: CaptureFrameState,
696        mut event_queue: EventQueue<CaptureFrameState>,
697        frame: ZwlrScreencopyFrameV1,
698        frame_format: FrameFormat,
699        fd: T,
700    ) -> Result<FrameGuard> {
701        // Connecting to wayland environment.
702        let qh = event_queue.handle();
703
704        // Instantiate shm global.
705        let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ())?;
706        let shm_pool = shm.create_pool(
707            fd.as_fd(),
708            frame_format
709                .byte_size()
710                .try_into()
711                .map_err(|_| Error::BufferTooSmall)?,
712            &qh,
713            (),
714        );
715        let buffer = shm_pool.create_buffer(
716            0,
717            frame_format.size.width as i32,
718            frame_format.size.height as i32,
719            frame_format.stride as i32,
720            frame_format.format,
721            &qh,
722            (),
723        );
724
725        // Copy the pixel data advertised by the compositor into the buffer we just created.
726        frame.copy(&buffer);
727        // On copy the Ready / Failed events are fired by the frame object, so here we check for them.
728        loop {
729            // Basically reads, if frame state is not None then...
730            if let Some(state) = state.state {
731                match state {
732                    FrameState::Failed => {
733                        tracing::error!("Frame copy failed");
734                        return Err(Error::FramecopyFailed);
735                    }
736                    FrameState::Finished => {
737                        tracing::trace!("Frame copy finished");
738                        return Ok(FrameGuard { buffer, shm_pool });
739                    }
740                }
741            }
742
743            event_queue.blocking_dispatch(&mut state)?;
744        }
745    }
746
747    /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
748    #[tracing::instrument(skip_all, fields(output = format!("{output_info}"), region = capture_region.map(|r| format!("{:}", r)).unwrap_or("fullscreen".to_string())))]
749    fn capture_frame_copy(
750        &self,
751        cursor_overlay: bool,
752        output_info: &OutputInfo,
753        capture_region: Option<EmbeddedRegion>,
754    ) -> Result<(FrameCopy, FrameGuard)> {
755        // Create an in memory file and return it's file descriptor.
756        let fd = create_shm_fd()?;
757        // Create a writeable memory map backed by a mem_file.
758        let mem_file = File::from(fd);
759
760        let (frame_format, frame_guard) = self.capture_output_frame_shm_from_file(
761            cursor_overlay,
762            &output_info.wl_output,
763            &mem_file,
764            capture_region,
765        )?;
766
767        let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? };
768        let data = &mut *frame_mmap;
769        let frame_color_type = match create_converter(frame_format.format) {
770            Some(converter) => converter.convert_inplace(data),
771            _ => {
772                tracing::error!("Unsupported buffer format: {:?}", frame_format.format);
773                tracing::error!(
774                    "You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."
775                );
776                return Err(Error::NoSupportedBufferFormat);
777            }
778        };
779        let rotated_physical_size = match output_info.transform {
780            Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => {
781                Size {
782                    width: frame_format.size.height,
783                    height: frame_format.size.width,
784                }
785            }
786            _ => frame_format.size,
787        };
788        let frame_copy = FrameCopy {
789            frame_format,
790            frame_color_type,
791            frame_data: FrameData::Mmap(frame_mmap),
792            transform: output_info.transform,
793            logical_region: capture_region
794                .map(|capture_region| capture_region.logical())
795                .unwrap_or(output_info.logical_region),
796            physical_size: rotated_physical_size,
797        };
798        tracing::debug!("Created frame copy: {:#?}", frame_copy);
799        Ok((frame_copy, frame_guard))
800    }
801
802    pub fn capture_frame_copies(
803        &self,
804        output_capture_regions: &[(OutputInfo, Option<EmbeddedRegion>)],
805        cursor_overlay: bool,
806    ) -> Result<Vec<(FrameCopy, FrameGuard, OutputInfo)>> {
807        output_capture_regions
808            .iter()
809            .map(|(output_info, capture_region)| {
810                self.capture_frame_copy(cursor_overlay, output_info, *capture_region)
811                    .map(|(frame_copy, frame_guard)| (frame_copy, frame_guard, output_info.clone()))
812            })
813            .collect()
814    }
815
816    /// Create a layer shell surface for each output,
817    /// render the screen captures on them and use the callback to select a region from them
818    fn overlay_frames_and_select_region<F>(
819        &self,
820        frames: &[(FrameCopy, FrameGuard, OutputInfo)],
821        callback: F,
822    ) -> Result<LogicalRegion>
823    where
824        F: Fn(&WayshotConnection) -> Result<LogicalRegion, Error>,
825    {
826        let mut state = LayerShellState {
827            configured_outputs: HashSet::new(),
828        };
829        let mut event_queue: EventQueue<LayerShellState> =
830            self.conn.new_event_queue::<LayerShellState>();
831        let qh = event_queue.handle();
832
833        let compositor = match self.globals.bind::<WlCompositor, _, _>(&qh, 3..=3, ()) {
834            Ok(x) => x,
835            Err(e) => {
836                tracing::error!(
837                    "Failed to create compositor. Does your compositor implement WlCompositor?"
838                );
839                tracing::error!("err: {e}");
840                return Err(Error::ProtocolNotFound(
841                    "WlCompositor not found".to_string(),
842                ));
843            }
844        };
845        let layer_shell = match self.globals.bind::<ZwlrLayerShellV1, _, _>(&qh, 1..=1, ()) {
846            Ok(x) => x,
847            Err(e) => {
848                tracing::error!(
849                    "Failed to create layer shell. Does your compositor implement WlrLayerShellV1?"
850                );
851                tracing::error!("err: {e}");
852                return Err(Error::ProtocolNotFound(
853                    "WlrLayerShellV1 not found".to_string(),
854                ));
855            }
856        };
857        let viewporter = self.globals.bind::<WpViewporter, _, _>(&qh, 1..=1, ()).ok();
858        if viewporter.is_none() {
859            tracing::info!(
860                "Compositor does not support wp_viewporter, display scaling may be inaccurate."
861            );
862        }
863
864        let mut layer_shell_surfaces = Vec::with_capacity(frames.len());
865
866        for (frame_copy, frame_guard, output_info) in frames {
867            tracing::span!(
868                tracing::Level::DEBUG,
869                "overlay_frames::surface",
870                output = format!("{output_info}")
871            )
872            .in_scope(|| -> Result<()> {
873                let surface = compositor.create_surface(&qh, ());
874
875                let layer_surface = layer_shell.get_layer_surface(
876                    &surface,
877                    Some(&output_info.wl_output),
878                    Layer::Top,
879                    "wayshot".to_string(),
880                    &qh,
881                    output_info.wl_output.clone(),
882                );
883
884                layer_surface.set_exclusive_zone(-1);
885                layer_surface.set_anchor(Anchor::Top | Anchor::Left);
886                layer_surface.set_size(
887                    frame_copy.frame_format.size.width,
888                    frame_copy.frame_format.size.height,
889                );
890
891                debug!("Committing surface creation changes.");
892                surface.commit();
893
894                debug!("Waiting for layer surface to be configured.");
895                while !state.configured_outputs.contains(&output_info.wl_output) {
896                    event_queue.blocking_dispatch(&mut state)?;
897                }
898
899                surface.set_buffer_transform(output_info.transform);
900                // surface.set_buffer_scale(output_info.scale());
901                surface.attach(Some(&frame_guard.buffer), 0, 0);
902
903                if let Some(viewporter) = viewporter.as_ref() {
904                    let viewport = viewporter.get_viewport(&surface, &qh, ());
905                    viewport.set_destination(
906                        output_info.logical_region.inner.size.width as i32,
907                        output_info.logical_region.inner.size.height as i32,
908                    );
909                }
910
911                debug!("Committing surface with attached buffer.");
912                surface.commit();
913                layer_shell_surfaces.push((surface, layer_surface));
914                event_queue.blocking_dispatch(&mut state)?;
915
916                Ok(())
917            })?;
918        }
919
920        let callback_result = callback(self);
921
922        debug!("Unmapping and destroying layer shell surfaces.");
923        for (surface, layer_shell_surface) in layer_shell_surfaces.iter() {
924            surface.attach(None, 0, 0);
925            surface.commit(); //unmap surface by committing a null buffer
926            layer_shell_surface.destroy();
927        }
928        event_queue.roundtrip(&mut state)?;
929
930        callback_result
931    }
932
933    /// Take a screenshot from the specified region.
934    #[tracing::instrument(skip_all, fields(max_scale = tracing::field::Empty))]
935    fn screenshot_region_capturer(
936        &self,
937        region_capturer: RegionCapturer,
938        cursor_overlay: bool,
939    ) -> Result<DynamicImage> {
940        let outputs_capture_regions: Vec<(OutputInfo, Option<EmbeddedRegion>)> =
941            match region_capturer {
942                RegionCapturer::Outputs(ref outputs) => outputs
943                    .iter()
944                    .map(|output_info| (output_info.clone(), None))
945                    .collect(),
946                RegionCapturer::Region(capture_region) => self
947                    .get_all_outputs()
948                    .iter()
949                    .filter_map(|output_info| {
950                        tracing::span!(
951                            tracing::Level::DEBUG,
952                            "filter_map",
953                            output = format!(
954                                "{output_info} at {region}",
955                                output_info = format!("{output_info}"),
956                                region = LogicalRegion::from(output_info),
957                            ),
958                            capture_region = format!("{}", capture_region),
959                        )
960                        .in_scope(|| {
961                            if let Some(relative_region) =
962                                EmbeddedRegion::new(capture_region, output_info.into())
963                            {
964                                tracing::debug!("Intersection found: {}", relative_region);
965                                Some((output_info.clone(), Some(relative_region)))
966                            } else {
967                                tracing::debug!("No intersection found");
968                                None
969                            }
970                        })
971                    })
972                    .collect(),
973                RegionCapturer::Freeze(_) => self
974                    .get_all_outputs()
975                    .iter()
976                    .map(|output_info| (output_info.clone(), None))
977                    .collect(),
978            };
979
980        let frames = self.capture_frame_copies(&outputs_capture_regions, cursor_overlay)?;
981
982        let capture_region: LogicalRegion = match region_capturer {
983            RegionCapturer::Outputs(outputs) => outputs.as_slice().try_into()?,
984            RegionCapturer::Region(region) => region,
985            RegionCapturer::Freeze(callback) => {
986                self.overlay_frames_and_select_region(&frames, callback)?
987            }
988        };
989
990        // TODO When freeze was used, we can still further remove the outputs
991        // that don't intersect with the capture region.
992
993        thread::scope(|scope| {
994            let max_scale = outputs_capture_regions
995                .iter()
996                .map(|(output_info, _)| output_info.scale())
997                .fold(1.0, f64::max);
998
999            tracing::Span::current().record("max_scale", max_scale);
1000
1001            let rotate_join_handles = frames
1002                .into_iter()
1003                .map(|(frame_copy, _, _)| {
1004                    scope.spawn(move || {
1005                        let image = (&frame_copy).try_into()?;
1006                        Ok((
1007                            image_util::rotate_image_buffer(
1008                                image,
1009                                frame_copy.transform,
1010                                frame_copy.logical_region.inner.size,
1011                                max_scale,
1012                            ),
1013                            frame_copy,
1014                        ))
1015                    })
1016                })
1017                .collect::<Vec<_>>();
1018
1019            rotate_join_handles
1020                .into_iter()
1021                .flat_map(|join_handle| join_handle.join())
1022                .fold(
1023                    None,
1024                    |composite_image: Option<Result<_>>, image: Result<_>| {
1025                        // Default to a transparent image.
1026                        let composite_image = composite_image.unwrap_or_else(|| {
1027                            Ok(DynamicImage::new_rgba8(
1028                                (capture_region.inner.size.width as f64 * max_scale) as u32,
1029                                (capture_region.inner.size.height as f64 * max_scale) as u32,
1030                            ))
1031                        });
1032
1033                        Some(|| -> Result<_> {
1034                            let mut composite_image = composite_image?;
1035                            let (image, frame_copy) = image?;
1036                            let (x, y) = (
1037                                ((frame_copy.logical_region.inner.position.x as f64
1038                                    - capture_region.inner.position.x as f64)
1039                                    * max_scale) as i64,
1040                                ((frame_copy.logical_region.inner.position.y as f64
1041                                    - capture_region.inner.position.y as f64)
1042                                    * max_scale) as i64,
1043                            );
1044                            tracing::span!(
1045                                tracing::Level::DEBUG,
1046                                "replace",
1047                                frame_copy_region = format!("{}", frame_copy.logical_region),
1048                                capture_region = format!("{}", capture_region),
1049                                x = x,
1050                                y = y,
1051                            )
1052                            .in_scope(|| {
1053                                tracing::debug!("Replacing parts of the final image");
1054                                replace(&mut composite_image, &image, x, y);
1055                            });
1056
1057                            Ok(composite_image)
1058                        }())
1059                    },
1060                )
1061                .ok_or_else(|| {
1062                    tracing::error!("Provided capture region doesn't intersect with any outputs!");
1063                    Error::NoOutputs
1064                })?
1065        })
1066    }
1067
1068    /// Take a screenshot from the specified region.
1069    pub fn screenshot(
1070        &self,
1071        capture_region: LogicalRegion,
1072        cursor_overlay: bool,
1073    ) -> Result<DynamicImage> {
1074        self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay)
1075    }
1076
1077    /// Take a screenshot, overlay the screenshot, run the callback, and then
1078    /// unfreeze the screenshot and return the selected region.
1079    pub fn screenshot_freeze<F>(&self, callback: F, cursor_overlay: bool) -> Result<DynamicImage>
1080    where
1081        F: Fn(&WayshotConnection) -> Result<LogicalRegion> + 'static,
1082    {
1083        self.screenshot_region_capturer(RegionCapturer::Freeze(Box::new(callback)), cursor_overlay)
1084    }
1085
1086    /// Take a screenshot from one output
1087    pub fn screenshot_single_output(
1088        &self,
1089        output_info: &OutputInfo,
1090        cursor_overlay: bool,
1091    ) -> Result<DynamicImage> {
1092        let (frame_copy, _) = self.capture_frame_copy(cursor_overlay, output_info, None)?;
1093        (&frame_copy).try_into()
1094    }
1095
1096    /// Take a screenshot from all of the specified outputs.
1097    pub fn screenshot_outputs(
1098        &self,
1099        outputs: &[OutputInfo],
1100        cursor_overlay: bool,
1101    ) -> Result<DynamicImage> {
1102        if outputs.is_empty() {
1103            return Err(Error::NoOutputs);
1104        }
1105
1106        self.screenshot_region_capturer(RegionCapturer::Outputs(outputs.to_owned()), cursor_overlay)
1107    }
1108
1109    /// Take a screenshot from all accessible outputs.
1110    pub fn screenshot_all(&self, cursor_overlay: bool) -> Result<DynamicImage> {
1111        self.screenshot_outputs(self.get_all_outputs(), cursor_overlay)
1112    }
1113}
libwayshot_xcap/lib.rs

libwayshot_xcap/
lib.rs
