From 8e480ae4cb520e81cc7966cb5ebed37d3483f824 Mon Sep 17 00:00:00 2001 From: Charles Johnson Date: Fri, 28 Jul 2023 21:34:14 +0100 Subject: [PATCH] [libwayshot] Mutli-monitor handling built into libwayshot Signed-off-by: c-h-johnson Signed-off-by: decodetalkers Signed-off-by: Shinyzenith --- Cargo.lock | 34 +- libwayshot/Cargo.toml | 3 +- libwayshot/README.md | 12 + libwayshot/src/dispatch.rs | 223 +++++++ libwayshot/src/error.rs | 24 + {wayshot => libwayshot}/src/image_util.rs | 0 libwayshot/src/lib.rs | 768 ++++++++++------------ libwayshot/src/output.rs | 20 + libwayshot/src/screencopy.rs | 137 ++++ wayshot/Cargo.toml | 5 +- wayshot/src/output.rs | 174 ----- wayshot/src/utils.rs | 43 +- wayshot/src/wayshot.rs | 218 +----- 13 files changed, 884 insertions(+), 777 deletions(-) create mode 100644 libwayshot/src/dispatch.rs create mode 100644 libwayshot/src/error.rs rename {wayshot => libwayshot}/src/image_util.rs (100%) create mode 100644 libwayshot/src/output.rs create mode 100644 libwayshot/src/screencopy.rs delete mode 100644 wayshot/src/output.rs diff --git a/Cargo.lock b/Cargo.lock index 69a9f6f2..7e1b7d49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -352,6 +352,7 @@ dependencies = [ "log", "memmap2", "nix", + "thiserror", "wayland-client", "wayland-protocols", "wayland-protocols-wlr", @@ -592,6 +593,17 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" +[[package]] +name = "syn" +version = "2.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04361975b3f5e348b2189d8dc55bc942f278b2d482a6a0365de5bdd62d351567" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "tempfile" version = "3.6.0" @@ -615,6 +627,26 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "thiserror" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "611040a08a0439f8248d1990b111c95baa9c704c805fa1f62104b39655fd7f90" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "090198534930841fab3a5d1bb637cde49e339654e606195f8d9c76eeb081dc96" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "thread_local" version = "1.1.7" @@ -728,8 +760,6 @@ dependencies = [ "image", "libwayshot", "log", - "wayland-client", - "wayland-protocols", ] [[package]] diff --git a/libwayshot/Cargo.toml b/libwayshot/Cargo.toml index a49e9073..c447fb1a 100644 --- a/libwayshot/Cargo.toml +++ b/libwayshot/Cargo.toml @@ -9,10 +9,11 @@ version = "0.1.2" edition = "2021" [dependencies] -image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm", "qoi"] } +image = { version = "0.24", default-features = false } log = "0.4.19" memmap2 = "0.7.1" nix = "0.26.2" +thiserror = "1" wayland-client = "0.30.2" wayland-protocols = { version = "0.30.0", features=["client", "unstable"] } wayland-protocols-wlr = { version = "0.1.0", features = ["client"] } diff --git a/libwayshot/README.md b/libwayshot/README.md index 5936005d..a18983c7 100644 --- a/libwayshot/README.md +++ b/libwayshot/README.md @@ -8,3 +8,15 @@

