diff --git a/.gitignore b/.gitignore index 5a4f6354..2622eacf 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,8 @@ target *.gz *.out .direnv +*.jpg +*.jpeg +*.png +*.ppm +*.qoi diff --git a/Cargo.lock b/Cargo.lock index 2d5bff5d..9244fa4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,9 +10,9 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" [[package]] name = "anstream" -version = "0.6.5" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d664a92ecae85fd0a7392615844904654d1d5f5514837f471ddef4a057aba1b6" +checksum = "6e2e1ebcb11de5c03c67de28a7df593d32191b44939c482e97702baaaa6ab6a5" dependencies = [ "anstyle", "anstyle-parse", @@ -103,18 +103,19 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "clap" -version = "4.4.11" +version = "4.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfaff671f6b22ca62406885ece523383b9b64022e341e53e009a62ebc47a45f2" +checksum = "1e578d6ec4194633722ccf9544794b71b1385c3c027efe0c55db226fc880865c" dependencies = [ "clap_builder", + "clap_derive", ] [[package]] name = "clap_builder" -version = "4.4.11" +version = "4.4.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a216b506622bb1d316cd51328dce24e07bdff4a6128a47c7e7fad11878d5adbb" +checksum = "4df4df40ec50c46000231c914968278b1eb05098cf8f1b3a518a95030e71d1c7" dependencies = [ "anstream", "anstyle", @@ -122,6 +123,18 @@ dependencies = [ "strsim", ] +[[package]] +name = "clap_derive" +version = "4.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "clap_lex" version = "0.6.0" @@ -207,6 +220,16 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "eyre" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c2b6b5a29c02cdc822728b7d7b8ae1bab3e3b05d44522770ddd49722eeac7eb" +dependencies = [ + "indenter", + "once_cell", +] + [[package]] name = "fastrand" version = "2.0.1" @@ -241,6 +264,12 @@ dependencies = [ "thread_local", ] +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "image" version = "0.24.7" @@ -257,6 +286,12 @@ dependencies = [ "qoi", ] +[[package]] +name = "indenter" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce23b50ad8242c51a442f3ff322d56b02f08852c77e4c0b4d3fd684abc89c683" + [[package]] name = "jpeg-decoder" version = "0.3.0" @@ -756,6 +791,7 @@ version = "1.3.2-dev" dependencies = [ "clap", "dialoguer", + "eyre", "flate2", "image", "libwayshot", diff --git a/libwayshot/README.md b/libwayshot/README.md index a18983c7..a3bf2ab2 100644 --- a/libwayshot/README.md +++ b/libwayshot/README.md @@ -17,6 +17,6 @@ ```rust use libwayshot::WayshotConnection; -let wayshot_connection = WayshotConnection::new().unwrap(); -let image_buffer = wayshot_connection.screenshot_all().unwrap(); +let wayshot_connection = WayshotConnection::new()?; +let image_buffer = wayshot_connection.screenshot_all()?; ``` diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 9b760c21..574fc813 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -1,13 +1,14 @@ use std::{ - process::exit, + collections::HashSet, sync::atomic::{AtomicBool, Ordering}, }; use wayland_client::{ delegate_noop, globals::GlobalListContents, protocol::{ - wl_buffer::WlBuffer, wl_output, wl_output::WlOutput, wl_registry, wl_registry::WlRegistry, - wl_shm::WlShm, wl_shm_pool::WlShmPool, + wl_buffer::WlBuffer, wl_compositor::WlCompositor, wl_output, wl_output::WlOutput, + wl_registry, wl_registry::WlRegistry, wl_shm::WlShm, wl_shm_pool::WlShmPool, + wl_surface::WlSurface, }, Connection, Dispatch, QueueHandle, WEnum, WEnum::Value, @@ -15,21 +16,28 @@ use wayland_client::{ use wayland_protocols::xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1, }; +use wayland_protocols_wlr::layer_shell::v1::client::{ + zwlr_layer_shell_v1::ZwlrLayerShellV1, + zwlr_layer_surface_v1::{self, ZwlrLayerSurfaceV1}, +}; use wayland_protocols_wlr::screencopy::v1::client::{ zwlr_screencopy_frame_v1, zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, }; use crate::{ - output::{OutputInfo, OutputPositioning, WlOutputMode}, + output::OutputInfo, + region::{Position, Region, Size}, screencopy::FrameFormat, }; +#[derive(Debug)] pub struct OutputCaptureState { pub outputs: Vec, } impl Dispatch for OutputCaptureState { + #[tracing::instrument(skip(wl_registry, qh), ret, level = "trace")] fn event( state: &mut Self, wl_registry: &WlRegistry, @@ -56,16 +64,8 @@ impl Dispatch for OutputCaptureState { name: "".to_string(), description: String::new(), transform: wl_output::Transform::Normal, - dimensions: OutputPositioning { - x: 0, - y: 0, - width: 0, - height: 0, - }, - mode: WlOutputMode { - width: 0, - height: 0, - }, + scale: 1, + region: Region::default(), }); } else { tracing::error!("Ignoring a wl_output with version < 4."); @@ -76,6 +76,7 @@ impl Dispatch for OutputCaptureState { } impl Dispatch for OutputCaptureState { + #[tracing::instrument(skip(wl_output), ret, level = "trace")] fn event( state: &mut Self, wl_output: &WlOutput, @@ -84,11 +85,13 @@ impl Dispatch for OutputCaptureState { _: &Connection, _: &QueueHandle, ) { - let output: &mut OutputInfo = state - .outputs - .iter_mut() - .find(|x| x.wl_output == *wl_output) - .unwrap(); + let output: &mut OutputInfo = + if let Some(output) = state.outputs.iter_mut().find(|x| x.wl_output == *wl_output) { + output + } else { + tracing::error!("Received event for an output that is not registered: {event:#?}"); + return; + }; match event { wl_output::Event::Name { name } => { @@ -97,16 +100,18 @@ impl Dispatch for OutputCaptureState { wl_output::Event::Description { description } => { output.description = description; } - wl_output::Event::Mode { width, height, .. } => { - output.mode = WlOutputMode { width, height }; - } + wl_output::Event::Mode { .. } => {} wl_output::Event::Geometry { transform: WEnum::Value(transform), .. } => { output.transform = transform; } - _ => (), + wl_output::Event::Scale { factor } => { + output.scale = factor; + } + wl_output::Event::Done => {} + _ => {} } } } @@ -114,6 +119,7 @@ impl Dispatch for OutputCaptureState { delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1); impl Dispatch for OutputCaptureState { + #[tracing::instrument(ret, level = "trace")] fn event( state: &mut Self, _: &ZxdgOutputV1, @@ -122,25 +128,34 @@ impl Dispatch for OutputCaptureState { _: &Connection, _: &QueueHandle, ) { - let output_info = state.outputs.get_mut(*index).unwrap(); + let output_info = if let Some(output_info) = state.outputs.get_mut(*index) { + output_info + } else { + tracing::error!( + "Received event for output index {index} that is not registered: {event:#?}" + ); + return; + }; match event { zxdg_output_v1::Event::LogicalPosition { x, y } => { - output_info.dimensions.x = x; - output_info.dimensions.y = y; - tracing::debug!("Logical position event fired!"); + output_info.region.position = Position { x, y }; } zxdg_output_v1::Event::LogicalSize { width, height } => { - output_info.dimensions.width = width; - output_info.dimensions.height = height; - tracing::debug!("Logical size event fired!"); + output_info.region.size = Size { + width: width as u32, + height: height as u32, + }; } + zxdg_output_v1::Event::Done => {} + zxdg_output_v1::Event::Name { .. } => {} + zxdg_output_v1::Event::Description { .. } => {} _ => {} }; } } -/// State of the frame after attemting to copy it's data to a wl_buffer. +/// State of the frame after attempting to copy it's data to a wl_buffer. #[derive(Debug, Copy, Clone, PartialEq)] pub enum FrameState { /// Compositor returned a failed event on calling `frame.copy`. @@ -156,6 +171,7 @@ pub struct CaptureFrameState { } impl Dispatch for CaptureFrameState { + #[tracing::instrument(skip(frame), ret, level = "trace")] fn event( frame: &mut Self, _: &ZwlrScreencopyFrameV1, @@ -171,43 +187,30 @@ impl Dispatch for CaptureFrameState { height, stride, } => { - tracing::debug!("Received Buffer event"); if let Value(f) = format { frame.formats.push(FrameFormat { format: f, - width, - height, + size: Size { width, height }, stride, }) } else { tracing::debug!("Received Buffer event with unidentified format"); - exit(1); } } - zwlr_screencopy_frame_v1::Event::Flags { .. } => { - tracing::debug!("Received Flags event"); - } zwlr_screencopy_frame_v1::Event::Ready { .. } => { // If the frame is successfully copied, a “flags” and a “ready” events are sent. Otherwise, a “failed” event is sent. // This is useful when we call .copy on the frame object. - tracing::debug!("Received Ready event"); frame.state.replace(FrameState::Finished); } zwlr_screencopy_frame_v1::Event::Failed => { - tracing::debug!("Received Failed event"); frame.state.replace(FrameState::Failed); } - zwlr_screencopy_frame_v1::Event::Damage { .. } => { - tracing::debug!("Received Damage event"); - } - zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => { - tracing::debug!("Received LinuxDmaBuf event"); - } + zwlr_screencopy_frame_v1::Event::Damage { .. } => {} + zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => {} zwlr_screencopy_frame_v1::Event::BufferDone => { - tracing::debug!("Received bufferdone event"); frame.buffer_done.store(true, Ordering::SeqCst); } - _ => unreachable!(), + _ => {} }; } } @@ -232,3 +235,43 @@ impl wayland_client::Dispatch for W ) { } } + +pub struct LayerShellState { + pub configured_outputs: HashSet, +} + +delegate_noop!(LayerShellState: ignore WlCompositor); +delegate_noop!(LayerShellState: ignore WlShm); +delegate_noop!(LayerShellState: ignore WlShmPool); +delegate_noop!(LayerShellState: ignore WlBuffer); +delegate_noop!(LayerShellState: ignore ZwlrLayerShellV1); +delegate_noop!(LayerShellState: ignore WlSurface); + +impl wayland_client::Dispatch for LayerShellState { + // No need to instrument here, span from lib.rs is automatically used. + fn event( + state: &mut Self, + proxy: &ZwlrLayerSurfaceV1, + event: ::Event, + data: &WlOutput, + _conn: &Connection, + _qhandle: &QueueHandle, + ) { + match event { + zwlr_layer_surface_v1::Event::Configure { + serial, + width: _, + height: _, + } => { + tracing::debug!("Acking configure"); + state.configured_outputs.insert(data.clone()); + proxy.ack_configure(serial); + tracing::trace!("Acked configure"); + } + zwlr_layer_surface_v1::Event::Closed => { + tracing::debug!("Closed") + } + _ => {} + } + } +} diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs index ddeff581..f3ad8523 100644 --- a/libwayshot/src/error.rs +++ b/libwayshot/src/error.rs @@ -1,7 +1,10 @@ use std::{io, result}; use thiserror::Error; -use wayland_client::{globals::GlobalError, ConnectError, DispatchError}; +use wayland_client::{ + globals::{BindError, GlobalError}, + ConnectError, DispatchError, +}; pub type Result = result::Result; @@ -17,6 +20,8 @@ pub enum Error { Io(#[from] io::Error), #[error("dispatch error: {0}")] Dispatch(#[from] DispatchError), + #[error("bind error: {0}")] + Bind(#[from] BindError), #[error("global error: {0}")] Global(#[from] GlobalError), #[error("connect error: {0}")] @@ -27,4 +32,6 @@ pub enum Error { NoSupportedBufferFormat, #[error("Cannot find required wayland protocol")] ProtocolNotFound(String), + #[error("error occurred in freeze callback")] + FreezeCallbackError, } diff --git a/libwayshot/src/image_util.rs b/libwayshot/src/image_util.rs index 92eb3967..197e8f11 100644 --- a/libwayshot/src/image_util.rs +++ b/libwayshot/src/image_util.rs @@ -7,6 +7,14 @@ pub(crate) fn rotate_image_buffer( width: u32, height: u32, ) -> DynamicImage { + // TODO Better document whether width and height are before or after the transform. + // Perhaps this should be part of a cleanup of the FrameCopy struct. + let (width, height) = match transform { + Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => { + (height, width) + } + _ => (width, height), + }; let final_image = match transform { Transform::_90 => image::imageops::rotate90(&image).into(), Transform::_180 => image::imageops::rotate180(&image).into(), diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 6c003320..beda48fe 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -8,22 +8,27 @@ mod dispatch; mod error; mod image_util; pub mod output; +pub mod region; mod screencopy; use std::{ - cmp, + collections::HashSet, fs::File, os::fd::AsFd, - process::exit, sync::atomic::{AtomicBool, Ordering}, thread, }; -use image::{imageops::overlay, DynamicImage}; +use dispatch::LayerShellState; +use image::{imageops::replace, DynamicImage}; use memmap2::MmapMut; +use region::{EmbeddedRegion, RegionCapturer}; +use screencopy::FrameGuard; +use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ + wl_compositor::WlCompositor, wl_output::{Transform, WlOutput}, wl_shm::{self, WlShm}, }, @@ -32,15 +37,22 @@ use wayland_client::{ use wayland_protocols::xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, }; -use wayland_protocols_wlr::screencopy::v1::client::{ - zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, - zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, +use wayland_protocols_wlr::{ + layer_shell::v1::client::{ + zwlr_layer_shell_v1::{Layer, ZwlrLayerShellV1}, + zwlr_layer_surface_v1::Anchor, + }, + screencopy::v1::client::{ + zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, + zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, + }, }; use crate::{ convert::create_converter, dispatch::{CaptureFrameState, FrameState, OutputCaptureState, WayshotState}, output::OutputInfo, + region::{LogicalRegion, Region, Size}, screencopy::{create_shm_fd, FrameCopy, FrameFormat}, }; @@ -51,34 +63,12 @@ pub mod reexport { pub use wl_output::{Transform, WlOutput}; } -type Frame = (Vec, (i32, i32)); - -/// Struct to store region capture details. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct CaptureRegion { - /// X coordinate of the area to capture. - pub x_coordinate: i32, - /// y coordinate of the area to capture. - pub y_coordinate: i32, - /// Width of the capture area. - pub width: i32, - /// Height of the capture area. - pub height: i32, -} - -#[derive(Debug)] -struct IntersectingOutput { - output: WlOutput, - region: CaptureRegion, - transform: Transform, -} - /// Struct to store wayland connection and globals list. /// # Example usage /// /// ``` -/// let wayshot_connection = WayshotConnection::new().unwrap(); -/// let image_buffer = wayshot_connection.screenshot_all().unwrap(); +/// let wayshot_connection = WayshotConnection::new()?; +/// let image_buffer = wayshot_connection.screenshot_all()?; /// ``` #[derive(Debug)] pub struct WayshotConnection { @@ -159,9 +149,9 @@ impl WayshotConnection { if state.outputs.is_empty() { tracing::error!("Compositor did not advertise any wl_output devices!"); - exit(1); + return Err(Error::NoOutputs); } - tracing::debug!("Outputs detected: {:#?}", state.outputs); + tracing::trace!("Outputs detected: {:#?}", state.outputs); self.output_infos = state.outputs; Ok(()) @@ -174,18 +164,21 @@ impl WayshotConnection { cursor_overlay: i32, output: &WlOutput, fd: T, - capture_region: Option, - ) -> Result { + 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_inner(state, event_queue, frame, frame_format, fd) + 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( &self, cursor_overlay: i32, output: &WlOutput, - capture_region: Option, + capture_region: Option, ) -> Result<( CaptureFrameState, EventQueue, @@ -216,15 +209,15 @@ impl WayshotConnection { } }; - // Capture output. - let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { + debug!("Capturing output..."); + let frame = if let Some(embedded_region) = capture_region { screencopy_manager.capture_output_region( cursor_overlay, output, - region.x_coordinate, - region.y_coordinate, - region.width, - region.height, + 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, (), ) @@ -238,7 +231,7 @@ impl WayshotConnection { event_queue.blocking_dispatch(&mut state)?; } - tracing::debug!( + tracing::trace!( "Received compositor frame buffer formats: {:#?}", state.formats ); @@ -258,7 +251,7 @@ impl WayshotConnection { ) }) .copied(); - tracing::debug!("Selected frame buffer format: {:#?}", frame_format); + tracing::trace!("Selected frame buffer format: {:#?}", frame_format); // Check if frame format exists. let frame_format = match frame_format { @@ -278,20 +271,25 @@ impl WayshotConnection { frame: ZwlrScreencopyFrameV1, frame_format: FrameFormat, fd: T, - ) -> Result { + ) -> Result { // Connecting to wayland environment. let qh = event_queue.handle(); - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - // 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 shm = self.globals.bind::(&qh, 1..=1, ())?; + let shm_pool = shm.create_pool( + fd.as_fd(), + frame_format + .byte_size() + .try_into() + .map_err(|_| Error::BufferTooSmall)?, + &qh, + (), + ); let buffer = shm_pool.create_buffer( 0, - frame_format.width as i32, - frame_format.height as i32, + frame_format.size.width as i32, + frame_format.size.height as i32, frame_format.stride as i32, frame_format.format, &qh, @@ -310,9 +308,7 @@ impl WayshotConnection { return Err(Error::FramecopyFailed); } FrameState::Finished => { - buffer.destroy(); - shm_pool.destroy(); - return Ok(frame_format); + return Ok(FrameGuard { buffer, shm_pool }); } } } @@ -326,34 +322,35 @@ impl WayshotConnection { cursor_overlay: bool, output: &WlOutput, file: &File, - capture_region: Option, - ) -> Result { + 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)?; - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - file.set_len(frame_bytes as u64)?; + file.set_len(frame_format.byte_size())?; + + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; - 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. - fn capture_output_frame( + #[tracing::instrument(skip_all, fields(output = format!("{output_info}"), region = capture_region.map(|r| format!("{:}", r))))] + fn capture_frame_copy( &self, cursor_overlay: bool, - output: &WlOutput, - transform: Transform, - capture_region: Option, - ) -> Result { + output_info: &OutputInfo, + capture_region: Option, + ) -> Result<(FrameCopy, FrameGuard)> { // 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 = File::from(fd); - let frame_format = self.capture_output_frame_shm_from_file( + let (frame_format, frame_guard) = self.capture_output_frame_shm_from_file( cursor_overlay, - output, + &output_info.wl_output, &mem_file, capture_region, )?; @@ -367,64 +364,53 @@ impl WayshotConnection { tracing::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); }; - Ok(FrameCopy { + let rotated_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_mmap, - transform, - }) + transform: output_info.transform, + region: LogicalRegion { + inner: Region { + position: capture_region + .map(|capture_region: EmbeddedRegion| { + capture_region.logical().inner.position + }) + .unwrap_or_else(|| output_info.region.position), + size: rotated_size, + }, + }, + }; + tracing::debug!("Created frame copy: {:#?}", frame_copy); + Ok((frame_copy, frame_guard)) } - fn create_frame_copy( + pub fn capture_frame_copies( &self, - capture_region: CaptureRegion, + output_capture_regions: &Vec<(OutputInfo, Option)>, cursor_overlay: bool, - ) -> Result { + ) -> Result> { let frame_copies = thread::scope(|scope| -> Result<_> { - let join_handles = self - .get_all_outputs() + let join_handles = output_capture_regions .into_iter() - .filter_map(|output| { - let x1: i32 = cmp::max(output.dimensions.x, capture_region.x_coordinate); - let y1: i32 = cmp::max(output.dimensions.y, capture_region.y_coordinate); - let x2: i32 = cmp::min( - output.dimensions.x + output.dimensions.width, - capture_region.x_coordinate + capture_region.width, - ); - let y2: i32 = cmp::min( - output.dimensions.y + output.dimensions.height, - capture_region.y_coordinate + capture_region.height, - ); - - let width = x2 - x1; - let height = y2 - y1; - - if width <= 0 || height <= 0 { - return None; - } - - let true_x = capture_region.x_coordinate - output.dimensions.x; - let true_y = capture_region.y_coordinate - output.dimensions.y; - let true_region = CaptureRegion { - x_coordinate: true_x, - y_coordinate: true_y, - width: capture_region.width, - height: capture_region.height, - }; - Some(IntersectingOutput { - output: output.wl_output.clone(), - region: true_region, - transform: output.transform, - }) - }) - .map(|intersecting_output| { + .map(|(output_info, capture_region)| { scope.spawn(move || { - self.capture_output_frame( + self.capture_frame_copy( cursor_overlay, - &intersecting_output.output, - intersecting_output.transform, - Some(intersecting_output.region), + &output_info, + capture_region.clone(), ) + .map(|(frame_copy, frame_guard)| { + (frame_copy, frame_guard, output_info.clone()) + }) }) }) .collect::>(); @@ -436,30 +422,160 @@ impl WayshotConnection { .collect::>() })?; - Ok((frame_copies, (capture_region.width, capture_region.height))) + Ok(frame_copies) + } + + fn overlay_frames(&self, frames: &Vec<(FrameCopy, FrameGuard, OutputInfo)>) -> Result<()> { + let mut state = LayerShellState { + configured_outputs: HashSet::new(), + }; + let mut event_queue: EventQueue = + self.conn.new_event_queue::(); + let qh = event_queue.handle(); + + let compositor = match self.globals.bind::(&qh, 3..=3, ()) { + Ok(x) => x, + Err(e) => { + tracing::error!( + "Failed to create compositor Does your compositor implement WlCompositor?" + ); + tracing::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "WlCompositor not found".to_string(), + )); + } + }; + let layer_shell = match self.globals.bind::(&qh, 1..=1, ()) { + Ok(x) => x, + Err(e) => { + tracing::error!( + "Failed to create layer shell. Does your compositor implement WlrLayerShellV1?" + ); + tracing::error!("err: {e}"); + return Err(Error::ProtocolNotFound( + "WlrLayerShellV1 not found".to_string(), + )); + } + }; + + for (frame_copy, frame_guard, output_info) in frames { + tracing::span!( + tracing::Level::DEBUG, + "overlay_frames::surface", + output = format!("{output_info}") + ) + .in_scope(|| -> Result<()> { + let surface = compositor.create_surface(&qh, ()); + + let layer_surface = layer_shell.get_layer_surface( + &surface, + Some(&output_info.wl_output), + Layer::Top, + "wayshot".to_string(), + &qh, + output_info.wl_output.clone(), + ); + + layer_surface.set_exclusive_zone(-1); + layer_surface.set_anchor(Anchor::Top | Anchor::Left); + layer_surface.set_size( + frame_copy.frame_format.size.width, + frame_copy.frame_format.size.height, + ); + + debug!("Committing surface creation changes."); + surface.commit(); + + debug!("Waiting for layer surface to be configured."); + while !state.configured_outputs.contains(&output_info.wl_output) { + event_queue.blocking_dispatch(&mut state)?; + } + + surface.set_buffer_transform(output_info.transform); + surface.set_buffer_scale(output_info.scale); + surface.attach(Some(&frame_guard.buffer), 0, 0); + + debug!("Committing surface with attached buffer."); + surface.commit(); + + event_queue.blocking_dispatch(&mut state)?; + + Ok(()) + })?; + } + Ok(()) } /// Take a screenshot from the specified region. - pub fn screenshot( + fn screenshot_region_capturer( &self, - capture_region: CaptureRegion, + region_capturer: RegionCapturer, cursor_overlay: bool, ) -> Result { - let (frame_copies, (width, height)) = - self.create_frame_copy(capture_region, cursor_overlay)?; + let outputs_capture_regions: &Vec<(OutputInfo, Option)> = + &match region_capturer { + RegionCapturer::Outputs(ref outputs) => outputs + .into_iter() + .map(|output_info| (output_info.clone(), None)) + .collect(), + RegionCapturer::Region(capture_region) => self + .get_all_outputs() + .into_iter() + .filter_map(|output_info| { + tracing::span!( + tracing::Level::DEBUG, + "filter_map", + output = format!( + "{output_info} at {region}", + output_info = format!("{output_info}"), + region = LogicalRegion::from(output_info), + ), + capture_region = format!("{}", capture_region), + ) + .in_scope(|| { + if let Some(relative_region) = + EmbeddedRegion::new(capture_region, output_info.into()) + { + tracing::debug!("Intersection found: {}", relative_region); + Some((output_info.clone(), Some(relative_region))) + } else { + tracing::debug!("No intersection found"); + None + } + }) + }) + .collect(), + RegionCapturer::Freeze(_) => self + .get_all_outputs() + .into_iter() + .map(|output_info| (output_info.clone(), None)) + .collect(), + }; + + let frames = self.capture_frame_copies(outputs_capture_regions, cursor_overlay)?; + + let capture_region: LogicalRegion = match region_capturer { + RegionCapturer::Outputs(ref outputs) => outputs.try_into()?, + RegionCapturer::Region(region) => region, + RegionCapturer::Freeze(callback) => { + self.overlay_frames(&frames).and_then(|_| callback())? + } + }; thread::scope(|scope| { - let rotate_join_handles = frame_copies + let rotate_join_handles = frames .into_iter() - .map(|frame_copy| { + .map(|(frame_copy, _, _)| { scope.spawn(move || { - let transform = frame_copy.transform; - let image = frame_copy.try_into()?; - Ok(image_util::rotate_image_buffer( - image, - transform, - width as u32, - height as u32, + let image = (&frame_copy).try_into()?; + Ok(( + image_util::rotate_image_buffer( + image, + frame_copy.transform, + frame_copy.frame_format.size.width, + frame_copy.frame_format.size.height, + ), + frame_copy, )) }) }) @@ -471,21 +587,39 @@ impl WayshotConnection { .flatten() .fold( None, - |possible_overlayed_image_or_error: Option>, image: Result<_>| { - if let Some(overlayed_image_or_error) = possible_overlayed_image_or_error { - if let Ok(mut overlayed_image) = overlayed_image_or_error { - if let Ok(image) = image { - overlay(&mut overlayed_image, &image, 0, 0); - Some(Ok(overlayed_image)) - } else { - Some(image) - } - } else { - Some(image) - } - } else { - Some(image) - } + |composite_image: Option>, image: Result<_>| { + // Default to a transparent image. + let composite_image = composite_image.unwrap_or_else(|| { + Ok(DynamicImage::new_rgba8( + capture_region.inner.size.width, + capture_region.inner.size.height, + )) + }); + + Some(|| -> Result<_> { + let mut composite_image = composite_image?; + let (image, frame_copy) = image?; + let (x, y) = ( + frame_copy.region.inner.position.x as i64 + - capture_region.inner.position.x as i64, + frame_copy.region.inner.position.y as i64 + - capture_region.inner.position.y as i64, + ); + tracing::span!( + tracing::Level::DEBUG, + "replace", + frame_copy_region = format!("{}", frame_copy.region), + capture_region = format!("{}", capture_region), + x = x, + y = y, + ) + .in_scope(|| { + tracing::debug!("Replacing parts of the final image"); + replace(&mut composite_image, &image, x, y); + }); + + Ok(composite_image) + }()) }, ) .ok_or_else(|| { @@ -495,19 +629,32 @@ impl WayshotConnection { }) } + /// Take a screenshot from the specified region. + pub fn screenshot( + &self, + capture_region: LogicalRegion, + cursor_overlay: bool, + ) -> Result { + self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay) + } + + /// Take a screenshot, overlay the screenshot, run the callback, and then + /// unfreeze the screenshot and return the selected region. + pub fn screenshot_freeze( + &self, + callback: Box Result>, + cursor_overlay: bool, + ) -> Result { + self.screenshot_region_capturer(RegionCapturer::Freeze(callback), cursor_overlay) + } /// shot one ouput pub fn screenshot_single_output( &self, output_info: &OutputInfo, cursor_overlay: bool, ) -> Result { - let frame_copy = self.capture_output_frame( - cursor_overlay, - &output_info.wl_output, - output_info.transform, - None, - )?; - frame_copy.try_into() + let (frame_copy, _) = self.capture_frame_copy(cursor_overlay, output_info, None)?; + (&frame_copy).try_into() } /// Take a screenshot from all of the specified outputs. @@ -520,33 +667,7 @@ impl WayshotConnection { return Err(Error::NoOutputs); } - let x1 = outputs - .iter() - .map(|output| output.dimensions.x) - .min() - .unwrap(); - let y1 = outputs - .iter() - .map(|output| output.dimensions.y) - .min() - .unwrap(); - let x2 = outputs - .iter() - .map(|output| output.dimensions.x + output.dimensions.width) - .max() - .unwrap(); - let y2 = outputs - .iter() - .map(|output| output.dimensions.y + output.dimensions.height) - .max() - .unwrap(); - let capture_region = CaptureRegion { - x_coordinate: x1, - y_coordinate: y1, - width: x2 - x1, - height: y2 - y1, - }; - self.screenshot(capture_region, cursor_overlay) + self.screenshot_region_capturer(RegionCapturer::Outputs(outputs.clone()), cursor_overlay) } /// Take a screenshot from all accessible outputs. diff --git a/libwayshot/src/output.rs b/libwayshot/src/output.rs index ccca1c80..5449f4f7 100644 --- a/libwayshot/src/output.rs +++ b/libwayshot/src/output.rs @@ -1,28 +1,29 @@ +use std::fmt::Display; + use wayland_client::protocol::{wl_output, wl_output::WlOutput}; +use crate::region::Region; + /// Represents an accessible wayland output. /// /// Do not instantiate, instead use [`crate::WayshotConnection::get_all_outputs`]. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct OutputInfo { pub wl_output: WlOutput, pub name: String, pub description: String, pub transform: wl_output::Transform, - pub dimensions: OutputPositioning, - pub mode: WlOutputMode, -} - -#[derive(Default, Debug, Clone, PartialEq, Eq)] -pub struct WlOutputMode { - pub width: i32, - pub height: i32, + pub scale: i32, + pub region: Region, } -#[derive(Default, Debug, Clone, PartialEq, Eq)] -pub struct OutputPositioning { - pub x: i32, - pub y: i32, - pub width: i32, - pub height: i32, +impl Display for OutputInfo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{name} ({description})", + name = self.name, + description = self.description + ) + } } diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs new file mode 100644 index 00000000..bbd81d39 --- /dev/null +++ b/libwayshot/src/region.rs @@ -0,0 +1,229 @@ +use std::cmp; + +use crate::error::{Error, Result}; +use crate::output::OutputInfo; + +/// Ways to say how a region for a screenshot should be captured. +pub enum RegionCapturer { + /// Capture all of the given outputs. + Outputs(Vec), + /// Capture an already known `LogicalRegion`. + Region(LogicalRegion), + /// The outputs will be "frozen" to the user at which point the given + /// callback is called to get the region to capture. This callback is often + /// a user interaction to let the user select a region. + Freeze(Box Result>), +} + +/// `Region` where the coordinate system is the logical coordinate system used +/// in Wayland to position outputs. Top left is (0, 0) and any transforms and +/// scaling have been applied. +#[derive(Debug, Copy, Clone)] +pub struct LogicalRegion { + pub inner: Region, +} + +/// An embedded region is a region entirely inside of another (often an output). +/// +/// It can only be contained inside of another and cannot exceed its bounds. +/// +/// Example of what +/// +/// ┌─────────────┐ +/// │ │ +/// │ ┌──────────┼──────┐ +/// │ │ │ ├──► Viewport +/// │ │ │ │ +/// │ │ ├──────┼─────────────────┐ +/// │ │ │xxxxxx│ │ +/// │ │ │xxxxx◄├─── Inner region │ +/// │ └──────────┼──────┘ │ +/// │ │ │ +/// │ │ Screen 2 ├──► Relative to +/// │ ├────────────────────────┘ +/// │ │ +/// │ Screen 1 │ +/// └─────────────┘ +#[derive(Debug, Copy, Clone)] +pub struct EmbeddedRegion { + /// The coordinate sysd + pub relative_to: LogicalRegion, + pub inner: Region, +} + +/// Rectangle area in an unspecified coordinate system. +/// +/// Use `LogicalRegion` or `EmbeddedRegion` instead as they convey the +/// coordinate system used. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +pub struct Region { + /// Position of the region. + pub position: Position, + /// Size of the region. + pub size: Size, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +pub struct Position { + /// X coordinate. + pub x: i32, + /// Y coordinate. + pub y: i32, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash, Default)] +pub struct Size { + /// Width. + pub width: u32, + /// Height. + pub height: u32, +} + +impl EmbeddedRegion { + /// Given two `LogicalRegion`s, one seen as the `viewport` and the other + /// `relative_to` (think the output we want to capture), create an + /// embedded region that is entirely inside of the `relative_to` region. + /// + /// See `EmbeddedRegion` for an example ASCII visualisation. + #[tracing::instrument(ret, level = "debug")] + pub fn new(viewport: LogicalRegion, relative_to: LogicalRegion) -> Option { + let x_relative: i32 = viewport.inner.position.x - relative_to.inner.position.x; + let y_relative = viewport.inner.position.y - relative_to.inner.position.y; + + let x1 = cmp::max(x_relative, 0); + let x2 = cmp::min( + x_relative + viewport.inner.size.width as i32, + relative_to.inner.size.width as i32, + ); + let width = if let Ok(width) = (x2 - x1).try_into() { + width + } else { + return None; + }; + + let y1 = cmp::max(y_relative, 0); + let y2 = cmp::min( + y_relative + viewport.inner.size.height as i32, + relative_to.inner.size.height as i32, + ); + let height = if let Ok(height) = (y2 - y1).try_into() { + height + } else { + return None; + }; + + Some(Self { + relative_to: relative_to, + inner: Region { + position: Position { x: x1, y: y1 }, + size: Size { width, height }, + }, + }) + } + + /// Return the `LogicalRegion` of the embedded region. + /// + /// Note that this remains a region of the same size, it's not the inverse + /// of `EmbeddedRegion::new` which removes the parts that are outside of + /// the `relative_to` region. + pub fn logical(&self) -> LogicalRegion { + LogicalRegion { + inner: Region { + position: Position { + x: self.relative_to.inner.position.x + self.inner.position.x, + y: self.relative_to.inner.position.y + self.inner.position.y, + }, + size: self.inner.size, + }, + } + } +} + +impl std::fmt::Display for EmbeddedRegion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{region} relative to {relative_to}", + region = self.inner, + relative_to = self.relative_to, + ) + } +} + +impl std::fmt::Display for Position { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({x}, {y})", x = self.x, y = self.y,) + } +} + +impl std::fmt::Display for Size { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "({width}x{height})", + width = self.width, + height = self.height, + ) + } +} + +impl std::fmt::Display for Region { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "({position}) ({size})", + position = self.position, + size = self.size, + ) + } +} + +impl std::fmt::Display for LogicalRegion { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{inner}", inner = self.inner) + } +} + +impl From<&OutputInfo> for LogicalRegion { + fn from(output_info: &OutputInfo) -> Self { + LogicalRegion { + inner: output_info.region, + } + } +} + +impl TryFrom<&Vec> for LogicalRegion { + type Error = Error; + + fn try_from(output_info: &Vec) -> std::result::Result { + let x1 = output_info + .iter() + .map(|output| output.region.position.x) + .min() + .ok_or(Error::NoOutputs)?; + let y1 = output_info + .iter() + .map(|output| output.region.position.y) + .min() + .ok_or(Error::NoOutputs)?; + let x2 = output_info + .iter() + .map(|output| output.region.position.x + output.region.size.width as i32) + .max() + .ok_or(Error::NoOutputs)?; + let y2 = output_info + .iter() + .map(|output| output.region.position.y + output.region.size.height as i32) + .max() + .ok_or(Error::NoOutputs)?; + Ok(LogicalRegion { + inner: Region { + position: Position { x: x1, y: y1 }, + size: Size { + width: (x2 - x1) as u32, + height: (y2 - y1) as u32, + }, + }, + }) + } +} diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index e2713582..734b6bfd 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -1,5 +1,5 @@ use std::{ - ffi::CStr, + ffi::CString, os::fd::{AsRawFd, IntoRawFd, OwnedFd}, time::{SystemTime, UNIX_EPOCH}, }; @@ -11,20 +11,49 @@ use nix::{ sys::{memfd, mman, stat}, unistd, }; -use wayland_client::protocol::{wl_output, wl_shm::Format}; +use wayland_client::protocol::{ + wl_buffer::WlBuffer, wl_output, wl_shm::Format, wl_shm_pool::WlShmPool, +}; + +use crate::{ + region::{LogicalRegion, Size}, + Error, Result, +}; + +pub struct FrameGuard { + pub buffer: WlBuffer, + pub shm_pool: WlShmPool, +} -use crate::{Error, Result}; +impl Drop for FrameGuard { + fn drop(&mut self) { + self.buffer.destroy(); + self.shm_pool.destroy(); + } +} /// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and /// Xbgr8888. +/// +/// See `zwlr_screencopy_frame_v1::Event::Buffer` as it's retrieved from there. #[derive(Debug, Copy, Clone, PartialEq)] pub struct FrameFormat { pub format: Format, - pub width: u32, - pub height: 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, + /// Stride is the number of bytes between the start of a row and the start of the next row. pub stride: u32, } +impl FrameFormat { + /// Returns the size of the frame in bytes, which is the stride * height. + pub fn byte_size(&self) -> u64 { + self.stride as u64 * self.size.height as u64 + } +} + +#[tracing::instrument(skip(frame_mmap))] fn create_image_buffer

( frame_format: &FrameFormat, frame_mmap: &MmapMut, @@ -32,8 +61,13 @@ fn create_image_buffer

( where P: Pixel, { - ImageBuffer::from_vec(frame_format.width, frame_format.height, frame_mmap.to_vec()) - .ok_or(Error::BufferTooSmall) + tracing::debug!("Creating image buffer"); + ImageBuffer::from_vec( + frame_format.size.width, + frame_format.size.height, + frame_mmap.to_vec(), + ) + .ok_or(Error::BufferTooSmall) } /// The copied frame comprising of the FrameFormat, ColorType (Rgba8), and a memory backed shm @@ -44,12 +78,14 @@ pub struct FrameCopy { pub frame_color_type: ColorType, pub frame_mmap: MmapMut, pub transform: wl_output::Transform, + /// Logical region with the transform already applied. + pub region: LogicalRegion, } -impl TryFrom for DynamicImage { +impl TryFrom<&FrameCopy> for DynamicImage { type Error = Error; - fn try_from(value: FrameCopy) -> Result { + 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)?) @@ -62,6 +98,16 @@ impl TryFrom for DynamicImage { } } +fn get_mem_file_handle() -> String { + format!( + "/libwayshot-{}", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|time| time.subsec_nanos().to_string()) + .unwrap_or("unknown".into()) + ) +} + /// Return a RawFd to a shm file. We use memfd create on linux and shm_open for BSD support. /// You don't need to mess around with this function, it is only used by /// capture_output_frame. @@ -71,7 +117,7 @@ pub fn create_shm_fd() -> std::io::Result { loop { // Create a file that closes on succesful execution and seal it's operations. match memfd::memfd_create( - CStr::from_bytes_with_nul(b"libwayshot\0").unwrap(), + CString::new("libwayshot")?.as_c_str(), memfd::MemFdCreateFlag::MFD_CLOEXEC | memfd::MemFdCreateFlag::MFD_ALLOW_SEALING, ) { Ok(fd) => { @@ -93,11 +139,7 @@ pub fn create_shm_fd() -> std::io::Result { } // Fallback to using shm_open. - let sys_time = SystemTime::now(); - let mut mem_file_handle = format!( - "/libwayshot-{}", - sys_time.duration_since(UNIX_EPOCH).unwrap().subsec_nanos() - ); + let mut mem_file_handle = get_mem_file_handle(); loop { match mman::shm_open( // O_CREAT = Create file if does not exist. @@ -122,10 +164,7 @@ pub fn create_shm_fd() -> std::io::Result { }, Err(nix::errno::Errno::EEXIST) => { // If a file with that handle exists then change the handle - mem_file_handle = format!( - "/libwayshot-{}", - sys_time.duration_since(UNIX_EPOCH).unwrap().subsec_nanos() - ); + mem_file_handle = get_mem_file_handle(); continue; } Err(nix::errno::Errno::EINTR) => continue, diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 8d4aa528..4594d0b4 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -12,13 +12,15 @@ repository.workspace = true [build-dependencies] flate2 = "1.0.27" +eyre = "0.6.8" + [dependencies] tracing.workspace = true libwayshot.workspace = true -clap = "4.4.6" +clap = { version = "4.4.18", features = ["derive"] } tracing-subscriber = "0.3.17" image = { version = "0.24", default-features = false, features = [ @@ -29,6 +31,7 @@ image = { version = "0.24", default-features = false, features = [ ] } dialoguer = { version = "0.11.0", features = ["fuzzy-select"] } +eyre = "0.6.8" [[bin]] name = "wayshot" diff --git a/wayshot/build.rs b/wayshot/build.rs index 14152acb..37104611 100644 --- a/wayshot/build.rs +++ b/wayshot/build.rs @@ -1,13 +1,14 @@ extern crate flate2; +use eyre::{ContextCompat, Result}; use flate2::{write::GzEncoder, Compression}; use std::{ fs::{read_dir, File, OpenOptions}, io::{copy, BufReader, ErrorKind}, path::Path, - process::{exit, Command, Stdio}, + process::{Command, Stdio}, }; -fn main() { +fn main() -> Result<()> { if let Err(e) = Command::new("scdoc") .stdin(Stdio::null()) .stdout(Stdio::null()) @@ -15,50 +16,60 @@ fn main() { .spawn() { if let ErrorKind::NotFound = e.kind() { - exit(0); + return Ok(()); } } // We just append "out" so it's easy to find all the scdoc output later in line 38. - let man_pages: Vec<(String, String)> = read_and_replace_by_ext("./docs", ".scd", ".out"); + let man_pages: Vec<(String, String)> = read_and_replace_by_ext("./docs", ".scd", ".out")?; for man_page in man_pages { let output = OpenOptions::new() .write(true) .create(true) - .open(Path::new(&man_page.1)) - .unwrap(); + .open(Path::new(&man_page.1))?; _ = Command::new("scdoc") - .stdin(Stdio::from(File::open(man_page.0).unwrap())) + .stdin(Stdio::from(File::open(man_page.0)?)) .stdout(output) .spawn(); } // Gzipping the man pages let scdoc_output_files: Vec<(String, String)> = - read_and_replace_by_ext("./docs", ".out", ".gz"); + read_and_replace_by_ext("./docs", ".out", ".gz")?; for scdoc_output in scdoc_output_files { - let mut input = BufReader::new(File::open(scdoc_output.0).unwrap()); + let mut input = BufReader::new(File::open(scdoc_output.0)?); let output = OpenOptions::new() .write(true) .create(true) - .open(Path::new(&scdoc_output.1)) - .unwrap(); + .open(Path::new(&scdoc_output.1))?; let mut encoder = GzEncoder::new(output, Compression::default()); - copy(&mut input, &mut encoder).unwrap(); - encoder.finish().unwrap(); + copy(&mut input, &mut encoder)?; + encoder.finish()?; } + + Ok(()) } -fn read_and_replace_by_ext(path: &str, search: &str, replace: &str) -> Vec<(String, String)> { +fn read_and_replace_by_ext( + path: &str, + search: &str, + replace: &str, +) -> Result> { let mut files: Vec<(String, String)> = Vec::new(); - for path in read_dir(path).unwrap() { - let path = path.unwrap(); - if path.file_type().unwrap().is_dir() { + for path in read_dir(path)? { + let path = path?; + if path.file_type()?.is_dir() { continue; } if let Some(file_name) = path.path().to_str() { - if *path.path().extension().unwrap().to_str().unwrap() != search[1..] { + if *path + .path() + .extension() + .wrap_err_with(|| format!("no extension found for {}", path.path().display()))? + .to_string_lossy() + != search[1..] + { continue; } @@ -66,5 +77,5 @@ fn read_and_replace_by_ext(path: &str, search: &str, replace: &str) -> Vec<(Stri files.push((file_name.to_string(), file)); } } - files + Ok(files) } diff --git a/wayshot/src/clap.rs b/wayshot/src/clap.rs deleted file mode 100644 index 37319334..00000000 --- a/wayshot/src/clap.rs +++ /dev/null @@ -1,67 +0,0 @@ -use clap::{arg, ArgAction, Command}; - -pub fn set_flags() -> Command { - Command::new("wayshot") - .version(env!("CARGO_PKG_VERSION")) - .author(env!("CARGO_PKG_AUTHORS")) - .about("Screenshot tool for compositors implementing zwlr_screencopy_v1.") - .arg( - arg!(-d - -debug) - .required(false) - .action(ArgAction::SetTrue) - .help("Enable debug mode"), - ) - .arg( - arg!(-s --slurp ) - .required(false) - .action(ArgAction::Set) - .help("Choose a portion of your display to screenshot using slurp"), - ) - .arg( - arg!(-f - -file ) - .required(false) - .conflicts_with("stdout") - .action(ArgAction::Set) - .help("Mention a custom file path"), - ) - .arg( - arg!(-c - -cursor) - .required(false) - .action(ArgAction::SetTrue) - .help("Enable cursor in screenshots"), - ) - .arg( - arg!(--stdout) - .required(false) - .conflicts_with("file") - .action(ArgAction::SetTrue) - .help("Output the image data to standard out"), - ) - .arg( - arg!(-e --extension ) - .required(false) - .action(ArgAction::Set) - .help("Set image encoder (Png is default)"), - ) - .arg( - arg!(-l - -listoutputs) - .required(false) - .action(ArgAction::SetTrue) - .help("List all valid outputs"), - ) - .arg( - arg!(-o --output ) - .required(false) - .action(ArgAction::Set) - .conflicts_with("slurp") - .help("Choose a particular display to screenshot"), - ) - .arg( - arg!(--chooseoutput) - .required(false) - .action(ArgAction::SetTrue) - .conflicts_with("slurp") - .conflicts_with("output") - .help("Present a fuzzy selector for outputs"), - ) -} diff --git a/wayshot/src/cli.rs b/wayshot/src/cli.rs new file mode 100644 index 00000000..418ac8d8 --- /dev/null +++ b/wayshot/src/cli.rs @@ -0,0 +1,46 @@ +use std::path::PathBuf; + +use clap::arg; + +use clap::Parser; +use eyre::WrapErr; + +use crate::utils::EncodingFormat; +use clap::builder::TypedValueParser; + +#[derive(Parser)] +#[command(version, about)] +pub struct Cli { + /// Where to save the screenshot, "-" for stdout. Defaults to "$UNIX_TIMESTAMP-wayshot.$EXTENSION". + #[arg(value_name = "OUTPUT")] + pub file: Option, + + /// Log level to be used for printing to stderr + #[arg(long, default_value = "info", value_parser = clap::builder::PossibleValuesParser::new(["trace", "debug", "info", "warn", "error"]).map(|s| -> tracing::Level{ s.parse().wrap_err_with(|| format!("Failed to parse log level: {}", s)).unwrap()}))] + pub log_level: tracing::Level, + + /// Arguments to call slurp with for selecting a region + #[arg(short, long, value_name = "SLURP_ARGS")] + pub slurp: Option, + + /// Enable cursor in screenshots + #[arg(short, long)] + pub cursor: bool, + + /// Set image encoder, by default uses the file extension from the OUTPUT + /// positional argument. Otherwise defaults to png. + #[arg(long, visible_aliases = ["extension", "format", "output-format"], value_name = "FILE_EXTENSION")] + pub encoding: Option, + + /// List all valid outputs + #[arg(short, long, alias = "listoutputs")] + pub list_outputs: bool, + + /// Choose a particular output/display to screenshot + #[arg(short, long, conflicts_with = "slurp")] + pub output: Option, + + /// Present a fuzzy selector for output/display selection + #[arg(long, alias = "chooseoutput", conflicts_with_all = ["slurp", "output"])] + pub choose_output: bool, +} diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index e24caf74..117c1f0c 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -1,58 +1,75 @@ +use clap::ValueEnum; +use eyre::{bail, ContextCompat, Error, Result}; + use std::{ - process::exit, + fmt::Display, + path::PathBuf, + str::FromStr, time::{SystemTime, UNIX_EPOCH}, }; -use libwayshot::CaptureRegion; +use libwayshot::region::{LogicalRegion, Position, Region, Size}; -pub fn parse_geometry(g: &str) -> Option { +pub fn parse_geometry(g: &str) -> Result { let tail = g.trim(); let x_coordinate: i32; let y_coordinate: i32; - let width: i32; - let height: i32; + let width: u32; + let height: u32; + + let validation_error = + "Invalid geometry provided.\nValid geometries:\n1) %d,%d %dx%d\n2) %d %d %d %d"; if tail.contains(',') { // this accepts: "%d,%d %dx%d" - let (head, tail) = tail.split_once(',')?; - x_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - y_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once('x')?; - width = head.parse::().ok()?; - height = tail.parse::().ok()?; + let (head, tail) = tail.split_once(',').wrap_err(validation_error)?; + x_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + y_coordinate = head.parse::()?; + let (head, tail) = tail.split_once('x').wrap_err(validation_error)?; + width = head.parse::()?; + height = tail.parse::()?; } else { // this accepts: "%d %d %d %d" - let (head, tail) = tail.split_once(' ')?; - x_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - y_coordinate = head.parse::().ok()?; - let (head, tail) = tail.split_once(' ')?; - width = head.parse::().ok()?; - height = tail.parse::().ok()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + x_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + y_coordinate = head.parse::()?; + let (head, tail) = tail.split_once(' ').wrap_err(validation_error)?; + width = head.parse::()?; + height = tail.parse::()?; } - Some(CaptureRegion { - x_coordinate, - y_coordinate, - width, - height, + Ok(LogicalRegion { + inner: Region { + position: Position { + x: x_coordinate, + y: y_coordinate, + }, + size: Size { width, height }, + }, }) } /// Supported image encoding formats. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, ValueEnum)] pub enum EncodingFormat { - /// Jpeg / jpg encoder. + /// JPG/JPEG encoder. Jpg, - /// Png encoder. + /// PNG encoder. Png, - /// Ppm encoder. + /// PPM encoder. Ppm, - /// Qoi encoder. + /// Qut encoder. Qoi, } +impl Default for EncodingFormat { + fn default() -> Self { + Self::Png + } +} + impl From for image::ImageOutputFormat { fn from(format: EncodingFormat) -> Self { match format { @@ -64,6 +81,33 @@ impl From for image::ImageOutputFormat { } } +impl TryFrom<&PathBuf> for EncodingFormat { + type Error = Error; + + fn try_from(value: &PathBuf) -> std::result::Result { + value + .extension() + .wrap_err_with(|| { + format!( + "no extension in {} to deduce encoding format", + value.display() + ) + }) + .and_then(|ext| { + ext.to_str().wrap_err_with(|| { + format!("extension in {} is not valid unicode", value.display()) + }) + }) + .and_then(|ext| ext.parse()) + } +} + +impl Display for EncodingFormat { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", Into::<&str>::into(*self)) + } +} + impl From for &str { fn from(format: EncodingFormat) -> Self { match format { @@ -75,14 +119,25 @@ impl From for &str { } } -pub fn get_default_file_name(extension: EncodingFormat) -> String { - let time = match SystemTime::now().duration_since(UNIX_EPOCH) { - Ok(n) => n.as_secs().to_string(), - Err(_) => { - tracing::error!("SystemTime before UNIX EPOCH!"); - exit(1); - } - }; +impl FromStr for EncodingFormat { + type Err = Error; + + fn from_str(s: &str) -> std::result::Result { + Ok(match s { + "jpg" | "jpeg" => Self::Jpg, + "png" => Self::Png, + "ppm" => Self::Ppm, + "qoi" => Self::Qoi, + _ => bail!("unsupported extension '{s}'"), + }) + } +} + +pub fn get_default_file_name(extension: EncodingFormat) -> PathBuf { + let time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|time| time.as_secs().to_string()) + .unwrap_or("unknown".into()); - time + "-wayshot." + extension.into() + format!("{time}-wayshot.{extension}").into() } diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b4b3c01f..37152a79 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,18 +1,17 @@ use std::{ - error::Error, io::{stdout, BufWriter, Cursor, Write}, - process::exit, + process::Command, }; -use libwayshot::WayshotConnection; +use clap::Parser; +use eyre::{bail, Result}; +use libwayshot::{region::LogicalRegion, WayshotConnection}; -mod clap; +mod cli; mod utils; use dialoguer::{theme::ColorfulTheme, FuzzySelect}; -use tracing::Level; - -use crate::utils::EncodingFormat; +use utils::EncodingFormat; fn select_ouput(ouputs: &[T]) -> Option where @@ -29,103 +28,99 @@ where Some(selection) } -fn main() -> Result<(), Box> { - let args = clap::set_flags().get_matches(); - let level = if args.get_flag("debug") { - Level::TRACE - } else { - Level::INFO - }; +fn main() -> Result<()> { + let cli = cli::Cli::parse(); tracing_subscriber::fmt() - .with_max_level(level) + .with_max_level(cli.log_level) .with_writer(std::io::stderr) .init(); - let extension = if let Some(extension) = args.get_one::("extension") { - let ext = extension.trim().to_lowercase(); - tracing::debug!("Using custom extension: {:#?}", ext); - - match ext.as_str() { - "jpeg" | "jpg" => EncodingFormat::Jpg, - "png" => EncodingFormat::Png, - "ppm" => EncodingFormat::Ppm, - "qoi" => EncodingFormat::Qoi, - _ => { - tracing::error!("Invalid extension provided.\nValid extensions:\n1) jpeg\n2) jpg\n3) png\n4) ppm\n5) qoi"); - exit(1); + let input_encoding = cli + .file + .as_ref() + .and_then(|pathbuf| pathbuf.try_into().ok()); + let requested_encoding = cli + .encoding + .or(input_encoding) + .unwrap_or(EncodingFormat::default()); + + if let Some(input_encoding) = input_encoding { + if input_encoding != requested_encoding { + tracing::warn!( + "The encoding requested '{requested_encoding}' does not match the output file's encoding '{input_encoding}'. Still using the requested encoding however.", + ); + } + } + + let file = match cli.file { + Some(pathbuf) => { + if pathbuf.to_string_lossy() == "-" { + None + } else { + Some(pathbuf) } } - } else { - EncodingFormat::Png + None => Some(utils::get_default_file_name(requested_encoding)), }; - let mut file_is_stdout: bool = false; - let mut file_path: Option = None; - - if args.get_flag("stdout") { - file_is_stdout = true; - } else if let Some(filepath) = args.get_one::("file") { - file_path = Some(filepath.trim().to_string()); - } else { - file_path = Some(utils::get_default_file_name(extension)); - } - let wayshot_conn = WayshotConnection::new()?; - if args.get_flag("listoutputs") { + if cli.list_outputs { let valid_outputs = wayshot_conn.get_all_outputs(); for output in valid_outputs { tracing::info!("{:#?}", output.name); } - exit(1); - } - - let mut cursor_overlay = false; - if args.get_flag("cursor") { - cursor_overlay = true; + return Ok(()); } - let image_buffer = if let Some(slurp_region) = args.get_one::("slurp") { - if let Some(region) = utils::parse_geometry(slurp_region) { - wayshot_conn.screenshot(region, cursor_overlay)? - } else { - tracing::error!("Invalid geometry specification"); - exit(1); - } - } else if let Some(output_name) = args.get_one::("output") { + let image_buffer = if let Some(slurp_region) = cli.slurp { + let slurp_region = slurp_region.clone(); + wayshot_conn.screenshot_freeze( + Box::new(move || { + || -> Result { + let slurp_output = Command::new("slurp") + .args(slurp_region.split(" ")) + .output()? + .stdout; + + utils::parse_geometry(&String::from_utf8(slurp_output)?) + }() + .map_err(|_| libwayshot::Error::FreezeCallbackError) + }), + cli.cursor, + )? + } else if let Some(output_name) = cli.output { let outputs = wayshot_conn.get_all_outputs(); - if let Some(output) = outputs.iter().find(|output| &output.name == output_name) { - wayshot_conn.screenshot_single_output(output, cursor_overlay)? + if let Some(output) = outputs.iter().find(|output| output.name == output_name) { + wayshot_conn.screenshot_single_output(output, cli.cursor)? } else { - tracing::error!("No output found!\n"); - exit(1); + bail!("No output found!"); } - } else if args.get_flag("chooseoutput") { + } else if cli.choose_output { let outputs = wayshot_conn.get_all_outputs(); let output_names: Vec = outputs .iter() .map(|display| display.name.to_string()) .collect(); if let Some(index) = select_ouput(&output_names) { - wayshot_conn.screenshot_single_output(&outputs[index], cursor_overlay)? + wayshot_conn.screenshot_single_output(&outputs[index], cli.cursor)? } else { - tracing::error!("No output found!\n"); - exit(1); + bail!("No output found!"); } } else { - wayshot_conn.screenshot_all(cursor_overlay)? + wayshot_conn.screenshot_all(cli.cursor)? }; - if file_is_stdout { + if let Some(file) = file { + image_buffer.save(file)?; + } else { let stdout = stdout(); let mut buffer = Cursor::new(Vec::new()); let mut writer = BufWriter::new(stdout.lock()); - image_buffer.write_to(&mut buffer, extension)?; + image_buffer.write_to(&mut buffer, requested_encoding)?; writer.write_all(buffer.get_ref())?; - } else { - image_buffer.save(file_path.unwrap())?; } Ok(())