Merge pull request #7273 from RenjiSann/tee-fix-p-broken-stdout

tee: fix -p behavior upon broken pipe stdout

Cargo.lock

@@ -3320,7 +3320,7 @@ name = "uu_tee"
 version = "0.0.29"
 dependencies = [
  "clap",
- "libc",
+ "nix",
  "uucore",
 ]
 

src/uu/tee/Cargo.toml

@@ -18,7 +18,7 @@ path = "src/tee.rs"
 
 [dependencies]
 clap = { workspace = true }
-libc = { workspace = true }
+nix = { workspace = true, features = ["poll", "fs"] }
 uucore = { workspace = true, features = ["libc", "signals"] }
 
 [[bin]]

src/uu/tee/src/tee.rs

@@ -3,6 +3,8 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
+// cSpell:ignore POLLERR POLLRDBAND pfds revents
+
 use clap::{builder::PossibleValue, crate_version, Arg, ArgAction, Command};
 use std::fs::OpenOptions;
 use std::io::{copy, stdin, stdout, Error, ErrorKind, Read, Result, Write};
@@ -33,15 +35,20 @@ mod options {
 struct Options {
     append: bool,
     ignore_interrupts: bool,
+    ignore_pipe_errors: bool,
     files: Vec<String>,
     output_error: Option<OutputErrorMode>,
 }
 
 #[derive(Clone, Debug)]
 enum OutputErrorMode {
+    /// Diagnose write error on any output
     Warn,
+    /// Diagnose write error on any output that is not a pipe
     WarnNoPipe,
+    /// Exit upon write error on any output
     Exit,
+    /// Exit upon write error on any output that is not a pipe
     ExitNoPipe,
 }
 
@@ -49,32 +56,39 @@ enum OutputErrorMode {
 pub fn uumain(args: impl uucore::Args) -> UResult<()> {
     let matches = uu_app().try_get_matches_from(args)?;
 
-    let options = Options {
-        append: matches.get_flag(options::APPEND),
-        ignore_interrupts: matches.get_flag(options::IGNORE_INTERRUPTS),
-        files: matches
-            .get_many::<String>(options::FILE)
-            .map(|v| v.map(ToString::to_string).collect())
-            .unwrap_or_default(),
-        output_error: {
-            if matches.get_flag(options::IGNORE_PIPE_ERRORS) {
-                Some(OutputErrorMode::WarnNoPipe)
-            } else if matches.contains_id(options::OUTPUT_ERROR) {
-                if let Some(v) = matches.get_one::<String>(options::OUTPUT_ERROR) {
-                    match v.as_str() {
-                        "warn" => Some(OutputErrorMode::Warn),
-                        "warn-nopipe" => Some(OutputErrorMode::WarnNoPipe),
-                        "exit" => Some(OutputErrorMode::Exit),
-                        "exit-nopipe" => Some(OutputErrorMode::ExitNoPipe),
+    let append = matches.get_flag(options::APPEND);
+    let ignore_interrupts = matches.get_flag(options::IGNORE_INTERRUPTS);
+    let ignore_pipe_errors = matches.get_flag(options::IGNORE_PIPE_ERRORS);
+    let output_error = if matches.contains_id(options::OUTPUT_ERROR) {
+        match matches
+            .get_one::<String>(options::OUTPUT_ERROR)
+            .map(String::as_str)
+        {
+            Some("warn") => Some(OutputErrorMode::Warn),
+            // If no argument is specified for --output-error,
+            // defaults to warn-nopipe
+            None | Some("warn-nopipe") => Some(OutputErrorMode::WarnNoPipe),
+            Some("exit") => Some(OutputErrorMode::Exit),
+            Some("exit-nopipe") => Some(OutputErrorMode::ExitNoPipe),
             _ => unreachable!(),
         }
-                } else {
+    } else if ignore_pipe_errors {
         Some(OutputErrorMode::WarnNoPipe)
-                }
     } else {
         None
-            }
-        },
+    };
+
+    let files = matches
+        .get_many::<String>(options::FILE)
+        .map(|v| v.map(ToString::to_string).collect())
+        .unwrap_or_default();
+
+    let options = Options {
+        append,
+        ignore_interrupts,
+        ignore_pipe_errors,
+        files,
+        output_error,
     };
 
     match tee(&options) {
@@ -140,7 +154,6 @@ pub fn uu_app() -> Command {
                         .help("exit on write errors to any output that are not pipe errors (equivalent to exit on non-unix platforms)"),
                 ]))
                 .help("set write error behavior")
-                .conflicts_with(options::IGNORE_PIPE_ERRORS),
         )
 }
 
@@ -177,6 +190,11 @@ fn tee(options: &Options) -> Result<()> {
         inner: Box::new(stdin()) as Box<dyn Read>,
     };
 
+    #[cfg(target_os = "linux")]
+    if options.ignore_pipe_errors && !ensure_stdout_not_broken()? && output.writers.len() == 1 {
+        return Ok(());
+    }
+
     let res = match copy(input, &mut output) {
         // ErrorKind::Other is raised by MultiWriter when all writers
         // have exited, so that copy will abort. It's equivalent to
@@ -367,3 +385,44 @@ impl Read for NamedReader {
         }
     }
 }
+
+/// Check that if stdout is a pipe, it is not broken.
+#[cfg(target_os = "linux")]
+pub fn ensure_stdout_not_broken() -> Result<bool> {
+    use nix::{
+        poll::{PollFd, PollFlags, PollTimeout},
+        sys::stat::{fstat, SFlag},
+    };
+    use std::os::fd::{AsFd, AsRawFd};
+
+    let out = stdout();
+
+    // First, check that stdout is a fifo and return true if it's not the case
+    let stat = fstat(out.as_raw_fd())?;
+    if !SFlag::from_bits_truncate(stat.st_mode).contains(SFlag::S_IFIFO) {
+        return Ok(true);
+    }
+
+    // POLLRDBAND is the flag used by GNU tee.
+    let mut pfds = [PollFd::new(out.as_fd(), PollFlags::POLLRDBAND)];
+
+    // Then, ensure that the pipe is not broken
+    let res = nix::poll::poll(&mut pfds, PollTimeout::NONE)?;
+
+    if res > 0 {
+        // poll succeeded;
+        let error = pfds.iter().any(|pfd| {
+            if let Some(revents) = pfd.revents() {
+                revents.contains(PollFlags::POLLERR)
+            } else {
+                true
+            }
+        });
+        return Ok(!error);
+    }
+
+    // if res == 0, it means that timeout was reached, which is impossible
+    // because we set infinite timeout.
+    // And if res < 0, the nix wrapper should have sent back an error.
+    unreachable!();
+}

tests/by-util/test_tee.rs

@@ -165,6 +165,7 @@ mod linux_only {
     use std::fmt::Write;
     use std::fs::File;
     use std::process::{Output, Stdio};
+    use std::time::Duration;
 
     fn make_broken_pipe() -> File {
         use libc::c_int;
@@ -183,6 +184,22 @@ mod linux_only {
         unsafe { File::from_raw_fd(fds[1]) }
     }
 
+    fn make_hanging_read() -> File {
+        use libc::c_int;
+        use std::os::unix::io::FromRawFd;
+
+        let mut fds: [c_int; 2] = [0, 0];
+        assert!(
+            (unsafe { libc::pipe(std::ptr::from_mut::<c_int>(&mut fds[0])) } == 0),
+            "Failed to create pipe"
+        );
+
+        // PURPOSELY leak the write end of the pipe, so the read end hangs.
+
+        // Return the read end of the pipe
+        unsafe { File::from_raw_fd(fds[0]) }
+    }
+
     fn run_tee(proc: &mut UCommand) -> (String, Output) {
         let content = (1..=100_000).fold(String::new(), |mut output, x| {
             let _ = writeln!(output, "{x}");
@@ -535,4 +552,31 @@ mod linux_only {
         expect_failure(&output, "No space left");
         expect_short(file_out_a, &at, content.as_str());
     }
+
+    #[test]
+    fn test_pipe_mode_broken_pipe_only() {
+        new_ucmd!()
+            .timeout(Duration::from_secs(1))
+            .arg("-p")
+            .set_stdin(make_hanging_read())
+            .set_stdout(make_broken_pipe())
+            .succeeds();
+    }
+
+    #[test]
+    fn test_pipe_mode_broken_pipe_file() {
+        let (at, mut ucmd) = at_and_ucmd!();
+
+        let file_out_a = "tee_file_out_a";
+
+        let proc = ucmd
+            .arg("-p")
+            .arg(file_out_a)
+            .set_stdout(make_broken_pipe());
+
+        let (content, output) = run_tee(proc);
+
+        expect_success(&output);
+        expect_correct(file_out_a, &at, content.as_str());
+    }
 }