+# `libwayshot` + +`libwayshot` is a convenient wrapper over the wlroots screenshot protocol that provides a simple API to take screenshots with. + +# Example usage + +```rust +use libwayshot::WayshotConnection; + +let wayshot_connection = WayshotConnection::new().unwrap(); +let image_buffer = wayshot_connection.screenshot_all().unwrap(); +``` diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs new file mode 100644 index 00000000..078a77e9 --- /dev/null +++ b/libwayshot/src/dispatch.rs @@ -0,0 +1,223 @@ +use std::{ + process::exit, + 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, + }, + Connection, Dispatch, QueueHandle, WEnum, + WEnum::Value, +}; +use wayland_protocols::xdg::xdg_output::zv1::client::{ + zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1, +}; +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}, + screencopy::FrameFormat, +}; + +pub struct OutputCaptureState { + pub outputs: Vec, +} + +impl Dispatch for OutputCaptureState { + fn event( + state: &mut Self, + wl_registry: &WlRegistry, + event: wl_registry::Event, + _: &(), + _: &Connection, + qh: &QueueHandle, + ) { + /* > The name event is sent after binding the output object. This event + * is only sent once per output object, and the name does not change + * over the lifetime of the wl_output global. */ + + if let wl_registry::Event::Global { + name, + interface, + version, + } = event + { + if interface == "wl_output" { + if version >= 4 { + let output = wl_registry.bind::(name, 4, qh, ()); + state.outputs.push(OutputInfo { + wl_output: output, + name: "".to_string(), + transform: wl_output::Transform::Normal, + dimensions: OutputPositioning { + x: 0, + y: 0, + width: 0, + height: 0, + }, + }); + } else { + log::error!("Ignoring a wl_output with version < 4."); + } + } + } + } +} + +impl Dispatch for OutputCaptureState { + fn event( + state: &mut Self, + wl_output: &WlOutput, + event: wl_output::Event, + _: &(), + _: &Connection, + _: &QueueHandle, + ) { + let output: &mut OutputInfo = state + .outputs + .iter_mut() + .find(|x| x.wl_output == *wl_output) + .unwrap(); + + match event { + wl_output::Event::Name { name } => { + output.name = name; + } + wl_output::Event::Geometry { + transform: WEnum::Value(transform), + .. + } => { + output.transform = transform; + } + _ => (), + } + } +} + +delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1); + +impl Dispatch for OutputCaptureState { + fn event( + state: &mut Self, + _: &ZxdgOutputV1, + event: zxdg_output_v1::Event, + index: &usize, + _: &Connection, + _: &QueueHandle, + ) { + let output_info = state.outputs.get_mut(*index).unwrap(); + + match event { + zxdg_output_v1::Event::LogicalPosition { x, y } => { + output_info.dimensions.x = x; + output_info.dimensions.y = y; + log::debug!("Logical position event fired!"); + } + zxdg_output_v1::Event::LogicalSize { width, height } => { + output_info.dimensions.width = width; + output_info.dimensions.height = height; + log::debug!("Logical size event fired!"); + } + _ => {} + }; + } +} + +/// State of the frame after attemting 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`. + Failed, + /// Compositor sent a Ready event on calling `frame.copy`. + Finished, +} + +pub struct CaptureFrameState { + pub formats: Vec, + pub state: Option, + pub buffer_done: AtomicBool, +} + +impl Dispatch for CaptureFrameState { + fn event( + frame: &mut Self, + _: &ZwlrScreencopyFrameV1, + event: zwlr_screencopy_frame_v1::Event, + _: &(), + _: &Connection, + _: &QueueHandle, + ) { + match event { + zwlr_screencopy_frame_v1::Event::Buffer { + format, + width, + height, + stride, + } => { + log::debug!("Received Buffer event"); + if let Value(f) = format { + frame.formats.push(FrameFormat { + format: f, + width, + height, + stride, + }) + } else { + log::debug!("Received Buffer event with unidentified format"); + exit(1); + } + } + zwlr_screencopy_frame_v1::Event::Flags { .. } => { + log::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. + log::debug!("Received Ready event"); + frame.state.replace(FrameState::Finished); + } + zwlr_screencopy_frame_v1::Event::Failed => { + log::debug!("Received Failed event"); + frame.state.replace(FrameState::Failed); + } + zwlr_screencopy_frame_v1::Event::Damage { .. } => { + log::debug!("Received Damage event"); + } + zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => { + log::debug!("Received LinuxDmaBuf event"); + } + zwlr_screencopy_frame_v1::Event::BufferDone => { + log::debug!("Received bufferdone event"); + frame.buffer_done.store(true, Ordering::SeqCst); + } + _ => unreachable!(), + }; + } +} + +delegate_noop!(CaptureFrameState: ignore WlShm); +delegate_noop!(CaptureFrameState: ignore WlShmPool); +delegate_noop!(CaptureFrameState: ignore WlBuffer); +delegate_noop!(CaptureFrameState: ignore ZwlrScreencopyManagerV1); + +// TODO: Create a xdg-shell surface, check for the enter event, grab the output from it. + +pub struct WayshotState {} + +impl wayland_client::Dispatch for WayshotState { + fn event( + _: &mut WayshotState, + _: &wl_registry::WlRegistry, + _: wl_registry::Event, + _: &GlobalListContents, + _: &Connection, + _: &QueueHandle, + ) { + } +} diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs new file mode 100644 index 00000000..883ca98e --- /dev/null +++ b/libwayshot/src/error.rs @@ -0,0 +1,24 @@ +use std::{io, result}; + +use thiserror::Error; +use wayland_client::{globals::GlobalError, ConnectError, DispatchError}; + +pub type Result = result::Result; + +#[derive(Error, Debug)] +pub enum Error { + #[error("no outputs supplied")] + NoOutputs, + #[error("image buffer is not big enough")] + BufferTooSmall, + #[error("image color type not supported")] + InvalidColor, + #[error("IO error: {0}")] + Io(#[from] io::Error), + #[error("dispatch error: {0}")] + Dispatch(#[from] DispatchError), + #[error("global error: {0}")] + Global(#[from] GlobalError), + #[error("connect error: {0}")] + Connect(#[from] ConnectError), +} diff --git a/wayshot/src/image_util.rs b/libwayshot/src/image_util.rs similarity index 100% rename from wayshot/src/image_util.rs rename to libwayshot/src/image_util.rs diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index e5c2b213..f51fb90d 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -1,84 +1,51 @@ +//! `libwayshot` is a convenient wrapper over the wlroots screenshot protocol +//! that provides a simple API to take screenshots with. +//! +//! To get started, look at [`WayshotConnection`]. + mod convert; +mod dispatch; +mod error; +mod image_util; +pub mod output; +mod screencopy; use std::{ - error::Error, - ffi::CStr, + cmp, fs::File, - io::Write, os::unix::prelude::FromRawFd, - os::unix::prelude::RawFd, process::exit, sync::atomic::{AtomicBool, Ordering}, - time::{SystemTime, UNIX_EPOCH}, -}; - -use nix::{ - fcntl, - sys::{memfd, mman, stat}, - unistd, -}; - -use image::{ - codecs::{ - jpeg::JpegEncoder, - png::PngEncoder, - pnm::{self, PnmEncoder}, - qoi::QoiEncoder, - }, - ColorType, ImageEncoder, }; +use image::{imageops::overlay, DynamicImage, RgbaImage}; use memmap2::MmapMut; - use wayland_client::{ - delegate_noop, - globals::GlobalList, + globals::{registry_queue_init, GlobalList}, protocol::{ - wl_buffer::WlBuffer, - wl_output::{self, WlOutput}, - wl_shm, - wl_shm::Format, - wl_shm::WlShm, - wl_shm_pool::WlShmPool, + wl_output::{Transform, WlOutput}, + wl_shm::{self, WlShm}, }, - Connection, Dispatch, QueueHandle, - WEnum::Value, + Connection, +}; +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, zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, + zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1, zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, }; -use crate::convert::create_converter; - -/// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and -/// Xbgr8888. -#[derive(Debug, Copy, Clone, PartialEq)] -pub struct FrameFormat { - pub format: Format, - pub width: u32, - pub height: u32, - pub stride: u32, -} +use crate::{ + convert::create_converter, + dispatch::{CaptureFrameState, FrameState, OutputCaptureState, WayshotState}, + output::OutputInfo, + screencopy::{create_shm_fd, FrameCopy}, +}; -/// State of the frame after attemting to copy it's data to a wl_buffer. -#[derive(Debug, Copy, Clone, PartialEq)] -enum FrameState { - /// Compositor returned a failed event on calling `frame.copy`. - Failed, - /// Compositor sent a Ready event on calling `frame.copy`. - Finished, -} +pub use crate::error::{Error, Result}; -/// The copied frame comprising of the FrameFormat, ColorType (Rgba8), and a memory backed shm -/// file that holds the image data in it. -#[derive(Debug)] -pub struct FrameCopy { - pub frame_format: FrameFormat, - pub frame_color_type: ColorType, - pub frame_mmap: MmapMut, - pub transform: wl_output::Transform, -} +type Frame = (Vec, Option<(i32, i32)>); /// Struct to store region capture details. #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -93,383 +60,374 @@ pub struct CaptureRegion { pub height: i32, } -/// Supported image encoding formats. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub enum EncodingFormat { - /// Jpeg / jpg encoder. - Jpg, - /// Png encoder. - Png, - /// Ppm encoder. - Ppm, - /// Qoi encoder. - Qoi, +struct IntersectingOutput { + output: WlOutput, + region: CaptureRegion, + transform: Transform, } -impl From for image::ImageOutputFormat { - fn from(format: EncodingFormat) -> Self { - match format { - EncodingFormat::Jpg => image::ImageFormat::Jpeg.into(), - EncodingFormat::Png => image::ImageFormat::Png.into(), - EncodingFormat::Ppm => image::ImageFormat::Pnm.into(), - EncodingFormat::Qoi => image::ImageFormat::Qoi.into(), - } - } +/// Struct to store wayland connection and globals list. +/// # Example usage +/// +/// ``` +/// let wayshot_connection = WayshotConnection::new().unwrap(); +/// let image_buffer = wayshot_connection.screenshot_all().unwrap(); +/// ``` +#[derive(Debug)] +pub struct WayshotConnection { + pub conn: Connection, + pub globals: GlobalList, } -impl From for &str { - fn from(format: EncodingFormat) -> Self { - match format { - EncodingFormat::Jpg => "jpg", - EncodingFormat::Png => "png", - EncodingFormat::Ppm => "ppm", - EncodingFormat::Qoi => "qoi", - } +impl WayshotConnection { + pub fn new() -> Result { + let conn = Connection::connect_to_env()?; + + Self::from_connection(conn) } -} -struct CaptureFrameState { - formats: Vec, - state: Option, - buffer_done: AtomicBool, -} + /// Recommended if you already have a [`wayland_client::Connection`]. + pub fn from_connection(conn: Connection) -> Result { + let (globals, _) = registry_queue_init::(&conn)?; -impl Dispatch for CaptureFrameState { - fn event( - frame: &mut Self, - _: &ZwlrScreencopyFrameV1, - event: zwlr_screencopy_frame_v1::Event, - _: &(), - _: &Connection, - _: &QueueHandle, - ) { - match event { - zwlr_screencopy_frame_v1::Event::Buffer { - format, - width, - height, - stride, - } => { - log::debug!("Received Buffer event"); - if let Value(f) = format { - frame.formats.push(FrameFormat { - format: f, - width, - height, - stride, - }) - } else { - log::debug!("Received Buffer event with unidentified format"); - exit(1); - } - } - zwlr_screencopy_frame_v1::Event::Flags { .. } => { - log::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. - log::debug!("Received Ready event"); - frame.state.replace(FrameState::Finished); - } - zwlr_screencopy_frame_v1::Event::Failed => { - log::debug!("Received Failed event"); - frame.state.replace(FrameState::Failed); - } - zwlr_screencopy_frame_v1::Event::Damage { .. } => { - log::debug!("Received Damage event"); - } - zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => { - log::debug!("Received LinuxDmaBuf event"); - } - zwlr_screencopy_frame_v1::Event::BufferDone => { - log::debug!("Received bufferdone event"); - frame.buffer_done.store(true, Ordering::SeqCst); + Ok(Self { conn, globals }) + } + + /// Fetch all accessible wayland outputs. + pub fn get_all_outputs(&self) -> Vec { + // Connecting to wayland environment. + let mut state = OutputCaptureState { + outputs: Vec::new(), + }; + let mut event_queue = self.conn.new_event_queue::(); + let qh = event_queue.handle(); + + // Bind to xdg_output global. + let zxdg_output_manager = match self.globals.bind::( + &qh, + 3..=3, + (), + ) { + Ok(x) => x, + Err(e) => { + log::error!("Failed to create ZxdgOutputManagerV1 version 3. Does your compositor implement ZxdgOutputManagerV1?"); + panic!("{:#?}", e); } - _ => unreachable!(), }; - } -} -delegate_noop!(CaptureFrameState: ignore WlShm); -delegate_noop!(CaptureFrameState: ignore WlShmPool); -delegate_noop!(CaptureFrameState: ignore WlBuffer); -delegate_noop!(CaptureFrameState: ignore ZwlrScreencopyManagerV1); + // Fetch all outputs; when their names arrive, add them to the list + let _ = self.conn.display().get_registry(&qh, ()); + event_queue.roundtrip(&mut state).unwrap(); + event_queue.roundtrip(&mut state).unwrap(); -/// Get a FrameCopy instance with screenshot pixel data for any wl_output object. -pub fn capture_output_frame( - globals: &mut GlobalList, - conn: &mut Connection, - cursor_overlay: i32, - output: WlOutput, - transform: wl_output::Transform, - capture_region: Option, -) -> Result> { - // Connecting to wayland environment. - let mut state = CaptureFrameState { - formats: Vec::new(), - state: None, - buffer_done: AtomicBool::new(false), - }; - let mut event_queue = conn.new_event_queue::(); - let qh = event_queue.handle(); - - // Instantiating screencopy manager. - let screencopy_manager = match globals.bind::(&qh, 3..=3, ()) { - Ok(x) => x, - Err(e) => { - log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"); - panic!("{:#?}", e); + let mut xdg_outputs: Vec = Vec::new(); + + // We loop over each output and request its position data. + for (index, output) in state.outputs.clone().iter().enumerate() { + let xdg_output = zxdg_output_manager.get_xdg_output(&output.wl_output, &qh, index); + xdg_outputs.push(xdg_output); + } + + event_queue.roundtrip(&mut state).unwrap(); + + for xdg_output in xdg_outputs { + xdg_output.destroy(); + } + + if state.outputs.is_empty() { + log::error!("Compositor did not advertise any wl_output devices!"); + exit(1); } - }; - - // Capture output. - let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { - screencopy_manager.capture_output_region( - cursor_overlay, - &output, - region.x_coordinate, - region.y_coordinate, - region.width, - region.height, + log::debug!("Outputs detected: {:#?}", state.outputs); + state.outputs + } + + /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. + fn capture_output_frame( + &self, + cursor_overlay: i32, + output: WlOutput, + transform: Transform, + capture_region: Option, + ) -> Result { + // Connecting to wayland environment. + let mut state = CaptureFrameState { + formats: Vec::new(), + state: None, + buffer_done: AtomicBool::new(false), + }; + let mut event_queue = self.conn.new_event_queue::(); + let qh = event_queue.handle(); + + // Instantiating screencopy manager. + let screencopy_manager = match self.globals.bind::( &qh, + 3..=3, (), - ) - } else { - screencopy_manager.capture_output(cursor_overlay, &output, &qh, ()) - }; - - // Empty internal event buffer until buffer_done is set to true which is when the Buffer done - // event is fired, aka the capture from the compositor is succesful. - while !state.buffer_done.load(Ordering::SeqCst) { - event_queue.blocking_dispatch(&mut state)?; - } + ) { + Ok(x) => x, + Err(e) => { + log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"); + panic!("{:#?}", e); + } + }; - log::debug!( - "Received compositor frame buffer formats: {:#?}", - state.formats - ); - // Filter advertised wl_shm formats and select the first one that matches. - let frame_format = state - .formats - .iter() - .find(|frame| { - matches!( - frame.format, - wl_shm::Format::Xbgr2101010 - | wl_shm::Format::Abgr2101010 - | wl_shm::Format::Argb8888 - | wl_shm::Format::Xrgb8888 - | wl_shm::Format::Xbgr8888 + // Capture output. + let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { + screencopy_manager.capture_output_region( + cursor_overlay, + &output, + region.x_coordinate, + region.y_coordinate, + region.width, + region.height, + &qh, + (), ) - }) - .copied(); - log::debug!("Selected frame buffer format: {:#?}", frame_format); - - // Check if frame format exists. - let frame_format = match frame_format { - Some(format) => format, - None => { - log::error!("No suitable frame format found"); - exit(1); + } else { + screencopy_manager.capture_output(cursor_overlay, &output, &qh, ()) + }; + + // Empty internal event buffer until buffer_done is set to true which is when the Buffer done + // event is fired, aka the capture from the compositor is succesful. + while !state.buffer_done.load(Ordering::SeqCst) { + event_queue.blocking_dispatch(&mut state)?; } - }; - - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - - // Create an in memory file and return it's file descriptor. - let mem_fd = create_shm_fd()?; - let mem_file = unsafe { File::from_raw_fd(mem_fd) }; - mem_file.set_len(frame_bytes as u64)?; - - // Instantiate shm global. - let shm = globals.bind::(&qh, 1..=1, ()).unwrap(); - let shm_pool = shm.create_pool(mem_fd, frame_bytes as i32, &qh, ()); - let buffer = shm_pool.create_buffer( - 0, - frame_format.width as i32, - frame_format.height as i32, - frame_format.stride as i32, - frame_format.format, - &qh, - (), - ); - - // Copy the pixel data advertised by the compositor into the buffer we just created. - frame.copy(&buffer); - - // On copy the Ready / Failed events are fired by the frame object, so here we check for them. - loop { - // Basically reads, if frame state is not None then... - if let Some(state) = state.state { - match state { - FrameState::Failed => { - log::error!("Frame copy failed"); - exit(1); - } - FrameState::Finished => { - // Create a writeable memory map backed by a mem_file. - let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? }; - let data = &mut *frame_mmap; - let frame_color_type = if let Some(converter) = - create_converter(frame_format.format) - { - converter.convert_inplace(data) - } else { - log::error!("Unsupported buffer format: {:?}", frame_format.format); - log::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."); + + log::debug!( + "Received compositor frame buffer formats: {:#?}", + state.formats + ); + // Filter advertised wl_shm formats and select the first one that matches. + let frame_format = state + .formats + .iter() + .find(|frame| { + matches!( + frame.format, + wl_shm::Format::Xbgr2101010 + | wl_shm::Format::Abgr2101010 + | wl_shm::Format::Argb8888 + | wl_shm::Format::Xrgb8888 + | wl_shm::Format::Xbgr8888 + ) + }) + .copied(); + log::debug!("Selected frame buffer format: {:#?}", frame_format); + + // Check if frame format exists. + let frame_format = match frame_format { + Some(format) => format, + None => { + log::error!("No suitable frame format found"); + exit(1); + } + }; + + // Bytes of data in the frame = stride * height. + let frame_bytes = frame_format.stride * frame_format.height; + + // Create an in memory file and return it's file descriptor. + let mem_fd = create_shm_fd()?; + let mem_file = unsafe { File::from_raw_fd(mem_fd) }; + mem_file.set_len(frame_bytes as u64)?; + + // Instantiate shm global. + let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); + let shm_pool = shm.create_pool(mem_fd, frame_bytes as i32, &qh, ()); + let buffer = shm_pool.create_buffer( + 0, + frame_format.width as i32, + frame_format.height as i32, + frame_format.stride as i32, + frame_format.format, + &qh, + (), + ); + + // Copy the pixel data advertised by the compositor into the buffer we just created. + frame.copy(&buffer); + + // On copy the Ready / Failed events are fired by the frame object, so here we check for them. + loop { + // Basically reads, if frame state is not None then... + if let Some(state) = state.state { + match state { + FrameState::Failed => { + log::error!("Frame copy failed"); exit(1); - }; - return Ok(FrameCopy { - frame_format, - frame_color_type, - frame_mmap, - transform, - }); + } + FrameState::Finished => { + // Create a writeable memory map backed by a mem_file. + let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? }; + let data = &mut *frame_mmap; + let frame_color_type = if let Some(converter) = + create_converter(frame_format.format) + { + converter.convert_inplace(data) + } else { + log::error!("Unsupported buffer format: {:?}", frame_format.format); + log::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."); + exit(1); + }; + return Ok(FrameCopy { + frame_format, + frame_color_type, + frame_mmap, + transform, + }); + } } } - } - event_queue.blocking_dispatch(&mut state)?; + event_queue.blocking_dispatch(&mut state)?; + } } -} -/// 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. -fn create_shm_fd() -> std::io::Result { - // Only try memfd on linux and freebsd. - #[cfg(any(target_os = "linux", target_os = "freebsd"))] - 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(), - memfd::MemFdCreateFlag::MFD_CLOEXEC | memfd::MemFdCreateFlag::MFD_ALLOW_SEALING, - ) { - Ok(fd) => { - // This is only an optimization, so ignore errors. - // F_SEAL_SRHINK = File cannot be reduced in size. - // F_SEAL_SEAL = Prevent further calls to fcntl(). - let _ = fcntl::fcntl( - fd, - fcntl::F_ADD_SEALS( - fcntl::SealFlag::F_SEAL_SHRINK | fcntl::SealFlag::F_SEAL_SEAL, - ), - ); - return Ok(fd); + fn create_frame_copy( + &self, + capture_region: CaptureRegion, + cursor_overlay: i32, + ) -> Result { + let mut framecopys: Vec = Vec::new(); + + let outputs = self.get_all_outputs(); + let mut intersecting_outputs: Vec = Vec::new(); + for output in outputs.iter() { + 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) { + 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, + }; + intersecting_outputs.push(IntersectingOutput { + output: output.wl_output.clone(), + region: true_region, + transform: output.transform, + }); } - Err(nix::errno::Errno::EINTR) => continue, - Err(nix::errno::Errno::ENOSYS) => break, - Err(errno) => return Err(std::io::Error::from(errno)), } + if intersecting_outputs.is_empty() { + log::error!("Provided capture region doesn't intersect with any outputs!"); + exit(1); + } + + for intersecting_output in intersecting_outputs { + framecopys.push(self.capture_output_frame( + cursor_overlay, + intersecting_output.output.clone(), + intersecting_output.transform, + Some(intersecting_output.region), + )?); + } + Ok(( + framecopys, + Some((capture_region.width, capture_region.height)), + )) } - // 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() - ); - loop { - match mman::shm_open( - // O_CREAT = Create file if does not exist. - // O_EXCL = Error if create and file exists. - // O_RDWR = Open for reading and writing. - // O_CLOEXEC = Close on succesful execution. - // S_IRUSR = Set user read permission bit . - // S_IWUSR = Set user write permission bit. - mem_file_handle.as_str(), - fcntl::OFlag::O_CREAT - | fcntl::OFlag::O_EXCL - | fcntl::OFlag::O_RDWR - | fcntl::OFlag::O_CLOEXEC, - stat::Mode::S_IRUSR | stat::Mode::S_IWUSR, - ) { - Ok(fd) => match mman::shm_unlink(mem_file_handle.as_str()) { - Ok(_) => return Ok(fd), - Err(errno) => match unistd::close(fd) { - Ok(_) => return Err(std::io::Error::from(errno)), - Err(errno) => return Err(std::io::Error::from(errno)), - }, - }, - 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() + /// Take a screenshot from the specified region. + pub fn screenshot( + &self, + capture_region: CaptureRegion, + cursor_overlay: bool, + ) -> Result { + let frame_copy = self.create_frame_copy(capture_region, cursor_overlay as i32)?; + + let mut composited_image; + + if frame_copy.0.len() == 1 { + let (width, height) = frame_copy.1.unwrap(); + let frame_copy = &frame_copy.0[0]; + + let image: DynamicImage = frame_copy.try_into()?; + composited_image = image_util::rotate_image_buffer( + &image, + frame_copy.transform, + width as u32, + height as u32, + ); + } else { + let mut images = Vec::new(); + let (frame_copy, region) = frame_copy; + let (width, height) = region.unwrap(); + for frame_copy in frame_copy { + let image: DynamicImage = (&frame_copy).try_into()?; + let image = image_util::rotate_image_buffer( + &image, + frame_copy.transform, + width as u32, + height as u32, ); - continue; + images.push(image); + } + composited_image = images[0].clone(); + for image in images { + overlay(&mut composited_image, &image, 0, 0); } - Err(nix::errno::Errno::EINTR) => continue, - Err(errno) => return Err(std::io::Error::from(errno)), } + + Ok(composited_image) } -} -/// Write an instance of FrameCopy to anything that implements Write trait. Eg: Stdout or a file -/// on the disk. -pub fn write_to_file( - mut output_file: &mut impl Write, - encoding_format: EncodingFormat, - frame_copy: &FrameCopy, -) -> Result<(), Box> { - log::debug!( - "Writing to disk with encoding format: {:#?}", - encoding_format - ); - match encoding_format { - EncodingFormat::Jpg => { - JpegEncoder::new(&mut output_file).write_image( - &frame_copy.frame_mmap, - frame_copy.frame_format.width, - frame_copy.frame_format.height, - frame_copy.frame_color_type, - )?; - output_file.flush()?; - } - EncodingFormat::Png => { - PngEncoder::new(&mut output_file).write_image( - &frame_copy.frame_mmap, - frame_copy.frame_format.width, - frame_copy.frame_format.height, - frame_copy.frame_color_type, - )?; - output_file.flush()?; - } - EncodingFormat::Qoi => { - QoiEncoder::new(&mut output_file).write_image( - &frame_copy.frame_mmap, - frame_copy.frame_format.width, - frame_copy.frame_format.height, - frame_copy.frame_color_type, - )?; - output_file.flush()?; - } - EncodingFormat::Ppm => { - let rgb8_data = if let ColorType::Rgba8 = frame_copy.frame_color_type { - let mut data = Vec::with_capacity( - (3 * frame_copy.frame_format.width * frame_copy.frame_format.height) as _, - ); - for chunk in frame_copy.frame_mmap.chunks_exact(4) { - data.extend_from_slice(&chunk[..3]); - } - data - } else { - unimplemented!("Currently only ColorType::Rgba8 is supported") - }; - - PnmEncoder::new(&mut output_file) - .with_subtype(pnm::PnmSubtype::Pixmap(pnm::SampleEncoding::Binary)) - .write_image( - &rgb8_data, - frame_copy.frame_format.width, - frame_copy.frame_format.height, - ColorType::Rgb8, - )?; - output_file.flush()?; + /// Take a screenshot from all of the specified outputs. + pub fn screenshot_outputs( + &self, + outputs: Vec, + cursor_overlay: bool, + ) -> Result { + if outputs.is_empty() { + 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) } - Ok(()) + /// Take a screenshot from all accessible outputs. + pub fn screenshot_all(&self, cursor_overlay: bool) -> Result { + self.screenshot_outputs(self.get_all_outputs(), cursor_overlay) + } } diff --git a/libwayshot/src/output.rs b/libwayshot/src/output.rs new file mode 100644 index 00000000..cf245750 --- /dev/null +++ b/libwayshot/src/output.rs @@ -0,0 +1,20 @@ +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)] +pub struct OutputInfo { + pub wl_output: WlOutput, + pub name: String, + pub transform: wl_output::Transform, + pub dimensions: OutputPositioning, +} + +#[derive(Default, Debug, Clone, PartialEq, Eq)] +pub struct OutputPositioning { + pub x: i32, + pub y: i32, + pub width: i32, + pub height: i32, +} diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs new file mode 100644 index 00000000..c00d4ef7 --- /dev/null +++ b/libwayshot/src/screencopy.rs @@ -0,0 +1,137 @@ +use std::{ + ffi::CStr, + os::unix::prelude::RawFd, + time::{SystemTime, UNIX_EPOCH}, +}; + +use image::{ColorType, DynamicImage, ImageBuffer, Pixel}; +use memmap2::MmapMut; +use nix::{ + fcntl, + sys::{memfd, mman, stat}, + unistd, +}; +use wayland_client::protocol::{wl_output, wl_shm::Format}; + +use crate::{Error, Result}; + +/// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and +/// Xbgr8888. +#[derive(Debug, Copy, Clone, PartialEq)] +pub struct FrameFormat { + pub format: Format, + pub width: u32, + pub height: u32, + pub stride: u32, +} + +fn create_image_buffer

