Skip to content

Commit

Permalink
Remove explicit panics and exits
Browse files Browse the repository at this point in the history
Let's bubble up errors instead of panicking and also not use `exit(...)` because it does not call destructors which might give issues.
  • Loading branch information
AndreasBackx committed Feb 2, 2024
1 parent 22147a9 commit 9e68016
Show file tree
Hide file tree
Showing 11 changed files with 89 additions and 67 deletions.
4 changes: 2 additions & 2 deletions libwayshot/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
```
23 changes: 15 additions & 8 deletions libwayshot/src/dispatch.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,5 @@
use std::{
collections::HashSet,
process::exit,
sync::atomic::{AtomicBool, Ordering},
};
use wayland_client::{
Expand Down Expand Up @@ -86,11 +85,13 @@ impl Dispatch<WlOutput, ()> for OutputCaptureState {
_: &Connection,
_: &QueueHandle<Self>,
) {
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 } => {
Expand Down Expand Up @@ -129,7 +130,14 @@ impl Dispatch<ZxdgOutputV1, usize> for OutputCaptureState {
_: &Connection,
_: &QueueHandle<Self>,
) {
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 } => {
Expand Down Expand Up @@ -189,7 +197,6 @@ impl Dispatch<ZwlrScreencopyFrameV1, ()> for CaptureFrameState {
})
} else {
tracing::debug!("Received Buffer event with unidentified format");
exit(1);
}
}
zwlr_screencopy_frame_v1::Event::Ready { .. } => {
Expand Down
7 changes: 6 additions & 1 deletion libwayshot/src/error.rs
Original file line number Diff line number Diff line change
@@ -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<T, E = Error> = result::Result<T, E>;

Expand All @@ -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}")]
Expand Down
9 changes: 4 additions & 5 deletions libwayshot/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@ use std::{
collections::HashSet,
fs::File,
os::fd::AsFd,
process::exit,
sync::atomic::{AtomicBool, Ordering},
thread,
};
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -280,7 +279,7 @@ impl WayshotConnection {
let frame_bytes = frame_format.stride * frame_format.height;

// Instantiate shm global.
let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ()).unwrap();
let shm = self.globals.bind::<WlShm, _, _>(&qh, 1..=1, ())?;
let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ());
let buffer = shm_pool.create_buffer(
0,
Expand Down
8 changes: 4 additions & 4 deletions libwayshot/src/region.rs
Original file line number Diff line number Diff line change
Expand Up @@ -174,22 +174,22 @@ impl TryFrom<&Vec<OutputInfo>> 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,
Expand Down
25 changes: 14 additions & 11 deletions libwayshot/src/screencopy.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use std::{
ffi::CStr,
ffi::CString,
os::fd::{AsRawFd, IntoRawFd, OwnedFd},
time::{SystemTime, UNIX_EPOCH},
};
Expand Down Expand Up @@ -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.
Expand All @@ -91,7 +101,7 @@ pub fn create_shm_fd() -> std::io::Result<OwnedFd> {
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) => {
Expand All @@ -113,11 +123,7 @@ pub fn create_shm_fd() -> std::io::Result<OwnedFd> {
}

// 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.
Expand All @@ -142,10 +148,7 @@ pub fn create_shm_fd() -> std::io::Result<OwnedFd> {
},
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,
Expand Down
2 changes: 2 additions & 0 deletions wayshot/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ repository.workspace = true

[build-dependencies]
flate2 = "1.0.27"
eyre = "0.6.8"


[dependencies]
tracing.workspace = true
Expand Down
49 changes: 30 additions & 19 deletions wayshot/build.rs
Original file line number Diff line number Diff line change
@@ -1,70 +1,81 @@
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())
.stderr(Stdio::null())
.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<Vec<(String, String)>> {
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;
}

let file = file_name.replace(search, replace);
files.push((file_name.to_string(), file));
}
}
files
Ok(files)
}
3 changes: 2 additions & 1 deletion wayshot/src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -15,7 +16,7 @@ pub struct Cli {
pub file: Option<PathBuf>,

/// 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
Expand Down
14 changes: 5 additions & 9 deletions wayshot/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ use eyre::{bail, ContextCompat, Error, Result};
use std::{
fmt::Display,
path::PathBuf,
process::exit,
str::FromStr,
time::{SystemTime, UNIX_EPOCH},
};
Expand Down Expand Up @@ -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()
}
12 changes: 5 additions & 7 deletions wayshot/src/wayshot.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -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();
Expand All @@ -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)?
Expand Down

0 comments on commit 9e68016

Please sign in to comment.