uucore: drop the pipe after a failed splice write

src/uucore/src/lib/features/pipes.rs

@@ -8,9 +8,9 @@
 use crate::io::{RawReader, RawWriter};
 use rustix::pipe::{SpliceFlags, fcntl_setpipe_size};
 use std::{
+    cell::Cell,
     io::{PipeReader, PipeWriter, Read, Write},
     os::fd::AsFd,
-    sync::OnceLock,
 };
 pub const MAX_ROOTLESS_PIPE_SIZE: usize = 1024 * 1024;
 const KERNEL_DEFAULT_PIPE_SIZE: usize = 64 * 1024;
@@ -32,6 +32,11 @@ impl From<PipeError> for std::io::Error {
 
 type PipeRes = std::io::Result<Result<(), ()>>;
 
+thread_local! {
+    /// Cache empty pipe pair to avoid calling `pipe2` at each copy.
+    static PIPE_CACHE: Cell<Option<(PipeReader, PipeWriter)>> = const { Cell::new(None) };
+}
+
 /// A type allows to
 /// - check that zero-copy succeed by ?.is_ok()
 /// - check that zero-copy failed, but read/write fallback succeed by ?.is_err()
@@ -71,6 +76,8 @@ pub fn splice(source: &impl AsFd, target: &impl AsFd, len: usize) -> rustix::io:
 }
 
 /// splice `len` bytes from `pipe` into `dest`.
