dd: refactor stdio handling

Author: xtqqczze

Cargo.lock

@@ -20,17 +20,18 @@ doctest = false
 
 [dependencies]
 clap = { workspace = true }
+fluent = { workspace = true }
 gcd = { workspace = true }
 libc = { workspace = true }
+rustix = { workspace = true, features = ["stdio"] }
+thiserror = { workspace = true }
 uucore = { workspace = true, features = [
   "format",
   "parser-size",
   "quoting-style",
   "fs",
   "signals",
 ] }
-thiserror = { workspace = true }
-fluent = { workspace = true }
 
 [target.'cfg(any(target_os = "linux", target_os = "android"))'.dependencies]
 nix = { workspace = true, features = ["fs", "signal"] }

src/uu/dd/Cargo.toml

@@ -12,8 +12,10 @@ mod datastructures;
 mod numbers;
 mod parseargs;
 mod progress;
+mod stdio;
 
 use crate::bufferedoutput::BufferedOutput;
+use crate::stdio::{StdinRaw, StdoutRaw};
 use blocks::conv_block_unblock_helper;
 use datastructures::{ConversionMode, IConvFlags, IFlags, OConvFlags, OFlags, options};
 #[cfg(any(target_os = "linux", target_os = "android"))]
@@ -25,7 +27,6 @@ use progress::ProgUpdateType;
 use progress::{ProgUpdate, ReadStat, StatusLevel, WriteStat, gen_prog_updater};
 #[cfg(target_os = "linux")]
 use progress::{check_and_reset_sigusr1, install_sigusr1_handler};
-use uucore::io::OwnedFileDescriptorOrHandle;
 use uucore::translate;
 
 use std::cmp;
@@ -37,15 +38,12 @@ use std::fs::{File, OpenOptions};
 use std::io::{self, Read, Seek, SeekFrom, Write};
 #[cfg(any(target_os = "linux", target_os = "android"))]
 use std::os::fd::AsFd;
+#[cfg(unix)]
+use std::os::unix::fs::FileTypeExt;
 #[cfg(any(target_os = "linux", target_os = "android"))]
 use std::os::unix::fs::OpenOptionsExt;
-#[cfg(unix)]
-use std::os::unix::{
-    fs::FileTypeExt,
-    io::{AsRawFd, FromRawFd},
-};
 #[cfg(windows)]
-use std::os::windows::{fs::MetadataExt, io::AsHandle};
+use std::os::windows::fs::MetadataExt;
 use std::path::Path;
 use std::sync::atomic::AtomicU8;
 use std::sync::{Arc, atomic::Ordering::Relaxed, mpsc};
@@ -178,20 +176,20 @@ impl Num {
     }
 }
 
-/// Read and discard `n` bytes from `reader` using a buffer of size `buf_size`.
+/// Read and discard `n` bytes from `file` using a buffer of size `buf_size`.
 ///
 /// This is more efficient than `io::copy` with `BufReader` because it reads
 /// directly in `buf_size`-sized chunks, matching GNU dd's behavior.
 /// Returns the total number of bytes actually read.
