diff --git a/Cargo.lock b/Cargo.lock index 2d5bff5d..500045a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -207,6 +207,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" @@ -257,6 +267,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 +772,7 @@ version = "1.3.2-dev" dependencies = [ "clap", "dialoguer", + "eyre", "flate2", "image", "libwayshot", diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 170e1026..f4eedb06 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashSet, process::exit, sync::atomic::{AtomicBool, Ordering}, }; @@ -6,8 +7,9 @@ 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,6 +17,10 @@ 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, @@ -58,16 +64,9 @@ 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, + dimensions: OutputPositioning::default(), + mode: WlOutputMode::default(), }); } else { tracing::error!("Ignoring a wl_output with version < 4."); @@ -109,7 +108,11 @@ impl Dispatch for OutputCaptureState { } => { output.transform = transform; } - _ => (), + wl_output::Event::Scale { factor } => { + output.scale = factor; + } + wl_output::Event::Done => {} + _ => {} } } } @@ -227,3 +230,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..ab0e1ec3 100644 --- a/libwayshot/src/error.rs +++ b/libwayshot/src/error.rs @@ -27,4 +27,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/lib.rs b/libwayshot/src/lib.rs index c93cb335..2128ed55 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -12,6 +12,7 @@ pub mod region; mod screencopy; use std::{ + collections::HashSet, fs::File, os::fd::AsFd, process::exit, @@ -19,6 +20,7 @@ use std::{ thread, }; +use dispatch::LayerShellState; use image::{imageops::replace, DynamicImage}; use memmap2::MmapMut; use region::{EmbeddedRegion, RegionCapturer}; @@ -27,6 +29,7 @@ use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ + wl_compositor::WlCompositor, wl_output::WlOutput, wl_shm::{self, WlShm}, }, @@ -35,9 +38,15 @@ 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::{ @@ -406,6 +415,87 @@ impl WayshotConnection { 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 = output_info.name.as_str() + ) + .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.width, + frame_copy.frame_format.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. fn screenshot_region_capturer( &self, @@ -445,6 +535,11 @@ impl WayshotConnection { }) }) .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)?; @@ -452,6 +547,9 @@ impl WayshotConnection { 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| { @@ -529,6 +627,15 @@ impl WayshotConnection { 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, diff --git a/libwayshot/src/output.rs b/libwayshot/src/output.rs index ccca1c80..986f74df 100644 --- a/libwayshot/src/output.rs +++ b/libwayshot/src/output.rs @@ -3,23 +3,24 @@ use wayland_client::protocol::{wl_output, wl_output::WlOutput}; /// 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 scale: i32, pub dimensions: OutputPositioning, pub mode: WlOutputMode, } -#[derive(Default, Debug, Clone, PartialEq, Eq)] +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] pub struct WlOutputMode { pub width: i32, pub height: i32, } -#[derive(Default, Debug, Clone, PartialEq, Eq)] +#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)] pub struct OutputPositioning { pub x: i32, pub y: i32, diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs index a82e0cdc..490e6082 100644 --- a/libwayshot/src/region.rs +++ b/libwayshot/src/region.rs @@ -2,7 +2,7 @@ use std::cmp; use wayland_client::protocol::wl_output::Transform; -use crate::error::Error; +use crate::error::{Error, Result}; use crate::output::OutputInfo; use crate::screencopy::FrameCopy; @@ -12,6 +12,10 @@ pub enum RegionCapturer { 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 diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 8d4aa528..e4ac7af3 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -29,6 +29,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/src/clap.rs b/wayshot/src/clap.rs index 37319334..dfed67b4 100644 --- a/wayshot/src/clap.rs +++ b/wayshot/src/clap.rs @@ -12,10 +12,10 @@ pub fn set_flags() -> Command { .help("Enable debug mode"), ) .arg( - arg!(-s --slurp ) + arg!(-s --slurp ) .required(false) .action(ArgAction::Set) - .help("Choose a portion of your display to screenshot using slurp"), + .help("Arguments to call slurp with for selecting a region"), ) .arg( arg!(-f - -file ) diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index c6cd60fa..cea4217d 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -1,3 +1,4 @@ +use eyre::{ContextCompat, Result}; use std::{ process::exit, time::{SystemTime, UNIX_EPOCH}, @@ -5,34 +6,37 @@ use std::{ use libwayshot::region::{LogicalRegion, Region}; -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 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(LogicalRegion { + Ok(LogicalRegion { inner: Region { x: x_coordinate, y: y_coordinate, diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b4b3c01f..4353d579 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,10 +1,10 @@ use std::{ - error::Error, io::{stdout, BufWriter, Cursor, Write}, - process::exit, + process::{exit, Command}, }; -use libwayshot::WayshotConnection; +use eyre::Result; +use libwayshot::{region::LogicalRegion, WayshotConnection}; mod clap; mod utils; @@ -29,7 +29,7 @@ where Some(selection) } -fn main() -> Result<(), Box> { +fn main() -> Result<()> { let args = clap::set_flags().get_matches(); let level = if args.get_flag("debug") { Level::TRACE @@ -86,12 +86,21 @@ fn main() -> Result<(), Box> { } 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); - } + 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) + }), + cursor_overlay, + )? } else if let Some(output_name) = args.get_one::("output") { let outputs = wayshot_conn.get_all_outputs(); if let Some(output) = outputs.iter().find(|output| &output.name == output_name) {