( + frame_format: &FrameFormat, + frame_mmap: &MmapMut, +) -> Result>> +where + P: Pixel, +{ + ImageBuffer::from_vec(frame_format.width, frame_format.height, frame_mmap.to_vec()) + .ok_or(Error::BufferTooSmall) +} + +/// The copied frame comprising of the FrameFormat, ColorType (Rgba8), and a memory backed shm +/// file that holds the image data in it. +#[derive(Debug)] +pub struct FrameCopy { + pub frame_format: FrameFormat, + pub frame_color_type: ColorType, + pub frame_mmap: MmapMut, + pub transform: wl_output::Transform, +} + +impl TryFrom<&FrameCopy> for DynamicImage { + type Error = Error; + + fn try_from(value: &FrameCopy) -> Result { + Ok(match value.frame_color_type { + ColorType::Rgb8 => DynamicImage::ImageRgb8(create_image_buffer( + &value.frame_format, + &value.frame_mmap, + )?), + ColorType::Rgba8 => DynamicImage::ImageRgba8(create_image_buffer( + &value.frame_format, + &value.frame_mmap, + )?), + _ => return Err(Error::InvalidColor), + }) + } +} + +/// 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. +pub fn create_shm_fd() -> std::io::Result { + // Only try memfd on linux and freebsd. + #[cfg(any(target_os = "linux", target_os = "freebsd"))] + 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(), + memfd::MemFdCreateFlag::MFD_CLOEXEC | memfd::MemFdCreateFlag::MFD_ALLOW_SEALING, + ) { + Ok(fd) => { + // This is only an optimization, so ignore errors. + // F_SEAL_SRHINK = File cannot be reduced in size. + // F_SEAL_SEAL = Prevent further calls to fcntl(). + let _ = fcntl::fcntl( + fd, + fcntl::F_ADD_SEALS( + fcntl::SealFlag::F_SEAL_SHRINK | fcntl::SealFlag::F_SEAL_SEAL, + ), + ); + return Ok(fd); + } + Err(nix::errno::Errno::EINTR) => continue, + Err(nix::errno::Errno::ENOSYS) => break, + Err(errno) => return Err(std::io::Error::from(errno)), + } + } + + // 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() + ); + loop { + match mman::shm_open( + // O_CREAT = Create file if does not exist. + // O_EXCL = Error if create and file exists. + // O_RDWR = Open for reading and writing. + // O_CLOEXEC = Close on succesful execution. + // S_IRUSR = Set user read permission bit . + // S_IWUSR = Set user write permission bit. + mem_file_handle.as_str(), + fcntl::OFlag::O_CREAT + | fcntl::OFlag::O_EXCL + | fcntl::OFlag::O_RDWR + | fcntl::OFlag::O_CLOEXEC, + stat::Mode::S_IRUSR | stat::Mode::S_IWUSR, + ) { + Ok(fd) => match mman::shm_unlink(mem_file_handle.as_str()) { + Ok(_) => return Ok(fd), + Err(errno) => match unistd::close(fd) { + Ok(_) => return Err(std::io::Error::from(errno)), + Err(errno) => return Err(std::io::Error::from(errno)), + }, + }, + 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() + ); + continue; + } + Err(nix::errno::Errno::EINTR) => continue, + Err(errno) => return Err(std::io::Error::from(errno)), + } + } +} diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 13f5d8e4..2314d026 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -19,12 +19,9 @@ clap = "4.3.11" env_logger = { version = "0.10.0", default-features = false, features = ["auto-color", "color"] } log = "0.4.19" -wayland-client = "0.30.2" -wayland-protocols = { version = "0.30.0", features=["client", "unstable"] } - libwayshot = { version="0.1.2", path = "../libwayshot" } -image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm"] } +image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm", "qoi"] } dialoguer = { version = "0.10.4", features = ["fuzzy-select"] } diff --git a/wayshot/src/output.rs b/wayshot/src/output.rs deleted file mode 100644 index 7b629a75..00000000 --- a/wayshot/src/output.rs +++ /dev/null @@ -1,174 +0,0 @@ -use std::process::exit; -use wayland_client::{ - delegate_noop, - globals::GlobalList, - protocol::{wl_output, wl_output::WlOutput, wl_registry, wl_registry::WlRegistry}, - Connection, Dispatch, QueueHandle, WEnum, -}; -use wayland_protocols::xdg::xdg_output::zv1::client::{ - zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1, -}; - -#[derive(Debug, Clone)] -pub struct OutputInfo { - pub wl_output: WlOutput, - pub name: String, - pub transform: wl_output::Transform, - pub dimensions: OutputPositioning, -} - -#[derive(Default, Debug, Clone, PartialEq, Eq)] -pub struct OutputPositioning { - pub x: i32, - pub y: i32, - pub width: i32, - pub height: i32, -} - -struct OutputCaptureState { - outputs: Vec, -} - -impl Dispatch for OutputCaptureState { - fn event( - state: &mut Self, - wl_registry: &WlRegistry, - event: wl_registry::Event, - _: &(), - _: &Connection, - qh: &QueueHandle, - ) { - /* > The name event is sent after binding the output object. This event - * is only sent once per output object, and the name does not change - * over the lifetime of the wl_output global. */ - - if let wl_registry::Event::Global { - name, - interface, - version, - } = event - { - if interface == "wl_output" { - if version >= 4 { - let output = wl_registry.bind::(name, 4, qh, ()); - state.outputs.push(OutputInfo { - wl_output: output, - name: "".to_string(), - transform: wl_output::Transform::Normal, - dimensions: OutputPositioning { - x: 0, - y: 0, - width: 0, - height: 0, - }, - }); - } else { - log::error!("Ignoring a wl_output with version < 4."); - } - } - } - } -} - -impl Dispatch for OutputCaptureState { - fn event( - state: &mut Self, - wl_output: &WlOutput, - event: wl_output::Event, - _: &(), - _: &Connection, - _: &QueueHandle, - ) { - let output: &mut OutputInfo = state - .outputs - .iter_mut() - .find(|x| x.wl_output == *wl_output) - .unwrap(); - - match event { - wl_output::Event::Name { name } => { - output.name = name; - } - wl_output::Event::Geometry { - transform: WEnum::Value(transform), - .. - } => { - output.transform = transform; - } - _ => (), - } - } -} - -delegate_noop!(OutputCaptureState: ignore ZxdgOutputManagerV1); - -impl Dispatch for OutputCaptureState { - fn event( - state: &mut Self, - _: &ZxdgOutputV1, - event: zxdg_output_v1::Event, - index: &usize, - _: &Connection, - _: &QueueHandle, - ) { - let output_info = state.outputs.get_mut(*index).unwrap(); - - match event { - zxdg_output_v1::Event::LogicalPosition { x, y } => { - output_info.dimensions.x = x; - output_info.dimensions.y = y; - log::debug!("Logical position event fired!"); - } - zxdg_output_v1::Event::LogicalSize { width, height } => { - output_info.dimensions.width = width; - output_info.dimensions.height = height; - log::debug!("Logical size event fired!"); - } - _ => {} - }; - } -} - -pub fn get_all_outputs(globals: &mut GlobalList, conn: &mut Connection) -> Vec { - // Connecting to wayland environment. - let mut state = OutputCaptureState { - outputs: Vec::new(), - }; - let mut event_queue = conn.new_event_queue::(); - let qh = event_queue.handle(); - - // Bind to xdg_output global. - let zxdg_output_manager = match globals.bind::(&qh, 3..=3, ()) { - Ok(x) => x, - Err(e) => { - log::error!("Failed to create ZxdgOutputManagerV1 version 3. Does your compositor implement ZxdgOutputManagerV1?"); - panic!("{:#?}", e); - } - }; - - // Fetch all outputs; when their names arrive, add them to the list - let _ = conn.display().get_registry(&qh, ()); - event_queue.roundtrip(&mut state).unwrap(); - event_queue.roundtrip(&mut state).unwrap(); - - let mut xdg_outputs: Vec = Vec::new(); - - // We loop over each output and request its position data. - for (index, output) in state.outputs.clone().iter().enumerate() { - let xdg_output = zxdg_output_manager.get_xdg_output(&output.wl_output, &qh, index); - xdg_outputs.push(xdg_output); - } - - event_queue.roundtrip(&mut state).unwrap(); - - for xdg_output in xdg_outputs { - xdg_output.destroy(); - } - - if state.outputs.is_empty() { - log::error!("Compositor did not advertise any wl_output devices!"); - exit(1); - } - log::debug!("Outputs detected: {:#?}", state.outputs); - state.outputs -} diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index 3c5ef86a..55fe34f8 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -3,7 +3,9 @@ use std::{ time::{SystemTime, UNIX_EPOCH}, }; -pub fn parse_geometry(g: &str) -> Option { +use libwayshot::CaptureRegion; + +pub fn parse_geometry(g: &str) -> Option { let tail = g.trim(); let x_coordinate: i32; let y_coordinate: i32; @@ -30,7 +32,7 @@ pub fn parse_geometry(g: &str) -> Option { height = tail.parse::().ok()?; } - Some(libwayshot::CaptureRegion { + Some(CaptureRegion { x_coordinate, y_coordinate, width, @@ -38,7 +40,42 @@ pub fn parse_geometry(g: &str) -> Option { }) } -pub fn get_default_file_name(extension: libwayshot::EncodingFormat) -> String { +/// Supported image encoding formats. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum EncodingFormat { + /// Jpeg / jpg encoder. + Jpg, + /// Png encoder. + Png, + /// Ppm encoder. + Ppm, + /// Qoi encoder. + Qoi, +} + +impl From for image::ImageOutputFormat { + fn from(format: EncodingFormat) -> Self { + match format { + EncodingFormat::Jpg => image::ImageFormat::Jpeg.into(), + EncodingFormat::Png => image::ImageFormat::Png.into(), + EncodingFormat::Ppm => image::ImageFormat::Pnm.into(), + EncodingFormat::Qoi => image::ImageFormat::Qoi.into(), + } + } +} + +impl From for &str { + fn from(format: EncodingFormat) -> Self { + match format { + EncodingFormat::Jpg => "jpg", + EncodingFormat::Png => "png", + EncodingFormat::Ppm => "ppm", + EncodingFormat::Qoi => "qoi", + } + } +} + +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(_) => { diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index eb89a4ff..4a657e3b 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,49 +1,18 @@ use std::{ - cmp, env, + env, error::Error, io::{stdout, BufWriter, Cursor, Write}, process::exit, }; -use image::imageops::overlay; -use libwayshot::CaptureRegion; -use wayland_client::{ - globals::{registry_queue_init, GlobalListContents}, - protocol::{ - wl_output::{self, WlOutput}, - wl_registry, - }, - Connection, QueueHandle, -}; +use libwayshot::WayshotConnection; mod clap; -mod image_util; -mod output; mod utils; use dialoguer::{theme::ColorfulTheme, FuzzySelect}; -// TODO: Create a xdg-shell surface, check for the enter event, grab the output from it. - -struct WayshotState {} - -impl wayland_client::Dispatch for WayshotState { - fn event( - _: &mut WayshotState, - _: &wl_registry::WlRegistry, - _: wl_registry::Event, - _: &GlobalListContents, - _: &Connection, - _: &QueueHandle, - ) { - } -} - -struct IntersectingOutput { - output: WlOutput, - region: CaptureRegion, - transform: wl_output::Transform, -} +use crate::utils::EncodingFormat; fn select_ouput(ouputs: &[T]) -> Option where @@ -75,17 +44,17 @@ fn main() -> Result<(), Box> { log::debug!("Using custom extension: {:#?}", ext); match ext.as_str() { - "jpeg" | "jpg" => libwayshot::EncodingFormat::Jpg, - "png" => libwayshot::EncodingFormat::Png, - "ppm" => libwayshot::EncodingFormat::Ppm, - "qoi" => libwayshot::EncodingFormat::Qoi, + "jpeg" | "jpg" => EncodingFormat::Jpg, + "png" => EncodingFormat::Png, + "ppm" => EncodingFormat::Ppm, + "qoi" => EncodingFormat::Qoi, _ => { log::error!("Invalid extension provided.\nValid extensions:\n1) jpeg\n2) jpg\n3) png\n4) ppm\n5) qoi"); exit(1); } } } else { - libwayshot::EncodingFormat::Png + EncodingFormat::Png }; let mut file_is_stdout: bool = false; @@ -99,192 +68,65 @@ fn main() -> Result<(), Box> { file_path = Some(utils::get_default_file_name(extension)); } - let mut conn = Connection::connect_to_env().unwrap(); - let (mut globals, _) = registry_queue_init::(&conn).unwrap(); + let wayshot_conn = WayshotConnection::new()?; if args.get_flag("listoutputs") { - let valid_outputs = output::get_all_outputs(&mut globals, &mut conn); + let valid_outputs = wayshot_conn.get_all_outputs(); for output in valid_outputs { log::info!("{:#?}", output.name); } exit(1); } - let mut cursor_overlay: i32 = 0; + let mut cursor_overlay = false; if args.get_flag("cursor") { - cursor_overlay = 1; + cursor_overlay = true; } - let capture_area = if let Some(slurpregion) = args.get_one::("slurp") { - match utils::parse_geometry(slurpregion) { - Some(region) => (wl_output::Transform::Normal, region), - None => { - log::error!("Invalid geometry specification"); - exit(1); - } + 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 { + log::error!("Invalid geometry specification"); + exit(1); } } else if let Some(output_name) = args.get_one::("output") { - let outputs = output::get_all_outputs(&mut globals, &mut conn); - let mut capture_info = None; - for output in outputs { - if &output.name == output_name { - capture_info = Some(( - output.transform, - CaptureRegion { - x_coordinate: output.dimensions.x, - y_coordinate: output.dimensions.y, - width: output.dimensions.width, - height: output.dimensions.height, - }, - )) - } - } - - if capture_info.is_none() { + let outputs = wayshot_conn.get_all_outputs(); + if let Some(output) = outputs + .into_iter() + .find(|output| &output.name == output_name) + { + wayshot_conn.screenshot_outputs(vec![output], cursor_overlay)? + } else { log::error!("No output found!\n"); exit(1); } - - capture_info.unwrap() } else if args.get_flag("chooseoutput") { - let outputs = output::get_all_outputs(&mut globals, &mut conn); + 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) { - ( - outputs[index].transform, - CaptureRegion { - x_coordinate: outputs[index].dimensions.x, - y_coordinate: outputs[index].dimensions.y, - width: outputs[index].dimensions.width, - height: outputs[index].dimensions.height, - }, - ) + wayshot_conn.screenshot_outputs(vec![outputs[index].clone()], cursor_overlay)? } else { log::error!("No output found!\n"); exit(1); } } else { - let output = &output::get_all_outputs(&mut globals, &mut conn)[0]; - ( - output.transform, - CaptureRegion { - x_coordinate: output.dimensions.x, - y_coordinate: output.dimensions.y, - width: output.dimensions.width, - height: output.dimensions.height, - }, - ) - }; - - let frame_copy: (Vec, Option<(i32, i32)>) = { - let transform = capture_area.0; - let region = capture_area.1; - - let mut framecopys: Vec = Vec::new(); - - let outputs = output::get_all_outputs(&mut globals, &mut conn); - let mut intersecting_outputs: Vec = Vec::new(); - for output in outputs.iter() { - let x1: i32 = cmp::max(output.dimensions.x, region.x_coordinate); - let y1: i32 = cmp::max(output.dimensions.y, region.y_coordinate); - let x2: i32 = cmp::min( - output.dimensions.x + output.dimensions.width, - region.x_coordinate + region.width, - ); - let y2: i32 = cmp::min( - output.dimensions.y + output.dimensions.height, - region.y_coordinate + region.height, - ); - - let width = x2 - x1; - let height = y2 - y1; - - if !(width <= 0 || height <= 0) { - let true_x = region.x_coordinate - output.dimensions.x; - let true_y = region.y_coordinate - output.dimensions.y; - let true_region = CaptureRegion { - x_coordinate: true_x, - y_coordinate: true_y, - width: region.width, - height: region.height, - }; - intersecting_outputs.push(IntersectingOutput { - output: output.wl_output.clone(), - region: true_region, - transform, - }); - } - } - if intersecting_outputs.is_empty() { - log::error!("Provided capture region doesn't intersect with any outputs!"); - exit(1); - } - - for intersecting_output in intersecting_outputs { - framecopys.push(libwayshot::capture_output_frame( - &mut globals, - &mut conn, - cursor_overlay, - intersecting_output.output.clone(), - intersecting_output.transform, - Some(intersecting_output.region), - )?); - } - (framecopys, Some((region.width, region.height))) + wayshot_conn.screenshot_all(cursor_overlay)? }; - let mut composited_image; - let mut buffer; - - if frame_copy.0.len() == 1 { - let (width, height) = frame_copy.1.unwrap(); - let frame_copy = &frame_copy.0[0]; - - buffer = Cursor::new(Vec::new()); - libwayshot::write_to_file(&mut buffer, extension, frame_copy)?; - - let image = image::load_from_memory(buffer.get_ref())?; - composited_image = image_util::rotate_image_buffer( - &image, - frame_copy.transform, - width as u32, - height as u32, - ); - } else { - let mut images = Vec::new(); - let (frame_copy, region) = frame_copy; - let (width, height) = region.unwrap(); - for frame_copy in frame_copy { - buffer = Cursor::new(Vec::new()); - libwayshot::write_to_file(&mut buffer, extension, &frame_copy)?; - let image = image::load_from_memory(buffer.get_ref())?; - let image = image_util::rotate_image_buffer( - &image, - frame_copy.transform, - width as u32, - height as u32, - ); - images.push(image); - } - composited_image = images[0].clone(); - for image in images { - overlay(&mut composited_image, &image, 0, 0); - } - } - if file_is_stdout { let stdout = stdout(); let mut buffer = Cursor::new(Vec::new()); let mut writer = BufWriter::new(stdout.lock()); - composited_image.write_to(&mut buffer, extension)?; + image_buffer.write_to(&mut buffer, extension)?; writer.write_all(buffer.get_ref())?; } else { - composited_image.save(file_path.unwrap())?; + image_buffer.save(file_path.unwrap())?; } Ok(())