-fn read_and_discard<R: Read>(reader: &mut R, n: u64, buf_size: usize) -> io::Result<u64> {
+fn read_and_discard(file: &mut File, n: u64, buf_size: usize) -> io::Result<u64> {
     // todo: consider splice()ing to /dev/null on Linux
     let mut buf = Vec::with_capacity(buf_size);
     let mut total = 0u64;
     let mut remaining = n;
     while remaining > 0 {
         let to_read = cmp::min(remaining, buf_size as u64);
         buf.clear();
-        match reader.by_ref().take(to_read).read_to_end(&mut buf) {
+        match file.take(to_read).read_to_end(&mut buf) {
             Ok(0) => break, // EOF
             Ok(bytes_read) => {
                 total += bytes_read as u64;
@@ -207,44 +205,25 @@ fn read_and_discard<R: Read>(reader: &mut R, n: u64, buf_size: usize) -> io::Res
 
 /// Data sources.
 ///
-/// Use [`Source::stdin_as_file`] if available to enable more
+/// Use [`Source::StdinRaw`] if available to enable more
 /// fine-grained access to reading from stdin.
 enum Source {
-    /// Input from stdin.
-    #[cfg(not(unix))]
-    Stdin(io::Stdin),
-
     /// Input from a file.
     File(File),
 
     /// Input from stdin, opened from its file descriptor.
-    #[cfg(unix)]
-    StdinFile(File),
+    StdinRaw(StdinRaw),
 
     /// Input from a named pipe, also known as a FIFO.
     #[cfg(unix)]
     Fifo(File),
 }
 
 impl Source {
-    /// Create a source from stdin using its raw file descriptor.
-    ///
-    /// This returns an instance of the `Source::StdinFile` variant,
-    /// using the raw file descriptor of [`io::Stdin`] to create
-    /// the [`File`] parameter. You can use this instead of
-    /// `Source::Stdin` to allow reading from stdin without consuming
-    /// the entire contents of stdin when this process terminates.
-    #[cfg(unix)]
-    fn stdin_as_file() -> Self {
-        let fd = io::stdin().as_raw_fd();
-        let f = unsafe { File::from_raw_fd(fd) };
-        Self::StdinFile(f)
-    }
-
     fn skip(&mut self, n: u64, ibs: usize) -> io::Result<u64> {
         match self {
             #[cfg(not(unix))]
-            Self::Stdin(stdin) => {
+            Self::StdinRaw(stdin) => {
                 let m = read_and_discard(stdin, n, ibs)?;
                 if m < n {
                     show_error!(
@@ -255,8 +234,8 @@ impl Source {
                 Ok(m)
             }
             #[cfg(unix)]
-            Self::StdinFile(f) => {
-                if let Ok(Some(len)) = try_get_len_of_block_device(f)
+            Self::StdinRaw(stdin) => {
+                if let Ok(Some(len)) = try_get_len_of_block_device(stdin)
                     && len < n
                 {
                     // GNU compatibility:
@@ -269,9 +248,9 @@ impl Source {
                     return Ok(len);
                 }
                 // Get file length before seeking to avoid race condition
-                let file_len = f.metadata().as_ref().map_or(u64::MAX, Metadata::len);
+                let file_len = stdin.metadata().as_ref().map_or(u64::MAX, Metadata::len);
                 // Try seek first; fall back to read if not seekable
-                match n.try_into().ok().map(|n| f.seek(SeekFrom::Current(n))) {
+                match n.try_into().ok().map(|n| stdin.seek(SeekFrom::Current(n))) {
                     Some(Ok(pos)) => {
                         if pos > file_len {
                             show_error!(
@@ -284,7 +263,7 @@ impl Source {
                     // ESPIPE means the file descriptor is not seekable (e.g., a pipe),
                     // so fall back to reading and discarding bytes using ibs-sized buffer
                     Some(Err(e)) if e.raw_os_error() == Some(libc::ESPIPE) => {
-                        let m = read_and_discard(f, n, ibs)?;
+                        let m = read_and_discard(stdin, n, ibs)?;
                         if m < n {
                             show_error!(
                                 "{}",
@@ -331,11 +310,8 @@ impl Source {
 impl Read for Source {
     fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
         match self {
-            #[cfg(not(unix))]
-            Self::Stdin(stdin) => stdin.read(buf),
             Self::File(f) => f.read(buf),
-            #[cfg(unix)]
-            Self::StdinFile(f) => f.read(buf),
+            Self::StdinRaw(f) => f.read(buf),
             #[cfg(unix)]
             Self::Fifo(f) => f.read(buf),
         }
@@ -359,10 +335,11 @@ struct Input<'a> {
 impl<'a> Input<'a> {
     /// Instantiate this struct with stdin as a source.
     fn new_stdin(settings: &'a Settings) -> UResult<Self> {
+        let stdin = StdinRaw::new();
+
         #[cfg(not(unix))]
         let mut src = {
-            let f = File::from(io::stdin().as_handle().try_clone_to_owned()?);
-            let is_file = if let Ok(metadata) = f.metadata() {
+            let is_file = if let Ok(metadata) = stdin.metadata() {
                 // this hack is needed as there is no other way on windows
                 // to differentiate between the case where `seek` works
                 // on a file handle or not. i.e. when the handle is no real
@@ -373,15 +350,15 @@ impl<'a> Input<'a> {
                 false
             };
             if is_file {
-                Source::File(f)
+                Source::File(stdin.try_clone()?)
             } else {
-                Source::Stdin(io::stdin())
+                Source::StdinRaw(stdin)
             }
         };
         #[cfg(unix)]
-        let mut src = Source::stdin_as_file();
+        let mut src = Source::StdinRaw(stdin);
         #[cfg(unix)]
-        if let Source::StdinFile(f) = &src
+        if let Source::StdinRaw(f) = &src
             && settings.iflags.directory
             && !f.metadata()?.is_dir()
         {
@@ -599,7 +576,7 @@ enum Density {
 /// Data destinations.
 enum Dest {
     /// Output to stdout.
-    Stdout(File),
+    StdoutRaw(StdoutRaw),
 
     /// Output to a file.
     ///
@@ -619,7 +596,7 @@ enum Dest {
 impl Dest {
     fn fsync(&mut self) -> io::Result<()> {
         match self {
-            Self::Stdout(stdout) => stdout.flush(),
+            Self::StdoutRaw(stdout) => stdout.flush(),
             Self::File(f, _) => {
                 f.flush()?;
                 f.sync_all()
@@ -636,7 +613,7 @@ impl Dest {
 
     fn fdatasync(&mut self) -> io::Result<()> {
         match self {
-            Self::Stdout(stdout) => stdout.flush(),
+            Self::StdoutRaw(stdout) => stdout.flush(),
             Self::File(f, _) => {
                 f.flush()?;
                 f.sync_data()
@@ -654,7 +631,7 @@ impl Dest {
     #[cfg_attr(not(unix), allow(unused_variables))]
     fn seek(&mut self, n: u64, obs: usize) -> io::Result<u64> {
         match self {
-            Self::Stdout(stdout) => io::copy(&mut io::repeat(0).take(n), stdout),
+            Self::StdoutRaw(stdout) => io::copy(&mut io::repeat(0).take(n), &mut **stdout),
             Self::File(f, _) => {
                 #[cfg(unix)]
                 if let Ok(Some(len)) = try_get_len_of_block_device(f)
@@ -790,7 +767,7 @@ impl Write for Dest {
                     Err(e) => Err(e),
                 }
             }
-            Self::Stdout(stdout) => stdout.write(buf),
+            Self::StdoutRaw(stdout) => stdout.write(buf),
             #[cfg(unix)]
             Self::Fifo(f) => f.write(buf),
             #[cfg(unix)]
@@ -800,7 +777,7 @@ impl Write for Dest {
 
     fn flush(&mut self) -> io::Result<()> {
         match self {
-            Self::Stdout(stdout) => stdout.flush(),
+            Self::StdoutRaw(stdout) => stdout.flush(),
             Self::File(f, _) => f.flush(),
             #[cfg(unix)]
             Self::Fifo(f) => f.flush(),
@@ -827,8 +804,8 @@ struct Output<'a> {
 impl<'a> Output<'a> {
     /// Instantiate this struct with stdout as a destination.
     fn new_stdout(settings: &'a Settings) -> UResult<Self> {
-        let fx = OwnedFileDescriptorOrHandle::from(io::stdout())?;
-        let mut dst = Dest::Stdout(fx.into_file());
+        let stdout = StdoutRaw::new();
+        let mut dst = Dest::StdoutRaw(stdout);
         dst.seek(settings.seek, settings.obs)
             .map_err_context(|| translate!("dd-error-write-error"))?;
         Ok(Self { dst, settings })
@@ -888,16 +865,17 @@ impl<'a> Output<'a> {
     /// already opened by the system (stdout) and has a state
     /// (current position) that shall be used.
     fn new_file_from_stdout(settings: &'a Settings) -> UResult<Self> {
-        let fx = OwnedFileDescriptorOrHandle::from(io::stdout())?;
+        let stdout = StdoutRaw::new();
         #[cfg(any(target_os = "linux", target_os = "android"))]
         if let Some(libc_flags) = make_linux_oflags(&settings.oflags) {
             nix::fcntl::fcntl(
-                fx.as_raw().as_fd(),
+                stdout.as_fd(),
                 FcntlArg::F_SETFL(OFlag::from_bits_retain(libc_flags)),
             )?;
         }
 
-        Self::prepare_file(fx.into_file(), settings)
+        // TODO: avoid cloning the underlying file descriptor here
+        Self::prepare_file(stdout.try_clone()?, settings)
     }
 
     /// Instantiate this struct with the given named pipe as a destination.

src/uu/dd/src/dd.rs

@@ -0,0 +1,81 @@
+// This file is part of the uutils coreutils package.
+//
+// For the full copyright and license information, please view the LICENSE
+// file that was distributed with this source code.
+
+//! Abstractions for raw access to stdin and stdout, without buffering.
+
+use core::ops::{Deref, DerefMut};
+#[cfg(windows)]
+use std::io::{stdin, stdout};
+#[cfg(windows)]
+use std::os::windows::io::{AsRawHandle as _, FromRawHandle as _};
+use std::{fs::File, mem::ManuallyDrop};
+
+pub(crate) struct StdinRaw(ManuallyDrop<File>);
+
+impl StdinRaw {
+    #[cfg(not(windows))]
+    pub fn new() -> Self {
+        // SAFETY: We ensure that the file descriptor is never closed by
+        // wrapping the `File` in `ManuallyDrop`.
+        let fd = unsafe { rustix::stdio::take_stdin() };
+        let f = File::from(fd);
+        Self(ManuallyDrop::new(f))
+    }
+
+    #[cfg(windows)]
+    pub fn new() -> Self {
+        let handle = stdin().as_raw_handle();
+        let f = unsafe { File::from_raw_handle(handle) };
+        Self(ManuallyDrop::new(f))
+    }
+}
+
+impl Deref for StdinRaw {
+    type Target = File;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for StdinRaw {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+pub(crate) struct StdoutRaw(ManuallyDrop<File>);
+
+impl StdoutRaw {
+    #[cfg(not(windows))]
+    pub fn new() -> Self {
+        // SAFETY: We ensure that the file descriptor is never closed by
+        // wrapping the `File` in `ManuallyDrop`.
+        let fd = unsafe { rustix::stdio::take_stdout() };
+        let f = File::from(fd);
+        Self(ManuallyDrop::new(f))
+    }
+
+    #[cfg(windows)]
+    pub fn new() -> Self {
+        let handle = stdout().as_raw_handle();
+        let f = unsafe { File::from_raw_handle(handle) };
+        Self(ManuallyDrop::new(f))
+    }
+}
+
+impl Deref for StdoutRaw {
+    type Target = File;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for StdoutRaw {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}