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 f4eedb06..5352deb4 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -1,6 +1,5 @@ use std::{ collections::HashSet, - process::exit, sync::atomic::{AtomicBool, Ordering}, }; use wayland_client::{ @@ -86,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 } => { @@ -129,7 +130,14 @@ 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 } => { @@ -189,7 +197,6 @@ impl Dispatch for CaptureFrameState { }) } else { tracing::debug!("Received Buffer event with unidentified format"); - exit(1); } } zwlr_screencopy_frame_v1::Event::Ready { .. } => { diff --git a/libwayshot/src/error.rs b/libwayshot/src/error.rs index ab0e1ec3..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}")] diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 764cbf7d..870fe233 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -15,7 +15,6 @@ use std::{ collections::HashSet, fs::File, os::fd::AsFd, - process::exit, sync::atomic::{AtomicBool, Ordering}, thread, }; @@ -68,8 +67,8 @@ pub mod reexport { /// # 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 { @@ -150,7 +149,7 @@ impl WayshotConnection { if state.outputs.is_empty() { tracing::error!("Compositor did not advertise any wl_output devices!"); - exit(1); + return Err(Error::NoOutputs); } tracing::trace!("Outputs detected: {:#?}", state.outputs); self.output_infos = state.outputs; @@ -280,7 +279,7 @@ impl WayshotConnection { let frame_bytes = frame_format.stride * frame_format.height; // Instantiate shm global. - let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); + let shm = self.globals.bind::(&qh, 1..=1, ())?; let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ()); let buffer = shm_pool.create_buffer( 0, diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs index 490e6082..27a35693 100644 --- a/libwayshot/src/region.rs +++ b/libwayshot/src/region.rs @@ -174,22 +174,22 @@ impl TryFrom<&Vec> for LogicalRegion { .iter() .map(|output| output.dimensions.x) .min() - .unwrap(); + .ok_or(Error::NoOutputs)?; let y1 = output_info .iter() .map(|output| output.dimensions.y) .min() - .unwrap(); + .ok_or(Error::NoOutputs)?; let x2 = output_info .iter() .map(|output| output.dimensions.x + output.dimensions.width) .max() - .unwrap(); + .ok_or(Error::NoOutputs)?; let y2 = output_info .iter() .map(|output| output.dimensions.y + output.dimensions.height) .max() - .unwrap(); + .ok_or(Error::NoOutputs)?; Ok(LogicalRegion { inner: Region { x: x1, diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index 2df15601..8d51b2ce 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}, }; @@ -82,6 +82,16 @@ impl TryFrom<&FrameCopy> 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. @@ -91,7 +101,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) => { @@ -113,11 +123,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. @@ -142,10 +148,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 344654f9..4594d0b4 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -12,6 +12,8 @@ repository.workspace = true [build-dependencies] flate2 = "1.0.27" +eyre = "0.6.8" + [dependencies] tracing.workspace = true 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/cli.rs b/wayshot/src/cli.rs index a076af5c..418ac8d8 100644 --- a/wayshot/src/cli.rs +++ b/wayshot/src/cli.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use clap::arg; use clap::Parser; +use eyre::WrapErr; use crate::utils::EncodingFormat; use clap::builder::TypedValueParser; @@ -15,7 +16,7 @@ pub struct Cli { 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().unwrap()}))] + #[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 diff --git a/wayshot/src/utils.rs b/wayshot/src/utils.rs index 1e394452..bb9a7a96 100644 --- a/wayshot/src/utils.rs +++ b/wayshot/src/utils.rs @@ -4,7 +4,6 @@ use eyre::{bail, ContextCompat, Error, Result}; use std::{ fmt::Display, path::PathBuf, - process::exit, str::FromStr, time::{SystemTime, UNIX_EPOCH}, }; @@ -134,13 +133,10 @@ impl FromStr for EncodingFormat { } pub fn get_default_file_name(extension: EncodingFormat) -> PathBuf { - 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); - } - }; + let time = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|time| time.as_secs().to_string()) + .unwrap_or("unknown".into()); - (time + "-wayshot." + extension.into()).into() + format!("{time}-wayshot.{extension}").into() } diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b5632caf..37152a79 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -1,10 +1,10 @@ use std::{ io::{stdout, BufWriter, Cursor, Write}, - process::{exit, Command}, + process::Command, }; use clap::Parser; -use eyre::Result; +use eyre::{bail, Result}; use libwayshot::{region::LogicalRegion, WayshotConnection}; mod cli; @@ -70,7 +70,7 @@ fn main() -> Result<()> { for output in valid_outputs { tracing::info!("{:#?}", output.name); } - exit(1); + return Ok(()); } let image_buffer = if let Some(slurp_region) = cli.slurp { @@ -94,8 +94,7 @@ fn main() -> Result<()> { 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 cli.choose_output { let outputs = wayshot_conn.get_all_outputs(); @@ -106,8 +105,7 @@ fn main() -> Result<()> { if let Some(index) = select_ouput(&output_names) { 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(cli.cursor)?