pub mod error;
pub use error::Error;
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CompatibilityMode {
#[default]
Default,
Strict,
}
pub const BUFSIZ: usize = 64 * 1024;
#[non_exhaustive]
pub struct Pee {
sinks: Vec<Box<dyn std::io::Write + Send>>,
compat: CompatibilityMode,
capture: bool,
ignore_write_errors: bool,
}
impl std::fmt::Debug for Pee {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Pee")
.field("sinks_count", &self.sinks.len())
.field("compat", &self.compat)
.field("capture", &self.capture)
.field("ignore_write_errors", &self.ignore_write_errors)
.finish()
}
}
#[non_exhaustive]
pub struct PeeBuilder {
sinks: Vec<Box<dyn std::io::Write + Send>>,
compat: CompatibilityMode,
capture: bool,
ignore_write_errors: bool,
}
impl std::fmt::Debug for PeeBuilder {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PeeBuilder")
.field("sinks_count", &self.sinks.len())
.field("compat", &self.compat)
.field("capture", &self.capture)
.field("ignore_write_errors", &self.ignore_write_errors)
.finish()
}
}
impl Default for PeeBuilder {
fn default() -> Self {
Self::new()
}
}
impl PeeBuilder {
#[must_use]
pub fn new() -> Self {
Self {
sinks: Vec::new(),
compat: CompatibilityMode::Default,
capture: false,
ignore_write_errors: true,
}
}
#[must_use]
pub fn sink(mut self, sink: Box<dyn std::io::Write + Send>) -> Self {
self.sinks.push(sink);
self
}
#[must_use]
pub fn compat(mut self, compat: CompatibilityMode) -> Self {
self.compat = compat;
self
}
#[must_use]
pub fn capture(mut self, capture: bool) -> Self {
self.capture = capture;
self
}
#[must_use]
pub fn ignore_write_errors(mut self, ignore: bool) -> Self {
self.ignore_write_errors = ignore;
self
}
pub fn build(self) -> Result<Pee, Error> {
if self.compat == CompatibilityMode::Strict && self.capture {
return Err(Error::CompatibilityViolation(
"--capture not honored in Strict mode",
));
}
Ok(Pee {
sinks: self.sinks,
compat: self.compat,
capture: self.capture,
ignore_write_errors: self.ignore_write_errors,
})
}
}
impl Pee {
pub fn run<R: std::io::Read>(&mut self, mut reader: R) -> Result<(), Error> {
let _ = (self.compat, self.capture);
if self.sinks.is_empty() {
let mut buf = [0u8; BUFSIZ];
loop {
let n = reader.read(&mut buf)?;
if n == 0 {
break;
}
}
return Ok(());
}
let mut buf = vec![0u8; BUFSIZ];
let mut live: Vec<usize> = (0..self.sinks.len()).collect();
loop {
let n = reader.read(&mut buf)?;
if n == 0 {
break;
}
let mut to_drop: Vec<usize> = Vec::new();
for &idx in &live {
use std::io::Write;
match self.sinks[idx].write_all(&buf[..n]) {
Ok(()) => {}
Err(e)
if e.kind() == std::io::ErrorKind::BrokenPipe
&& self.ignore_write_errors =>
{
to_drop.push(idx);
}
Err(e) if self.ignore_write_errors => {
to_drop.push(idx);
let _ = e;
}
Err(e) => {
return Err(Error::SinkWriteFailed {
sink_index: idx,
source: e,
});
}
}
}
live.retain(|i| !to_drop.contains(i));
if live.is_empty() {
loop {
let n = reader.read(&mut buf)?;
if n == 0 {
break;
}
}
break;
}
}
for &idx in &live {
use std::io::Write;
let _ = self.sinks[idx].flush();
}
Ok(())
}
}
#[cfg(feature = "cli")]
pub mod aggregate;
#[cfg(feature = "cli")]
pub mod capture;
#[cfg(feature = "cli")]
pub mod cli;
#[cfg(feature = "cli")]
pub mod fanout;
#[cfg(feature = "cli")]
pub mod mode;
#[cfg(feature = "cli")]
pub mod spawner;
#[cfg(feature = "cli")]
pub mod strict;
#[cfg(feature = "cli")]
pub fn run() -> std::process::ExitCode {
use clap::Parser;
use std::ffi::OsString;
use std::process::ExitCode;
let raw_argv: Vec<OsString> = std::env::args_os().collect();
let pre_strict = strict::pre_scan_strict_flag(&raw_argv);
let env_strict = std::env::var_os("RUSTY_PEE_STRICT");
let argv0 = raw_argv.first().cloned();
let resolved_mode = mode::resolve(pre_strict, env_strict.as_deref(), argv0.as_deref());
if resolved_mode == CompatibilityMode::Strict {
return strict::run(&raw_argv);
}
let cli_args = match cli::Cli::try_parse() {
Ok(args) => args,
Err(e) => {
e.print().ok();
return match e.kind() {
clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion => {
ExitCode::SUCCESS
}
_ => ExitCode::from(2),
};
}
};
if let Some(cli::Subcommand::Completions { shell }) = cli_args.command {
use clap::CommandFactory;
let mut cmd = cli::Cli::command();
let name = cmd.get_name().to_string();
clap_complete::generate(shell, &mut cmd, name, &mut std::io::stdout());
return ExitCode::SUCCESS;
}
if cli_args.commands.is_empty() {
use std::io::Read;
let stdin = std::io::stdin();
let mut handle = stdin.lock();
let mut buf = [0u8; BUFSIZ];
loop {
match handle.read(&mut buf) {
Ok(0) => break,
Ok(_) => {}
Err(e) => {
eprintln!("rusty-pee: stdin read error: {e}");
return ExitCode::from(1);
}
}
}
return ExitCode::SUCCESS;
}
let statuses = if cli_args.capture {
let children = match capture::spawn_all_piped(&cli_args.commands) {
Ok(c) => c,
Err(e) => {
eprintln!("rusty-pee: failed to spawn child: {e}");
return ExitCode::from(127);
}
};
let stdin = std::io::stdin();
let stdout = std::io::stdout();
let mut out = stdout.lock();
match capture::run_with_capture(stdin.lock(), children, &mut out) {
Ok(s) => s,
Err(e) => {
eprintln!("rusty-pee: capture error: {e}");
return ExitCode::from(1);
}
}
} else {
let mut children = Vec::with_capacity(cli_args.commands.len());
for cmd in &cli_args.commands {
match spawner::spawn_one(cmd) {
Ok(c) => children.push(c),
Err(e) => {
eprintln!("rusty-pee: failed to spawn child '{cmd}': {e}");
for mut c in children.into_iter() {
let _ = c.kill();
let _ = c.wait();
}
return ExitCode::from(127);
}
}
}
let stdin = std::io::stdin();
match fanout::run(stdin.lock(), children) {
Ok(s) => s,
Err(e) => {
eprintln!("rusty-pee: fan-out error: {e}");
return ExitCode::from(1);
}
}
};
let codes: Vec<i32> = statuses.iter().map(|s| s.code().unwrap_or(1)).collect();
let aggregated = aggregate::default_max(&codes);
let byte = if (0..=255).contains(&aggregated) {
aggregated as u8
} else {
1u8
};
ExitCode::from(byte)
}