+///
+/// On error, `pipe` may still hold bytes that were not written: don't reuse it.
 #[inline]
 pub fn drain_pipe(pipe: &PipeReader, dest: &impl AsFd, len: usize) -> PipeRes {
     debug_assert!(len <= MAX_ROOTLESS_PIPE_SIZE, "unexpected RAM usage");
@@ -98,8 +105,7 @@ pub fn drain_pipe(pipe: &PipeReader, dest: &impl AsFd, len: usize) -> PipeRes {
 /// This includes read ahead and optimization for stdout's pipe size
 #[inline]
 pub fn splice_unbounded_auto(source: &impl AsFd, dest: &mut impl AsFd) -> SpliceRes {
-    static PIPE_CACHE: OnceLock<Option<(PipeReader, PipeWriter)>> = OnceLock::new();
-    let Some((pipe_rd, pipe_wr)) = PIPE_CACHE.get_or_init(|| pipe::<false>().ok()) else {
+    let Ok((pipe_rd, pipe_wr)) = PIPE_CACHE.take().map_or_else(pipe::<false>, Ok) else {
         return Ok(Err(()));
     };
 
@@ -110,41 +116,45 @@ pub fn splice_unbounded_auto(source: &impl AsFd, dest: &mut impl AsFd) -> Splice
     // pre-generate page caches for splice
     let _ = rustix::fs::fadvise(source, 0, None, rustix::fs::Advice::Sequential);
     // 1st error is used to detect missing support for splice
-    match splice(&source, &pipe_wr, MAX_ROOTLESS_PIPE_SIZE) {
-        Ok(0) => return Ok(Ok(())),
+    let res = match splice(&source, &pipe_wr, MAX_ROOTLESS_PIPE_SIZE) {
+        Ok(0) => Ok(()),
         Ok(n) => {
-            if drain_pipe(pipe_rd, dest, n)
+            if drain_pipe(&pipe_rd, dest, n)
                 .map_err(PipeError::Write)?
                 .is_err()
             {
-                return Ok(Err(()));
+                Err(())
+            } else {
+                // GNU cat catches all strace injections for 2nd+ splice
+                while let mut n @ 1.. = splice(&source, &pipe_wr, MAX_ROOTLESS_PIPE_SIZE)
+                    .map_err(|error| PipeError::Read(error.into()))?
+                {
+                    while n > 0 {
+                        n -= splice(&pipe_rd, dest, n)
+                            .map_err(|error| PipeError::Write(error.into()))?;
+                    }
+                }
+                Ok(())
             }
         }
-        Err(_) => return Ok(Err(())),
-    }
-    // GNU cat catches all strace injections for 2nd+ splice
-    while let mut n @ 1.. = splice(&source, &pipe_wr, MAX_ROOTLESS_PIPE_SIZE)
-        .map_err(|error| PipeError::Read(error.into()))?
-    {
-        while n > 0 {
-            n -= splice(pipe_rd, dest, n).map_err(|error| PipeError::Write(error.into()))?;
-        }
-    }
-    Ok(Ok(()))
+        Err(_) => Err(()),
+    };
+    // pipe is empty. Then we can reuse it at next splice.
+    PIPE_CACHE.set(Some((pipe_rd, pipe_wr)));
+    Ok(res)
 }
 
 /// splice `n` bytes with read/write fallback
 /// return actually sent bytes
 #[inline]
 pub fn send_n_bytes(input: impl AsFd, target: impl AsFd, n: u64) -> std::io::Result<u64> {
-    static PIPE_CACHE: OnceLock<Option<(PipeReader, PipeWriter)>> = OnceLock::new();
     let pipe_size = n.min(MAX_ROOTLESS_PIPE_SIZE as u64) as usize;
     // improve throughput if output is pipe
     // expected that input is already extended if it is coming from splice
     if pipe_size > KERNEL_DEFAULT_PIPE_SIZE {
         let _ = fcntl_setpipe_size(&target, pipe_size);
     }
-    let Some((broker_r, broker_w)) = PIPE_CACHE.get_or_init(|| {
+    let Some((broker_r, broker_w)) = PIPE_CACHE.take().or_else(|| {
         // use std::io::pipe to avoid unnecessary fcntl
         let pair = std::io::pipe().ok()?;
         if pipe_size > KERNEL_DEFAULT_PIPE_SIZE {
@@ -158,17 +168,22 @@ pub fn send_n_bytes(input: impl AsFd, target: impl AsFd, n: u64) -> std::io::Res
     let mut bytes_written: u64 = 0;
     while n > 0 {
         match splice(&input, &broker_w, usize::try_from(n).unwrap_or(usize::MAX)) {
-            Ok(0) => return Ok(bytes_written),
+            Ok(0) => {
+                PIPE_CACHE.set(Some((broker_r, broker_w)));
+                return Ok(bytes_written);
+            }
             Ok(s) => {
                 n -= s as u64;
                 bytes_written += s as u64;
-                if drain_pipe(broker_r, &target, s)?.is_err() {
+                if drain_pipe(&broker_r, &target, s)?.is_err() {
                     break;
                 }
             }
             Err(_) => break,
         }
     }
+    // pipe is empty. Then we can reuse it at next splice.
+    PIPE_CACHE.set(Some((broker_r, broker_w)));
     // remove buffering from this fallback by RawReader, or order of output would be wrong with multiple input
     bytes_written += std::io::copy(&mut RawReader(input).take(n), &mut RawWriter(target))?;
     Ok(bytes_written)

tests/by-util/test_cp.rs

@@ -3,7 +3,7 @@
 // For the full copyright and license information, please view the LICENSE
 // file that was distributed with this source code.
 
-// spell-checker:ignore (flags) reflink (fs) tmpfs (linux) filefrag rlimit Rlim Nofile clob btrfs neve ROOTDIR USERDIR outfile subvolume uufs xattrs ELOOP
+// spell-checker:ignore (flags) reflink (fs) tmpfs (linux) filefrag rlimit Rlim Nofile Fsize SIGXFSZ clob btrfs neve ROOTDIR USERDIR outfile subvolume uufs xattrs ELOOP
 // spell-checker:ignore bdfl hlsl IRWXO IRWXG nconfined matchpathcon libselinux-devel prwx doesnotexist reftests subdirs mksocket srwx dstlink mcstransd
 
 #[cfg(unix)]
@@ -3472,6 +3472,32 @@ fn test_closes_file_descriptors() {
         .succeeds();
 }
 
+// When writing one file fails part-way, the rest of its data must not end up
+// in the files copied after it.
+#[cfg(any(target_os = "linux", target_os = "android"))]
+#[cfg(not(wasi_runner))] // linux specific
+#[test]
+fn test_cp_failed_write_does_not_spill_into_next_file() {
+    use rustix::process::Resource;
+
+    const CAP: u64 = 100 * 1024;
+
+    let (at, mut ucmd) = at_and_ucmd!();
+    at.write_bytes("big", &vec![b'x'; 3 * CAP as usize]);
+    at.write("small", "small file\n");
+    at.write("small2", "another small file\n");
+    at.mkdir("out");
+    ucmd.args(&["--reflink=never", "big", "small", "small2", "out"])
+        .limit(Resource::Fsize, CAP, CAP)
+        .ignore_sigxfsz()
+        .fails()
+        .stderr_contains("'big'");
+
+    assert_eq!(at.metadata("out/big").len(), CAP);
+    assert_eq!(at.read("out/small"), "small file\n");
+    assert_eq!(at.read("out/small2"), "another small file\n");
+}
+
 #[cfg(any(target_os = "linux", target_os = "android"))]
 #[test]
 #[cfg_attr(wasi_runner, ignore = "WASI: --sparse is only supported on linux")]